global simulations

Author: kmdeck

src/Artifacts.jl

@@ -60,18 +60,36 @@ end
 """
     era5_land_forcing_data2008_path(; context, lowres=false)
 
-Return the path to the directory that contains the ERA5 forcing data for 2008.
+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.
+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 
+available locally, we also return the path to the low res data.
 """
 function era5_land_forcing_data2008_folder_path(;
     context = nothing,
     lowres = false,
 )
+    lowres_path = joinpath(
+        @clima_artifact("era5_land_forcing_data2008_lowres", context),
+        "era5_2008_1.0x1.0_lowres.nc",
+    )
     if lowres
-        return @clima_artifact("era5_land_forcing_data2008_lowres", context)
+        return lowres_path
     else
-        return @clima_artifact("era5_land_forcing_data2008", context)
+        try
+            hires_path = joinpath(
+                @clima_artifact("era5_land_forcing_data2008", context),
+                "era5_2008_1.0x1.0.nc",
+            )
+
+            return hires_path
+        catch
+            @warn(
+                "High resolution ERA5 forcing not available locally; downloading and using low resolution data instead."
+            )
+            return lowres_path
+        end
     end
 end
 

src/shared_utilities/utils.jl

@@ -8,6 +8,7 @@ using ClimaUtilities.ClimaArtifacts
 import Interpolations
 import ClimaUtilities.SpaceVaryingInputs: SpaceVaryingInput
 import ClimaUtilities.Regridders: InterpolationsRegridder
+import ClimaUtilities.OnlineLogging: WallTimeInfo, report_walltime
 
 export FTfromY, call_count_nans_state
 
@@ -716,3 +717,14 @@ function landsea_mask(domain::Domains.AbstractDomain; kwargs...)
     )
     return landsea_mask(domain.space.surface; filepath, kwargs...)
 end
+
+
+function ReportCallback(Nsteps::Int)
+    walltime_info = WallTimeInfo()
+    everyNsteps(u, t, integrator) = mod(integrator.step, Nsteps) == 0
+    report = let wt = walltime_info
+        (integrator) -> report_walltime(wt, integrator)
+    end
+    report_cb = SciMLBase.DiscreteCallback(everyNsteps, report)
+    return report_cb
+end

src/simulations/Simulations.jl

