Export function data to vtk or h5

demos/gray_scott.py

@@ -159,18 +159,11 @@ def qoi():
 
 # Finally, plot the outputs to be viewed in Paraview. ::
 
-ic = mesh_seq.get_initial_condition()
-for field, sols in solutions.items():
-    fwd_outfile = VTKFile(f"gray_scott/{field}_forward.pvd")
-    adj_outfile = VTKFile(f"gray_scott/{field}_adjoint.pvd")
-    fwd_outfile.write(*ic[field].subfunctions)
-    for i in range(num_subintervals):
-        for sol in sols["forward"][i]:
-            fwd_outfile.write(*sol.subfunctions)
-        for sol in sols["adjoint"][i]:
-            adj_outfile.write(*sol.subfunctions)
-    adj_end = Function(ic[field]).assign(0.0)
-    adj_outfile.write(*adj_end.subfunctions)
+solutions.export(
+    "gray_scott/solutions.pvd",
+    export_field_types=["forward", "adjoint"],
+    initial_condition=mesh_seq.get_initial_condition(),
+)
 
 # In the `next demo <./gray_scott_split.py.html>`__, we consider solving the same
 # problem, but splitting the solution field into multiple components.

demos/gray_scott_split.py

@@ -161,16 +161,10 @@ def qoi():
 )
 solutions = mesh_seq.solve_adjoint()
 
-ic = mesh_seq.get_initial_condition()
-for field, sols in solutions.items():
-    fwd_outfile = VTKFile(f"gray_scott_split/{field}_forward.pvd")
-    adj_outfile = VTKFile(f"gray_scott_split/{field}_adjoint.pvd")
-    fwd_outfile.write(ic[field])
-    for i in range(num_subintervals):
-        for sol in sols["forward"][i]:
-            fwd_outfile.write(sol)
-        for sol in sols["adjoint"][i]:
-            adj_outfile.write(sol)
-    adj_outfile.write(Function(ic[field]).assign(0.0))
+solutions.export(
+    "gray_scott_split/solutions.pvd",
+    export_field_types=["forward", "adjoint"],
+    initial_condition=mesh_seq.get_initial_condition(),
+)
 
 # This tutorial can be dowloaded as a `Python script <gray_scott_split.py>`__.

demos/solid_body_rotation.py

@@ -237,14 +237,9 @@ def qoi():
 # use Paraview. To save all adjoint solution components in Paraview format, use the
 # following. ::
 
-for field, sols in solutions.items():
-    fwd_outfile = VTKFile(f"solid_body_rotation/{field}_forward.pvd")
-    adj_outfile = VTKFile(f"solid_body_rotation/{field}_adjoint.pvd")
-    for i in range(len(mesh_seq)):
-        for sol in sols["forward"][i]:
-            fwd_outfile.write(sol)
-        for sol in sols["adjoint"][i]:
-            adj_outfile.write(sol)
+solutions.export(
+    "solid_body_rotation/solutions.pvd", export_field_types=["forward", "adjoint"]
+)
 
 # In the `next demo <./gray_scott.py.html>`__, we increase the complexity by considering
 # two concentration fields in an advection-diffusion-reaction problem.

goalie/function_data.py

@@ -6,6 +6,8 @@
 
 import firedrake.function as ffunc
 import firedrake.functionspace as ffs
+from firedrake.checkpointing import CheckpointFile
+from firedrake.output.vtk_output import VTKFile
 
 from .utility import AttrDict
 
@@ -147,6 +149,138 @@ def extract(self, layout="field"):
         else:
             raise ValueError(f"Layout type '{layout}' not recognised.")
 
