Update coupler stepping logic so that fluxes are exchanged as often as possible

Project.toml

@@ -58,7 +58,7 @@ ClimaCore = "0.14.46"
 ClimaCoreMakie = "0.4"
 ClimaDiagnostics = "0.2.6, 0.3"
 ClimaLand = "1.6"
-ClimaOcean = "0.9"
+ClimaOcean = "0.9.5"
 ClimaParams = "1.0"
 ClimaSeaIce = "0.4.0 - 0.4.5" # [0.4.0, 0.4.5]
 ClimaTimeSteppers = "0.7, 0.8"

config/ci_configs/cmip_oceananigans_climaseaice.yml

@@ -3,7 +3,7 @@ albedo_model: "CouplerAlbedo"
 atmos_config_file: "config/atmos_configs/climaatmos_edonly.yml"
 coupler_toml: ["toml/amip_edonly.toml"]
 dt_atmos: "120secs" # 2 minutes
-dt_cpl: "1800secs" # 30 minutes
+dt_cpl: "360secs" # 6 minutes
 dt_land: "360secs" # 6 minutes
 dt_ocean: "1800secs" # 30 minutes
 dt_seaice: "1800secs" # 30 minutes

config/ci_configs/cmip_oceananigans_climaseaice_bucket.yml

@@ -4,7 +4,7 @@ atmos_config_file: "config/atmos_configs/climaatmos_edonly.yml"
 bucket_albedo_type: "map_temporal"
 coupler_toml: ["toml/amip_edonly.toml"]
 dt_atmos: "120secs" # 2 minutes
-dt_cpl: "1800secs" # 30 minutes
+dt_cpl: "360secs" # 6 minutes
 dt_land: "360secs" # 6 minutes
 dt_ocean: "1800secs" # 30 minutes
 dt_seaice: "1800secs" # 30 minutes

ext/ClimaCouplerCMIPExt/clima_seaice.jl

@@ -5,6 +5,7 @@ import ClimaOcean.EN4: download_dataset
 import SurfaceFluxes as SF
 import SurfaceFluxes.Parameters as SFP
 import Thermodynamics as TD
+import Dates
 using StaticArrays
 
 # Rename ECCO password env variable to match ClimaOcean.jl
@@ -26,12 +27,14 @@ It contains the following objects:
 - `ice_properties::IP`: A NamedTuple of sea ice properties, including melting speed, Stefan-Boltzmann constant,
     and the Celsius to Kelvin conversion constant.
 """
-struct ClimaSeaIceSimulation{SIM, A, REMAP, NT, IP} <: Interfacer.AbstractSeaIceSimulation
+struct ClimaSeaIceSimulation{SIM, A, REMAP, NT, IP, MDT} <:
+       Interfacer.AbstractSeaIceSimulation
     ice::SIM
     area_fraction::A
     remapping::REMAP
     ocean_ice_interface::NT
     ice_properties::IP
+    model_Δt::MDT
 end
 
 """
@@ -77,13 +80,22 @@ function ClimaSeaIceSimulation(
     arch = OC.Architectures.architecture(grid)
 
     advection = ocean.ocean.model.advection.T
-    ice = CO.SeaIces.sea_ice_simulation(grid, ocean.ocean; Δt = float(dt), advection)
+    ice = CO.SeaIces.sea_ice_simulation(
+        grid,
+        ocean.ocean;
+        clock = deepcopy(ocean.ocean.model.clock),
+        Δt = float(dt),
+        advection,
+    )
 
     ocean_ice_flux_formulation =
         CO.OceanSeaIceModels.InterfaceComputations.ThreeEquationHeatFlux(ice)
     interface_temperature = OC.Field{OC.Center, OC.Center, Nothing}(grid)
     interface_salinity = OC.Field{OC.Center, OC.Center, Nothing}(grid)
 
+    # Initialize model_Δt so that time stepping works properly
+    model_Δt = dt
+
     # Initialize nonzero sea ice if start date provided
     if !isnothing(start_date)
         sic_metadata = CO.DataWrangling.Metadatum(
@@ -161,6 +173,7 @@ function ClimaSeaIceSimulation(
         remapping,
         ocean_ice_interface,
         ice_properties,
+        model_Δt,
     )
 
     # Ensure ocean temperature is above freezing where there is sea ice
@@ -173,8 +186,22 @@ end
 ###############################################################################
 
 # Timestep the simulation forward to time `t`
-Interfacer.step!(sim::ClimaSeaIceSimulation, t) =
-    OC.time_step!(sim.ice, float(t) - sim.ice.model.clock.time)
+function Interfacer.step!(sim::ClimaSeaIceSimulation, t::Float64)
+    Δt = t - sim.ice.model.clock.time
+    if isapprox(Δt, sim.model_Δt) || Δt > sim.model_Δt
+        OC.time_step!(sim.ice, Δt)
+    end
+    return nothing
+end
+
+function Interfacer.step!(sim::ClimaSeaIceSimulation, t::ITime)
+    Δt_msec = Dates.DateTime(t) - sim.ice.model.clock.time
+    model_Δt_msec = Dates.DateTime(sim.model_Δt) - sim.model_Δt.epoch
+    if Δt_msec >= model_Δt_msec
+        OC.time_step!(sim.ice, float(sim.model_Δt))
+    end
+    return nothing
+end
 
 Interfacer.get_field(sim::ClimaSeaIceSimulation, ::Val{:area_fraction}) = sim.area_fraction
 Interfacer.get_field(sim::ClimaSeaIceSimulation, ::Val{:ice_concentration}) =

ext/ClimaCouplerCMIPExt/oceananigans.jl

@@ -2,6 +2,7 @@ import ClimaComms
 import SurfaceFluxes as SF
 import Thermodynamics as TD
 import ClimaOcean.EN4: download_dataset
+import ClimaUtilities.TimeManager: ITime
 import Dates
 
 """
@@ -18,13 +19,14 @@ It contains the following objects:
 - `remapping::REMAP`: Objects needed to remap from the exchange (spectral) grid to Oceananigans spaces.
 - `ice_concentration::SIC`: An Oceananigans Field representing the sea ice concentration on the ocean/sea ice grid.
 """
-struct OceananigansSimulation{SIM, A, OPROP, REMAP, SIC} <:
+struct OceananigansSimulation{SIM, A, OPROP, REMAP, SIC, MDT} <:
        Interfacer.AbstractOceanSimulation
     ocean::SIM
     area_fraction::A
     ocean_properties::OPROP
     remapping::REMAP
     ice_concentration::SIC
+    model_Δt::MDT
 end
 
 """
@@ -67,7 +69,7 @@ function OceananigansSimulation(
     tspan,
     output_dir,
     simple_ocean = false,
-    dt = nothing,
+    dt = 1800.0, # 30 minutes
     comms_ctx = ClimaComms.context(),
     coupled_param_dict = CP.create_toml_dict(FT),
     extra_kwargs...,
@@ -156,10 +158,22 @@ function OceananigansSimulation(
         closure = (horizontal_viscosity, vertical_mixing)
     end
 
-    Δt = isnothing(dt) ? CO.OceanSimulations.estimate_maximum_Δt(grid) : float(dt)
+    if tspan[1] isa ITime
+        # create a model clock that uses DateTime, for compatibility with ITime.
+        model_clock = OC.TimeSteppers.Clock(time = start_date)
+        stop_time = Dates.DateTime(tspan[2])
+    elseif tspan[1] isa Float64
+        model_clock = OC.TimeSteppers.Clock{Float64}(time = tspan[1])
+        stop_time = tspan[2]
+    else
+        error("Unsupported time type: $(typeof(tspan[1]))")
+    end
+    model_Δt = dt
     ocean = CO.ocean_simulation(
         grid;
-        Δt,
+        clock = model_clock,
+        stop_time,
+        Δt = float(model_Δt),
         timestepper = :SplitRungeKutta3,
         momentum_advection,
         tracer_advection,
@@ -287,6 +301,7 @@ function OceananigansSimulation(
         ocean_properties,
         remapping,
         ice_concentration,
+        model_Δt,
     )
 end
 
@@ -335,9 +350,23 @@ end
 ### Functions required by ClimaCoupler.jl for a AbstractSurfaceSimulation
 ###############################################################################
 
-# Timestep the simulation forward to time `t`
-Interfacer.step!(sim::OceananigansSimulation, t) =
-    OC.time_step!(sim.ocean, float(t) - sim.ocean.model.clock.time)
+# Timestep the simulation forward to time `t`. This may not actually do anything.
+function Interfacer.step!(sim::OceananigansSimulation, t::Float64)
+    Δt = t - sim.ocean.model.clock.time
+    if isapprox(Δt, sim.model_Δt) || Δt > sim.model_Δt
+        OC.time_step!(sim.ocean, Δt)
+    end
+    return nothing
+end
+
+function Interfacer.step!(sim::OceananigansSimulation, t::ITime)
+    Δt_msec = Dates.DateTime(t) - sim.ocean.model.clock.time
+    model_Δt_msec = Dates.DateTime(sim.model_Δt) - sim.model_Δt.epoch
+    if Δt_msec >= model_Δt_msec
+        OC.time_step!(sim.ocean, float(sim.model_Δt))
+    end
+    return nothing
+end
 
 Interfacer.get_field(sim::OceananigansSimulation, ::Val{:area_fraction}) = sim.area_fraction
 

ext/ClimaCouplerClimaAtmosExt.jl

@@ -513,13 +513,26 @@ function FluxCalculator.update_turbulent_fluxes!(sim::ClimaAtmosSimulation, fiel
 end
 
 # extensions required by FieldExchanger
-Interfacer.step!(sim::ClimaAtmosSimulation, t::Real) =
-    Interfacer.step!(sim.integrator, t - sim.integrator.t, true)
+function Interfacer.step!(sim::ClimaAtmosSimulation, t::Float64)
+    model_t = Float64(sim.integrator.t)
+    Δt = t - model_t
+    model_dt = Float64(sim.integrator.dt)
+    if isapprox(Δt, model_dt) || Δt > model_dt
+        while Float64(sim.integrator.t) < t
+            Interfacer.step!(sim.integrator, Δt, true)
+        end
+    end
+end
+
 function Interfacer.step!(sim::ClimaAtmosSimulation, t::ITime)
-    while sim.integrator.t < t
-        Interfacer.step!(sim.integrator)
+    # Don't step until we've reached a step boundary
+    # (This can happen if the coupler dt is less than this model's)
+    Δt = t - sim.integrator.t
+    if Δt >= sim.integrator.dt
+        while sim.integrator.t < t
+            Interfacer.step!(sim.integrator)
+        end
     end
-    return nothing
 end
 
 """

ext/ClimaCouplerClimaLandExt/climaland_bucket.jl

@@ -311,12 +311,30 @@ function Interfacer.update_field!(
     Interfacer.remap!(sim.integrator.p.bucket.turbulent_fluxes.vapor_flux, field ./ ρ_liq) # TODO: account for sublimation
 end
 
-function Interfacer.step!(sim::BucketSimulation, t)
-    while float(sim.integrator.t) < float(t)
-        Interfacer.step!(sim.integrator)
+# Don't step if we haven't reached a step boundary
+# (This can happen if the coupler dt is less than this model's)
+function Interfacer.step!(sim::BucketSimulation, t::Float64)
+    model_t = Float64(sim.integrator.t)
+    Δt = t - model_t
+    model_dt = Float64(sim.integrator.dt)
+    if isapprox(Δt, model_dt) || Δt > model_dt
+        while Float64(sim.integrator.t) < t
+            Interfacer.step!(sim.integrator)
+        end
     end
     return nothing
 end
+
+function Interfacer.step!(sim::BucketSimulation, t::ITime)
+    Δt = t - sim.integrator.t
+    if Δt >= sim.integrator.dt
+        while sim.integrator.t < t
+            Interfacer.step!(sim.integrator)
+        end
+    end
+    return nothing
+end
+
 Interfacer.close_output_writers(sim::BucketSimulation) =
     isnothing(sim.output_writer) || close(sim.output_writer)
 

ext/ClimaCouplerClimaLandExt/climaland_integrated.jl

@@ -378,12 +378,29 @@ function Interfacer.update_field!(
         StaticArrays.SVector.(sim.integrator.p.scratch1, sim.integrator.p.scratch2)
 end
 
-function Interfacer.step!(sim::ClimaLandSimulation, t)
-    while float(sim.integrator.t) < float(t)
-        Interfacer.step!(sim.integrator)
+# Don't step if we haven't reached a step boundary
+# (This can happen if the coupler dt is less than this model's)
+function Interfacer.step!(sim::ClimaLandSimulation, t::ITime)
+    Δt = t - sim.integrator.t
+    if Δt >= sim.integrator.dt
+        while sim.integrator.t < t
+            Interfacer.step!(sim.integrator)
+        end
     end
     return nothing
 end
+function Interfacer.step!(sim::ClimaLandSimulation, t::Float64)
+    model_t = Float64(sim.integrator.t)
+    Δt = t - model_t
+    model_dt = Float64(sim.integrator.dt)
+    if isapprox(Δt, model_dt) || Δt > model_dt
+        while Float64(sim.integrator.t) < t
+            Interfacer.step!(sim.integrator)
+        end
+    end
+    return nothing
+end
+
 Interfacer.close_output_writers(sim::ClimaLandSimulation) =
     isnothing(sim.output_writer) || close(sim.output_writer)
 

src/Input.jl

@@ -678,7 +678,23 @@ function parse_component_dts!(config_dict)
         end
         for key in component_dt_names
             component_dt = Float64(Utilities.time_to_seconds(config_dict[key]))
-            @assert isapprox(Δt_cpl % component_dt, 0.0) "Coupler dt must be divisible by all component dt's\n dt_cpl = $Δt_cpl\n $key = $component_dt"
+            if key == "dt_atmos"
+                # ensure either that the coupler dt is an integer multiple of the atmos dt
+                # or that the atmos dt is an integer multiple of the coupler dt,
+                # to ensure consistent coupling time steps and compability with legacy configs
+                assertion =
+                    isapprox(Δt_cpl % component_dt, 0.0) ||
+                    isapprox(component_dt % Δt_cpl, 0.0)
+                @assert assertion "Coupler time step and atmosphere time step must be integer multiples of each other\n dt_cpl = $Δt_cpl\n $key = $component_dt"
+            else
+                # all other (surface) model dts must be divisible by the coupler dt,
+                # or the coupler time step must be divisible by the surface model dt,
+                # to ensure consistent coupling time steps and compability with legacy configs
+                assertion =
+                    isapprox(Δt_cpl % component_dt, 0.0) ||
+                    isapprox(component_dt % Δt_cpl, 0.0)
+                @assert assertion "All surface component dts and coupler dtmust be integer multiples of each other\n $key = $component_dt\n dt_cpl = $Δt_cpl"
+            end
             component_dt_dict[key] = component_dt
         end
     else