Merge pull request #1254 from CliMA/js/atmos-land-sim

run AMIP with the integrated land model

Author: juliasloan25

.buildkite/pipeline.yml

@@ -176,6 +176,20 @@ steps:
         agents:
           slurm_mem: 20GB
 
+      - label: "AMIP: DecayWithHeight + integrated land"
+        key: "amip_integrated_land"
+        command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_decay_with_height_integrated_land.yml --job_id amip_integrated_land"
+        artifact_paths: "experiments/ClimaEarth/output/amip_integrated_land/artifacts/*"
+        agents:
+          slurm_mem: 20GB
+
+      - label: "AMIP: DecayWithHeight + bucket"
+        key: "amip_decay_with_height_bucket"
+        command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_decay_with_height_bucket.yml --job_id amip_decay_with_height_bucket"
+        artifact_paths: "experiments/ClimaEarth/output/amip_decay_with_height_bucket/artifacts/*"
+        agents:
+          slurm_mem: 20GB
+
       - label: "AMIP: bucket initial condition test"
         command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_bucket_ic.yml --job_id amip_bucket_ic"
         artifact_paths: "experiments/ClimaEarth/output/amip_bucket_ic/artifacts/*"
@@ -283,6 +297,26 @@ steps:
           slurm_mem: 20GB
           slurm_gpus: 1
 
+      - label: "GPU AMIP: DecayWithHeight + integrated land"
+        key: "gpu_amip_integrated_land"
+        command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_decay_with_height_integrated_land.yml --job_id gpu_amip_integrated_land"
+        artifact_paths: "experiments/ClimaEarth/output/gpu_amip_integrated_land/artifacts/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_mem: 20GB
+          slurm_gpus: 1
+
+      - label: "GPU AMIP: DecayWithHeight + bucket"
+        key: "gpu_amip_decay_with_height_bucket"
+        command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_decay_with_height_bucket.yml --job_id gpu_amip_decay_with_height_bucket"
+        artifact_paths: "experiments/ClimaEarth/output/gpu_amip_decay_with_height_bucket/artifacts/*"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+        agents:
+          slurm_mem: 20GB
+          slurm_gpus: 1
+
       - label: "GPU AMIP test: albedo from function"
         key: "gpu_amip_albedo_function"
         command: "julia --color=yes --project=experiments/ClimaEarth/ experiments/ClimaEarth/run_amip.jl --config_file $CONFIG_PATH/amip_albedo_function.yml --job_id gpu_amip_albedo_function"

NEWS.md

@@ -6,7 +6,18 @@ ClimaCoupler.jl Release Notes
 
 ### ClimaCoupler features
 
+#### Test AMIP with integrated land model. PR[#1254](https://github.com/CliMA/ClimaCoupler.jl/pull/1254)
+
+The integrated ClimaLand model can now be used in coupled simulations.
+A short run of AMIP using the full land model is now tested in our regular
+Buildkite pipeline, and the restarts test uses the full land model
+rather than the bucket.
+
+This PR adds the config option `land_model`, which can be set to either
+`bucket` or `integrated` to choose which land model to run with.
+
 #### Remove `nans_to_zero`. PR[#1278](https://github.com/CliMA/ClimaCoupler.jl/pull/1278)
+
 Instead of zeroing out all NaNs in a surface field, we zero out all values
 where the area fraction is 0, which is logically what we want to do.
 

config/ci_configs/amip_decay_with_height_bucket.yml

@@ -0,0 +1,36 @@
+FLOAT_TYPE: "Float32"
+albedo_model: "CouplerAlbedo"
+atmos_config_file: "config/longrun_configs/amip_target_edonly.yml"
+checkpoint_dt: "720hours"
+coupler_toml: ["toml/amip.toml"]
+dt: "240secs"
+dt_cpl: "240secs"
+dt_save_state_to_disk: "12hours"
+dt_save_to_sol: "12hours"
+dz_bottom: 100.0
+energy_check: false
+h_elem: 8
+land_albedo_type: "map_temporal"
+land_model: "bucket"
+mode_name: "amip"
+moist: equil
+netcdf_interpolation_num_points: [90, 45, 31]
+netcdf_output_at_levels: true
+output_default_diagnostics: true
+rad: allskywithclear
+radiation_reset_rng_seed: true
+rayleigh_sponge: true
+start_date: "20100101"
+surface_setup: "PrescribedSurface"
+t_end: "480secs"
+topography: "Earth"
+topo_smoothing: true
+turbconv: ~
+vert_diff: "DecayWithHeightDiffusion"
+viscous_sponge: true
+z_elem: 39
+z_max: 60000.0
+extra_atmos_diagnostics:
+  - short_name: [pr, rhoa, uas, vas, thetaa, hussfc, ts]
+    reduction_time: average
+    period: 1days

config/ci_configs/amip_decay_with_height_integrated_land.yml

