Merge pull request #523 from gdmcbain/522-xdmf-timeseries-vectors

WIP: extend AttributeType to XdmfTimeSeries
nschloe · Nov 19, 2019 · eea9293 · eea9293
2 parents deab277 + 88666d3
commit eea9293
Show file tree

Hide file tree

Showing 6 changed files with 43 additions and 33 deletions.
diff --git a/meshio/_cli.py b/meshio/_cli.py
@@ -134,7 +134,7 @@ def info(argv=None):
 
 def _get_info_parser():
     parser = argparse.ArgumentParser(
-        description=("Print mesh info."), formatter_class=argparse.RawTextHelpFormatter,
+        description=("Print mesh info."), formatter_class=argparse.RawTextHelpFormatter
     )
 
     parser.add_argument("infile", type=str, help="mesh file to be read from")

diff --git a/meshio/_flac3d.py b/meshio/_flac3d.py
@@ -8,19 +8,10 @@
 from .__about__ import __version__ as version
 from ._mesh import Mesh
 
-meshio_only = {
-    "tetra",
-    "pyramid",
-    "wedge",
-    "hexahedron",
-}
+meshio_only = {"tetra", "pyramid", "wedge", "hexahedron"}
 
 
-meshio_data = {
-    "flac3d:zone",
-    "gmsh:physical",
-    "medit:ref",
-}
+meshio_data = {"flac3d:zone", "gmsh:physical", "medit:ref"}
 
 
 flac3d_to_meshio_type = {

diff --git a/meshio/_xdmf/common.py b/meshio/_xdmf/common.py
@@ -132,3 +132,20 @@ def translate_mixed_cells(data):
         cells[meshio_type] = data[indices]
 
     return cells
+
+
+def attribute_type(data):
+    # <http://www.xdmf.org/index.php/XDMF_Model_and_Format#Attribute>
+    if len(data.shape) == 1 or (len(data.shape) == 2 and data.shape[1] == 1):
+        return "Scalar"
+    elif len(data.shape) == 2 and data.shape[1] in [2, 3]:
+        return "Vector"
+    elif (len(data.shape) == 2 and data.shape[1] == 9) or (
+        len(data.shape) == 3 and data.shape[1] == 3 and data.shape[2] == 3
+    ):
+        return "Tensor"
+    elif len(data.shape) == 2 and data.shape[1] == 6:
+        return "Tensor6"
+
+    assert len(data.shape) == 3
+    return "Matrix"
diff --git a/meshio/_xdmf/main.py b/meshio/_xdmf/main.py
@@ -10,6 +10,7 @@
 from .._mesh import Mesh
 from .._vtk import raw_from_cell_data
 from .common import (
+    attribute_type,
     dtype_to_format_string,
     meshio_to_xdmf_type,
     meshio_type_to_xdmf_index,
@@ -423,22 +424,6 @@ def cells(self, cells, grid):
             data_item.text = self.numpy_to_xml_string(cd)
         return
 
-    def _attribute_type(self, data):
-        # <http://www.xdmf.org/index.php/XDMF_Model_and_Format#Attribute>
-        if len(data.shape) == 1 or (len(data.shape) == 2 and data.shape[1] == 1):
-            return "Scalar"
-        elif len(data.shape) == 2 and data.shape[1] in [2, 3]:
-            return "Vector"
-        elif (len(data.shape) == 2 and data.shape[1] == 9) or (
-            len(data.shape) == 3 and data.shape[1] == 3 and data.shape[2] == 3
-        ):
-            return "Tensor"
-        elif len(data.shape) == 2 and data.shape[1] == 6:
-            return "Tensor6"
-
-        assert len(data.shape) == 3
-        return "Matrix"
-
     def point_data(self, point_data, grid):
         from lxml import etree as ET
 
@@ -447,7 +432,7 @@ def point_data(self, point_data, grid):
                 grid,
                 "Attribute",
                 Name=name,
-                AttributeType=self._attribute_type(data),
+                AttributeType=attribute_type(data),
                 Center="Node",
             )
             dt, prec = numpy_to_xdmf_dtype[data.dtype.name]
@@ -472,7 +457,7 @@ def cell_data(self, cell_data, grid):
                 grid,
                 "Attribute",
                 Name=name,
-                AttributeType=self._attribute_type(data),
+                AttributeType=attribute_type(data),
                 Center="Cell",
             )
             dt, prec = numpy_to_xdmf_dtype[data.dtype.name]

diff --git a/meshio/_xdmf/time_series.py b/meshio/_xdmf/time_series.py
@@ -5,6 +5,7 @@
 from .._common import cell_data_from_raw, write_xml
 from .._vtk import raw_from_cell_data
 from .common import (
+    attribute_type,
     dtype_to_format_string,
     meshio_to_xdmf_type,
     meshio_type_to_xdmf_index,
@@ -408,7 +409,13 @@ def point_data(self, point_data, grid):
         from lxml import etree as ET
 
         for name, data in point_data.items():
-            att = ET.SubElement(grid, "Attribute", Name=name, Center="Node")
+            att = ET.SubElement(
+                grid,
+                "Attribute",
+                Name=name,
+                AttributeType=attribute_type(data),
+                Center="Node",
+            )
             dt, prec = numpy_to_xdmf_dtype[data.dtype.name]
             dim = " ".join([str(s) for s in data.shape])
             data_item = ET.SubElement(
@@ -427,7 +434,11 @@ def cell_data(self, cell_data, grid):
         raw = raw_from_cell_data(cell_data)
         for name, data in raw.items():
             att = ET.SubElement(
-                grid, "Attribute", Name=name, Type="None", Center="Cell"
+                grid,
+                "Attribute",
+                Name=name,
+                AttributeType=attribute_type(data),
+                Center="Cell",
             )
             dt, prec = numpy_to_xdmf_dtype[data.dtype.name]
             dim = " ".join([str(s) for s in data.shape])

diff --git a/test/test_xdmf.py b/test/test_xdmf.py
@@ -59,7 +59,13 @@ def test_time_series():
         n = helpers.tri_mesh_2d.points.shape[0]
 
         times = numpy.linspace(0.0, 1.0, 5)
-        point_data = [{"phi": numpy.full(n, t)} for t in times]
+        point_data = [
+            {
+                "phi": numpy.full(n, t),
+                "u": numpy.full(helpers.tri_mesh_2d.points.shape, t),
+            }
+            for t in times
+        ]
         for t, pd in zip(times, point_data):
             writer.write_data(
                 t, point_data=pd, cell_data={"triangle": {"a": [3.0, 4.2]}}