Merge pull request nschloe#254 from hillyuan/master

Add node sets read & write function

meshio/exodus_io.py

@@ -80,6 +80,9 @@ def read(filename):
     point_data_names = []
     pd = []
     cells = {}
+    ns_names = []
+    ns = []
+    node_sets = {}
     for key, value in nc.variables.items():
         if key[:7] == "connect":
             meshio_type = exodus_to_meshio_type[value.elem_type.upper()]
@@ -100,11 +103,17 @@ def read(filename):
             point_data_names = [b"".join(c).decode("UTF-8") for c in value[:]]
         elif key == "vals_nod_var":
             pd = value[0, :]
+        elif key == "ns_names":
+            value.set_auto_mask(False)
+            ns_names = [b"".join(c).decode("UTF-8") for c in value[:]]
+        elif key == "node_ns":
+            ns = value
 
     point_data = {name: dat for name, dat in zip(point_data_names, pd)}
+    node_sets = {name: dat for name, dat in zip(ns_names, ns)}
 
     nc.close()
-    return Mesh(points, cells, point_data=point_data)
+    return Mesh(points, cells, point_data=point_data, node_sets=node_sets)
 
 
 numpy_to_exodus_dtype = {
@@ -140,6 +149,7 @@ def write(filename, mesh):
     rootgrp.createDimension("num_dim", 3)
     rootgrp.createDimension("num_elem", total_num_elems)
     rootgrp.createDimension("num_el_blk", len(mesh.cells))
+    rootgrp.createDimension("num_node_sets", len(mesh.node_sets))
     rootgrp.createDimension("len_string", 33)
     rootgrp.createDimension("len_line", 81)
     rootgrp.createDimension("four", 4)
@@ -202,5 +212,24 @@ def write(filename, mesh):
         for k, (name, data) in enumerate(mesh.point_data.items()):
             node_data[0, k] = data
 
+    # node sets
+    num_node_sets = len(mesh.node_sets)
+    if num_node_sets > 0:
+        data = rootgrp.createVariable("ns_prop1", "i4", "num_node_sets")
+        data_names = rootgrp.createVariable(
+            "ns_names", "S1", ("num_node_sets", "len_string")
+        )
+        for k, name in enumerate(mesh.node_sets.keys()):
+            data[k] = k
+            for i, letter in enumerate(name):
+                data_names[k, i] = letter.encode("utf-8")
+        for k, (key, values) in enumerate(mesh.node_sets.items()):
+            dim1 = "num_nod_ns{}".format(k + 1)
+            rootgrp.createDimension(dim1, values.shape[0])
+            dtype = numpy_to_exodus_dtype[values.dtype.name]
+            data = rootgrp.createVariable("node_ns{}".format(k + 1), dtype, (dim1,))
+            # Exodus is 1-based
+            data[:] = values + 1
+
     rootgrp.close()
     return

meshio/mesh.py

@@ -4,12 +4,21 @@
 
 
 class Mesh(object):
-    def __init__(self, points, cells, point_data=None, cell_data=None, field_data=None):
+    def __init__(
+        self,
+        points,
+        cells,
+        point_data=None,
+        cell_data=None,
+        field_data=None,
+        node_sets=None,
+    ):
         self.points = points
         self.cells = cells
         self.point_data = point_data if point_data else {}
         self.cell_data = cell_data if cell_data else {}
         self.field_data = field_data if field_data else {}
+        self.node_sets = node_sets if node_sets else {}
         return
 
     def __repr__(self):
@@ -19,6 +28,9 @@ def __repr__(self):
         for tpe, elems in self.cells.items():
             lines.append("  Number of {}s: {}".format(tpe, len(elems)))
 
+        if self.node_sets:
+            lines.append("Node sets: {}".format(", ".join(self.node_sets.keys())))
+
         if self.point_data:
             lines.append("Point data: {}".format(", ".join(self.point_data.keys())))
 

test/helpers.py

@@ -216,6 +216,12 @@ def add_field_data(mesh, value, dtype):
     return mesh2
 
 
+def add_node_sets(mesh):
+    mesh2 = copy.deepcopy(mesh)
+    mesh2.node_sets = {"fixed": numpy.array([1, 2])}
+    return mesh2
+
+
 def write_read(writer, reader, input_mesh, atol):
     """Write and read a file, and make sure the data is the same as before.
     """

test/test_exodus.py

@@ -23,6 +23,8 @@
     helpers.add_point_data(helpers.tri_mesh, 1),
     # helpers.add_point_data(helpers.tri_mesh, 2),
     # helpers.add_point_data(helpers.tri_mesh, 3),
+    helpers.add_node_sets(helpers.tri_mesh),
+    helpers.add_node_sets(helpers.tet_mesh),
 ]