wip

Author: bischtob

src/Radiation/gray_breezy.jl

@@ -1,30 +1,64 @@
+using RRTMGP.Optics: GrayOpticalThicknessSchneider2004
+using RRTMGP.Parameters: RRTMGPParameters
+using RRTMGP.AtmosphericStates: GrayAtmosphericState
+using RRTMGP.RTE: TwoStreamLWRTE
+using RRTMGP.RTESolver: solve_lw!
+
 using Oceananigans
 using Oceananigans: field
 using Oceananigans.Architectures: array_type
 using Oceananigans: RectilinearGrid
 
-arch = CPU()
+include("rrtmgp_interface.jl")
+
+FT = Float64
+architecture = CPU()
+DA = array_type(architecture)
+Nx, Ny, Nz = 1, 1, 64
+Lx, Ly, Lz = 1, 20_000_000, 1_000 # domain size in meters
+Nbnd, Ngpt = 1, 1 # gray model
+lw_inc_flux = nothing # no incoming longwave flux
+sfc_emissivity = 1
+planet_radius = 6_371_000
+p_surface = 100_000 # surface pressure (Pa)
+p_top = 9_000 # top of atmosphere pressure / emission level (Pa)
+lat_center = 0
+optical_properties = GrayOpticalThicknessSchneider2004(FT)
+
+# Grid setup, we also need to say where on the planet our box is located
 grid = RectilinearGrid(
-    CPU(), 
-    size=(1, 1, 64), 
-    x=(0, 2π), 
-    y=(0, 2π), 
-    z=(0, 1), 
+    architecture,
+    FT,
+    size=(Nx, Ny, Nz,), 
+    x=(-Lx/2, Lx/2),
+    y=(-Ly/2, Ly/2), 
+    z=(0, Lz), 
     topology=(Periodic, Periodic, Bounded)
 )
 
+# Set up 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
+    pressure, temperature = gray_test_t_p_profiles(grid; p0=p_surface, pe=p_top)
+end
 
+# Set up atmospheric state
+atmospheric_state = GrayAtmosphericState(
+    grid; 
+    temperature, 
+    pressure, 
+    otp=optical_properties, 
+    lat_center=lat_center
+)
 
-function t_p_profiles(grid::RectilinearGrid; lat, p0, pe, otp)
-
-    context = context_from_arch(grid.architecture)
-    nlay = grid.Nz
-    param_set = RRTMGPParameters(eltype(grid))
-    DA = array_type(grid.architecture)
-    gray_as = setup_gray_as_pr_grid(context, nlay, lat, p0, pe, otp, param_set, DA)
+# Set up radiation model
+sfc_emission = DA{FT}(undef, Nbnd, Nx, Ny)
+fill!(sfc_emission, FT(sfc_emissivity))
+SLVLW = TwoStreamLWRTE
+slv_lw = SLVLW(grid; sfc_emission, lw_inc_flux)
 
-    pressure = field((Center(), Center(), Face()), gray_as.p_lay, grid)
-    temperature = field((Center(), Center(), Face()), gray_as.t_lay, grid)
+# solve the radiation model
+solve_lw!(slv_lw, atmospheric_state)
 
-    return pressure, temperature
-end

src/Radiation/proposed_design.jl

@@ -3,54 +3,84 @@ using RRTMGP: RRTMGPGridParams
 using RRTMGP.AngularDiscretizations
 using RRTMGP.Fluxes
 using RRTMGP.Optics
-using RRTMGP.GrayUtils
+using RRTMGP.GrayUtils: setup_gray_as_pr_grid
 using RRTMGP.Sources
 using RRTMGP.BCs
+using RRTMGP.Parameters: RRTMGPParameters
+using RRTMGP.Optics: GrayOpticalThicknessSchneider2004
+
+using Oceananigans
+using Oceananigans: RectilinearGrid
+using Oceananigans: field
 using ClimaParams
+using ClimaComms
 
 import RRTMGP.Parameters.RRTMGPParameters
 import RRTMGP.AtmosphericStates
 import RRTMGP.RTE: NoScatLWRTE, TwoStreamLWRTE, NoScatSWRTE, TwoStreamSWRTE
 