@@ -0,0 +1,36 @@
+FLOAT_TYPE: "Float32"
+albedo_model: "CouplerAlbedo"
+atmos_config_file: "config/longrun_configs/amip_target_edonly.yml"
+checkpoint_dt: "720hours"
+coupler_toml: ["toml/amip.toml"]
+dt: "240secs"
+dt_cpl: "240secs"
+dt_save_state_to_disk: "12hours"
+dt_save_to_sol: "12hours"
+dz_bottom: 100.0
+energy_check: false
+h_elem: 8
+land_albedo_type: "map_temporal"
+land_model: "integrated"
+mode_name: "amip"
+moist: equil
+netcdf_interpolation_num_points: [90, 45, 31]
+netcdf_output_at_levels: true
+output_default_diagnostics: true
+rad: allskywithclear
+radiation_reset_rng_seed: true
+rayleigh_sponge: true
+start_date: "20100101"
+surface_setup: "PrescribedSurface"
+t_end: "480secs"
+topography: "Earth"
+topo_smoothing: true
+turbconv: ~
+vert_diff: "DecayWithHeightDiffusion"
+viscous_sponge: true
+z_elem: 39
+z_max: 60000.0
+extra_atmos_diagnostics:
+  - short_name: [pr, rhoa, uas, vas, thetaa, hussfc, ts]
+    reduction_time: average
+    period: 1days

docs/src/interfacer.md

@@ -53,9 +53,17 @@ after solving of the simulation has been paused or interrupted. Like
 of SciMLBase.jl.
 
 ### ComponentModelSimulation - optional functions
-- `get_model_prog_state(::ComponentModelSimulation)`: A function that
-returns the state vector of the simulation at its current state. This
-is used for checkpointing the simulation.
+- `Checkpointer.get_model_prog_state(::ComponentModelSimulation)`:
+A function that returns the state vector of the simulation at its current
+state. This is used for checkpointing the simulation.
+
+- `Checkpointer.get_model_cache(::ComponentModelSimulation)`:
+A function that returns the cache of the simulation at its current state.
+This is used for checkpointing the simulation.
+
+- `Checkpointer.restore_cache(::ComponentModelSimulation, new_cache)`:
+A function that updates the cache of the simulation with the provided
+`new_cache`. This is used for restarting the simulation.
 
 - `get_field(::ComponentModelSimulation, ::Val{property})`:
 Default `get_field` functions are provided for `energy` and `water` fields,

experiments/ClimaEarth/Manifest-v1.11.toml

@@ -344,9 +344,9 @@ version = "0.2.12"
 
 [[deps.ClimaLand]]
 deps = ["ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaUtilities", "Dates", "DocStringExtensions", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "NCDatasets", "SciMLBase", "StaticArrays", "SurfaceFluxes", "Thermodynamics"]
-git-tree-sha1 = "309507cd078678f291e1b8698b64b8200a78ab52"
+git-tree-sha1 = "08d08423e0a955180d3eea39c340b81603a00585"
 uuid = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
-version = "0.15.13"
+version = "0.15.14"
 
     [deps.ClimaLand.extensions]
     CreateParametersExt = "ClimaParams"

experiments/ClimaEarth/Manifest.toml

@@ -341,9 +341,9 @@ version = "0.2.12"
 
 [[deps.ClimaLand]]
 deps = ["ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaUtilities", "Dates", "DocStringExtensions", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "NCDatasets", "SciMLBase", "StaticArrays", "SurfaceFluxes", "Thermodynamics"]
-git-tree-sha1 = "309507cd078678f291e1b8698b64b8200a78ab52"
+git-tree-sha1 = "08d08423e0a955180d3eea39c340b81603a00585"
 uuid = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
-version = "0.15.13"
+version = "0.15.14"
 
     [deps.ClimaLand.extensions]
     CreateParametersExt = "ClimaParams"

experiments/ClimaEarth/cli_options.jl

@@ -131,6 +131,10 @@ function argparse_settings()
         arg_type = Vector{Dict{Any, Any}}
         default = []
         ### ClimaLand specific
+        "--land_model"
+        help = "Land model to use. [`bucket` (default), `integrated`]"
+        arg_type = String
+        default = "bucket"
         "--land_albedo_type"
         help = "Access land surface albedo information from data file. [`map_static` (default), `function`, `map_temporal`]"
         arg_type = String

experiments/ClimaEarth/components/land/climaland_integrated.jl

