Merge ef3da3e into 277162e

adcc/AdcMatrix.py

@@ -32,6 +32,38 @@
 from .AmplitudeVector import AmplitudeVector
 
 
+class AdcExtraTerm:
+    def __init__(self, matrix, blocks):
+        """Initialise an AdcExtraTerm.
+        This class can be used to add customs terms
+        to an existing :py:class:`AdcMatrix`
+
+        Parameters
+        ----------
+        matrix : AdcMatrix
+            The matrix for which the extra term
+            should be created.
+        blocks : dict
+            A dictionary where the key labels the matrix block
+            and the item denotes a callable to construct
+            an :py:class:`AdcBlock`
+        """
+        self.ground_state = matrix.ground_state
+        self.reference_state = matrix.reference_state
+        self.intermediates = matrix.intermediates
+        self.blocks = {}
+        if not isinstance(blocks, dict):
+            raise TypeError("blocks needs to be a dict.")
+        for space in blocks:
+            block_fun = blocks[space]
+            if not callable(block_fun):
+                raise TypeError("Items in additional_blocks must be callable.")
+            block = block_fun(
+                self.reference_state, self.ground_state, self.intermediates
+            )
+            self.blocks[space] = block
+
+
 class AdcMatrixlike:
     """
     Base class marker for all objects like ADC matrices.
@@ -50,7 +82,8 @@ class AdcMatrix(AdcMatrixlike):
         "adc3":  dict(ph_ph=3, ph_pphh=2,    pphh_ph=2,    pphh_pphh=1),     # noqa: E501
     }
 
-    def __init__(self, method, hf_or_mp, block_orders=None, intermediates=None):
+    def __init__(self, method, hf_or_mp, block_orders=None, intermediates=None,
+                 diagonal_precomputed=None):
         """
         Initialise an ADC matrix.
 
@@ -65,6 +98,8 @@ def __init__(self, method, hf_or_mp, block_orders=None, intermediates=None):
             If not set, defaults according to the selected ADC method are chosen.
         intermediates : adcc.Intermediates or NoneType
             Allows to pass intermediates to re-use to this class.
