Merge pull request #4022 from CliMA/cc/test_wxquest_ic

Update WeatherModel initial conditions

Author: costachris

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: ~

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]
     period: 60mins
-  - short_name: [rsut, rlut, ts, pr, mslp]
+  - short_name: [rsut, rlut, tas, pr, mslp]
     period: 60mins

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)
@@ -377,13 +378,139 @@ function overwrite_initial_conditions!(
     thermo_params,
 )
     extrapolation_bc = (Intp.Periodic(), Intp.Flat(), Intp.Flat())
-    file_path = weather_model_data_path(initial_condition.start_date)
-    return _overwrite_initial_conditions_from_file!(
+
+    # Extract face coordinates and compute center midpoints
+    # Compute target levels on CPU to avoid GPU reductions
+    z_arr_cpu = Array(Fields.field2array(Fields.coordinate_field(Y.c).z))
+    icol = argmin(z_arr_cpu[1, :])
+    target_levels = z_arr_cpu[:, icol]
+
+    file_path = weather_model_data_path(
+        initial_condition.start_date,
+        target_levels,
+        initial_condition.era5_initial_condition_dir,
+    )
+
+    regridder_kwargs = (; extrapolation_bc)
+    isfile(file_path) || error("$(file_path) is not a file")
+    @info "Overwriting initial conditions with data from file $(file_path)"
+
+    center_space = Fields.axes(Y.c)
+    face_space = Fields.axes(Y.f)
+
+    # Using surface pressure, air temperature and specific humidity
+    # from the dataset, compute air pressure.
+    p_sfc = Fields.level(
+        SpaceVaryingInputs.SpaceVaryingInput(
+            file_path,
+            "p",
+            face_space,
+            regridder_kwargs = regridder_kwargs,
+        ),
+        Fields.half,
+    )
+    ᶜT = SpaceVaryingInputs.SpaceVaryingInput(
         file_path,
-        extrapolation_bc,
-        Y,
-        thermo_params,
+        "t",
+        center_space,
+        regridder_kwargs = regridder_kwargs,
+    )
+    ᶜq_tot = SpaceVaryingInputs.SpaceVaryingInput(
+        file_path,
+        "q",
+        center_space,
+        regridder_kwargs = regridder_kwargs,
     )
+
+    # With the known temperature (ᶜT) and moisture (ᶜq_tot) profile,
+    # recompute the pressure levels assuming hydrostatic balance is maintained.
+    # Uses the ClimaCore `column_integral_indefinite!` function to solve
+    # ∂(ln𝑝)/∂z = -g/(Rₘ(q)T), where
+    # p is the local pressure
+    # g is the gravitational constant
+    # q is the specific humidity
+    # Rₘ is the gas constant for moist air
+    # T is the air temperature
+    # p is then updated with the integral result, given p_sfc,
+    # following which the thermodynamic state is constructed.
+    ᶜ∂lnp∂z = @. -thermo_params.grav /
+       (TD.gas_constant_air(thermo_params, TD.PhasePartition(ᶜq_tot)) * ᶜT)
+    ᶠ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)
+
+    # 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 =
+        Geometry.UVWVector.(
+            SpaceVaryingInputs.SpaceVaryingInput(
+                file_path,
+                "u",
+                center_space,
+                regridder_kwargs = regridder_kwargs,
+            ),
+            SpaceVaryingInputs.SpaceVaryingInput(
+                file_path,
+                "v",
+                center_space,
+                regridder_kwargs = regridder_kwargs,
+            ),
+            SpaceVaryingInputs.SpaceVaryingInput(
+                file_path,
+                "w",
+                center_space,
+                regridder_kwargs = regridder_kwargs,
+            ),
+        )
+    Y.c.uₕ .= C12.(Geometry.UVVector.(vel))
+    Y.f.u₃ .= ᶠinterp.(C3.(Geometry.WVector.(vel)))
+    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
+    Y.c.ρe_tot .= TD.total_energy.(thermo_params, ᶜts, e_kin, e_pot) .* Y.c.ρ
+    # Initialize prognostic EDMF 0M subdomains if present
+    if hasproperty(Y.c, :sgsʲs)
+        ᶜ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_tot)
+        end
+    end
+    if hasproperty(Y.c, :ρq_tot)
+        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.ρ
+    end
+
+    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
+
+    return nothing
 end
 
 """

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"])",

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
@@ -118,54 +139,49 @@ function to_z_levels(source_file, target_file, target_levels, FT)
     z_var = defVar(ncout, "z", FT, ("z",), attrib = z_attrib)
     z_var[:] = target_levels
 
-    # Interpolate and write 3D variables
-    u_var =
-        defVar(ncout, "u", FT, ("lon", "lat", "z"), attrib = ncin["u"].attrib)
-    u_var[:, :, :] =
-        interpz_3d(target_levels, source_z, FT.(ncin["u"][:, :, :, 1]))
-
-    v_var =
-        defVar(ncout, "v", FT, ("lon", "lat", "z"), attrib = ncin["v"].attrib)
-    v_var[:, :, :] =
-        interpz_3d(target_levels, source_z, FT.(ncin["v"][:, :, :, 1]))
-
+    # Interpolate and write required 3D variables via loop
     # ERA5 w is from a hydrostatic model and so isn't meaningful for ClimaAtmos
     # See https://agupubs.onlinelibrary.wiley.com/doi/full/10.1002/2017MS001059
-    w_var =
-        defVar(ncout, "w", FT, ("lon", "lat", "z"), attrib = ncin["w"].attrib)
-    w_var[:, :, :] = zeros(FT, length(lon), length(lat), length(target_levels))
-
-    t_var =
-        defVar(ncout, "t", FT, ("lon", "lat", "z"), attrib = ncin["t"].attrib)
-    t_var[:, :, :] =
-        interpz_3d(target_levels, source_z, FT.(ncin["t"][:, :, :, 1]))
-
-    q_var =
-        defVar(ncout, "q", FT, ("lon", "lat", "z"), attrib = ncin["q"].attrib)
-    q_var[:, :, :] =
-        max.(
-            interpz_3d(target_levels, source_z, FT.(ncin["q"][:, :, :, 1])),
-            FT(0),
-        )
+    req3d = ["u", "v", "t", "q", "w"]
+    for var_name in req3d
+        var_obj =
+            defVar(ncout, var_name, FT, ("lon", "lat", "z"), attrib = ncin[var_name].attrib)
+        if var_name == "w"
+            var_obj[:, :, :] = zeros(FT, length(lon), length(lat), length(target_levels))
+        else
+            data = interpz_3d(target_levels, source_z, FT.(ncin[var_name][:, :, :, 1]))
+            if var_name == "q"
+                data = max.(data, FT(0))
+            end
+            var_obj[:, :, :] = data
+        end
+    end
 
     # Write 2D surface variables - extend to all levels (TODO: accept 2D variables in atmos)
     # Simply repeat the surface values for all levels
-    skt_var = defVar(
-        ncout,
-        "skt",
-        FT,
-        ("lon", "lat", "z"),
-        attrib = ncin["skt"].attrib,
-    )
-    for k in 1:length(target_levels)
-        skt_var[:, :, k] = FT.(ncin["skt"][:, :, 1])
+    surf_map = Dict("skt" => "skt", "sp" => "p")
+    for (src_name, dst_name) in surf_map
+        var_obj =
+            defVar(ncout, dst_name, FT, ("lon", "lat", "z"), attrib = ncin[src_name].attrib)
+        for k in 1:length(target_levels)
+            var_obj[:, :, k] = FT.(ncin[src_name][:, :, 1])
+        end
     end
 
-    # TODO: rename p to sp in MoistFromFile for consistency
-    sp_var =
-        defVar(ncout, "p", FT, ("lon", "lat", "z"), attrib = ncin["sp"].attrib)
-    for k in 1:length(target_levels)
-        sp_var[:, :, k] = FT.(ncin["sp"][:, :, 1])
+    # 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