diff --git a/Project.toml b/Project.toml index 2a230fff..7227f599 100644 --- a/Project.toml +++ b/Project.toml @@ -6,20 +6,26 @@ version = "0.1.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" AtmosphericProfilesLibrary = "86bc3604-9858-485a-bdbe-831ec50de11d" +ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d" +ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c" CloudMicrophysics = "6a9e3e04-43cd-43ba-94b9-e8782df3c71b" JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" Oceananigans = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" +RRTMGP = "a01a1ee8-cea4-48fc-987c-fc7878d79da1" RootSolvers = "7181ea78-2dcb-4de3-ab41-2b8ab5a31e74" [compat] Adapt = "4.3.0" AtmosphericProfilesLibrary = "0.1.7" +ClimaComms = "0.6.9" +ClimaParams = "0.10.35" CloudMicrophysics = "0.22.13" JLD2 = "0.5.13" Oceananigans = "0.99" Printf = "1" +RRTMGP = "0.21.4" RootSolvers = "0.4.4" [extras] diff --git a/src/Breeze.jl b/src/Breeze.jl index 548a9180..db4cda9b 100644 --- a/src/Breeze.jl +++ b/src/Breeze.jl @@ -20,8 +20,11 @@ export using Oceananigans using Oceananigans.Grids: znode +using Oceananigans.Architectures: array_type, CPU, GPU +using Oceananigans: field export + array_type, CPU, GPU, Center, Face, Periodic, Bounded, Flat, RectilinearGrid, @@ -35,6 +38,7 @@ export FluxBoundaryCondition, ValueBoundaryCondition, GradientBoundaryCondition, OpenBoundaryCondition, PerturbationAdvection, FieldBoundaryConditions, Field, CenterField, XFaceField, YFaceField, ZFaceField, + field, Average, Integral, BackgroundField, interior, set!, compute!, regrid!, Forcing, @@ -49,6 +53,10 @@ export ∂x, ∂y, ∂z, @at, KernelFunctionOperation, prettytime +include("utils_grid.jl") +export + ncols + include("Thermodynamics/Thermodynamics.jl") using .Thermodynamics @@ -58,4 +66,14 @@ using .MoistAirBuoyancies include("AtmosphereModels/AtmosphereModels.jl") using .AtmosphereModels +include("Radiation/Radiation.jl") +using .Radiation + +export + AbstractRadiationModel, + RRTMGPModel, + initialize_rrtmgp_model, + compute_vertical_fluxes!, + flux_results + end # module Breeze diff --git a/src/Radiation/Radiation.jl b/src/Radiation/Radiation.jl new file mode 100644 index 00000000..24347a14 --- /dev/null +++ b/src/Radiation/Radiation.jl @@ -0,0 +1,11 @@ +module Radiation + +export + AbstractRadiationModel, + GrayRadiationModel + +include("rrtmgp_interface.jl") +include("radiation_model.jl") +include("radiation_model_gray.jl") + +end diff --git a/src/Radiation/radiation_model.jl b/src/Radiation/radiation_model.jl new file mode 100644 index 00000000..dc0b7316 --- /dev/null +++ b/src/Radiation/radiation_model.jl @@ -0,0 +1,12 @@ +abstract type AbstractRadiationModel end + +""" + update_radative_fluxes!(model) + +Update radiative fluxes for the given `GrayRadiationModel` by running the +longwave and shortwave two-stream solvers with the current atmospheric state +and boundary conditions. +""" +function update_radative_fluxes!(::AbstractRadiationModel) + throw("Not implemented") +end diff --git a/src/Radiation/radiation_model_gray.jl b/src/Radiation/radiation_model_gray.jl new file mode 100644 index 00000000..c41b0dfc --- /dev/null +++ b/src/Radiation/radiation_model_gray.jl @@ -0,0 +1,129 @@ +using RRTMGP.AtmosphericStates: GrayAtmosphericState +using RRTMGP.RTE: TwoStreamLWRTE, TwoStreamSWRTE +using RRTMGP.RTESolver: solve_lw!, solve_sw! + +using Breeze: field, ZFaceField, RectilinearGrid, set!, ncols, array_type +using Breeze.Radiation: AbstractRadiationModel, update_atmospheric_state! + +""" +GrayRadiationModel stores state and solver handles for a gray-band radiative +transfer setup. + +Field conventions (array shapes): +- cos_zenith_angle: (Nx, Ny) — cos of zenith angle on the horizontal grid +- sfc_emissivity: (Nbnd=1, Nx, Ny) — surface emissivity per band +- sfc_albedo_direct: (Nbnd=1, Nx, Ny) — direct-beam surface albedo per band +- sfc_albedo_diffuse: (Nbnd=1, Nx, Ny) — diffuse surface albedo per band +- toa_sw_flux_inc: (Nx, Ny) — incoming shortwave flux at TOA +""" +struct GrayRadiationModel{FT, AS, SLVLW, SLVSW} <: AbstractRadiationModel + fluxes_net_lw :: Field + fluxes_net_sw :: Field + cos_zenith_angle :: AbstractArray{FT} + sfc_emissivity :: AbstractArray{FT} + sfc_albedo_direct :: AbstractArray{FT} + sfc_albedo_diffuse :: AbstractArray{FT} + toa_sw_toa_flux :: AbstractArray{FT} + rrtmgp_atmospheric_state :: AS # GrayAtmosphericState + rrtmgp_solver_lw :: SLVLW + rrtmgp_solver_sw :: SLVSW +end + +""" + GrayRadiationModel(grid; + temperature, + pressure, + zenith_angle, + sfc_emissivity, + sfc_albedo_direct, + sfc_albedo_diffuse, + toa_sw_flux_inc) + +Construct a gray-band model using a precomputed atmospheric state. +Inputs may be scalars or arrays and are normalized to device arrays. +""" +function GrayRadiationModel( + grid; + zenith_angle, + sfc_emissivity, + sfc_albedo_direct, + sfc_albedo_diffuse, + toa_sw_flux_inc, + latitude, + isothermal_boundary_layer=false, +) + DA = array_type(grid.architecture) + FT = eltype(grid) + + # Create atmospheric state from provided Fields (generic library behavior) + # Allocated the required memory for the atmospheric state inside the + # RRTMGP.jl atmospheric state + rrtmgp_atmospheric_state = GrayAtmosphericState( + grid; + latitude=latitude, + ) + + # Build solvers using high-level wrappers + # They reshape internally for RRTMGP.jl compatibility to match the column layout + # used by RRTMGP.jl + SLVLW = TwoStreamLWRTE + SLVSW = TwoStreamSWRTE + lw_params = ( + sfc_emission = DA(sfc_emissivity), + lw_inc_flux = nothing, + isothermal_boundary_layer = isothermal_boundary_layer, + ) + sw_params = ( + cos_zenith = DA(cos.(FT(π / 180) .* zenith_angle)), + toa_flux = DA(toa_sw_flux_inc), + sfc_alb_direct = DA(sfc_albedo_direct), + inc_flux_diffuse = nothing, + sfc_alb_diffuse = DA(sfc_albedo_diffuse), + isothermal_boundary_layer = isothermal_boundary_layer, + ) + rrtmgp_solver_lw = SLVLW(grid; lw_params...) + rrtmgp_solver_sw = SLVSW(grid; sw_params...) + + # Create Fields for radiative fluxes to store the radiative fluxes from + # the RRTMGP.jl solvers for native compatibility with Oceananigans operators + fluxes_net_lw = ZFaceField(grid) + fluxes_net_sw = ZFaceField(grid) + + return GrayRadiationModel( + fluxes_net_lw, + fluxes_net_sw, + sw_params.cos_zenith, + lw_params.sfc_emission, + sw_params.sfc_alb_direct, + sw_params.sfc_alb_diffuse, + sw_params.toa_flux, + rrtmgp_atmospheric_state, + rrtmgp_solver_lw, + rrtmgp_solver_sw, + ) +end + +""" + (model::GrayRadiationModel)(temperature::Field, pressure::Field) + +Update the radiative fluxes for the given `GrayRadiationModel` by running the +longwave and shortwave two-stream solvers with the current atmospheric state +and boundary conditions. +""" +function (model::GrayRadiationModel)(temperature::Field, pressure::Field) + # Update the atmospheric state inside the RRTMGP.jl atmospheric state + # This is needed before running the solvers to make sure the atmospheric state is updated + # before computing the radiative fluxes + update_atmospheric_state!(model.rrtmgp_atmospheric_state, temperature, pressure) + + # Compute the radiative fluxes inside the RRTMGP.jl solvers + solve_lw!(model.rrtmgp_solver_lw, model.rrtmgp_atmospheric_state) + solve_sw!(model.rrtmgp_solver_sw, model.rrtmgp_atmospheric_state) + + # Update the radiative flux Field objects with the radiative fluxes from the RRTMGP.jl solvers + Nx, Ny, Nz = size(model.fluxes_net_lw.grid) + set!(model.fluxes_net_lw, permutedims(reshape(model.rrtmgp_solver_lw.flux.flux_net, Nz+1, Nx, Ny), (2, 3, 1))) + set!(model.fluxes_net_sw, permutedims(reshape(model.rrtmgp_solver_sw.flux.flux_net, Nz+1, Nx, Ny), (2, 3, 1))) + + return model.fluxes_net_lw, model.fluxes_net_sw +end diff --git a/src/Radiation/rrtmgp_interface.jl b/src/Radiation/rrtmgp_interface.jl new file mode 100644 index 00000000..f81ac6a8 --- /dev/null +++ b/src/Radiation/rrtmgp_interface.jl @@ -0,0 +1,188 @@ +using RRTMGP +using RRTMGP: RRTMGPGridParams +using RRTMGP.AngularDiscretizations +using RRTMGP.Fluxes +using RRTMGP.Optics +using RRTMGP.GrayUtils: setup_gray_as_pr_grid +using RRTMGP.Sources +using RRTMGP.BCs +using RRTMGP.Parameters: RRTMGPParameters +using RRTMGP.Optics: GrayOpticalThicknessOGorman2008 + +using Oceananigans +using Oceananigans: fill_halo_regions! +using Oceananigans: RectilinearGrid +using Oceananigans: field +using ClimaParams +using ClimaComms + +using Breeze: ncols, array_type, CPU, GPU + +# Import RRTMGP types for extensions +import RRTMGP.Parameters.RRTMGPParameters +import RRTMGP.AtmosphericStates +import RRTMGP.RTE: TwoStreamLWRTE, TwoStreamSWRTE + +""" + RRTMGPGridParams(grid; isothermal_boundary_layer=false) + +Construct RRTMGP grid parameters and a compute context compatible with the +grid's architecture. `Ncols = Nx * Ny` follows RRTMGP's column-major layout. +""" +function RRTMGP.RRTMGPGridParams(grid::RectilinearGrid; isothermal_boundary_layer::Bool = false) + FT = eltype(grid) + Nz = size(grid, 3) + context = context_from_arch(grid.architecture) + return RRTMGPGridParams{FT, typeof(context)}(context, Nz, ncols(grid), isothermal_boundary_layer) +end + +""" + AtmosphericStates.GrayAtmosphericState(grid; temperature, pressure, otp, + lat_center=0, planet_radius=6_371_000) + +Build a gray-atmosphere state from Oceananigans `Field`s, reshaping data into +the column-major layout used by RRTMGP. Shapes: +- `lat`: (Ncols) +- `p_lay`, `t_lay`: (Nz, Ncols) +- `p_lev`, `t_lev`, `z_lev`: (Nz+1, Ncols) +""" +function RRTMGP.AtmosphericStates.GrayAtmosphericState( + grid :: RectilinearGrid; + latitude :: AbstractArray, +) + DA = array_type(grid.architecture) + FT = eltype(grid) + Nz = size(grid, 3) + Ncols = ncols(grid) + + # Default optical thickness parameters from O'Gorman 2008 + otp = GrayOpticalThicknessOGorman2008(FT) + + # Default: reshape the provided fields into RRTMGP's column layout + lat = reshape(latitude, Ncols) + p_lay = DA{FT}(undef, Nz, Ncols) + p_lev = DA{FT}(undef, Nz + 1, Ncols) + t_lay = DA{FT}(undef, Nz, Ncols) + t_lev = DA{FT}(undef, Nz + 1, Ncols) + z_lev = DA{FT}(undef, Nz + 1, Ncols) + t_sfc = DA{FT}(undef, Ncols) + + return GrayAtmosphericState{FT, AbstractVector{FT}, AbstractMatrix{FT}, typeof(otp)}( + lat, + p_lay, + p_lev, + t_lay, + t_lev, + z_lev, + t_sfc, + otp + ) +end + +""" + RTE.TwoStreamLWRTE(grid; sfc_emission, lw_inc_flux) + +Construct the two-stream longwave RTE. Arrays are reshaped to RRTMGP's column +layout internally: +- `sfc_emission`: (Nbnd, Nx, Ny) → (Nbnd, Ncols) +- `lw_inc_flux`: nothing or (Nbnd, Nx, Ny) → (Nbnd, Ncols) +""" +function RRTMGP.RTE.TwoStreamLWRTE( + grid::RectilinearGrid; + sfc_emission, + lw_inc_flux, + isothermal_boundary_layer +) + Ncols = ncols(grid) + Nbnd = size(sfc_emission, 1) + grid_params = RRTMGPGridParams(grid; isothermal_boundary_layer) + params = RRTMGPParameters(eltype(grid)) + # TODO!: need helpers here for reshaping fields to column layout + sfc_emis = reshape(sfc_emission, Nbnd, Ncols) + lw_inc_flux = lw_inc_flux isa Nothing ? nothing : reshape(lw_inc_flux, Nbnd, Ncols) + + return TwoStreamLWRTE(grid_params; params, sfc_emis, inc_flux=lw_inc_flux) +end + +""" + RTE.TwoStreamSWRTE(grid; + cos_zenith, + toa_flux, + sfc_alb_direct, + inc_flux_diffuse, + sfc_alb_diffuse, + inc_flux_diffuse, + isothermal_boundary_layer) + +Construct the two-stream shortwave RTE. Arrays are reshaped to column layout: +- `cos_zenith`, `toa_flux`: (Nx, Ny) → (Ncols) +- `sfc_alb_direct`, `sfc_alb_diffuse`: (Nbnd, Nx, Ny) → (Nbnd, Ncols) +- `inc_flux_diffuse`: nothing or (Nbnd, Nx, Ny) → (Nbnd, Ncols) +- `isothermal_boundary_layer`: Bool +""" +function RRTMGP.RTE.TwoStreamSWRTE( + grid::RectilinearGrid; + cos_zenith, + toa_flux, + sfc_alb_direct, + sfc_alb_diffuse, + inc_flux_diffuse, + isothermal_boundary_layer +) + Ncols = ncols(grid) + Nbnd = size(sfc_alb_direct, 1) + grid_params = RRTMGPGridParams(grid; isothermal_boundary_layer) + + # TODO!: need helpers here for reshaping fields to column layout + cosz = reshape(cos_zenith, Ncols) + toa = reshape(toa_flux, Ncols) + alb_dir = reshape(sfc_alb_direct, Nbnd, Ncols) + alb_diff = reshape(sfc_alb_diffuse, Nbnd, Ncols) + inc_flux_diffuse = inc_flux_diffuse isa Nothing ? nothing : reshape(inc_flux_diffuse, Nbnd, Ncols) + + return TwoStreamSWRTE( + grid_params; + cos_zenith=cosz, + toa_flux=toa, + sfc_alb_direct=alb_dir, + inc_flux_diffuse=inc_flux_diffuse, + sfc_alb_diffuse=alb_diff, + ) +end + +""" + update_atmospheric_state!(as::GrayAtmosphericState, temperature::Field, pressure::Field) + +Update the atmospheric state with new temperature and pressure fields. +""" +function update_atmospheric_state!(as::RRTMGP.AtmosphericStates.GrayAtmosphericState, temperature::Field, pressure::Field) + grid = temperature.grid + Nz = size(grid, 3) + Ncols = ncols(grid) + + # TODO!: need helpers here for reshaping fields to column layout + as.p_lay .= reshape(PermutedDimsArray(interior(pressure), (3, 1, 2)), Nz, Ncols) + as.p_lev .= reshape(PermutedDimsArray(interior(Field(@at (Center, Center, Face) pressure)), (3, 1, 2)), Nz + 1, Ncols) + as.t_lay .= reshape(PermutedDimsArray(interior(temperature), (3, 1, 2)), Nz, Ncols) + as.t_lev .= reshape(PermutedDimsArray(interior(Field(@at (Center, Center, Face) temperature)), (3, 1, 2)), Nz + 1, Ncols) + as.t_sfc .= view(as.t_lev, 1, :) + + return nothing +end + +""" + context_from_arch(arch) + +Return a `ClimaComms` execution context compatible with the provided +Oceananigans architecture (CPU single/multithreaded or GPU). +""" +function context_from_arch(arch) + if arch isa CPU + dev = Threads.nthreads() > 1 ? + ClimaComms.CPUMultiThreaded() : + ClimaComms.CPUSingleThreaded() + return ClimaComms.context(dev) + elseif arch isa GPU + return ClimaComms.context(ClimaComms.CUDADevice()) + end +end diff --git a/src/utils_grid.jl b/src/utils_grid.jl new file mode 100644 index 00000000..1f06432e --- /dev/null +++ b/src/utils_grid.jl @@ -0,0 +1,10 @@ +using Oceananigans: RectilinearGrid + +""" + ncols(grid::RectilinearGrid) + +Return the number of columns in the grid. +""" +ncols(grid::RectilinearGrid) = grid.Nx * grid.Ny + + diff --git a/test/test_breeze_interface.jl b/test/test_breeze_interface.jl new file mode 100644 index 00000000..78790cbe --- /dev/null +++ b/test/test_breeze_interface.jl @@ -0,0 +1,65 @@ +using Oceananigans +using Oceananigans: field +using Oceananigans.Architectures: array_type +using Oceananigans: RectilinearGrid + +using Breeze.Radiation: GrayRadiationModel + +include("utils_test.jl") + +# Parameters +FT = Float64 +architecture = CPU() +DA = array_type(architecture) +Nx, Ny, Nz = 1, 1, 60 +Lx, Ly, Lz = 1, 20_000_000, 1_000 # domain size in meters + +# Set up grid +grid = RectilinearGrid( + architecture, + FT, + size=(Nx, Ny, Nz,), + x=(-Lx/2, Lx/2), + y=(-Ly/2, Ly/2), + z=(0, Lz), + topology=(Periodic, Periodic, Bounded) +) + +# Radiation parameters for Gray Radiation Model +Nbnd = 1 # This is gray band +zenith_angle = DA{FT}(undef, Nx, Ny) +zenith_angle .= FT(52.95) # zenith angle in degrees +sw_inc_flux = DA{FT}(undef, Nx, Ny) +sw_inc_flux .= FT(1407.679) # incoming shortwave flux +albedo_direct = DA{FT}(undef, Nbnd, Nx, Ny) +albedo_direct .= FT(0.1) # surface albedo (direct) +albedo_diffuse = DA{FT}(undef, Nbnd, Nx, Ny) +albedo_diffuse .= FT(0.1) # surface albedo (diffuse) +sfc_emissivity = DA{FT}(undef, Nbnd, Nx, Ny) +sfc_emissivity .= FT(1) # surface emissivity +latitude = latitude_from_grid(grid; lat_center=0) + +# construct the radiation model +radiation_model = GrayRadiationModel( + grid; + zenith_angle=zenith_angle, + sfc_emissivity=sfc_emissivity, + sfc_albedo_direct=albedo_direct, + sfc_albedo_diffuse=albedo_diffuse, + toa_sw_flux_inc=sw_inc_flux, + latitude=latitude, + isothermal_boundary_layer=false +) + +# Set up test case temperature and pressure profiles for the atmosphere +if grid.Nx > 1 + throw("Grid must not be wider than 1 point in the x direction") +else + # assumes only 1 point in the x direction + p_surface = 100_000 # surface pressure (Pa) + p_top = 9_000 # top of atmosphere pressure / emission level (Pa) + pressure, temperature = gray_test_t_p_profiles(grid; p0=p_surface, pe=p_top) +end + +# compute the radiative fluxes +fluxes_net_lw, fluxes_net_sw = radiation_model(temperature, pressure) diff --git a/test/test_rrtmgp_gray.jl b/test/test_rrtmgp_gray.jl new file mode 100644 index 00000000..92003857 --- /dev/null +++ b/test/test_rrtmgp_gray.jl @@ -0,0 +1,184 @@ +import ClimaComms + +using RRTMGP +using RRTMGP: RRTMGPGridParams +using RRTMGP.AngularDiscretizations +using RRTMGP.Fluxes +using RRTMGP.RTE +using RRTMGP.RTESolver +using RRTMGP.Optics +using RRTMGP.GrayUtils +using RRTMGP.AtmosphericStates +using RRTMGP.Sources +using RRTMGP.BCs + +import RRTMGP.Parameters.RRTMGPParameters +import ClimaParams as CP + +FT = Float64 +context = ClimaComms.context() + +device = ClimaComms.device(context) +param_set = RRTMGPParameters(FT) +ncol = 1 +# ncol = if device isa ClimaComms.CUDADevice +# 4096 +# else +# 9 +# end + +DA = ClimaComms.array_type(device) + +nlay = 60 # number of layers +p0 = FT(100000) # surface pressure (Pa) +pe = FT(9000) # TOA pressure (Pa) +nbnd, ngpt = 1, 1 # # of nbands/g-points (=1 for gray radiation) +nlev = nlay + 1 # # of layers +tb = FT(320) # surface temperature +# tstep = 6.0 # timestep in hours +# Δt = FT(60 * 60 * tstep) # timestep in seconds +# ndays = 365 * 40 # # of simulation days +# nsteps = ndays * (24 / tstep) # number of timesteps +nsteps = 1 +temp_toler = FT(0.1) # tolerance for temperature (Kelvin) +flux_grad_toler = FT(1e-5) # tolerance for flux gradient +n_gauss_angles = 1 # for non-scattering calculation +sfc_emis = DA{FT}(undef, nbnd, ncol) # surface emissivity +sfc_emis .= FT(1.0) +inc_flux = nothing # incoming flux + +grid_params = RRTMGPGridParams(FT; context, nlay, ncol) + +if ncol == 1 + lat = DA{FT}([0]) # latitude +else + lat = DA{FT}(range(FT(-90), FT(90), length = ncol)) # latitude +end + +otp = GrayOpticalThicknessOGorman2008(FT) # optical thickness parameters +gray_as = setup_gray_as_pr_grid(context, nlay, lat, p0, pe, otp, param_set, DA) + +SLVLW = TwoStreamLWRTE +slv_lw = SLVLW(grid_params; params = param_set, sfc_emis, inc_flux) + +(; flux_up, flux_dn, flux_net) = slv_lw.flux +(; t_lay, p_lay, t_lev, p_lev) = gray_as +# sbc = FT(RRTMGP.Parameters.Stefan(param_set)) +# cp_d_ = FT(RRTMGP.Parameters.cp_d(param_set)) +# grav_ = FT(RRTMGP.Parameters.grav(param_set)) +# #---------------------------------------------------- +# hr_lay = DA{FT}(undef, nlay, ncol) +# T_ex_lev = DA{FT}(undef, nlev, ncol) +# flux_grad = DA{FT}(undef, nlay, ncol) +# flux_grad_err = FT(0) + +#for i in 1:nsteps + # calling the long wave gray radiation solver + solve_lw!(slv_lw, gray_as) + # computing heating rate +# compute_gray_heating_rate!(device, hr_lay, p_lev, ncol, nlay, flux_net, cp_d_, grav_) + # updating t_lay and t_lev based on heating rate +# update_profile_lw!( +# device, +# sbc, +# t_lay, +# t_lev, +# flux_dn, +# flux_net, +# hr_lay, +# flux_grad, +# T_ex_lev, +# Δt, +# nlay, +# nlev, +# ncol, +# ) +# flux_grad_err = maximum(flux_grad) +# if flux_grad_err < flux_grad_toler +# break +# end +# end + +# t_error = maximum(abs.(T_ex_lev .- gray_as.t_lev)) +# color2 = :cyan + +# printstyled("\nGray atmosphere longwave test with ncol = $ncol, nlev = $nlev, solver = $SLVLW\n", color = color2) +# printstyled("device = $device\n", color = color2) +# printstyled("Integration time = $(FT(nsteps)/FT(24.0/tstep) / FT(365.0)) years\n\n", color = color2) + +# println("t_error = $(t_error); flux_grad_err = $(flux_grad_err)\n") +# end + +# function gray_atmos_sw_test( +# context, +# ::Type{SLVSW}, +# ::Type{FT}, +# ncol::Int; +# exfiltrate = false, +# ) where {FT <: AbstractFloat, SLVSW} +# param_set = RRTMGPParameters(FT) +# device = ClimaComms.device(context) +# DA = ClimaComms.array_type(device) +# nlay = 60 # number of layers +# p0 = FT(100000) # surface pressure (Pa) +# pe = FT(9000) # TOA pressure (Pa) +# nbnd, ngpt = 1, 1 # # of nbands/g-points (=1 for gray radiation) +# nlev = nlay + 1 # # of layers +# tb = FT(320) # surface temperature +# tstep = 6.0 # timestep in hours +# Δt = FT(60 * 60 * tstep) # timestep in seconds +# ndays = 365 * 40 # # of simulation days +# nsteps = ndays * (24 / tstep) # number of timesteps +# n_gauss_angles = 1 # for non-scattering calculation +# sfc_emis = Array{FT}(undef, nbnd, ncol) # surface emissivity +# sfc_emis .= FT(1.0) +# deg2rad = FT(π) / FT(180) +# grid_params = RRTMGPGridParams(FT; context, nlay, ncol) + +# cos_zenith = DA{FT, 1}(undef, ncol) # cosine of solar zenith angle +# toa_flux = DA{FT, 1}(undef, ncol) # top of atmosphere flux +# sfc_alb_direct = DA{FT, 2}(undef, nbnd, ncol) # surface albedo (direct) +# sfc_alb_diffuse = DA{FT, 2}(undef, nbnd, ncol) # surface albedo (diffuse) +# inc_flux_diffuse = nothing +# otp = GrayOpticalThicknessOGorman2008(FT) # optical thickness parameters + +# top_at_1 = false # Top-of-atmos at pt# 1 (true/false) + +# if ncol == 1 +# lat = DA{FT}([0]) # latitude +# else +# lat = DA{FT}(range(FT(-90), FT(90), length = ncol)) # latitude +# end + +# cos_zenith .= cos(deg2rad * 52.95) # corresponding to ~52.95 deg zenith angle +# toa_flux .= FT(1407.679) +# sfc_alb_direct .= FT(0.1) +# sfc_alb_diffuse .= FT(0.1) +# inc_flux_diffuse = nothing + +# as = setup_gray_as_pr_grid(context, nlay, lat, p0, pe, otp, param_set, DA) # init gray atmos state + +# swbcs = (; cos_zenith, toa_flux, sfc_alb_direct, inc_flux_diffuse, sfc_alb_diffuse) +# slv_sw = SLVSW(grid_params; swbcs...) + +# solve_sw!(slv_sw, as) + +# τ = Array(slv_sw.op.τ) +# cos_zenith = Array(cos_zenith) +# flux_dn_dir = Array(slv_sw.flux.flux_dn_dir) +# toa_flux = Array(toa_flux) + +# # testing with exact solution +# ot_tot = sum(τ[:, 1]) / cos_zenith[1] +# exact = (toa_flux[1] * cos_zenith[1]) * exp(-ot_tot) + +# rel_toler = FT(0.001) +# rel_error = abs(flux_dn_dir[1] - exact) / exact + +# color2 = :cyan + +# printstyled("\nGray atmosphere shortwave test with ncol = $ncol, nlev = $nlev, solver = $SLVSW\n", color = color2) +# printstyled("device = $device\n\n", color = color2) + +# println("relative error = $rel_error") +# end diff --git a/test/utils_test.jl b/test/utils_test.jl new file mode 100644 index 00000000..fa7ed899 --- /dev/null +++ b/test/utils_test.jl @@ -0,0 +1,82 @@ +using Breeze.Radiation: context_from_arch, ncols + +using RRTMGP.Parameters: RRTMGPParameters +using RRTMGP.GrayUtils: GrayOpticalThicknessOGorman2008 +using RRTMGP.GrayUtils: setup_gray_as_pr_grid + +using Oceananigans: RectilinearGrid +using Oceananigans: field, ynodes, Center, CenterField, FieldBoundaryConditions +using Oceananigans.Architectures: array_type +using Oceananigans: fill_halo_regions! +using Oceananigans: set! +using Oceananigans.BoundaryConditions: ValueBoundaryCondition + +""" + latitude(grid; lat_center=0, planet_radius=6_371_000, unit=:degrees) + +Return a device array of latitudes on the grid's horizontal layout with shape +`(Nx, Ny)`. Values are centered around `lat_center` and computed from the +grid's y-nodes and `planet_radius`. `unit` can be `:degrees` or `:radians`. +""" +function latitude_from_grid(grid::RectilinearGrid; lat_center=0, planet_radius=6_371_000, unit=:degrees) + FT = eltype(grid) + DA = array_type(grid.architecture) + + y = ynodes(grid, Center()) + lat = lat_center .+ y / planet_radius * (unit == :degrees ? FT(180/π) : 1) + lat = DA{FT}(repeat(lat', grid.Nx, 1)) + + return lat +end + +""" + gray_test_t_p_profiles(grid; p0, pe) + +Generate test-case pressure and temperature `Field`s for a gray atmosphere over +the provided `grid`. Profiles are derived from RRTMGP's gray utilities and +returned as Oceananigans fields shaped `(Nx, Ny, Nz)`. + +Returns `(pressure, temperature)`. +""" +function gray_test_t_p_profiles(grid::RectilinearGrid; p0, pe) + # get data types and context from grid + FT = eltype(grid) + context = context_from_arch(grid.architecture) + DA = array_type(grid.architecture) + + # derive latitude from grid + Nx, Ny, Nz = size(grid) + Ncols = ncols(grid) + if Ncols == 1 + lat = DA{FT}([0]) + else + lat = DA{FT}(range(FT(-90), FT(90), length = Ncols)) + end + + # get test case temperature and pressure profiles + param_set = RRTMGPParameters(FT) + otp = GrayOpticalThicknessOGorman2008(FT) + gray_as = setup_gray_as_pr_grid(context, Nz, lat, FT(p0), FT(pe), otp, param_set, DA) + p_lay_as_array = reshape(gray_as.p_lay', Nx, Ny, Nz) + t_lay_as_array = reshape(gray_as.t_lay', Nx, Ny, Nz) + + # create Fields from test case temperature and pressure profiles + p_bcs = FieldBoundaryConditions( + grid, (Center(), Center(), Center()); + bottom = ValueBoundaryCondition(gray_as.p_lev[1]), + top = ValueBoundaryCondition(gray_as.p_lev[end]), + ) + T_bcs = FieldBoundaryConditions( + grid, (Center(), Center(), Center()); + bottom = ValueBoundaryCondition(gray_as.t_lev[1]), + top = ValueBoundaryCondition(gray_as.t_lev[end]), + ) + pressure = CenterField(grid; boundary_conditions = p_bcs) + temperature = CenterField(grid; boundary_conditions = T_bcs) + set!(pressure, p_lay_as_array) + fill_halo_regions!(pressure) # ensure halos reflect BCs + set!(temperature, t_lay_as_array) + fill_halo_regions!(temperature) # ensure halos reflect BCs + + return pressure, temperature, gray_as +end