move restore into Checkpointer

Author: juliasloan25

docs/src/checkpointer.md

@@ -12,7 +12,7 @@ Checkpoints are a mix of HDF5 and JLD2 files and are typically saved in a
 !!! warning "Known limitations"
 
     - The number of MPI processes has to remain the same across checkpoints
-    - Restart files are generally not portable across machines, julia versions, 
+    - Restart files are generally not portable across machines, julia versions,
       and package versions
     - Adding/changing new component models will probably require adding/changing code
 
@@ -56,7 +56,7 @@ If the model does not support directly reading a checkpoint, the `Checkpointer`
 module provides a straightforward way to add this feature.
 [`Checkpointer.restart!`](@ref) takes a coupled simulation, a `restart_dir`, and
 a `restart_t` and overwrites the content of the coupled simulation with what is
-in the checkpoint. 
+in the checkpoint.
 
 ## Developer notes
 
@@ -135,4 +135,5 @@ Types to watch for:
     ClimaCoupler.Checkpointer.restart!
     ClimaCoupler.Checkpointer.checkpoint_sims
     ClimaCoupler.Checkpointer.t_start_from_checkpoint
+    ClimaCoupler.Checkpointer.restore!
 ```

experiments/ClimaEarth/components/atmosphere/climaatmos.jl

@@ -20,7 +20,6 @@ if pkgversion(CA) < v"0.28.6"
     CC.Adapt.@adapt_structure CA.RRTMGPInterface.RRTMGPModel
 end
 
-include("../shared/restore.jl")
 
 ###
 ### Functions required by ClimaCoupler.jl for an AtmosModelSimulation
@@ -140,7 +139,7 @@ end
 
 function Checkpointer.restore_cache!(sim::ClimaAtmosSimulation, new_cache)
     comms_ctx = ClimaComms.context(sim.integrator.u.c)
-    restore!(
+    Checkpointer.restore!(
         Checkpointer.get_model_cache(sim),
         new_cache,
         comms_ctx;
@@ -409,14 +408,14 @@ Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:height_sfc}) =
 function Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:u_int})
     # NOTE: This calculation is copied from ClimaAtmos (and is allocating! Fix me if you can!)
     int_local_geometry_values =
-        Fields.level(Fields.local_geometry_field(sim.integrator.u.c), 1)
+        CC.Fields.level(CC.Fields.local_geometry_field(sim.integrator.u.c), 1)
     int_u_values = CC.Spaces.level(sim.integrator.p.precomputed.ᶜu, 1)
     return CA.projected_vector_data.(CA.CT1, int_u_values, int_local_geometry_values)
 end
 function Interfacer.get_field(sim::ClimaAtmosSimulation, ::Val{:v_int})
     # NOTE: This calculation is copied from ClimaAtmos (and is allocating! Fix me if you can!)
     int_local_geometry_values =
-        Fields.level(Fields.local_geometry_field(sim.integrator.u.c), 1)
+        CC.Fields.level(CC.Fields.local_geometry_field(sim.integrator.u.c), 1)
     int_u_values = CC.Spaces.level(sim.integrator.p.precomputed.ᶜu, 1)
     return CA.projected_vector_data.(CA.CT2, int_u_values, int_local_geometry_values)
 end

experiments/ClimaEarth/components/land/climaland_bucket.jl

@@ -13,7 +13,6 @@ import ClimaCoupler: Checkpointer, FluxCalculator, Interfacer, FieldExchanger
 using NCDatasets
 include("climaland_helpers.jl")
 
-include("../shared/restore.jl")
 
 ###
 ### Functions required by ClimaCoupler.jl for a SurfaceModelSimulation
@@ -411,7 +410,7 @@ end
 function Checkpointer.restore_cache!(sim::BucketSimulation, new_cache)
     old_cache = Checkpointer.get_model_cache(sim)
     comms_ctx = ClimaComms.context(sim.model)
-    restore!(
+    Checkpointer.restore!(
         old_cache,
         new_cache,
         comms_ctx,

experiments/ClimaEarth/components/land/climaland_helpers.jl

@@ -75,7 +75,7 @@ function make_land_domain(
 
     vertmesh = CC.Meshes.IntervalMesh(
         vertdomain,
-        ClimaCore.Meshes.GeneralizedExponentialStretching{FT}(dz_tuple[1], dz_tuple[2]);
+        CC.Meshes.GeneralizedExponentialStretching{FT}(dz_tuple[1], dz_tuple[2]);
         nelems = nelements_vert,
         reverse_mode = true,
     )

experiments/ClimaEarth/components/land/climaland_integrated.jl

@@ -533,7 +533,7 @@ end
 function Checkpointer.restore_cache!(sim::ClimaLandSimulation, new_cache)
     old_cache = Checkpointer.get_model_cache(sim)
     comms_ctx = ClimaComms.context(sim.model.soil)
-    restore!(
+    Checkpointer.restore!(
         old_cache,
         new_cache,
         comms_ctx,

experiments/ClimaEarth/components/shared/restore.jl

@@ -3,12 +3,7 @@
 
 import ClimaComms
 import ClimaAtmos as CA
-import ClimaCore
-import ClimaCore: DataLayouts, Fields, Geometry
-import ClimaCore.Fields: Field, FieldVector, field_values
-import ClimaCore.DataLayouts: AbstractData
-import ClimaCore.Geometry: AxisTensor
-import ClimaCore.Spaces: AbstractSpace
+import ClimaCore as CC
 import NCDatasets
 
 """
@@ -55,8 +50,8 @@ function compare(
     name = "",
     ignore = Set([:rc]),
 ) where {
-    T1 <: Union{FieldVector, AbstractSpace, NamedTuple, CA.AtmosCache},
-    T2 <: Union{FieldVector, AbstractSpace, NamedTuple, CA.AtmosCache},
+    T1 <: Union{CC.Fields.FieldVector, CC.Spaces.AbstractSpace, NamedTuple, CA.AtmosCache},
+    T2 <: Union{CC.Fields.FieldVector, CC.Spaces.AbstractSpace, NamedTuple, CA.AtmosCache},
 }
     pass = true
     return _compare(pass, v1, v2; name, ignore)
@@ -100,11 +95,16 @@ function _compare(pass, v1::T, v2::T; name, ignore) where {T <: NCDatasets.NCDat
     return pass
 end
 
-function _compare(v1::T, v2::T; name, ignore) where {T <: Field{<:AbstractData{<:Real}}}
+function _compare(
+    v1::T,
+    v2::T;
+    name,
+    ignore,
+) where {T <: CC.Fields.Field{<:CC.DataLayouts.AbstractData{<:Real}}}
     return _compare(parent(v1), parent(v2); name, ignore)
 end
 
-function _compare(pass, v1::T, v2::T; name, ignore) where {T <: AbstractData}
+function _compare(pass, v1::T, v2::T; name, ignore) where {T <: CC.DataLayouts.AbstractData}
     return pass && _compare(parent(v1), parent(v2); name, ignore)
 end