Tr/use land simulation struct

.buildkite/Manifest.toml

@@ -469,7 +469,7 @@ uuid = "1ecacbb8-0713-4841-9a07-eb5aa8a2d53f"
 version = "0.2.14"
 
 [[deps.ClimaLand]]
-deps = ["ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaTimeSteppers", "ClimaUtilities", "Dates", "DocStringExtensions", "Insolation", "Interpolations", "LazyArtifacts", "LinearAlgebra", "NCDatasets", "SciMLBase", "StaticArrays", "SurfaceFluxes", "Thermodynamics"]
+deps = ["ClimaComms", "ClimaCore", "ClimaDiagnostics", "ClimaTimeSteppers", "ClimaUtilities", "Dates", "DocStringExtensions", "Insolation", "Interpolations", "LazyArtifacts", "LazyBroadcast", "LinearAlgebra", "NCDatasets", "SciMLBase", "StaticArrays", "SurfaceFluxes", "Thermodynamics"]
 path = ".."
 uuid = "08f4d4ce-cf43-44bb-ad95-9d2d5f413532"
 version = "0.16.2"

experiments/integrated/global/global_soil_canopy.jl

@@ -15,6 +15,7 @@ using ClimaLand.Soil
 using ClimaLand.Canopy
 import ClimaLand
 import ClimaLand.Parameters as LP
+import ClimaLand.Simulations: LandSimulation, solve!
 
 import CairoMakie
 import GeoMakie
@@ -43,9 +44,8 @@ surface_space = domain.space.surface
 subsurface_space = domain.space.subsurface
 
 start_date = DateTime(2008);
-t0 = 0.0
 dt = 450.0
-tf = 3600
+stop_date = start_date + Dates.Second(3600)
 
 # Forcing data
 era5_ncdata_path =
