Merge #442

442: add dss_state! function r=juliasloan25 a=juliasloan25



Co-authored-by: Julia Sloan <jsloan@caltech.edu>

experiments/AMIP/modular/components/atmosphere/climaatmos_init.jl

@@ -64,8 +64,11 @@ function atmos_init(::Type{FT}, atmos_config_dict::Dict) where {FT}
     integrator.p.col_integrated_rain .= FT(0)
     integrator.p.col_integrated_snow .= FT(0)
 
+    sim = ClimaAtmosSimulation(integrator.p.params, Y, spaces, integrator)
 
-    ClimaAtmosSimulation(integrator.p.params, Y, spaces, integrator)
+    # DSS state to ensure we have continuous fields
+    dss_state!(sim)
+    return sim
 end
 
 # extensions required by the Interfacer
@@ -300,3 +303,20 @@ end
 get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:energy}) = atmos_sim.integrator.u.c.ρe_tot
 
 get_field(atmos_sim::ClimaAtmosSimulation, ::Val{:water}) = atmos_sim.integrator.u.c.ρq_tot
+
+"""
+    dss_state!(sim::ClimaAtmosSimulation)
+
+Perform DSS on the state of a component simulation, intended to be used
+before the initial step of a run. This method acts on atmosphere simulations.
+These sims don't store a dss buffer in their cache, so we must allocate
+one here.
+"""
+function dss_state!(sim::ClimaAtmosSimulation)
+    Y = sim.integrator.u
+    for key in propertynames(Y)
+        field = getproperty(Y, key)
+        buffer = Spaces.create_dss_buffer(field)
+        Spaces.weighted_dss!(field, buffer)
+    end
+end

experiments/AMIP/modular/components/land/bucket_init.jl

@@ -4,6 +4,7 @@ using ClimaLSM
 import ClimaLSM
 import ClimaTimeSteppers as CTS
 import Thermodynamics as TD
+using Dates: DateTime
 
 include(joinpath(pkgdir(ClimaLSM), "parameters", "create_parameters.jl"))
 using ClimaLSM.Bucket: BucketModel, BucketModelParameters, AbstractAtmosphericDrivers, AbstractRadiativeDrivers
