diff --git a/vayesta/core/types/wf/ccsd.py b/vayesta/core/types/wf/ccsd.py index 036514291..77eed2751 100644 --- a/vayesta/core/types/wf/ccsd.py +++ b/vayesta/core/types/wf/ccsd.py @@ -9,7 +9,7 @@ from vayesta.core import spinalg from vayesta.core.util import NotCalculatedError, Object, callif, einsum, dot from vayesta.core.types import wf as wf_types -from vayesta.core.types.orbitals import SpatialOrbitals +from vayesta.core.types.orbitals import SpatialOrbitals, SpinOrbitals from vayesta.core.types.wf.project import ( project_c1, project_c2, @@ -398,27 +398,29 @@ def rotate_ov(self, to, tv, inplace=False): wf.l2 = transform_uc2(wf.l2, to, tv) return wf - # def pack(self, dtype=float): - # """Pack into a single array of data type `dtype`. - - # Useful for communication via MPI.""" - # mo = self.mo.pack(dtype=dtype) - # l1 = self.l1 is not None else [None, None] - # l2 = self.l2 is not None else len(self.t2)*[None] - # projector = self.projector is not None else [None] - # data = (mo, *self.t1, *self.t2, *l1, *l2, *projector) - # pack = pack_arrays(*data, dtype=dtype) - # return pack - - # @classmethod - # def unpack(cls, packed): - # """Unpack from a single array of data type `dtype`. - - # Useful for communication via MPI.""" - # mo, *unpacked = unpack_arrays(packed) - # mo = SpinOrbitals.unpack(mo) - # t1a, t1b, t2, l1, l2, projector = - # wf = cls(mo, t1, t2, l1=l1, l2=l2) - # if projector is not None: - # wf.projector = projector - # return wf + def pack(self, dtype=float): + """Pack into a single array of data type `dtype`. + Useful for communication via MPI.""" + mo = self.mo.pack(dtype=dtype) + l1 = self.l1 if self.l1 is not None else [None, None] + l2 = self.l2 if self.l2 is not None else len(self.t2)*[None] + projector=self.projector + data = (mo, *self.t1, *self.t2, *l1, *l2, *projector) + pack = pack_arrays(*data, dtype=dtype) + return pack + + @classmethod + def unpack(cls, packed): + """Unpack from a single array of data type `dtype`. + Useful for communication via MPI.""" + mo, t1a, t1b, t2aa, t2ab, t2ba, t2bb, l1a, l1b, l2aa, l2ab, l2ba, l2bb, proja, projb = unpack_arrays(packed) + t1 = (t1a, t1b) + t2 = (t2aa, t2ab, t2ba, t2bb) + l1 = (l1a, l1b) + l2 = (l2aa, l2ab, l2ba, l2bb) + projector = (proja, projb) + mo = SpinOrbitals.unpack(mo) + wf = cls(mo, t1, t2, l1=l1, l2=l2) + if projector is not None: + wf.projector = projector + return wf