@@ -107,7 +107,7 @@ photosynthesis_args =
 # Set up plant hydraulics
 modis_lai_ncdata_path = ClimaLand.Artifacts.modis_lai_single_year_path(;
     context = nothing,
-    year = Dates.year(Second(t0) + start_date),
+    year = Dates.year(start_date),
 )
 LAIfunction = ClimaLand.prescribed_lai_modis(
     modis_lai_ncdata_path,
@@ -175,22 +175,8 @@ land = SoilCanopyModel{FT}(;
     canopy_model_args = canopy_model_args,
 )
 
-Y, p, cds = initialize(land)
-
-t0 = 0.0
-dt = 450.0
-tf = 3600
-
 ic_path = ClimaLand.Artifacts.soil_ic_2008_50m_path(; context = context)
-set_initial_state! =
-    ClimaLand.Simulations.make_set_initial_state_from_file(ic_path, land)
-set_initial_cache! = make_set_initial_cache(land)
-exp_tendency! = make_exp_tendency(land)
-imp_tendency! = ClimaLand.make_imp_tendency(land)
-jacobian! = ClimaLand.make_jacobian(land)
-
-set_initial_state!(Y, p, t0, land)
-set_initial_cache!(p, Y, t0)
+set_ic! = ClimaLand.Simulations.make_set_initial_state_from_file(ic_path, land)
 
 stepper = CTS.ARS343()
 ode_algo = CTS.IMEXAlgorithm(
@@ -201,21 +187,6 @@ ode_algo = CTS.IMEXAlgorithm(
     ),
 )
 
-# set up jacobian info
-jac_kwargs =
-    (; jac_prototype = ClimaLand.FieldMatrixWithSolver(Y), Wfact = jacobian!)
-
-prob = SciMLBase.ODEProblem(
-    CTS.ClimaODEFunction(
-        T_exp! = exp_tendency!,
-        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
-        dss! = ClimaLand.dss!,
-    ),
-    Y,
-    (t0, tf),
-    p,
-)
-
 # ClimaDiagnostics
 nc_writer =
     ClimaDiagnostics.Writers.NetCDFWriter(subsurface_space, outdir; start_date)
@@ -227,23 +198,18 @@ diags = ClimaLand.default_diagnostics(
     average_period = :hourly,
 )
 
-diagnostic_handler =
-    ClimaDiagnostics.DiagnosticsHandler(diags, Y, p, t0; dt = dt)
-
-diag_cb = ClimaDiagnostics.DiagnosticsCallback(diagnostic_handler)
-
-updateat = Array(t0:(3600 * 3):tf)
-drivers = ClimaLand.get_drivers(land)
-updatefunc = ClimaLand.make_update_drivers(drivers)
-driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
-
-sol = ClimaComms.@time ClimaComms.device() SciMLBase.solve(
-    prob,
-    ode_algo;
-    dt = dt,
-    callback = SciMLBase.CallbackSet(driver_cb, diag_cb),
-)
-
+simulation = LandSimulation(
+    FT,
+    start_date,
+    stop_date,
+    dt,
+    land;
+    outdir,
+    diagnostics = diags,
+    timestepper = ode_algo,
+    set_ic!,
+)
+ClimaLand.Simulations.solve!(simulation)
 # ClimaAnalysis
 if ClimaComms.iamroot(context)
     simdir = ClimaAnalysis.SimDir(outdir)

experiments/integrated/performance/integrated_timestep_test.jl

@@ -56,6 +56,7 @@ using ClimaLand.Canopy
 using ClimaLand.Canopy.PlantHydraulics
 import ClimaLand
 import ClimaLand.Parameters as LP
+import ClimaLand.Simulations: LandSimulation, solve!
 import ClimaParams
 
 
@@ -67,8 +68,7 @@ try
 catch
 end
 
-function set_initial_conditions(land, t0)
-    Y, p, cds = initialize(land)
+function set_ic!(Y, p, t0, model)
     FT = eltype(Y.soil.ϑ_l)
     set_initial_cache! = make_set_initial_cache(land)
 
@@ -111,7 +111,7 @@ function set_initial_conditions(land, t0)
 
     Y.canopy.energy.T = FT(297.5)
     set_initial_cache!(p, Y, t0)
-    return Y, p
+    return
 end
 
 context = ClimaComms.context()
@@ -179,7 +179,7 @@ function zenith_angle(
     insol_params::Insolation.Parameters.InsolationParameters{_FT} = earth_param_set.insol_params,
 ) where {_FT}
     # This should be time in UTC
-    current_datetime = start_date + Dates.Second(round(t))
+    current_datetime = date(t)
 
     # Orbital Data uses Float64, so we need to convert to our sim FT
     d, δ, η_UTC =
@@ -346,11 +346,6 @@ land = SoilCanopyModel{FT}(;
     canopy_model_args = canopy_model_args,
 )
 
-# Define explicit and implicit tendencies, and the jacobian
-exp_tendency! = make_exp_tendency(land)
-imp_tendency! = make_imp_tendency(land);
-jacobian! = make_jacobian(land);
-
 # Timestepping information
 N_hours = 8
 tf = Float64(t0 + N_hours * 3600.0)
@@ -365,10 +360,6 @@ ode_algo = CTS.IMEXAlgorithm(
     ),
 );
 
-# Set up simulation callbacks
-drivers = ClimaLand.get_drivers(land)
-updatefunc = ClimaLand.make_update_drivers(drivers)
-
 # Choose timesteps and set up arrays to store information
 ref_dt = 50.0
 dts = [ref_dt, 100.0, 200.0, 450.0, 900.0, 1800.0, 3600.0]
@@ -384,36 +375,18 @@ times = []
 for dt in dts
     @info dt
     # Initialize model and set initial conditions before each simulation
