wip: interface shaping up for gray radiation model

Author: bischtob

src/Radiation/gray_breezy.jl

@@ -1,5 +1,13 @@
 abstract type AbstractRadiationModel end
 
-function update_radative_fluxes!(model::AbstractRadiationModel) 
-    throw("Not implemented for $(typeof(model)).")
+"""
+    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!(model::AbstractRadiationModel)
+    solve_lw!(model.slv_lw, model.atmospheric_state)
+    solve_sw!(model.slv_sw, model.atmospheric_state)
 end

src/Radiation/radiation_model.jl

@@ -7,116 +7,121 @@ using Oceananigans
 using Oceananigans: field
 using Oceananigans.Architectures: array_type
 using Oceananigans: RectilinearGrid
+ 
 
-mutable struct GrayRadiationModel{FT, DA, SLVLW, SLVSW, AS, OTP} <: AbstractRadiationModel
+"""
+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, Nx, Ny) — surface emissivity per band
+- sfc_albedo_direct: (Nbnd, Nx, Ny) — direct-beam surface albedo per band
+- sfc_albedo_diffuse: (Nbnd, Nx, Ny) — diffuse surface albedo per band
+- sw_toa_flux_inc: (Nx, Ny) — incoming shortwave flux at TOA
+- sw_inc_flux_diffuse: (Nx, Ny) or nothing — optional diffuse SW incoming flux
+- lw_toa_inc_flux: (Nx, Ny) or nothing — optional incoming LW flux at TOA
+"""
+mutable struct GrayRadiationModel{FT, AS, SLVLW, SLVSW, OTP} <: AbstractRadiationModel
     atmospheric_state :: AS
     slv_lw :: SLVLW
     slv_sw :: SLVSW
     optical_properties :: OTP
-    cos_zenith_angle :: DA{FT}
-    sfc_emissivity :: DA{FT}
-    sfc_albedo_direct :: DA{FT}
-    sfc_albedo_diffuse :: DA{FT}
-    sw_toa_flux_inc :: DA{FT}
-    sw_inc_flux_diffuse :: Union{Nothing, DA{FT}}
-    lw_toa_inc_flux :: Union{Nothing, DA{FT}}
+    cos_zenith_angle :: AbstractArray{FT}
+    sfc_emissivity :: AbstractArray{FT}
+    sfc_albedo_direct :: AbstractArray{FT}
+    sfc_albedo_diffuse :: AbstractArray{FT}
+    sw_toa_flux_inc :: AbstractArray{FT}
+    sw_inc_flux_diffuse :: Union{Nothing, AbstractArray{FT}}
+    lw_toa_inc_flux :: Union{Nothing, AbstractArray{FT}}
 end
 
+"""
+    GrayRadiationModel(grid; atmospheric_state,
+                       zenith_angle, sfc_emissivity, sfc_albedo_direct,
+                       sfc_albedo_diffuse, sw_flux_inc_toa;
+                       sw_flux_inc_toa_diffusive=nothing,
+                       lw_flux_inc_toa=nothing,
+                       optical_properties=nothing)
+
+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; 
-    temperature, 
+    grid;
+    temperature,
     pressure,
-    cos_zenith_angle,
+    zenith_angle,
     sfc_emissivity,
     sfc_albedo_direct,
     sfc_albedo_diffuse,
     sw_flux_inc_toa,
-    sw_flux_inc_toa_diffusive,
-    lw_flux_inc_toa,
+    sw_flux_inc_toa_diffusive=nothing,
+    lw_flux_inc_toa=nothing,
     optical_properties=GrayOpticalThicknessSchneider2004(FT),
-    lat_center=0
-)  
-    # We assemble objects required for RRTMGP to work.
+    lat_center=0,
+)
+    DA = array_type(grid.architecture)
+    FT = eltype(grid)
+    Nx, Ny = grid.Nx, grid.Ny
+    Nbnd = 1
+
+    # Create atmospheric state from provided Fields (generic library behavior)
     atmospheric_state = GrayAtmosphericState(
-        grid; 
-        temperature, 
-        pressure, 
-        otp=optical_properties, 
-        lat_center=lat_center
+        grid;
+        temperature,
+        pressure,
+        otp=optical_properties,
+        lat_center=lat_center,
     )
