Extension to RRTMGP.jl and implement RadiativeTransferModel

docs/make.jl

@@ -44,6 +44,7 @@ makedocs(
             "Overview" => "microphysics/microphysics_overview.md",
             "Warm phase saturation adjustment" => "microphysics/saturation_adjustment.md",
         ],
+        "Radiative Transfer" => "radiative_transfer.md",
         "Dycore equations and algorithms" => "dycore_equations_algorithms.md",
         "Appendix" => Any[
             "Notation" => "appendix/notation.md",

docs/src/radiative_transfer.md

@@ -0,0 +1,173 @@
+# Radiative Transfer
+
+Breeze.jl integrates with [RRTMGP.jl](https://github.com/NumericalEarth/RRTMGP.jl) to provide radiative transfer capabilities for atmospheric simulations. The radiative transfer model computes longwave and shortwave radiative fluxes and heating rates, which are then incorporated into the moist static energy tendency equation.
+
+## Overview
+
+The radiative transfer implementation in Breeze uses a column-based approach, where each horizontal column of the 3D grid is treated independently. This allows efficient computation of radiative fluxes using RRTMGP's optimized column radiation solver.
+
+## Basic Usage
+
+To use radiative transfer in a Breeze simulation, create a `RadiativeTransferModel` and pass it to the `AtmosphereModel` constructor:
+
+```julia
+using Breeze
+using Oceananigans
+
+# Create grid
+grid = RectilinearGrid(CPU(), Float64; size=(32, 32, 64), x=(0, 10_000), y=(0, 10_000), z=(0, 20_000))
+
+# Create radiative transfer model
+rtm = RadiativeTransferModel(
+    grid;
+    surface_emissivity = 0.98,
+    surface_albedo_direct = 0.1,
+    surface_albedo_diffuse = 0.1,
+    cos_zenith = 0.5,
+    toa_solar_flux = 1360.0,
+    toa_longwave_flux = 0.0
+)
+
+# Create atmosphere model with radiative transfer
+model = AtmosphereModel(
+    grid;
+    radiative_transfer = rtm
+)
+```
+
+## Column Model Example
+
+Here we demonstrate a simple column model calculation of radiative fluxes. This example sets up a single-column atmosphere and computes radiative heating rates.
+
+```julia
+using Breeze
+using Oceananigans
+using CairoMakie
+
+# Create a single-column grid (1x1 horizontal, 64 vertical levels)
+grid = RectilinearGrid(CPU(), Float64; 
+    size=(1, 1, 64), 
+    x=(0, 1), 
+    y=(0, 1), 
+    z=(0, 20_000)
+)
+
+# Thermodynamic constants
+thermo = ThermodynamicConstants(Float64)
+
+# Reference state
+reference_state = ReferenceState(grid, thermo, 
+    base_pressure = 101325.0,
+    potential_temperature = 288.0
+)
+
+formulation = AnelasticFormulation(reference_state)
+
+# Create radiative transfer model
+rtm = RadiativeTransferModel(
+    grid;
+    surface_emissivity = 0.98,
+    surface_albedo_direct = 0.1,
+    surface_albedo_diffuse = 0.1,
+    cos_zenith = 0.5,  # Solar zenith angle cosine
+    toa_solar_flux = 1360.0,  # Top-of-atmosphere solar flux [W/m²]
+    toa_longwave_flux = 0.0
+)
+
+# Create atmosphere model
+model = AtmosphereModel(
+    grid;
+    thermodynamics = thermo,
+    formulation = formulation,
+    radiative_transfer = rtm
+)
+
+# Initialize with a temperature profile
+# Simple linear temperature profile decreasing with height
+set!(model.temperature, (x, y, z) -> 288.0 - 0.0065 * z)
+
+# Update model state
+update_state!(model)
+
+# Update radiative fluxes
+Breeze.AtmosphereModels._update_radiative_fluxes!(rtm, model)
+
+# Extract heating rates
+nz = size(grid, 3)
+heating_rate = zeros(Float64, nz)
+
+using Oceananigans: interior
+using Oceananigans.Grids: znode
+using GPUArraysCore: @allowscalar
+
+@allowscalar begin
+    ρᵣ = model.formulation.reference_state.density
+    for k in 1:nz
+        hr = Breeze.AtmosphereModels._radiative_heating_rate(
+            1, 1, k, grid, rtm, 
+            ρᵣ,
+            thermo
+        )
+        # Convert from energy density tendency to heating rate per unit mass
+        heating_rate[k] = hr / ρᵣ[1, 1, k]
+    end
+end
+
+# Extract height levels for plotting
+z_lev = [znode(1, 1, k, grid, Center(), Center(), Center()) for k in 1:nz]
+
+# Plot heating rate profile
+using CairoMakie
+
+fig = Figure(resolution = (600, 400))
+ax = Axis(fig[1, 1], 
+    xlabel = "Heating Rate [K/day]",
+    ylabel = "Height [m]",
+    title = "Radiative Heating Rate Profile"
+)
+
+# Convert to K/day
+cp = thermo.dry_air.heat_capacity
+seconds_per_day = 86400.0
+heating_rate_k_per_day = heating_rate .* seconds_per_day ./ cp
+
+lines!(ax, heating_rate_k_per_day, z_lev, linewidth = 2)
+fig
+```
+
+## Surface Properties
+
+The `RadiativeTransferModel` requires several surface properties that are not part of the `AtmosphereModel`:
+
+- **Surface temperature**: Temperature at the surface (can be extracted from model or specified)
+- **Surface emissivity**: Longwave emissivity of the surface (typically 0.95-0.99)
+- **Surface albedo (direct)**: Albedo for direct solar radiation
+- **Surface albedo (diffuse)**: Albedo for diffuse solar radiation
+- **Cosine of solar zenith angle**: Determines solar insolation
+- **TOA solar flux**: Top-of-atmosphere solar flux [W/m²]
+- **TOA longwave flux**: Top-of-atmosphere longwave flux [W/m²] (usually 0)
+
+These properties are stored in the `RadiativeTransferModel` and can be updated as needed during the simulation.
+
+## Reference Pressure
+
+Breeze uses the **reference pressure** from the anelastic formulation for radiative transfer calculations, not the total pressure. This is consistent with the anelastic approximation where pressure perturbations are small compared to the reference state.
+
+## Integration with Dynamics
+
+Radiative heating is automatically added to the moist static energy tendency equation when a `RadiativeTransferModel` is provided. The heating rate is computed from flux differences:
+
+```math
+\frac{\partial (\rho e)}{\partial t} = \ldots + \rho_r \frac{g}{c_p} \frac{F_{k+1} - F_k}{\Delta p}
+```
+
+where ``F_k`` is the net radiative flux at level ``k``, ``\Delta p`` is the pressure difference across the layer, ``g`` is gravitational acceleration, ``c_p`` is specific heat capacity, and ``\rho_r`` is the reference density.
+
+## Gray Atmosphere Model
+
+The current implementation uses a gray atmosphere radiation model, which treats the atmosphere as having a single effective absorption coefficient. This is suitable for initial testing and development. Future versions will support full spectral radiation using RRTMGP's band-by-band calculations.
+
+## Architecture Support
+
+The radiative transfer implementation supports both CPU and GPU architectures. The grid conversion utilities automatically handle the conversion between Oceananigans' 3D grid format and RRTMGP's column-based format, including proper handling of CPU and GPU arrays.
+

ext/RRTMGPExt.jl

@@ -0,0 +1,27 @@
+module RRTMGPExt
+
+import Breeze
+import RRTMGP
+
+# Load the RadiativeTransfer module when RRTMGP is available
+include(joinpath(@__DIR__, "..", "src", "RadiativeTransfer", "RadiativeTransfer.jl"))
+
+using .RadiativeTransfer
+
+# Re-export RadiativeTransferModel in Breeze module
+const RadiativeTransferModel = RadiativeTransfer.RadiativeTransferModel
+
+# Make functions available to AtmosphereModels module by adding them to the module
+# This allows the functions to be called from dynamics_kernel_functions.jl
+Breeze.AtmosphereModels.eval(quote
+    const _radiative_heating_rate = $(RadiativeTransfer._radiative_heating_rate)
+    const _update_radiative_fluxes! = $(RadiativeTransfer._update_radiative_fluxes!)
+end)
+
+# Export RadiativeTransferModel from Breeze
+Breeze.eval(quote
+    export RadiativeTransferModel
+end)
+
+end # module
+

src/AtmosphereModels/atmosphere_model.jl

@@ -33,7 +33,7 @@ end
 formulation_pressure_solver(formulation, grid) = nothing
 
 mutable struct AtmosphereModel{Frm, Arc, Tst, Grd, Clk, Thm, Den, Mom, Eng, Mse, Moi, Mfr,
-                               Tmp, Prs, Ppa, Sol, Vel, Trc, Adv, Cor, Frc, Mic, Cnd, Cls, Cfs} <: AbstractModel{Tst, Arc}
+                               Tmp, Prs, Ppa, Sol, Vel, Trc, Adv, Cor, Frc, Mic, Cnd, Cls, Cfs, Rad} <: AbstractModel{Tst, Arc}
     architecture :: Arc
     grid :: Grd
     clock :: Clk
@@ -59,6 +59,7 @@ mutable struct AtmosphereModel{Frm, Arc, Tst, Grd, Clk, Thm, Den, Mom, Eng, Mse,
     timestepper :: Tst
     closure :: Cls
     closure_fields :: Cfs
+    radiative_transfer :: Rad
 end
 
 function default_formulation(grid, thermo)
@@ -108,7 +109,8 @@ function AtmosphereModel(grid;
                          advection = Centered(order=2),
                          closure = nothing,
                          microphysics = nothing, # WarmPhaseSaturationAdjustment(),
-                         timestepper = :RungeKutta3)
+                         timestepper = :RungeKutta3,
+                         radiative_transfer = nothing)
 
     arch = grid.architecture
     tracers = tupleit(tracers) # supports tracers=:c keyword argument (for example)
@@ -185,7 +187,8 @@ function AtmosphereModel(grid;
                             microphysical_fields,
                             timestepper,
                             closure,
-                            closure_fields)
+                            closure_fields,
+                            radiative_transfer)
 
     update_state!(model)
 

