feat: add exodus reader

Author: gilrrei

src/fourc_webviewer/read_geometry_from_file.py

@@ -0,0 +1,306 @@
+"""
+I/O for Exodus II. This file is based on the meshio project see:
+
+https://github.com/nschloe/meshio/blob/main/src/meshio/exodus/_exodus.py
+"""
+
+import re
+
+import numpy as np
+
+from meshio.__about__ import __version__
+from meshio._common import warn
+from meshio._exceptions import ReadError
+from meshio._mesh import Mesh
+
+exodus_to_meshio_type = {
+    "SPHERE": "vertex",
+    # curves
+    "BEAM": "line",
+    "BEAM2": "line",
+    "BEAM3": "line3",
+    "BAR2": "line",
+    # surfaces
+    "SHELL": "quad",
+    "SHELL4": "quad",
+    "SHELL8": "quad8",
+    "SHELL9": "quad9",
+    "QUAD": "quad",
+    "QUAD4": "quad",
+    "QUAD5": "quad5",
+    "QUAD8": "quad8",
+    "QUAD9": "quad9",
+    #
+    "TRI": "triangle",
+    "TRIANGLE": "triangle",
+    "TRI3": "triangle",
+    "TRI6": "triangle6",
+    "TRI7": "triangle7",
+    # 'TRISHELL': 'triangle',
+    # 'TRISHELL3': 'triangle',
+    # 'TRISHELL6': 'triangle6',
+    # 'TRISHELL7': 'triangle',
+    #
+    # volumes
+    "HEX": "hexahedron",
+    "HEXAHEDRON": "hexahedron",
+    "HEX8": "hexahedron",
+    "HEX9": "hexahedron9",
+    "HEX20": "hexahedron20",
+    "HEX27": "hexahedron27",
+    #
+    "TETRA": "tetra",
+    "TETRA4": "tetra4",
+    "TET4": "tetra4",
+    "TETRA8": "tetra8",
+    "TETRA10": "tetra10",
+    "TETRA14": "tetra14",
+    #
+    "PYRAMID": "pyramid",
+    "WEDGE": "wedge",
+}
+meshio_to_exodus_type = {v: k for k, v in exodus_to_meshio_type.items()}
+
+
+def _categorize(names):
+    # Check if there are any <name>R, <name>Z tuples or <name>X, <name>Y, <name>Z
+    # triplets in the point data. If yes, they belong together.
+    single = []
+    double = []
+    triple = []
+    is_accounted_for = [False] * len(names)
+    k = 0
+    while True:
+        if k == len(names):
+            break
+        if is_accounted_for[k]:
+            k += 1
+            continue
+        name = names[k]
+        if name[-1] == "X":
+            ix = k
+            try:
+                iy = names.index(name[:-1] + "Y")
+            except ValueError:
+                iy = None
+            try:
+                iz = names.index(name[:-1] + "Z")
+            except ValueError:
+                iz = None
+            if iy and iz:
+                triple.append((name[:-1], ix, iy, iz))
+                is_accounted_for[ix] = True
+                is_accounted_for[iy] = True
+                is_accounted_for[iz] = True
+            else:
+                single.append((name, ix))
+                is_accounted_for[ix] = True
+        elif name[-2:] == "_R":
+            ir = k
+            try:
+                iz = names.index(name[:-2] + "_Z")
+            except ValueError:
+                iz = None
+            if iz:
+                double.append((name[:-2], ir, iz))
+                is_accounted_for[ir] = True
+                is_accounted_for[iz] = True
+            else:
+                single.append((name, ir))
+                is_accounted_for[ir] = True
+        else:
+            single.append((name, k))
+            is_accounted_for[k] = True
+
+        k += 1
+
+    if not all(is_accounted_for):
+        raise ReadError()
+    return single, double, triple
+
+
+def read_exodus(filename, use_set_names=False):  # noqa: C901
+    import netCDF4
+
+    with netCDF4.Dataset(filename) as nc:
+        # assert nc.version == np.float32(5.1)
+        # assert nc.api_version == np.float32(5.1)
+        # assert nc.floating_point_word_size == 8
+
+        # assert b''.join(nc.variables['coor_names'][0]) == b'X'
+        # assert b''.join(nc.variables['coor_names'][1]) == b'Y'
+        # assert b''.join(nc.variables['coor_names'][2]) == b'Z'
+
+        points = np.zeros((len(nc.dimensions["num_nodes"]), 3))
+        point_data_names = []
+        cell_data_names = []
+        pd = {}
+        cd = {}
+        cells = []
+        ns_names = []
+        eb_names = []
+        ns = []
+        point_sets = {}
+        info = []
+
+        cell_sets = {}
+
+        element_running_index = 0
+
+        for key, value in nc.variables.items():
+            if key == "info_records":
+                value.set_auto_mask(False)
+                for c in value[:]:
+                    info += [b"".join(c).decode("UTF-8")]
+            elif key == "qa_records":
+                value.set_auto_mask(False)
+                for val in value:
+                    info += [b"".join(c).decode("UTF-8") for c in val[:]]
+            elif key[:7] == "connect":
+                meshio_type = exodus_to_meshio_type[value.elem_type.upper()]
+                cell_sets[str(len(cell_sets) + 1)] = np.arange(
+                    element_running_index, element_running_index + len(value[:])
+                )
+                cells.append((meshio_type, value[:] - 1))
+            elif key == "coord":
+                points = nc.variables["coord"][:].T
+            elif key == "coordx":
+                points[:, 0] = value[:]
+            elif key == "coordy":
+                points[:, 1] = value[:]
+            elif key == "coordz":
+                points[:, 2] = value[:]
+            elif key == "name_nod_var":
+                value.set_auto_mask(False)
+                point_data_names = [b"".join(c).decode("UTF-8") for c in value[:]]
+            elif key[:12] == "vals_nod_var":
+                idx = 0 if len(key) == 12 else int(key[12:]) - 1
+                value.set_auto_mask(False)
+                # For now only take the first value
+                pd[idx] = value[0]
+                if len(value) > 1:
+                    warn("Skipping some time data")
+            elif key == "name_elem_var":
+                value.set_auto_mask(False)
+                cell_data_names = [b"".join(c).decode("UTF-8") for c in value[:]]
+            elif key[:13] == "vals_elem_var":
+                # eb: element block
+                m = re.match("vals_elem_var(\\d+)?(?:eb(\\d+))?", key)
+                idx = 0 if m.group(1) is None else int(m.group(1)) - 1
+                block = 0 if m.group(2) is None else int(m.group(2)) - 1
+
+                value.set_auto_mask(False)
+                # For now only take the first value
+                if idx not in cd:
+                    cd[idx] = {}
+                cd[idx][block] = value[0]
+
+                if len(value) > 1:
+                    warn("Skipping some time data")
+            elif key == "ns_names":
+                value.set_auto_mask(False)
+                ns_names = [b"".join(c).decode("UTF-8") for c in value[:]]
+            elif key == "eb_names":
+                value.set_auto_mask(False)
+                eb_names = [b"".join(c).decode("UTF-8") for c in value[:]]
+            elif key.startswith("node_ns"):  # Expected keys: node_ns1, node_ns2
+                ns.append(value[:] - 1)  # Exodus is 1-based
+
+        # merge element block data; can't handle blocks yet
+        for k, value in cd.items():
+            cd[k] = np.concatenate(list(value.values()))
+
+        # Check if there are any <name>R, <name>Z tuples or <name>X, <name>Y, <name>Z
+        # triplets in the point data. If yes, they belong together.
+        single, double, triple = _categorize(point_data_names)
+
+        point_data = {}
+        for name, idx in single:
+            point_data[name] = pd[idx]
+        for name, idx0, idx1 in double:
+            point_data[name] = np.column_stack([pd[idx0], pd[idx1]])
+        for name, idx0, idx1, idx2 in triple:
+            point_data[name] = np.column_stack([pd[idx0], pd[idx1], pd[idx2]])
+
+        cell_data = {}
+        k = 0
+        for _, cell in cells:
+            n = len(cell)
+            for name, data in zip(cell_data_names, cd.values()):
+                if name not in cell_data:
+                    cell_data[name] = []
+                cell_data[name].append(data[k : k + n])
+            k += n
+
+    point_sets = {str(i + 1): dat.tolist() for i, dat in enumerate(ns)}
+
+    if use_set_names:
+        point_sets = {name: dat for name, dat in zip(ns_names, ns)}
+        cell_sets = {
+            name: cell_set for name, cell_set in zip(eb_names, cell_sets.values())
+        }
+
+    return Mesh(
+        points,
+        cells,
+        point_data=point_data,
+        cell_data=cell_data,
+        point_sets=point_sets,
+        info=info,
+        cell_sets=cell_sets,
+    )
+
+
+class FourCGeometry:
+
+    def __init__(
+        self,
+        mesh: Mesh,
+        element_blocks: dict[str, np.array],
+        node_sets: dict[str, np.array],
+    ):
+        self.mesh = mesh
+        self.element_blocks = element_blocks
+        self.node_sets = node_sets
+
+    def get_element_ids_of_block(self, element_block_id):
+        return self.element_blocks[element_block_id]
+
+    def get_node_sets(self, node_set_id):
+        return self.node_sets[node_set_id]
+
+    @classmethod
+    def from_exodus(cls, file_path):
+        mesh = read_exodus(file_path, False)
+        element_blocks = mesh.cell_sets.copy()
+        node_sets = mesh.point_sets.copy()
+        mesh.cell_sets = {}
+        mesh.point_sets = {}
+
+        return cls(mesh, element_blocks, node_sets)
+
+    def __str__(self):
+        string = "4C Geometry"
+        if self.mesh.info:
+            string += "\n  Info"
+            for s in self.mesh.info:
+                string += "\n    " + s
+        string += f"\n Nodes: {len(self.mesh.points)}"
+        string += f"\n Cells: {sum([len(c) for c in self.mesh.cells])}"
+        if self.element_blocks:
+            string += f"\n Element Blocks"
+            for e, v in self.element_blocks.items():
+                string += f"\n    {e}: {len(v)} cells"
+        if self.node_sets:
+            string += f"\n Node Sets"
+            for e, v in self.node_sets.items():
+                string += f"\n    {e}: {len(v)} nodes"
+
+        return string
+
+
+from pathlib import Path
+
+mesh = FourCGeometry.from_exodus(
+    Path(__file__).parents[2] / "tests/files/tutorial_solid_geo.e"
+)