set! for time step

CHANGELOG.md

@@ -2,6 +2,7 @@
 
 ## Unreleased
 
+- set! for time step [#650](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/650)
 - Bug: ZonalWind had sqrt of negative for uncommon resolutions [#649](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/649)
 - Spectral Gradients are now differentiable, with correctness check in extended CI [#638](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/638)[#651](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/651)
 - Tracer advection [#579](https://github.com/SpeedyWeather/SpeedyWeather.jl/pull/579)

src/dynamics/time_integration.jl

@@ -114,14 +114,37 @@ function initialize!(L::Leapfrog, model::AbstractModel)
         L.Δt = L.Δt_sec/L.radius
     end
 
-    # check how time stepping time step and output time step align
+    # check how time steps from time integration and output align
     n = round(Int, Millisecond(output_dt).value/L.Δt_millisec.value)
     nΔt = n*L.Δt_millisec
     if nΔt != output_dt
         @warn "$n steps of Δt = $(L.Δt_millisec.value)ms yield output every $(nΔt.value)ms (=$(nΔt.value/1000)s), but output_dt = $(output_dt.value)s"
     end
 end
 
+"""$(TYPEDSIGNATURES)
+Change time step of timestepper `L` to `Δt` (unscaled)
+and disables adjustment to output frequency."""
+function set!(
+    L::AbstractTimeStepper,
+    Δt::Period,                 # unscaled time step in Second, Minute, ...
+)
+    L.Δt_millisec = Millisecond(Δt)         # recalculate all Δt fields
+    L.Δt_sec = L.Δt_millisec.value/1000
+    L.Δt = L.Δt_sec/L.radius
+
+    # recalculate the default time step at resolution T31 to be consistent
+    resolution_factor = (L.trunc+1)/(DEFAULT_TRUNC+1)
+    L.Δt_at_T31 = Second(round(Int, L.Δt_sec*resolution_factor))
+
+    # given Δt was manually set disallow adjustment to output frequency
+    L.adjust_with_output = false
+    return L
+end
+
+# also allow for keyword arguments
+set!(L::AbstractTimeStepper; Δt::Period) = set!(L, Δt)
+
 """
 $(TYPEDSIGNATURES)
 Performs one leapfrog time step with (`lf=2`) or without (`lf=1`) Robert+Williams filter

test/time_stepping.jl

@@ -114,4 +114,16 @@ end
         @test M_Ronly < 1
         @test M_Ronly <= M_RAW
     end
+end
+
+@testset "Set timestep manually" begin
+    @testset for trunc in (31, 63, 127)
+        spectral_grid = SpectralGrid(; trunc)
+        time_stepping = Leapfrog(spectral_grid)
+        set!(time_stepping, Minute(10))
+        @test time_stepping.Δt_sec == 10*60
+
+        set!(time_stepping, Δt=Minute(20))
+        @test time_stepping.Δt_sec == 20*60
+    end
 end