Merge pull request #68 from BoothGroup/namespace_eom

Use `Namespace` for EOM amplitudes
BoothGroup · Aug 7, 2024 · fbde438 · fbde438
2 parents 3ae4415 + 02ce987
commit fbde438
Show file tree

Hide file tree

Showing 10 changed files with 303 additions and 262 deletions.
diff --git a/ebcc/cc/base.py b/ebcc/cc/base.py
@@ -475,6 +475,7 @@ def _load_function(
         eris: Optional[Union[ERIsInputType, Literal[False]]] = False,
         amplitudes: Optional[Union[Namespace[SpinArrayType], Literal[False]]] = False,
         lambdas: Optional[Union[Namespace[SpinArrayType], Literal[False]]] = False,
+        excitations: Optional[Namespace[SpinArrayType]] = None,
         **kwargs: Any,
     ) -> tuple[Callable[..., Any], dict[str, Any]]:
         """Load a function from the generated code, and return the arguments."""
@@ -497,6 +498,10 @@ def _load_function(
             else:
                 dicts.append(dict(self._get_lams(lambdas=lambdas)))
 
+        # Get the excitation amplitudes:
+        if excitations:
+            dicts.append(dict(excitations))
+
         # Get the function:
         func = getattr(self._eqns, name, None)
         if func is None:
@@ -791,89 +796,83 @@ def make_eb_coup_rdm(
 
     def hbar_matvec_ip(
         self,
-        *excitations: SpinArrayType,
+        excitations: Namespace[SpinArrayType],
         eris: Optional[ERIsInputType] = None,
         amplitudes: Optional[Namespace[SpinArrayType]] = None,
-    ) -> tuple[SpinArrayType, SpinArrayType]:
+    ) -> Namespace[SpinArrayType]:
         """Compute the product between a state vector and the IP-EOM Hamiltonian.
 
         Args:
-            r1: State vector (single excitations).
-            r2: State vector (double excitations).
+            excitations: State vector as a set of excitation amplitudes.
             eris: Electron repulsion integrals.
             amplitudes: Cluster amplitudes.
 
         Returns:
             Products between the state vectors and the IP-EOM Hamiltonian for the singles and
             doubles.
         """
-        r1, r2 = excitations  # FIXME
         func, kwargs = self._load_function(
             "hbar_matvec_ip",
             eris=eris,
             amplitudes=amplitudes,
-            r1=r1,
-            r2=r2,
+            excitations=excitations,
         )
-        res: tuple[SpinArrayType, SpinArrayType] = func(**kwargs)
+        res: Namespace[SpinArrayType] = func(**kwargs)
+        res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()})
         return res
 
     def hbar_matvec_ea(
         self,
-        *excitations: SpinArrayType,
+        excitations: Namespace[SpinArrayType],
         eris: Optional[ERIsInputType] = None,
         amplitudes: Optional[Namespace[SpinArrayType]] = None,
-    ) -> tuple[SpinArrayType, SpinArrayType]:
+    ) -> Namespace[SpinArrayType]:
         """Compute the product between a state vector and the EA-EOM Hamiltonian.
 
         Args:
-            r1: State vector (single excitations).
-            r2: State vector (double excitations).
+            excitations: State vector as a set of excitation amplitudes.
             eris: Electron repulsion integrals.
             amplitudes: Cluster amplitudes.
 
         Returns:
             Products between the state vectors and the EA-EOM Hamiltonian for the singles and
             doubles.
         """
-        r1, r2 = excitations  # FIXME
         func, kwargs = self._load_function(
             "hbar_matvec_ea",
             eris=eris,
             amplitudes=amplitudes,
-            r1=r1,
-            r2=r2,
+            excitations=excitations,
         )
-        res: tuple[SpinArrayType, SpinArrayType] = func(**kwargs)
+        res: Namespace[SpinArrayType] = func(**kwargs)
+        res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()})
         return res
 
     def hbar_matvec_ee(
         self,
-        *excitations: SpinArrayType,
+        excitations: Namespace[SpinArrayType],
         eris: Optional[ERIsInputType] = None,
         amplitudes: Optional[Namespace[SpinArrayType]] = None,
-    ) -> tuple[SpinArrayType, SpinArrayType]:
+    ) -> Namespace[SpinArrayType]:
         """Compute the product between a state vector and the EE-EOM Hamiltonian.
 
         Args:
-            r1: State vector (single excitations).
-            r2: State vector (double excitations).
+            excitations: State vector as a set of excitation amplitudes.
             eris: Electron repulsion integrals.
             amplitudes: Cluster amplitudes.
 
         Returns:
             Products between the state vectors and the EE-EOM Hamiltonian for the singles and
             doubles.
         """
