initialize! for ReactantSimulation plus better find_time_index (#4215)

* Add support for Nothing callbacks, output writers, diagnostics

* Add a test

* Bump

* Bugfix

* Extend initialize! for ReactantSimulation

* Add find_time_index for StepRangeLen

* Silly reordering

* Try to fix compile problem

* Revert changes to Simulation construcotr

* Bump version

* Fix

* Make linear extrapolation work

* Uncomment tests

* Fix time step import

---------

Co-authored-by: Navid C. Constantinou <[email protected]>

Author: glwagner

Project.toml

@@ -1,7 +1,7 @@
 name = "Oceananigans"
 uuid = "9e8cae18-63c1-5223-a75c-80ca9d6e9a09"
 authors = ["Climate Modeling Alliance and contributors"]
-version = "0.95.27"
+version = "0.95.28"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"

ext/OceananigansReactantExt/OceananigansReactantExt.jl

@@ -22,15 +22,15 @@ using .Fields
 include("TimeSteppers.jl")
 using .TimeSteppers
 
+include("TurbulenceClosures.jl")
+using .TurbulenceClosures
+
 include("Models.jl")
 using .Models
 
 include("Simulations/Simulations.jl")
 using .Simulations
 
-include("TurbulenceClosures.jl")
-using .TurbulenceClosures
-
 #####
 ##### Telling Reactant how to construct types
 #####

ext/OceananigansReactantExt/Simulations/run.jl

@@ -1,4 +1,5 @@
-import ..TimeSteppers: first_time_step!
+import Oceananigans: initialize!
+import Oceananigans.TimeSteppers: first_time_step!
 
 """ Step `sim`ulation forward by one time step. """
 time_step!(sim::ReactantSimulation) = time_step!(sim.model, sim.Δt)
@@ -18,3 +19,11 @@ function time_step_for!(sim::ReactantSimulation, Nsteps)
     return nothing
 end
 
+function initialize!(sim::ReactantSimulation)
+    model = sim.model
+    initialize!(model)
+    update_state!(model)
+    sim.initialized = true
+    return nothing
+end
+

src/OutputReaders/field_time_series_indexing.jl

@@ -68,24 +68,23 @@ end
 ##### fine_time_index
 #####
 
-#=
 @inline function find_time_index(times::StepRangeLen, t)
     n₂ = searchsortedfirst(times, t)
-    n₁ = max(1, n₂ - 1)
 
     Nt = length(times)
     n₂ = min(Nt, n₂) # cap
+    n₁ = max(1, n₂ - 1)
 
     @inbounds begin
         t₁ = times[n₁]
         t₂ = times[n₂]
     end
 
     ñ = (t - t₁) / (t₂ - t₁)
+    ñ = ifelse(n₂ == n₁, zero(ñ), ñ)
 
     return ñ, n₁, n₂
 end
-=#
 
 @inline function find_time_index(times, t)
     Nt = length(times)

test/test_output_readers.jl

@@ -2,7 +2,7 @@ include("dependencies_for_runtests.jl")
 
 using Oceananigans.Utils: Time
 using Oceananigans.Fields: indices, interpolate!
-using Oceananigans.OutputReaders: Cyclical, Clamp
+using Oceananigans.OutputReaders: Cyclical, Clamp, Linear
 
 function generate_some_interesting_simulation_data(Nx, Ny, Nz; architecture=CPU())
     grid = RectilinearGrid(architecture, size=(Nx, Ny, Nz), extent=(64, 64, 32))
@@ -485,9 +485,9 @@ end
 
         sinf(t) = sin(2π * t / 3)
 
-        f   = CenterField(grid) 
         fts = FieldTimeSeries{Center, Center, Center}(grid, times; backend=OnDisk(), path=filepath_sine, name="f")
-        
+
+        f = CenterField(grid) 
         for (i, time) in enumerate(fts.times)
             set!(f, (x, y, z) -> sinf(time))
             set!(fts, f, i)
@@ -496,9 +496,9 @@ end
         # Now we load the FTS partly in memory
         # using different time indexing strategies
         M = 5
-        fts1 = FieldTimeSeries(filepath_sine, "f"; backend = InMemory(M))
-        fts2 = FieldTimeSeries(filepath_sine, "f"; backend = InMemory(M), time_indexing = Cyclical())
-        fts3 = FieldTimeSeries(filepath_sine, "f"; backend = InMemory(M), time_indexing = Clamp())
+        fts_lin = FieldTimeSeries(filepath_sine, "f"; backend = InMemory(M), time_indexing = Linear())
+        fts_cyc = FieldTimeSeries(filepath_sine, "f"; backend = InMemory(M), time_indexing = Cyclical())
+        fts_clp = FieldTimeSeries(filepath_sine, "f"; backend = InMemory(M), time_indexing = Clamp())
 
         # Test that linear interpolation is correct within the time domain
         for time in 0:0.01:last(fts.times)
@@ -509,24 +509,29 @@ end
 
                 Δt⁺ = (time - t⁻) / (t⁺ - t⁻)
 
-                @test fts1[Time(time)][1, 1, 1] ≈ (sinf(t⁻) * (1 - Δt⁺) + sinf(t⁺) * Δt⁺) 
-                @test fts2[Time(time)][1, 1, 1] ≈ (sinf(t⁻) * (1 - Δt⁺) + sinf(t⁺) * Δt⁺) 
-                @test fts3[Time(time)][1, 1, 1] ≈ (sinf(t⁻) * (1 - Δt⁺) + sinf(t⁺) * Δt⁺) 
+                @test fts_lin[Time(time)][1, 1, 1] ≈ (sinf(t⁻) * (1 - Δt⁺) + sinf(t⁺) * Δt⁺) 
+                @test fts_cyc[Time(time)][1, 1, 1] ≈ (sinf(t⁻) * (1 - Δt⁺) + sinf(t⁺) * Δt⁺) 
+                @test fts_clp[Time(time)][1, 1, 1] ≈ (sinf(t⁻) * (1 - Δt⁺) + sinf(t⁺) * Δt⁺) 
             end
         end
 
-        Δt = fts.times[end] - fts.times[end - 1]        
-
         # Test that the time interpolation is correct outside the time domain
-        for time in last(fts.times) + 1 : 0.01 :  last(fts.times) * 2
-            @test fts1[Time(time)][1, 1, 1] ≈ (fts1[end][1, 1, 1] - fts1[end-1][1, 1, 1]) / Δt * (time - last(fts.times))
-            @test fts3[Time(time)][1, 1, 1] ≈ fts3[end][1, 1, 1]
+        Δt = fts.times[end] - fts.times[end-1]        
+        Tf = last(fts.times)
+        from = Tf+1
+        to = 2Tf
+        
+        for t in from:0.01:to
+            dfdt = (fts_lin[end][1, 1, 1] - fts_lin[end-1][1, 1, 1]) / Δt
+            extrapolated = (t - Tf) * dfdt
+            @test fts_lin[Time(t)][1, 1, 1] ≈ extrapolated 
+            @test fts_clp[Time(t)][1, 1, 1] ≈ fts_clp[end][1, 1, 1]
         end
-
     end
 
     rm(filepath1d)
     rm(filepath2d)
     rm(filepath3d)
     rm(filepath_sine)
 end
+