Merge pull request #2 from BauerMarco/qed_merge_check

adcc/AmplitudeVector.py

@@ -22,6 +22,8 @@
 ## ---------------------------------------------------------------------
 import warnings
 
+import numpy as np
+
 
 class AmplitudeVector(dict):
     def __init__(self, *args, **kwargs):
@@ -217,3 +219,140 @@ def __radd__(self, other):
         else:
             ret = {k: other + tensor for k, tensor in self.items()}
             return AmplitudeVector(**ret)
+
+
+class QED_AmplitudeVector:
+
+    # TODO: Initialize this class with **kwargs, and think of a functionality,
+    # which then eases up e.g. the matvec function in the AdcMatrix.py for 
+    # arbitrarily large QED vectors. However, this is not necessarily required, 
+    # since e.g. QED-ADC(3) is very complicated to derive and QED-ADC(1) (with
+    # just the single dispersion mode) is purely for academic purposes and
+    # hence not required to provide optimum performance.
+
+    def __init__(self, ph=None, pphh=None, gs1=None, ph1=None, pphh1=None,
+                gs2=None, ph2=None, pphh2=None):
+
+        if pphh != None:
+            self.elec = AmplitudeVector(ph=ph, pphh=pphh)
+            self.phot = AmplitudeVector(ph=ph1, pphh=pphh1)
+            self.phot2 = AmplitudeVector(ph=ph2, pphh=pphh2)
+        else:
+            self.elec = AmplitudeVector(ph=ph)
+            self.phot = AmplitudeVector(ph=ph1)
+            self.phot2 = AmplitudeVector(ph=ph2)
+
+        self.gs1 = gs1
+        self.gs2 = gs2
+        self.ph = ph
+        self.ph1 = ph1
+        self.ph2 = ph2
+        self.pphh = pphh
+        self.pphh1 = pphh1
+        self.pphh2 = pphh2
+
+
+    def dot(self, invec):
+        def dot_(self, invec):
+            if "pphh" in self.elec.blocks_ph:
+                return (self.elec.ph.dot(invec.elec.ph) + self.elec.pphh.dot(invec.elec.pphh)
+                        + self.gs1 * invec.gs1 + self.phot.ph.dot(invec.phot.ph) 
+                        + self.phot.pphh.dot(invec.phot.pphh)
+                        + self.gs2 * invec.gs2 + self.phot2.ph.dot(invec.phot2.ph) 
+                        + self.phot2.pphh.dot(invec.phot2.pphh))
+            else:
+                return (self.elec.ph.dot(invec.elec.ph) 
+                        + self.gs1 * invec.gs1 + self.phot.ph.dot(invec.phot.ph)
+                        + self.gs2 * invec.gs2 + self.phot2.ph.dot(invec.phot2.ph) )
+        if isinstance(invec, list):
+            return np.array([dot_(self, elem) for elem in invec])
+        else:
+           return dot_(self, invec)
+
+    def __matmul__(self, other):
+        if isinstance(other, QED_AmplitudeVector):
+            return self.dot(other)
+        if isinstance(other, list):
+            if all(isinstance(elem, QED_AmplitudeVector) for elem in other):
+                return self.dot(other)
+        return NotImplemented
+
+
+    def __sub__(self, invec):
+        if isinstance(invec, (float, int)):
+            if "pphh" in self.elec.blocks_ph:
+                return QED_AmplitudeVector(ph=self.elec.ph.__sub__(invec), 
+                        pphh=self.elec.pphh.__sub__(invec),
+                        gs1=self.gs1 - invec, ph1=self.phot.ph.__sub__(invec), 
+                        pphh1=self.phot.pphh.__sub__(invec),
+                        gs2=self.gs2 - invec, ph2=self.phot2.ph.__sub__(invec), 
+                        pphh2=self.phot2.pphh.__sub__(invec))
+            else:
+                return QED_AmplitudeVector(ph=self.elec.ph.__sub__(invec), 
+                        gs1=self.gs1 - invec, ph1=self.phot.ph.__sub__(invec),
+                        gs2=self.gs2 - invec, ph2=self.phot2.ph.__sub__(invec))
+
+
+    def __truediv__(self, other):
+        if isinstance(other, QED_AmplitudeVector):
+            if "pphh" in self.elec.blocks_ph:
+                return QED_AmplitudeVector(ph=self.elec.ph.__truediv__(other.elec.ph),
+                        pphh=self.elec.pphh.__truediv__(other.elec.pphh),
+                        gs1=self.gs1 / other.gs1, ph1=self.phot.ph.__truediv__(other.phot.ph),
+                        pphh1=self.phot.pphh.__truediv__(other.phot.pphh),
+                        gs2=self.gs2 / other.gs2, ph2=self.phot2.ph.__truediv__(other.phot2.ph), 
+                        pphh2=self.phot2.pphh.__truediv__(other.phot2.pphh))
+            else:
+                return QED_AmplitudeVector(ph=self.elec.ph.__truediv__(other.elec.ph), 
+                        gs1=self.gs1 / other.gs1, ph1=self.phot.ph.__truediv__(other.phot.ph),
+                        gs2=self.gs2 / other.gs2, ph2=self.phot2.ph.__truediv__(other.phot2.ph))
+        elif isinstance(other, (float, int)):
+            if "pphh" in self.elec.blocks_ph:
+                return QED_AmplitudeVector(ph=self.elec.ph.__truediv__(other), 
+                        pphh=self.elec.pphh.__truediv__(other),
+                        gs1=self.gs1 / other, ph1=self.phot.ph.__truediv__(other), 
+                        pphh1=self.phot.pphh.__truediv__(other),
+                        gs2=self.gs2 / other, ph2=self.phot2.ph.__truediv__(other), 
+                        pphh2=self.phot2.pphh.__truediv__(other))
+            else:
+                return QED_AmplitudeVector(ph=self.elec.ph.__truediv__(other), 
+                        gs1=self.gs1 / other, ph1=self.phot.ph.__truediv__(other),
+                        gs2=self.gs2 / other, ph2=self.phot2.ph.__truediv__(other))
+
+    def zeros_like(self):
+        if "pphh" in self.elec.blocks_ph:
+            return QED_AmplitudeVector(ph=self.elec.zeros_like().ph, 
+                        pphh=self.elec.zeros_like().pphh, 
+                        gs1=0, ph1=self.phot.zeros_like().ph, pphh1=self.phot.zeros_like().pphh, 
+                        gs2=0, ph2=self.phot2.zeros_like().ph, pphh2=self.phot2.zeros_like().pphh)
+        else:
+            return QED_AmplitudeVector(ph=self.elec.zeros_like().ph, pphh=None,
+                        gs1=0, ph1=self.phot.zeros_like().ph, pphh1=None, 
+                        gs2=0, ph2=self.phot2.zeros_like().ph, pphh2=None)
+
+    def empty_like(self):
+        if "pphh" in self.elec.blocks_ph:
+            return QED_AmplitudeVector(ph=self.elec.empty_like().ph, 
+                        pphh=self.elec.empty_like().pphh,
+                        gs1=0, ph1=self.phot.empty_like().ph, pphh1=self.phot.empty_like().pphh,
+                        gs2=0, ph2=self.phot2.empty_like().ph, pphh2=self.phot2.empty_like().pphh)
+        else:
+            QED_AmplitudeVector(ph=self.elec.empty_like().ph, pphh=None,
+                        gs1=0, ph1=self.phot.empty_like().ph, pphh1=None,
+                        gs2=0, ph2=self.phot2.empty_like().ph, pphh2=None)
+
+    def copy(self):
+        if "pphh" in self.elec.blocks_ph:
+            return QED_AmplitudeVector(ph=self.elec.copy().ph, pphh=self.elec.copy().pphh,
+                        gs1=self.gs1, ph1=self.phot.copy().ph, pphh1=self.phot.copy().pphh,
+                        gs2=self.gs2, ph2=self.phot2.copy().ph, pphh2=self.phot2.copy().pphh)
+        else:
+            QED_AmplitudeVector(ph=self.elec.copy().ph, pphh=None,
+                        gs1=self.gs1, ph1=self.phot.copy().ph, pphh1=None,
+                        gs2=self.gs2, ph2=self.phot2.copy().ph, pphh2=None)
+
+    def evaluate(self):
+        self.elec.evaluate()
+        self.phot.evaluate()
+        self.phot2.evaluate()
+        return self

adcc/ElectronicTransition.py

@@ -28,6 +28,10 @@
 from .visualisation import ExcitationSpectrum
 from .OneParticleOperator import product_trace
 from .AdcMethod import AdcMethod
+from adcc.functions import einsum
+from adcc import block as b
+from adcc.adc_pp.state2state_transition_dm import state2state_transition_dm
+from adcc.adc_pp.transition_dm import transition_dm
 
 from scipy import constants
 from .Excitation import mark_excitation_property
@@ -165,6 +169,126 @@ def transition_dipole_moment(self):
             for tdm in self.transition_dm
         ])
 
+
+    @cached_property
+    @mark_excitation_property()
+    @timed_member_call(timer="_property_timer")
+    def transition_dipole_moments_qed(self):
+        """
+        List of transition dipole moments of all computed states
+        to build the QED-matrix in the basis of the diagonal
+        purely electric subblock
+        """
+        if hasattr(self.reference_state, "approx"):
+
+            dipole_integrals = self.operators.electric_dipole
+            def tdm(i, prop_level):
+                self.ground_state.tdm_contribution = prop_level
+                return transition_dm(self.method, self.ground_state, 
+                                     self.excitation_vector[i])
+
+            if hasattr(self.reference_state, "first_order_coupling"):
+
+                return np.array([
+                    [product_trace(comp, tdm(i, "adc0")) for comp in dipole_integrals]
+                    for i in np.arange(len(self.excitation_energy))
+                ])
+            else:
+                prop_level = "adc" + str(self.property_method.level - 1)
+                return np.array([
+                    [product_trace(comp, tdm(i, prop_level)) for comp in dipole_integrals]
+                    for i in np.arange(len(self.excitation_energy))
+                ])
+        else:
+            return ("transition_dipole_moments_qed are only calculated,"
+                    "if reference_state contains 'approx' attribute")
+
+    @cached_property
+    @mark_excitation_property()
+    @timed_member_call(timer="_property_timer")
+    def s2s_dipole_moments_qed(self):
+        """
+        List of s2s transition dipole moments of all computed states
+        to build the QED-matrix in the basis of the diagonal
+        purely electric subblock
+        """
+        if hasattr(self.reference_state, "approx"):
+            dipole_integrals = self.operators.electric_dipole
+            print("note, that only the z coordinate of the dipole integrals is calculated")
+            n_states = len(self.excitation_energy)
+
+            def s2s(i, f, s2s_contribution):
+                self.ground_state.s2s_contribution = s2s_contribution
+                vec = self.excitation_vector
+                return state2state_transition_dm(self.method, self.ground_state, 
+                                                 vec[i], vec[f])
+
+            def final_block(name):
+                return np.array([[product_trace(dipole_integrals[2], s2s(i, j, name)) 
+                        for j in np.arange(n_states)] for i in np.arange(n_states)])
+
+            block_dict = {}
+
+            block_dict["qed_adc1_off_diag"] = final_block("adc1")
+
+            if self.method.name == "adc2" and not hasattr(self.reference_state, "first_order_coupling"): # noqa: E501
+
+                block_dict["qed_adc2_diag"] = final_block("qed_adc2_diag")     
+                block_dict["qed_adc2_edge_couple"] = final_block("qed_adc2_edge_couple")
+                block_dict["qed_adc2_edge_phot_couple"] = final_block("qed_adc2_edge_phot_couple") # noqa: E501
+                block_dict["qed_adc2_ph_pphh"] = final_block("qed_adc2_ph_pphh")
+                block_dict["qed_adc2_pphh_ph"] = final_block("qed_adc2_pphh_ph")
+            return block_dict
+        else:
+            return ("s2s_dipole_moments_qed are only calculated,"
+                    "if reference_state contains 'approx' attribute")
+
+    @cached_property
+    @mark_excitation_property()
+    @timed_member_call(timer="_property_timer")
+    def qed_second_order_ph_ph_couplings(self):
+        """
+        List of blocks containing the expectation value of the perturbation
+        of the Hamiltonian for all computed states required
+        to build the QED-matrix in the basis of the diagonal
+        purely electric subblock
+        """
+        if hasattr(self.reference_state, "approx"):
+            qed_t1 = self.ground_state.qed_t1(b.ov)
+
+            def couple(qed_t1, ul, ur):
+                return {
+                    b.ooov: einsum("kc,ia,ja->kjic", qed_t1, ul, ur) + \
+                            einsum("ka,ia,jb->jkib", qed_t1, ul, ur),
+                    b.ovvv: einsum("kc,ia,ib->kacb", qed_t1, ul, ur) + \
+                            einsum("ic,ia,jb->jabc", qed_t1, ul, ur) 
+                }
+
+            def phot_couple(qed_t1, ul, ur):
+                return {
+                    b.ooov: einsum("kc,ia,ja->kijc", qed_t1, ul, ur) + \
+                            einsum("kb,ia,jb->ikja", qed_t1, ul, ur),
+                    b.ovvv: einsum("kc,ia,ib->kbca", qed_t1, ul, ur) + \
+                            einsum("jc,ia,jb->ibac", qed_t1, ul, ur) 
+                }
+
+            def prod_sum(hf, two_p_op):
+                return + (einsum("ijka,ijka->", hf.ooov, two_p_op[b.ooov]) 
+                                + einsum("iabc,iabc->", hf.ovvv, two_p_op[b.ovvv]))
+
+            def final_block(func):
+                return np.array([[prod_sum(self.reference_state, func(qed_t1, i.ph, j.ph)) 
+                                for i in self.excitation_vector] for j in self.excitation_vector])
+
+            block_dict = {}
+            block_dict["couple"] = final_block(couple)
+            block_dict["phot_couple"] = final_block(phot_couple)
+
+            return block_dict
+        else:
+            return ("qed_second_order_ph_ph_couplings are only calculated,"
+                    "if reference_state contains 'approx' attribute")
+
     @cached_property
     @mark_excitation_property()
     @timed_member_call(timer="_property_timer")

