Fix extruded checkpointing

firedrake/checkpointing.py

@@ -1496,10 +1496,7 @@ def load_function(self, mesh, name, idx=None):
                 path = self._path_to_function_embedded(tmesh_name, mesh.name, V_name, name)
                 _name = self.get_attr(path, PREFIX_EMBEDDED + "_function")
                 _f = self.load_function(mesh, _name, idx=idx)
-                element = V.ufl_element()
-                _element = get_embedding_element_for_checkpointing(element, V.value_shape)
-                method = get_embedding_method_for_checkpointing(element)
-                assert _element == _f.function_space().ufl_element()
+                method = get_embedding_method_for_checkpointing(V.ufl_element())
                 f = Function(V, name=name)
                 self._project_function_for_checkpointing(f, _f, method)
                 return f

pyop2/types/mat.py

@@ -1173,7 +1173,7 @@ def mult(self, mat, x, y):
                     else:
                         x.array_r
                     with mpi.temp_internal_comm(x.comm) as comm:
-                        comm.bcast(a)
+                        a = comm.bcast(a)
                     return y.scale(a)
                 else:
                     return v.pointwiseMult(x, y)

tests/firedrake/output/test_io_backward_compat.py

@@ -1,36 +1,18 @@
-import pytest
-from os.path import abspath, dirname, join, exists
-from firedrake import *
-from firedrake.mesh import make_mesh_from_coordinates
-from firedrake.utils import IntType
-import shutil
+"""Backwards compatibility tests for CheckpointFile.
 
+New save files for the current version of Firedrake can be generated by
+running this file as a script.
 
-test_version = "2024_01_27"
+"""
+import importlib.metadata
+from os.path import abspath, dirname, join, exists
+import shutil
 
+import pytest
 
-"""
-2024_01_27:
----------------------------------------------------------------------------
-|Package             |Branch                        |Revision  |Modified  |
----------------------------------------------------------------------------
-|COFFEE              |master                        |70c1e66   |False     |
-|FInAT               |master                        |e2805c4   |False     |
-|PyOP2               |master                        |e0a4d3a9  |False     |
-|fiat                |master                        |e7b2909   |False     |
-|firedrake           |master                        |393f82f85 |False     |
-|h5py                |firedrake                     |4c01efa9  |False     |
-|libspatialindex     |master                        |4768bf3   |True      |
-|libsupermesh        |master                        |84becef   |False     |
-|loopy               |main                          |8158afdb  |False     |
-|petsc               |firedrake                     |09f36907a6e|False     |
-|pyadjoint           |master                        |2c6614d   |False     |
-|pytest-mpi          |main                          |a478bc8   |False     |
-|slepc               |firedrake                     |a3f39c853 |False     |
-|tsfc                |master                        |799191d   |False     |
-|ufl                 |master                        |054b0617  |False     |
----------------------------------------------------------------------------
-"""
+from firedrake import *
+from firedrake.mesh import make_mesh_from_coordinates
+from firedrake.utils import IntType, complex_mode
 
 
 cwd = abspath(dirname(__file__))
@@ -42,6 +24,15 @@
 stokes_control_mesh_file = join(cwd, "..", "..", "..", "docs", "notebooks/stokes-control.msh")
 
 
+SAVED_VERSIONS = ["2024_01_27", "2025.10.3.dev0"]
+"""Firedrake versions used to generate historic checkpoint files."""
+
+
+@pytest.fixture(params=SAVED_VERSIONS)
+def version(request):
+    return request.param
+
+
 def _initialise_function(f, _f):
     f.project(_f, solver_parameters={"ksp_type": "cg", "pc_type": "jacobi", "ksp_rtol": 1.e-16})
 
@@ -119,7 +110,10 @@ def _get_mesh_and_V(params):
             elem = TensorProductElement(helem, velem)
             V = FunctionSpace(mesh, elem)
         else:
