adds era5 driven to pr

hpc edits and formatter

fix naming

.buildkite/pipeline.yml

@@ -831,6 +831,15 @@ steps:
           slurm_mem: 20GB
         soft_fail: true
 
+      - label: ":genie: Prognostic EDMFX ERA5 driven in a column"
+        command: >
+          julia --color=yes --project=examples examples/hybrid/driver.jl
+          --config_file $CONFIG_PATH/prognostic_edmfx_era5driven_column.yml
+          --job_id prognostic_edmfx_era5driven_column
+        artifact_paths: "prognostic_edmfx_era5driven_column/output_active/*"
+        agents:
+          slurm_mem: 20GB
+
       - label: ":genie: Prognostic EDMFX Bomex in a box"
         command: >
           julia --color=yes --project=examples examples/hybrid/driver.jl

config/model_configs/prognostic_edmfx_era5driven_column.yml

@@ -0,0 +1,50 @@
+initial_condition: "ERA5"
+external_forcing: "ERA5"
+external_forcing_file: "/central/groups/esm/jschmitt/era5/forcing/era5_monthly_forcing_1.nc"
+cfsite_number: "site23"
+surface_setup: "ERA5"
+turbconv: "prognostic_edmfx"
+implicit_diffusion: true
+implicit_sgs_advection: false
+approximate_linear_solve_iters: 2
+edmfx_upwinding: first_order
+rayleigh_sponge: true
+edmfx_entr_model: "PiGroups"
+edmfx_detr_model: "PiGroups"
+edmfx_sgs_mass_flux: true
+edmfx_sgs_diffusive_flux: true
+edmfx_nh_pressure: true
+edmfx_filter: true
+prognostic_tke: true
+moist: "equil"
+config: "column"
+z_max: 70e3
+truncation: 40000
+z_elem: 200
+z_stretch: true
+dz_bottom: 30
+perturb_initstate: false
+dt: "10secs"
+dt_rad: "30mins"
+t_end: "8hours"
+cloud_model: "quadrature_sgs"
+call_cloud_diagnostics_per_stage : true
+toml: [toml/prognostic_edmfx_era5driven_calibrated.toml]
+netcdf_output_at_levels: true
+netcdf_interpolation_num_points: [2, 2, 60]
+output_default_diagnostics: false
+rad: allskywithclear
+insolation: "era5driven"
+diagnostics:
+  - short_name: [ts, ta, thetaa, ha, pfull, rhoa, ua, va, wa, hur, hus, cl, clw, cli, hussfc, evspsbl, pr]
+    period: 10mins
+  - short_name: [arup, waup, taup, thetaaup, haup, husup, hurup, clwup, cliup, waen, taen, thetaaen, haen, husen, huren, clwen, clien, tke]
+    period: 10mins
+  - short_name: [entr, detr, lmix, bgrad, strain, edt, evu]
+    period: 10mins
+  - short_name: [rlut, rlutcs, rsut, rsutcs, clwvi, lwp, clivi, dsevi, clvi, prw, hurvi, husv]
+    period: 10mins
+  - reduction_time: max
+    short_name: tke
+    period: 10mins
+ode_algo: ARS343

src/callbacks/callbacks.jl

@@ -288,7 +288,12 @@ function set_insolation_variables!(Y, p, t, ::RCEMIPIIInsolation)
     rrtmgp_model.weighted_irradiance .= FT(551.58)
 end
 
-function set_insolation_variables!(Y, p, t, ::GCMDrivenInsolation)
+function set_insolation_variables!(
+    Y,
+    p,
+    t,
+    ::Union{GCMDrivenInsolation, ERA5DrivenInsolation},
+)
     (; rrtmgp_model) = p.radiation
     rrtmgp_model.cos_zenith .= Fields.field2array(p.external_forcing.cos_zenith)
     rrtmgp_model.weighted_irradiance .=

src/initial_conditions/InitialConditions.jl

@@ -20,6 +20,8 @@ import ..n_mass_flux_subdomains
 import ..gcm_driven_profile
 import ..gcm_height
 import ..gcm_driven_profile_tmean
+import ..era5_driven_profile
+import ..era5_height
 
 import Thermodynamics.TemperatureProfiles:
     DecayingTemperatureProfile, DryAdiabaticProfile

src/initial_conditions/initial_conditions.jl

@@ -1331,6 +1331,62 @@ function gcm_initial_conditions(external_forcing_file, cfsite_number)
     end
 end
 