@@ -0,0 +1,186 @@
+module Simulations
+using ClimaTimeSteppers
+using SciMLBase
+using Dates
+import ClimaUtilities.TimeVaryingInputs:
+    TimeVaryingInput, LinearInterpolation, PeriodicCalendar
+import ClimaUtilities.ClimaArtifacts: @clima_artifact
+import ClimaUtilities.TimeManager: ITime, date
+import ClimaDiagnostics
+import ClimaParams as CP
+using ClimaLand
+import ClimaLand.Parameters as LP
+
+include("domains.jl")
+include("initial_conditions.jl")
+include("spatial_parameters.jl")
+include("model_setup.jl")
+
+struct LandSimulation{
+    C,
+    P,
+    M <: ClimaLand.AbstractLandModel,
+    D <: ClimaLand.Domains.AbstractDomain,
+    T <: ClimaTimeSteppers.IMEXAlgorithm,
+    UC,
+    DI,
+    RC,
+    CA <: SciMLBase.CallbackSet,
+    P <: SciMLBase.ODEProblem,
+}
+    context::C
+    params::P
+    model::M
+    domain::D
+    timestepper::T
+    user_callbacks::UC
+    diagnostics::DI
+    required_callbacks::RC
+    callbacks::CA
+    problem::P
+end
+
+function GlobalLandSimulation(
+    FT,
+    context,
+    start_date,
+    t0,
+    tf,
+    Δt;
+    params = LP.LandParameters(FT),
+    domain = ClimaLand.global_domain(FT; comms_ctx = context),
+    model = LandModel(domain, start_date, params, FT), # add assert
+    set_ic! = ClimaLand.set_ic_from_file(
+        ClimaLand.Artifacts.soil_ic_2008_50m_path(; context = context),
+    ), # could also support functions
+    timestepper = ClimaTimeSteppers.IMEXAlgorithm(
+        ClimaTimeSteppers.ARS111(),
+        ClimaTimeSteppers.NewtonsMethod(
+            max_iters = 3,
+            update_j = ClimaTimeSteppers.UpdateEvery(
+                ClimaTimeSteppers.NewNewtonIteration,
+            ),
+        ),
+    ),
+    user_callbacks = (
+        ClimaLand.NaNCheckCallback(
+            Dates.Month(1),
+            start_date,
+            t0,# needs tobe ITime? promote
+            Δt; # needs to be ITime? promote
+            mask = ClimaLand.landsea_mask(domain),
+        ),
+        ClimaLand.ReportCallback(1000),
+    ),
+    diagnostics = (;
+        output_vars = :short,
+        average_period = :monthly,
+        outdir = "",
+    ), # need to generalize
+)
+    # convert times to Itime
+    t0 = ITime(t0, epoch = start_date)
+    tf = ITime(tf, epoch = start_date)
+    Δt = ITime(Δt, epoch = start_date)
+    t0, tf, Δt = promote(t0, tf, Δt)
+
+    # set initial conditions
+    Y, p, cds = initialize(model)
+    set_ic!(Y, p, model, t0)
+    set_initial_cache! = make_set_initial_cache(model)
+    set_initial_cache!(p, Y, t0)
+
+    # Create tendencies and jacobian update function
+    exp_tendency! = make_exp_tendency(model)
+    imp_tendency! = ClimaLand.make_imp_tendency(model)
+    jacobian! = ClimaLand.make_jacobian(model)
+    jac_kwargs = (;
+        jac_prototype = ClimaLand.FieldMatrixWithSolver(Y),
+        Wfact = jacobian!,
+    )
+
+    # Create SciML ODE Problem
+    problem = SciMLBase.ODEProblem(
+        ClimaTimeSteppers.ClimaODEFunction(
+            T_exp! = exp_tendency!,
+            T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
+            dss! = ClimaLand.dss!,
+        ),
+        Y,
+        (t0, tf),
+        p,
+    )
+
+    # Required callbacks
+    updateat = [promote(t0:(ITime(3600 * 3)):tf...)...]
+    drivers = ClimaLand.get_drivers(model)
+    updatefunc = ClimaLand.make_update_drivers(drivers)
+    driver_cb = ClimaLand.DriverUpdateCallback(updateat, updatefunc)
+    required_callbacks = (driver_cb,) # can we update each step?
+
+    # Diagnostics callbacks - can be generalized in the future
+    if !(diagnostics isa Nothing)
+        nc_writer = ClimaDiagnostics.Writers.NetCDFWriter(
+            domain.space.subsurface,
+            diagnostics.outdir;
+            start_date,
+        )
+
+        diags = ClimaLand.default_diagnostics(
+            model,
+            start_date;
+            output_writer = nc_writer,
+            output_vars = diagnostics.output_vars,
+            average_period = diagnostics.average_period,
+        )
+
+        diagnostic_handler =
+            ClimaDiagnostics.DiagnosticsHandler(diags, Y, p, t0; dt = Δt)
+        diag_cb = (ClimaDiagnostics.DiagnosticsCallback(diagnostic_handler),)
+    else
+        diag_cb = ()
+    end
+
+
+    # Collect all callbacks
+    callbacks = SciMLBase.CallbackSet(
+        user_callbacks...,
+        diag_cb...,
+        required_callbacks...,
+    )
+
+    #_integrator = SciMLBase.init(problem, callbacks)
+    
+    return LandSimulation(
+        context,
+        params,
+        model,
+        domain,
+        timestepper,
+        user_callbacks,
+        diagnostics,
+        _required_callbacks,
+        callbacks,
+        problem,
+        t
+        Y,
+        __p
+        #_integrator
+    )
+end
+
+function step!(landsimulation)
+    integrator = SciMLBase.integrator(Y, p, t, callbacks)
+    SciMLBase.step!(integrator)
+    simulation.t[] = integrator.t
+end
+
+function solve!(simulation)
+    SciMLBase.solve!(problem, ode_algo, callbacks)
+end
+
+    function plot(landsimulation)
+
+        end
+    
+end#module

src/simulations/domains.jl

@@ -44,6 +44,7 @@ function global_domain(
     dz_tuple = (10.0, 0.05),
     depth = 50.0,
     npolynomial = 0,
+    comms_ctx = ClimaComms.context(),
 )
     if pkgversion(ClimaCore) < v"0.14.30" && apply_mask
         @warn "The land mask cannot be applied with ClimaCore < v0.14.30. Update ClimaCore for significant performance gains."
@@ -59,6 +60,7 @@ function global_domain(
         nelements,
         npolynomial,
         dz_tuple,
+        comms_ctx,
     )
     if apply_mask
         surface_space = domain.space.surface # 2d space

src/simulations/global_simulations_notes.txt

@@ -0,0 +1,29 @@
+function GlobalDomain(...)
+    return SphericalShell()
+end
+
+context = ;
+FT = ;
+param_set = ...earth_param_set(FT)
+parameter_maps = spatially_varying_maps(FT, ...)
+
+domain = GlobalDomain(context, )
+land = LandModel(FT, domain, params, component_choices; parameter_maps)
+# if hires fails to download, use lowres (era5 forcing?)
+
+
+#set_ic!(Y, land, t0)
+timestepper = (alg::ImEx = ARS111, t0, dt, tf; start_date = nothing) # handled internally with itime, e.g. pass in seconds
+timestepper = (alg::ImEx = ARS111, t0::DateTime, dt::Seconds, duration::Period) # handled internally with itime
+user_callbacks # all SciML callbacks, but we provide helper functions for specific instances (nancheck, report, checkpoint)
+diagnostics # Tuple of scheduled diagnostics, diagnostics handler is constructed internally
+# driver is combine with user_callbacks and diagnostics callback internally to make the CB set
+output_dir # will have default
+#context  # will have default, this includes the device, or take the device and create the context
+config = # specifies a default postprocessing/plots to make
+    simulation = Simulation(land, set_ic!, domain, timestepper, user_callbacks, diagnostics, config)
+
+
+
+   
+