Fix consistency of precomputed quantities for Rosenbrock

Project.toml

@@ -1,7 +1,7 @@
 name = "ClimaTimeSteppers"
 uuid = "595c0a79-7f3d-439a-bc5a-b232dc3bde79"
 authors = ["Climate Modeling Alliance"]
-version = "0.7.32"
+version = "0.7.33"
 
 [deps]
 ClimaComms = "3a4d1b5c-c61d-41fd-a00a-5873ba7a1b0d"

docs/src/plotting_utils.jl

@@ -143,6 +143,7 @@ function verify_convergence(
     float_str(x) = @sprintf "%.4f" x
     pow_str(x) = "10^{$(@sprintf "%.1f" log10(x))}"
     function si_str(x)
+        x in (0, Inf, -Inf, NaN) && return string(x)
         exponent = floor(Int, log10(x))
         mantissa = x / 10.0^exponent
         return "$(float_str(mantissa)) \\times 10^{$exponent}"

src/solvers/rosenbrock.jl

@@ -146,9 +146,13 @@ function step_u!(int, cache::RosenbrockCache{Nstages}) where {Nstages}
             U .+= A[i, j] .* k[j]
         end
 
+        # NOTE: post_implicit! is a misnomer
         if !isnothing(post_implicit!)
-            # NOTE: post_implicit! is a misnomer
-            post_implicit!(U, p, t + αi * dt)
+            # We update p on every stage but the first, and at the end of each
+            # timestep. Since the first stage is unchanged from the end of the
+            # previous timestep, this order of operations ensures that p is
+            # always consistent with the state, including between timesteps.
+            (i != 1) && post_implicit!(U, p, t + αi * dt)
         end
 
         if !isnothing(T_imp!)
@@ -197,6 +201,7 @@ function step_u!(int, cache::RosenbrockCache{Nstages}) where {Nstages}
     end
 
     dss!(u, p, t + dt)
+    post_implicit!(u, p, t + dt)
     return nothing
 end