add in ability to save as numpy archive (#35)

mapdl_archive/archive.py

@@ -37,7 +37,8 @@ class Archive(Mesh):
     Parameters
     ----------
     filename : string, pathlib.Path
-        Filename of block formatted cdb file
+        Supports blocked MAPDL archive .cdb files or .npz written from this
+        class using :class:`Archive.save_as_numpy`.
 
     read_parameters : bool, default: False
         Optionally read parameters from the archive file.
@@ -123,12 +124,31 @@ def __init__(
         self._read_parameters: bool = read_parameters
         self._filename: pathlib.Path = pathlib.Path(filename)
         self._name: str = name
-        self._raw = _reader.read(
-            self.filename,
-            read_parameters=read_parameters,
-            debug=verbose,
-            read_eblock=read_eblock,
-        )
+        if self.filename.endswith("npz"):
+            npz_file = np.load(self.filename, allow_pickle=True)
+
+            self._raw: _reader.ReadReturnDict = {
+                "rnum": npz_file["rnum"],
+                "rdat": npz_file["rdat"].tolist(),
+                "ekey": npz_file["ekey"],
+                "nnum": npz_file["nnum"],
+                "nodes": npz_file["nodes"],
+                "elem": npz_file["elem"],
+                "elem_off": npz_file["elem_off"],
+                "node_comps": npz_file["node_comps"].item(),
+                "elem_comps": npz_file["elem_comps"].item(),
+                "keyopt": npz_file["keyopt"].item(),
+                "parameters": npz_file["parameters"].item(),
+                "nblock_start": npz_file["nblock_start"],
+                "nblock_end": npz_file["nblock_end"],
+            }
+        else:
+            self._raw = _reader.read(
+                self.filename,
+                read_parameters=read_parameters,
+                debug=verbose,
+                read_eblock=read_eblock,
+            )
         super().__init__(
             self._raw["nnum"],
             self._raw["nodes"],
@@ -289,6 +309,18 @@ def overwrite_nblock(self, filename, node_id, pos, angles=None, sig_digits=13):
 
             shutil.copyfileobj(src_file, dest_file)
 
+    def save_as_numpy(self, filename: str):
+        """Save this archive as a numpy "npz" file.
+
+        This reduces the file size by around 50% compared with the Ansys
+        blocked file format.
+
+        """
+        if not filename.endswith("npz"):
+            raise ValueError("Filename should end with 'npz'")
+
+        np.savez(filename, **self._raw)  # type: ignore
+
 
 def save_as_archive(
     filename: Union[pathlib.Path, str],

tests/test_archive.py

@@ -38,6 +38,15 @@
 DAT_FILE = os.path.join(TESTFILES_PATH, "Panel_Transient.dat")
 
 
+def compare_dicts_with_arrays(dict1, dict2):
+    if dict1.keys() != dict2.keys():
+        return False
+    for key in dict1:
+        if not np.array_equal(dict1[key], dict2[key]):
+            return False
+    return True
+
+
 def proto_cmblock(array):
     """prototype cmblock code"""
     items = np.zeros_like(array)
@@ -592,3 +601,21 @@ def test_filename_setter_pathlib(self, pathlib_archive):
     def test_filename_setter_string(self, pathlib_archive):
         with pytest.raises(AttributeError):
             pathlib_archive.filename = "dummy2"
+
+
+def test_save_as_numpy(tmpdir, hex_archive):
+    """Test saving to and loading from a npz file."""
+    path = str(tmpdir.join("data.npz"))
+    hex_archive.save_as_numpy(path)
+    hex_in = Archive(path)
+
+    assert hex_archive._raw.keys() == hex_in._raw.keys()
+
+    for key, value in hex_archive._raw.items():
+        if isinstance(value, np.ndarray):
+            assert np.allclose(value, hex_in._raw[key])
+        else:
+            if isinstance(value, dict):
+                compare_dicts_with_arrays(value, hex_in._raw[key])
+            else:
+                assert value == hex_in._raw[key]