CheckpointFile should store MeshHierarchy information.

[JHopeCollins]

Describe the current issue
The CheckpointFile doesn't retain any information on mesh hierarchies, so if you generate a mesh hiearchy you can't save it to a file and load it up later to use with a multigrid solver.

Describe the solution you'd like
I would like to be able to use checkpoint.save_mesh(mesh) with a mesh from a MeshHierarchy and then have the MeshHierarchy still be there when I use mesh = checkpoint.load_mesh(mesh_name)

Describe alternatives you've considered
Here's an MFE:

from firedrake import *

# make a refined mesh
coarse_mesh = UnitIntervalMesh(1, name="coarse_mesh")
mesh_hierarchy = MeshHierarchy(coarse_mesh, 1)
fine_mesh = mesh_hierarchy[-1]
fine_mesh.name = "fine_mesh"

# some data on the fine mesh
Vf = FunctionSpace(fine_mesh, "CG", 1)
uf = Function(Vf, name="fine_function").assign(2)

# save both meshes and the fine function
with CheckpointFile("checkpoint.h5", "w") as checkpoint:
    print("saved mesh has hierarchy: ", hasattr(fine_mesh.topological, "__level_info__"))  # True
    checkpoint.save_mesh(coarse_mesh)
    checkpoint.save_mesh(fine_mesh)
    checkpoint.save_function(uf)

# try to read from the mesh and retain the hiearchy information
with CheckpointFile("checkpoint.h5", "r") as checkpoint:

    # 1. Ideally the loaded mesh would just remember the hiearchy, but it doesn't
    fmesh0 = checkpoint.load_mesh(name="fine_mesh")
    uf0 = checkpoint.load_function(fmesh0, "fine_function")
    print("loaded mesh has hierarchy: ", hasattr(fmesh0.topological, "__level_info__"))  # False

    # 2. Regenerating the fine mesh wouldn't be ideal but would probably do for a lot of cases using simple meshes
    cmesh = checkpoint.load_mesh(name="coarse_mesh")
    fmesh1 = MeshHierarchy(cmesh, 1)[-1]
    fmesh1.name = "fine_mesh"
    # This line fails with the following error:
    # uf1 = checkpoint.load_function(fmesh1, "fine_function")
    #
    # File "/home/jhc/codes/firedrake/src/firedrake/firedrake/checkpointing.py", line 1095, in _load_function_space
    # RuntimeError: 
    #     FunctionSpace (firedrake_function_space_fine_mesh_CG1(None,None)) not found in either of the following path in checkpoint.h5:
    #         topologies/firedrake_mixed_meshes/fine_mesh/firedrake_mixed_function_spaces/firedrake_function_space_fine_mesh_CG1(None,None)
    #         topologies/DM_0x84000003_1/firedrake_meshes/fine_mesh/firedrake_function_spaces/firedrake_function_space_fine_mesh_CG1(None,None)


    # 3. Current workaround is to manually copy the raw data. Breaks abstractions, not ideal at all.
    Vf1 = FunctionSpace(fmesh1, uf0.ufl_element())
    uf1 = Function(Vf1)
    for dat0, dat1 in zip(uf0.dat, uf1.dat):
        dat1.data[:] = dat0.data[:]

Additional info
Output of firedrake-status when I ran the MFE:

Firedrake Configuration:
    package_manager: True
    minimal_petsc: False
    mpicc: None
    mpicxx: None
    mpif90: None
    mpiexec: None
    disable_ssh: False
    honour_petsc_dir: False
    with_parmetis: True
    slepc: True
    packages: ['git+ssh://github.com/firedrakeproject/Irksome.git#egg=Irksome', 'git+ssh://github.com/firedrakeproject/gusto.git@main#egg=gusto']
    honour_pythonpath: False
    opencascade: False
    tinyasm: True
    torch: False
    petsc_int_type: int32
    cache_dir: /home/jhc/codes/firedrake/.cache
    complex: False
    remove_build_files: False
    with_blas: None
    netgen: False
Additions:
    None
Environment:
    PYTHONPATH: None
    PETSC_ARCH: None
    PETSC_DIR: None
Status of components:
---------------------------------------------------------------------------
|Package             |Branch                        |Revision  |Modified  |
---------------------------------------------------------------------------
|FInAT               |master                        |e2805c4   |False     |
|Irksome             |master                        |05e1220   |False     |
|PyOP2               |master                        |d017d594  |False     |
|TinyASM             |master                        |015a89a   |False     |
|fiat                |master                        |501504d   |False     |
|firedrake           |master                        |8b906be1f |False     |
|gusto               |main                          |1708e31f  |False     |
|h5py                |firedrake                     |6cc4c912  |False     |
|libspatialindex     |master                        |4768bf3   |True      |
|libsupermesh        |master                        |dbe226b   |False     |
|loopy               |main                          |8158afdb  |False     |
|petsc               |firedrake                     |5029d32bf6|False     |
|pyadjoint           |master                        |f194553   |False     |
|pytest-mpi          |main                          |a478bc8   |False     |
|slepc               |firedrake                     |8242f3c0b |False     |
|tsfc                |master                        |799191d   |False     |
|ufl                 |master                        |054b0617  |False     |
---------------------------------------------------------------------------

[ksagiyam]

As pointed out, CheckpointFile is currently unable to save/load mesh hierarchies.

I can not come up with a nice workaround, as MeshHierarchy is expecting uninitialised mesh while CheckpointFile.load_mesh() internally calls mesh.init(). This means that we can not save the coarse mesh, load it, and pass that to MeshHierarcy.

So, we might need to remember how we construct the coarse mesh (e.g., with UnitSquareMesh). The following for example works:

from firedrake import *
  

def get_mh():
    nref = 1
    coarse_mesh = UnitSquareMesh(8, 8, name="coarse_mesh")
    return MeshHierarchy(coarse_mesh, nref)


fine_mesh = get_mh()[-1]
fine_mesh.name = "fine_mesh"
Vf = FunctionSpace(fine_mesh, "CG", 1)
x, y = SpatialCoordinate(fine_mesh)
uf = Function(Vf, name="fine_function").interpolate(x * y)
print(assemble(uf * dx))
with CheckpointFile("checkpoint.h5", "w") as checkpoint:
    checkpoint.save_function(uf)
with CheckpointFile("checkpoint.h5", "r") as checkpoint:
    fine_mesh_ = checkpoint.load_mesh(name="fine_mesh")
    uf_ = checkpoint.load_function(fine_mesh_, "fine_function")
    Vf_ = uf_.function_space()
    # Need to do this.
    fine_mesh = get_mh()[-1]
    Vf = FunctionSpace(fine_mesh, Vf_.ufl_element())
    uf = Function(Vf)
    with uf.dat.vec as vec, uf_.dat.vec as vec_:
        vec_.copy(result=vec)
    print(assemble(uf * dx))

[JHopeCollins]

Ok, not ideal but doable. I can use the checkpoint.set_attr to save some information on how to regenerate the mesh at least.