wip: cleaner interface

Author: bischtob

src/Radiation/Radiation.jl

@@ -5,9 +5,9 @@ export
     GrayRadiationModel,
     update_radative_fluxes!
 
+include("utils_grid.jl")
 include("rrtmgp_interface.jl")
 include("radiation_model.jl")
-include("radiation_model_utils.jl")
 include("radiation_model_gray.jl")
 
 end

src/Radiation/_rrtmgp_design.jl

@@ -7,7 +7,6 @@ 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)
+function update_radative_fluxes!(::AbstractRadiationModel)
+    throw("Not implemented")
 end

src/Radiation/radiation_model.jl

@@ -1,4 +1,3 @@
-using RRTMGP.Optics: GrayOpticalThicknessSchneider2004
 using RRTMGP.AtmosphericStates: GrayAtmosphericState
 using RRTMGP.RTE: TwoStreamLWRTE, TwoStreamSWRTE
 using RRTMGP.RTESolver: solve_lw!, solve_sw!
@@ -8,7 +7,6 @@ using Oceananigans: field
 using Oceananigans.Architectures: array_type
 using Oceananigans: RectilinearGrid
 
-
 """
 GrayRadiationModel stores state and solver handles for a gray-band radiative
 transfer setup.
@@ -18,110 +16,87 @@ Field conventions (array shapes):
 - 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
+- toa_sw_flux_inc: (Nx, Ny) — incoming shortwave flux at TOA
 """
-mutable struct GrayRadiationModel{FT, AS, SLVLW, SLVSW, OTP} <: AbstractRadiationModel
+struct GrayRadiationModel{FT, AS, SLVLW, SLVSW} <: AbstractRadiationModel
     atmospheric_state :: AS
-    slv_lw :: SLVLW
-    slv_sw :: SLVSW
-    optical_properties :: OTP
     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}}
+    toa_sw_toa_flux :: AbstractArray{FT}
+    solver_lw :: SLVLW
+    solver_sw :: SLVSW
 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)
+    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;
-    temperature,
-    pressure,
+    temperature :: Field,
+    pressure :: Field,
     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=GrayOpticalThicknessSchneider2004(FT),
-    lat_center=0,
+    toa_sw_flux_inc,
+    latitude,
+    isothermal_boundary_layer=false,
 )
     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,
+        temperature=temperature,
+        pressure=pressure,
+        latitude=latitude,
     )
 
-    # 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)
+    # Build solvers using the high-level wrappers (they reshape internally for RRTMGP.jl compatibility)
     SLVLW = TwoStreamLWRTE
     SLVSW = TwoStreamSWRTE
     lw_params = (
-        sfc_emission = sfc_emission,
-        lw_inc_flux = lw_toa_inc_flux,
+        sfc_emission = DA(sfc_emissivity),
+        lw_inc_flux = nothing,
+        isothermal_boundary_layer = isothermal_boundary_layer,
     )
     sw_params = (
-        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,
+        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,
     )
-    slv_lw = SLVLW(grid; lw_params...)
-    slv_sw = SLVSW(grid; sw_params...)
+    solver_lw = SLVLW(grid; lw_params...)
+    solver_sw = SLVSW(grid; sw_params...)
 
     return GrayRadiationModel(
         atmospheric_state,
-        slv_lw,
-        slv_sw,
-        optical_properties,
-        cos_zenith,
-        sfc_emission,
-        sfc_alb_direct,
-        sfc_alb_diffuse,
-        toa_flux,
-        inc_flux_diffuse,
-        lw_toa_inc_flux,
+        sw_params.cos_zenith,
+        lw_params.sfc_emission,
+        sw_params.sfc_alb_direct,
+        sw_params.sfc_alb_diffuse,
+        sw_params.toa_flux,
+        solver_lw,
+        solver_sw,
     )
 end
+
+function update_radative_fluxes!(model::GrayRadiationModel)
+    solve_lw!(model.solver_lw, model.atmospheric_state)
+    solve_sw!(model.solver_sw, model.atmospheric_state)
+end