+"""
+    ERA5Driven <: InitialCondition
+The `InitialCondition` from a provided ERA5 forcing file, with data type `DType`.
+"""
+struct ERA5Driven <: InitialCondition
+    external_forcing_file::String
+    cfsite_number::String
+end
+
+function (initial_condition::ERA5Driven)(params)
+    (; external_forcing_file, cfsite_number) = initial_condition
+    thermo_params = CAP.thermodynamics_params(params)
+
+    # Read forcing file
+    z_era5 = NC.NCDataset(external_forcing_file) do ds
+        vec(era5_height(ds.group[cfsite_number]))
+    end
+    vars = era5_initial_conditions(external_forcing_file, cfsite_number)
+    T, u, v, q_tot, ρ₀ = map(vars) do value
+        Intp.extrapolate(
+            Intp.interpolate((z_era5,), value, Intp.Gridded(Intp.Linear())),
+            Intp.Flat(),
+        )
+    end
+
+    function local_state(local_geometry)
+        (; z) = local_geometry.coordinates
+        FT = typeof(z)
+        return LocalState(;
+            params,
+            geometry = local_geometry,
+            thermo_state = ts = TD.PhaseEquil_ρTq(
+                thermo_params,
+                FT(ρ₀(z)),
+                FT(T(z)),
+                FT(q_tot(z)),
+            ),
+            velocity = Geometry.UVVector(FT(u(z)), FT(v(z))),
+            turbconv_state = EDMFState(; tke = FT(0)),
+        )
+    end
+    return local_state
+end
+
+function era5_initial_conditions(external_forcing_file, cfsite_number)
+    NC.NCDataset(external_forcing_file) do ds
+        (  # TODO: Cast to CuVector for GPU compatibility
+            era5_driven_profile(ds.group[cfsite_number], "ta"),
+            era5_driven_profile(ds.group[cfsite_number], "ua"),
+            era5_driven_profile(ds.group[cfsite_number], "va"),
+            era5_driven_profile(ds.group[cfsite_number], "hus"),
+            ds.group[cfsite_number]["rho"][:], # convert alpha to rho using rho=1/alpha, take average profile
+        )
+    end
+end
+
 Base.@kwdef struct ISDAC <: InitialCondition
     prognostic_tke::Bool = false
     perturb::Bool = false

src/prognostic_equations/forcing/external_forcing.jl

@@ -297,3 +297,200 @@ function external_forcing_tendency!(Yₜ, Y, p, t, ::ISDACForcing)
     # total specific humidity
     @. Yₜ.c.ρq_tot += Y.c.ρ * ᶜdqtdt_nudging
 end
+
+# For ERA5 reanalysis forcing data
+function external_forcing_cache(Y, external_forcing::ERA5Forcing, params)
+    FT = Spaces.undertype(axes(Y.c))
+    ᶜdTdt_fluc = similar(Y.c, FT)
+    ᶜdqtdt_fluc = similar(Y.c, FT)
+    ᶜdTdt_hadv = similar(Y.c, FT)
+    ᶜdqtdt_hadv = similar(Y.c, FT)
+    ᶜdTdt_rad = similar(Y.c, FT)
+    ᶜT_nudge = similar(Y.c, FT)
+    ᶜqt_nudge = similar(Y.c, FT)
+    ᶜu_nudge = similar(Y.c, FT)
+    ᶜv_nudge = similar(Y.c, FT)
+    ᶜinv_τ_wind = similar(Y.c, FT)
+    ᶜinv_τ_scalar = similar(Y.c, FT)
+    ᶜls_subsidence = similar(Y.c, FT)
+    insolation = similar(Fields.level(Y.c.ρ, 1), FT)
+    cos_zenith = similar(Fields.level(Y.c.ρ, 1), FT)
+
+    (; external_forcing_file, cfsite_number) = external_forcing
+
+    NC.Dataset(external_forcing_file, "r") do ds
+
+        function setvar!(cc_field, varname, colidx, zc_gcm, zc_forcing)
+            parent(cc_field[colidx]) .= interp_vertical_prof(
+                zc_gcm,
+                zc_forcing,
+                era5_driven_profile(ds.group[cfsite_number], varname),
+            )
+        end
+
+        function setvar_subsidence!(
+            cc_field,
+            varname,
+            colidx,
+            zc_gcm,
+            zc_forcing,
+            params,
+        )
+            parent(cc_field[colidx]) .= interp_vertical_prof(
+                zc_gcm,
+                zc_forcing,
+                era5_driven_profile(ds.group[cfsite_number], varname) ./
+                .-(era5_driven_profile(ds.group[cfsite_number], "rho")) ./
+                CAP.grav(params),
+            )
+        end
+
+        function set_insolation!(cc_field)
+            parent(cc_field) .= mean(
+                ds.group[cfsite_number]["rsdt"][:] ./
+                ds.group[cfsite_number]["coszen"][:],
+            )
+        end
+
+        function set_cos_zenith!(cc_field)
+            parent(cc_field) .= ds.group[cfsite_number]["coszen"][1]
+        end
+
+        zc_forcing = era5_height(ds.group[cfsite_number])
+        Fields.bycolumn(axes(Y.c)) do colidx
+
+            zc_gcm = Fields.coordinate_field(Y.c).z[colidx]
+
+            setvar!(ᶜdTdt_hadv, "tntha", colidx, zc_gcm, zc_forcing)
+            setvar!(ᶜdqtdt_hadv, "tnhusha", colidx, zc_gcm, zc_forcing)
+            #setvar!(ᶜdTdt_rad, "tntr", colidx, zc_gcm, zc_forcing)
+            setvar_subsidence!(
+                ᶜls_subsidence,
+                "wap",
+                colidx,
+                zc_gcm,
+                zc_forcing,
+                params,
+            )
+            # GCM states, used for nudging + vertical eddy advection
+            setvar!(ᶜT_nudge, "ta", colidx, zc_gcm, zc_forcing)
+            setvar!(ᶜqt_nudge, "hus", colidx, zc_gcm, zc_forcing)
+            setvar!(ᶜu_nudge, "ua", colidx, zc_gcm, zc_forcing)
+            setvar!(ᶜv_nudge, "va", colidx, zc_gcm, zc_forcing)
+
+            # vertical eddy advection (Shen et al., 2022; eqn. 9,10)
+            # sum of two terms to give total tendency. First term:
+            setvar!(ᶜdTdt_fluc, "tntva", colidx, zc_gcm, zc_forcing)
+            setvar!(ᶜdqtdt_fluc, "tnhusva", colidx, zc_gcm, zc_forcing)
+            # second term:
+            eddy_vert_fluctuation!(ᶜdTdt_fluc, ᶜT_nudge, ᶜls_subsidence)
+            eddy_vert_fluctuation!(ᶜdqtdt_fluc, ᶜqt_nudge, ᶜls_subsidence)
+
+            set_insolation!(insolation)
+            set_cos_zenith!(cos_zenith)
+
+            @. ᶜinv_τ_wind[colidx] = 1 / (6 * 3600)
+            @. ᶜinv_τ_scalar[colidx] = compute_gcm_driven_scalar_inv_τ(zc_gcm)
+        end
+    end
+
+    return (;
+        ᶜdTdt_fluc,
+        ᶜdqtdt_fluc,
+        ᶜdTdt_hadv,
+        ᶜdqtdt_hadv,
+        #ᶜdTdt_rad,
+        ᶜT_nudge,
+        ᶜqt_nudge,
+        ᶜu_nudge,
+        ᶜv_nudge,
+        ᶜinv_τ_wind,
+        ᶜinv_τ_scalar,
+        ᶜls_subsidence,
+        insolation,
+        cos_zenith,
+    )
+end
+
+function external_forcing_tendency!(Yₜ, Y, p, t, ::ERA5Forcing)
+    # horizontal advection, vertical fluctuation, nudging, subsidence (need to add),
+    (; params) = p
+    thermo_params = CAP.thermodynamics_params(params)
+    (; ᶜspecific, ᶜts, ᶜh_tot) = p.precomputed
+    (;
+        ᶜdTdt_fluc,
+        ᶜdqtdt_fluc,
+        ᶜdTdt_hadv,
+        ᶜdqtdt_hadv,
+        #ᶜdTdt_rad,
+        ᶜT_nudge,
+        ᶜqt_nudge,
+        ᶜu_nudge,
+        ᶜv_nudge,
+        ᶜls_subsidence,
+        ᶜinv_τ_wind,
+        ᶜinv_τ_scalar,
+    ) = p.external_forcing
+
+    ᶜlg = Fields.local_geometry_field(Y.c)
+    ᶜuₕ_nudge = p.scratch.ᶜtemp_C12
+    @. ᶜuₕ_nudge = C12(Geometry.UVVector(ᶜu_nudge, ᶜv_nudge), ᶜlg)
+    @. Yₜ.c.uₕ -= (Y.c.uₕ - ᶜuₕ_nudge) * ᶜinv_τ_wind
+
+    # nudging tendency
+    ᶜdTdt_nudging = p.scratch.ᶜtemp_scalar
+    ᶜdqtdt_nudging = p.scratch.ᶜtemp_scalar_2
+    @. ᶜdTdt_nudging =
+        -(TD.air_temperature(thermo_params, ᶜts) - ᶜT_nudge) * ᶜinv_τ_scalar
+    @. ᶜdqtdt_nudging = -(ᶜspecific.q_tot - ᶜqt_nudge) * ᶜinv_τ_scalar
+
+    ᶜdTdt_sum = p.scratch.ᶜtemp_scalar
+    ᶜdqtdt_sum = p.scratch.ᶜtemp_scalar_2
+    @. ᶜdTdt_sum = ᶜdTdt_hadv + ᶜdTdt_nudging + ᶜdTdt_fluc
+    @. ᶜdqtdt_sum = ᶜdqtdt_hadv + ᶜdqtdt_nudging + ᶜdqtdt_fluc
+
+    T_0 = TD.Parameters.T_0(thermo_params)
+    Lv_0 = TD.Parameters.LH_v0(thermo_params)
+    cv_v = TD.Parameters.cv_v(thermo_params)
+    R_v = TD.Parameters.R_v(thermo_params)
+    # total energy
+    @. Yₜ.c.ρe_tot +=
+        Y.c.ρ * (
+            TD.cv_m(thermo_params, ᶜts) * ᶜdTdt_sum +
+            (
+                cv_v * (TD.air_temperature(thermo_params, ᶜts) - T_0) + Lv_0 -
+                R_v * T_0
+            ) * ᶜdqtdt_sum
+        )
+    # total specific humidity
+    @. Yₜ.c.ρq_tot += Y.c.ρ * ᶜdqtdt_sum
+
+    ## subsidence -->
+    ᶠls_subsidence³ = p.scratch.ᶠtemp_CT3
+    @. ᶠls_subsidence³ =
+        ᶠinterp(ᶜls_subsidence * CT3(unit_basis_vector_data(CT3, ᶜlg)))
+    subsidence!(
+        Yₜ.c.ρe_tot,
+        Y.c.ρ,
+        ᶠls_subsidence³,
+        ᶜh_tot,
+        Val{:first_order}(),
+    )
+    subsidence!(
+        Yₜ.c.ρq_tot,
+        Y.c.ρ,
+        ᶠls_subsidence³,
+        ᶜspecific.q_tot,
+        Val{:first_order}(),
+    )
+
+    # needed to address top boundary condition for forcings. Otherwise upper portion of domain is anomalously cold
+    ρe_tot_top = Fields.level(Yₜ.c.ρe_tot, Spaces.nlevels(axes(Y.c)))
+    @. ρe_tot_top = 0.0
+
+    ρq_tot_top = Fields.level(Yₜ.c.ρq_tot, Spaces.nlevels(axes(Y.c)))
+    @. ρq_tot_top = 0.0
+    # <-- subsidence
+
+    return nothing
+end