-            raise NotImplementedError
+            assert cell_type == "quadrilateral"
+            base = UnitSquareMesh(10, 10, name=f"{mesh_name}_base")
+            mesh = ExtrudedMesh(base, layers=5, layer_height=1.0, name=mesh_name)
+            V = FunctionSpace(mesh, "P", 3)
     elif periodic:
         if cell_type == "triangle":
             mesh = PeriodicUnitSquareMesh(20, 20, name=mesh_name)
@@ -175,6 +169,7 @@ def _get_expr(V):
     ("quadrilateral", False, False, False, False, False, False),
     ("hexahedron", False, False, False, False, False, False),
     ("triangle", False, True, False, False, False, False),  # extruded (variable layer)
+    ("quadrilateral", False, True, False, False, False, False),  # extruded (constant layer)
     ("triangle", True, False, False, False, False, False),  # periodic
     ("tetrahedron", True, False, False, False, False, False),  # periodic
     ("interval", False, True, True, False, False, False),  # extruded_periodic
@@ -210,24 +205,7 @@ def _test_io_backward_compat_base_idfunc(params):
 
 
 @pytest.mark.skipcomplex
-@pytest.mark.parallel(nprocs=3)
-@pytest.mark.parametrize('version', [test_version])
-@pytest.mark.parametrize('params', test_io_backward_compat_base_params, ids=_test_io_backward_compat_base_idfunc)
-@pytest.mark.skip(reason="Only run these tests to create test files.")
-def test_io_backward_compat_base_save(version, params):
-    filename = join(filedir, "_".join([basename, version, _make_name(params) + ".h5"]))
-    if exists(filename):
-        raise RuntimeError(f"path {filename} already exists.")
-    mesh, V = _get_mesh_and_V(params)
-    f = Function(V, name=func_name)
-    _initialise_function(f, _get_expr(V))
-    with CheckpointFile(filename, "w") as afile:
-        afile.save_function(f)
-
-
-@pytest.mark.skipcomplex
-@pytest.mark.parallel(nprocs=4)
-@pytest.mark.parametrize('version', ["2024_01_27"])
+@pytest.mark.parallel(4)
 @pytest.mark.parametrize('params', test_io_backward_compat_base_params, ids=_test_io_backward_compat_base_idfunc)
 def test_io_backward_compat_base_load(version, params):
     filename = join(filedir, "_".join([basename, version, _make_name(params) + ".h5"]))
@@ -251,28 +229,7 @@ def _get_expr_timestepping(V, i):
 
 
 @pytest.mark.skipcomplex
-@pytest.mark.parallel(nprocs=3)
-@pytest.mark.parametrize('version', [test_version])
-@pytest.mark.skip(reason="Only run these tests to create test files.")
-def test_io_backward_compat_timestepping_save(version):
-    filename = join(filedir, "_".join([basename, version, "timestepping" + ".h5"]))
-    if exists(filename):
-        raise RuntimeError(f"path {filename} already exists.")
-    mesh = UnitSquareMesh(8, 8, name=mesh_name)
-    BDM = FunctionSpace(mesh, "BDM", 1)
-    DG = FunctionSpace(mesh, "DG", 0)
-    R = FunctionSpace(mesh, "Real", 0)
-    V = BDM * DG * R
-    f = Function(V, name=func_name)
-    with CheckpointFile(filename, 'w') as afile:
-        for i in range(5):
-            _initialise_function(f, _get_expr_timestepping(V, i))
-            afile.save_function(f, idx=i)
-
-
-@pytest.mark.skipcomplex
-@pytest.mark.parallel(nprocs=4)
-@pytest.mark.parametrize('version', ["2024_01_27"])
+@pytest.mark.parallel(4)
 def test_io_backward_compat_timestepping_load(version):
     filename = join(filedir, "_".join([basename, version, "timestepping" + ".h5"]))
     with CheckpointFile(filename, "r") as afile:
@@ -286,31 +243,70 @@ def test_io_backward_compat_timestepping_load(version):
 
 
 @pytest.mark.skipcomplex
