Add functions to write CG1 and DG0 functions to VTKHDF

Author: mleoni-pf

cpp/dolfinx/io/VTKHDF.h

@@ -8,7 +8,9 @@
 #include <algorithm>
 #include <concepts>
 #include <dolfinx/common/IndexMap.h>
+#include <dolfinx/fem/Function.h>
 #include <dolfinx/io/cells.h>
+#include <dolfinx/io/utils.h>
 #include <dolfinx/mesh/Mesh.h>
 #include <dolfinx/mesh/Topology.h>
 #include <dolfinx/mesh/utils.h>
@@ -298,6 +300,68 @@ void write_data(std::string point_or_cell, std::string filename,
   hdf5::close_file(h5file);
 }
 
+/// @brief Write a CG1 function to VTKHDF.
+///
+/// Adds a function to an existing VTKHDF file, which already contains a mesh.
+///
+/// @tparam U Scalar type.
+/// @param[in] filename File for output.
+/// @param[in] mesh Mesh, which must be the same as the original mesh
+/// used in the file.
+/// @param[in] u Function to write to file.
+/// @param[in] time Timestamp.
+///
+/// @note Mesh must be written to file first using `VTKHDF::write_mesh`.
+/// @note Only one dataset "u" can be written per file at present, with
+/// multiple timesteps.
+/// @note Limited support for floating point types at present (no
+/// complex number support).
+template <std::floating_point U>
+void write_CG1_function(std::string filename, const mesh::Mesh<U>& mesh,
+                        const fem::Function<U>& u, double time)
+{
+  io::VTKHDF::write_data<U>(
+      "Point", filename, mesh,
+      std::vector<U>(std::begin(u.x()->array()),
+                     std::begin(u.x()->array())
+                         + mesh.geometry().index_map()->size_local()
+                               * u.function_space()->dofmaps(0)->bs()),
+      time);
+}
+
+/// @brief Write a DG0 function to VTKHDF.
+///
+/// Adds a function to an existing VTKHDF file, which already contains a mesh.
+///
+/// @tparam U Scalar type.
+/// @param[in] filename File for output.
+/// @param[in] mesh Mesh, which must be the same as the original mesh
+/// used in the file.
+/// @param[in] u Function to write to file.
+/// @param[in] time Timestamp.
+///
+/// @note Mesh must be written to file first using `VTKHDF::write_mesh`.
+/// @note Only one dataset "u" can be written per file at present, with
+/// multiple timesteps.
+/// @note Limited support for floating point types at present (no
+/// complex number support).
+template <std::floating_point U>
+void write_DG0_function(std::string filename, const mesh::Mesh<U>& mesh,
+                        const fem::Function<U>& u, double time)
+{
+  const auto index_maps = mesh.topology()->index_maps(mesh.topology()->dim());
+  int npoints
+      = std::accumulate(index_maps.begin(), index_maps.end(), 0,
+                        [](int a, auto im) { return a + im->size_local(); });
+
+  io::VTKHDF::write_data<U>(
+      "Cell", filename, mesh,
+      std::vector<U>(std::begin(u.x()->array()),
+                     std::begin(u.x()->array())
+                         + npoints * u.function_space()->dofmaps(0)->bs()),
+      time);
+}
+
 /// @brief Read a mesh from a VTKHDF format file.
 ///
 /// @tparam U Scalar type of mesh

python/dolfinx/io/vtkhdf.py

@@ -16,12 +16,22 @@
 from dolfinx.cpp.io import (
     read_vtkhdf_mesh_float32,
     read_vtkhdf_mesh_float64,
+    write_vtkhdf_cg1_function,
     write_vtkhdf_data,
+    write_vtkhdf_dg0_function,
     write_vtkhdf_mesh,
 )
+from dolfinx.fem import Function
 from dolfinx.mesh import Mesh
 
-__all__ = ["read_mesh", "write_cell_data", "write_mesh", "write_point_data"]
+__all__ = [
+    "read_mesh",
+    "write_cell_data",
+    "write_cg1_function",
+    "write_dg0_function",
+    "write_mesh",
+    "write_point_data",
+]
 
 
 def read_mesh(
@@ -100,3 +110,27 @@ def write_cell_data(filename: str | Path, mesh: Mesh, data: npt.NDArray, time: f
         time: Timestamp.
     """
     write_vtkhdf_data("Cell", filename, mesh._cpp_object, data, time)
+
+
+def write_cg1_function(filename: str | Path, mesh: Mesh, u: Function, time: float):
+    """Write a CG1 function to file.
+
+    Args:
+        filename: File to write to.
+        mesh: Mesh.
+        u: Function to write.
+        time: Timestamp.
+    """
+    write_vtkhdf_cg1_function(filename, mesh._cpp_object, u._cpp_object, time)
+
+
+def write_dg0_function(filename: str | Path, mesh: Mesh, u: Function, time: float):
+    """Write a DG0 function to file.
+
+    Args:
+        filename: File to write to.
+        mesh: Mesh.
+        u: Function to write.
+        time: Timestamp.
+    """
+    write_vtkhdf_dg0_function(filename, mesh._cpp_object, u._cpp_object, time)

python/dolfinx/wrappers/io.cpp

@@ -225,6 +225,10 @@ void io(nb::module_& m)
       .def("write_vtkhdf_mesh", &dolfinx::io::VTKHDF::write_mesh<float>);
   m.def("write_vtkhdf_data", &dolfinx::io::VTKHDF::write_data<double>);
   m.def("write_vtkhdf_data", &dolfinx::io::VTKHDF::write_data<float>);
+  m.def("write_vtkhdf_cg1_function", &dolfinx::io::VTKHDF::write_CG1_function<float>);
+  m.def("write_vtkhdf_cg1_function", &dolfinx::io::VTKHDF::write_CG1_function<double>);
+  m.def("write_vtkhdf_dg0_function", &dolfinx::io::VTKHDF::write_DG0_function<float>);
+  m.def("write_vtkhdf_dg0_function", &dolfinx::io::VTKHDF::write_DG0_function<double>);
   m.def("read_vtkhdf_mesh_float64",
         [](MPICommWrapper comm, const std::string& filename, std::size_t gdim,
            std::optional<std::int32_t> max_facet_to_cell_links)