src/AtmosphereModels/dynamics_kernel_functions.jl

@@ -129,15 +129,26 @@ end
                                               moisture_mass_fraction,
                                               thermo,
                                               microphysical_fields,
-                                              microphysics)
+                                              microphysics,
+                                              radiative_transfer)
 
     # Compute the buoyancy flux term, ρᵣ w b
     buoyancy_flux = ℑzᵃᵃᶜ(i, j, k, grid, ρ_w_bᶜᶜᶠ, velocities.w, reference_density,
                           temperature, moisture_mass_fraction, formulation, thermo)
 
+    # Compute radiative heating if radiative transfer is enabled
+    radiative_heating = if radiative_transfer === nothing
+        zero(eltype(energy_density))
+    else
+        # Call radiative heating rate function
+        # Note: This will be available when RRTMGP extension is loaded
+        _radiative_heating_rate(i, j, k, grid, radiative_transfer, reference_density, thermo)
+    end
+
     return ( - div_Uc(i, j, k, grid, advection, velocities, energy_density)
              + buoyancy_flux
              - ∇_dot_Jᶜ(i, j, k, grid, reference_density, closure, closure_fields, id, energy, clock, model_fields, nothing)
              # + microphysical_energy_tendency(i, j, k, grid, formulation, microphysics, microphysical_fields)
-             + ρe_forcing(i, j, k, grid, clock, model_fields))
+             + ρe_forcing(i, j, k, grid, clock, model_fields)
+             + radiative_heating)
 end

