Add get_mesh_dimension() and check_mesh_shape().

src/tetrax/helpers/io.py

+"""
+This submodule provides some helper functions for input and output of files,
+as well as some mesh checkers.
+"""
+
 import numpy as np
 import meshio
 
+# from meshio
+topological_dimension = {
+    "line": 1,
+    "polygon": 2,
+    "triangle": 2,
+    "quad": 2,
+    "tetra": 3,
+    "hexahedron": 3,
+    "wedge": 3,
+    "pyramid": 3,
+    "line3": 1,
+    "triangle6": 2,
+    "quad9": 2,
+    "tetra10": 3,
+    "hexahedron27": 3,
+    "wedge18": 3,
+    "pyramid14": 3,
+    "vertex": 0,
+    "quad8": 2,
+    "hexahedron20": 3,
+    "triangle10": 2,
+    "triangle15": 2,
+    "triangle21": 2,
+    "line4": 1,
+    "line5": 1,
+    "line6": 1,
+    "tetra20": 3,
+    "tetra35": 3,
+    "tetra56": 3,
+    "quad16": 2,
+    "quad25": 2,
+    "quad36": 2,
+    "triangle28": 2,
+    "triangle36": 2,
+    "triangle45": 2,
+    "triangle55": 2,
+    "triangle66": 2,
+    "quad49": 2,
+    "quad64": 2,
+    "quad81": 2,
+    "quad100": 2,
+    "quad121": 2,
+    "line7": 1,
+    "line8": 1,
+    "line9": 1,
+    "line10": 1,
+    "line11": 1,
+    "tetra84": 3,
+    "tetra120": 3,
+    "tetra165": 3,
+    "tetra220": 3,
+    "tetra286": 3,
+    "wedge40": 3,
+    "wedge75": 3,
+    "hexahedron64": 3,
+    "hexahedron125": 3,
+    "hexahedron216": 3,
+    "hexahedron343": 3,
+    "hexahedron512": 3,
+    "hexahedron729": 3,
+    "hexahedron1000": 3,
+    "wedge126": 3,
+    "wedge196": 3,
+    "wedge288": 3,
+    "wedge405": 3,
+    "wedge550": 3,
+    "VTK_LAGRANGE_CURVE": 1,
+    "VTK_LAGRANGE_TRIANGLE": 2,
+    "VTK_LAGRANGE_QUADRILATERAL": 2,
+    "VTK_LAGRANGE_TETRAHEDRON": 3,
+    "VTK_LAGRANGE_HEXAHEDRON": 3,
+    "VTK_LAGRANGE_WEDGE": 3,
+    "VTK_LAGRANGE_PYRAMID": 3,
+}
+
+
 def read_mode_from_vtk(fname, as_flattened=False):
     """Reads a complex-valued spatial mode profile from a ´vtk` file.
 
@@ -14,7 +95,8 @@ def read_mode_from_vtk(fname, as_flattened=False):
     Returns
     -------
     numpy.array
-        Mode profile as mesh vector of shape `(N, 3)` if `as_flattened=False` or as flattened mesh vector of shape `(3*N,)` if `as_flattened=True`.
+        Mode profile as mesh vector of shape `(N, 3)` if `as_flattened=False` or
+        as flattened mesh vector of shape `(3*N,)` if `as_flattened=True`.
 
     Notes
     -----
@@ -25,4 +107,67 @@ def read_mode_from_vtk(fname, as_flattened=False):
     my = mode.point_data["Re(m_y)"] + 1j * mode.point_data["Im(m_y)"]
     mz = mode.point_data["Re(m_z)"] + 1j * mode.point_data["Im(m_z)"]
 
-    return np.array([mx, my, mz]).flatten() if as_flattened else np.array([mx, my, mz]).T
\ No newline at end of file
+    return np.array([mx, my, mz]).flatten() if as_flattened else np.array([mx, my, mz]).T
+
+
+def get_mesh_dimension(mesh: meshio.Mesh) -> int:
+    """Obtain the dimension of a mesh.
+
+    The dimension of the mesh is obtained as the highest topoligical dimension
+    of the different cell types in the mesh.
+
+    Parameters
+    ----------
+    mesh : meshio.Mesh
+        Mesh object.
+
+    Returns
+    -------
+    int
+        Dimension of the mesh (0, 1, 2 or 3).
+    """
+
+    mesh_dim = max([topological_dimension[cellblock.type] for cellblock in mesh.cells])
+
+    return mesh_dim
+
+
+def check_mesh_shape(mesh: meshio.Mesh, dim: int) -> bool:
+    """Check if mesh is properly shaped.
+
+    - `dim=3`: always `True`
+    - `dim=2`: `True` if mesh is in `xy` plane.
+    - `dim=1`: `True` if mesh is on `y` axis.
+
+    Parameters
+    ----------
+    mesh : meshio.Mesh
+        Mesh object.
+    dim : int
+        Dimension of mesh.
+
+    Returns
+    -------
+    bool
+    """
+
+    if dim == 3:
+        properly_shaped = True
+
+    elif dim == 2:
+        # check if all z coordinates are 0.
+        z = mesh.points.T[2]
+        properly_shaped = np.any(z) == 0
+
+    elif dim == 1:
+        # check if all z and x coordinates are 0.
+        x = mesh.points.T[0]
+        z = mesh.points.T[2]
+        properly_shaped = np.any(z) == 0 and np.any(x) == 0
+
+    else:
+        raise ValueError(
+            'Illegal value for "dim", should be 1, 2, or 3.'
+        )
+
+    return properly_shaped