adcc/ExcitedStates.py

@@ -35,6 +35,7 @@
 from .OneParticleOperator import product_trace
 from .ElectronicTransition import ElectronicTransition
 from .FormatDominantElements import FormatDominantElements
+from .AmplitudeVector import QED_AmplitudeVector
 
 
 class EnergyCorrection:
@@ -410,12 +411,18 @@ def describe_amplitudes(self, tolerance=0.01, index_format=None):
             to index relative on the HOMO / LUMO / HOCO orbitals.
             If ``None`` an automatic selection will be made.
         """
-        eV = constants.value("Hartree energy in eV")
+        eV = constants.value("Hartree energy in eV") 
         vector_format = FormatExcitationVector(self.matrix, tolerance=tolerance,
                                                index_format=index_format)
 
         # Optimise the formatting by pre-inspecting all tensors
         for tensor in self.excitation_vector:
+            if isinstance(tensor, QED_AmplitudeVector):
+                # TODO: Implement tdm and s2s_tdm for QED, so properties can also
+                # be evaluated for QED_AmplitudeVector objects. For now only
+                # use AmplitudeVector describing the electric part, so the
+                # formatting does not have to be adapted.
+                tensor = tensor.elec
             vector_format.optimise_formatting(tensor)
 
         # Determine width of a line
@@ -424,6 +431,13 @@ def describe_amplitudes(self, tolerance=0.01, index_format=None):
 
         ret = separator
         for i, vec in enumerate(self.excitation_vector):
+            if isinstance(vec, QED_AmplitudeVector):
+                # TODO: Implement tdm and s2s_tdm for QED, so properties can also
+                # be evaluated for QED_AmplitudeVector objects. For now only
+                # use AmplitudeVector describing the electric part, since 
+                # most low-energy states are almost purely electric, so the
+                # formatting does not have to be adapted.
+                vec = vec.elec
             ene = self.excitation_energy[i]
             eev = ene * eV
             head = f"State {i:3d} , {ene:13.7g} au"