+
+    # Helper functions
+    to_grid_array(val) = val isa AbstractArray ? DA(val) : (A = DA{FT}(undef, Nx, Ny); fill!(A, FT(val)))
+    to_banded_array(val) = val isa AbstractArray ? DA(val) : (A = DA{FT}(undef, Nbnd, Nx, Ny); fill!(A, FT(val)))
+    to_optional_grid_array(val) = val === nothing ? nothing : to_grid_array(val)
+
+    # Bring arrays to the correct shape for RRTMGP.jl interface
+    cos_zenith = if zenith_angle isa AbstractArray
+        DA(cos.(FT(π / 180) .* zenith_angle))
+    else
+        A = DA{FT}(undef, Nx, Ny)
+        fill!(A, FT(cos(FT(π / 180) * zenith_angle)))
+    end
+    sfc_emission = to_banded_array(sfc_emissivity)
+    sfc_alb_direct = to_banded_array(sfc_albedo_direct)
+    sfc_alb_diffuse = to_banded_array(sfc_albedo_diffuse)
+    toa_flux = to_grid_array(sw_flux_inc_toa)
+    inc_flux_diffuse = to_optional_grid_array(sw_flux_inc_toa_diffusive)
+    lw_toa_inc_flux = to_optional_grid_array(lw_flux_inc_toa)
+
+    # Build solvers using the high-level wrappers (they reshape internally)
     SLVLW = TwoStreamLWRTE
     SLVSW = TwoStreamSWRTE
     lw_params = (
-        sfc_emission = sfc_emissivity, 
-        lw_inc_flux = lw_flux_inc_toa,
+        sfc_emission = sfc_emission,
+        lw_inc_flux = lw_toa_inc_flux,
     )
     sw_params = (
-        cos_zenith = cos_zenith_angle, 
-        toa_flux = sw_flux_inc_toa,
-        sfc_alb_direct = sfc_albedo_diffuse,
-        inc_flux_diffuse = sw_flux_inc_toa_diffusive,
-        sfc_alb_diffuse = sfc_albedo_diffuse,
+        cos_zenith = cos_zenith,
+        toa_flux = toa_flux,
+        sfc_alb_direct = sfc_alb_direct,
+        inc_flux_diffuse = inc_flux_diffuse,
+        sfc_alb_diffuse = sfc_alb_diffuse,
     )
     slv_lw = SLVLW(grid; lw_params...)
     slv_sw = SLVSW(grid; sw_params...)
-    
-    if sfc_emissivity isa DA
-
-    elseif sfc_emissivity isa AbstractFloat
-
-    end
 
     return GrayRadiationModel(
         atmospheric_state,
-        slv_lw, 
+        slv_lw,
         slv_sw,
         optical_properties,
-        cos_zenith_angle,
+        cos_zenith,
         sfc_emission,
-        sfc_albedo_direct,
-        sfc_albedo_diffuse,
-        sw_flux_inc_toa, 
-        sw_flux_inc_toa_diffusive,
-        lw_flux_inc_toa,
+        sfc_alb_direct,
+        sfc_alb_diffuse,
+        toa_flux,
+        inc_flux_diffuse,
+        lw_toa_inc_flux,
     )
 end
-
-function GrayRadiationModel(
-    grid; 
-    temperature :: Field, 
-    pressure :: Field,
-    zenith_angle :: FT,
-    sfc_emissivity :: FT,
-    sfc_albedo_direct :: FT,
-    sfc_albedo_diffuse :: FT,
-    sw_flux_inc_toa :: FT,
-    sw_flux_inc_toa_diffusive :: Union{Nothing, FT},
-    lw_flux_inc_toa :: Union{Nothing, FT},
-    optical_properties=GrayOpticalThicknessSchneider2004(FT),
-    lat_center=0
-)
-    FT = eltype(grid)
-    DA = device_array(grid.architecture)
-    Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz
-    Nbnd = 1
-
-
-    sfc_emission = DA{FT}(undef, Nbnd, Nx, Ny)
-    sfc_alb_direct = DA{FT}(undef, Nbnd, Nx, Ny)
-    sfc_alb_diffuse = DA{FT}(undef, Nbnd, Nx, Ny)
-    cos_zenith = DA{FT}(undef, Nx, Ny)
-    toa_flux = DA{FT}(undef, Nx, Ny)
-    lw_toa_inc_flux = nothing
-    inc_flux_diffuse = nothing
-    fill!(sfc_emission, FT(sfc_emissivity))
-    fill!(sfc_alb_direct, FT(albedo_direct))
-    fill!(sfc_alb_diffuse, FT(albedo_diffuse))
-    fill!(cos_zenith, FT(cos(deg2rad * zenith_angle)))
-    fill!(toa_flux, FT(sw_inc_flux))
-end
-
-function (rad::GrayRadiationModel)(::Val{:ρe}, temperature, pressure)
-    update_atmospheric_state(rad, temperature, pressure)
-    solve_lw!(rad.slv_lw, rad.atmospheric_state)
-    solve_sw!(rad.slv_sw, rad.atmospheric_state)    
-end