+        diagonal_precomputed: adcc.AmplitudeVector
+            Allows to pass a pre-computed diagonal, for internal use only.
         """
         if isinstance(hf_or_mp, (libadcc.ReferenceState,
                                  libadcc.HartreeFockSolution_i)):
@@ -84,6 +119,7 @@ def __init__(self, method, hf_or_mp, block_orders=None, intermediates=None):
         self.mospaces = hf_or_mp.reference_state.mospaces
         self.is_core_valence_separated = method.is_core_valence_separated
         self.ndim = 2
+        self.extra_terms = []
 
         self.intermediates = intermediates
         if self.intermediates is None:
@@ -112,19 +148,84 @@ def __init__(self, method, hf_or_mp, block_orders=None, intermediates=None):
         # Build the blocks and diagonals
         with self.timer.record("build"):
             variant = None
-            if method.is_core_valence_separated:
+            if self.is_core_valence_separated:
                 variant = "cvs"
-            self.blocks_ph = {  # TODO Rename to self.block in 0.16.0
+            blocks = {
                 block: ppmatrix.block(self.ground_state, block.split("_"),
                                       order=order, intermediates=self.intermediates,
                                       variant=variant)
-                for block, order in block_orders.items() if order is not None
+                for block, order in self.block_orders.items() if order is not None
             }
-            self.__diagonal = sum(bl.diagonal for bl in self.blocks_ph.values()
-                                  if bl.diagonal)
-            self.__diagonal.evaluate()
+            # TODO Rename to self.block in 0.16.0
+            self.blocks_ph = {bl: blocks[bl].apply for bl in blocks}
+            if diagonal_precomputed:
+                if not isinstance(diagonal_precomputed, AmplitudeVector):
+                    raise TypeError("diagonal_precomputed needs to be"
+                                    " an AmplitudeVector.")
+                if diagonal_precomputed.needs_evaluation:
+                    raise ValueError("diagonal_precomputed must already"
+                                     " be evaluated.")
+                self.__diagonal = diagonal_precomputed
+            else:
+                self.__diagonal = sum(bl.diagonal for bl in blocks.values()
+                                      if bl.diagonal)
+                self.__diagonal.evaluate()
             self.__init_space_data(self.__diagonal)
 
+    def __iadd__(self, other):
+        """In-place addition of an :py:class:`AdcExtraTerm`
+
+        Parameters
+        ----------
+        other : AdcExtraTerm
+            the extra term to be added
+        """
+        if not isinstance(other, AdcExtraTerm):
+            return NotImplemented
+        if not all(k in self.blocks_ph for k in other.blocks):
+            raise ValueError("Can only add to blocks of"
+                             " AdcMatrix that already exist.")
+        for sp in other.blocks:
+            orig_app = self.blocks_ph[sp]
+            other_app = other.blocks[sp].apply
+
+            def patched_apply(ampl, original=orig_app, other=other_app):
+                return sum(app(ampl) for app in (original, other))
+            self.blocks_ph[sp] = patched_apply
+        other_diagonal = sum(bl.diagonal for bl in other.blocks.values()
+                             if bl.diagonal)
+        # iadd does not work with numbers
+        self.__diagonal = self.__diagonal + other_diagonal
+        self.__diagonal.evaluate()
+        self.extra_terms.append(other)
+        return self
+
+    def __add__(self, other):
+        """Addition of an :py:class:`AdcExtraTerm`, creating
+        a copy of self and adding the term to the new matrix
+
+        Parameters
+        ----------
+        other : AdcExtraTerm
+            the extra term to be added
+
+        Returns
+        -------
+        AdcMatrix
+            a copy of the AdcMatrix with the extra term added
+        """
+        if not isinstance(other, AdcExtraTerm):
+            return NotImplemented
+        ret = AdcMatrix(self.method, self.ground_state,
+                        block_orders=self.block_orders,
+                        intermediates=self.intermediates,
+                        diagonal_precomputed=self.diagonal())
+        ret += other
+        return ret
+
+    def __radd__(self, other):
+        return self.__add__(other)
+
     def __init_space_data(self, diagonal):
         """Update the cached data regarding the spaces of the ADC matrix"""
         self.axis_spaces = {}
@@ -208,7 +309,7 @@ def block_apply(self, block, tensor):
         with self.timer.record(f"apply/{block}"):
             outblock, inblock = block.split("_")
             ampl = AmplitudeVector(**{inblock: tensor})
-            ret = self.blocks_ph[block].apply(ampl)
+            ret = self.blocks_ph[block](ampl)
             return getattr(ret, outblock)
 
     @timed_member_call()
@@ -217,7 +318,7 @@ def matvec(self, v):
         Compute the matrix-vector product of the ADC matrix
         with an excitation amplitude and return the result.
         """
-        return sum(block.apply(v) for block in self.blocks_ph.values())
+        return sum(block(v) for block in self.blocks_ph.values())
 
     def rmatvec(self, v):
         # ADC matrix is symmetric

adcc/AmplitudeVector.py

@@ -110,6 +110,10 @@ def evaluate(self):
             t.evaluate()
         return self
 
+    @property
+    def needs_evaluation(self):
+        return any(t.needs_evaluation for k, t in self.items())
+
     def ones_like(self):
         """Return an empty AmplitudeVector of the same shape and symmetry"""
         return AmplitudeVector(**{k: t.ones_like() for k, t in self.items()})

adcc/ExcitedStates.py

@@ -38,6 +38,32 @@
 from .FormatDominantElements import FormatDominantElements
 
 
+class EnergyCorrection:
+    def __init__(self, name, function):
+        """A helper class to represent excitation energy
+        corrections.
+
+        Parameters
+        ----------
+        name : str
+            descriptive name of the energy correction
+        function : callable
+            function that takes a :py:class:`Excitation`
+            as single argument and returns the energy
+            correction as a float
+        """
+        if not isinstance(name, str):
+            raise TypeError("name needs to be a string.")
+        if not callable(function):
+            raise TypeError("function needs to be callable.")
+        self.name = name
+        self.function = function
+
+    def __call__(self, excitation):
+        assert isinstance(excitation, Excitation)
+        return self.function(excitation)
+
+
 class FormatExcitationVector:
     def __init__(self, matrix, tolerance=0.01, index_format=None):
         """
@@ -127,8 +153,7 @@ def format(self, vector):
 
 
 class ExcitedStates(ElectronicTransition):
-    def __init__(self, data, method=None, property_method=None,
-                 excitation_energy_corrections={}):
+    def __init__(self, data, method=None, property_method=None):
         """Construct an ExcitedStates class from some data obtained
         from an interative solver or another :class:`ExcitedStates`
         object.
