Add npz export for gustaf entities

examples/export_npz.py

@@ -0,0 +1,115 @@
+import numpy as np
+
+import gustaf
+
+if __name__ == "__main__":
+    # Define coordinates
+    vertices = np.array(
+        [
+            [0.0, 0.0, 0.0],
+            [1.0, 0.0, 0.0],
+            [0.0, 1.0, 0.0],
+            [1.0, 1.0, 0.0],
+            [0.0, 0.0, 1.0],
+            [1.0, 0.0, 1.0],
+            [0.0, 1.0, 1.0],
+            [1.0, 1.0, 1.0],
+        ]
+    )
+
+    # Define hexa connectivity
+    hv = np.array([[0, 1, 3, 2, 4, 5, 7, 6]])
+
+    # Init hexa elements
+    hexa = gustaf.Volumes(
+        vertices=vertices,
+        volumes=hv,
+    )
+
+    # Define tet connectivity
+    tv = np.array(
+        [
+            [0, 2, 7, 3],
+            [0, 2, 6, 7],
+            [0, 6, 4, 7],
+            [5, 0, 4, 7],
+            [5, 0, 7, 1],
+            [7, 0, 3, 1],
+        ]
+    )
+
+    # Create tetra elements and set vertex data
+    tet = gustaf.Volumes(vertices=vertices, volumes=tv)
+    tet.vertex_data["arange"] = -np.arange(len(vertices))
+    tet.BC = {"min_z": np.array([0, 1, 2, 3]), "max_z": np.array([4, 5, 6, 7])}
+
+    # Create hexa elements and set vertex data
+    hexa.vertex_data["arange"] = np.arange(len(vertices))
+    hexa.vertex_data["quad_x"] = [v[0] ** 2 for v in vertices]
+    hexa.BC = {
+        "min_z": np.array([0, 1, 2, 3]),
+        "max_z": np.array([4, 5, 6, 7]),
+    }
+
+    hexa.show_options["data_name"] = "arange"
+
+    # Export hexa, edges and vertices
+    gustaf.io.npz.export(hexa, "export/hexa.npz")
+    gustaf.io.npz.export(hexa.to_edges(), "export/hexa_edges.npz")
+    gustaf.io.npz.export(hexa.to_vertices(), "export/hexa_vertices.npz")
+
+    # Load hexa files
+    load_hexa = gustaf.io.npz.load("export/hexa.npz")
+    load_edges = gustaf.io.npz.load("export/hexa_edges.npz")
+    load_verts = gustaf.io.npz.load("export/hexa_vertices.npz")
+
+    # Show original and loaded file
+    # Show options are not exported and have to be reset
+    gustaf.show(
+        ["Original hexa", hexa],
+        ["Loaded hexa (show options must be set again)", load_hexa],
+    )
+
+    # Show the expored edges and vertices
+    gustaf.show(
+        ["Exported edges", load_edges],
+        ["Exported vertices", load_verts],
+        [
+            "Exported hexa with indicated boundary nodes",
+            load_hexa,
+            gustaf.Vertices(load_hexa.vertices[load_hexa.BC["min_z"]]),
+            gustaf.Vertices(load_hexa.vertices[load_hexa.BC["max_z"]]),
+        ],
+    )
+
+    # Vertex data has to be set for edges and vertices as to_edges() and
+    # to_vertices() does not copy it
+    load_hexa.show_options["data_name"] = "arange"
+
+    load_edges.vertex_data["quad_x"] = hexa.vertex_data["quad_x"]
+    load_edges.show_options["data_name"] = "quad_x"
+
+    load_verts.vertex_data["arange"] = hexa.vertex_data["arange"]
+    load_verts.show_options["data_name"] = "arange"
+
+    # Show colored vertices
+    gustaf.show(
+        ["Show entities with copied data", load_hexa], load_edges, load_verts
+    )
+
+    # Export tetra files
+    gustaf.io.npz.export(tet, "export/tet.npz")
+    gustaf.io.npz.export(tet.to_edges(), "export/tet_edges.npz")
+    gustaf.io.npz.export(tet.to_vertices(), "export/tet_vertices.npz")
+
+    # Load tetra files with default load function
+    load_tet = gustaf.io.load("export/tet.npz")
+    load_tet_edges = gustaf.io.load("export/tet_edges.npz")
+    load_tet_verts = gustaf.io.load("export/tet_vertices.npz")
+
+    # Show tetras
+    gustaf.show(
+        ["Tetra", load_tet],
+        ["Tetra edges", load_tet_edges],
+        ["Tetra vertices", load_verts],
+    )

gustaf/io/__init__.py

