From 19579784ceb70a9de7a47731caa8e1ce78e5a340 Mon Sep 17 00:00:00 2001
From: Jaewook Lee <jaewooklee042@gmail.com>
Date: Fri, 22 Sep 2023 10:39:23 +0200
Subject: [PATCH] update meshio.load

---
 examples/show_gmsh.py |  7 +++---
 gustaf/io/meshio.py   | 52 ++++++++++++++++++++++++-------------------
 2 files changed, 33 insertions(+), 26 deletions(-)

diff --git a/examples/show_gmsh.py b/examples/show_gmsh.py
index 66ea1d2f1..38d3e5e63 100644
--- a/examples/show_gmsh.py
+++ b/examples/show_gmsh.py
@@ -7,6 +7,7 @@
 
 import load_sample_file
 
+import gustaf
 from gustaf import io
 
 if __name__ == "__main__":
@@ -20,15 +21,15 @@
     # load the .msh file directly with the correct io module (meshio)
     loaded_mesh_tri = io.meshio.load(base_samples_path / mesh_file_tri)
 
-    loaded_mesh_tri.show()
+    gustaf.show(loaded_mesh_tri)
 
     # load the .msh file directly with the correct io module (meshio)
     loaded_mesh_quad = io.meshio.load(base_samples_path / mesh_file_quad)
 
-    loaded_mesh_quad.show()
+    gustaf.show(loaded_mesh_quad)
 
     # load the .msh file with the default load function which needs to find out
     # it self which module is the correct one.
     loaded_mesh_default = io.load(base_samples_path / mesh_file_tri)
 
-    loaded_mesh_default.show()
+    gustaf.show(loaded_mesh_default)
diff --git a/gustaf/io/meshio.py b/gustaf/io/meshio.py
index 7dfb2499a..7027adf4f 100644
--- a/gustaf/io/meshio.py
+++ b/gustaf/io/meshio.py
@@ -7,15 +7,25 @@
 
 import numpy as np
 
+from gustaf.edges import Edges
 from gustaf.faces import Faces
 from gustaf.helpers.raise_if import ModuleImportRaiser
+from gustaf.utils import log
+from gustaf.vertices import Vertices
 from gustaf.volumes import Volumes
 
 try:
     import meshio
 except ModuleNotFoundError as err:
     meshio = ModuleImportRaiser("meshio", err)
-    # from meshio import Mesh as MeshioMesh
+
+_meshio2gus = {
+    "hexahedron": Volumes,
+    "tetra": Volumes,
+    "quad": Faces,
+    "triangle": Faces,
+    "line": Edges,
+}
 
 
 def load(fname):
@@ -36,8 +46,7 @@ def load(fname):
     --------
     MESH_TYPES
     """
-    mesh_type = Faces
-
+    # fname sanity check
     fname = pathlib.Path(fname)
     if not (fname.exists() and fname.is_file()):
         raise ValueError(
@@ -45,31 +54,28 @@ def load(fname):
             f"{fname.resolve()}"
         )
 
+    # load
     meshio_mesh: meshio.Mesh = meshio.read(fname)
+
+    # first get vertices
     vertices = meshio_mesh.points
 
-    # check for 2D mesh
-    # Try for triangle grid
-    cells = meshio_mesh.get_cells_type("triangle")
-    # If  no triangle elements, try for square
-    if len(cells) == 0:
-        cells = meshio_mesh.get_cells_type("quad")
-    if not len(cells) > 0:
-        # 3D mesh
-        mesh_type = Volumes
-        cells = meshio_mesh.get_cells_type("tetra")
-        if len(cells) == 0:
-            cells = meshio_mesh.get_cells_type("hexahedron")
-
-    for i, cell in enumerate(cells):
-        if i == 0:
-            elements = cell.data
-        else:
-            elements = np.vstack((elements, cell.data))
+    # early exit if cells doesn't exist
+    if len(meshio_mesh.cells_dict) == 0:
+        return Vertices(vertices)
+
+    meshes = []
+    for element_type, elements in meshio_mesh.cells_dict.items():
+        # skip unsupported
+        if element_type not in _meshio2gus:
+            log.warning(
+                f"`{element_type}`-elements are not supported in gustaf"
+            )
+            continue
 
-    mesh = mesh_type(vertices=vertices, elements=elements)
+        meshes.append(_meshio2gus[element_type](vertices, elements=elements))
 
-    return mesh
+    return meshes[0] if len(meshes) == 1 else meshes
 
 
 def export(mesh, fname, submeshes=None, **kwargs):