@@ -153,33 +178,45 @@ def __init__(self, data, method=None, property_method=None,
         property_method : str, optional
             Provide an explicit method for property calculations to
             override the automatic selection.
-        excitation_energy_corrections : dict, optional
-            Provide a dictionary of functions to compute corrections for
-            excitation energies, called for each excitation individually
         """
         super().__init__(data, method, property_method)
-
-        if hasattr(data, "excitation_energy_corrections"):
-            merged = data.excitation_energy_corrections.copy()
-            merged.update(excitation_energy_corrections)
-            excitation_energy_corrections = merged
-
-        if hasattr(self.reference_state, "excitation_energy_corrections"):
-            merged = self.reference_state.excitation_energy_corrections.copy()
-            merged.update(excitation_energy_corrections)
-            excitation_energy_corrections = merged
-
-        self.excitation_energy_corrections = excitation_energy_corrections
-        for key in self.excitation_energy_corrections:
-            corr_function = self.excitation_energy_corrections[key]
-            if not callable(corr_function):
-                raise TypeError("Elements in excitation_energy_corrections "
-                                "must be callable.")
-            # call the function for exc. energy correction
-            correction = np.array([corr_function(exci)
-                                   for exci in self.excitations])
-            setattr(self, key, correction)
-            self._excitation_energy += correction
+        self._excitation_energy_corrections = []
+        # copy energy corrections if possible
+        # and avoids re-computation of the corrections
+        if hasattr(data, "_excitation_energy_corrections"):
+            for eec in data._excitation_energy_corrections:
+                correction_energy = getattr(data, eec.name)
+                setattr(self, eec.name, correction_energy)
+                self._excitation_energy += correction_energy
+                self._excitation_energy_corrections.append(eec)
+
+    def __add_energy_correction(self, correction):
+        assert isinstance(correction, EnergyCorrection)
+        if hasattr(self, correction.name):
+            raise ValueError("ExcitedStates already has a correction"
+                             f" with the name '{correction.name}'")
+        correction_energy = np.array([correction(exci)
+                                      for exci in self.excitations])
+        setattr(self, correction.name, correction_energy)
+        self._excitation_energy += correction_energy
+        self._excitation_energy_corrections.append(correction)
+
+    def __iadd__(self, other):
+        if isinstance(other, EnergyCorrection):
+            self.__add_energy_correction(other)
+        elif isinstance(other, list):
+            for k in other:
+                self += k
+        else:
+            raise TypeError("Can only add EnergyCorrection (or list"
+                            " of EnergyCorrection) to"
+                            f" ExcitedState, not '{type(other)}'")
+        return self
+
+    def __add__(self, other):
+        ret = ExcitedStates(self, self.method, self.property_method)
+        ret += other
+        return ret
 
     @cached_property
     @mark_excitation_property(transform_to_ao=True)
@@ -338,6 +375,16 @@ def describe(self, oscillator_strengths=True, rotatory_strengths=False,
             text += body.format(i=i, ene=self.excitation_energy[i],
                                 ev=self.excitation_energy[i] * eV, **fields)
         text += separator + "\n"
+        if len(self._excitation_energy_corrections):
+            head_corr = "|  excitation energy corrections included:"
+            text += head_corr
+            nspace = len(separator) - len(head_corr) - 1
+            text += nspace * " " + "|\n"
+            maxlen = len(separator) - 8
+            for eec in self._excitation_energy_corrections:
+                label = f"|    - {eec.name:{maxlen}.{maxlen}}|\n"
+                text += label
+            text += separator + "\n"
         return text
 
     def _repr_pretty_(self, pp, cycle):
@@ -398,7 +445,7 @@ def to_dataframe(self):
         Atomic units are used for all values.
         """
         propkeys = self.excitation_property_keys
-        propkeys.extend(self.excitation_energy_corrections.keys())
+        propkeys.extend([k.name for k in self._excitation_energy_corrections])
         data = {
             "excitation": np.arange(0, self.size, dtype=int),
             "kind": np.tile(self.kind, self.size)