debug remaining runs and add global to buildkite

.buildkite/pipeline.yml

@@ -120,6 +120,21 @@ steps:
         env:
           CLIMACOMMS_DEVICE: "CUDA"
 
+      - label: "Global Run CPU"
+        command: "julia --color=yes --project=.buildkite experiments/integrated/global/global_soil_canopy.jl"
+        artifact_paths: "experiments/integrated/global/plots/*cpu*png"
+        agents:
+          slurm_mem: 16G
+
+    steps:
+      - label: "Global run GPU"
+        command: "julia --color=yes --project=.buildkite experiments/integrated/global/global_soil_canopy.jl"
+        agents:
+          slurm_ntasks: 1
+          slurm_gres: "gpu:1"
+        env:
+          CLIMACOMMS_DEVICE: "CUDA"
+
   - group: "ClimaLandSimulations"
     steps:
       - label: "Ozark figures Makie"

experiments/integrated/fluxnet/ozark_pft.jl

@@ -271,14 +271,15 @@ T_0 =
     drivers.TS.values[1 + Int(round(t0 / DATA_DT))] :
     drivers.TA.values[1 + Int(round(t0 / DATA_DT))] + 40# Get soil temperature at t0
 ρc_s =
-    volumetric_heat_capacity.(Y.soil.ϑ_l, Y.soil.θ_i, Ref(land.soil.parameters))
-Y.soil.ρe_int =
-    volumetric_internal_energy.(
+    volumetric_heat_capacity.(
+        Y.soil.ϑ_l,
         Y.soil.θ_i,
-        ρc_s,
-        T_0,
-        Ref(land.soil.parameters),
+        land.soil.parameters.ρc_ds,
+        earth_param_set,
     )
+Y.soil.ρe_int =
+    volumetric_internal_energy.(Y.soil.θ_i, ρc_s, T_0, earth_param_set)
+
 Y.soilco2.C .= FT(0.000412) # set to atmospheric co2, mol co2 per mol air
 ψ_stem_0 = FT(-1e5 / 9800) # pressure in the leaf divided by rho_liquid*gravitational acceleration [m]
 ψ_leaf_0 = FT(-2e5 / 9800)

experiments/integrated/global/global_soil_canopy.jl

@@ -1,7 +1,6 @@
-using CairoMakie
-using Statistics
-using Dates
 import SciMLBase
+import ClimaComms
+@static pkgversion(ClimaComms) >= v"0.6" && ClimaComms.@import_required_backends
 import ClimaTimeSteppers as CTS
 using ClimaCore
 using ClimaUtilities.ClimaArtifacts
@@ -12,22 +11,30 @@ import ClimaUtilities.TimeVaryingInputs: TimeVaryingInput
 import ClimaUtilities.SpaceVaryingInputs: SpaceVaryingInput
 import ClimaUtilities.Regridders: InterpolationsRegridder
 import ClimaUtilities.ClimaArtifacts: @clima_artifact
-import NCDatasets
 import ClimaParams as CP
-using ClimaComms
 
 using ClimaLand
 using ClimaLand.Soil
 using ClimaLand.Canopy
 import ClimaLand
 import ClimaLand.Parameters as LP
 
+using CairoMakie
+using Statistics
+using Dates
+import NCDatasets
+
 regridder_type = :InterpolationsRegridder
 extrapolation_bc =
     (Interpolations.Periodic(), Interpolations.Flat(), Interpolations.Flat())
 context = ClimaComms.context()
 outdir = joinpath(pkgdir(ClimaLand), "experiments/integrated/global/plots")
 !ispath(outdir) && mkpath(outdir)
+
+device_suffix =
+    typeof(ClimaComms.context().device) <: ClimaComms.CPUSingleThreaded ?
+    "cpu" : "gpu"
+
 FT = Float64
 radius = FT(6378.1e3);
 depth = FT(50)
@@ -325,21 +332,36 @@ Y.canopy.hydraulics.ϑ_l.:1 .= plant_ν
 evaluate!(Y.canopy.energy.T, atmos.T, t0)
 
 set_initial_cache! = make_set_initial_cache(land)
-exp_tendency! = make_exp_tendency(land)
+exp_tendency! = make_exp_tendency(land);
+imp_tendency! = ClimaLand.make_imp_tendency(land);
+tendency_jacobian! = ClimaLand.make_tendency_jacobian(land);
 set_initial_cache!(p, Y, t0)
-timestepper = CTS.SSP22Heuns()
-ode_algo = CTS.ExplicitAlgorithm(timestepper)
+stepper = CTS.ARS343()
+norm_condition = CTS.MaximumError(FT(1e-8))
+conv_checker = CTS.ConvergenceChecker(; norm_condition = norm_condition)
+ode_algo = CTS.IMEXAlgorithm(
+    stepper,
+    CTS.NewtonsMethod(
+        max_iters = 2,
+        update_j = CTS.UpdateEvery(CTS.NewNewtonIteration),
+        convergence_checker = conv_checker,
+    ),
+)
+
+# set up jacobian info
+jac_kwargs =
+    (; jac_prototype = ImplicitEquationJacobian(Y), Wfact = tendency_jacobian!)
 
 prob = SciMLBase.ODEProblem(
     CTS.ClimaODEFunction(
         T_exp! = exp_tendency!,
+        T_imp! = SciMLBase.ODEFunction(imp_tendency!; jac_kwargs...),
         dss! = ClimaLand.dss!,
-        T_imp! = nothing,
     ),
     Y,
     (t0, tf),
     p,
-);
+)
 
 saveat = Array(t0:(3600 * 3):tf)
 sv = (;
@@ -360,55 +382,57 @@ cb = SciMLBase.CallbackSet(driver_cb, saving_cb)
     saveat = saveat,
 )
 
