kernel: move packing/unpacking into class methods

miniff/kernel.py

@@ -9,7 +9,8 @@ from scipy.sparse import coo_matrix, csr_matrix
 from scipy.spatial import cKDTree
 
 from itertools import product
-from collections import namedtuple, OrderedDict
+from collections import OrderedDict
+from typing import NamedTuple
 try:
     from functools import cached_property
 except ImportError:
@@ -712,78 +713,56 @@ class SnapshotHistory(list):
         self.append(cell)
 
 
-sf_parameters = namedtuple("sf_parameters", ("cartesian", "vectors"))
-
+class sf_parameters(NamedTuple):
+    cartesian: np.ndarray
+    vectors: np.ndarray
 
-def pack_coordinates(c=None, v=None, vc=None, vv=None):
-    """
-    A convention to pack coordinates.
+    def copy(self, c=True, v=True):
+        """
+        Makes a copy.
 
-    Parameters
-    ----------
-    c : np.ndarray
-        Atomic coordinates array.
-    v : np.ndarray
-        Cell vectors array.
-    vc : np.ndarray
-        Atomic velocities array.
-    vv : np.ndarray
-        (Fake) Cell vector velocities.
+        Parameters
+        ----------
+        c : bool
+            Copy coordinates, if available.
+        v : bool
+            Copy vectors, if available.
 
-    Returns
-    -------
-    result : np.ndarray
-        A 2D array with packed coordinates.
-    """
-    if vc is not None:
-        assert c is not None
-    if vv is not None:
-        assert v is not None
-    args = (c, v, vc, vv)
-    assert sum(i is not None for i in args) != 3
-    args = tuple(i for i in args if i is not None)
-    return np.concatenate(args, axis=0)
+        Returns
+        -------
+        result : sf_parameters
+            The resulting parameters.
+        """
+        return sf_parameters(self.cartesian if c else None, self.vectors if v else None)
 
+    def pack(self):
+        return np.concatenate(tuple(i for i in self if i is not None))
 
-def unpack_coordinates(params, c=False, v=False, vc=False, vv=False, nv=3):
-    """
-    A convention to unpack coordinates.
+    @staticmethod
+    def unpack(data, c=False, v=False, nv=3):
+        """
+        Unpacks the data into parameters.
 
-    Parameters
-    ----------
-    params : np.ndarray
-        The packed coordinates.
-    c : bool
-        Indicates the presence of atomic coordinates.
-    v : bool
-        Indicates the presence of vectors.
-    vc : bool
-        Indicates the presence of atomic velocities.
-    vv : bool
-        Indicates the presence of (fake) cell vector velocities.
-    nv : int
-        Number of vectors.
+        Parameters
+        ----------
+        data : np.ndarray
+            The packed parameters.
+        c : bool
+            Indicates the presence of atomic coordinates.
+        v : bool
+            Indicates the presence of vectors.
+        nv : int
+            Number of vectors.
 
-    Returns
-    -------
-    result : tuple
-        Coordinates and velocities as separate arrays.
-    """
-    assert (c, vc) != (False, True)
-    assert (v, vv) != (False, True)
-    args = (c, v, vc, vv)
-    assert sum(args) != 3
-    if vc or vv:
-        params_c, params_v = params[:len(params) // 2], params[len(params) // 2:]
-    else:
-        params_c = params
-        params_v = None
-    return (
-        params_c[:-nv] if c and v else params_c if c else None,
-        params_c[-nv:] if v else None,
-        params_v[:-nv] if vc and vv else params_v if vc else None,
-        params_v[-nv:] if vv else None,
-    )
+        Returns
+        -------
+        result : sf_parameters
+            The resulting parameters.
+        """
+        return sf_parameters(
+            data[:-nv] if c and v else data if c else None,
+            data[-nv:] if v else None,
+        )
 
 
 class ScalarFunctionWrapper:
@@ -834,18 +813,18 @@ class ScalarFunctionWrapper:
 
     def p2c(self, parameters: np.ndarray) -> sf_parameters:
         """Unpacks parameters."""
-        return sf_parameters(*unpack_coordinates(
+        return sf_parameters.unpack(
             parameters.reshape(-1, 3),
             c=self.include["cartesian"],
             v=self.include["vectors"],
-        )[:2])
+        )
 
     def c2p(self, coordinates: sf_parameters) -> np.ndarray:
         """Packs parameters."""
-        return pack_coordinates(
-            c=coordinates.cartesian if self.include["cartesian"] else None,
-            v=coordinates.vectors if self.include["vectors"] else None,
-        ).reshape(-1)
+        return coordinates.copy(
+            c=self.include["cartesian"],
+            v=self.include["vectors"],
+        ).pack().reshape(-1)
 
     def start_recording(self):
         """Starts recording of coordinates passed."""