src/AtmosphereModels/update_atmosphere_model_state.jl

@@ -130,9 +130,15 @@ function compute_tendencies!(model::AnelasticModel)
     ρe = model.energy_density
     e = model.moist_static_energy
     Fρe = model.forcing.ρe
+    # Update radiative fluxes if radiative transfer is enabled
+    if model.radiative_transfer !== nothing
+        # Call update function - will be available when RRTMGP extension is loaded
+        _update_radiative_fluxes!(model.radiative_transfer, model)
+    end
+
     ρe_args = tuple(ρe, Val(1), e, Fρe, scalar_args...,
                     model.formulation, model.temperature,
-                    model.moisture_mass_fraction, model.thermodynamics, model.microphysical_fields, model.microphysics)
+                    model.moisture_mass_fraction, model.thermodynamics, model.microphysical_fields, model.microphysics, model.radiative_transfer)
     launch!(arch, grid, :xyz, compute_moist_static_energy_tendency!, Gρe, grid, ρe_args)
 
     ρqᵗ = model.moisture_density

src/Breeze.jl

@@ -87,4 +87,7 @@ using .Microphysics
 include("TurbulenceClosures/TurbulenceClosures.jl")
 using .TurbulenceClosures
 
+# RadiativeTransfer extension will be loaded via ext/RRTMGPExt.jl when RRTMGP is available
+# Export will be added by the extension
+
 end # module Breeze