-    Y, p = set_initial_conditions(land, t0)
-    jac_kwargs = (;
-        jac_prototype = ClimaLand.FieldMatrixWithSolver(Y),
-        Wfact = jacobian!,
-    )
-
-    # Create callback for driver updates
-    saveat = vcat(Array(t0:(3 * 3600):tf), tf)
-    updateat = vcat(Array(t0:(3 * 3600):tf), tf)
-    driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
-
-    # Solve simulation
-    prob = SciMLBase.ODEProblem(
-        CTS.ClimaODEFunction(
-            T_exp! = exp_tendency!,
-            T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
-            dss! = ClimaLand.dss!,
-        ),
-        Y,
-        (t0, tf),
-        p,
+    simulation = LandSimulation(
+        FT,
+        start_date,
+        start_date + Dates.Second(round(tf)),
+        dt,
+        land;
+        diagnostics = [],
+        timestepper = ode_algo,
+        set_ic!,
+        sol_save_interval = 3 * 3600,
     )
-    @time sol = SciMLBase.solve(
-        prob,
-        ode_algo;
-        dt = dt,
-        callback = driver_cb,
-        saveat = saveat,
-    )
-
+    sol = ClimaLand.Simulations.solve!(simulation)
     # Save results for comparison
     if dt == ref_dt
         global ref_T =
@@ -474,9 +447,13 @@ ax3 = Axis(
     ylabel = "T (K)",
     title = "T throughout simulation; length = $(sim_time) hours, dts in [$(dts[1]), $(dts[end])]",
 )
-times = times ./ 3600.0 # hours
 for i in 1:length(times)
-    lines!(ax3, times[i], T_states[i], label = "dt $(dts[i]) s")
+    lines!(
+        ax3,
+        float.(times[i]) ./ 3600.0,
+        T_states[i],
+        label = "dt $(dts[i]) s",
+    )
 end
 axislegend(ax3, position = :rt)
 save(joinpath(savedir, "states.png"), fig3)

src/Artifacts.jl

@@ -63,7 +63,7 @@ end
 Return the path to the file that contains the ERA5 forcing data for 2008.
 
 Optionally, you can pass the lowres=true keyword to download a lower spatial resolution version of the data and return the path to that file.
- If the high resolution data is not 
+ If the high resolution data is not
 available locally, we also return the path to the low res data.
 """
 function era5_land_forcing_data2008_path(; context = nothing, lowres = false)
@@ -83,7 +83,7 @@ function era5_land_forcing_data2008_path(; context = nothing, lowres = false)
             return hires_path
         catch
             @warn(
-                "High resolution ERA5 forcing not available locally; downloading and using low resolution data instead."
+                "High resolution ERA5 forcing not available locally; using low resolution data instead."
             )
             return lowres_path
         end

src/simulations/Simulations.jl

@@ -22,7 +22,7 @@ include("initial_conditions.jl")
         I <: SciMLBase.DEIntegrator,
     }
 
-the ClimaLand LandSimulation struct, which specifies 
+the ClimaLand LandSimulation struct, which specifies
 - the discrete set of equations to solve (defined by the `model`);
 - the timestepping algorithm;
 - user callbacks (passed as a tuple) to be executed at specific times in the simulations;
@@ -32,13 +32,13 @@ User callbacks are optional: examples currently include callbacks that estimate
 to solution and SYPD of the simulation as it runs, checkpoint the state, or check the solution
 for NaNs. Others can be added here.
 
-Diagnostics are implemented as callbacks, and are also optional. 
-However, a default is provided. `diagnostics` is expected to be a 
+Diagnostics are implemented as callbacks, and are also optional.
+However, a default is provided. `diagnostics` is expected to be a
 list of `ClimaDiagnostics.ScheduledDiagnostics`.
 
 Finally, the private field _required_callbacks consists of callbacks that are required for the
 simulation to run correctly. Currently, this includes the callbacks which update the atmospheric
-forcing and update the LAI using prescribed data. 
+forcing and update the LAI using prescribed data.
 """
 struct LandSimulation{
     M <: ClimaLand.AbstractModel,
@@ -58,6 +58,7 @@ struct LandSimulation{
     _integrator::I
 end
 
+# TODO: Add doc string
 function LandSimulation(
     FT,
     start_date::Dates.DateTime,
@@ -98,16 +99,19 @@ function LandSimulation(
             start_date,
         ),
     ),
+    sol_save_interval = nothing,
 )
     if !isnothing(diagnostics) &&
        !isempty(diagnostics) &&
        first(diagnostics).output_writer.output_dir != outdir
         @warn "Note that the kwarg outdir and outdir used in diagnostics are inconsistent; using $(first(diagnostics).output_writer.output_dir)"
     end
 
-    domain = ClimaLand.get_domain(model)
     t0 = ITime(0, Dates.Second(1), start_date)
-    tf = ITime(Dates.seconds(stop_date - start_date), epoch = start_date)
+    tf = ITime(
+        Dates.value(convert(Dates.Second, stop_date - start_date)),
+        epoch = start_date,
+    )
     Δt = ITime(Δt, epoch = start_date)
     t0, tf, Δt = promote(t0, tf, Δt)
 
@@ -145,6 +149,9 @@ function LandSimulation(
 
     # Required callbacks
     updateat = [promote(t0:(ITime(3600 * 3)):tf...)...]
+    saveat =
+        isnothing(sol_save_interval) ? [] :
+        [promote(t0:(ITime(sol_save_interval)):tf...)...]
     drivers = ClimaLand.get_drivers(model)
     updatefunc = ClimaLand.make_update_drivers(drivers)
     driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
@@ -159,13 +166,13 @@ function LandSimulation(
     # Collect all callbacks
     callbacks =
         SciMLBase.CallbackSet(user_callbacks..., required_callbacks..., diag_cb)
-
     _integrator = SciMLBase.init(
         problem,
         timestepper;
         dt = Δt,
         callback = callbacks,
         adaptive = false,
+        saveat,
     )
     return LandSimulation(
         model,

src/standalone/Soil/energy_hydrology.jl

@@ -384,9 +384,11 @@ function ClimaLand.make_compute_jacobian(model::EnergyHydrology{FT}) where {FT}
             @. p.soil.full_bidiag_matrix_scratch +=
                 MatrixFields.LowerDiagonalMatrixRow(p.soil.topBC_scratch)
         end
-
+        # dtγ can be an ITime or a float
         @. ∂ϑres∂ϑ =
-            -dtγ * (divf2c_matrix() ⋅ p.soil.full_bidiag_matrix_scratch) - (I,)
+            FT(-1) *
+            float(dtγ) *
+            (divf2c_matrix() ⋅ p.soil.full_bidiag_matrix_scratch) - (I,)
 
         # Now create the flux term for ∂ρe∂ϑ using bidiag_matrix_scratch
         # This overwrites full_bidiag_matrix_scratch
@@ -400,7 +402,9 @@ function ClimaLand.make_compute_jacobian(model::EnergyHydrology{FT}) where {FT}
                 ),
             ) ⋅ p.soil.bidiag_matrix_scratch
         @. ∂ρeres∂ϑ =
-            -dtγ * (divf2c_matrix() ⋅ p.soil.full_bidiag_matrix_scratch) - (I,)
+            FT(-1) *
+            float(dtγ) *
+            (divf2c_matrix() ⋅ p.soil.full_bidiag_matrix_scratch) - (I,)
 
         # Now overwrite bidiag_matrix_scratch and full_bidiag scratch for the ρe ρe bidiagonal
         @. p.soil.bidiag_matrix_scratch =
@@ -416,7 +420,9 @@ function ClimaLand.make_compute_jacobian(model::EnergyHydrology{FT}) where {FT}
             MatrixFields.DiagonalMatrixRow(interpc2f_op(-p.soil.κ)) ⋅
             p.soil.bidiag_matrix_scratch
         @. ∂ρeres∂ρe =
-            -dtγ * (divf2c_matrix() ⋅ p.soil.full_bidiag_matrix_scratch) - (I,)
+            FT(-1) *
+            float(dtγ) *
+            (divf2c_matrix() ⋅ p.soil.full_bidiag_matrix_scratch) - (I,)
     end
     return compute_jacobian!
 end

src/standalone/Soil/rre.jl

@@ -438,7 +438,9 @@ function ClimaLand.make_compute_jacobian(model::RichardsModel{FT}) where {FT}
         end
         # Compute divergence matrix
         @. ∂ϑres∂ϑ =
-            -dtγ * (divf2c_matrix() ⋅ p.soil.full_bidiag_matrix_scratch) - (I,)
+            FT(-1) *
+            float(dtγ) *
+            (divf2c_matrix() ⋅ p.soil.full_bidiag_matrix_scratch) - (I,)
 
 
     end