-        r1, r2 = excitations  # FIXME
         func, kwargs = self._load_function(
             "hbar_matvec_ee",
             eris=eris,
             amplitudes=amplitudes,
-            r1=r1,
-            r2=r2,
+            excitations=excitations,
         )
-        res: tuple[SpinArrayType, SpinArrayType] = func(**kwargs)
+        res: Namespace[SpinArrayType] = func(**kwargs)
+        res = util.Namespace(**{key.rstrip("new"): val for key, val in res.items()})
         return res
 
     def make_ip_mom_bras(
@@ -1089,7 +1088,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[SpinArrayType]:
         pass
 
     @abstractmethod
-    def excitations_to_vector_ip(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]:
+    def excitations_to_vector_ip(self, excitations: Namespace[SpinArrayType]) -> NDArray[float]:
         """Construct a vector containing all of the IP-EOM excitations.
 
         Args:
@@ -1101,7 +1100,7 @@ def excitations_to_vector_ip(self, *excitations: Namespace[SpinArrayType]) -> ND
         pass
 
     @abstractmethod
-    def excitations_to_vector_ea(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]:
+    def excitations_to_vector_ea(self, excitations: Namespace[SpinArrayType]) -> NDArray[float]:
         """Construct a vector containing all of the EA-EOM excitations.
 
         Args:
@@ -1113,7 +1112,7 @@ def excitations_to_vector_ea(self, *excitations: Namespace[SpinArrayType]) -> ND
         pass
 
     @abstractmethod
-    def excitations_to_vector_ee(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]:
+    def excitations_to_vector_ee(self, excitations: Namespace[SpinArrayType]) -> NDArray[float]:
         """Construct a vector containing all of the EE-EOM excitations.
 
         Args:
@@ -1125,9 +1124,7 @@ def excitations_to_vector_ee(self, *excitations: Namespace[SpinArrayType]) -> ND
         pass
 
     @abstractmethod
-    def vector_to_excitations_ip(
-        self, vector: NDArray[float]
-    ) -> tuple[Namespace[SpinArrayType], ...]:
+    def vector_to_excitations_ip(self, vector: NDArray[float]) -> Namespace[SpinArrayType]:
         """Construct a namespace of IP-EOM excitations from a vector.
 
         Args:
@@ -1139,9 +1136,7 @@ def vector_to_excitations_ip(
         pass
 
     @abstractmethod
-    def vector_to_excitations_ea(
-        self, vector: NDArray[float]
-    ) -> tuple[Namespace[SpinArrayType], ...]:
+    def vector_to_excitations_ea(self, vector: NDArray[float]) -> Namespace[SpinArrayType]:
         """Construct a namespace of EA-EOM excitations from a vector.
 
         Args:
@@ -1153,9 +1148,7 @@ def vector_to_excitations_ea(
         pass
 
     @abstractmethod
-    def vector_to_excitations_ee(
-        self, vector: NDArray[float]
-    ) -> tuple[Namespace[SpinArrayType], ...]:
+    def vector_to_excitations_ee(self, vector: NDArray[float]) -> Namespace[SpinArrayType]:
         """Construct a namespace of EE-EOM excitations from a vector.
 
         Args:

diff --git a/ebcc/cc/gebcc.py b/ebcc/cc/gebcc.py
@@ -815,7 +815,7 @@ def vector_to_lambdas(self, vector: NDArray[float]) -> Namespace[SpinArrayType]:
 
         return lambdas
 
-    def excitations_to_vector_ip(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]:
+    def excitations_to_vector_ip(self, excitations: Namespace[SpinArrayType]) -> NDArray[float]:
         """Construct a vector containing all of the IP-EOM excitations.
 
         Args:
@@ -825,12 +825,10 @@ def excitations_to_vector_ip(self, *excitations: Namespace[SpinArrayType]) -> ND
             IP-EOM excitations as a vector.
         """
         vectors = []
-        m = 0
 
         for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type):
             key = key[:-1]
-            vectors.append(util.compress_axes(key, excitations[m]).ravel())
-            m += 1
+            vectors.append(util.compress_axes(key, excitations[f"r{n}"]).ravel())
 
         for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type):
             raise util.ModelNotImplemented
@@ -840,7 +838,7 @@ def excitations_to_vector_ip(self, *excitations: Namespace[SpinArrayType]) -> ND
 
         return np.concatenate(vectors)
 
-    def excitations_to_vector_ea(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]:
+    def excitations_to_vector_ea(self, excitations: Namespace[SpinArrayType]) -> NDArray[float]:
         """Construct a vector containing all of the EA-EOM excitations.
 
         Args:
@@ -849,9 +847,9 @@ def excitations_to_vector_ea(self, *excitations: Namespace[SpinArrayType]) -> ND
         Returns:
             EA-EOM excitations as a vector.
         """
-        return self.excitations_to_vector_ip(*excitations)
+        return self.excitations_to_vector_ip(excitations)
 
