From 66afff6fcb141227f1ed64ce2eb9fa5db78f83a2 Mon Sep 17 00:00:00 2001
From: Abhishek Khedkar <abhishek.khedkar@kcl.ac.uk>
Date: Tue, 5 Sep 2023 18:39:33 +0100
Subject: [PATCH] The unpack, pack functions needed for MPI communications were
 not overridden for the unrestricted case. This lead to the errors seen in
 emb.make_rdm1() calls with MPI.

---
 vayesta/core/types/wf/ccsd.py | 52 ++++++++++++++++++-----------------
 1 file changed, 27 insertions(+), 25 deletions(-)

diff --git a/vayesta/core/types/wf/ccsd.py b/vayesta/core/types/wf/ccsd.py
index a4ac69082..9f1630c6c 100644
--- a/vayesta/core/types/wf/ccsd.py
+++ b/vayesta/core/types/wf/ccsd.py
@@ -7,7 +7,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, project_uc1, project_uc2, symmetrize_c2,
                                            symmetrize_uc2, transform_c1, transform_c2, transform_uc1, transform_uc2)
 from vayesta.core.helper import pack_arrays, unpack_arrays
@@ -383,27 +383,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