-longpts = range(-180.0, 180.0, 101)
-latpts = range(-90.0, 90.0, 101)
-hcoords = [
-    ClimaCore.Geometry.LatLongPoint(lat, long) for long in longpts,
-    lat in reverse(latpts)
-]
-remapper = ClimaCore.Remapping.Remapper(surface_space, hcoords)
-S = ClimaLand.Domains.top_center_to_surface(
-    (sol.u[4].soil.ϑ_l .- θ_r) ./ (ν .- θ_r),
-)
-S_ice = ClimaLand.Domains.top_center_to_surface(sol.u[4].soil.θ_i ./ ν)
-T_soil = ClimaLand.Domains.top_center_to_surface(sv.saveval[4].soil.T)
-SW = sv.saveval[4].drivers.SW_d
-
-GPP = sv.saveval[4].canopy.photosynthesis.GPP .* 1e6
-T_canopy = sol.u[4].canopy.energy.T
-fields = [S, S_ice, T_soil, GPP, T_canopy, SW]
-titles = [
-    "Effective saturation",
-    "Effective ice saturation",
-    "Temperature (K) - Soil",
-    "GPP",
-    "Temperature (K) - Canopy",
-    "Incident SW",
-]
-plotnames = ["S", "Sice", "temp", "gpp", "temp_canopy", "sw"]
-mask_top = ClimaLand.Domains.top_center_to_surface(soil_params_mask)
-mask_remap = ClimaCore.Remapping.interpolate(remapper, mask_top)
-for (id, x) in enumerate(fields)
-    title = titles[id]
-    plotname = plotnames[id]
-    x_remap = ClimaCore.Remapping.interpolate(remapper, x)
-
-    fig = Figure(size = (600, 400))
-    ax = Axis(
-        fig[1, 1],
-        xlabel = "Longitude",
-        ylabel = "Latitude",
-        title = title,
-    )
-    clims = extrema(x_remap)
-    CairoMakie.heatmap!(
-        ax,
-        longpts,
-        latpts,
-        oceans_to_value.(x_remap, mask_remap, 0.0),
-        colorrange = clims,
+if context.device isa ClimaComms.CPUSingleThreaded
+    longpts = range(-180.0, 180.0, 101)
+    latpts = range(-90.0, 90.0, 101)
+    hcoords = [
+        ClimaCore.Geometry.LatLongPoint(lat, long) for long in longpts,
+        lat in reverse(latpts)
+    ]
+    remapper = ClimaCore.Remapping.Remapper(surface_space, hcoords)
+    S = ClimaLand.Domains.top_center_to_surface(
+        (sol.u[4].soil.ϑ_l .- θ_r) ./ (ν .- θ_r),
     )
-    Colorbar(fig[:, end + 1], colorrange = clims)
-    outfile = joinpath(outdir, "$plotname.png")
-    CairoMakie.save(outfile, fig)
+    S_ice = ClimaLand.Domains.top_center_to_surface(sol.u[4].soil.θ_i ./ ν)
+    T_soil = ClimaLand.Domains.top_center_to_surface(sv.saveval[4].soil.T)
+    SW = sv.saveval[4].drivers.SW_d
+
+    GPP = sv.saveval[4].canopy.photosynthesis.GPP .* 1e6
+    T_canopy = sol.u[4].canopy.energy.T
+    fields = [S, S_ice, T_soil, GPP, T_canopy, SW]
+    titles = [
+        "Effective saturation",
+        "Effective ice saturation",
+        "Temperature (K) - Soil",
+        "GPP",
+        "Temperature (K) - Canopy",
+        "Incident SW",
+    ]
+    plotnames = ["S", "Sice", "temp", "gpp", "temp_canopy", "sw"]
+    mask_top = ClimaLand.Domains.top_center_to_surface(soil_params_mask)
+    mask_remap = ClimaCore.Remapping.interpolate(remapper, mask_top)
+    for (id, x) in enumerate(fields)
+        title = titles[id]
+        plotname = plotnames[id]
+        x_remap = ClimaCore.Remapping.interpolate(remapper, x)
+
+        fig = Figure(size = (600, 400))
+        ax = Axis(
+            fig[1, 1],
+            xlabel = "Longitude",
+            ylabel = "Latitude",
+            title = title,
+        )
+        clims = extrema(x_remap)
+        CairoMakie.heatmap!(
+            ax,
+            longpts,
+            latpts,
+            oceans_to_value.(x_remap, mask_remap, 0.0),
+            colorrange = clims,
+        )
+        Colorbar(fig[:, end + 1], colorrange = clims)
+        outfile = joinpath(outdir, "$plotname.png")
+        CairoMakie.save(outfile, fig)
+    end
 end

experiments/standalone/Biogeochemistry/experiment.jl

@@ -133,9 +133,19 @@ for (FT, tf) in ((Float32, 2 * dt), (Float64, tf))
         Y.soil.ϑ_l .= FT(0.33)
         Y.soil.θ_i .= FT(0.1)
         T = FT(279.85)
-        ρc_s = Soil.volumetric_heat_capacity(FT(0.33), FT(0.1), params)
+        ρc_s = Soil.volumetric_heat_capacity(
+            FT(0.33),
+            FT(0.1),
+            params.ρc_ds,
+            params.earth_param_set,
+        )
         Y.soil.ρe_int .=
-            Soil.volumetric_internal_energy.(FT(0.0), ρc_s, T, Ref(params))
+            Soil.volumetric_internal_energy.(
+                FT(0.0),
+                ρc_s,
+                T,
+                params.earth_param_set,
+            )
     end
 
     function init_co2!(Y, z)

lib/ClimaLandSimulations/src/Fluxnet/run_fluxnet.jl

@@ -249,14 +249,15 @@ function run_fluxnet(
         volumetric_heat_capacity.(
             Y.soil.ϑ_l,
             Y.soil.θ_i,
-            Ref(land.soil.parameters),
+            land.soil.parameters.ρc_ds,
+            land.soil.parameters.earth_param_set,
         )
     Y.soil.ρe_int =
         volumetric_internal_energy.(
             Y.soil.θ_i,
             ρc_s,
             T_0,
-            Ref(land.soil.parameters),
+            land.soil.parameters.earth_param_set,
         )
     Y.soilco2.C .= FT(0.000412) # set to atmospheric co2, mol co2 per mol air
     ψ_stem_0 = FT(-1e5 / 9800) # pressure in the leaf divided by rho_liquid*gravitational acceleration [m] 

src/standalone/Soil/soil_hydrology_parameterizations.jl

@@ -336,9 +336,9 @@ function soil_resistance(
 ) where {FT, EP, C}
     (; S_c) = hydrology_cm
     _D_vapor = FT(LP.D_vapor(earth_param_set))
-    S_l = effective_saturation(ν, ϑ_l, θ_r)
+    S_w = effective_saturation(ν, θ_l + θ_i, θ_r)
     τ_a = soil_tortuosity(θ_l, θ_i, ν)
-    dsl = dry_soil_layer_thickness(S_l, S_c, d_ds)
+    dsl::FT = dry_soil_layer_thickness(S_w, S_c, d_ds)
     r_soil = dsl / (_D_vapor * τ_a) # [s\m]
     return r_soil
 end