+    def export(self, output_fpath, export_field_types=None, initial_condition=None):
+        """
+        Export field data to a file. The file format is determined by the extension of
+        the output file path. Supported formats are '.pvd' and '.h5'.
+
+        If the output file format is '.pvd', the data is exported as a series of VTK
+        files using Firedrake's :class:`~.VTKFile`. Since mixed function spaces are not
+        supported by VTK, each subfunction of a mixed function is exported separately.
+        If initial conditions are provided and fields other than 'forward' are to be
+        exported, the initial values of these other fields are set to 'nan' since they
+        are not defined at the initial time (e.g., 'adjoint' fields).
+
+        If the output file format is '.h5', the data is exported as a single HDF5 file
+        using Firedrake's :class:`~.CheckpointFile`. If names of meshes in the mesh
+        sequence are not unique, they are renamed to ``"mesh_i"``, where ``i`` is the
+        subinterval index. Functions are saved with names of the form
+        ``"field_label_i_j"``, where ``i`` is the subinterval index and ``j`` is the
+        export index. Initial conditions are named in the form ``"field_initial"``.
+        The exported data may then be loaded using, for example,
+
+        .. code-block:: python
+
+            with CheckpointFile(output_fpath, "r") as afile:
+                first_mesh = afile.load_mesh("mesh_0")
+                initial_condition = afile.load_function(first_mesh, "u_initial")
+                first_export = afile.load_function(first_mesh, "u_forward_0_0")
+
+        :arg output_fpath: the path to the output file
+        :type output_fpath: :class:`str`
+        :kwarg export_field_types: the field types to export; defaults to all available
+            field types
+        :type export_field_types: :class:`str` or :class:`list` of :class:`str`
+        :kwarg initial_condition: if provided, exports the provided initial condition
+            for 'forward' fields.
+        :type initial_condition: :class:`dict` of :class:`~.Function`
+        """
+        if export_field_types is None:
+            default_export_types = {"forward", "adjoint", "error_indicator"}
+            export_field_types = list(set(self.labels) & default_export_types)
+        if isinstance(export_field_types, str):
+            export_field_types = [export_field_types]
+        if not all(field_type in self.labels for field_type in export_field_types):
+            raise ValueError(
+                f"Field types {export_field_types} not recognised."
+                f" Available types are {self.labels}."
+            )
+
+        if output_fpath.endswith(".pvd"):
+            self._export_vtk(output_fpath, export_field_types, initial_condition)
+        elif output_fpath.endswith(".h5"):
+            self._export_h5(output_fpath, export_field_types, initial_condition)
+        else:
+            raise ValueError(
+                f"Output file format not recognised: '{output_fpath}'."
+                " Supported formats are '.pvd' and '.h5'."
+            )
+
+    def _export_vtk(self, output_fpath, export_field_types, initial_condition=None):
+        """
+        Export field data to a series of VTK files. Arguments are the same as for
+        :meth:`~.export`.
+        """
+        tp = self.time_partition
+        outfile = VTKFile(output_fpath, adaptive=True)
+        if initial_condition is not None:
+            ics = []
+            for field_type in export_field_types:
+                for field, ic in initial_condition.items():
+                    ic = ic.copy(deepcopy=True)
+                    # If the function space is mixed, rename and append each
+                    # subfunction separately
+                    if hasattr(ic.function_space(), "num_sub_spaces"):
+                        for idx, sf in enumerate(ic.subfunctions):
+                            if field_type != "forward":
+                                sf = sf.copy(deepcopy=True)
+                                sf.assign(float("nan"))
+                            sf.rename(f"{field}[{idx}]_{field_type}")
+                            ics.append(sf)
+                    else:
+                        if field_type != "forward":
+                            ic.assign(float("nan"))
+                        ic.rename(f"{field}_{field_type}")
+                        ics.append(ic)
+            outfile.write(*ics, time=tp.subintervals[0][0])
+
+        for i in range(tp.num_subintervals):
+            for j in range(tp.num_exports_per_subinterval[i] - 1):
+                time = (
+                    tp.subintervals[i][0]
+                    + (j + 1) * tp.timesteps[i] * tp.num_timesteps_per_export[i]
+                )
+                fs = []
+                for field in tp.field_names:
+                    mixed = hasattr(self.function_spaces[field][0], "num_sub_spaces")
+                    for field_type in export_field_types:
+                        f = self._data[field][field_type][i][j].copy(deepcopy=True)
+                        if mixed:
+                            for idx, sf in enumerate(f.subfunctions):
+                                sf.rename(f"{field}[{idx}]_{field_type}")
+                                fs.append(sf)
+                        else:
+                            f.rename(f"{field}_{field_type}")
+                            fs.append(f)
+                outfile.write(*fs, time=time)
+
+    def _export_h5(self, output_fpath, export_field_types, initial_condition=None):
+        """
+        Export field data to an HDF5 file. Arguments are the same as for
+        :meth:`~.export`.
+        """
+        tp = self.time_partition
+
+        # Mesh names must be unique
+        mesh_names = [fs.mesh().name for fs in self.function_spaces[tp.field_names[0]]]
+        rename_meshes = len(set(mesh_names)) != len(mesh_names)
+        with CheckpointFile(output_fpath, "w") as outfile:
+            if initial_condition is not None:
+                for field, ic in initial_condition.items():
+                    outfile.save_function(ic, name=f"{field}_initial")
+            for i in range(tp.num_subintervals):
+                if rename_meshes:
+                    mesh_name = f"mesh_{i}"
+                    mesh = self.function_spaces[tp.field_names[0]][i].mesh()
+                    mesh.name = mesh_name
+                    mesh.topology_dm.name = mesh_name
+                for j in range(tp.num_exports_per_subinterval[i] - 1):
+                    for field in tp.field_names:
+                        for field_type in export_field_types:
+                            f = self._data[field][field_type][i][j]
+                            name = f"{field}_{field_type}_{i}_{j}"
+                            outfile.save_function(f, name=name)
+
 
 class ForwardSolutionData(FunctionData):
     """

test/test_function_data.py

@@ -4,6 +4,7 @@
 
 import abc
 import unittest
+from tempfile import TemporaryDirectory
 
 from firedrake import *
 
@@ -213,5 +214,58 @@ def test_extract_by_subinterval(self):
                 self.assertTrue(isinstance(f, Function))
 
 
+class TestExportFunctionData(BaseTestCases.TestFunctionData):
+    """
+    Unit tests for exporting and checkpointing :class:`~.FunctionData`.
+    """
+
+    def setUp(self):
+        super().setUpUnsteady()
+        self.labels = ("forward", "forward_old")
+
+    def _create_function_data(self):
+        self.solution_data = ForwardSolutionData(
+            self.time_partition, self.function_spaces
+        )
+        self.solution_data._create_data()
+
+    def test_export_extension_error(self):
+        with self.assertRaises(ValueError) as cm:
+            self.solution_data.export("test.ext")
+        msg = (
+            "Output file format not recognised: 'test.ext'."
+            + " Supported formats are '.pvd' and '.h5'."
+        )
+        self.assertEqual(str(cm.exception), msg)
+
+    def test_export_field_error(self):
+        with self.assertRaises(ValueError) as cm:
+            self.solution_data.export("test.pvd", export_field_types="test")
+        msg = (
+            "Field types ['test'] not recognised."
+            + f" Available types are {self.solution_data.labels}."
+        )
+        self.assertEqual(str(cm.exception), msg)
+
+    def test_export_pvd(self):
+        with TemporaryDirectory() as tmpdir:
+            export_filepath = os.path.join(tmpdir, "test.pvd")
+            self.solution_data.export(export_filepath)
+            self.assertTrue(os.path.exists(export_filepath))
+
+    def test_export_pvd_ic(self):
+        ic = {field: Function(fs[0]) for field, fs in self.function_spaces.items()}
+        with TemporaryDirectory() as tmpdir:
+            export_filepath = os.path.join(tmpdir, "test.pvd")
+            self.solution_data.export(export_filepath, initial_condition=ic)
+            self.assertTrue(os.path.exists(export_filepath))
+
+    def test_export_h5(self):
+        with TemporaryDirectory() as tmpdir:
+            export_filepath = os.path.join(tmpdir, "test.h5")
+            self.solution_data.export(export_filepath)
+            self.assertTrue(os.path.exists(export_filepath))
+
+
 if __name__ == "__main__":
     unittest.main()

test_adjoint/test_demos.py

@@ -25,16 +25,16 @@
     "burgers-hessian.py": {""""maxiter": 35""": """"maxiter": 3"""},
     "gray_scott.py": {
         "end_time = 2000.0": "end_time = 10.0",
-        r"ic = mesh_seq.get_initial_condition\(\)\n.*?adj_outfile.write\(\*adj_end.subfunctions\)": "",
+        r"solutions\.export\([\s\S]*?\)\s*,?\s*\)?\n?": "",
     },
     "gray_scott_split.py": {
         "end_time = 2000.0": "end_time = 10.0",
-        r"ic = mesh_seq.get_initial_condition\(\)\n.*?adj_outfile.write\(Function\(ic\[field\]\).assign\(0.0\)\)": "",
+        r"solutions\.export\([\s\S]*?\)\s*,?\s*\)?\n?": "",
     },
     "point_discharge2d-hessian.py": {""""maxiter": 35""": """"maxiter": 3"""},
     "point_discharge2d-goal_oriented.py": {""""maxiter": 35""": """"maxiter": 3"""},
     "solid_body_rotation.py": {
-        r"for field, sols in solutions.items\(\):\n.*?adj_outfile.write\(sol\)": "",
+        r"solutions\.export\((.*?)\)": "",
     },
 }