diff --git a/cpp/dolfinx/io/VTKHDF.h b/cpp/dolfinx/io/VTKHDF.h index 29f63a5071..4ef10b8720 100644 --- a/cpp/dolfinx/io/VTKHDF.h +++ b/cpp/dolfinx/io/VTKHDF.h @@ -6,9 +6,13 @@ #include "HDF5Interface.h" #include +#include +#include #include #include +#include #include +#include #include #include #include @@ -298,6 +302,131 @@ void write_data(std::string point_or_cell, std::string filename, hdf5::close_file(h5file); } +/// @brief Write a 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). This function only supports DG0 and CG1 functions. +template +void write_function(std::string filename, const mesh::Mesh& mesh, + const fem::Function& u, double time) +{ + auto dofmap = u.function_space()->dofmap(); + assert(dofmap); + const int bs = dofmap->bs(); + + auto map_c = mesh.topology()->index_map(mesh.topology()->dim()); + assert(map_c); + + std::shared_ptr> element + = u.function_space()->element(); + assert(element); + + std::span value_shape = element->value_shape(); + int rank = value_shape.size(); + std::int32_t num_components = std::reduce( + value_shape.begin(), value_shape.end(), 1, std::multiplies{}); + + std::span x = u.x()->array(); + + // Check that it is a Lagrange family element + if (element->basix_element().family() != basix::element::family::P) + { + throw std::runtime_error("Unsupported function space. Only DG0 and CG1 are " + "supported at the moment."); + } + + // DG0 + if (element->basix_element().degree() == 0) + { + const std::int32_t num_local_cells = map_c->size_local(); + std::vector data(num_local_cells * num_components); + + for (std::int32_t c = 0; c < num_local_cells; ++c) + { + auto dofs = dofmap->cell_dofs(c); + assert(dofs.size() == 1); + for (std::size_t i = 0; i < dofs.size(); ++i) + { + std::copy_n(std::cbegin(x) + bs * dofs[i], bs, + std::begin(data) + num_components * c); + } + } + + io::VTKHDF::write_data("Cell", filename, mesh, data, time); + } + // CG1 + else if (element->basix_element().discontinuous() == false + and element->basix_element().degree() == 1) + { + auto map_x = mesh.geometry().index_map(); + assert(map_x); + + auto& geometry = mesh.geometry(); + auto& cmap = geometry.cmap(); + int cmap_dim = cmap.dim(); + int cell_dim = element->space_dimension() / element->block_size(); + if (cmap_dim != cell_dim) + { + throw std::runtime_error("Degree of output Function must be the same as " + "mesh degree. Maybe the " + "Function needs to be interpolated?"); + } + + // Check that dofmap layouts are equal and check Lagrange variants + if (dofmap->element_dof_layout() != cmap.create_dof_layout()) + { + throw std::runtime_error("Function and Mesh dof layouts do not match. " + "Maybe the Function needs to be interpolated?"); + } + if (cmap.degree() > 2 + and element->basix_element().lagrange_variant() != cmap.variant()) + { + throw std::runtime_error("Mismatch in Lagrange family. Maybe the " + "Function needs to be interpolated?"); + } + + std::int32_t num_cells = map_c->size_local() + map_c->num_ghosts(); + std::int32_t num_local_points = map_x->size_local(); + + // Get dof array and pack into array (padded where appropriate) + auto dofmap_x = geometry.dofmap(); + std::vector data(num_local_points * num_components); + for (std::int32_t c = 0; c < num_cells; ++c) + { + auto dofs = dofmap->cell_dofs(c); + auto dofs_x = md::submdspan(dofmap_x, c, md::full_extent); + assert(dofs.size() == dofs_x.size()); + for (std::size_t i = 0; i < dofs.size(); ++i) + { + if (dofs_x[i] < num_local_points) + { + std::copy_n(std::cbegin(x) + bs * dofs[i], bs, + std::begin(data) + num_components * dofs_x[i]); + } + } + } + + io::VTKHDF::write_data("Point", filename, mesh, data, time); + } + else + { + throw std::runtime_error("Unsupported function space. Only DG0 and CG1 are " + "supported at the moment."); + } +} + /// @brief Read a mesh from a VTKHDF format file. /// /// @tparam U Scalar type of mesh @@ -473,4 +602,45 @@ mesh::Mesh read_mesh(MPI_Comm comm, std::string filename, points_pruned, {(std::size_t)x_shape[0], gdim}, part, max_facet_to_cell_links); } + +/// @brief Read data from a VTKHDF format file. +/// +/// @tparam U Scalar type of mesh +/// @param[in] point_or_cell String "Point" or "Cell" determining data +/// location. +/// @param filename Name of the file to read from. +/// @param mesh Mesh previously read from the same file. +/// @param range The local range of data to read. +/// @param timestep The time step to read for time-dependent data. +/// @return The data read from file. +template +std::vector read_data(std::string point_or_cell, std::string filename, + const mesh::Mesh& mesh, + std::array range, int timestep = 0) +{ + hid_t h5file = hdf5::open_file(mesh.comm(), filename, "r", true); + std::string dataset_name = "/VTKHDF/" + point_or_cell + "Data/u"; + + std::int64_t data_offset = 0; + // Read the offset for the requested timestep + std::string offset_path = "/VTKHDF/Steps/" + point_or_cell + "DataOffsets/u"; + hid_t offset_dset = hdf5::open_dataset(h5file, offset_path); + std::vector offsets = hdf5::read_dataset( + offset_dset, {timestep, timestep + 1}, true); + H5Dclose(offset_dset); + data_offset = offsets[0]; + + // Adjust range to account for timestep offset + range[0] += data_offset; + range[1] += data_offset; + + // Read data using HDF5 + hid_t dset_id = hdf5::open_dataset(h5file, dataset_name); + std::vector values = hdf5::read_dataset(dset_id, range, true); + H5Dclose(dset_id); + + hdf5::close_file(h5file); + + return values; +} } // namespace dolfinx::io::VTKHDF diff --git a/python/dolfinx/io/vtkhdf.py b/python/dolfinx/io/vtkhdf.py index 958648db58..75b66d7478 100644 --- a/python/dolfinx/io/vtkhdf.py +++ b/python/dolfinx/io/vtkhdf.py @@ -17,11 +17,19 @@ read_vtkhdf_mesh_float32, read_vtkhdf_mesh_float64, write_vtkhdf_data, + write_vtkhdf_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_function", + "write_mesh", + "write_point_data", +] def read_mesh( @@ -100,3 +108,15 @@ 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_function(filename: str | Path, mesh: Mesh, u: Function, time: float): + """Write a function to file. Only CG1 and DG0 functions are supported. + + Args: + filename: File to write to. + mesh: Mesh. + u: Function to write. + time: Timestamp. + """ + write_vtkhdf_function(filename, mesh._cpp_object, u._cpp_object, time) diff --git a/python/dolfinx/wrappers/io.cpp b/python/dolfinx/wrappers/io.cpp index 582996c9ad..89044dfd4c 100644 --- a/python/dolfinx/wrappers/io.cpp +++ b/python/dolfinx/wrappers/io.cpp @@ -225,6 +225,8 @@ void io(nb::module_& m) .def("write_vtkhdf_mesh", &dolfinx::io::VTKHDF::write_mesh); m.def("write_vtkhdf_data", &dolfinx::io::VTKHDF::write_data); m.def("write_vtkhdf_data", &dolfinx::io::VTKHDF::write_data); + m.def("write_vtkhdf_function", &dolfinx::io::VTKHDF::write_function); + m.def("write_vtkhdf_function", &dolfinx::io::VTKHDF::write_function); m.def("read_vtkhdf_mesh_float64", [](MPICommWrapper comm, const std::string& filename, std::size_t gdim, std::optional max_facet_to_cell_links)