Add WeatherModel initial conditions for prognostic EDMF with support for

costachris · costachris · commit be82aa33b92b · 2025-09-23T10:31:59.000-07:00
water species. Interpolate IC to arbitrary vertical grid. Generalize
paths for IC.
diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml
@@ -172,6 +172,9 @@ bubble:
 start_date:
   help: "Start date and time of the simulation. Specified as either yyyymmdd (defaults to midnight) or yyyymmdd-HHMM. Examples: [`20100101`, `20100101-0000`]"
   value: "20100101"
+era5_initial_condition_dir:
+  help: "Directory containing ERA5 initial condition files. Filenames inferred from start_date [none (default)]. Generated with `https://github.com/CliMA/WeatherQuest`."
+  value: ~
 forcing:
   help: "Forcing [`nothing` (default), `held_suarez`]"
   value: ~
diff --git a/config/model_configs/diagnostic_edmfx_era5_initial_condition.yml b/config/model_configs/diagnostic_edmfx_era5_initial_condition.yml
@@ -15,7 +15,6 @@ hyperdiff: ClimaHyperdiffusion
 divergence_damping_factor: 50.0 # large value to damp IC waves
 scalar_hyperdiffusion_coefficient: 0.929738
 vorticity_hyperdiffusion_coefficient: 0.1857
-
 insolation: "timevarying"
 surface_setup: DefaultMoninObukhov
 rad: allskywithclear
@@ -30,14 +29,14 @@ edmfx_detr_model: "Generalized"
 edmfx_nh_pressure: true
 edmfx_sgs_mass_flux: true
 edmfx_sgs_diffusive_flux: true
-moist: equil
-cloud_model: "grid_scale"
+cloud_model: "quadrature_sgs"
 precip_model: 0M
+moist: equil
 toml: [toml/diagnostic_edmfx_era5_ic.toml]
 use_itime: true
 output_default_diagnostics: false
 diagnostics:
-  - short_name: [ta, ua, wa, va, rhoa, hur, hus, clw, cli]
+  - short_name: [ts, ta, thetaa, ha, pfull, rhoa, ua, va, wa, hur, hus, cl, clw, cli, hussfc, evspsbl, pr]
     period: 60mins
   - short_name: [rsut, rlut, ts, pr, mslp]
     period: 60mins