@@ -54,13 +54,13 @@ soil CO2 model. Specific details about the complexity of the model
 can be found in the ClimaLand.jl documentation.
 """
 function ClimaLandSimulation(
-    ::Type{FT},
+    ::Type{FT};
     dt::TT,
     tspan::Tuple{TT, TT},
     start_date::Dates.DateTime,
     output_dir::String,
     boundary_space,
-    area_fraction;
+    area_fraction,
     saveat::Vector{TT} = [tspan[1], tspan[2]],
     surface_elevation = CC.Fields.zeros(boundary_space),
     land_temperature_anomaly::String = "amip",
@@ -160,7 +160,6 @@ function ClimaLandSimulation(
     exp_tendency! = CL.make_exp_tendency(model)
     imp_tendency! = CL.make_imp_tendency(model)
     jacobian! = CL.make_jacobian(model)
-    set_initial_cache!(p, Y, tspan[1])
 
     # set up jacobian info
     jac_kwargs = (; jac_prototype = CL.FieldMatrixWithSolver(Y), Wfact = jacobian!)
@@ -179,7 +178,7 @@ function ClimaLandSimulation(
             model,
             start_date,
             output_writer = output_writer,
-            output_vars = :short,
+            output_vars = :long,
             average_period = :monthly,
         )
         diagnostic_handler = CD.DiagnosticsHandler(scheduled_diagnostics, Y, p, tspan[1]; dt = dt)
@@ -430,7 +429,7 @@ function Interfacer.update_field!(sim::ClimaLandSimulation, ::Val{:cos_zenith_an
 end
 
 Interfacer.step!(sim::ClimaLandSimulation, t) = Interfacer.step!(sim.integrator, t - sim.integrator.t, true)
-Interfacer.reinit!(sim::ClimaLandSimulation, t) = Interfacer.reinit!(sim.integrator, t)
+Interfacer.reinit!(sim::ClimaLandSimulation) = Interfacer.reinit!(sim.integrator)
 
 function FieldExchanger.update_sim!(sim::ClimaLandSimulation, csf, area_fraction)
     # update fields for radiative transfer
@@ -486,7 +485,22 @@ function Interfacer.add_coupler_fields!(coupler_field_names, ::ClimaLandSimulati
 end
 
 function Checkpointer.get_model_prog_state(sim::ClimaLandSimulation)
-    error("get_model_prog_state not implemented")
+    return sim.integrator.u
+end
+
+function Checkpointer.get_model_cache(sim::ClimaLandSimulation)
+    return sim.integrator.p
+end
+
+function Checkpointer.restore_cache!(sim::ClimaLandSimulation, new_cache)
+    old_cache = Checkpointer.get_model_cache(sim)
+    comms_ctx = ClimaComms.context(sim.model.soil)
+    restore!(
+        old_cache,
+        new_cache,
+        comms_ctx,
+        ignore = Set([:dss_buffer_2d, :dss_buffer_3d, :scratch1, :scratch2, :scratch3, :sfc_scratch, :subsfc_scratch]),
+    )
 end
 
 Interfacer.name(::ClimaLandSimulation) = "ClimaLandSimulation"

experiments/ClimaEarth/components/shared/restore.jl

@@ -25,11 +25,7 @@ function restore!(v1::T1, v2::T2, comms_ctx; name = "", ignore) where {T1, T2}
     # to CuArray)
     fields = filter(x -> !(x in ignore), fieldnames(T2))
     if isempty(fields)
-        if !Base.issingletontype(typeof(v1))
-            restore!(v1, v2, comms_ctx; name, ignore)
-        else
-            v1 == v2 || error("$v1 != $v2")
-        end
+        v1 == v2 || error("$v1 != $v2")
     else
         # Recursive case
         for p in fields
@@ -52,12 +48,12 @@ function restore!(
 end
 
 function restore!(
-    v1::T1,
-    v2::T2,
+    v1::Union{AbstractData, AbstractArray},
+    v2::Union{AbstractData, AbstractArray},
     comms_ctx;
     name,
     ignore,
-) where {T1 <: Union{AbstractData, AbstractArray}, T2 <: Union{AbstractData, AbstractArray}}
+)
     ArrayType = parent(v1) isa Array ? Array : ClimaComms.array_type(ClimaComms.device(comms_ctx))
     moved_to_device = ArrayType(parent(v2))
 
@@ -66,20 +62,17 @@ function restore!(
 end
 
 function restore!(
-    v1::T1,
-    v2::T2,
+    v1::Union{StaticArrays.StaticArray, Number, UnitRange, LinRange, Symbol},
+    v2::Union{StaticArrays.StaticArray, Number, UnitRange, LinRange, Symbol},
     comms_ctx;
     name,
     ignore,
-) where {
-    T1 <: Union{StaticArrays.StaticArray, Number, UnitRange, Symbol},
-    T2 <: Union{StaticArrays.StaticArray, Number, UnitRange, Symbol},
-}
+)
     v1 == v2 || error("$name is a immutable but it inconsistent ($(v1) != $(v2))")
     return nothing
 end
 
-function restore!(v1::T1, v2::T2, comms_ctx; name, ignore) where {T1 <: Dict, T2 <: Dict}
+function restore!(v1::Dict, v2::Dict, comms_ctx; name, ignore)
     # RRTGMP has some internal dictionaries
     v1 == v2 || error("$name is inconsistent")
     return nothing