Merge pull request #3883 from CliMA/j/ecmwf_init

Julians42 · web-flow · commit 96845801aaf4 · 2025-07-09T19:11:45.000Z
Add WeatherQuest Atmos IC
diff --git a/.buildkite/Manifest-v1.11.toml b/.buildkite/Manifest-v1.11.toml
@@ -363,7 +363,7 @@ version = "0.5.18"
     Makie = "ee78f7c6-11fb-53f2-987a-cfe4a2b5a57a"
 
 [[deps.ClimaAtmos]]
-deps = ["Adapt", "ArgParse", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "ForwardDiff", "Insolation", "Interpolations", "LazyArtifacts", "LazyBroadcast", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "NullBroadcasts", "RRTMGP", "Random", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
+deps = ["Adapt", "ArgParse", "Artifacts", "AtmosphericProfilesLibrary", "ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaInterpolations", "ClimaParams", "ClimaTimeSteppers", "ClimaUtilities", "CloudMicrophysics", "Dates", "ForwardDiff", "Insolation", "Interpolations", "LazyArtifacts", "LazyBroadcast", "LinearAlgebra", "Logging", "NCDatasets", "NVTX", "NullBroadcasts", "RRTMGP", "Random", "SciMLBase", "StaticArrays", "Statistics", "SurfaceFluxes", "Thermodynamics", "UnrolledUtilities", "YAML"]
 path = ".."
 uuid = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 version = "0.30.4"
@@ -408,6 +408,16 @@ git-tree-sha1 = "5c031f4fc0c044e3da65c7cbc173962cc1489074"
 uuid = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f"
 version = "0.2.15"
 
+[[deps.ClimaInterpolations]]
+deps = ["DocStringExtensions"]
+git-tree-sha1 = "e62662cfbf368f7d9cdb1876dcb6e328437cb676"
+uuid = "dd0f122e-fa3b-47f3-bcf0-93bbc60d885e"
+version = "0.1.0"
+weakdeps = ["CUDA"]
+
+    [deps.ClimaInterpolations.extensions]
+    ClimaInterpolationsCUDAExt = "CUDA"
+
 [[deps.ClimaParams]]
 deps = ["TOML"]
 git-tree-sha1 = "8f76fb84975762ac8e12fadcb87992b2a92e3966"
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -1035,6 +1035,18 @@ steps:
         agents:
           slurm_gpus: 1
           slurm_mem: 16GB
+      
+      - label: ":earth_americas: GPU: Diagnostic EDMFX ERA5 Weather Model Initial Condition"
+        command: > 
+          julia --color=yes --project=.buildkite .buildkite/ci_driver.jl
+          --config_file $CONFIG_PATH/diagnostic_edmfx_era5_initial_condition.yml
+          --job_id diagnostic_edmfx_era5_initial_condition
+        artifact_paths: "diagnostic_edmfx_era5_initial_condition/output_active/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_gpus: 1
+          slurm_mem: 16GB
 
       - label: "GPU: EDOnly EDMFX aquaplanet"
         command: >
diff --git a/Artifacts.toml b/Artifacts.toml
@@ -77,3 +77,6 @@ lazy = true
 
 [DYAMOND_SUMMER_ICS_p14deg]
 git-tree-sha1 = "789c8dd160b022a6c714a102c1495530c7a9d20c"
+
+[weather_model_ic]
+git-tree-sha1 = "1c05bc305e79ceaab8c537be4f8c4a26465996a3"
diff --git a/NEWS.md b/NEWS.md
@@ -4,6 +4,8 @@ ClimaAtmos.jl Release Notes
 main
 -------
 
+PR [#3883](https://github.com/CliMA/ClimaAtmos.jl/pull/3883) adds atmosphere initial condition for WeatherQuest from ERA5
+
 PR [#3870](https://github.com/CliMA/ClimaAtmos.jl/pull/3870) adds externally driven SCM models forced by monthly-averaged ERA5
 reanalysis which are less data intensive and more stable. The simulation also employs the new implicit solver yielding a several 
 factor speedup. The associated documentation section, "Single Column Model", was also improved.
diff --git a/Project.toml b/Project.toml
@@ -11,6 +11,7 @@ AtmosphericProfilesLibrary = "86bc3604-9858-485a-bdbe-831ec50de11d"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 ClimaDiagnostics = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f"
+ClimaInterpolations = "dd0f122e-fa3b-47f3-bcf0-93bbc60d885e"
 ClimaParams = "5c42b081-d73a-476f-9059-fd94b934656c"
 ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 ClimaUtilities = "b3f4f4ca-9299-4f7f-bd9b-81e1242a7513"
@@ -44,6 +45,7 @@ AtmosphericProfilesLibrary = "0.1.7"
 ClimaComms = "0.6.8"
 ClimaCore = "0.14.35"
 ClimaDiagnostics = "0.2.12"
+ClimaInterpolations = "0.1.0"
 ClimaParams = "0.10.35"
 ClimaTimeSteppers = "0.8.2"
 ClimaUtilities = "0.1.22"
diff --git a/config/default_configs/default_config.yml b/config/default_configs/default_config.yml
@@ -161,7 +161,7 @@ bubble:
   help: "Enable bubble correction for more accurate surface areas"
   value: false
 start_date:
-  help: "Start date of the simulation"
+  help: "Start date and time of the simulation. Specified as either yyyymmdd (defaults to midnight) or yyyymmdd-HHMM. Examples: [`20100101`, `20100101-0000`]"
   value: "20100101"
 forcing:
   help: "Forcing [`nothing` (default), `held_suarez`]"
diff --git a/config/model_configs/diagnostic_edmfx_era5_initial_condition.yml b/config/model_configs/diagnostic_edmfx_era5_initial_condition.yml
@@ -0,0 +1,43 @@
+z_max: 48000.0
+z_elem: 63
+h_elem: 15
+dz_bottom: 30.0
+start_date: "20250627-0000"
+initial_condition: "WeatherModel"
+discrete_hydrostatic_balance: false # TODO: change to true once we fix discrete hydrostatic balance to work with orography
+topo_smoothing: true
+topography: "Earth"
+rayleigh_sponge: true
+
+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
+co2_model: fixed
+turbconv: diagnostic_edmfx
+implicit_diffusion: true
+approximate_linear_solve_iters: 2
+prognostic_tke: true
+edmfx_upwinding: first_order
+edmfx_entr_model: "Generalized"
+edmfx_detr_model: "Generalized"
+edmfx_nh_pressure: true
+edmfx_sgs_mass_flux: true
+edmfx_sgs_diffusive_flux: true
+moist: equil
+cloud_model: "grid_scale"
+precip_model: 0M
+dt: 120secs
+t_end: 12hours
+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]
+    period: 60mins
+  - short_name: [rsut, rlut, ts, pr]
+    period: 60mins
diff --git a/src/ClimaAtmos.jl b/src/ClimaAtmos.jl
@@ -23,6 +23,7 @@ include(joinpath("utils", "debug_utils.jl"))
 include(joinpath("utils", "variable_manipulations.jl"))
 include(joinpath("utils", "read_gcm_driven_scm_data.jl"))
 include(joinpath("utils", "era5_observations_to_forcing_file.jl"))
+include(joinpath("utils", "weather_model.jl"))
 
 include(joinpath("utils", "AtmosArtifacts.jl"))
 import .AtmosArtifacts as AA
diff --git a/src/callbacks/callbacks.jl b/src/callbacks/callbacks.jl
@@ -215,7 +215,7 @@ function set_insolation_variables!(Y, p, t, tvi::TimeVaryingInsolation)
     max_zenith_angle = FT(π) / 2 - eps(FT)
     irradiance = FT(CAP.tot_solar_irrad(params))
     au = FT(CAP.astro_unit(params))
-    # TODO: Where does this date0 come from?
+    # date0 references the start of the astronomial epoch J2000
     date0 = DateTime("2000-01-01T11:58:56.816")
     d, δ, η_UTC =
         FT.(
diff --git a/src/initial_conditions/InitialConditions.jl b/src/initial_conditions/InitialConditions.jl
@@ -23,6 +23,7 @@ import ..gcm_driven_profile
 import ..gcm_height
 import ..gcm_driven_profile_tmean
 import ..constant_buoyancy_frequency_initial_state
+import ..weather_model_data_path
 
 import Thermodynamics.TemperatureProfiles:
     DecayingTemperatureProfile, DryAdiabaticProfile
diff --git a/src/initial_conditions/initial_conditions.jl b/src/initial_conditions/initial_conditions.jl
@@ -192,7 +192,23 @@ struct MoistFromFile <: InitialCondition
     file_path::String
 end
 
-function (initial_condition::MoistFromFile)(params)
+"""
+    WeatherModel(start_date)
+
+An `InitialCondition` that initializes the model with an empty state, and then overwrites
+it with the content of a NetCDF file that contains the initial conditions, stored in the 
+artifact `weather_model_ic`/raw/era5_raw_YYYYMMDD_HHMM.nc. We interpolate the initial 
+conditions from ERA5 pressure level grid to a z grid, saving to the artifact 
+weather_model_ic/init/era5_init_YYYYMMDD_HHMM.nc. It is then interpolated to the model
+grid in `_overwrite_initial_conditions_from_file!`, which documents the required variables.
+Recall running `ClimaUtilities.ClimaArtiffacts.@clima_artifact("weather_model_ic")` gets 
+the artifact path.
+"""
+struct WeatherModel <: InitialCondition
+    start_date::String
+end
+
+function (initial_condition::Union{MoistFromFile, WeatherModel})(params)
     function local_state(local_geometry)
         FT = eltype(params)
         grav = CAP.grav(params)
@@ -340,8 +356,38 @@ function overwrite_initial_conditions!(
     return nothing
 end
 
+# Restored original MoistFromFile function behavior
+function overwrite_initial_conditions!(
+    initial_condition::MoistFromFile,
+    Y,
+    thermo_params,
+)
+    return _overwrite_initial_conditions_from_file!(
+        initial_condition.file_path,
+        nothing, # use default extrapolation bc
+        Y,
+        thermo_params,
+    )
+end
+
+# WeatherModel function using the shared implementation
+function overwrite_initial_conditions!(
+    initial_condition::WeatherModel,
+    Y,
+    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!(
+        file_path,
+        extrapolation_bc,
+        Y,
+        thermo_params,
+    )
+end
+
 """
-    overwrite_initial_conditions!(initial_condition::MoistFromFile, Y, thermo_params, config)
+    _overwrite_initial_conditions_from_file!(file_path::String, Y, thermo_params, config)
 
 Given a prognostic state `Y`, an `initial condition` (specifically, where initial values are
 assigned from interpolations of existing datasets), a `thermo_state`, this function
@@ -357,24 +403,40 @@ We expect the file to contain the following variables:
 - `u, v, w`, for velocity,
 - `cswc, crwc` for snow and rain water content (for 1 moment microphysics).
 """
-function overwrite_initial_conditions!(
-    initial_conditions::MoistFromFile,
+function _overwrite_initial_conditions_from_file!(
+    file_path::String,
+    extrapolation_bc,
     Y,
     thermo_params,
 )
-    file_path = initial_conditions.file_path
+    regridder_kwargs = isnothing(extrapolation_bc) ? () : (; 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),
+        SpaceVaryingInputs.SpaceVaryingInput(
+            file_path,
+            "p",
+            face_space,
+            regridder_kwargs = regridder_kwargs,
+        ),
         Fields.half,
     )
-    ᶜT = SpaceVaryingInputs.SpaceVaryingInput(file_path, "t", center_space)
-    ᶜq_tot = SpaceVaryingInputs.SpaceVaryingInput(file_path, "q", center_space)
+    ᶜT = SpaceVaryingInputs.SpaceVaryingInput(
+        file_path,
+        "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.
@@ -400,9 +462,24 @@ function overwrite_initial_conditions!(
     # cell faces.
     vel =
         Geometry.UVWVector.(
-            SpaceVaryingInputs.SpaceVaryingInput(file_path, "u", center_space),
-            SpaceVaryingInputs.SpaceVaryingInput(file_path, "v", center_space),
-            SpaceVaryingInputs.SpaceVaryingInput(file_path, "w", center_space),
+            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)))
@@ -423,12 +500,14 @@ function overwrite_initial_conditions!(
                 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
 
diff --git a/src/solver/model_getters.jl b/src/solver/model_getters.jl
@@ -43,7 +43,7 @@ function get_insolation_form(parsed_args)
     elseif insolation == "timevarying"
         # TODO: Remove this argument once we have support for integer time and
         # we can easily convert from time to date
-        start_date = DateTime(parsed_args["start_date"], dateformat"yyyymmdd")
+        start_date = parse_date(parsed_args["start_date"])
         TimeVaryingInsolation(start_date)
     elseif insolation == "rcemipii"
         RCEMIPIIInsolation()
diff --git a/src/solver/type_getters.jl b/src/solver/type_getters.jl
@@ -398,6 +398,8 @@ function get_initial_condition(parsed_args, atmos)
             atmos.external_forcing.external_forcing_file,
             parsed_args["start_date"],
         )
+    elseif parsed_args["initial_condition"] == "WeatherModel"
+        return ICs.WeatherModel(parsed_args["start_date"])
     else
         error(
             "Unknown `initial_condition`: $(parsed_args["initial_condition"])",
@@ -630,7 +632,7 @@ function get_sim_info(config::AtmosConfig)
 
     isnothing(restart_file) ||
         @info "Restarting simulation from file $restart_file"
-    epoch = DateTime(parsed_args["start_date"], dateformat"yyyymmdd")
+    epoch = parse_date(parsed_args["start_date"])
     t_start_int = time_to_seconds(parsed_args["t_start"])
     if !isnothing(restart_file) && t_start_int != 0
         @warn "Non zero `t_start` passed with a restarting simulation. The provided `t_start` will be ignored."
diff --git a/src/utils/utilities.jl b/src/utils/utilities.jl
@@ -4,6 +4,7 @@
 import ClimaComms
 import ClimaCore: Spaces, Topologies, Fields, Geometry
 import LinearAlgebra: norm_sqr
+using Dates: DateTime, @dateformat_str
 
 is_energy_var(symbol) = symbol in (:ρe_tot, :ρae_tot)
 is_momentum_var(symbol) = symbol in (:uₕ, :ρuₕ, :u₃, :ρw)
@@ -479,6 +480,28 @@ function promote_period(period::Dates.OtherPeriod)
     return period
 end
 
+"""
+    parse_date(date_str)
+
+Parse a date string into a `DateTime` object. Currently, only the following formats are supported:
+- yyyymmdd
+- yyyymmdd-HHMM
+"""
+function parse_date(date_str)
+    # Define a mapping between allowed formats and corresponding date format 
+    date_format_mapping = Dict(
+        r"^\d{8}$" => dateformat"yyyymmdd",
+        r"^\d{8}-\d{4}$" => dateformat"yyyymmdd-HHMM",
+    )
+    for (pattern, format) in date_format_mapping
+        !isnothing(match(pattern, date_str)) &&
+            return DateTime(date_str, format)
+    end
+    error(
+        "Date string $date_str does not match any of the allowed formats: yyyymmdd or yyyymmdd-HHMM",
+    )
+end
+
 function iscolumn(space)
     # TODO: Our columns are 2+1D boxes with one element at the base. Fix this
     isbox =
diff --git a/src/utils/weather_model.jl b/src/utils/weather_model.jl
diff --git a/test/dependencies.jl b/test/dependencies.jl
diff --git a/test/utilities.jl b/test/utilities.jl
diff --git a/toml/diagnostic_edmfx_era5_ic.toml b/toml/diagnostic_edmfx_era5_ic.toml
