Add XDMF writer; resolves issues #8 and #10.

flowtorch/data/__init__.py

@@ -1,5 +1,5 @@
 from .foam_dataloader import FOAMDataloader, FOAMCase, FOAMMesh
-from .hdf5_file import HDF5Dataloader, HDF5Writer, FOAM2HDF5
+from .hdf5_file import HDF5Dataloader, HDF5Writer, FOAM2HDF5, XDMFWriter
 from .flexi_dataloader import FLEXIDataloader
 from .psp_dataloader import PSPDataloader
 from .tau_dataloader import TAUDataloader

flowtorch/data/hdf5_file.py

@@ -9,7 +9,6 @@ class that allows efficient batch access to simulation data stored in the
 
 from .dataloader import Dataloader
 from .foam_dataloader import FOAMCase, FOAMMesh, FOAMDataloader, POLYMESH_PATH, MAX_LINE_HEADER, FIELD_TYPE_DIMENSION
-from .xdmf import XDMFWriter
 from .mpi_tools import main_only, main_bcast, job_conditional, log_message
 from os.path import exists
 from os import remove
@@ -24,6 +23,17 @@ class that allows efficient batch access to simulation data stored in the
 CONNECTIVITY_DS = "connectivity"
 CENTERS_DS = "centers"
 VOLUMES_DS = "volumes"
+XDMF_HEADER = """<source lang="xml" line="1">
+    <Domain>
+        <Grid Name="flowTorch" GridType="Collection" CollectionType="Temporal">
+"""
+XDMF_FOOTER = """
+        </Grid>
+    </Domain>
+</source> 
+"""
+TOPOLOGY = "topology"
+GEOMETRY = "geometry"
 
 dtype_conversion = {
     pt.float32: "f4",
@@ -36,6 +46,25 @@ class that allows efficient batch access to simulation data stored in the
     "int64": "i8"
 }
 
