refactor pe induction operator

adcc/OperatorIntegrals.py

@@ -32,16 +32,21 @@
 
 def transform_operator_ao2mo(tensor_bb, tensor_ff, coefficients,
                              conv_tol=1e-14):
-    """
-    Take a block-diagonal tensor in the atomic orbital basis
+    """Take a block-diagonal tensor in the atomic orbital basis
     and transform it into the molecular orbital basis in the
     convention used by adcc.
 
-    @param tensor_bb  Block-diagonal tensor in the atomic orbital basis
-    @param tensor_ff  Output tensor with the symmetry set-up to contain
-                      the operator in the molecular orbital representation
-    @param coefficients    Function providing coefficient blocks
-    @param conv_tol   SCF convergence tolerance
+    Parameters
+    ----------
+    tensor_bb : Tensor
+        Block-diagonal tensor in the atomic orbital basis
+    tensor_ff : Tensor
+        Output tensor with the symmetry set-up to contain
+        the operator in the molecular orbital representation
+    coefficients : callable
+        Function providing coefficient blocks
+    conv_tol : float, optional
+        SCF convergence tolerance, by default 1e-14
     """
     for blk in tensor_ff.blocks:
         assert len(blk) == 4
@@ -93,8 +98,13 @@ def provider_ao(self):
     @cached_property
     def available(self):
         """Which integrals are available in the underlying backend"""
-        return [integral
-                for integral in ("electric_dipole", "magnetic_dipole", "nabla")
+        integrals = (
+            "electric_dipole",
+            "magnetic_dipole",
+            "nabla",
+            "pe_induction_elec"
+        )
+        return [integral for integral in integrals
                 if hasattr(self.provider_ao, integral)]
 
     def import_dipole_like_operator(self, integral, is_symmetric=True):
@@ -137,25 +147,45 @@ def nabla(self):
         """
         return self.import_dipole_like_operator("nabla", is_symmetric=False)
 
+    def __import_density_dependent_operator(self, ao_callback, is_symmetric=True):
+        """Returns a function that imports a density-dependent operator.
+        The returned function imports the operator and transforms it to the
+        molecular orbital basis.
+
+        Parameters
+        ----------
+        ao_callback : callable
+            Function that computes the operator in atomic orbitals using
+            a :py:class:`OneParticleOperator` (the density matrix
+            in atomic orbitals) as single argument
+        is_symmetric : bool, optional
+            if the imported operator is symmetric, by default True
+        """
+        def process_operator(dm, callback=ao_callback, is_symmetric=is_symmetric):
+            dm_ao = sum(dm.to_ao_basis())
+            v_ao = callback(dm_ao)
+            v_bb = replicate_ao_block(
+                self.mospaces, v_ao, is_symmetric=is_symmetric
+            )
+            v_ff = OneParticleOperator(self.mospaces, is_symmetric=is_symmetric)
+            transform_operator_ao2mo(
+                v_bb, v_ff, self.__coefficients, self.__conv_tol
+            )
+            return v_ff
+        return process_operator
+
     @property
-    def density_dependent_operators(self):
-        if not hasattr(self.provider_ao, "density_dependent_operators"):
-            return {}
-        ret = {}
-        for op in self.provider_ao.density_dependent_operators:
-            ao_callback = self.provider_ao.density_dependent_operators[op]
-
-            def process_operator(dm, callback=ao_callback):
-                dm_ao = sum(dm.to_ao_basis())
-                v_ao = callback(dm_ao)
-                v_bb = replicate_ao_block(self.mospaces, v_ao, is_symmetric=True)
-                v_ff = OneParticleOperator(self.mospaces, is_symmetric=True)
-                transform_operator_ao2mo(
-                    v_bb, v_ff, self.__coefficients, self.__conv_tol
-                )
-                return v_ff
-            ret[op] = process_operator
-        return ret
+    def pe_induction_elec(self):
+        """
+        Returns a function to obtain the (density-dependent)
+        PE electronic induction operator in the molecular orbital basis
+        """
+        if "pe_induction_elec" not in self.available:
+            raise NotImplementedError("PE electronic induction operator "
+                                      "not implemented "
+                                      f"in {self.provider_ao.backend} backend.")
+        callback = self.provider_ao.pe_induction_elec
+        return self.__import_density_dependent_operator(callback)
 
     @property
     def timer(self):

adcc/adc_pp/solvent.py

@@ -37,6 +37,6 @@ def block_ph_ph_0_pe(hf, mp, intermediates):
     def apply(ampl):
         tdm = OneParticleOperator(mp, is_symmetric=False)
         tdm.vo = ampl.ph.transpose()
-        vpe = op.density_dependent_operators["pe_induction_elec"](tdm)
+        vpe = op.pe_induction_elec(tdm)
         return AmplitudeVector(ph=vpe.ov)
     return AdcBlock(apply, 0)

adcc/backends/psi4.py

@@ -55,15 +55,14 @@ def nabla(self):
         return [-1.0 * np.asarray(comp) for comp in self.mints.ao_nabla()]
 
     @property
-    def density_dependent_operators(self):
-        ret = {}
+    def pe_induction_elec(self):
         if hasattr(self.wfn, "pe_state"):
-            ret["pe_induction_elec"] = lambda dm: \
-                self.wfn.pe_state.get_pe_contribution(
+            def pe_induction_elec_ao(dm):
+                return self.wfn.pe_state.get_pe_contribution(
                     psi4.core.Matrix.from_array(dm.to_ndarray()),
                     elec_only=True
-            )[1]
-        return ret
+                )[1]
+            return pe_induction_elec_ao
 
 
 class Psi4EriBuilder(EriBuilder):

adcc/backends/pyscf.py

@@ -57,14 +57,14 @@ def nabla(self):
             )
 
     @property
-    def density_dependent_operators(self):
-        ret = {}
+    def pe_induction_elec(self):
         if hasattr(self.scfres, "with_solvent"):
             if isinstance(self.scfres.with_solvent, solvent.pol_embed.PolEmbed):
-                ret["pe_induction_elec"] = lambda dm: \
-                    self.scfres.with_solvent._exec_cppe(dm.to_ndarray(),
-                                                        elec_only=True)[1]
-        return ret
+                def pe_induction_elec_ao(dm):
+                    return self.scfres.with_solvent._exec_cppe(
+                        dm.to_ndarray(), elec_only=True
+                    )[1]
+                return pe_induction_elec_ao
 
 
 # TODO: refactor ERI builder to be more general