@@ -5,14 +5,7 @@
 O - `export`.
 """
 
-from gustaf.io import ioutils, meshio, mfem, mixd, nutils
+from gustaf.io import ioutils, meshio, mfem, mixd, npz, nutils
 from gustaf.io.default import load
 
-__all__ = [
-    "ioutils",
-    "mfem",
-    "meshio",
-    "mixd",
-    "nutils",
-    "load",
-]
+__all__ = ["ioutils", "mfem", "meshio", "mixd", "nutils", "load", "npz"]

gustaf/io/default.py

@@ -1,6 +1,6 @@
 import pathlib
 
-from gustaf.io import meshio, mfem, mixd
+from gustaf.io import meshio, mfem, mixd, npz
 
 
 def load(fname):
@@ -22,6 +22,7 @@ def load(fname):
         ".mixd": mixd.load,
         ".mfem": mfem.load,
         ".msh": meshio.load,
+        ".npz": npz.load,
     }
 
     fname = pathlib.Path(fname).resolve()

gustaf/io/npz.py

@@ -0,0 +1,147 @@
+import numpy as np
+
+from gustaf.edges import Edges
+from gustaf.faces import Faces
+from gustaf.vertices import Vertices
+from gustaf.volumes import Volumes
+
+
+def load(fname, **kwargs):
+    """
+    Read gus_object from `.npz` file.
+
+    .. code-block:: python
+
+        tet = gustaf.io.npz.load('export/tet.npz')
+        tet.show()
+
+    Returns
+    --------
+    gus_object: gustaf entity
+        Can be vertices, Edges, Faces or Volumes
+    """
+
+    # Load the file and check entity type
+    loaded = np.load(fname, **kwargs)
+    dict_keys = list(loaded.keys())
+    whatami = str(loaded["whatami"])
+
+    gus_object_types = {
+        "edges": Edges,
+        "tri": Faces,
+        "quad": Faces,
+        "tet": Volumes,
+        "hexa": Volumes,
+    }
+
+    # Vertices
+    if whatami == "vertices":
+        gus_object = Vertices(vertices=loaded["vertices"])
+    # Meshes
+    elif whatami in gus_object_types.keys():
+        gus_object = gus_object_types[whatami](
+            vertices=loaded["vertices"], elements=loaded["elements"]
+        )
+    # Unknown types
+    else:
+        raise ValueError(f"Type {whatami} is not supported in gustaf.")
+
+    for key in dict_keys:
+        # Load vertex data
+        if key.startswith("vertex_data-"):
+            gus_object.vertex_data[key.removeprefix("vertex_data-")] = loaded[
+                key
+            ]
+        # Load boundaries
+        elif key.startswith("BC-"):
+            gus_object.BC[key.removeprefix("BC-")] = loaded[key]
+
+    return gus_object
+
+
+def export(gus_object, fname):
+    """
+    Save a gustaf object (Vertices, Edges, Faces or Volumes) as a `npz`-file.
+    The export file contains (if applicable): vertices, elements, vertex data,
+    BC.
+    The `npz`-format is a compressed numpy format
+    https://numpy.org/doc/stable/reference/generated/numpy.savez.html .
+
+    .. code-block:: python
+
+        import gustaf
+        import numpy as np
+
+        # Define coordinates
+        vertices = np.array(
+            [
+                [0.0, 0.0, 0.0],
+                [1.0, 0.0, 0.0],
+                [0.0, 1.0, 0.0],
+                [1.0, 1.0, 0.0],
+                [0.0, 0.0, 1.0],
+                [1.0, 0.0, 1.0],
+                [0.0, 1.0, 1.0],
+                [1.0, 1.0, 1.0],
+            ]
+        )
+
+        # Define hexa connectivity
+        hv = np.array([[0, 1, 3, 2, 4, 5, 7, 6]])
+
+        # Init hexa elements
+        hexa = gustaf.Volumes(
+            vertices=vertices,
+            volumes=hv,
+        )
+
+        # Create hexa elements and set vertex data
+        hexa.vertex_data['arange'] = np.arange(len(vertices))
+        hexa.vertex_data['quad_x'] = [v[0]**2 for v in vertices]
+        hexa.BC = {'min_z' : np.array([0, 1, 2, 3]),
+        'max_z' : np.array([4, 5, 6, 7])}
+
+        # Export hexa
+        gustaf.io.npz.export(hexa, 'export/hexa.npz')
+
+    Parameters
+    -----------
+    gus_object: gustaf entity
+    fname: str
+       Export filename with `npz` suffix
+
+    Returns
+    --------
+    None
+    """
+
+    if not fname.endswith(".npz"):
+        raise ValueError("The filename must end with .npz.")
+
+    property_dicts = {
+        "whatami": gus_object.whatami,
+        "vertices": gus_object.vertices,
+    }
+
+    # Include elements for meshes
+    if type(gus_object) in [Edges, Faces, Volumes]:
+        property_dicts["elements"] = gus_object.elements
+
+    # Export vertex data
+    property_dicts.update(
+        {
+            f"vertex_data-{key}": item
+            for key, item in gus_object.vertex_data.items()
+        }
+    )
+
+    # In case of faces and volumes, export the boundaries
+    if type(gus_object) in [Faces, Volumes]:
+        property_dicts.update(
+            {f"BC-{key}": item for key, item in gus_object.BC.items()}
+        )
+
+    np.savez(
+        fname,
+        **property_dicts,
+    )