diff --git a/src/initial_conditions/initial_conditions.jl b/src/initial_conditions/initial_conditions.jl
@@ -206,6 +206,7 @@ the artifact path.
 """
 struct WeatherModel <: InitialCondition
     start_date::String
+    era5_initial_condition_dir::Union{Nothing, String}
 end
 
 function (initial_condition::Union{MoistFromFile, WeatherModel})(params)
@@ -214,7 +215,9 @@ function (initial_condition::Union{MoistFromFile, WeatherModel})(params)
         grav = CAP.grav(params)
         thermo_params = CAP.thermodynamics_params(params)
 
-        T, p = FT(NaN), FT(NaN) # placeholder values
+        # Use safe defaults to avoid transient NaNs before overwrite
+        T = FT(CAP.T_0(params))
+        p = FT(CAP.MSLP(params))
 
         return LocalState(;
             params,
@@ -377,7 +380,24 @@ function overwrite_initial_conditions!(
     thermo_params,
 )
     extrapolation_bc = (Intp.Periodic(), Intp.Flat(), Intp.Flat())
-    file_path = weather_model_data_path(initial_condition.start_date)
+
+    # Extract face coordinates and compute center midpoints
+    z_coords = Fields.axes(Y.c).grid.vertical_grid.topology.mesh.faces
+
+    # Compute center coordinates as midpoints between faces
+    face_z_values = [z.z for z in z_coords]
+    center_z_values = [
+        (face_z_values[i] + face_z_values[i + 1]) / 2 for
+        i in 1:(length(face_z_values) - 1)
+    ]
+
+    target_levels = Array(center_z_values)
+
+    file_path = weather_model_data_path(
+        initial_condition.start_date,
+        target_levels,
+        initial_condition.era5_initial_condition_dir,
+    )
     return _overwrite_initial_conditions_from_file!(
         file_path,
         extrapolation_bc,
@@ -454,10 +474,9 @@ function _overwrite_initial_conditions_from_file!(
     ᶠlnp_over_psfc = zeros(face_space)
     Operators.column_integral_indefinite!(ᶠlnp_over_psfc, ᶜ∂lnp∂z)
     ᶠp = p_sfc .* exp.(ᶠlnp_over_psfc)
-    ᶜts = TD.PhaseEquil_pTq.(thermo_params, ᶜinterp.(ᶠp), ᶜT, ᶜq_tot)
+    # Compute center pressure for thermodynamic state construction
+    ᶜp = ᶜinterp.(ᶠp)
 
-    # Assign prognostic variables from equilibrium moisture models
-    Y.c.ρ .= TD.air_density.(thermo_params, ᶜts)
     # Velocity is first assigned on cell-centers and then interpolated onto
     # cell faces.
     vel =
@@ -483,36 +502,125 @@ function _overwrite_initial_conditions_from_file!(
         )
     Y.c.uₕ .= C12.(Geometry.UVVector.(vel))
     Y.f.u₃ .= ᶠinterp.(C3.(Geometry.WVector.(vel)))
+
+    # Load cloud/precip water fields if present (fallback to zeros)
+    has_liq_ice = hasproperty(Y.c, :ρq_liq) && hasproperty(Y.c, :ρq_ice)
+    has_precip = hasproperty(Y.c, :ρq_rai) && hasproperty(Y.c, :ρq_sno)
+    # TEMP: diagnose 1M condensed/precip species at runtime (initialize to zero here)
+    diagnose_1m = true
+
+    ᶜq_liq = zeros(center_space)
+    ᶜq_ice = zeros(center_space)
+    ᶜq_rai = zeros(center_space)
+    ᶜq_sno = zeros(center_space)
+
+    if has_liq_ice && !diagnose_1m
+        try
+            ᶜq_liq = SpaceVaryingInputs.SpaceVaryingInput(
+                file_path,
+                "clwc",
+                center_space,
+                regridder_kwargs = regridder_kwargs,
+            )
+        catch
+            @info "Variable clwc not found in $(file_path); initializing liquid cloud to zero"
+        end
+        try
+            ᶜq_ice = SpaceVaryingInputs.SpaceVaryingInput(
+                file_path,
+                "ciwc",
+                center_space,
+                regridder_kwargs = regridder_kwargs,
+            )
+        catch
+            @info "Variable ciwc not found in $(file_path); initializing ice cloud to zero"
+        end
+    end
+    if has_precip && !diagnose_1m
+        try
+            ᶜq_rai = SpaceVaryingInputs.SpaceVaryingInput(
+                file_path,
+                "crwc",
+                center_space,
+                regridder_kwargs = regridder_kwargs,
+            )
+        catch
+            @info "Variable crwc not found in $(file_path); initializing rain water to zero"
+        end
+        try
+            ᶜq_sno = SpaceVaryingInputs.SpaceVaryingInput(
+                file_path,
+                "cswc",
+                center_space,
+                regridder_kwargs = regridder_kwargs,
+            )
+        catch
+            @info "Variable cswc not found in $(file_path); initializing snow water to zero"
+        end
+    end
+
+    # Build thermodynamic state consistently with moisture partition
+    is_nonequil = hasproperty(Y.c, :ρq_liq) || hasproperty(Y.c, :ρq_ice)
+    # Use vapor from ERA5 q as first component; add cloud+precip to condensed partitions (zeros if not provided)
+    ᶜq_vap = max.(ᶜq_tot, zero.(ᶜq_tot))
+    ᶜq_liq .= ifelse.(isfinite.(ᶜq_liq), ᶜq_liq, zero.(ᶜq_liq))
+    ᶜq_liq .= max.(ᶜq_liq, zero.(ᶜq_liq))
+    ᶜq_ice .= ifelse.(isfinite.(ᶜq_ice), ᶜq_ice, zero.(ᶜq_ice))
+    ᶜq_ice .= max.(ᶜq_ice, zero.(ᶜq_ice))
+    ᶜq_rai .= ifelse.(isfinite.(ᶜq_rai), ᶜq_rai, zero.(ᶜq_rai))
+    ᶜq_rai .= max.(ᶜq_rai, zero.(ᶜq_rai))
+    ᶜq_sno .= ifelse.(isfinite.(ᶜq_sno), ᶜq_sno, zero.(ᶜq_sno))
+    ᶜq_sno .= max.(ᶜq_sno, zero.(ᶜq_sno))
+    ᶜpartition = TD.PhasePartition.(ᶜq_vap, ᶜq_liq .+ ᶜq_rai, ᶜq_ice .+ ᶜq_sno)
+    if is_nonequil
+        ᶜts = TD.PhaseNonEquil_pTq.(thermo_params, ᶜp, ᶜT, ᶜpartition)
+    else
+        ᶜts = TD.PhaseEquil_pTq.(thermo_params, ᶜp, ᶜT, ᶜq_tot)
+    end
+
+    # Density
+    Y.c.ρ .= TD.air_density.(thermo_params, ᶜts)
+
+    # Kinetic and potential energy
     e_kin = similar(ᶜT)
     e_kin .= compute_kinetic(Y.c.uₕ, Y.f.u₃)
     e_pot = Fields.coordinate_field(Y.c).z .* thermo_params.grav
+
+    # Total energy
     Y.c.ρe_tot .= TD.total_energy.(thermo_params, ᶜts, e_kin, e_pot) .* Y.c.ρ
+
+    # Moisture mass densities
     if hasproperty(Y.c, :ρq_tot)
+        # Keep ρq_tot equal to ERA5 vapor (do not add precipitating species)
         Y.c.ρq_tot .= ᶜq_tot .* Y.c.ρ
     else
         error(
             "`dry` configurations are incompatible with the interpolated initial conditions.",
         )
     end
-    if hasproperty(Y.c, :ρq_sno) && hasproperty(Y.c, :ρq_rai)
-        Y.c.ρq_sno .=
-            SpaceVaryingInputs.SpaceVaryingInput(
-                file_path,
-                "cswc",
-                center_space,
-                regridder_kwargs = regridder_kwargs,
-            ) .* Y.c.ρ
-        Y.c.ρq_rai .=
-            SpaceVaryingInputs.SpaceVaryingInput(
-                file_path,
-                "crwc",
-                center_space,
-                regridder_kwargs = regridder_kwargs,
-            ) .* Y.c.ρ
+
+    if hasproperty(Y.c, :ρq_rai)
+        Y.c.ρq_rai .= ᶜq_rai .* Y.c.ρ
+        Y.c.ρq_sno .= ᶜq_sno .* Y.c.ρ
+    end
+    if hasproperty(Y.c, :ρq_liq)
+        Y.c.ρq_liq .= ᶜq_liq .* Y.c.ρ
+        Y.c.ρq_ice .= ᶜq_ice .* Y.c.ρ
+    end
+
+    # Set prognostic EDMF draft subdomains to consistent values
+    if hasproperty(Y.c, :sgsʲs) # if prognostic EDMF
+        ᶜmse = TD.specific_enthalpy.(thermo_params, ᶜts) .+ e_pot
+        for name in propertynames(Y.c.sgsʲs)
+            s = getproperty(Y.c.sgsʲs, name)
+            hasproperty(s, :ρa) && fill!(s.ρa, 0)
+            hasproperty(s, :mse) && (s.mse .= ᶜmse)
+            hasproperty(s, :q_tot) && (s.q_tot .= ᶜq_vap)
+        end
     end
 
+    # NOTE: This is not the most consistent, but it is better than NaNs
     if hasproperty(Y.c, :sgs⁰) && hasproperty(Y.c.sgs⁰, :ρatke)
-        # NOTE: This is not the most consistent, but it is better than NaNs
         fill!(Y.c.sgs⁰.ρatke, 0)
     end
 
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
@@ -423,7 +423,10 @@ function get_initial_condition(parsed_args, atmos)
             parsed_args["start_date"],
         )
     elseif parsed_args["initial_condition"] == "WeatherModel"
-        return ICs.WeatherModel(parsed_args["start_date"])
+        return ICs.WeatherModel(
+            parsed_args["start_date"],
+            parsed_args["era5_initial_condition_dir"],
+        )
     else
         error(
             "Unknown `initial_condition`: $(parsed_args["initial_condition"])",
diff --git a/src/utils/weather_model.jl b/src/utils/weather_model.jl
@@ -3,23 +3,60 @@ using Dates
 import ClimaInterpolations.Interpolation1D: interpolate1d!, Linear, Flat
 import ..parse_date
 
+
 """
-    weather_model_data_path(start_date)
+    weather_model_data_path(start_date, target_levels, era5_initial_condition_dir])
 
 Get the path to the weather model data for a given start date and time.
 If the data is not found, will attempt to generate it from raw data. If 
 the raw data is not found, throw an error
 
 # Arguments
 - `start_date`: Start date as string yyyymmdd or yyyymmdd-HHMM
+- `target_levels`: Vector of target altitude levels (in meters) to interpolate to
+- `era5_initial_condition_dir`: Optional directory containing preprocessed ERA5
+
+- If `era5_initial_condition_dir` is provided, use
+  `era5_init_processed_internal_YYYYMMDD_0000.nc` from that directory.
+- Otherwise, use the `weather_model_ic` artifact.
 """
-function weather_model_data_path(start_date)
+function weather_model_data_path(
+    start_date,
+    target_levels,
+    era5_initial_condition_dir = nothing,
+)
     # Parse the date using the existing parse_date function
     dt = parse_date(start_date)
 
     # Extract components for filename generation
     start_date_str = Dates.format(dt, "yyyymmdd")
-    start_time = Dates.format(dt, "HHMM")
+    start_time = Dates.format(dt, "HHMM") # Note: this is not the same as `start_time` in the coupler!
+
+    # If user provided a directory with preprocessed initial conditions, use it
+    if !isnothing(era5_initial_condition_dir)
+        ic_data_path = joinpath(
+            era5_initial_condition_dir,
+            "era5_init_processed_internal_$(start_date_str)_0000.nc", # TODO: generalize for all times once Coupler supports HHMM specification
+        )
+        raw_data_path = joinpath(
+            era5_initial_condition_dir,
+            "era5_raw_$(start_date_str)_0000.nc",
+        )
+        if !isfile(ic_data_path)
+            if !isfile(raw_data_path)
+                error(
+                    "Neither preprocessed nor raw initial condition file exist in $(era5_initial_condition_dir).  Please run `python get_initial_conditions.py` in the WeatherQuest repository to download the data.",
+                )
+            end
+            @info "Interpolating raw weather model data onto z-levels from user-provided directory"
+            to_z_levels(raw_data_path, ic_data_path, target_levels, Float32)
+        else
+            @info "Using existing interpolated IC file: $ic_data_path"
+        end
+        return ic_data_path
+    end
+
+    # Otherwise, use artifact-based paths and generate if needed
     ic_data_path = joinpath(
         @clima_artifact("weather_model_ic"),
         "init",
@@ -31,23 +68,6 @@ function weather_model_data_path(start_date)
         "era5_raw_$(start_date_str)_$(start_time).nc",
     )
 
-    if !isfile(ic_data_path)
-        @info "Initial condition file $ic_data_path does not exist. Attempting to generate it now..."
-        if !isfile(raw_data_path)
-            day = Dates.format(dt, "yyyy-mm-dd")
-            time = Dates.format(dt, "HH:MM")
-            error(
-                "Source file $(raw_data_path) does not exist. Please run `python get_initial_conditions.py --output-dir $(@clima_artifact("weather_model_ic"))/raw --date $(day) --time $(time)` to download the data.",
-            )
-        end
-        # get target levels - TODO: make this more flexible
-        target_level_path =
-            joinpath(@clima_artifact("weather_model_ic"), "target_levels.txt")
-        target_levels = parse.(Float32, readlines(target_level_path))
-
-        @info "Interpolating raw weather model data onto z-levels"
-        to_z_levels(raw_data_path, ic_data_path, target_levels, Float32)
-    end
     return ic_data_path
 end
 
@@ -74,6 +94,7 @@ function to_z_levels(source_file, target_file, target_levels, FT)
 
     # assert ncin has required variables
     req_vars = ["u", "v", "w", "t", "q", "skt", "sp"]
+    opt_vars = ["crwc", "cswc", "clwc", "ciwc"]
     @assert all(map(x -> x in (keys(ncin)), req_vars)) "Source file $source_file is missing subset of the required variables: $req_vars"
 
     # Read and cast coordinates to FT type
@@ -168,6 +189,22 @@ function to_z_levels(source_file, target_file, target_levels, FT)
         sp_var[:, :, k] = FT.(ncin["sp"][:, :, 1])
     end
 
+    # Interpolate optional cloud water content variables if available
+    for var_name in opt_vars
+        if haskey(ncin, var_name)
+            @info "Interpolating optional variable: $var_name"
+            var_data = ncin[var_name][:, :, :, 1]
+            var_var = defVar(
+                ncout,
+                var_name,
+                FT,
+                ("lon", "lat", "z"),
+                attrib = ncin[var_name].attrib,
+            )
+            var_var[:, :, :] = interpz_3d(target_levels, source_z, FT.(var_data))
+        end
+    end
+
     # Close files
     close(ncin)
     close(ncout)
