Skeleton ideas for radiation interfaces

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]

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

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

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

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