src/RadiativeTransfer/RadiativeTransfer.jl

@@ -0,0 +1,10 @@
+module RadiativeTransfer
+
+export RadiativeTransferModel, _radiative_heating_rate, _update_radiative_fluxes!
+
+include("radiative_transfer_model.jl")
+include("grid_conversion.jl")
+include("atmosphere_model_integration.jl")
+
+end # module
+

src/RadiativeTransfer/atmosphere_model_integration.jl

@@ -0,0 +1,61 @@
+"""
+Integration of RadiativeTransferModel with AtmosphereModel.
+"""
+
+using Oceananigans: CenterField
+using Oceananigans.Operators: ℑzᵃᵃᶜ
+
+import ..AtmosphereModels: AtmosphereModel, AnelasticFormulation
+import ..grid_conversion: reshape_from_columns
+
+"""
+    _radiative_heating_rate(i, j, k, grid, rtm::RadiativeTransferModel, 
+                          reference_density, thermo)
+
+Compute the radiative heating rate contribution to moist static energy tendency
+at grid point (i, j, k).
+
+The heating rate is computed from flux differences:
+    hr = g * (flux_net[k+1] - flux_net[k]) / (cp * Δp)
+
+and the contribution to energy density tendency is:
+    ρᵣ * hr
+
+where ρᵣ is the reference density.
+"""
+@inline function _radiative_heating_rate(i, j, k, grid, rtm, reference_density, thermo)
+    # Get net fluxes
+    flux_net_lw = rtm.flux_lw.flux_net
+    flux_net_sw = rtm.flux_sw.flux_net
+
+    # Combine longwave and shortwave
+    flux_net_total = flux_net_lw .+ flux_net_sw
+
+    # Convert grid indices to column index
+    nx, ny, nz = size(grid)
+    icol = (j - 1) * nx + i
+
+    # Get flux at levels k and k+1 (fluxes are at levels, k is cell index)
+    # Level k corresponds to bottom of cell k
+    # Level k+1 corresponds to top of cell k
+    flux_bottom = @inbounds flux_net_total[k, icol]
+    flux_top = @inbounds flux_net_total[k+1, icol]
+
+    # Get pressure difference from atmospheric state
+    p_lev = rtm.atmospheric_state.p_lev
+    Δp = @inbounds p_lev[k+1, icol] - p_lev[k, icol]
+
+    # Get constants from thermodynamics
+    g = thermo.gravitational_acceleration
+    cp = thermo.dry_air.heat_capacity
+
+    # Compute heating rate per unit mass
+    hr = g * (flux_top - flux_bottom) / (cp * Δp)
+
+    # Get reference density at this point
+    ρᵣ = @inbounds reference_density[i, j, k]
+
+    # Return contribution to energy density tendency
+    return ρᵣ * hr
+end
+