Merge pull request #79 from TUM-PIK-ESM/bg/modify-timestep-interface

Refactor timestep! interface to improve flexibility

Author: bgroenks96

src/Terrarium.jl

@@ -100,10 +100,6 @@ export InputSource, InputSources, FieldInputSource, FieldTimeSeriesInputSource
 export update_inputs!
 include("input_output/input_sources.jl")
 
-# timestepping
-export timestep!, default_dt, is_adaptive
-include("timesteppers/abstract_timestepper.jl")
-
 # process/model interface
 export get_grid, get_initializer, variables, processes, compute_auxiliary!, compute_tendencies!
 include("abstract_model.jl")
@@ -124,6 +120,10 @@ include("boundary_conditions.jl")
 export Forcings
 include("forcings.jl")
 
+# timestepping
+export timestep!, default_dt, is_adaptive
+include("timesteppers/abstract_timestepper.jl")
+
 # abstract model types
 include("models/abstract_types.jl")
 
@@ -133,14 +133,14 @@ include("processes/processes.jl")
 # concrete model implementations
 include("models/models.jl")
 
-# timestepper implementations
+# model integrator/simulation types and methods
+export ModelIntegrator, initialize, current_time, iteration
+include("timesteppers/model_integrator.jl")
+
+# Concrete timestepper implementations
 export ForwardEuler
 include("timesteppers/forward_euler.jl")
 export Heun
 include("timesteppers/heun.jl")
 
-# model integrator/simulation types and methods
-export ModelIntegrator, initialize, current_time, iteration
-include("timesteppers/model_integrator.jl")
-
 end # module Terrarium

src/timesteppers/abstract_timestepper.jl

@@ -32,12 +32,20 @@ Return `true` if the given time stepper is adaptive, false otherwise.
 function is_adaptive end
 
 """
-    timestep!(state, timestepper::AbstractTimeStepper, model::AbstractModel, inputs::InputSources, Δt)
+    timestep!(integrator::ModelIntegrator, timestepper::AbstractTimeStepper, Δt)
 
-Advance prognostic variables by one time step based on the current state, or by `Δt` units of time.
+Advance prognostic variables of the `integrator` model by one time step based on the current state, or by `Δt` units of time.
 """
 function timestep! end
 
+"""
+    timestep!(state, model::AbstractModel, timestepper::AbstractTimeStepper, Δt)
+
+Apply any necessary corrections or model-specific time stepping logic after applying `timestepper` to the prognostic state
+variables defined by `model`.
+"""
+timestep!(state, model::AbstractModel, timestepper::AbstractTimeStepper, Δt) = nothing
+
 """
     initialize(::AbstractTimeStepper, model, state) where {NF}
 

src/timesteppers/forward_euler.jl

@@ -16,11 +16,16 @@ is_adaptive(euler::ForwardEuler) = false
 
 is_initialized(euler::ForwardEuler) = true
 
-function timestep!(state, timestepper::ForwardEuler, model::AbstractModel, inputs::InputSources, Δt)
+function timestep!(integrator::ModelIntegrator, timestepper::ForwardEuler, Δt)
+    # Compute auxiliaries and tendencies
+    update_state!(integrator, compute_tendencies = true)
     # Euler step
-    explicit_step!(state, get_grid(model), timestepper, Δt)
+    explicit_step!(integrator.state, get_grid(integrator.model), timestepper, Δt)
+    # Call timestep! on model
+    timestep!(integrator.state, integrator.model, timestepper, Δt)
     # Apply closure relations
-    closure!(state, model)
+    closure!(integrator.state, integrator.model)
     # Update clock
-    return tick!(state.clock, Δt)
+    tick!(integrator.state.clock, Δt)
+    return nothing
 end

src/timesteppers/heun.jl

@@ -31,26 +31,35 @@ function average_tendencies!(
     return
 end
 
-function timestep!(state, timestepper::Heun, model::AbstractModel, inputs::InputSources, Δt = default_dt(timestepper))
+function timestep!(integrator::ModelIntegrator, timestepper::Heun, Δt = default_dt(timestepper))
     @assert is_initialized(timestepper)
 
+    (; model, state, inputs) = integrator
     grid = get_grid(model)
 
+    # Update current state
+    update_state!(state, model, inputs, compute_tendencies = true)
+
     # Copy current state to stage
     stage = timestepper.stage
     copyto!(stage, state)
 
     # Compute stage
     explicit_step!(stage, grid, timestepper, Δt)
+    # Call timestep! on model
+    timestep!(stage, model, timestepper, Δt)
     # Apply closure relations
     closure!(stage, model)
     # Update clock
     tick!(stage.clock, Δt)
+    # Recompute tendencies after timestep
     update_state!(stage, model, inputs, compute_tendencies = true)
 
     # Final improved Euler step call that steps `state` forward but averages `state.tendencies`
     average_tendencies!(state, stage)
     explicit_step!(state, grid, timestepper, Δt)
+    # Call timestep! on model
+    timestep!(state, model, timestepper, Δt)
     # Apply closure relations
     closure!(state, model)
     # Update clock

src/timesteppers/model_integrator.jl

@@ -123,8 +123,7 @@ variables.
 """
 timestep!(integrator::ModelIntegrator; finalize = true) = timestep!(integrator, default_dt(timestepper(integrator)); finalize)
 function timestep!(integrator::ModelIntegrator, Δt; finalize = true)
