Update saveat (#2174)

ph-kev · charleskawczynski · web-flow · commit c6fac9bd81bb · 2025-02-10T13:34:11.000-08:00
* Update saveat

With ClimaTimeSteppers 0.8.2, you cannot pass saveat as a scalar.

* Define zero for SchurComplementW

* Fix

---------

Co-authored-by: Charles Kawczynski &lt;kawczynski.charles@gmail.com&gt;
diff --git a/docs/tutorials/introduction.jl b/docs/tutorials/introduction.jl
@@ -390,7 +390,8 @@ function heat_fd_tendency!(dydt, y, α, t)
 end
 
 heat_fd_prob = ODEProblem(heat_fd_tendency!, y0, (0.0, 5.0), 0.1)
-heat_fd_sol = solve(heat_fd_prob, SSPRK33(), dt = 0.1, saveat = 0.25)
+heat_fd_sol =
+    solve(heat_fd_prob, SSPRK33(), dt = 0.1, saveat = collect(0.0:0.25:5.0))
 #----------------------------------------------------------------------------
 
 anim = Plots.@animate for u in heat_fd_sol.u
@@ -417,7 +418,8 @@ end
 y0 = exp.(.-(coord.y .^ 2 .+ coord.x .^ 2) ./ 2)
 
 heat_cg_prob = ODEProblem(heat_cg_tendency!, y0, (0.0, 5.0), 0.1)
-heat_cg_sol = solve(heat_cg_prob, SSPRK33(), dt = 0.1, saveat = 0.5)
+heat_cg_sol =
+    solve(heat_cg_prob, SSPRK33(), dt = 0.1, saveat = collect(0.0:0.5:5.0))
 #----------------------------------------------------------------------------
 
 anim = Plots.@animate for u in heat_cg_sol.u
@@ -535,8 +537,12 @@ end
 #----------------------------------------------------------------------------
 
 shallow_water_prob = ODEProblem(shallow_water_tendency!, y0, (0.0, 20.0))
-@time shallow_water_sol =
-    solve(shallow_water_prob, SSPRK33(), dt = 0.05, saveat = 1.0)
+@time shallow_water_sol = solve(
+    shallow_water_prob,
+    SSPRK33(),
+    dt = 0.05,
+    saveat = collect(0.0:1.0:20.0),
+)
 anim = Plots.@animate for u in shallow_water_sol.u
     Plots.plot(u.ρθ, clim = (-1, 1))
 end
diff --git a/examples/bickleyjet/bickleyjet_cg.jl b/examples/bickleyjet/bickleyjet_cg.jl
@@ -127,7 +127,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = 0.02,
-    saveat = 1.0,
+    saveat = collect(0.0:1.0:80.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 )
diff --git a/examples/bickleyjet/bickleyjet_cg_invariant_hypervisc.jl b/examples/bickleyjet/bickleyjet_cg_invariant_hypervisc.jl
@@ -154,7 +154,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = 0.02,
-    saveat = 1.0,
+    saveat = collect(0.0:1.0:80.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 )
diff --git a/examples/bickleyjet/bickleyjet_cg_unsmesh.jl b/examples/bickleyjet/bickleyjet_cg_unsmesh.jl
@@ -127,7 +127,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = 0.02,
-    saveat = 1.0,
+    saveat = collect(0.0:1.0:80.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 )
diff --git a/examples/bickleyjet/bickleyjet_dg.jl b/examples/bickleyjet/bickleyjet_dg.jl
@@ -226,7 +226,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = 0.02,
-    saveat = 1.0,
+    saveat = collect(0.0:1.0:200.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 )
diff --git a/examples/column/advect.jl b/examples/column/advect.jl
@@ -119,31 +119,31 @@ sol1 = solve(
     prob1,
     SSPRK33(),
     dt = Δt,
-    saveat = 10 * Δt,
+    saveat = collect(0.0:(10 * Δt):10.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
 sol2 = solve(
     prob2,
     SSPRK33(),
     dt = Δt,
-    saveat = 10 * Δt,
+    saveat = collect(0.0:(10 * Δt):10.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
 sol3 = solve(
     prob3,
     SSPRK33(),
     dt = Δt,
-    saveat = 10 * Δt,
+    saveat = collect(0.0:(10 * Δt):10.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
 sol4 = solve(
     prob4,
     SSPRK33(),
     dt = Δt,
-    saveat = 10 * Δt,
+    saveat = collect(0.0:(10 * Δt):10.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/column/advect_diffusion.jl b/examples/column/advect_diffusion.jl
@@ -95,7 +95,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 10000 * Δt,
+    saveat = collect(t₀:(10000 * Δt):t₁),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/column/bb_fct_advection.jl b/examples/column/bb_fct_advection.jl
@@ -115,7 +115,7 @@ for (i, stretch_fn) in enumerate(stretch_fns)
         prob,
         ExplicitAlgorithm(SSP33ShuOsher()),
         dt = Δt,
-        saveat = Δt,
+        saveat = collect(t₀:Δt:t₁),
         progress = true,
         adaptive = false,
         progress_message = (dt, u, p, t) -> t,
diff --git a/examples/column/ekman.jl b/examples/column/ekman.jl
@@ -118,7 +118,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 600, # save 10 min
+    saveat = collect(0.0:600:(60 * 60 * 50)), # save 10 min
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/column/fct_advection.jl b/examples/column/fct_advection.jl
@@ -116,7 +116,7 @@ for (i, stretch_fn) in enumerate(stretch_fns)
         prob,
         SSPRK33(),
         dt = Δt,
-        saveat = Δt,
+        saveat = collect(t₀:Δt:t₁),
         progress = true,
         adaptive = false,
         progress_message = (dt, u, p, t) -> t,
diff --git a/examples/column/heat.jl b/examples/column/heat.jl
@@ -56,7 +56,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 10 * Δt,
+    saveat = collect(0.0:(10 * Δt):10.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/column/hydrostatic.jl b/examples/column/hydrostatic.jl
@@ -134,7 +134,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 60 * 60, # save every hour
+    saveat = collect(0.0:(60 * 60):(60 * 60 * 24 * ndays)), # save every hour
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/column/hydrostatic_discrete.jl b/examples/column/hydrostatic_discrete.jl
@@ -146,7 +146,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 60 * 60, # save every hour
+    saveat = collect(0.0:(60 * 60):(60 * 60 * 24 * ndays)), # save every hour
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/column/hydrostatic_ekman.jl b/examples/column/hydrostatic_ekman.jl
@@ -180,7 +180,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 600, # save every hour
+    saveat = collect(0.0:600:(60 * 60)), # save every hour
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/column/step.jl b/examples/column/step.jl
@@ -121,31 +121,31 @@ sol1 = solve(
     prob1,
     SSPRK33(),
     dt = Δt,
-    saveat = 10 * Δt,
+    saveat = collect(0.0:(10 * Δt):5.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
 sol2 = solve(
     prob2,
     SSPRK33(),
     dt = Δt,
-    saveat = 10 * Δt,
+    saveat = collect(0.0:(10 * Δt):5.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
 sol3 = solve(
     prob3,
     SSPRK33(),
     dt = Δt,
-    saveat = 10 * Δt,
+    saveat = collect(0.0:(10 * Δt):5.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
 sol4 = solve(
     prob4,
     SSPRK33(),
     dt = Δt,
-    saveat = 10 * Δt,
+    saveat = collect(0.0:(10 * Δt):5.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/column/tvd_advection.jl b/examples/column/tvd_advection.jl
@@ -139,7 +139,7 @@ for (i, stretch_fn) in enumerate(stretch_fns)
             prob,
             ExplicitAlgorithm(SSP33ShuOsher()),
             dt = Δt,
-            saveat = Δt,
+            saveat = [t₀:Δt:t₁..., t₁],
         )
 
         q_final = sol.u[end].q
diff --git a/examples/column/vanleer_advection.jl b/examples/column/vanleer_advection.jl
@@ -113,7 +113,7 @@ for (i, stretch_fn) in enumerate(stretch_fns)
             prob,
             ExplicitAlgorithm(SSP33ShuOsher()),
             dt = Δt,
-            saveat = Δt,
+            saveat = [t₀:Δt:t₁..., t₁],
         )
 
         q_final = sol.u[end].q
diff --git a/examples/column/wave.jl b/examples/column/wave.jl
@@ -62,7 +62,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 10 * Δt,
+    saveat = [0.0:(10 * Δt):(4 * pi)..., 4 * pi],
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/column/zalesak_fct_advection.jl b/examples/column/zalesak_fct_advection.jl
@@ -112,7 +112,12 @@ for (i, stretch_fn) in enumerate(stretch_fns)
         (t₀, t₁),
         parameters,
     )
-    sol = solve(prob, ExplicitAlgorithm(SSP33ShuOsher()), dt = Δt, saveat = Δt)
+    sol = solve(
+        prob,
+        ExplicitAlgorithm(SSP33ShuOsher()),
+        dt = Δt,
+        saveat = collect(t₀:Δt:t₁),
+    )
 
     q_final = sol.u[end].q
     q_analytic = pulse.(z, t₁, z₀, zₕ, z₁)
diff --git a/examples/hybrid/box/bubble_3d_invariant_rhoe.jl b/examples/hybrid/box/bubble_3d_invariant_rhoe.jl
@@ -451,7 +451,7 @@ function bubble_3d_invariant_ρe(ARGS, comms_ctx, ::Type{FT}) where {FT}
         prob,
         SSPRK33(),
         dt = Δt,
-        saveat = 10.0,
+        saveat = [0.0:10.0:(sim_params.t_int)..., sim_params.t_int],
         progress = true,
         progress_message = (dt, u, p, t) -> t,
         internalnorm = (u, t) -> norm(u),
diff --git a/examples/hybrid/box/bubble_3d_invariant_rhotheta.jl b/examples/hybrid/box/bubble_3d_invariant_rhotheta.jl
@@ -311,7 +311,7 @@ sol = solve(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 1.0,
+    saveat = [0.0, 1.0],
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/hybrid/box/bubble_3d_rhotheta.jl b/examples/hybrid/box/bubble_3d_rhotheta.jl
@@ -312,7 +312,7 @@ integrator = init(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 50.0,
+    saveat = [0.0, 1.0],
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/hybrid/box/limiters_advection.jl b/examples/hybrid/box/limiters_advection.jl
@@ -320,7 +320,7 @@ for (k, horz_ne) in enumerate(horz_ne_seq)
         prob,
         CTS.ExplicitAlgorithm(CTS.SSP33ShuOsher()),
         dt = dt,
-        saveat = 0.99 * 80 * dt,
+        saveat = [0.0:(0.99 * 80 * dt):end_time..., end_time],
     )
 
     q_final = sol.u[end].ρq ./ sol.u[end].ρ
diff --git a/examples/hybrid/driver.jl b/examples/hybrid/driver.jl
@@ -177,7 +177,7 @@ problem = SciMLBase.ODEProblem(
 integrator = SciMLBase.init(
     problem,
     ode_algo;
-    saveat = dt_save_to_sol == 0 ? [] : dt_save_to_sol,
+    saveat = dt_save_to_sol == 0 ? [] : collect(t_start:dt_save_to_sol:t_end),
     callback = callback,
     dt = dt,
     adaptive = false,
diff --git a/examples/hybrid/plane/bubble_2d_invariant_rhoe.jl b/examples/hybrid/plane/bubble_2d_invariant_rhoe.jl
@@ -260,7 +260,7 @@ integrator = init(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 10.0,
+    saveat = collect(0.0:10.0:1200.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/hybrid/plane/density_current_2d_flux_form.jl b/examples/hybrid/plane/density_current_2d_flux_form.jl
@@ -303,7 +303,7 @@ integrator = init(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 50.0,
+    saveat = collect(0.0:50.0:900.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/hybrid/plane/density_current_2dinvariant_rhoe.jl b/examples/hybrid/plane/density_current_2dinvariant_rhoe.jl
@@ -280,7 +280,7 @@ integrator = init(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 10.0,
+    saveat = collect(0.0:10.0:timeend),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/hybrid/plane/isothermal_channel.jl b/examples/hybrid/plane/isothermal_channel.jl
@@ -284,7 +284,7 @@ integrator = init(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 30.0,
+    saveat = collect(0.0:30.0:15000.0),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
 );
diff --git a/examples/hybrid/plane/topo_agnesi_nh.jl b/examples/hybrid/plane/topo_agnesi_nh.jl
@@ -351,7 +351,7 @@ integrator = init(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 1800.0,
+    saveat = collect(0.0:1800.0:timeend),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
     callback = dss_callback,
diff --git a/examples/hybrid/plane/topo_schar_nh.jl b/examples/hybrid/plane/topo_schar_nh.jl
@@ -415,7 +415,7 @@ integrator = init(
     prob,
     SSPRK33(),
     dt = Δt,
-    saveat = 500.0,
+    saveat = collect(0.0:500.0:timeend),
     progress = true,
     progress_message = (dt, u, p, t) -> t,
     callback = dss_callback,
diff --git a/examples/hybrid/schur_complement_W.jl b/examples/hybrid/schur_complement_W.jl
@@ -34,6 +34,24 @@ struct SchurComplementW{F, FT, J1, J2, J3, J4, S, T}
     temp2::T
 end
 
+function Base.zero(jac::SchurComplementW)
+    return SchurComplementW(
+        jac.transform,
+        jac.flags,
+        jac.dtγ_ref,
+        Base.zero(jac.∂ᶜρₜ∂ᶠ𝕄),
+        Base.zero(jac.∂ᶜ𝔼ₜ∂ᶠ𝕄),
+        Base.zero(jac.∂ᶠ𝕄ₜ∂ᶜ𝔼),
+        Base.zero(jac.∂ᶠ𝕄ₜ∂ᶜρ),
+        Base.zero(jac.∂ᶠ𝕄ₜ∂ᶠ𝕄),
+        Base.zero(jac.S),
+        jac.test,
+        Base.zero(jac.temp1),
+        Base.zero(jac.temp2),
+    )
+end
+
+
 function SchurComplementW(Y, transform, flags, test = false)
     FT = eltype(Y)
     dtγ_ref = Ref(zero(FT))
diff --git a/examples/hybrid/sphere/deformation_flow.jl b/examples/hybrid/sphere/deformation_flow.jl
@@ -305,7 +305,12 @@ function run_deformation_flow(use_limiter, fct_op, dt)
         (0, t_end),
         cache,
     )
-    sol = solve(problem, ode_algorithm; dt, saveat = t_end / 2)
+    sol = solve(
+        problem,
+        ode_algorithm;
+        dt,
+        saveat = collect(0.0:(t_end / 2):t_end),
+    )
     if !(cache.limiter isa Nothing)
         @show cache.limiter.rtol
         Limiters.print_convergence_stats(cache.limiter)
diff --git a/examples/hybrid/sphere/hadley_circulation.jl b/examples/hybrid/sphere/hadley_circulation.jl
diff --git a/examples/hybrid/sphere/nonhydrostatic_gravity_wave.jl b/examples/hybrid/sphere/nonhydrostatic_gravity_wave.jl
diff --git a/examples/hybrid/sphere/solid_body_rotation_3d.jl b/examples/hybrid/sphere/solid_body_rotation_3d.jl
diff --git a/examples/plane/limiters_advection.jl b/examples/plane/limiters_advection.jl
diff --git a/examples/sphere/limiters_advection.jl b/examples/sphere/limiters_advection.jl
diff --git a/examples/sphere/shallow_water.jl b/examples/sphere/shallow_water.jl
diff --git a/examples/sphere/solidbody.jl b/examples/sphere/solidbody.jl

Original file line number	Diff line number	Diff line change
`@@ -127,7 +127,7 @@ sol = solve(`
`127`	`127`	`prob,`
`128`	`128`	`SSPRK33(),`
`129`	`129`	`dt = 0.02,`
`130`		`- saveat = 1.0,`
	`130`	`+ saveat = collect(0.0:1.0:80.0),`
`131`	`131`	`progress = true,`
`132`	`132`	`progress_message = (dt, u, p, t) -> t,`
`133`	`133`	`)`
Original file line number	Diff line number	Diff line change
`@@ -154,7 +154,7 @@ sol = solve(`
`154`	`154`	`prob,`
`155`	`155`	`SSPRK33(),`
`156`	`156`	`dt = 0.02,`
`157`		`- saveat = 1.0,`
	`157`	`+ saveat = collect(0.0:1.0:80.0),`
`158`	`158`	`progress = true,`
`159`	`159`	`progress_message = (dt, u, p, t) -> t,`
`160`	`160`	`)`
Original file line number	Diff line number	Diff line change
`@@ -226,7 +226,7 @@ sol = solve(`
`226`	`226`	`prob,`
`227`	`227`	`SSPRK33(),`
`228`	`228`	`dt = 0.02,`
`229`		`- saveat = 1.0,`
	`229`	`+ saveat = collect(0.0:1.0:200.0),`
`230`	`230`	`progress = true,`
`231`	`231`	`progress_message = (dt, u, p, t) -> t,`
`232`	`232`	`)`