Use VTX instead of XDMF when saving functions to file, and use adios4…

…dolfinx for checkpoints instead of PETSc views

pyproject.toml

@@ -53,6 +53,7 @@ funding = "https://github.com/sponsors/francesco-ballarin"
 
 [project.optional-dependencies]
 backends = [
+    "adios4dolfinx @ git+https://github.com/jorgensd/adios4dolfinx.git",
     "fenics-dolfinx >=0.8.0.dev0, <0.9.0",
     "multiphenicsx @ git+https://github.com/multiphenics/multiphenicsx.git",
     "ufl4rom @ git+https://github.com/RBniCS/ufl4rom.git"

rbnicsx/backends/export.py

@@ -8,6 +8,7 @@
 import os
 import typing
 
+import adios4dolfinx
 import dolfinx.fem
 import dolfinx.io
 import numpy as np
@@ -17,6 +18,7 @@
 from rbnicsx._backends.export import (
     export_matrices as export_matrices_super, export_matrix as export_matrix_super,
     export_vector as export_vector_super, export_vectors as export_vectors_super)
+from rbnicsx.io import on_rank_zero
 
 
 def export_function(function: dolfinx.fem.Function, directory: str, filename: str) -> None:
@@ -32,14 +34,18 @@ def export_function(function: dolfinx.fem.Function, directory: str, filename: st
     filename
         Name of the file where to export the function.
     """
-    os.makedirs(directory, exist_ok=True)
-    # Export to XDMF file for visualization
-    mesh = function.function_space.mesh
-    with dolfinx.io.XDMFFile(mesh.comm, os.path.join(directory, filename + ".xdmf"), "w") as xdmf_file:
-        xdmf_file.write_mesh(mesh)
-        xdmf_file.write_function(function)
-    # Export the underlying vector for restart purposes
-    export_vector(function.vector, directory, filename)
+    comm = function.function_space.mesh.comm
+    visualization_directory = os.path.join(directory, filename + ".bp")
+    checkpointing_directory = os.path.join(directory, filename + ".bp", ".checkpoint")
+    os.makedirs(visualization_directory, exist_ok=True)
+    os.makedirs(checkpointing_directory, exist_ok=True)
+
+    # Export for visualization
+    with dolfinx.io.VTXWriter(comm, visualization_directory, function, "bp4") as vtx_file:
+        vtx_file.write(0)
+
+    # Export for checkpointing
+    adios4dolfinx.write_function(function, checkpointing_directory, "bp4")
 
 
 def export_functions(
@@ -59,15 +65,31 @@ def export_functions(
     filename
         Name of the file where to export the function.
     """
-    os.makedirs(directory, exist_ok=True)
-    # Export to XDMF file for visualization
-    mesh = functions[0].function_space.mesh
-    with dolfinx.io.XDMFFile(mesh.comm, os.path.join(directory, filename + ".xdmf"), "w") as xdmf_file:
-        xdmf_file.write_mesh(mesh)
+    comm = functions[0].function_space.mesh.comm
+    visualization_directory = os.path.join(directory, filename + ".bp")
+    checkpointing_directory = os.path.join(directory, filename + ".bp", ".checkpoint")
+    os.makedirs(visualization_directory, exist_ok=True)
+    os.makedirs(checkpointing_directory, exist_ok=True)
+
+    # Export for visualization
+    output = functions[0].copy()
+    with dolfinx.io.VTXWriter(comm, visualization_directory, output, "bp4") as vtx_file:
         for (function, index) in zip(functions, indices):
-            xdmf_file.write_function(function, index)
-    # Export the underlying vectors for restart purposes
-    export_vectors([function.vector for function in functions], directory, filename)
+            output.x.array[:] = function.x.array
+            output.x.scatter_forward()
+            vtx_file.write(index)
+    del output
+
+    # Export for checkpointing: write out length of the list
+    def write_length() -> None:
+        with open(os.path.join(checkpointing_directory, "length.dat"), "w") as length_file:
+            length_file.write(str(len(functions)))
+    on_rank_zero(comm, write_length)
+
+    # Export for checkpointing: write out the list
+    # Note that here index is an integer counter, rather than an entry of the input array indices.
+    for (index, function) in enumerate(functions):
+        adios4dolfinx.write_function(function, os.path.join(checkpointing_directory, str(index)), "bp4")
 
 
 def export_matrix(  # type: ignore[no-any-unimported]

rbnicsx/backends/import_.py

@@ -8,6 +8,7 @@
 import os
 import typing
 
+import adios4dolfinx
 import dolfinx.fem
 import dolfinx.fem.petsc
 import mpi4py.MPI
@@ -39,11 +40,9 @@ def import_function(
     :
         Function imported from file.
     """
-    comm = function_space.mesh.comm
     function = dolfinx.fem.Function(function_space)
-    viewer = petsc4py.PETSc.Viewer().createBinary(os.path.join(directory, filename + ".dat"), "r", comm)
-    function.vector.load(viewer)
-    viewer.destroy()
+    checkpointing_directory = os.path.join(directory, filename + ".bp", ".checkpoint")
+    adios4dolfinx.read_function(function, checkpointing_directory, "bp4")
     return function
 
 
@@ -68,23 +67,22 @@ def import_functions(
         Functions imported from file.
     """
     comm = function_space.mesh.comm
+    checkpointing_directory = os.path.join(directory, filename + ".bp", ".checkpoint")
 
     # Read in length of the list
     def read_length() -> int:
-        with open(os.path.join(directory, filename, "length.dat"), "r") as length_file:
+        with open(os.path.join(checkpointing_directory, "length.dat"), "r") as length_file:
             return int(length_file.readline())
     length = on_rank_zero(comm, read_length)
 
     # Read in the list
     function_placeholder = dolfinx.fem.Function(function_space)
     functions = list()
     for index in range(length):
-        viewer = petsc4py.PETSc.Viewer().createBinary(
-            os.path.join(directory, filename, str(index) + ".dat"), "r", comm)
         function = function_placeholder.copy()
-        function.vector.load(viewer)
+        adios4dolfinx.read_function(function, os.path.join(checkpointing_directory, str(index)), "bp4")
         functions.append(function)
-        viewer.destroy()
+    del function_placeholder
     return functions
 
 

tests/unit/backends/test_backends_export_import.py

@@ -19,6 +19,9 @@
 
 import rbnicsx.backends
 
+all_families = ["Lagrange", "Discontinuous Lagrange"]
+all_degrees = [1, 2]
+
 
 @pytest.fixture
 def mesh() -> dolfinx.mesh.Mesh:
@@ -27,9 +30,11 @@ def mesh() -> dolfinx.mesh.Mesh:
     return dolfinx.mesh.create_unit_square(comm, 2 * comm.size, 2 * comm.size)
 
 
-def test_backends_export_import_function(mesh: dolfinx.mesh.Mesh) -> None:
+@pytest.mark.parametrize("family", all_families)
+@pytest.mark.parametrize("degree", all_degrees)
+def test_backends_export_import_function(mesh: dolfinx.mesh.Mesh, family: str, degree: str) -> None:
     """Check I/O for a dolfinx.fem.Function."""
-    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+    V = dolfinx.fem.functionspace(mesh, (family, degree))
     function = dolfinx.fem.Function(V)
     function.vector.set(1.0)
 
@@ -40,9 +45,11 @@ def test_backends_export_import_function(mesh: dolfinx.mesh.Mesh) -> None:
         assert np.allclose(function2.vector.array, function.vector.array)
 
 
-def test_backends_export_import_functions(mesh: dolfinx.mesh.Mesh) -> None:
+@pytest.mark.parametrize("family", all_families)
+@pytest.mark.parametrize("degree", all_degrees)
+def test_backends_export_import_functions(mesh: dolfinx.mesh.Mesh, family: str, degree: str) -> None:
     """Check I/O for a list of dolfinx.fem.Function."""
-    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+    V = dolfinx.fem.functionspace(mesh, (family, degree))
     functions = list()
     indices = list()
     for i in range(2):
@@ -60,9 +67,11 @@ def test_backends_export_import_functions(mesh: dolfinx.mesh.Mesh) -> None:
             assert np.allclose(function2.vector.array, function.vector.array)
 
 
-def test_backends_export_import_vector(mesh: dolfinx.mesh.Mesh) -> None:
+@pytest.mark.parametrize("family", all_families)
+@pytest.mark.parametrize("degree", all_degrees)
+def test_backends_export_import_vector(mesh: dolfinx.mesh.Mesh, family: str, degree: str) -> None:
     """Check I/O for a petsc4py.PETSc.Vec."""
-    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+    V = dolfinx.fem.functionspace(mesh, (family, degree))
     v = ufl.TestFunction(V)
     linear_form = ufl.inner(1, v) * ufl.dx
     linear_form_cpp = dolfinx.fem.form(linear_form)
@@ -76,9 +85,11 @@ def test_backends_export_import_vector(mesh: dolfinx.mesh.Mesh) -> None:
         assert np.allclose(vector2.array, vector.array)
 
 
-def test_backends_export_import_vectors(mesh: dolfinx.mesh.Mesh) -> None:
+@pytest.mark.parametrize("family", all_families)
+@pytest.mark.parametrize("degree", all_degrees)
+def test_backends_export_import_vectors(mesh: dolfinx.mesh.Mesh, family: str, degree: str) -> None:
     """Check I/O for a list of petsc4py.PETSc.Vec."""
-    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+    V = dolfinx.fem.functionspace(mesh, (family, degree))
     v = ufl.TestFunction(V)
     linear_forms = [ufl.inner(i + 1, v) * ufl.dx for i in range(2)]
     linear_forms_cpp = dolfinx.fem.form(linear_forms)
@@ -95,12 +106,15 @@ def test_backends_export_import_vectors(mesh: dolfinx.mesh.Mesh) -> None:
             assert np.allclose(vector2.array, vector.array)
 
 
+@pytest.mark.parametrize("family", all_families)
+@pytest.mark.parametrize("degree", all_degrees)
 def test_backends_export_import_matrix(  # type: ignore[no-any-unimported]
     mesh: dolfinx.mesh.Mesh,
-    to_dense_matrix: typing.Callable[[petsc4py.PETSc.Mat], np.typing.NDArray[petsc4py.PETSc.ScalarType]]
+    to_dense_matrix: typing.Callable[[petsc4py.PETSc.Mat], np.typing.NDArray[petsc4py.PETSc.ScalarType]],
+    family: str, degree: str
 ) -> None:
     """Check I/O for a petsc4py.PETSc.Mat."""
-    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+    V = dolfinx.fem.functionspace(mesh, (family, degree))
     u = ufl.TrialFunction(V)
     v = ufl.TestFunction(V)
     bilinear_form = ufl.inner(u, v) * ufl.dx
@@ -115,12 +129,15 @@ def test_backends_export_import_matrix(  # type: ignore[no-any-unimported]
         assert np.allclose(to_dense_matrix(matrix2), to_dense_matrix(matrix))
 
 
+@pytest.mark.parametrize("family", all_families)
+@pytest.mark.parametrize("degree", all_degrees)
 def test_backends_export_import_matrices(  # type: ignore[no-any-unimported]
     mesh: dolfinx.mesh.Mesh,
-    to_dense_matrix: typing.Callable[[petsc4py.PETSc.Mat], np.typing.NDArray[petsc4py.PETSc.ScalarType]]
+    to_dense_matrix: typing.Callable[[petsc4py.PETSc.Mat], np.typing.NDArray[petsc4py.PETSc.ScalarType]],
+    family: str, degree: str
 ) -> None:
     """Check I/O for a list of petsc4py.PETSc.Mat."""
-    V = dolfinx.fem.functionspace(mesh, ("Lagrange", 1))
+    V = dolfinx.fem.functionspace(mesh, (family, degree))
     u = ufl.TrialFunction(V)
     v = ufl.TestFunction(V)
     bilinear_forms = [(i + 1) * ufl.inner(u, v) * ufl.dx for i in range(2)]