Looks like tests will all pass 🎉 I'll start testing the checkpointing of increasingly complex simulations while iterating on the design! This way we'll be able to weed out most bugs and issues.

Ok so I think checkpointing now works for all turbulence closures!

In general, we shouldn't be calling initialize!(model) when restoring from a checkpoint because this could overwrite/corrupt closure fields. And we should use update_state!(model; compute_tendencies=false) when restoring from a checkpoint.

I also had to add a very small absolute tolerance to model equality tests when testing checkpointing turbulence closures, usually for fields like e and ε that may accumulate floating point differences due to differing order of operations (reductions?). But it's tiny, like atol = 1e-20.

Smagorinsky closures

Smagorinsky and SmagorinskyLilly don't need checkpointing.

In DirectionallyAveragedDynamicSmagorinsky we technically may not need to checkpoint 𝒥ᴸᴹ and 𝒥ᴹᴹ since they're recomputed from LM and MM but I did it to ensure exact reproducibility.

LagrangianAveragedDynamicSmagorinsky needed some changes to be checkpointed because it maintains time history for Lagrangian trajectory averaging. I had to also checkpoint previous_compute_time along with 𝒥ᴸᴹ⁻ and 𝒥ᴹᴹ⁻.

When using a multi-stage timestepper like RK3, I think it should only compute coefficients at the final stage. Lagrangian averaging was being applied 3 times with small fractional Δt values instead of once with the full Δt. Not doing so leaves the closure in an inconsistent state so when you pickup you can't recreate the same sequence of intermediate states. I think the main issue is that t⁻ is updated at every substep. The checkpoint captures t⁻ from stage=1, but the next computation in the continuous run happens at stage=2.

But please let me know if this makes sense for LagrangianAveragedDynamicSmagorinsky + RK3. If not, I can revert the changes and we can fix this properly in another PR. This is not an issue with QAB2. cc @tomchor @glwagner @simone-silvestri

RiBasedVerticalDiffusivity

This was relatively easy to checkpoint. Gotta checkpoint some closure fields but otherwise just needed to keep track of the previous_compute_time.

CATKE

I had to add some checks like making sure we only call time_step_catke_equation! and compute_average_surface_buoyancy_flux! on new iterations (Δt > 0) and making sure they work with split RK3. I also added a _skip_next_compute property to CATKEDiffusivityFields because we need to

TKEDissipationVerticalDiffusivity

Also similar to CATKE, I had to introduce a previous_compute_time and _skip_next_compute. Should also work once TKE-dissipation supports RK3 (#5127).

I think this is great!

We could consider a refactor in the future --- for example adding a new function that's something like time_step_closure_fields!(closure, model). With such an interface, we can distinguish between true "auxiliary state updates" (via update_state!) and prognostic state updates.

@simone-silvestri ZStarCoordinate checkpointing works now (and is tested)! I didn't have to checkpoint grids, just model.vertical_coordinate and also pass the grid to prognostic_state and restore_prognostic_state!. But if we want to checkpoint grids in the future it'll be easy to do.

I think this PR is ready to be merged now so just re-requesting review. I'd also like to wait to hear from @tomchor about whether the changes I made to DynamicSmagorinsky are okay.

Codecov Report

❌ Patch coverage is 83.80282% with 69 lines in your changes missing coverage. Please review.
✅ Project coverage is 73.22%. Comparing base (5a0553f) to head (9f0e493).
⚠️ Report is 1 commits behind head on main.

@@             Coverage Diff             @@
##             main    #4892       +/-   ##
===========================================
+ Coverage   32.87%   73.22%   +40.35%     
===========================================
  Files         383      391        +8     
  Lines       21105    21897      +792     
===========================================
+ Hits         6938    16034     +9096     
+ Misses      14167     5863     -8304     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

Smagorinsky and SmagorinskyLilly don't need checkpointing.

In DirectionallyAveragedDynamicSmagorinsky we technically may not need to checkpoint 𝒥ᴸᴹ and 𝒥ᴹᴹ since they're recomputed from LM and MM but I did it to ensure exact reproducibility.

LagrangianAveragedDynamicSmagorinsky needed some changes to be checkpointed because it maintains time history for Lagrangian trajectory averaging. I had to also checkpoint previous_compute_time along with 𝒥ᴸᴹ⁻ and 𝒥ᴹᴹ⁻.

When using a multi-stage timestepper like RK3, I think it should only compute coefficients at the final stage. Lagrangian averaging was being applied 3 times with small fractional Δt values instead of once with the full Δt. Not doing so leaves the closure in an inconsistent state so when you pickup you can't recreate the same sequence of intermediate states. I think the main issue is that t⁻ is updated at every substep. The checkpoint captures t⁻ from stage=1, but the next computation in the continuous run happens at stage=2.

But please let me know if this makes sense for LagrangianAveragedDynamicSmagorinsky + RK3. If not, I can revert the changes and we can fix this properly in another PR. This is not an issue with QAB2. cc @tomchor @glwagner @simone-silvestri

I tried to find some info on this online but couldn't. What I can say is that the models (I'm aware of) will usually update the closure every n time steps (usually 5), and when comes time to update the closure, they compute the averaging at every substep. Given that, and given the fact that substeps are very short (and perhaps not a "real" solution of the model), I think your changes are okay. Perhaps even an improvement! Great work!