-    def excitations_to_vector_ee(self, *excitations: Namespace[SpinArrayType]) -> NDArray[float]:
+    def excitations_to_vector_ee(self, excitations: Namespace[SpinArrayType]) -> NDArray[float]:
         """Construct a vector containing all of the EE-EOM excitations.
 
         Args:
@@ -861,11 +859,9 @@ def excitations_to_vector_ee(self, *excitations: Namespace[SpinArrayType]) -> ND
             EE-EOM excitations as a vector.
         """
         vectors = []
-        m = 0
 
         for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type):
-            vectors.append(util.compress_axes(key, excitations[m]).ravel())
-            m += 1
+            vectors.append(util.compress_axes(key, excitations[f"r{n}"]).ravel())
 
         for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type):
             raise util.ModelNotImplemented
@@ -875,9 +871,7 @@ def excitations_to_vector_ee(self, *excitations: Namespace[SpinArrayType]) -> ND
 
         return np.concatenate(vectors)
 
-    def vector_to_excitations_ip(
-        self, vector: NDArray[float]
-    ) -> tuple[Namespace[SpinArrayType], ...]:
+    def vector_to_excitations_ip(self, vector: NDArray[float]) -> Namespace[SpinArrayType]:
         """Construct a namespace of IP-EOM excitations from a vector.
 
         Args:
@@ -886,7 +880,7 @@ def vector_to_excitations_ip(
         Returns:
             IP-EOM excitations.
         """
-        excitations = []
+        excitations: Namespace[SpinArrayType] = util.Namespace()
         i0 = 0
 
         for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type):
@@ -895,7 +889,7 @@ def vector_to_excitations_ip(
             shape = tuple(self.space.size(k) for k in key)
             vn_tril = vector[i0 : i0 + size]
             vn = util.decompress_axes(key, vn_tril, shape=shape)
-            excitations.append(vn)
+            excitations[f"r{n}"] = vn
             i0 += size
 
         for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type):
@@ -904,11 +898,9 @@ def vector_to_excitations_ip(
         for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type):
             raise util.ModelNotImplemented
 
-        return tuple(excitations)
+        return excitations
 
-    def vector_to_excitations_ea(
-        self, vector: NDArray[float]
-    ) -> tuple[Namespace[SpinArrayType], ...]:
+    def vector_to_excitations_ea(self, vector: NDArray[float]) -> Namespace[SpinArrayType]:
         """Construct a namespace of EA-EOM excitations from a vector.
 
         Args:
@@ -917,7 +909,7 @@ def vector_to_excitations_ea(
         Returns:
             EA-EOM excitations.
         """
-        excitations = []
+        excitations: Namespace[SpinArrayType] = util.Namespace()
         i0 = 0
 
         for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type):
@@ -926,7 +918,7 @@ def vector_to_excitations_ea(
             shape = tuple(self.space.size(k) for k in key)
             vn_tril = vector[i0 : i0 + size]
             vn = util.decompress_axes(key, vn_tril, shape=shape)
-            excitations.append(vn)
+            excitations[f"r{n}"] = vn
             i0 += size
 
         for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type):
@@ -935,11 +927,9 @@ def vector_to_excitations_ea(
         for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type):
             raise util.ModelNotImplemented
 
-        return tuple(excitations)
+        return excitations
 
-    def vector_to_excitations_ee(
-        self, vector: NDArray[float]
-    ) -> tuple[Namespace[SpinArrayType], ...]:
+    def vector_to_excitations_ee(self, vector: NDArray[float]) -> Namespace[SpinArrayType]:
         """Construct a namespace of EE-EOM excitations from a vector.
 
         Args:
@@ -948,15 +938,15 @@ def vector_to_excitations_ee(
         Returns:
             EE-EOM excitations.
         """
-        excitations = []
+        excitations: Namespace[SpinArrayType] = util.Namespace()
         i0 = 0
 
         for name, key, n in self.ansatz.fermionic_cluster_ranks(spin_type=self.spin_type):
             size = util.get_compressed_size(key, **{k: self.space.size(k) for k in set(key)})
             shape = tuple(self.space.size(k) for k in key)
             vn_tril = vector[i0 : i0 + size]
             vn = util.decompress_axes(key, vn_tril, shape=shape)
-            excitations.append(vn)
+            excitations[f"r{n}"] = vn
             i0 += size
 
         for name, key, n in self.ansatz.bosonic_cluster_ranks(spin_type=self.spin_type):
@@ -965,7 +955,7 @@ def vector_to_excitations_ee(
         for name, key, nf, nb in self.ansatz.coupling_cluster_ranks(spin_type=self.spin_type):
             raise util.ModelNotImplemented
 
-        return tuple(excitations)
+        return excitations
 
     def get_mean_field_G(self) -> NDArray[float]:
         """Get the mean-field boson non-conserving term.