-    update_state!(integrator, compute_tendencies = true)
-    timestep!(integrator.state, integrator.timestepper, integrator.model, integrator.inputs, convert_dt(Δt))
+    timestep!(integrator, integrator.timestepper, convert_dt(Δt))
     if finalize
         compute_auxiliary!(integrator.state, integrator.model)
     end

test/differentiability/soil_energy_diff.jl

@@ -14,9 +14,10 @@ function build_soil_energy_model(arch, ::Type{NF}) where {NF}
     return model
 end
 
-function mean_soil_temperature_step!(state, timestepper, model, inputs, Δt)
-    timestep!(state, timestepper, model, inputs, Δt)
-    return mean(interior(state.temperature))
+function mean_soil_temperature_step!(integrator, timestepper, Δt)
+
+    timestep!(integrator, timestepper, Δt)
+    return mean(interior(integrator.state.temperature))
     # TODO: Figure out why this is segfaulting in Enzyme
     # Answer: Average operator is not type inferrable, see:
     # https://github.com/CliMA/Oceananigans.jl/issues/4869
@@ -67,11 +68,9 @@ end
 @testset "Soil energy model: timestep!" begin
     model = build_soil_energy_model(CPU(), Float64)
     integrator = initialize(model, ForwardEuler())
-    inputs = integrator.inputs
-    state = integrator.state
-    dstate = make_zero(state)
+    dintegrator = make_zero(integrator)
     stepper = integrator.timestepper
     dstepper = make_zero(stepper)
-    @time Enzyme.autodiff(set_runtime_activity(Reverse), mean_soil_temperature_step!, Active, Duplicated(state, dstate), Duplicated(stepper, dstepper), Const(model), Const(integrator.inputs), Const(integrator.timestepper.Δt))
-    @test all(isfinite.(dstate.temperature))
+    @time Enzyme.autodiff(set_runtime_activity(Reverse), mean_soil_temperature_step!, Active, Duplicated(integrator, dintegrator), Duplicated(stepper, dstepper), Const(integrator.timestepper.Δt))
+    @test all(isfinite.(dintegrator.state.temperature))
 end

test/differentiability/soil_hydrology_diff.jl

@@ -136,13 +136,13 @@ end
     hydraulic_properties = ConstantSoilHydraulics(Float64)
     model = build_soil_energy_hydrology_model(CPU(), Float64; hydraulic_properties)
     integrator = initialize(model, ForwardEuler())
-    inputs = integrator.inputs
-    state = integrator.state
-    dstate = make_zero(state)
+    dintegrator = make_zero(integrator)
+    # set a seed for the temperature
+    dintegrator.state.temperature .= 1.0
     stepper = integrator.timestepper
     dstepper = make_zero(stepper)
     Δt = 60.0
-    @time Enzyme.autodiff(set_runtime_activity(Reverse), timestep!, Const, Duplicated(state, dstate), Duplicated(stepper, dstepper), Const(model), Const(inputs), Const(Δt))
-    @test all(isfinite.(dstate.temperature))
-    @test all(isfinite.(dstate.pressure_head))
+    @time Enzyme.autodiff(set_runtime_activity(Reverse), timestep!, Const, Duplicated(integrator, dintegrator), Duplicated(stepper, dstepper), Const(Δt))
+    @test all(isfinite.(dintegrator.state.temperature))
+    @test all(isfinite.(dintegrator.state.pressure_head))
 end

test/inputs/input_forcing.jl

@@ -23,11 +23,6 @@ function Terrarium.compute_tendencies!(state, model::TestModel)
     return state.tendencies.x .= state.F
 end
 
-function Terrarium.timestep!(state, model::TestModel, euler::ForwardEuler, Δt)
-    Terrarium.compute_tendencies!(state, model)
-    return @. state.x += Δt * state.tendencies.x
-end
-
 @testset "Forcing input" begin
     grid = ColumnGrid(CPU(), DEFAULT_NF, ExponentialSpacing())
     model = TestModel(grid)

test/timestepping/heun.jl

@@ -47,3 +47,21 @@ end
     dt_heun = default_dt(integrator_heun.timestepper)
     @test integrator_heun.state.u[2] == (0.1 * dt_heun + (0.1 * dt_heun + 0.1) * dt_heun) / 2
 end
+
+# Use timestep!(state, model, timestepper, Δt) to clip negative values in an super simple example sim
+@testset "ExpModel: clip negative values" begin
+    grid = ColumnGrid(CPU(), Float64, UniformSpacing(N = 1))
+    model = ExpModel(grid)
+
+    Terrarium.timestep!(state, model::ExpModel, timestepper::ForwardEuler, Δt) = begin
+        state.u[2] = max(state.u[2], 0.0)
+    end
+
+    initializers = (u = -20, v = -5.0)
+    integrator = initialize(model, ForwardEuler(); initializers)
+
+    # Test that timestep! clips negative values
+    timestep!(integrator)
+
+    @test integrator.state.u[2] >= 0.0
+end