16384966787

Committed 03 Jul 2025 11:29AM UTC coverage: 73.883% (+0.3%) from 73.602%

Build # 16384966787

Build Type

push

github

Committed by

web-flow

Commit Message

Restructure ExcitedStates etc for IP/EA Integration (#195)

* allow timer customization in cached_member_function

* add option to exclude the args in the timer task

* timer description if nothing was recorded

* basic restructuring of ElectronicTransition, ExcitedStates and State2States

* basic restructuring of Excitation

* add _module member variable to dispatch to the corresponding working equations

* basic restructuring for excited state properties

* enable excitation for s2s

* port dataframe export

* add state_dm

* port the remaining properties - tests passing

* * cache the MO integrals instead of the AO integrals
* remove unnecessary property decorated returns of callbacks
* implement available integrals directly in the backends

* raise Exception directly in the backends for PE and PCM

* add to_qcvars back in

* split plot_spectrum in common function on base class and adc type dependent function on child classes

* invert order of unit conversion and broadening for plot_spectrum

* add input options for lower and upper bounds of the broadened spectrum

* remove broadening min/max and allow variable width_units

* refactor the ExcitedStates.describe method

* determine the column width automatically

* move describe_amplitudes to ElectronicStates and adapt for IP/EA

* store operators in tuple on IsrMatrix

* make flake happy and remove the cache for now since it is not used anyway

* reintroduce the cache for available backends and cache them lazily

* explicitly install setupstools

* enable libtensorlight download for macos arm64

* naively update CI to macos-latest

* install llvm and use arm64 compilers

* explicitly only build for the current platform

* only set architecture for macos

* generic block_orders for secular and isr matrix

* move init complete back to secular and isr matrix

* raise more explicit exception for not implemented methods

* patch coverage not required for status check

* add token for github api request... (continued)

Run Details

1176 of 1884 branches covered (62.42%)

Branch coverage included in aggregate %.

700 of 965 new or added lines in 21 files covered. (72.54%)

6 existing lines in 6 files now uncovered.

7376 of 9691 relevant lines covered (76.11%)

175662.79 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

85.5

/adcc/AdcMatrix.py

#!/usr/bin/env python3
## vi: tabstop=4 shiftwidth=4 softtabstop=4 expandtab
## ---------------------------------------------------------------------
##
## Copyright (C) 2018 by the adcc authors
##
## This file is part of adcc.
##
## adcc is free software: you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published
## by the Free Software Foundation, either version 3 of the License, or
## (at your option) any later version.
##
## adcc is distributed in the hope that it will be useful,
## but WITHOUT ANY WARRANTY; without even the implied warranty of
## MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
## GNU General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with adcc. If not, see <http://www.gnu.org/licenses/>.
##
## ---------------------------------------------------------------------
import itertools
import numpy as np

import libadcc

from .LazyMp import LazyMp
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


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.
    """

    _special_block_orders = {
        "adc2x": {"ph_ph": 2, "ph_pphh": 1, "pphh_ph": 1, "pphh_pphh": 1},
    }

    @classmethod
    def _default_block_orders(cls, method: AdcMethod) -> dict[str, int]:
        """
        Determines the default block orders for the given adc method.
        """
        # check if we have a special method like adc2x
        # I guess base_method should also contain the adc_type prefix so
        # we don't need to separate different adc_types
        block_orders = cls._special_block_orders.get(method.base_method.name, None)
        if block_orders is not None:
            return block_orders.copy()
        # otherwise assume that we have a "normal" PP/IP/...-ADC(n) method
        # - determine which spaces are available in the ADC(n) matrix
        #   starting from the given minimal space
        min_space = {
            "pp": "ph"
        }.get(method.adc_type, None)
        if min_space is None:
            raise ValueError(f"Unknown adc type {method.adc_type} for method "
                             f"{method.name}. Can not determine default block "
                             "orders.")
        spaces = [
            "p" * i + min_space + "h" * i for i in range(0, (method.level // 2) + 1)
        ]
        # exploit the fact that the spaces are sorted from small to high:
        # If we walk the adc matrix in any direction we always have to subtract 1!
        # Therefore, we can determine the order according to the position of the
        # spaces: maxorder - 0 for singles, maxorder - 2 for doubles and
        # maxorder - 3 for the doubles/triples coupling
        ret = {}
        for ((i1, bra), (i2, ket)) in \
                itertools.product(enumerate(spaces), repeat=2):
            order = method.level - i1 - i2
            assert order >= 0
            ret[f"{bra}_{ket}"] = order
        return ret

    @classmethod
    def _validate_block_orders(cls, block_orders: dict[str, int],
                               method: AdcMethod,
                               allow_missing_diagonal_blocks: bool = False) -> None:
        """
        Validates that the given block_orders form a valid adc matrix for the given
        adc method.

        Parameters
        ----------
        block_orders: dict[str, int]
            The block orders to validate. Block orders should be of the form
            {'ph_ph': 2, 'ph_pphh': 1, ...}
        method: AdcMethod
            The adc method/adc type (PP-ADC, ...) for which to validate
            the block_orders.
        allow_missing_diagonal_blocks: bool, optional
            If set, couplings between missing diagonal blocks are allowed, e.g.,
            {'ph_ph': 1, 'ph_pphh': 0, 'pphh_ph': 0}
            will be valid although there is a coupling between the 'ph_ph' block
            and the missing 'pphh_pphh' block. (default: False)
        """
        for block, order in block_orders.items():
            if order is None:
                continue
            # ensure that the block is valid for the given adc type
            bra, ket = block.split("_")
            if not cls._is_valid_space(bra, method) or \
                    not cls._is_valid_space(ket, method):
                raise ValueError(f"Invalid block {block} for a "
                                 f"{method.adc_type} ADC matrix.")
            if bra == ket:  # done for diagonal blocks
                continue
            # ensure that the matrix is symmetric
            inv_block = f"{ket}_{bra}"
            inv_order = block_orders.get(inv_block, None)
            if inv_order is None or inv_order != order:
                raise ValueError(f"{block} and {inv_block} should always have "
                                 "the same order.")
            if allow_missing_diagonal_blocks:
                continue
            # ensure that we have no coupling between missing diagonal blocks
            bra_diag = f"{bra}_{bra}"
            ket_diag = f"{ket}_{ket}"
            bra_diag_order = block_orders.get(bra_diag, None)
            ket_diag_order = block_orders.get(ket_diag, None)
            if bra_diag_order is None or ket_diag_order is None:
                raise ValueError(f"Can only have a couling block {block} "
                                 f"if both diagonal blocks {bra_diag} and "
                                 f"{ket_diag} are in the matrix too.")

    @classmethod
    def _is_valid_space(cls, space: str, method: AdcMethod) -> bool:
        """
        Checks whether the given space ('ph' for instance) is valid for the given
        adc method. Thereby we only verify that the space matches the adc_type of
        adc method!
        """
        n_particle, n_hole = space.count("p"), space.count("h")
        # ensure that the space is of the form pp...hh...
        if ("p" * n_particle + "h" * n_hole) != space:
            return False
        # depending on the adc type n_particle and n_hole have to
        # be equal or differ e.g. by +-1 (IP/EA)
        if method.adc_type == "pp":
            return n_particle == n_hole
        raise ValueError(f"Unknown adc type {method.adc_type} for method "
                         f"{method.name}. Can not validate space.")


class AdcMatrix(AdcMatrixlike):

    def __init__(self, method, hf_or_mp, block_orders=None, intermediates=None,
                 diagonal_precomputed=None):
        """
        Initialise an ADC matrix.

        Parameters
        ----------
        method : str or AdcMethod
            Method to use.
        hf_or_mp : adcc.ReferenceState or adcc.LazyMp
            HF reference or MP ground state
        block_orders : optional
            The order of perturbation theory to employ for each matrix block.
            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)):
            hf_or_mp = LazyMp(hf_or_mp)
        if not isinstance(hf_or_mp, LazyMp):
            raise TypeError("hf_or_mp is not a valid object. It needs to be "
                            "either a LazyMp, a ReferenceState or a "
                            "HartreeFockSolution_i.")

        if not isinstance(method, AdcMethod):
            method = AdcMethod(method)

        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.timer = Timer()
        self.method = method
        self.ground_state = hf_or_mp
        self.reference_state = hf_or_mp.reference_state
        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:
            self.intermediates = Intermediates(self.ground_state)

        self.block_orders = self._default_block_orders(self.method)
        if block_orders is None:
            if method.level > 3:
                raise NotImplementedError("The ADC secular matrix is not "
                                          f"implemented for method {method.name}.")
        else:
            self.block_orders.update(block_orders)
        self._validate_block_orders(
            block_orders=self.block_orders, method=self.method,
            allow_missing_diagonal_blocks=False
        )

        # Build the blocks and diagonals
        with self.timer.record("build"):
            variant = None
            if self.is_core_valence_separated:
                variant = "cvs"
            blocks = {
                block: ppmatrix.block(self.ground_state, block.split("_"),
                                      order=order, intermediates=self.intermediates,
                                      variant=variant)
                for block, order in self.block_orders.items() if order is not None
            }
            self.blocks = {bl: blocks[bl].apply for bl in blocks}
            if diagonal_precomputed:
                self._diagonal: AmplitudeVector = diagonal_precomputed
            else:
                self._diagonal: AmplitudeVector = 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 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[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[sp] = patched_apply
        other_diagonal = sum(bl.diagonal for bl in other.blocks.values()
                             if bl.diagonal)
        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 = {}
        self.axis_lengths = {}
        for block in diagonal.blocks:
            self.axis_spaces[block] = getattr(diagonal, block).subspaces
            self.axis_lengths[block] = np.prod([
                self.mospaces.n_orbs(sp) for sp in self.axis_spaces[block]
            ])
        self.shape = (sum(self.axis_lengths.values()),
                      sum(self.axis_lengths.values()))

    def __repr__(self):
        ret = f"AdcMatrix({self.method.name}, "
        for b, o in self.block_orders.items():
            ret += f"{b}={o}, "
        return ret + ")"

    def __len__(self):
        return self.shape[0]

    @property
    def axis_blocks(self):
        """
        Return the blocks used along one of the axes of the ADC matrix
        (e.g. ['ph', 'pphh']).
        """
        # sort the keys by length to ensure that we always get
        # [singles, doubles, ...]
        return sorted(self.axis_spaces, key=len)

    def diagonal(self):
        """Return the diagonal of the ADC matrix"""
        return self._diagonal

    def block_apply(self, block, tensor):
        """
        Compute the application of a block of the ADC matrix
        with another AmplitudeVector or Tensor. Non-matching blocks
        in the AmplitudeVector will be ignored.
        """
        if not isinstance(tensor, libadcc.Tensor):
            raise TypeError("tensor should be an adcc.Tensor")

        with self.timer.record(f"apply/{block}"):
            outblock, inblock = block.split("_")
            ampl = AmplitudeVector(**{inblock: tensor})
            ret = self.blocks[block](ampl)
            return getattr(ret, outblock)

    @timed_member_call()
    def matvec(self, v):
        """
        Compute the matrix-vector product of the ADC matrix
        with an excitation amplitude and return the result.
        """
        return sum(block(v) for block in self.blocks.values())

    def rmatvec(self, v):
        # ADC matrix is symmetric
        return self.matvec(v)

    def __matmul__(self, other):
        if isinstance(other, AmplitudeVector):
            return self.matvec(other)
        if isinstance(other, list):
            if all(isinstance(elem, AmplitudeVector) for elem in other):
                return [self.matvec(ov) for ov in other]
        return NotImplemented

    def block_view(self, block):
        """
        Return a view into the AdcMatrix that represents a single
        block of the matrix. Currently only diagonal blocks are supported.
        """
        b1, b2 = block.split("_")
        if b1 != b2:
            raise NotImplementedError("Off-diagonal block views not yet "
                                      "implemented.")
            # TODO For off-diagonal blocks we probably need a different
            #      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 statistics like
            #      spaces and sizes etc.
        block_orders = {bl: None for bl in self.block_orders.keys()}
        block_orders[block] = self.block_orders[block]
        return AdcMatrix(self.method, self.ground_state,
                         block_orders=block_orders,
                         intermediates=self.intermediates)

    def construct_symmetrisation_for_blocks(self):
        """
        Construct the symmetrisation functions, which need to be
        applied to relevant blocks of an AmplitudeVector in order
        to symmetrise it to the right symmetry in order to be used
        with the various matrix-vector-products of this function.

        Most importantly the returned functions antisymmetrise
        the occupied and virtual parts of the doubles parts
        if this is sensible for the method behind this adcmatrix.

        Returns a dictionary block identifier -> function
        """
        ret = {}
        if self.is_core_valence_separated:
            # CVS doubles part is antisymmetric wrt. (i,K,a,b) <-> (i,K,b,a)
            ret["pphh"] = lambda v: v.antisymmetrise([(2, 3)])
        else:
            def symmetrise_generic_adc_doubles(invec):
                # doubles part is antisymmetric wrt. (i,j,a,b) <-> (i,j,b,a)
                scratch = invec.antisymmetrise([(2, 3)])
                # doubles part is symmetric wrt. (i,j,a,b) <-> (j,i,b,a)
                return scratch.symmetrise([(0, 1), (2, 3)])
            ret["pphh"] = symmetrise_generic_adc_doubles
        return ret

    def dense_basis(self, axis_blocks=None, ordering="adcc"):
        """
        Return the list of indices and their values
        of the dense basis representation

        ordering: adcc, spin, spatial
        """
        ret = []
        if axis_blocks is None:
            axis_blocks = self.axis_blocks
        if not isinstance(axis_blocks, list):
            axis_blocks = [axis_blocks]

        # Define function to impose the order in the basis
        if ordering == "adcc":
            def reduce_index(n_orbsa, idx):
                return idx, idx
        elif ordering == "spin":
            def reduce_index(n_orbsa, idx):
                is_beta = [idx[i] >= n_orbsa[i] for i in range(len(idx))]
                spatial = [idx[i] - n_orbsa[i] if is_beta[i] else idx[i]
                           for i in range(len(idx))]
                # Sort first by spin, then by spatial
                return (is_beta, spatial)
        elif ordering == "spatial":
            def reduce_index(n_orbsa, idx):
                is_beta = [idx[i] >= n_orbsa[i] for i in range(len(idx))]
                spatial = [idx[i] - n_orbsa[i] if is_beta[i] else idx[i]
                           for i in range(len(idx))]
                # Sort first by spatial, then by spin
                return (spatial, is_beta)

        if "ph" in axis_blocks:
            ret_s = []
            sp_s = self.axis_spaces["ph"]
            n_orbs_s = [self.mospaces.n_orbs(sp) for sp in sp_s]
            n_orbsa_s = [self.mospaces.n_orbs_alpha(sp) for sp in sp_s]
            for i in range(n_orbs_s[0]):
                for a in range(n_orbs_s[1]):
                    ret_s.append([((i, a), 1)])

            def sortfctn(x):
                return min(reduce_index(n_orbsa_s, idx) for idx, factor in x)
            ret_s.sort(key=sortfctn)
            ret_s.sort(key=sortfctn)
            ret.extend(ret_s)

        if "pphh" in axis_blocks:
            ret_d = []
            sp_d = self.axis_spaces["pphh"]
            n_orbsa_d = [self.mospaces.n_orbs_alpha(sp) for sp in sp_d]

            if sp_d[0] == sp_d[1] and sp_d[2] == sp_d[3]:
                nso = self.mospaces.n_orbs(sp_d[0])
                nsv = self.mospaces.n_orbs(sp_d[2])
                ret_d.extend([[((i, j, a, b), +1 / 2),
                               ((j, i, a, b), -1 / 2),
                               ((i, j, b, a), -1 / 2),
                               ((j, i, b, a), +1 / 2)]
                              for i in range(nso) for j in range(i)
                              for a in range(nsv) for b in range(a)])
            elif sp_d[2] == sp_d[3]:
                nso = self.mospaces.n_orbs(sp_d[0])
                nsc = self.mospaces.n_orbs(sp_d[1])
                nsv = self.mospaces.n_orbs(sp_d[2])
                ret_d.extend([[((i, j, a, b), +1 / np.sqrt(2)),
                               ((i, j, b, a), -1 / np.sqrt(2))]
                              for i in range(nso) for j in range(nsc)
                              for a in range(nsv) for b in range(a)])
            else:
                nso = self.mospaces.n_orbs(sp_d[0])
                nsc = self.mospaces.n_orbs(sp_d[1])
                nsv = self.mospaces.n_orbs(sp_d[2])
                nsw = self.mospaces.n_orbs(sp_d[3])
                ret_d.append([((i, j, b, a), 1)
                              for i in range(nso) for j in range(nsc)
                              for a in range(nsv) for b in range(nsw)])

            def sortfctn(x):
                return min(reduce_index(n_orbsa_d, idx) for idx, factor in x)
            ret_d.sort(key=sortfctn)
            ret_d.sort(key=sortfctn)
            ret.extend(ret_d)

        if any(b not in ("ph", "pphh") for b in self.axis_blocks):
            raise NotImplementedError("Blocks other than ph and pphh "
                                      "not implemented")
        return ret

    def to_ndarray(self, out=None):
        """
        Return the ADC matrix object as a dense numpy array. Converts the sparse
        internal representation of the ADC matrix to a dense matrix and return
        as a numpy array.

        Notes
        -----

        This method is only intended to be used for debugging and
        visualisation purposes as it involves computing a large amount of
        matrix-vector products and the returned array consumes a considerable
        amount of memory.

        The resulting matrix has no spin symmetry imposed, which means that
        its eigenspectrum may contain non-physical excitations (e.g. with linear
        combinations of α->β and α->α components in the excitation vector).

        This function has not been sufficiently tested to be considered stable.
        """
        # TODO Update to ph / pphh
        # TODO Still uses deprecated functions
        import tqdm

        from adcc import guess_zero

        # Get zero amplitude of the appropriate symmetry
        # (TODO: Only true for C1, where there is only a single irrep)
        ampl_zero = guess_zero(self)
        assert self.mospaces.point_group == "C1"

        # Build the shape of the returned array
        # Since the basis of the doubles block is not the unit vectors
        # this *not* equal to the shape of the AdcMatrix object
        basis = {b: self.dense_basis(b) for b in self.axis_blocks}
        mat_len = sum(len(basis[b]) for b in basis)

        if out is None:
            out = np.zeros((mat_len, mat_len))
        else:
            if out.shape != (mat_len, mat_len):
                raise ValueError("Output array has shape ({0:}, {1:}), but "
                                 "shape ({2:}, {2:}) is required."
                                 "".format(*out.shape, mat_len))
            out[:] = 0  # Zero all data in out.

        # Check for the cases actually implemented
        if any(b not in ("ph", "pphh") for b in self.axis_blocks):
            raise NotImplementedError("Blocks other than ph and pphh "
                                      "not implemented")
        if "ph" not in self.axis_blocks:
            raise NotImplementedError("Block 'ph' needs to be present")

        # Extract singles-singles block (contiguous)
        assert "ph" in self.axis_blocks
        n_orbs_ph = [self.mospaces.n_orbs(sp) for sp in self.axis_spaces["ph"]]
        n_ph = np.prod(n_orbs_ph)
        assert len(basis["ph"]) == n_ph
        view_ss = out[:n_ph, :n_ph].reshape(*n_orbs_ph, *n_orbs_ph)
        for i in range(n_orbs_ph[0]):
            for a in range(n_orbs_ph[1]):
                ampl = ampl_zero.copy()
                ampl.ph[i, a] = 1
                view_ss[:, :, i, a] = (self @ ampl).ph.to_ndarray()

        # Extract singles-doubles and doubles-doubles block
        if "pphh" in self.axis_blocks:
            assert self.axis_blocks == ["ph", "pphh"]
            view_sd = out[:n_ph, n_ph:].reshape(*n_orbs_ph, len(basis["pphh"]))
            view_dd = out[n_ph:, n_ph:]
            for j, bas1 in tqdm.tqdm(enumerate(basis["pphh"]),
                                     total=len(basis["pphh"])):
                ampl = ampl_zero.copy()
                for idx, val in bas1:
                    ampl.pphh[idx] = val
                ret_ampl = self @ ampl
                view_sd[:, :, j] = ret_ampl.ph.to_ndarray()

                for i, bas2 in enumerate(basis["pphh"]):
                    view_dd[i, j] = sum(val * ret_ampl.pphh[idx]
                                        for idx, val in bas2)

            out[n_ph:, :n_ph] = np.transpose(out[:n_ph, n_ph:])
        return out


class AdcMatrixShifted(AdcMatrix):
    def __init__(self, matrix, shift=0.0):
        """
        Initialise a shifted ADC matrix. Applying this class to a vector ``v``
        represents an efficient version of ``matrix @ v + shift * v``.

        Parameters
        ----------
        matrix : AdcMatrix
            Matrix which is shifted
        shift : float
            Value by which to shift the matrix
        """
        super().__init__(matrix.method, matrix.ground_state,
                         block_orders=matrix.block_orders,
                         intermediates=matrix.intermediates)
        self.shift = shift

    def matvec(self, in_ampl):
        out = super().matvec(in_ampl)
        out = out + self.shift * in_ampl
        return out

    def to_ndarray(self, out=None):
        super().to_ndarray(self, out)
        out = out + self.shift * np.eye(*out.shape)
        return out

    def block_apply(self, block, in_vec):
        ret = super().block_apply(block, in_vec)
        inblock, outblock = block.split("_")
        if inblock == outblock:
            ret += self.shift * in_vec
        return ret

    def diagonal(self):
        out = super().diagonal()
        out = out + self.shift  # Shift the diagonal
        return out

    def block_view(self, block):
        raise NotImplementedError("Block-view not yet implemented for "
                                  "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.

adc-connect / adcc / 16384966787

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous