Projectors and projected ADC matrix (#149)

adcc/AdcMatrix.py

@@ -28,6 +28,7 @@
 from .adc_pp import matrix as ppmatrix
 from .timings import Timer, timed_member_call
 from .AdcMethod import AdcMethod
+from .functions import ones_like
 from .Intermediates import Intermediates
 from .AmplitudeVector import AmplitudeVector
 
@@ -355,7 +356,7 @@ def block_view(self, block):
             #      data structure as the AdcMatrix class as these block
             #      are inherently different than an AdcMatrix (no Hermiticity
             #      for example) and basically they only need to support some
-            #      form of matrix-vector product and some stastics like
+            #      form of matrix-vector product and some statistics like
             #      spaces and sizes etc.
         block_orders = {bl: None for bl in self.block_orders.keys()}
         block_orders[block] = self.block_orders[block]
@@ -628,3 +629,90 @@ def block_view(self, block):
                                   "shifted ADC matrices.")
         # TODO The way to implement this is to ask the inner matrix to
         #      a block_view and then wrap that in an AdcMatrixShifted.
+
+
+class AdcMatrixProjected(AdcMatrix):
+    def __init__(self, matrix, excitation_blocks, core_orbitals=None,
+                 outer_virtuals=None):
+        """
+        Initialise a projected ADC matrix, i.e. represents the expression
+        ``P @ M @ P`` where ``P`` is a projector onto a subset of
+        ``excitation_blocks``.
+
+        The ``excitation_blocks`` are defined by partitioning the ``o1`` occupied
+        and ``v1`` virtual space of the ``matrix.mospaces`` into a core-occupied
+        ``c``, valence-occupied ``o``, inner-virtual ``v`` and outer-virtual ``w``.
+        This matrix will only keep selected blocks in the amplitudes non-zero, which
+        are selected in the ``excitation_blocks`` list
+        (e.g. ``["cv", "ccvv", "ocvv"]``).
+
+        For details on the option how to select the spaces, see the documentation
+        in :py:`adcc.ReferenceState.__init__` (``outer_virtuals`` follows the same
+        rules as ``frozen_virtuals``).
+
+        Parameters
+        ----------
+        matrix : AdcMatrix
+            Matrix which is projected
+        excitation_blocks : list
+            Excitation blocks to keep in the Amplitudes.
+        core_orbitals : int or list or tuple, optional
+            The orbitals to be put into the ``c`` core space.
+        outer_virtuals : int or list or tuple, optional
+            The virtuals to be put into the ``w`` outer-virtual space.
+        """
+        from .projection import Projector, SubspacePartitioning
+
+        for sp in excitation_blocks:
+            if not any(len(sp) == len(ax_sp)
+                       for ax_sp in matrix.axis_spaces.values()):
+                raise ValueError(f"Invalid partition block {sp}.")
+
+        super().__init__(matrix.method, matrix.ground_state,
+                         block_orders=matrix.block_orders,
+                         intermediates=matrix.intermediates)
+        partitioning = SubspacePartitioning(matrix.mospaces, core_orbitals,
+                                            outer_virtuals)
+
+        projectors = {}
+        for block in matrix.axis_spaces.keys():
+            block_partitions = [sp for sp in excitation_blocks
+                                if len(sp) == len(matrix.axis_spaces[block])]
+            projectors[block] = Projector(matrix.axis_spaces[block],
+                                          partitioning, block_partitions)
+        self.projectors = projectors
+
+    def apply_projection(self, in_ampl):
+        return AmplitudeVector(**{
+            block: self.projectors[block] @ in_ampl[block]
+            for block in in_ampl.keys()
+        })
+
+    def matvec(self, in_ampl):
+        in_proj = self.apply_projection(in_ampl)
+        out = super().matvec(in_proj)
+        return self.apply_projection(out)
+
+    def block_apply(self, block, in_vec):
+        inblock, outblock = block.split("_")
+        in_proj = self.projectors[inblock].apply(in_vec)
+        ret = super().block_apply(block, in_proj)
+        return self.projectors[outblock].apply(ret)
+
+    def diagonal(self):
+        blocks = {}
+        for (block, diagblock) in super().diagonal().items():
+            # On the diagonal don't set the ignored amplitudes to zero,
+            # but instead set them to a very large value to (a) avoid these being
+            # selected for the initial guess and (b) ensure the preconditioning
+            # naturally reduces components along this direction.
+            P = self.projectors[block]
+            one = ones_like(diagblock)
+            blocks[block] = P @ diagblock + 100000 * (one - P @ one)
+        return AmplitudeVector(**blocks)
+
+    def block_view(self, block):
+        raise NotImplementedError("Block-view not yet implemented for "
+                                  "projected ADC matrices.")
+        # TODO The way to implement this is to ask the inner matrix to
+        #      a block_view and then wrap that in an AdcMatrixProjected.

adcc/MoSpaces.py

@@ -48,7 +48,14 @@ def split_spaces(spacestr):
         raise ValueError(f"Invalid space string '{spacestr}'.")
 
 
-def expand_spaceargs(hfdata, **spaceargs):
+def expand_spaceargs(hfdata_or_n_orbs, **spaceargs):
+    if isinstance(hfdata_or_n_orbs, tuple):
+        (noa, nob) = hfdata_or_n_orbs
+    else:
+        hfdata = hfdata_or_n_orbs
+        noa = hfdata.n_orbs_alpha
+        nob = hfdata.n_orbs_beta
+
     if isinstance(spaceargs.get("frozen_core", None), bool) \
        and spaceargs.get("frozen_core", None):
         # Determine number of frozen core electrons automatically
@@ -93,7 +100,6 @@ def expand_to_list(space, entry, from_min=None, from_max=None):
                                      "core_orbitals, frozen_virtual is an "
                                      "iterable, all must be.")
 
-    noa = hfdata.n_orbs_alpha
     n_orbs = [0, 0]
     for key in ["frozen_core", "core_orbitals"]:
         if key not in spaceargs or not spaceargs[key]:
@@ -110,7 +116,7 @@ def expand_to_list(space, entry, from_min=None, from_max=None):
     if key in spaceargs and spaceargs[key]:
         spaceargs[key] = np.concatenate((
             expand_to_list(key, spaceargs[key][0], from_max=noa),
-            expand_to_list(key, spaceargs[key][1], from_max=noa) + noa
+            expand_to_list(key, spaceargs[key][1], from_max=nob) + noa
         )).tolist()
     return spaceargs
 

adcc/guess/__init__.py

@@ -28,6 +28,23 @@
            "guesses_spin_flip", "guess_symmetries"]
 
 
+def guess_kwargs_kind(kind):
+    """
+    Return the kwargs required to be passed to `guesses_from_diagonal` to
+    computed states of the passed excitation `kind`.
+    """
+    kwargsmap = dict(
+        singlet=dict(spin_block_symmetrisation="symmetric", spin_change=0),
+        triplet=dict(spin_block_symmetrisation="antisymmetric", spin_change=0),
+        spin_flip=dict(spin_block_symmetrisation="none", spin_change=-1),
+        any=dict(spin_block_symmetrisation="none", spin_change=0),
+    )
+    try:
+        return kwargsmap[kind]
+    except KeyError:
+        raise ValueError(f"Kind not known: {kind}")
+
+
 def guesses_singlet(matrix, n_guesses, block="ph", **kwargs):
     """
     Obtain guesses for computing singlet states by inspecting the passed
@@ -41,8 +58,7 @@ def guesses_singlet(matrix, n_guesses, block="ph", **kwargs):
     kwargs       Any other argument understood by guesses_from_diagonal.
     """
     return guesses_from_diagonal(matrix, n_guesses, block=block,
-                                 spin_block_symmetrisation="symmetric",
-                                 spin_change=0, **kwargs)
+                                 **guess_kwargs_kind("singlet"), **kwargs)
 
 
 def guesses_triplet(matrix, n_guesses, block="ph", **kwargs):
@@ -58,8 +74,7 @@ def guesses_triplet(matrix, n_guesses, block="ph", **kwargs):
     kwargs       Any other argument understood by guesses_from_diagonal.
     """
     return guesses_from_diagonal(matrix, n_guesses, block=block,
-                                 spin_block_symmetrisation="antisymmetric",
-                                 spin_change=0, **kwargs)
+                                 **guess_kwargs_kind("triplet"), **kwargs)
 
 
 # guesses for computing any state (singlet or triplet)
@@ -79,4 +94,4 @@ def guesses_spin_flip(matrix, n_guesses, block="ph", **kwargs):
     kwargs       Any other argument understood by guesses_from_diagonal.
     """
     return guesses_from_diagonal(matrix, n_guesses, block=block,
-                                 spin_change=-1, **kwargs)
+                                 **guess_kwargs_kind("spin_flip"), **kwargs)

adcc/guess/guesses_from_diagonal.py

@@ -33,7 +33,7 @@
 
 def guesses_from_diagonal(matrix, n_guesses, block="ph", spin_change=0,
                           spin_block_symmetrisation="none",
-                          degeneracy_tolerance=1e-14):
+                          degeneracy_tolerance=1e-14, max_diagonal_value=1000):
     """
     Obtain guesses by inspecting a block of the diagonal of the passed ADC
     matrix. The symmetry of the returned vectors is already set-up properly.
@@ -59,6 +59,9 @@ def guesses_from_diagonal(matrix, n_guesses, block="ph", spin_change=0,
     degeneracy_tolerance
                  Tolerance for two entries of the diagonal to be considered
                  degenerate, i.e. identical.
+    max_diagonal_value
+                 Maximal diagonal value, which is considered as a valid candidate
+                 to form a guess.
     """
     if not isinstance(matrix, AdcMatrixlike):
         raise TypeError("matrix needs to be of type AdcMatrixlike")
@@ -88,7 +91,8 @@ def guesses_from_diagonal(matrix, n_guesses, block="ph", spin_change=0,
         raise ValueError(f"Don't know how to generate guesses for block {block}")
 
     return guessfunction(matrix, n_guesses, spin_change,
-                         spin_block_symmetrisation, degeneracy_tolerance)
+                         spin_block_symmetrisation, degeneracy_tolerance,
+                         max_diagonal_value)
 
 
 class TensorElement:
@@ -195,7 +199,8 @@ def telem_nospin(telem):
 
 def guesses_from_diagonal_singles(matrix, n_guesses, spin_change=0,
                                   spin_block_symmetrisation="none",
-                                  degeneracy_tolerance=1e-14):
+                                  degeneracy_tolerance=1e-14,
+                                  max_diagonal_value=1000):
     motrans = MoIndexTranslation(matrix.mospaces, matrix.axis_spaces["ph"])
     if n_guesses == 0:
         return []
@@ -207,10 +212,13 @@ def guesses_from_diagonal_singles(matrix, n_guesses, spin_change=0,
 
     # Search of the smallest elements
     # This predicate checks an index is an allowed element for the singles
-    # part of the guess vectors and has the requested spin-change
+    # part of the guess vectors and has the requested spin-change.
+    # Also it filters out too large diagonal entries (which are essentially
+    # hopeless to give useful excitations)
     def pred_singles(telem):
         return (ret[0].ph.is_allowed(telem.index)
-                and telem.spin_change == spin_change)
+                and telem.spin_change == spin_change
+                and abs(telem.value) <= max_diagonal_value)
 
     elements = find_smallest_matching_elements(
         pred_singles, matrix.diagonal().ph, motrans, n_guesses,
@@ -265,7 +273,8 @@ def telem_nospin(telem):
 
 def guesses_from_diagonal_doubles(matrix, n_guesses, spin_change=0,
                                   spin_block_symmetrisation="none",
-                                  degeneracy_tolerance=1e-14):
+                                  degeneracy_tolerance=1e-14,
+                                  max_diagonal_value=1000):
     if n_guesses == 0:
         return []
 
@@ -286,9 +295,14 @@ def guesses_from_diagonal_doubles(matrix, n_guesses, spin_change=0,
         guesses_d, matrix.mospaces, df02, df13,
         spin_change_twice, degeneracy_tolerance
     )
-
     # Resize in case less guesses found than requested
-    return ret[:n_found]
+    ret = ret[:n_found]
+
+    # Filter out elements above the noted diagonal value
+    diagonal_elements = [ret_d.pphh.dot(matrix.diagonal().pphh * ret_d.pphh)
+                         for ret_d in ret]
+    return [ret[i] for (i, elem) in enumerate(diagonal_elements)
+            if elem <= max_diagonal_value]
 
 # TODO Generic algorithm for forming orthogonal spin components of arbitrary
 #      size. Could be useful for the doubles guesses later on.