@@ -267,5 +268,9 @@ function bucket_init(
     prob = ODEProblem(bucket_ode_function, Y, tspan, p_new)
     integrator = init(prob, ode_algo; dt = dt, saveat = saveat, adaptive = false)
 
-    BucketSimulation(model, Y, (; domain = domain, soil_depth = d_soil), integrator, area_fraction)
+    sim = BucketSimulation(model, Y, (; domain = domain, soil_depth = d_soil), integrator, area_fraction)
+
+    # DSS state to ensure we have continuous fields
+    dss_state!(sim)
+    return sim
 end

experiments/AMIP/modular/components/land/bucket_utils.jl

@@ -149,3 +149,16 @@ function get_land_temp_from_state(land_sim, u)
     # required by viz_explorer.jl
     return ClimaLSM.surface_temperature(land_sim.model, u, land_sim.integrator.p, land_sim.integrator.t)
 end
+
+"""
+    dss_state!(sim::BucketSimulation)
+
+Perform DSS on the state of a component simulation, intended to be used
+before the initial step of a run. This method acts on bucket land simulations.
+The `dss!` function of ClimaLSM must be called because it uses either the 2D
+or 3D dss buffer stored in the cache depending on space of each variable in
+`sim.integrator.u`.
+"""
+function dss_state!(sim::BucketSimulation)
+    ClimaLSM.dss!(sim.integrator.u, sim.integrator.p, sim.integrator.t)
+end

experiments/AMIP/modular/components/ocean/slab_seaice_init.jl → experiments/AMIP/modular/components/ocean/prescr_seaice_init.jl

@@ -2,6 +2,8 @@ import ClimaTimeSteppers as CTS
 import ClimaCoupler.Interfacer: SeaIceModelSimulation, get_field, update_field!, name
 import ClimaCoupler.FieldExchanger: step!, reinit!
 import ClimaCoupler.FluxCalculator: update_turbulent_fluxes_point!
+using ClimaCoupler: Regridder
+import Thermodynamics as TD
 
 include("../slab_utils.jl")
 
@@ -111,7 +113,11 @@ function ice_init(::Type{FT}; tspan, saveat, dt, space, area_fraction, thermo_pa
     problem = ODEProblem(ode_function, Y, FT.(tspan), (; additional_cache..., params = params))
     integrator = init(problem, ode_algo, dt = FT(dt), saveat = FT(saveat), adaptive = false)
 
-    PrescribedIceSimulation(params, Y, space, integrator)
+    sim = PrescribedIceSimulation(params, Y, space, integrator)
+
+    # DSS state to ensure we have continuous fields
+    dss_state!(sim)
+    return sim
 end
 
 # file-specific
@@ -178,3 +184,19 @@ get_field(sim::PrescribedIceSimulation, ::Val{:energy}) =
     sim.integrator.p.params.ρ .* sim.integrator.p.params.c .* sim.integrator.u.T_sfc .* sim.integrator.p.params.h
 
 get_field(sim::PrescribedIceSimulation, ::Val{:water}) = nothing
+
+"""
+    dss_state!(sim::PrescribedIceSimulation)
+
+Perform DSS on the state of a component simulation, intended to be used
+before the initial step of a run. This method acts on prescribed ice simulations.
+"""
+function dss_state!(sim::PrescribedIceSimulation)
+    Y = sim.integrator.u
+    p = sim.integrator.p
+    for key in propertynames(Y)
+        field = getproperty(Y, key)
+        buffer = get_dss_buffer(axes(field), p)
+        Spaces.weighted_dss!(field, buffer)
+    end
+end

experiments/AMIP/modular/components/ocean/slab_ocean_init.jl

@@ -104,7 +104,11 @@ function ocean_init(::Type{FT}; tspan, dt, saveat, space, area_fraction, thermo_
     problem = ODEProblem(ode_function, Y, FT.(tspan), cache)
     integrator = init(problem, ode_algo, dt = FT(dt), saveat = FT(saveat), adaptive = false)
 
-    SlabOceanSimulation(params, Y, space, integrator)
+    sim = SlabOceanSimulation(params, Y, space, integrator)
+
+    # DSS state to ensure we have continuous fields
+    dss_state!(sim)
+    return sim
 end
 
 # file specific
@@ -168,3 +172,19 @@ get_field(sim::SlabOceanSimulation, ::Val{:energy}) =
     sim.integrator.p.params.ρ .* sim.integrator.p.params.c .* sim.integrator.u.T_sfc .* sim.integrator.p.params.h
 
 get_field(sim::SlabOceanSimulation, ::Val{:water}) = nothing
+
+"""
+    dss_state!(sim::SlabOceanSimulation)
+
+Perform DSS on the state of a component simulation, intended to be used
+before the initial step of a run. This method acts on slab ocean model sims.
+"""
+function dss_state!(sim::SlabOceanSimulation)
+    Y = sim.integrator.u
+    p = sim.integrator.p
+    for key in propertynames(Y)
+        field = getproperty(Y, key)
+        buffer = get_dss_buffer(axes(field), p)
+        Spaces.weighted_dss!(field, buffer)
+    end
+end

experiments/AMIP/modular/coupler_driver_modular.jl

@@ -110,7 +110,7 @@ include("components/atmosphere/climaatmos_init.jl")
 include("components/land/bucket_init.jl")
 include("components/land/bucket_utils.jl")
 include("components/ocean/slab_ocean_init.jl")
-include("components/ocean/slab_seaice_init.jl")
+include("components/ocean/prescr_seaice_init.jl")
 
 ## helpers for user-specified IO
 include("user_io/user_diagnostics.jl")

test/Project.toml

@@ -2,13 +2,15 @@
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 ArtifactWrappers = "a14bc488-3040-4b00-9dc1-f6467924858a"
 CLIMAParameters = "6eacf6c3-8458-43b9-ae03-caf5306d3d53"
+ClimaAtmos = "b2c96348-7fb7-4fe0-8da9-78d88439e717"
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"
 ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 ClimaCoupler = "4ade58fe-a8da-486c-bd89-46df092ec0c7"
 ClimaLSM = "7884a58f-fab6-4fd0-82bb-ecfedb2d8430"
 ClimaTimeSteppers = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Downloads = "f43a241f-c20a-4ad4-852c-f6b1247861c6"
+Insolation = "e98cc03f-d57e-4e3c-b70c-8d51efe9e0d8"
 IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 MPIPreferences = "3da0fdf6-3ccc-4f1b-acd9-58baa6c99267"

test/component_model_tests/bucket_tests.jl

@@ -0,0 +1,59 @@
+using Test
+import ClimaCoupler
+using ClimaCoupler.TestHelper: create_space
+using ClimaCore: Fields, Spaces
+
+include(pkgdir(ClimaCoupler, "experiments/AMIP/modular/components/land/bucket_init.jl"))
+include(pkgdir(ClimaCoupler, "experiments/AMIP/modular/components/land/bucket_utils.jl"))
+
+# struct DummySimulationBucket{I} <: BucketSimulation
+#     integrator::I
+# end
+
+# TODO bucket doesn't currently work with Float32, but we want to eventually test with both FTs
+for FT in (Float64,)
+    @testset "dss_state! BucketSimulation for FT=$FT" begin
+        # use TestHelper to create space, extract surface space
+        subsurface_space = create_space(FT, nz = 2)
+        surface_space = subsurface_space.horizontal_space
+
+        # set up objects for test
+        dss_buffer_3d = Spaces.create_dss_buffer(Fields.zeros(subsurface_space))
+        dss_buffer_2d = Spaces.create_dss_buffer(Fields.zeros(surface_space))
+
+        integrator = (;
+            u = Fields.FieldVector(
+                state_field1 = Fields.ones(surface_space),
+                state_field_2d = Fields.zeros(surface_space),
+                state_field_3d = Fields.zeros(subsurface_space),
+            ),
+            p = (;
+                cache_field = Fields.zeros(surface_space),
+                dss_buffer_2d = dss_buffer_2d,
+                dss_buffer_3d = dss_buffer_3d,
+            ),
+            t = FT(0),
+        )
+        integrator_copy = deepcopy(integrator)
+        # sim = DummySimulationBucket(integrator)
+        sim = BucketSimulation(nothing, nothing, nothing, integrator, nothing)
+
+        # make fields non-constant to check the impact of the dss step
+        for i in eachindex(parent(sim.integrator.u.state_field_2d))
+            parent(sim.integrator.u.state_field_2d)[i] = sin(i)
+        end
+        for i in eachindex(parent(sim.integrator.u.state_field_3d))
+            parent(sim.integrator.u.state_field_3d)[i] = sin(i)
+        end
+
+        # apply DSS
+        dss_state!(sim)
+
+        # test that uniform field and cache are unchanged, non-constant is changed
+        # note: uniform field is changed slightly by dss
+        @test sim.integrator.u.state_field1 ≈ integrator_copy.u.state_field1
+        @test sim.integrator.u.state_field_2d != integrator_copy.u.state_field_2d
+        @test sim.integrator.u.state_field_3d != integrator_copy.u.state_field_3d
+        @test sim.integrator.p.cache_field == integrator_copy.p.cache_field
+    end
+end

test/component_model_tests/climaatmos_tests.jl

@@ -0,0 +1,36 @@
+using Test
+import ClimaCoupler
+using ClimaCoupler.Interfacer: AtmosModelSimulation
+using ClimaCoupler.TestHelper: create_space
+using ClimaCore: Fields, Spaces
+
+include(pkgdir(ClimaCoupler, "experiments/AMIP/modular/components/atmosphere/climaatmos_init.jl"))
+
+for FT in (Float32, Float64)
+    @testset "dss_state! ClimaAtmosSimulation for FT=$FT" begin
+        # use TestHelper to create space
+        boundary_space = create_space(FT)
+
+        # set up objects for test
+        integrator = (;
+            u = (; state_field1 = FT.(Fields.ones(boundary_space)), state_field2 = FT.(Fields.zeros(boundary_space))),
+            p = (; cache_field = FT.(Fields.zeros(boundary_space))),
+        )
+        integrator_copy = deepcopy(integrator)
+        sim = ClimaAtmosSimulation(nothing, nothing, nothing, integrator)
+
+        # make field non-constant to check the impact of the dss step
+        for i in eachindex(parent(sim.integrator.u.state_field2))
+            parent(sim.integrator.u.state_field2)[i] = FT(sin(i))
+        end
+
+        # apply DSS
+        dss_state!(sim)
+
+        # test that uniform field and cache are unchanged, non-constant is changed
+        # note: uniform field is changed slightly by dss
+        @test sim.integrator.u.state_field1 ≈ integrator_copy.u.state_field1
+        @test sim.integrator.u.state_field2 != integrator_copy.u.state_field2
+        @test sim.integrator.p.cache_field == integrator_copy.p.cache_field
+    end
+end

test/component_model_tests.jl → test/component_model_tests/prescr_seaice_tests.jl

@@ -1,23 +1,18 @@
-# this folder contains temporary tests for component code that has not been moved to any src code (location TBD)
 using Test
+import ClimaCoupler
+using ClimaCoupler.Interfacer: SeaIceModelSimulation
+using ClimaCoupler.TestHelper: create_space
 using ClimaCore
-using ClimaCore: Fields
-using ClimaCoupler.Regridder
-import ClimaCoupler.Regridder: binary_mask
-import ClimaCoupler.Interfacer: AtmosModelSimulation, SurfaceModelSimulation, SurfaceStub, get_field
-import Thermodynamics as TD
+using ClimaCore: Fields, Spaces
 import CLIMAParameters as CP
 import Thermodynamics.Parameters as TDP
 
-include("TestHelper.jl")
-
-# sea ice
-include("../experiments/AMIP/modular/components/ocean/slab_seaice_init.jl")
+include(pkgdir(ClimaCoupler, "experiments/AMIP/modular/components/ocean/prescr_seaice_init.jl"))
 
 for FT in (Float32, Float64)
     @testset "test sea-ice energy slab for FT=$FT" begin
         function test_sea_ice_rhs(; F_radiative = 0.0, T_base = 271.2, global_mask = 1.0)
-            space = TestHelper.create_space(FT)
+            space = create_space(FT)
             params = IceSlabParameters(
                 FT(2),  # ice thickness
                 FT(900.0), # density of sea ice
@@ -80,4 +75,34 @@ for FT in (Float32, Float64)
         dY, Y, p = test_sea_ice_rhs(F_radiative = 0.0, T_base = 269.2, global_mask = 0.0)
         @test sum([i for i in extrema(dY)] .≈ [FT(0.0), FT(0.0)]) == 2
     end
+
+    @testset "dss_state! SeaIceModelSimulation for FT=$FT" begin
+        # use TestHelper to create space
+        boundary_space = create_space(FT)
+
+        # construct dss buffer to put in cache
+        dss_buffer = Spaces.create_dss_buffer(Fields.zeros(boundary_space))
+
+        # set up objects for test
+        integrator = (;
+            u = (; state_field1 = FT.(Fields.ones(boundary_space)), state_field2 = FT.(Fields.zeros(boundary_space))),
+            p = (; cache_field = FT.(Fields.zeros(boundary_space)), dss_buffer = dss_buffer),
+        )
+        integrator_copy = deepcopy(integrator)
+        sim = PrescribedIceSimulation(nothing, nothing, nothing, integrator)
+
+        # make field non-constant to check the impact of the dss step
+        for i in eachindex(parent(sim.integrator.u.state_field2))
+            parent(sim.integrator.u.state_field2)[i] = FT(sin(i))
+        end
+
+        # apply DSS
+        dss_state!(sim)
+
+        # test that uniform field and cache are unchanged, non-constant is changed
+        # note: uniform field is changed slightly by dss
+        @test sim.integrator.u.state_field1 ≈ integrator_copy.u.state_field1
+        @test sim.integrator.u.state_field2 != integrator_copy.u.state_field2
+        @test sim.integrator.p.cache_field == integrator_copy.p.cache_field
+    end
 end

test/component_model_tests/slab_ocean_tests.jl

@@ -0,0 +1,42 @@
+using Test
+import ClimaCoupler
+using ClimaCoupler.Interfacer: OceanModelSimulation
+using ClimaCoupler.TestHelper: create_space
+using ClimaCore
+using ClimaCore: Fields, Spaces
+import CLIMAParameters as CP
+import Thermodynamics.Parameters as TDP
+
+include(pkgdir(ClimaCoupler, "experiments/AMIP/modular/components/ocean/slab_ocean_init.jl"))
+
+for FT in (Float32, Float64)
+    @testset "dss_state! SlabOceanSimulation for FT=$FT" begin
+        # use TestHelper to create space
+        boundary_space = create_space(FT)
+
+        # construct dss buffer to put in cache
+        dss_buffer = Spaces.create_dss_buffer(Fields.zeros(boundary_space))
+
+        # set up objects for test
+        integrator = (;
+            u = (; state_field1 = FT.(Fields.ones(boundary_space)), state_field2 = FT.(Fields.zeros(boundary_space))),
+            p = (; cache_field = FT.(Fields.zeros(boundary_space)), dss_buffer = dss_buffer),
+        )
+        integrator_copy = deepcopy(integrator)
+        sim = SlabOceanSimulation(nothing, nothing, nothing, integrator)
+
+        # make field non-constant to check the impact of the dss step
+        for i in eachindex(parent(sim.integrator.u.state_field2))
+            parent(sim.integrator.u.state_field2)[i] = FT(sin(i))
+        end
+
+        # apply DSS
+        dss_state!(sim)
+
+        # test that uniform field and cache are unchanged, non-constant is changed
+        # note: uniform field is changed slightly by dss
+        @test sim.integrator.u.state_field1 ≈ integrator_copy.u.state_field1
+        @test sim.integrator.u.state_field2 != integrator_copy.u.state_field2
+        @test sim.integrator.p.cache_field == integrator_copy.p.cache_field
+    end
+end

test/interfacer_tests.jl

@@ -75,8 +75,6 @@ end
 end
 
 @testset "update_field! the SurfaceStub area_fraction" begin
-    boundary_space = TestHelper.create_space(FT)
-
     stub = SurfaceStub((; area_fraction = zeros(boundary_space), T_sfc = zeros(boundary_space)))
 
     update_field!(stub, Val(:area_fraction), ones(boundary_space))