-@pytest.mark.parallel(nprocs=3)
-@pytest.mark.parametrize('version', ["2024_01_27"])
+@pytest.mark.parallel(3)
 def test_io_backward_compat_timestepping_append(version, tmpdir):
     filename = join(filedir, "_".join([basename, version, "timestepping" + ".h5"]))
     copyname = join(str(tmpdir), "test_io_backward_compat_timestepping_append_dump.h5")
     copyname = COMM_WORLD.bcast(copyname, root=0)
     shutil.copyfile(filename, copyname)
     with CheckpointFile(copyname, "r") as afile:
-        version = afile.opts.parameters['dm_plex_view_hdf5_storage_version']
-        assert version == CheckpointFile.latest_version
         mesh = afile.load_mesh(mesh_name)
         f = afile.load_function(mesh, func_name, idx=0)
         V = f.function_space()
     with CheckpointFile(copyname, 'a') as afile:
-        version = afile.opts.parameters['dm_plex_view_hdf5_storage_version']
-        assert version == '2.1.0'
         for i in range(5, 10):
             _initialise_function(f, _get_expr_timestepping(V, i))
             afile.save_function(f, idx=i)
     with CheckpointFile(copyname, "r") as afile:
-        version = afile.opts.parameters['dm_plex_view_hdf5_storage_version']
-        assert version == CheckpointFile.latest_version
         for i in range(0, 10):
             f = afile.load_function(mesh, func_name, idx=i)
             V_ = f.function_space()
             f_ = Function(V_)
             _initialise_function(f_, _get_expr_timestepping(V_, i))
             assert assemble(inner(f - f_, f - f_) * dx) < 1.e-16
+
+
+if __name__ == "__main__":
+    if complex_mode:
+        print("Backwards compat checkpoint files should be generated in real mode")
+        exit(1)
+
+    version = importlib.metadata.version("firedrake")  # noqa: F811
+    if version in SAVED_VERSIONS:
+        print(f"Version {version} already has historic checkpoint file data saved")
+        exit(1)
+    print(f"Saving checkpoint files for Firedrake version {version}")
+
+    print("Saving files for 'test_io_backward_compat_base_load'")
+    for params in test_io_backward_compat_base_params:
+        test_id = _test_io_backward_compat_base_idfunc(params)
+        print(f"Saving function for '{test_id}'")
+        filename = join(filedir, "_".join([basename, version, _make_name(params) + ".h5"]))
+        if exists(filename):
+            raise RuntimeError(f"path {filename} already exists.")
+        mesh, V = _get_mesh_and_V(params)
+        f = Function(V, name=func_name)
+        _initialise_function(f, _get_expr(V))
+        with CheckpointFile(filename, "w") as afile:
+            afile.save_function(f)
+        print(f"Function for '{test_id}' saved")
+    print("Files for 'test_io_backward_compat_base_load' saved")
+
+    print("Saving files for 'test_io_backward_compat_timestepping_load'")
+    filename = join(filedir, "_".join([basename, version, "timestepping" + ".h5"]))
+    if exists(filename):
+        raise RuntimeError(f"path {filename} already exists.")
+    mesh = UnitSquareMesh(8, 8, name=mesh_name)
+    BDM = FunctionSpace(mesh, "BDM", 1)
+    DG = FunctionSpace(mesh, "DG", 0)
+    R = FunctionSpace(mesh, "Real", 0)
+    V = BDM * DG * R
+    f = Function(V, name=func_name)
+    with CheckpointFile(filename, 'w') as afile:
+        for i in range(5):
+            _initialise_function(f, _get_expr_timestepping(V, i))
+            afile.save_function(f, idx=i)
+    print("Files for 'test_io_backward_compat_timestepping_load' saved")
+
+    print("All checkpoint files saved, please add them to the Firedrake repository, "
+          f"making sure to include the new version '{version}' in 'SAVED_VERSIONS'")