src/solver/model_getters.jl

@@ -25,7 +25,8 @@ end
 
 function get_insolation_form(parsed_args)
     insolation = parsed_args["insolation"]
-    @assert insolation in ("idealized", "timevarying", "rcemipii", "gcmdriven")
+    @assert insolation in
+            ("idealized", "timevarying", "rcemipii", "gcmdriven", "era5driven")
     return if insolation == "idealized"
         IdealizedInsolation()
     elseif insolation == "timevarying"
@@ -37,6 +38,8 @@ function get_insolation_form(parsed_args)
         RCEMIPIIInsolation()
     elseif insolation == "gcmdriven"
         GCMDrivenInsolation()
+    elseif insolation == "era5driven"
+        ERA5DrivenInsolation()
     end
 end
 
@@ -407,7 +410,7 @@ end
 
 function get_external_forcing_model(parsed_args)
     external_forcing = parsed_args["external_forcing"]
-    @assert external_forcing in (nothing, "GCM", "ISDAC")
+    @assert external_forcing in (nothing, "GCM", "ISDAC", "ERA5")
     return if isnothing(external_forcing)
         nothing
     elseif external_forcing == "GCM"
@@ -416,6 +419,12 @@ function get_external_forcing_model(parsed_args)
             parsed_args["external_forcing_file"],
             parsed_args["cfsite_number"],
         )
+    elseif external_forcing == "ERA5"
+        DType = Float64  # TODO: Read from `parsed_args`
+        ERA5Forcing{DType}(
+            parsed_args["external_forcing_file"],
+            parsed_args["cfsite_number"],
+        )
     elseif external_forcing == "ISDAC"
         ISDACForcing()
     end