+xdmf_attributes = {
+    1: "Scalar",
+    2: "Vector",
+    3: "Vector",
+    6: "Tensor6",
+    9: "Tensor"
+}
+
+xdmf_dtype_conversion = {
+    "f4": ("Float", 4),
+    "f8": ("Float", 8),
+    "i4": ("Int", 4),
+    "i8": ("Int", 8),
+    "float32": ("Float", 4),
+    "float64": ("Float", 8),
+    "int32": ("Int", 4),
+    "int64": ("Int", 8)
+}
+
 
 class HDF5Dataloader(Dataloader):
     """
@@ -141,8 +170,10 @@ def write(self,
                     str(dtype), name)
             )
 
+    @main_only
     def write_xdmf(self):
-        pass
+        writer = XDMFWriter(self._file_path, self._file)
+        writer.create_xdmf("flowtorch.xdmf")
 
 
 class FOAM2HDF5(object):
@@ -179,6 +210,7 @@ def convert(self, filename, skip_zero=True):
             log_message("Converting fields.")
             self._convert_fields(writer, skip_zero)
             log_message("Conversion finished.")
+            writer.write_xdmf()
 
     @main_only
     def _remove_file_if_present(self, file_path):
@@ -195,13 +227,13 @@ def _convert_mesh(self, writer):
         mesh_path = self._loader._case._path + "/" + POLYMESH_PATH
         n_cells, n_points, n_top = self._gather_mesh_information(mesh_path)
         data = self._get_vertices(mesh_path, job=0)
-        writer.write(VERTICES_DS, (n_points, 3), data)
+        writer.write(VERTICES_DS, (n_points, 3), data, None, self._dtype)
         data = self._get_topology(job=0)
-        writer.write(CONNECTIVITY_DS, (n_top,), data)
+        writer.write(CONNECTIVITY_DS, (n_top,), data, None, pt.int32)
         data = self._get_cell_centers(job=1)
-        writer.write(CENTERS_DS, (n_cells, 3), data)
+        writer.write(CENTERS_DS, (n_cells, 3), data, None, self._dtype)
         data = self._get_cell_volumes(job=1)
-        writer.write(VOLUMES_DS, (n_cells,), data)
+        writer.write(VOLUMES_DS, (n_cells,), data, None, self._dtype)
 
     @main_bcast
     def _gather_mesh_information(self, mesh_path):
@@ -313,3 +345,130 @@ def load_n_lines(file_name, n):
     @job_conditional
     def _load_field(self, field, time, job=0):
         return self._loader.load_snapshot(field, time)
+
+
+class XDMFWriter(object):
+    def __init__(self, file_path: str, hdf5_file: File):
+        if "/" in file_path:
+            self._path = file_path[:file_path.rfind("/")]
+            self._hdf5_filename = file_path[file_path.rfind("/") + 1:]
+        else:
+            self._path = "."
+            self._hdf5_filename = file_path
+        self._file = hdf5_file
+        self._n_cells = self._get_n_cells()
+
+    @classmethod
+    def from_filepath(cls, file_path: str):
+        return cls(
+            file_path,
+            File(file_path, mode="a", driver="mpio", comm=MPI.COMM_WORLD)
+        )
+
+    def _get_n_cells(self) -> int:
+        n_cells = 0
+        location = "/{:s}/{:s}".format(CONST_GROUP, VOLUMES_DS)
+        if location in self._file:
+            n_cells = self._file[location].shape[0]
+        if n_cells == 0:
+            log_message("XDMF warning: could not determine number of cells.")
+        return n_cells
+
+    def _add_grid(self, time: str, offset: str = "") -> str:
+        """
+        """
+        grid = offset + "<Grid Name=\"Grid\" GridType=\"Uniform\">\n"
+        grid += self._add_topology(offset + " "*4)
+        grid += self._add_geometry(offset + " "*4)
+        if time is not None:
+            grid += offset + " "*4 + "<Time Value=\"{:s}\"/>\n".format(time)
+            for key in self._find_attributes(time):
+                grid += self._add_attribute(time, key, offset + " "*4)
+        grid += offset + "</Grid>\n"
+        return grid
+
+    def _find_attributes(self, time: str) -> list[str]:
+        location = "/{:s}/{:s}".format(VAR_GROUP, time)
+        keys = self._file[location].keys()
+        valid_attr = []
+        for key in keys:
+            loc = location + "/{:s}".format(key)
+            first_dim = self._file[loc].shape[0]
+            if first_dim == self._n_cells:
+                valid_attr.append(key)
+        return valid_attr
+
+    def _add_topology(self, offset: str = "") -> str:
+        topology = offset + \
+            "<Topology Name=\"{:s}\" TopologyType=\"Mixed\">\n".format(
+                TOPOLOGY)
+        location = self._hdf5_filename + \
+            ":/{:s}/{:s}".format(CONST_GROUP, CONNECTIVITY_DS)
+        topology += self._add_dataitem(location, offset + " "*4)
+        topology += offset + "</Topology>\n"
+        return topology
+
+    def _add_geometry(self, offset: str = "") -> str:
+        geometry = offset + "<Geometry GeometryType=\"XYZ\">\n"
+        location = self._hdf5_filename + \
+            ":/{:s}/{:s}".format(CONST_GROUP, VERTICES_DS)
+        geometry += self._add_dataitem(location, offset + " "*4)
+        geometry += offset + "</Geometry>\n"
+        return geometry
+
+    def _add_attribute(self, time: str, name: str, offset: str = "") -> str:
+        location = self._hdf5_filename + ":/{:s}/{:s}/{:s}".format(
+            VAR_GROUP, time, name
+        )
+        shape = self._file[location.split(":")[-1]].shape
+        tensor_type = xdmf_attributes[len(shape)]
+        attribute = offset + "<Attribute Name=\"{:s}\" AttributeType=\"{:s}\" Center=\"Cell\">\n".format(
+            name, tensor_type
+        )
+        attribute += self._add_dataitem(location, offset + " "*4)
+        attribute += offset + "</Attribute>\n"
+        return attribute
+
+    def _add_dataitem(self, location: str,  offset: str = "") -> str:
+        path_in_file = location.split(":")[-1]
+        shape = self._file[path_in_file].shape
+        dimensions = " ".join(["{:d}".format(i) for i in shape])
+        dtype, precision = xdmf_dtype_conversion[
+            str(self._file[path_in_file].dtype)
+        ]
+        dataitem = offset + "<DataItem Dimensions=\"{:s}\" NumberType=\"{:s}\" Precision=\"{:d}\" Format=\"HDF\">\n".format(
+            dimensions, dtype, precision
+        )
+        dataitem += offset + " "*4 + "{:s}\n".format(location)
+        dataitem += offset + "</DataItem>\n"
+        return dataitem
+
+    def create_xdmf(self, filename: str = None):
+        """
+
+        :param filename: [description]
+        :type filename: [type]
+        """
+        xdmf_str = XDMF_HEADER
+        times = list(self._file[VAR_GROUP].keys())
+        if len(times) > 0:
+            for time in times:
+                xdmf_str += self._add_grid(time, " "*12)
+        else:
+            xdmf_str += self._add_grid(None, " "*12)
+        xdmf_str = xdmf_str[:-1]  # remove last linebreak
+        xdmf_str += XDMF_FOOTER
+
+        if filename is None:
+            if "." in self._hdf5_filename:
+                filename = self._hdf5_filename[:self._hdf5_filename.rfind(
+                    ".")] + ".xdmf"
+            else:
+                filename = self._hdf5_filename + ".xdmf"
+        log_message(
+            "Writing file {:s} as wrapper for {:s} at location {:s}".format(
+                filename, self._hdf5_filename, self._path
+            )
+        )
+        with open(self._path + "/" + filename, "w") as file:
+            file.write(xdmf_str)

test/data/test_hdf5_file.py

@@ -4,7 +4,8 @@
 import torch as pt
 from h5py import File
 from mpi4py import MPI
-from flowtorch.data import HDF5Dataloader, HDF5Writer, FOAM2HDF5
+from flowtorch.data import HDF5Dataloader, HDF5Writer, FOAM2HDF5, XDMFWriter
+
 
 class HDF5TestData:
     def __init__(self):
@@ -15,7 +16,8 @@ def __init__(self):
             "of_cavity_ascii_parallel",
             "of_cavity_binary_parallel"
         ]
-        self.const_group = sorted(["vertices", "connectivity", "centers", "volumes"])
+        self.const_group = sorted(
+            ["vertices", "connectivity", "centers", "volumes"])
         self.var_group = ["0.1", "0.2", "0.3", "0.4", "0.5"]
 
 
@@ -32,23 +34,43 @@ def test_write(self, get_test_data):
         writer.write("twos", (3, 2), pt.ones(3, 2)*2, "0.01")
         writer.write("threes", (3, 2), pt.ones(3, 2)*3, "0.03")
         del writer
-        hdf5_file = File(file_path, mode="a", driver="mpio", comm=MPI.COMM_WORLD)
+        hdf5_file = File(file_path, mode="a",
+                         driver="mpio", comm=MPI.COMM_WORLD)
         assert os.path.isfile(file_path)
         assert list(hdf5_file["variable"].keys()) == ["0.01", "0.03"]
         hdf5_file.close()
 
     def test_write_const(self, get_test_data):
         file_path = get_test_data.test_path + "test_file.hdf5"
         writer = HDF5Writer(file_path)
-        writer.write("zeros", (3, 2), pt.zeros(3,2), dtype=pt.float64)
-        writer.write("zeros_single", (3, 2), pt.zeros(3,2), dtype=pt.float32)
+        writer.write("zeros", (3, 2), pt.zeros(
+            (3, 2), dtype=pt.float64), dtype=pt.float64)
+        writer.write("zeros_single", (3, 2), pt.zeros(
+            (3, 2), dtype=pt.float32), dtype=pt.float32)
+        writer.write("zeros_int", (3, 2), pt.zeros(
+            (3, 2), dtype=pt.int32), dtype=pt.int32)
         del writer
-        hdf5_file = File(file_path, mode="a", driver="mpio", comm=MPI.COMM_WORLD)
-        assert list(hdf5_file["constant"].keys()) == ["zeros", "zeros_single"]
+        hdf5_file = File(file_path, mode="a",
+                         driver="mpio", comm=MPI.COMM_WORLD)
+        assert list(hdf5_file["constant"].keys()) == [
+            "zeros", "zeros_int", "zeros_single"]
         assert hdf5_file["constant/zeros"].dtype == "float64"
         assert hdf5_file["constant/zeros_single"].dtype == "float32"
+        assert hdf5_file["constant/zeros_int"].dtype == "int32"
         hdf5_file.close()
 
+    def test_write_xdmf(self, get_test_data):
+        case = get_test_data.test_cases[0]
+        case_path = get_test_data.test_path + case
+        converter = FOAM2HDF5(case_path)
+        converter.convert("flowtorch.hdf5")
+        del converter
+        file_path = case_path + "/flowtorch.hdf5"
+        writer = XDMFWriter.from_filepath(file_path)
+        writer.create_xdmf()
+        assert os.path.isfile(case_path + "/flowtorch.xdmf")
+
+
 class TestFOAM2HDF5():
     def test_convert(self, get_test_data):
         for case in get_test_data.test_cases:
@@ -58,11 +80,12 @@ def test_convert(self, get_test_data):
             del converter
             filename = case_path + "/flowtorch.hdf5"
             if os.path.isfile(filename):
-                hdf5_file = File(filename, mode="a", driver="mpio", comm=MPI.COMM_WORLD)
+                hdf5_file = File(filename, mode="a",
+                                 driver="mpio", comm=MPI.COMM_WORLD)
                 const_keys = sorted(hdf5_file["constant"].keys())
                 assert const_keys == get_test_data.const_group
                 var_keys = sorted(hdf5_file["variable"].keys())
                 assert var_keys == get_test_data.var_group
                 assert hdf5_file["constant/volumes"].shape[0] == 400
                 assert hdf5_file["constant/centers"].shape == (400, 3)
-                hdf5_file.close()            
+                hdf5_file.close()