Add remapping to coupler puts & gets

Project.toml

@@ -4,4 +4,5 @@ authors = ["CliMA Contributors <[email protected]>"]
 version = "0.1.0"
 
 [deps]
+ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 PrettyTables = "08abe8d2-0d0c-5749-adfa-8a2ac140af0d"

experiments/ClimaCore/sea_breeze/run.jl

@@ -86,13 +86,17 @@ lnd_Y_default, lnd_domain = lnd_init(xmin = 0, xmax = 500, helem = 10, npoly = 0
 # Build remapping operators
 atm_boundary = Spaces.level(atm_domain.hv_face_space, PlusHalf(0))
 
-atm_to_ocn = Operators.LinearRemap(ocn_domain, atm_boundary)
-atm_to_lnd = Operators.LinearRemap(lnd_domain, atm_boundary)
-ocn_to_atm = Operators.LinearRemap(atm_boundary, ocn_domain)
-lnd_to_atm = Operators.LinearRemap(atm_boundary, lnd_domain)
+maps = (
+    atmos_to_ocean = Operators.LinearRemap(ocn_domain, atm_boundary),
+    atmos_to_land = Operators.LinearRemap(lnd_domain, atm_boundary),
+    ocean_to_atmos = Operators.LinearRemap(atm_boundary, ocn_domain),
+    land_to_atmos = Operators.LinearRemap(atm_boundary, lnd_domain),
+)
 
 # initialize coupling fields
-atm_T_sfc = Operators.remap(ocn_to_atm, ocn_Y_default.T_sfc) .+ Operators.remap(lnd_to_atm, lnd_Y_default.T_sfc) # masked arrays; regrid to atm grid
+atm_T_sfc =
+    Operators.remap(maps.ocean_to_atmos, ocn_Y_default.T_sfc) .+
+    Operators.remap(maps.land_to_atmos, lnd_Y_default.T_sfc) # masked arrays; regrid to atm grid
 atm_F_sfc = Fields.zeros(atm_boundary)
 ocn_F_sfc = Fields.zeros(ocn_domain)
 lnd_F_sfc = Fields.zeros(lnd_domain)
@@ -111,18 +115,35 @@ lnd_p = (cpl_parameters, F_sfc = lnd_F_sfc)
 land = LandSimulation(lnd_Y, t_start, cpl_Δt / lnd_nsteps, t_end, SSPRK33(), lnd_p, saveat)
 
 # coupled simulation
-struct AOLCoupledSimulation{FT, A <: AtmosSimulation, O <: OceanSimulation, L <: LandSimulation} <:
-       ClimaCoupler.AbstractCoupledSimulation
+struct AOLCoupledSimulation{
+    FT,
+    A <: AtmosSimulation,
+    O <: OceanSimulation,
+    L <: LandSimulation,
+    C <: ClimaCoupler.CouplerState,
+} <: ClimaCoupler.AbstractCoupledSimulation
     # Atmosphere Simulation
     atmos::A
     # Ocean Simulation
     ocean::O
     # Land Simulation
     land::L
+    # Coupler storage
+    coupler::C
     # The coupled time step size
     Δt::FT
 end
-sim = AOLCoupledSimulation(atmos, ocean, land, cpl_Δt)
+
+# init coupler fields and maps
+coupler = CouplerState()
+coupler_add_field!(coupler, :T_sfc_ocean, ocean.integrator.u.T_sfc; write_sim = ocean)
+coupler_add_field!(coupler, :T_sfc_land, land.integrator.u.T_sfc; write_sim = land)
+coupler_add_field!(coupler, :F_sfc, atmos.integrator.u.F_sfc; write_sim = atmos)
+for (name, map) in pairs(maps)
+    coupler_add_map!(coupler, name, map)
+end
+
+sim = AOLCoupledSimulation(atmos, ocean, land, coupler, cpl_Δt)
 
 # step for sims built on OrdinaryDiffEq
 function step!(sim::ClimaCoupler.AbstractSimulation, t_stop)
@@ -140,30 +161,35 @@ function cpl_run(simulation::AOLCoupledSimulation)
         ## Atmos
         # pre: reset flux accumulator
         atmos.integrator.u.F_sfc .= 0.0 # reset surface flux to be accumulated
+        # don't want to alloc here..
+        T_sfc_ocean = coupler_get(coupler, :T_sfc_ocean, atmos)
+        T_sfc_land = coupler_get(coupler, :T_sfc_land, atmos)
+        atmos.integrator.p.T_sfc .= T_sfc_land .+ T_sfc_ocean
 
-        # run 
+        # run
         # NOTE: use (t - integ_atm.t) here instead of Δt_cpl to avoid accumulating roundoff error in our timestepping.
         step!(atmos, t)
+        coupler_put!(coupler, :F_sfc, atmos.integrator.u.F_sfc, atmos)
 
         ## Ocean
         # pre: get accumulated flux from atmos
         ocn_F_sfc = ocean.integrator.p.F_sfc
-        Operators.remap!(ocn_F_sfc, atm_to_ocn, atmos.integrator.u.F_sfc ./ cpl_Δt)
+        ocn_F_sfc .= coupler_get(coupler, :F_sfc, ocean) ./ cpl_Δt
 
         # run
         step!(ocean, t)
         # post: send ocean surface temp to atmos
-        Operators.remap!(atmos.integrator.p.T_sfc, ocn_to_atm, ocean.integrator.u.T_sfc)
+        coupler_put!(coupler, :T_sfc_ocean, ocean.integrator.u.T_sfc, ocean)
 
         ## Land
         # pre: get accumulated flux from atmos
         lnd_F_sfc = land.integrator.p.F_sfc
-        Operators.remap!(lnd_F_sfc, atm_to_lnd, atmos.integrator.u.F_sfc ./ cpl_Δt)
+        lnd_F_sfc .= coupler_get(coupler, :F_sfc, land) ./ cpl_Δt
 
         # run
         step!(land, t)
-        # post: send ocean surface temp to atmos
-        atmos.integrator.p.T_sfc .+= Operators.remap(lnd_to_atm, land.integrator.u.T_sfc)
+        # post: send land surface temp to atmos
+        coupler_put!(coupler, :T_sfc_land, land.integrator.u.T_sfc, land)
     end
     @info "Simulation Complete"
 end

src/CouplerState/coupler_state.jl

@@ -1,16 +1,22 @@
+using ClimaCore.Operators
 using PrettyTables
 
 export CouplerState
 export coupler_push!, coupler_pull!, coupler_put!, coupler_get
-export coupler_add_field!
-
-mutable struct CplFieldInfo{AT, MD}
-    data::AT
+export coupler_add_field!, coupler_add_map!
+
+mutable struct CplFieldInfo{DT, MD}
+    # the coupled data
+    data::DT
+    # the name of the model that has permission to write to data
+    write_sim::Symbol
+    # catch-all metadata container
     metadata::MD
 end
 
-mutable struct CouplerState{DT}
-    coupled_fields::DT
+mutable struct CouplerState{CF, RO}
+    coupled_fields::CF
+    remap_operators::RO
 end
 
 _fields(coupler::CouplerState) = getfield(coupler, :coupled_fields)
@@ -26,7 +32,7 @@ etc... can be embeded in the intermdediate coupling layer.
 A field is exported by one component and imported by one or more other components.
 """
 function CouplerState()
-    return CouplerState(Dict{Symbol, CplFieldInfo}())
+    return CouplerState(Dict{Symbol, CplFieldInfo}(), Dict{Symbol, Operators.LinearRemap}())
 end
 
 """
@@ -43,17 +49,42 @@ Add a field to the coupler that is accessible with key `fieldname`.
 - `fieldname`: key to access the field in the coupler.
 - `fieldvalue`: data array of field values.
 """
-function coupler_add_field!(coupler::CouplerState, fieldname::Symbol, fieldvalue, metadata = nothing)
-    push!(coupler.coupled_fields, fieldname => CplFieldInfo(fieldvalue, metadata))
+function coupler_add_field!(
+    coupler::CouplerState,
+    fieldname::Symbol,
+    fieldvalue;
+    write_sim::AbstractSimulation,
+    metadata = nothing,
+)
+    push!(coupler.coupled_fields, fieldname => CplFieldInfo(fieldvalue, name(write_sim), metadata))
+end
+
+"""
+    coupler_add_map!(
+            coupler::CouplerState,
+            map_name::Symbol,
+            map::Operators.LinearRemap
+        )
+
+Add a map to the coupler that is accessible with key `mapname`. 
+
+# Arguments
+- `coupler`: coupler object the field is added to.
+- `mapname`: key to access the map in the coupler's map list.
+- `map`: a remap operator.
+"""
+function coupler_add_map!(coupler::CouplerState, map_name::Symbol, map::Operators.LinearRemap)
+    push!(coupler.remap_operators, map_name => map)
 end
 
 """
     coupler_put!(coupler::CouplerState, fieldname::Symbol, fieldvalue)
 
 Sets coupler field `fieldname` to `fieldvalue`.
 """
-function coupler_put!(coupler::CouplerState, fieldname::Symbol, fieldvalue)
+function coupler_put!(coupler::CouplerState, fieldname::Symbol, fieldvalue, source_sim::AbstractSimulation)
     cplfield = coupler.coupled_fields[fieldname]
+    @assert cplfield.write_sim == name(source_sim) "$fieldname can only be written to by $(cplfield.write_sim)."
 
     cplfield.data .= fieldvalue
 
@@ -80,9 +111,15 @@ Retrieve data array corresponding to `fieldname`.
 Returns data on the grid specified by `gridinfo` and in the units of `units`. Checks that
 the coupler data field is the state at time `datetime`.
 """
-function coupler_get(coupler::CouplerState, fieldname::Symbol)
+function coupler_get(coupler::CouplerState, fieldname::Symbol, target_sim::AbstractSimulation)
     cplfield = coupler.coupled_fields[fieldname]
 
+    map = get_remap_operator(coupler, name(target_sim), cplfield.write_sim)
+    return Operators.remap(map, cplfield.data)
+end
+
+function coupler_get(coupler::CouplerState, fieldname::Symbol)
+    cplfield = coupler.coupled_fields[fieldname]
     return cplfield.data
 end
 
@@ -98,6 +135,11 @@ them for use in the component model.
 """
 function coupler_pull!(model, coupler::CouplerState) end
 
+function get_remap_operator(coupler, target_sim_name::Symbol, source_sim_name::Symbol)
+    op_name = Symbol(source_sim_name, "_to_", target_sim_name)
+    return coupler.remap_operators[op_name]
+end
+
 # display table of registered fields & their info
 function Base.show(io::IO, coupler::CouplerState)
     #=

test/CouplerState/cplstate_interface.jl

@@ -1,37 +1,77 @@
 using Test
 using Random
 using ClimaCoupler, Dates, Unitful
+using IntervalSets
+using ClimaCore: Domains, Meshes, Geometry, Topologies, Spaces, Fields, Operators
+
+struct SimulationA{T} <: ClimaCoupler.AbstractSimulation
+    data::T
+end
+ClimaCoupler.name(::SimulationA) = :simA
+
+struct SimulationB{T} <: ClimaCoupler.AbstractSimulation
+    data::T
+end
+ClimaCoupler.name(::SimulationB) = :simB
+
+function spectral_space_2D(; n1 = 1, n2 = 1, Nij = 4)
+    domain = Domains.RectangleDomain(
+        Geometry.XPoint(-1.0) .. Geometry.XPoint(1.0),
+        Geometry.YPoint(-1.0) .. Geometry.YPoint(1.0),
+        x1periodic = false,
+        x2periodic = false,
+        x1boundary = (:east, :west),
+        x2boundary = (:south, :north),
+    )
+    mesh = Meshes.RectilinearMesh(domain, n1, n2)
+    grid_topology = Topologies.Topology2D(mesh)
+
+    quad = Spaces.Quadratures.GLL{Nij}()
+    space = Spaces.SpectralElementSpace2D(grid_topology, quad)
+    return space
+end
 
 @testset "Coupler Interface" begin
-    Random.seed!(26)
 
+    spaceA = spectral_space_2D()
+    spaceB = spectral_space_2D(n1 = 2, n2 = 2)
+
+    simA = SimulationA(ones(spaceA))
+    simB = SimulationB(zeros(spaceB))
     coupler = CouplerState()
 
-    data = rand(10, 10)
-    coupler_add_field!(coupler, :test1, data)
-    coupler_add_field!(coupler, :test2, data)
-    coupler_add_field!(coupler, :test3, data)
+    coupler_add_field!(coupler, :test1, simA.data; write_sim = simA)
+
+    map = Operators.LinearRemap(spaceB, spaceA)
+    coupler_add_map!(coupler, :simA_to_simB, map)
 
     @show coupler
 
     @testset "coupler_get" begin
-        @test data === coupler_get(coupler, :test1)
+        @test simA.data === coupler_get(coupler, :test1)
+
+        # test remapping
+        @test map === ClimaCoupler.get_remap_operator(coupler, ClimaCoupler.name(simB), ClimaCoupler.name(simA))
+        @test ones(spaceB) ≈ coupler_get(coupler, :test1, simB)
 
         # key not in coupler dict
         @test_throws KeyError coupler_get(coupler, :idontexist)
     end
 
     @testset "coupler_put!" begin
-        newdata = rand(10, 10)
-        coupler_put!(coupler, :test1, newdata)
-
+        newdata = zeros(spaceA)
+        # simA can write to :test1
+        coupler_put!(coupler, :test1, newdata, simA)
         @test newdata == coupler_get(coupler, :test1)
         # coupler_put! is in-place; original data array has been modified
-        @test data === coupler_get(coupler, :test1)
+        @test simA.data === coupler_get(coupler, :test1)
+
+        # simB cannot write to :test1
+        @test_throws AssertionError coupler_put!(coupler, :test1, newdata, simB)
 
-        # coupler_put! must be to a previously add_fielded field
-        @test_throws KeyError coupler_put!(coupler, :idontexist, newdata)
-        # incoming data must match dimensions of add_fielded field
-        @test_throws DimensionMismatch coupler_put!(coupler, :test1, rand(10, 5))
+        # coupler_put! must be to a previously added field
+        @test_throws KeyError coupler_put!(coupler, :idontexist, newdata, simA)
+        # incoming data must match dimensions/space of added field
+        @test_throws ErrorException coupler_put!(coupler, :test1, ones(spaceB), simA)
     end
 end

test/Project.toml

@@ -1,5 +1,7 @@
 [deps]
+ClimaCore = "d414da3d-4745-48bb-8d80-42e94e092884"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
+IntervalSets = "8197267c-284f-5f27-9208-e0e47529a953"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"