+function latitude(grid::RectilinearGrid; lat_center=0, planet_radius=6_371_000, unit=:degrees)
+    FT = eltype(grid)
+    DA = array_type(grid.architecture)
+    Nx = grid.Nx
+
+    y = ynodes(grid, Center())
+    lat = lat_center .+ y / planet_radius * (unit == :degrees ? 180/π : 1)
+    lat = DA{FT}(repeat(lat', Nx, 1))
+
+    return lat
+end
+
 function RRTMGP.RRTMGPGridParams(grid::RectilinearGrid; isothermal_boundary_layer::Bool = false)
-    nlay = grid.Nz
-    ncols = grid.Nx * grid.Ny
+    FT = eltype(grid)
+    Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz
+    Ncols = Nz * Ny
     context = context_from_arch(grid.architecture)
-    return RRTMGPGridParams(context, nlay, ncols, isothermal_boundary_layer)
+    return RRTMGPGridParams{FT, typeof(context)}(context, Nz, Ncols, isothermal_boundary_layer)
 end
 
 function RRTMGP.AtmosphericStates.GrayAtmosphericState(
     grid::RectilinearGrid;
     temperature::Field,
     pressure::Field,
-    lat::Field,
     otp,
+    lat_center=0, 
+    planet_radius=6_371_000,
 )
-    lat = grid.lat
-    nlay = grid.Nz
-    ncols = grid.Nx * grid.Ny
-    p_lay = reshape(interior(pressure), nlay, ncols)
-    p_lev = reshape(interior(Field(@at (Center, Center, Face) pressure)), nlay + 1, ncols)
-    t_lay = reshape(interior(temperature), nlay, ncols)
-    t_lev = reshape(interior(Field(@at (Center, Center, Face) temperature)), nlay + 1, ncols)
-    z_lev = znodes(grid, Face())
+    FT = eltype(grid)
+    Nx, Ny, Nz = grid.Nx, grid.Ny, grid.Nz
+    Ncols = Nx * Ny
+    lat = reshape(latitude(grid; lat_center, planet_radius), Ncols)
+    p_lay = reshape(interior(pressure), Nz, Ncols)
+    p_lev = reshape(interior(Field(@at (Center, Center, Face) pressure)), Nz + 1, Ncols)
+    t_lay = reshape(interior(temperature), Nz, Ncols)
+    t_lev = reshape(interior(Field(@at (Center, Center, Face) temperature)), Nz + 1, Ncols)
+    z_lev = reshape(view(repeat(znodes(grid, Face()), Nx, Ny, 1), :, :, :), Nz + 1, Ncols)
     t_sfc = t_lev[1, :]
-    otp = otp
 
-    return GrayAtmosphericState(lat, p_lay, p_lev, t_lay, t_lev, z_lev, t_sfc, otp)
+    return GrayAtmosphericState{FT, AbstractVector{FT}, AbstractMatrix{FT}, typeof(otp)}(
+        lat, 
+        p_lay, 
+        p_lev, 
+        t_lay, 
+        t_lev, 
+        z_lev, 
+        t_sfc, 
+        otp
+    )
 end
 
 function RRTMGP.RTE.NoScatLWRTE(
     grid::RectilinearGrid; 
-    sfc_emis, 
-    inc_flux
+    sfc_emission, 
+    lw_inc_flux
 )
     grid_params = RRTMGPGridParams(grid)
     params = RRTMGPParameters(eltype(grid))
     (; context) = grid_params
     op = OneScalar(grid_params)
     src = SourceLWNoScat(grid_params; params)
-    bcs = LwBCs(sfc_emis, inc_flux)
+    bcs = LwBCs(sfc_emission, lw_inc_flux)
     fluxb = FluxLW(grid_params)
     flux = FluxLW(grid_params)
     ad = AngularDiscretization(grid_params, 1)
@@ -59,15 +89,21 @@ end
 
 function RRTMGP.RTE.TwoStreamLWRTE(
     grid::RectilinearGrid; 
-    sfc_emis, 
-    inc_flux
+    sfc_emission, 
+    lw_inc_flux
 )
+    FT = eltype(grid)
+    Nx, Ny = grid.Nx, grid.Ny
+    Ncols = Nx * Ny
     grid_params = RRTMGPGridParams(grid)
     params = RRTMGPParameters(eltype(grid))
     (; context) = grid_params
     op = TwoStream(grid_params)
     src = SourceLW2Str(grid_params; params)
-    bcs = LwBCs(sfc_emis, inc_flux)
+    if lw_inc_flux !== nothing
+        lw_inc_flux = reshape(lw_inc_flux, 1, Ncols)
+    end
+    bcs = LwBCs(reshape(sfc_emission, 1, Ncols), lw_inc_flux)
     fluxb = FluxLW(grid_params)
     flux = FluxLW(grid_params)
     return TwoStreamLWRTE(context, op, src, bcs, fluxb, flux)
@@ -118,3 +154,36 @@ function context_from_arch(arch)
         return ClimaComms.context(ClimaComms.CUDADevice())
     end
 end
+
+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 = grid.Nx, grid.Ny, grid.Nz
+    ncol = grid.Nx * grid.Ny
+    if ncol == 1
+        lat = DA{FT}([0])                   
+    else
+        lat = DA{FT}(range(FT(-90), FT(90), length = ncol))
+    end
+    
+    # get test case temperature and pressure profiles
+    param_set = RRTMGPParameters(FT)
+    otp = GrayOpticalThicknessSchneider2004(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
+    pressure = CenterField(grid)
+    temperature = CenterField(grid)
+    set!(pressure, p_lay_as_array)
+    set!(temperature, t_lay_as_array)
+
+    return pressure, temperature
+end
+
+