Skip to content
Snippets Groups Projects
Commit 5b9b1a73 authored by Artem Pulkin's avatar Artem Pulkin
Browse files

kernel: move packing/unpacking into class methods

parent d22f394a
No related branches found
No related tags found
No related merge requests found
Pipeline #79998 passed
......@@ -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."""
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment