diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 381b2b1..7651699 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -4,7 +4,6 @@ on: push: branches: [master] pull_request: - branches: [master] jobs: build: @@ -15,9 +14,10 @@ jobs: fail-fast: false matrix: include: - - {python-version: "3.8", os: ubuntu-latest, documentation: True} - - {python-version: "3.9", os: ubuntu-latest, documentation: False} - - {python-version: "3.10", os: ubuntu-latest, documentation: False} + - {python-version: "3.10", os: ubuntu-latest} + - {python-version: "3.11", os: ubuntu-latest} + - {python-version: "3.12", os: ubuntu-latest} + - {python-version: "3.12", os: macos-latest} steps: - uses: actions/checkout@v2 @@ -34,12 +34,17 @@ jobs: python -m pip install .[dev] - name: Linting run: | - python -m black dyson/ --diff --check --verbose - python -m isort dyson/ --diff --check-only --verbose + ruff check + ruff format --check + mypy dyson/ tests/ - name: Run unit tests run: | - python -m pip install pytest pytest-cov - pytest --cov dyson/ + OMP_NUM_THREADS=1 pytest + - name: Run examples + env: + MPLBACKEND: Agg + run: | + find examples -name "*.py" -print0 | xargs -0 -n1 python - name: Upload coverage to Codecov uses: codecov/codecov-action@v3 with: diff --git a/.github/workflows/docs.yaml b/.github/workflows/docs.yaml new file mode 100644 index 0000000..d2ecd51 --- /dev/null +++ b/.github/workflows/docs.yaml @@ -0,0 +1,21 @@ +name: Documentation + +on: + push: + branches: [refactor] + +jobs: + pages: + runs-on: ubuntu-latest + environment: + name: github-pages + url: ${{ steps.deployment.outputs.page_url }} + permissions: + pages: write + id-token: write + steps: + - id: deployment + uses: sphinx-notes/pages@v3 + with: + documentation_path: docs/source + pyproject_extras: dev diff --git a/README.md b/README.md index a223d75..3d9ce1c 100644 --- a/README.md +++ b/README.md @@ -1,9 +1,11 @@ -# dyson: Dyson equation solvers for electron propagator methods +# `dyson`: Dyson equation solvers for Green's function methods -The `dyson` package implements various Dyson equation solvers, with a focus on those avoiding explicitly grid-resolved numerical procedures such as Fourier transforms and analytical continutation. +[![CI](https://github.com/BoothGroup/dyson/workflows/CI/badge.svg)](https://github.com/BoothGroup/dyson/actions?query=workflow%3ACI+branch%3Amaster) + +The `dyson` package implements various Dyson equation solvers, including novel approaches that avoiding explicitly grid-resolved numerical procedures such as Fourier transforms and analytical continuation. These include the moment-resolved block Lanczos methods for moments of the Green's function or self-energy. -### Installation: +## Installation From source: @@ -12,6 +14,6 @@ git clone https://github.com/BoothGroup/dyson pip install . ``` -### Usage +## Usage Examples are available in the `examples` directory. diff --git a/docs/Makefile b/docs/Makefile new file mode 100644 index 0000000..d0c3cbf --- /dev/null +++ b/docs/Makefile @@ -0,0 +1,20 @@ +# Minimal makefile for Sphinx documentation +# + +# You can set these variables from the command line, and also +# from the environment for the first two. +SPHINXOPTS ?= +SPHINXBUILD ?= sphinx-build +SOURCEDIR = source +BUILDDIR = build + +# Put it first so that "make" without argument is like "make help". +help: + @$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) + +.PHONY: help Makefile + +# Catch-all target: route all unknown targets to Sphinx using the new +# "make mode" option. $(O) is meant as a shortcut for $(SPHINXOPTS). +%: Makefile + @$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O) diff --git a/docs/make.bat b/docs/make.bat new file mode 100644 index 0000000..747ffb7 --- /dev/null +++ b/docs/make.bat @@ -0,0 +1,35 @@ +@ECHO OFF + +pushd %~dp0 + +REM Command file for Sphinx documentation + +if "%SPHINXBUILD%" == "" ( + set SPHINXBUILD=sphinx-build +) +set SOURCEDIR=source +set BUILDDIR=build + +%SPHINXBUILD% >NUL 2>NUL +if errorlevel 9009 ( + echo. + echo.The 'sphinx-build' command was not found. Make sure you have Sphinx + echo.installed, then set the SPHINXBUILD environment variable to point + echo.to the full path of the 'sphinx-build' executable. Alternatively you + echo.may add the Sphinx directory to PATH. + echo. + echo.If you don't have Sphinx installed, grab it from + echo.https://www.sphinx-doc.org/ + exit /b 1 +) + +if "%1" == "" goto help + +%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% +goto end + +:help +%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O% + +:end +popd diff --git a/tests/solvers/__init__.py b/docs/source/_static/.gitignore similarity index 100% rename from tests/solvers/__init__.py rename to docs/source/_static/.gitignore diff --git a/docs/source/_templates/.gitignore b/docs/source/_templates/.gitignore new file mode 100644 index 0000000..e69de29 diff --git a/docs/source/conf.py b/docs/source/conf.py new file mode 100644 index 0000000..80ac9cc --- /dev/null +++ b/docs/source/conf.py @@ -0,0 +1,64 @@ +"""Configuration file for the Sphinx documentation builder. + +For the full list of built-in configuration values, see the documentation: +https://www.sphinx-doc.org/en/master/usage/configuration.html +""" + + +# Project information + +project = "dyson" +copyright = "2025, Booth Group, King's College London" +author = "Oliver J. Backhouse, Basil Ibrahim, Marcus K. Allen, George H. Booth" + + +# General configuration + +extensions = [ + "sphinx.ext.autodoc", + "sphinx.ext.autosummary", + "sphinx.ext.coverage", + "sphinx.ext.doctest", + "sphinx.ext.intersphinx", + "sphinx.ext.napoleon", + "sphinx.ext.viewcode", + "sphinx_mdinclude", +] + +templates_path = ["_templates"] +exclude_patterns = [] + + +# Options for HTML output + +html_theme = "sphinx_book_theme" +html_static_path = ["_static"] +default_role = "autolink" + + +# Options for autosummary + +autosummary_generate = True +autodoc_default_options = { + "members": True, + "undoc-members": True, + "inherited-members": False, + "show-inheritance": True, + "member-order": "bysource", +} + + +# Options for intersphinx + +intersphinx_mapping = { + "python": ("https://docs.python.org/3/", None), + "numpy": ("https://numpy.org/doc/stable/", None), + "scipy": ("https://docs.scipy.org/doc/scipy/", None), + "pyscf": ("https://pyscf.org/", None), + "rich": ("https://rich.readthedocs.io/en/stable/", None), +} + + +# Options for napoleon + +napoleon_google_docstring = True diff --git a/docs/source/dyson.rst b/docs/source/dyson.rst new file mode 100644 index 0000000..50de07f --- /dev/null +++ b/docs/source/dyson.rst @@ -0,0 +1,4 @@ +dyson +===== + +.. automodule:: dyson diff --git a/docs/source/index.rst b/docs/source/index.rst new file mode 100644 index 0000000..9ba783a --- /dev/null +++ b/docs/source/index.rst @@ -0,0 +1,17 @@ +.. dyson documentation master file, created by + sphinx-quickstart on Sat Jul 19 19:21:09 2025. + You can adapt this file completely to your liking, but it should at least + contain the root `toctree` directive. + +.. mdinclude:: ../../README.md + +.. toctree:: + :maxdepth: 2 + :caption: Contents: + +.. toctree:: + :maxdepth: 1 + :hidden: + :caption: API Reference + + dyson diff --git a/dyson/__init__.py b/dyson/__init__.py index 13f6c6a..6f3ab4c 100644 --- a/dyson/__init__.py +++ b/dyson/__init__.py @@ -1,98 +1,142 @@ """ -************************************************************* -dyson: Dyson equation solvers for electron propagator methods -************************************************************* -""" - -__version__ = "0.0.0" - -import logging -import os -import subprocess -import sys - -# --- Logging: - - -def output(self, msg, *args, **kwargs): - if self.isEnabledFor(25): - self._log(25, msg, args, **kwargs) - +********************************************************** +dyson: Dyson equation solvers for Green's function methods +********************************************************** + +Dyson equation solvers in :mod:`dyson` are general solvers that accept a variety of inputs to +represent self-energies or existing Green's functions, and solve the Dyson equation in some fashion +to obtain either + +a) a static spectral representation that can be projected into a static Lehmann representation + of the Green's function or self-energy, or +b) a dynamic Green's function. + +The self-energy and Green's function are represented in the following ways: + +.. list-table:: + :header-rows: 1 + :widths: 20 80 + + * - Representation + - Description + * - :class:`~dyson.representations.spectral.Spectral` + - Eigenvalues and eigenvectors of the static self-energy supermatrix, from which the + Lehmann representation of the self-energy or Green's function can be constructed. + * - :class:`~dyson.representations.lehmann.Lehmann` + - The Lehmann representation of the self-energy or Green's function, consisting of pole + energies and their couplings to a physical space. + * - :class:`~dyson.representations.dynamic.Dynamic` + - The dynamic self-energy or Green's function, represented as a series of arrays at each + point on a grid of time or frequency points. + +The available static solvers are, along with their expected inputs: + +.. list-table:: + :header-rows: 1 + :widths: 20 80 + + * - Solver + - Inputs + * - :class:`~dyson.solvers.static.exact.Exact` + - Supermatrix of the static and dynamic self-energy. + * - :class:`~dyson.solvers.static.davidson.Davidson` + - Matrix-vector operation and diagonal of the supermatrix of the static and dynamic + self-energy. + * - :class:`~dyson.solvers.static.downfolded.Downfolded` + - Static self-energy and function returning the dynamic self-energy at a given frequency. + * - :class:`~dyson.solvers.static.mblse.MBLSE` + - Static self-energy and moments of the dynamic self-energy. + * - :class:`~dyson.solvers.static.mblgf.MBLGF` + - Moments of the dynamic Green's function. + * - :class:`~dyson.solvers.static.chempot.AufbauPrinciple` + - Static self-energy, Lehmann representation of the dynamic self-energy, and the target + number of electrons. + * - :class:`~dyson.solvers.static.chempot.AuxiliaryShift` + - Static self-energy, Lehmann representation of the dynamic self-energy, and the target + number of electrons. + * - :class:`~dyson.solvers.static.density.DensityRelaxation` + - Lehmann representation of the dynamic self-energy, function returning the Fock matrix at a + given density, and the target number of electrons. + +For dynamic solvers, all solvers require the grid parameters, along with: + +.. list-table:: + :header-rows: 1 + :widths: 20 80 + + * - Solver + - Inputs + * - :class:`~dyson.solvers.dynamic.corrvec.CorrectionVector` + - Matrix-vector operation and diagonal of the supermatrix of the static and dynamic + self-energy. + * - :class:`~dyson.solvers.dynamic.cpgf.CPGF` + - Chebyshev polynomial moments of the dynamic Green's function. + +For a full accounting of the inputs and their types, please see the documentation for each solver. + +A number of classes are provided to represent the expressions needed to construct these inputs at +different levels of theory. These expressions are all implemented for RHF references, with other +spin symmetries left to the user to implement as needed. The available expressions are: + +.. list-table:: + :header-rows: 1 + :widths: 20 80 + + * - Expression + - Description + * - :data:`~dyson.expressions.hf.HF` + - Hartree--Fock (mean-field) ground state, exploiting Koopmans' theorem for the excited + states. + * - :data:`~dyson.expressions.ccsd.CCSD` + - Coupled cluster singles and doubles ground state, and the respective equation-of-motion + method for the excited states. + * - :data:`~dyson.expressions.fci.FCI` + - Full configuration interaction (exact diagonalisation) ground and excited states. + * - :data:`~dyson.expressions.adc.ADC2` + - Algebraic diagrammatic construction second order excited states, based on a mean-field + ground state. + * - :data:`~dyson.expressions.adc.ADC2x` + - Algebraic diagrammatic construction extended second order excited states, based on a + mean-field ground state. + * - :data:`~dyson.expressions.gw.TDAGW` + - GW theory with the Tamm--Dancoff approximation for the excited states, based on a + mean-field ground state. + * - :data:`~dyson.expressions.hamiltonian.Hamiltonian` + - General Hamiltonian expression, which accepts an array representing the supermatrix of the + self-energy, and supports :mod:`scipy.sparse` matrices. + + +Submodules +---------- + +.. autosummary:: + :toctree: _autosummary + + dyson.expressions + dyson.grids + dyson.representations + dyson.solvers + dyson.util -default_log = logging.getLogger(__name__) -default_log.setLevel(logging.INFO) -default_log.addHandler(logging.StreamHandler(sys.stderr)) -logging.addLevelName(25, "OUTPUT") -logging.Logger.output = output - - -class NullLogger(logging.Logger): - def __init__(self, *args, **kwargs): - super().__init__("null") - - def _log(self, level, msg, args, **kwargs): - pass - - -HEADER = """ _ - | | - __| | _ _ ___ ___ _ __ - / _` || | | |/ __| / _ \ | '_ \ -| (_| || |_| |\__ \| (_) || | | | - \__,_| \__, ||___/ \___/ |_| |_| - __/ | - |___/ %s """ - -def init_logging(log): - """Initialise the logging with a header.""" - - if globals().get("_DYSON_LOG_INITIALISED", False): - return - - # Print header - header_size = max([len(line) for line in HEADER.split("\n")]) - log.info(HEADER % (" " * (18 - len(__version__)) + __version__)) - - # Print versions of dependencies and ebcc - def get_git_hash(directory): - git_directory = os.path.join(directory, ".git") - cmd = ["git", "--git-dir=%s" % git_directory, "rev-parse", "--short", "HEAD"] - try: - git_hash = subprocess.check_output( - cmd, universal_newlines=True, stderr=subprocess.STDOUT - ).rstrip() - except subprocess.CalledProcessError: - git_hash = "N/A" - return git_hash - - import numpy - import pyscf - - log.info("numpy:") - log.info(" > Version: %s" % numpy.__version__) - log.info(" > Git hash: %s" % get_git_hash(os.path.join(os.path.dirname(numpy.__file__), ".."))) - - log.info("pyscf:") - log.info(" > Version: %s" % pyscf.__version__) - log.info(" > Git hash: %s" % get_git_hash(os.path.join(os.path.dirname(pyscf.__file__), ".."))) - - log.info("dyson:") - log.info(" > Version: %s" % __version__) - log.info(" > Git hash: %s" % get_git_hash(os.path.join(os.path.dirname(__file__), ".."))) - - # Environment variables - log.info("OMP_NUM_THREADS = %s" % os.environ.get("OMP_NUM_THREADS", "")) - - log.info("") - - globals()["_DYSON_LOG_INITIALISED"] = True - - -# -- Other imports: - -from dyson.util import * -from dyson.lehmann import Lehmann -from dyson.solvers import * -from dyson.expressions import * +__version__ = "1.0.0" + +import numpy +import scipy + +from dyson.printing import console, quiet +from dyson.representations import Lehmann, Spectral, Dynamic +from dyson.solvers import ( + Exact, + Davidson, + Downfolded, + MBLSE, + MBLGF, + AufbauPrinciple, + AuxiliaryShift, + DensityRelaxation, + CorrectionVector, + CPGF, +) +from dyson.expressions import HF, CCSD, FCI, ADC2, ADC2x, TDAGW, Hamiltonian diff --git a/dyson/expressions/__init__.py b/dyson/expressions/__init__.py index 3df5b77..24d7b00 100644 --- a/dyson/expressions/__init__.py +++ b/dyson/expressions/__init__.py @@ -1,5 +1,94 @@ -from dyson.expressions.expression import BaseExpression -from dyson.expressions.fci import FCI -from dyson.expressions.mp2 import MP2 +r"""Expressions for constructing Green's functions and self-energies. + +Subclasses of :class:`~dyson.expressions.expression.BaseExpression` expose various methods which +provide different representations of the self-energy or Green's function for the given level of +theory. The Green's function is related to the resolvent + +.. math:: + \left[ \omega - \mathbf{H} \right]^{-1}, + +where :math:`\mathbf{H}` is the Hamiltonian, and in the presence of correlation, takes the form of a +self-energy supermatrix + +.. math:: + \mathbf{H} = \begin{bmatrix} \boldsymbol{\Sigma}(\omega) & \mathbf{v} \\ \mathbf{v}^\dagger & + \mathbf{K} + \mathbf{C} \end{bmatrix}, + +which possesses its own Lehmann representation. For more details on these representations, see the +:mod:`~dyson.representations` module. + +The :class:`~dyson.expressions.expression.BaseExpression` interface provides a +:func:`~dyson.expressions.expression.BaseExpression.from_mf` constructor to create an expression of +that level of theory from a mean-field object + +>>> from dyson import util, quiet, FCI +>>> quiet() # Suppress output +>>> mf = util.get_mean_field("H 0 0 0; H 0 0 1", "6-31g") +>>> fci = FCI.h.from_mf(mf) + +The :class:`~dyson.expressions.expression.BaseExpression` interface provides methods to compute the +matrix-vector operations and diagonal of the self-energy supermatrix + +>>> import numpy as np +>>> ham = fci.build_matrix() +>>> np.allclose(np.diag(ham), fci.diagonal()) +True +>>> vec = np.random.random(fci.shape[0]) +>>> np.allclose(fci.apply_hamiltonian(vec), ham @ vec) +True + +More precisely, the Green's function requires also the excitation operators to connect to the +ground state + +.. math:: + \mathbf{G}(\omega) = \left\langle \boldsymbol{\Psi}_0 \right| \hat{a}_p \left[ \omega - + \mathbf{H} \right]^{-1} \hat{a}_q^\dagger \left| \boldsymbol{\Psi}_0 \right\rangle, + +which may be a simple projection when the ground state is mean-field, or otherwise +in the case of correlated ground states. The interface can provide these vectors + +>>> bra = fci.get_excitation_bras() +>>> ket = fci.get_excitation_kets() + +which are vectors with shape ``(nphys, nconfig)`` where ``nphys`` is the number of physical states. + +These methods can be used to construct the moments of the Green's function + +.. math:: + \mathbf{G}_n = \left\langle \boldsymbol{\Psi}_0 \right| \hat{a}_p \mathbf{H}^n + \hat{a}_q^\dagger \left| \boldsymbol{\Psi}_0 \right\rangle, + +which are important for some of the novel approaches implemented in :mod:`dyson`. In the case of +some levels of theory, analytic expressions for the moments of the self-energy are also available. +These moments can be calculated using + +>>> gf_moments = fci.build_gf_moments(nmom=6) + +A list of available expressions is provided in the documentation of :mod:`dyson`. Each expression +is an instance of :class:`~dyson.expressions.expression.ExpressionCollection`, which provides the +subclasses of :class:`~dyson.expressions.expression.BaseExpression` for various sectors such as the +hole and particle. + + +Submodules +---------- + +.. autosummary:: + :toctree: + + expression + hf + ccsd + fci + adc + gw + hamiltonian + +""" + +from dyson.expressions.hf import HF from dyson.expressions.ccsd import CCSD -from dyson.expressions.gw import GW +from dyson.expressions.fci import FCI +from dyson.expressions.adc import ADC2, ADC2x +from dyson.expressions.gw import TDAGW +from dyson.expressions.hamiltonian import Hamiltonian diff --git a/dyson/expressions/adc.py b/dyson/expressions/adc.py new file mode 100644 index 0000000..8eb7c72 --- /dev/null +++ b/dyson/expressions/adc.py @@ -0,0 +1,406 @@ +"""Algebraic diagrammatic construction theory (ADC) expressions [1]_ [2]_. + +.. [1] Banerjee, S., & Sokolov, A. Y. (2019). Third-order algebraic diagrammatic + construction theory for electron attachment and ionization energies: Conventional and Green’s + function implementation. The Journal of Chemical Physics, 151(22). + https://doi.org/10.1063/1.5131771 + +.. [2] Schirmer, J., Cederbaum, L. S., & Walter, O. (1983). New approach to the + one-particle Green’s function for finite Fermi systems. Physical Review. A, General Physics, + 28(3), 1237–1259. https://doi.org/10.1103/physreva.28.1237 +""" + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +from pyscf import adc, ao2mo + +from dyson import numpy as np +from dyson import util +from dyson.expressions.expression import BaseExpression, ExpressionCollection +from dyson.representations.enums import Reduction + +if TYPE_CHECKING: + from types import ModuleType + from typing import Any + + from pyscf.gto.mole import Mole + from pyscf.scf.hf import RHF + + from dyson.typing import Array + + +class BaseADC(BaseExpression): + """Base class for ADC expressions.""" + + hermitian_downfolded = True + hermitian_upfolded = False + + PYSCF_ADC: ModuleType + SIGN: int + METHOD: str = "adc(2)" + METHOD_TYPE: str = "ip" + + def __init__(self, mol: Mole, adc_obj: adc.radc.RADC, imds: Any, eris: Any) -> None: + """Initialise the expression. + + Args: + mol: The molecule object. + adc_obj: PySCF ADC object. + imds: Intermediates from PySCF. + eris: Electron repulsion integrals from PySCF. + """ + self._mol = mol + self._adc_obj = adc_obj + self._imds = imds + self._eris = eris + + @classmethod + def from_adc(cls, adc_obj: adc.radc.RADC) -> BaseADC: + """Construct an MP2 expression from an ADC object. + + Args: + adc_obj: ADC object. + + Returns: + Expression object. + """ + if adc_obj.t1 is None or adc_obj.t2 is None: + warnings.warn("ADC object is not converged.", UserWarning, stacklevel=2) + eris = adc_obj.transform_integrals() + imds = cls.PYSCF_ADC.get_imds(adc_obj, eris) + return cls(adc_obj.mol, adc_obj, imds, eris) + + @classmethod + def from_mf(cls, mf: RHF) -> BaseADC: + """Create an expression from a mean-field object. + + Args: + mf: Mean-field object. + + Returns: + Expression object. + """ + adc_obj = adc.radc.RADC(mf) + adc_obj.method = cls.METHOD + adc_obj.method_type = cls.METHOD_TYPE + adc_obj.kernel_gs() + return cls.from_adc(adc_obj) + + def apply_hamiltonian(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector. + + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. + """ + if np.iscomplexobj(vector): + if np.max(np.abs(vector.imag)) > 1e-11: + raise ValueError("ADC does not support complex vectors.") + vector = vector.real + return self.PYSCF_ADC.matvec(self._adc_obj, self._imds, self._eris)(vector) * self.SIGN + + def apply_hamiltonian_left(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector on the left. + + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. + """ + raise NotImplementedError("Left application of Hamiltonian is not implemented for ADC.") + + def diagonal(self) -> Array: + """Get the diagonal of the Hamiltonian. + + Returns: + Diagonal of the Hamiltonian. + """ + return self.PYSCF_ADC.get_diag(self._adc_obj, self._imds, self._eris) * self.SIGN + + @property + def mol(self) -> Mole: + """Molecule object.""" + return self._mol + + @property + def non_dyson(self) -> bool: + """Whether the expression produces a non-Dyson Green's function.""" + return False + + +class BaseADC_1h(BaseADC): + """Base class for ADC expressions with one-hole Green's function.""" + + PYSCF_ADC = adc.radc_ip + SIGN = -1 + METHOD_TYPE = "ip" + + def get_excitation_vector(self, orbital: int) -> Array: + r"""Obtain the vector corresponding to a fermionic operator acting on the ground state. + + This vector is a generalisation of + + .. math:: + f_i^{\pm} \left| \Psi_0 \right> + + where :math:`f_i^{\pm}` is the fermionic creation or annihilation operator, or a product + thereof, depending on the particular expression and what Green's function it corresponds to. + + The vector defines the excitaiton manifold probed by the Green's function corresponding to + the expression. + + Args: + orbital: Orbital index. + + Returns: + Excitation vector. + """ + if orbital < self.nocc: + return util.unit_vector(self.shape[0], orbital) + return np.zeros(self.shape[0]) + + @property + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + return self.nocc + + +class BaseADC_1p(BaseADC): + """Base class for ADC expressions with one-particle Green's function.""" + + PYSCF_ADC = adc.radc_ea + SIGN = 1 + METHOD_TYPE = "ea" + + def get_excitation_vector(self, orbital: int) -> Array: + r"""Obtain the vector corresponding to a fermionic operator acting on the ground state. + + This vector is a generalisation of + + .. math:: + f_i^{\pm} \left| \Psi_0 \right> + + where :math:`f_i^{\pm}` is the fermionic creation or annihilation operator, or a product + thereof, depending on the particular expression and what Green's function it corresponds to. + + The vector defines the excitaiton manifold probed by the Green's function corresponding to + the expression. + + Args: + orbital: Orbital index. + + Returns: + Excitation vector. + """ + if orbital >= self.nocc: + return util.unit_vector(self.shape[0], orbital - self.nocc) + return np.zeros(self.shape[0]) + + @property + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + return self.nvir + + +class ADC2_1h(BaseADC_1h): + """ADC(2) expressions for the one-hole Green's function.""" + + METHOD = "adc(2)" + + def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) -> Array: + """Build the self-energy moments. + + Args: + nmom: Number of moments to compute. + reduction: Reduction type for the moments. + + Returns: + Moments of the self-energy. + """ + # Get the orbital energies and coefficients + eo = self._adc_obj.mo_energy[: self.nocc] + ev = self._adc_obj.mo_energy[self.nocc :] + co = self._adc_obj.mo_coeff[:, : self.nocc] + cv = self._adc_obj.mo_coeff[:, self.nocc :] + + # Rotate the two-electron integrals + ooov = ao2mo.kernel(self._adc_obj.mol, (co, co, co, cv), compact=False) + ooov = ooov.reshape(eo.size, eo.size, eo.size, ev.size) + left = ooov * 2 - ooov.swapaxes(1, 2) + + # Get the subscript based on the reduction + if Reduction(reduction) == Reduction.NONE: + subscript = "ikla,jkla->ij" + elif Reduction(reduction) == Reduction.DIAG: + subscript = "ikla,ikla->i" + elif Reduction(reduction) == Reduction.TRACE: + subscript = "ikla,ikla->" + else: + Reduction(reduction).raise_invalid_representation() + + # Recursively build the moments + moments_occ: list[Array] = [] + for i in range(nmom): + moments_occ.append(util.einsum(subscript, left, ooov.conj())) + if i < nmom - 1: + left = ( + +util.einsum("ikla,k->ikla", left, eo) + + util.einsum("ikla,l->ikla", left, eo) + - util.einsum("ikla,a->ikla", left, ev) + ) + + # Include the virtual contributions + if Reduction(reduction) == Reduction.NONE: + moments = np.array( + [ + util.block_diag(moment, np.zeros((self.nvir, self.nvir))) + for moment in moments_occ + ] + ) + elif Reduction(reduction) == Reduction.DIAG: + moments = np.array( + [np.concatenate((moment, np.zeros(self.nvir))) for moment in moments_occ] + ) + + return moments + + @property + def nconfig(self) -> int: + """Number of configurations.""" + return self.nocc * self.nocc * self.nvir + + +class ADC2_1p(BaseADC_1p): + """ADC(2) expressions for the one-particle Green's function.""" + + METHOD = "adc(2)" + + def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) -> Array: + """Build the self-energy moments. + + Args: + nmom: Number of moments to compute. + reduction: Reduction type for the moments. + + Returns: + Moments of the self-energy. + """ + # Get the orbital energies and coefficients + eo = self._adc_obj.mo_energy[: self.nocc] + ev = self._adc_obj.mo_energy[self.nocc :] + co = self._adc_obj.mo_coeff[:, : self.nocc] + cv = self._adc_obj.mo_coeff[:, self.nocc :] + + # Rotate the two-electron integrals + vvvo = ao2mo.kernel(self._adc_obj.mol, (cv, cv, cv, co), compact=False) + vvvo = vvvo.reshape(ev.size, ev.size, ev.size, eo.size) + left = vvvo * 2 - vvvo.swapaxes(1, 2) + + # Get the subscript based on the reduction + if Reduction(reduction) == Reduction.NONE: + subscript = "acdi,bcdi->ab" + elif Reduction(reduction) == Reduction.DIAG: + subscript = "acdi,acdi->a" + elif Reduction(reduction) == Reduction.TRACE: + subscript = "acdi,acdi->" + else: + Reduction(reduction).raise_invalid_representation() + + # Recursively build the moments + moments_vir: list[Array] = [] + for i in range(nmom): + moments_vir.append(util.einsum(subscript, left, vvvo.conj())) + if i < nmom - 1: + left = ( + +util.einsum("acdi,c->acdi", left, ev) + + util.einsum("acdi,d->acdi", left, ev) + - util.einsum("acdi,i->acdi", left, eo) + ) + + # Include the occupied contributions + if Reduction(reduction) == Reduction.NONE: + moments = np.array( + [ + util.block_diag(np.zeros((self.nocc, self.nocc)), moment) + for moment in moments_vir + ] + ) + elif Reduction(reduction) == Reduction.DIAG: + moments = np.array( + [np.concatenate((np.zeros(self.nocc), moment)) for moment in moments_vir] + ) + + return moments + + @property + def nconfig(self) -> int: + """Number of configurations.""" + return self.nvir * self.nvir * self.nocc + + +class ADC2x_1h(BaseADC_1h): + """ADC(2)-x expressions for the one-hole Green's function.""" + + METHOD = "adc(2)-x" + + def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) -> Array: + """Build the self-energy moments. + + Args: + nmom: Number of moments to compute. + reduction: Reduction type for the moments. + + Returns: + Moments of the self-energy. + """ + raise NotImplementedError("Self-energy moments not implemented for ADC(2)-x.") + + @property + def nconfig(self) -> int: + """Number of configurations.""" + return self.nocc * self.nocc * self.nvir + + +class ADC2x_1p(BaseADC_1p): + """ADC(2)-x expressions for the one-particle Green's function.""" + + METHOD = "adc(2)-x" + + def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) -> Array: + """Build the self-energy moments. + + Args: + nmom: Number of moments to compute. + reduction: Reduction type for the moments. + + Returns: + Moments of the self-energy. + """ + raise NotImplementedError("Self-energy moments not implemented for ADC(2)-x.") + + @property + def nconfig(self) -> int: + """Number of configurations.""" + return self.nvir * self.nvir * self.nocc + + +class ADC2(ExpressionCollection): + """Collection of ADC(2) expressions for different parts of the Green's function.""" + + _hole = ADC2_1h + _particle = ADC2_1p + _name = "ADC(2)" + + +class ADC2x(ExpressionCollection): + """Collection of ADC(2)-x expressions for different parts of the Green's function.""" + + _hole = ADC2x_1h + _particle = ADC2x_1p + _name = "ADC(2)-x" diff --git a/dyson/expressions/ccsd.py b/dyson/expressions/ccsd.py index b8152e5..ae85e11 100644 --- a/dyson/expressions/ccsd.py +++ b/dyson/expressions/ccsd.py @@ -1,228 +1,512 @@ -""" -EOM-CCSD expressions. +"""Coupled cluster singles and doubles (CCSD) expressions [1]_ [2]_. + +.. [1] Purvis, G. D., & Bartlett, R. J. (1982). A full coupled-cluster singles and doubles + model: The inclusion of disconnected triples. The Journal of Chemical Physics, 76(4), 1910–1918. + https://doi.org/10.1063/1.443164 + +.. [2] Stanton, J. F., & Bartlett, R. J. (1993). The equation of motion coupled-cluster + method. A systematic biorthogonal approach to molecular excitation energies, transition + probabilities, and excited state properties. The Journal of Chemical Physics, 98(9), 7029–7039. + https://doi.org/10.1063/1.464746 + """ -import numpy as np -from pyscf import ao2mo, cc, lib, pbc, scf +from __future__ import annotations +import warnings +from abc import abstractmethod +from typing import TYPE_CHECKING + +from pyscf import cc + +from dyson import numpy as np from dyson import util -from dyson.expressions import BaseExpression - - -class CCSD_1h(BaseExpression): - """ - IP-EOM-CCSD expressions. - """ - - hermitian = False - - def __init__(self, *args, ccsd=None, t1=None, t2=None, l1=None, l2=None, **kwargs): - BaseExpression.__init__(self, *args, **kwargs) - - if ccsd is None: - if isinstance(self.mf, scf.hf.RHF): - ccsd = cc.CCSD(self.mf, mo_coeff=self.mo_coeff, mo_occ=self.mo_occ) - elif isinstance(self.mf, pbc.scf.hf.RHF): - ccsd = pbc.cc.CCSD(self.mf, mo_coeff=self.mo_coeff, mo_occ=self.mo_occ) - else: - raise NotImplementedError( - "EOM-CCSD not implemented for this type of mean-field object." - ) - - ccsd.t1 = t1 - ccsd.t2 = t2 - ccsd.l1 = l1 - ccsd.l2 = l2 - - # Solve CCSD if amplitudes are not provided - if ccsd.t1 is None or ccsd.t2 is None: - ccsd.kernel() - self.t1 = ccsd.t1 - self.t2 = ccsd.t2 - else: - self.t1 = ccsd.t1 - self.t2 = ccsd.t2 +from dyson.expressions.expression import BaseExpression, ExpressionCollection +from dyson.representations.enums import Reduction + +if TYPE_CHECKING: + from typing import Any + + from pyscf.gto.mole import Mole + from pyscf.scf.hf import RHF + + from dyson.typing import Array + + +class BaseCCSD(BaseExpression): + """Base class for CCSD expressions.""" + + hermitian_downfolded = False + hermitian_upfolded = False + + partition: str | None = None + + PYSCF_EOM = cc.eom_rccsd + + def __init__( + self, + mol: Mole, + t1: Array, + t2: Array, + l1: Array, + l2: Array, + imds: Any, + ): + """Initialise the expression. + + Args: + mol: Molecule object. + t1: T1 amplitudes. + t2: T2 amplitudes. + l1: L1 amplitudes. + l2: L2 amplitudes. + imds: Intermediate integrals. + """ + self._mol = mol + self._t1 = t1 + self._t2 = t2 + self._l1 = l1 + self._l2 = l2 + self._imds = imds + self._precompute_imds() + + @abstractmethod + def _precompute_imds(self) -> None: + """Precompute intermediate integrals.""" + pass + + @classmethod + def from_ccsd(cls, ccsd: cc.CCSD) -> BaseCCSD: + """Create an expression from a CCSD object. + + Args: + ccsd: CCSD object. + + Returns: + Expression object. + """ + if not ccsd.converged: + warnings.warn("CCSD T amplitudes are not converged.", UserWarning, stacklevel=2) + if not ccsd.converged_lambda: + warnings.warn("CCSD L amplitudes are not converged.", UserWarning, stacklevel=2) + eris = ccsd.ao2mo() + imds = cls.PYSCF_EOM._IMDS(ccsd, eris=eris) # pylint: disable=protected-access + return cls( + mol=ccsd._scf.mol, # pylint: disable=protected-access + t1=ccsd.t1, + t2=ccsd.t2, + l1=ccsd.l1, + l2=ccsd.l2, + imds=imds, + ) + + @classmethod + def from_mf(cls, mf: RHF) -> BaseCCSD: + """Create an expression from a mean-field object. + + Args: + mf: Mean-field object. + + Returns: + Expression object. + """ + ccsd = cc.CCSD(mf) + ccsd.conv_tol_normt = 1e-9 + ccsd.kernel() + ccsd.solve_lambda() + return cls.from_ccsd(ccsd) + + @abstractmethod + def vector_to_amplitudes(self, vector: Array, *args: Any) -> tuple[Array, Array]: + """Convert a vector to amplitudes. + + Args: + vector: Vector to convert. + args: Additional arguments, redunantly passed during interoperation with PySCF. + + Returns: + Amplitudes. + """ + pass + + @abstractmethod + def amplitudes_to_vector(self, t1: Array, t2: Array) -> Array: + """Convert amplitudes to a vector. + + Args: + t1: T1 amplitudes. + t2: T2 amplitudes. + + Returns: + Vector. + """ + pass + + def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) -> Array: + """Build the self-energy moments. + + Args: + nmom: Number of moments to compute. + reduction: Reduction method to apply to the moments. + + Returns: + Moments of the self-energy. + """ + raise NotImplementedError("Self-energy moments not implemented for CCSD.") + + @property + def mol(self) -> Mole: + """Molecule object.""" + return self._mol + + @property + def t1(self) -> Array: + """T1 amplitudes.""" + return self._t1 + + @property + def t2(self) -> Array: + """T2 amplitudes.""" + return self._t2 + + @property + def l1(self) -> Array: + """L1 amplitudes.""" + return self._l1 + + @property + def l2(self) -> Array: + """L2 amplitudes.""" + return self._l2 + + @property + def non_dyson(self) -> bool: + """Whether the expression produces a non-Dyson Green's function.""" + return False + + # The following properties are for interoperability with PySCF: + + @property + def nmo(self) -> int: + """Get the number of molecular orbitals.""" + return self.nphys + + +class CCSD_1h(BaseCCSD): # pylint: disable=invalid-name + """IP-EOM-CCSD expressions.""" + + def _precompute_imds(self) -> None: + """Precompute intermediate integrals.""" + self._imds.make_ip() + + def vector_to_amplitudes(self, vector: Array, *args: Any) -> tuple[Array, Array]: + """Convert a vector to amplitudes. + + Args: + vector: Vector to convert. + args: Additional arguments, redunantly passed during interoperation with PySCF. + + Returns: + Amplitudes. + """ + return self.PYSCF_EOM.vector_to_amplitudes_ip(vector, self.nphys, self.nocc) + + def amplitudes_to_vector(self, t1: Array, t2: Array) -> Array: + """Convert amplitudes to a vector. + + Args: + t1: T1 amplitudes. + t2: T2 amplitudes. + + Returns: + Vector. + """ + return self.PYSCF_EOM.amplitudes_to_vector_ip(t1, t2) - if ccsd.l1 is None or ccsd.l2 is None: - self.l1, self.l2 = ccsd.solve_lambda() - else: - self.l1 = ccsd.l1 - self.l2 = ccsd.l2 - - self.eris = ccsd.ao2mo() - self.imds = cc.eom_rccsd._IMDS(ccsd, eris=self.eris) - self.imds.make_ip() - - self.eom = lambda: None - self.eom.nmo = self.nmo - self.eom.nocc = self.nocc - self.eom.vector_to_amplitudes = cc.eom_rccsd.vector_to_amplitudes_ip - self.eom.amplitudes_to_vector = cc.eom_rccsd.amplitudes_to_vector_ip - self.eom.partition = None - - def diagonal(self): - diag = -cc.eom_rccsd.ipccsd_diag(self.eom, imds=self.imds) - return diag - - def apply_hamiltonian(self, vector): - hvec = -cc.eom_rccsd.ipccsd_matvec(self.eom, vector, imds=self.imds) - return hvec - - def apply_hamiltonian_left(self, vector): - hvec = -cc.eom_rccsd.lipccsd_matvec(self.eom, vector, imds=self.imds) - return hvec - - def get_wavefunction_bra(self, orb): - t1 = self.t1 - t2 = self.t2 - l1 = self.l1 - l2 = self.l2 - - if orb < self.nocc: - v1 = np.eye(self.nocc)[orb] - v1 -= lib.einsum("ie,e->i", l1, t1[orb]) - tmp = t2[orb] * 2.0 - tmp -= t2[orb].swapaxes(1, 2) - v1 -= lib.einsum("imef,mef->i", l2, tmp) - - tmp = -lib.einsum("ijea,e->ija", l2, t1[orb]) - v2 = tmp * 2.0 - v2 -= tmp.swapaxes(0, 1) - tmp = lib.einsum("ja,i->ija", l1, np.eye(self.nocc)[orb]) - v2 += tmp * 2.0 - v2 -= tmp.swapaxes(0, 1) + def apply_hamiltonian_right(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector on the right. + + Args: + vector: Vector to apply Hamiltonian to. - else: - v1 = l1[:, orb - self.nocc].copy() - v2 = l2[:, :, orb - self.nocc] * 2.0 - v2 -= l2[:, :, :, orb - self.nocc] + Returns: + Output vector. - return self.eom.amplitudes_to_vector(v1, v2) + Notes: + The Hamiltonian is applied in the opposite direction compared to canonical IP-EOM-CCSD, + which reflects the opposite ordering of the excitation operators with respect to the + physical indices in the Green's function. This is only of consequence to non-Hermitian + Green's functions. + """ + return -self.PYSCF_EOM.lipccsd_matvec(self, vector, imds=self._imds) - def get_wavefunction_ket(self, orb): - t1 = self.t1 - t2 = self.t2 - l1 = self.l1 - l2 = self.l2 + def apply_hamiltonian_left(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector on the left. + + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. + + Notes: + The Hamiltonian is applied in the opposite direction compared to canonical IP-EOM-CCSD, + which reflects the opposite ordering of the excitation operators with respect to the + physical indices in the Green's function. This is only of consequence to non-Hermitian + Green's functions. + """ + return -self.PYSCF_EOM.ipccsd_matvec(self, vector, imds=self._imds) + + apply_hamiltonian = apply_hamiltonian_right + apply_hamiltonian.__doc__ = BaseCCSD.apply_hamiltonian.__doc__ + + def diagonal(self) -> Array: + """Get the diagonal of the Hamiltonian. + + Returns: + Diagonal of the Hamiltonian. + """ + return -self.PYSCF_EOM.ipccsd_diag(self, imds=self._imds) + + def get_excitation_bra(self, orbital: int) -> Array: + r"""Obtain the bra vector corresponding to a fermionic operator acting on the ground state. + + The bra vector is the excitation vector corresponding to the bra state, which may or may not + be the same as the ket state vector. + + Args: + orbital: Orbital index. + + Returns: + Bra excitation vector. + + Notes: + The bra and ket are defined in the opposite direction compared to canonical IP-EOM-CCSD, + which reflects the opposite ordering of the excitation operators with respect to the + physical indices in the Green's function. This is only of consequence to non-Hermitian + Green's functions. + + See Also: + :func:`get_excitation_vector`: Function to get the excitation vector when the bra and + ket are the same. + """ + r1: Array + r2: Array + + if orbital < self.nocc: + r1 = np.eye(self.nocc)[orbital] + r2 = np.zeros((self.nocc, self.nocc, self.nvir)) - if orb < self.nocc: - v1 = np.eye(self.nocc)[orb] - v2 = np.zeros((self.nocc, self.nocc, self.nvir)) - else: - v1 = t1[:, orb - self.nocc] - v2 = t2[:, :, orb - self.nocc] - - return self.eom.amplitudes_to_vector(v1, v2) - - -class CCSD_1p(BaseExpression): - """ - EA-EOM-CCSD expressions. - """ - - hermitian = False - - def __init__(self, *args, ccsd=None, t1=None, t2=None, l1=None, l2=None, **kwargs): - BaseExpression.__init__(self, *args, **kwargs) - - if ccsd is None: - if isinstance(self.mf, scf.hf.RHF): - ccsd = cc.CCSD(self.mf, mo_coeff=self.mo_coeff, mo_occ=self.mo_occ) - elif isinstance(self.mf, pbc.scf.hf.RHF): - ccsd = pbc.cc.CCSD(self.mf, mo_coeff=self.mo_coeff, mo_occ=self.mo_occ) - else: - raise NotImplementedError( - "momCCSD not implemented for this type of mean-field object." - ) - # ccsd = cc.CCSD(self.mf, mo_coeff=self.mo_coeff, mo_occ=self.mo_occ) - # Use provided amplitudes if available - ccsd.t1 = t1 - ccsd.t2 = t2 - ccsd.l1 = l1 - ccsd.l2 = l2 - - # Solve CCSD if amplitudes are not provided - if ccsd.t1 is None or ccsd.t2 is None: - ccsd.kernel() - self.t1 = ccsd.t1 - self.t2 = ccsd.t2 else: - self.t1 = ccsd.t1 - self.t2 = ccsd.t2 + r1 = self.t1[:, orbital - self.nocc] + r2 = self.t2[:, :, orbital - self.nocc] + + return self.amplitudes_to_vector(r1, r2) + + def get_excitation_ket(self, orbital: int) -> Array: + r"""Obtain the ket vector corresponding to a fermionic operator acting on the ground state. + + The ket vector is the excitation vector corresponding to the ket state, which may or may not + be the same as the bra state vector. + + Args: + orbital: Orbital index. + + Returns: + Ket excitation vector. + + Notes: + The bra and ket are defined in the opposite direction compared to canonical IP-EOM-CCSD, + which reflects the opposite ordering of the excitation operators with respect to the + physical indices in the Green's function. This is only of consequence to non-Hermitian + Green's functions. + + See Also: + :func:`get_excitation_vector`: Function to get the excitation vector when the bra and + ket are the same. + """ + if orbital < self.nocc: + r1 = np.eye(self.nocc)[orbital] + r1 -= util.einsum("ie,e->i", self.l1, self.t1[orbital]) + tmp = self.t2[orbital] * 2.0 + tmp -= self.t2[orbital].swapaxes(1, 2) + r1 -= util.einsum("imef,mef->i", self.l2, tmp) + + tmp = -util.einsum("ijea,e->ija", self.l2, self.t1[orbital]) + r2 = tmp * 2.0 + r2 -= tmp.swapaxes(0, 1) + tmp = util.einsum("ja,i->ija", self.l1, np.eye(self.nocc)[orbital]) + r2 += tmp * 2.0 + r2 -= tmp.swapaxes(0, 1) - if ccsd.l1 is None or ccsd.l2 is None: - self.l1, self.l2 = ccsd.solve_lambda() else: - self.l1 = ccsd.l1 - self.l2 = ccsd.l2 - - self.eris = ccsd.ao2mo() - self.imds = cc.eom_rccsd._IMDS(ccsd, eris=self.eris) - self.imds.make_ea() - - self.eom = lambda: None - self.eom.nmo = self.nmo - self.eom.nocc = self.nocc - self.eom.vector_to_amplitudes = cc.eom_rccsd.vector_to_amplitudes_ea - self.eom.amplitudes_to_vector = cc.eom_rccsd.amplitudes_to_vector_ea - self.eom.partition = None - - def diagonal(self): - diag = cc.eom_rccsd.eaccsd_diag(self.eom, imds=self.imds) - return diag - - def apply_hamiltonian(self, vector): - hvec = cc.eom_rccsd.eaccsd_matvec(self.eom, vector, imds=self.imds) - return hvec - - def apply_hamiltonian_left(self, vector): - hvec = cc.eom_rccsd.leaccsd_matvec(self.eom, vector, imds=self.imds) - return hvec - - def get_wavefunction_bra(self, orb): - t1 = self.t1 - t2 = self.t2 - l1 = self.l1 - l2 = self.l2 - - if orb < self.nocc: - v1 = -l1[orb] - v2 = -l2[orb] * 2.0 - v2 += l2[:, orb] + r1 = self.l1[:, orbital - self.nocc].copy() + r2 = self.l2[:, :, orbital - self.nocc] * 2.0 + r2 -= self.l2[:, :, :, orbital - self.nocc] + + return self.amplitudes_to_vector(r1, r2) + + get_excitation_vector = get_excitation_ket + get_excitation_vector.__doc__ = BaseCCSD.get_excitation_vector.__doc__ + + @property + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + return self.nocc + + @property + def nconfig(self) -> int: + """Number of configurations.""" + return self.nocc * self.nocc * self.nvir + + +class CCSD_1p(BaseCCSD): # pylint: disable=invalid-name + """EA-EOM-CCSD expressions.""" + + def _precompute_imds(self) -> None: + """Precompute intermediate integrals.""" + self._imds.make_ea() + + def vector_to_amplitudes(self, vector: Array, *args: Any) -> tuple[Array, Array]: + """Convert a vector to amplitudes. + + Args: + vector: Vector to convert. + args: Additional arguments, redunantly passed during interoperation with PySCF. + + Returns: + Amplitudes. + """ + return self.PYSCF_EOM.vector_to_amplitudes_ea(vector, self.nphys, self.nocc) + + def amplitudes_to_vector(self, t1: Array, t2: Array) -> Array: + """Convert amplitudes to a vector. + + Args: + t1: T1 amplitudes. + t2: T2 amplitudes. + + Returns: + Vector. + """ + return self.PYSCF_EOM.amplitudes_to_vector_ea(t1, t2) + + def apply_hamiltonian_right(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector on the right. + + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. + """ + return self.PYSCF_EOM.eaccsd_matvec(self, vector, imds=self._imds) + + def apply_hamiltonian_left(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector on the left. + + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. + """ + return self.PYSCF_EOM.leaccsd_matvec(self, vector, imds=self._imds) + + apply_hamiltonian = apply_hamiltonian_right + apply_hamiltonian.__doc__ = BaseCCSD.apply_hamiltonian.__doc__ + + def diagonal(self) -> Array: + """Get the diagonal of the Hamiltonian. + + Returns: + Diagonal of the Hamiltonian. + """ + return self.PYSCF_EOM.eaccsd_diag(self, imds=self._imds) + + def get_excitation_bra(self, orbital: int) -> Array: + r"""Obtain the bra vector corresponding to a fermionic operator acting on the ground state. + + The bra vector is the excitation vector corresponding to the bra state, which may or may not + be the same as the ket state vector. + + Args: + orbital: Orbital index. + + Returns: + Bra excitation vector. + + See Also: + :func:`get_excitation_vector`: Function to get the excitation vector when the bra and + ket are the same. + """ + if orbital < self.nocc: + r1 = -self.l1[orbital] + r2 = -self.l2[orbital] * 2.0 + r2 += self.l2[:, orbital] else: - v1 = np.eye(self.nvir)[orb - self.nocc] - v1 -= lib.einsum("mb,m->b", l1, t1[:, orb - self.nocc]) - tmp = t2[:, :, :, orb - self.nocc] * 2.0 - tmp -= t2[:, :, orb - self.nocc] - v1 -= lib.einsum("kmeb,kme->b", l2, tmp) - - tmp = -lib.einsum("ikba,k->iab", l2, t1[:, orb - self.nocc]) - v2 = tmp * 2.0 - v2 -= tmp.swapaxes(1, 2) - tmp = lib.einsum("ib,a->iab", l1, np.eye(self.nvir)[orb - self.nocc]) - v2 += tmp * 2.0 - v2 -= tmp.swapaxes(1, 2) - - return self.eom.amplitudes_to_vector(v1, v2) - - def get_wavefunction_ket(self, orb): - t1 = self.t1 - t2 = self.t2 - l1 = self.l1 - l2 = self.l2 - - if orb < self.nocc: - v1 = t1[orb] - v2 = t2[orb] + r1 = np.eye(self.nvir)[orbital - self.nocc] + r1 -= util.einsum("mb,m->b", self.l1, self.t1[:, orbital - self.nocc]) + tmp = self.t2[:, :, :, orbital - self.nocc] * 2.0 + tmp -= self.t2[:, :, orbital - self.nocc] + r1 -= util.einsum("kmeb,kme->b", self.l2, tmp) + + tmp = -util.einsum("ikba,k->iab", self.l2, self.t1[:, orbital - self.nocc]) + r2 = tmp * 2.0 + r2 -= tmp.swapaxes(1, 2) + tmp = util.einsum("ib,a->iab", self.l1, np.eye(self.nvir)[orbital - self.nocc]) + r2 += tmp * 2.0 + r2 -= tmp.swapaxes(1, 2) + + return self.amplitudes_to_vector(r1, r2) + + def get_excitation_ket(self, orbital: int) -> Array: + r"""Obtain the ket vector corresponding to a fermionic operator acting on the ground state. + + The ket vector is the excitation vector corresponding to the ket state, which may or may not + be the same as the bra state vector. + + Args: + orbital: Orbital index. + + Returns: + Ket excitation vector. + + See Also: + :func:`get_excitation_vector`: Function to get the excitation vector when the bra and + ket are the same. + """ + r1: Array + r2: Array + + if orbital < self.nocc: + r1 = self.t1[orbital] + r2 = self.t2[orbital] + else: - v1 = -np.eye(self.nvir)[orb - self.nocc] - v2 = np.zeros((self.nocc, self.nvir, self.nvir)) + r1 = -np.eye(self.nvir)[orbital - self.nocc] + r2 = np.zeros((self.nocc, self.nvir, self.nvir)) + + return -self.amplitudes_to_vector(r1, r2) + + get_excitation_vector = get_excitation_ket + get_excitation_vector.__doc__ = BaseCCSD.get_excitation_vector.__doc__ + + @property + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + return self.nvir + + @property + def nconfig(self) -> int: + """Number of configurations.""" + return self.nvir * self.nvir * self.nocc - return -self.eom.amplitudes_to_vector(v1, v2) +class CCSD(ExpressionCollection): + """Collection of CCSD expressions for different parts of the Green's function.""" -CCSD = { - "1h": CCSD_1h, - "1p": CCSD_1p, -} + _hole = CCSD_1h + _particle = CCSD_1p + _name = "CCSD" diff --git a/dyson/expressions/expression.py b/dyson/expressions/expression.py index 161ceed..c866748 100644 --- a/dyson/expressions/expression.py +++ b/dyson/expressions/expression.py @@ -1,287 +1,519 @@ -""" -Expression base class. -""" +"""Base class for expressions.""" -import numpy as np +from __future__ import annotations -from dyson import default_log, init_logging +import warnings +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING +from dyson import numpy as np +from dyson import util +from dyson.representations.enums import Reduction -class BaseExpression: - """ - Base class for all expressions. +if TYPE_CHECKING: + from typing import Callable + + from pyscf.gto.mole import Mole + from pyscf.scf.hf import RHF + + from dyson.typing import Array + + +class BaseExpression(ABC): + """Base class for expressions. + + Attributes: + hermitian_downfolded: Whether the expression is Hermitian when downfolded into the physical + space. + hermitian_upfolded: Whether the expression is Hermitian when upfolded as a supermatrix. """ - hermitian = True - - def __init__(self, mf, mo_energy=None, mo_coeff=None, mo_occ=None, log=None): - self.log = log or default_log - # init_logging(self.log) - # self.log.info("") - # self.log.info("%s", self.__class__.__name__) - # self.log.info("%s", "*" * len(self.__class__.__name__)) - - if mo_energy is None: - mo_energy = mf.mo_energy - if mo_coeff is None: - mo_coeff = mf.mo_coeff - if mo_occ is None: - mo_occ = mf.mo_occ - - self.mf = mf - self.mo_energy = mo_energy - self.mo_coeff = mo_coeff - self.mo_occ = mo_occ - - def apply_hamiltonian(self, vector): - """Apply the Hamiltonian to a trial vector. - - Parameters - ---------- - vector : numpy.ndarray - Vector to apply Hamiltonian to. - - Returns - ------- - output : numpy.ndarray - Output vector. + hermitian_downfolded: bool = True + hermitian_upfolded: bool = True + + @classmethod + @abstractmethod + def from_mf(cls, mf: RHF) -> BaseExpression: + """Create an expression from a mean-field object. + + Args: + mf: Mean-field object. + + Returns: + Expression object. """ + pass + + @abstractmethod + def apply_hamiltonian(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector. + + Args: + vector: Vector to apply Hamiltonian to. - raise NotImplementedError + Returns: + Output vector. + """ + pass - def apply_hamiltonian_left(self, vector): - """Apply the Hamiltonian to a trial vector on the left. + def apply_hamiltonian_left(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector on the left. - Parameters - ---------- - vector : numpy.ndarray - Vector to apply Hamiltonian to. + Args: + vector: Vector to apply Hamiltonian to. - Returns - ------- - output : numpy.ndarray + Returns: Output vector. """ + return self.apply_hamiltonian(vector) - raise NotImplementedError + def apply_hamiltonian_right(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector on the right. + + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. + """ + return self.apply_hamiltonian(vector) - def diagonal(self): + @abstractmethod + def diagonal(self) -> Array: """Get the diagonal of the Hamiltonian. - Returns - ------- - diag : numpy.ndarray + Returns: Diagonal of the Hamiltonian. """ + pass - raise NotImplementedError + def build_matrix(self) -> Array: + """Build the Hamiltonian matrix. - def get_wavefunction(self, orb): - """Obtain the wavefunction as a vector, for a given orbital. + Returns: + Hamiltonian matrix. - Parameters - ---------- - orb : int - Orbital index. - - Returns - ------- - wfn : numpy.ndarray - Wavefunction vector. + Notes: + This method uses :func:`apply_hamiltonian` to build the matrix by applying unit vectors, + it is not designed to be efficient. """ + size = self.diagonal().size + if size > 2048: + warnings.warn( + "The Hamiltonian matrix is large. This may take a while to compute.", + UserWarning, + 2, + ) + return np.array([self.apply_hamiltonian(util.unit_vector(size, i)) for i in range(size)]).T - raise NotImplementedError - - def get_wavefunction_bra(self, orb): - return self.get_wavefunction(orb) - - def get_wavefunction_ket(self, orb): - return self.get_wavefunction(orb) - - def build_gf_moments(self, nmom, store_vectors=True, left=False): - """Build moments of the Green's function. - - Parameters - ---------- - nmom : int or tuple of int - Number of moments to compute. - store_vectors : bool, optional - Store all vectors on disk rather than storing them all - ahead of time. With `store_vectors=True`, the memory - overhead of the vectors is O(N) larger. With - `store_vectors=False`, the CPU overhead of the vectors is - O(N) larger. Default value is `True`. - left : bool, optional - Use the left-handed Hamiltonian application instead of the - right-handed one. Default value is `False`. - - Returns - ------- - t : numpy.ndarray - Moments of the Green's function. + @abstractmethod + def get_excitation_vector(self, orbital: int) -> Array: + r"""Obtain the vector corresponding to a fermionic operator acting on the ground state. + + This vector is a generalisation of + + .. math:: + f_i^{\pm} \left| \Psi_0 \right> + + where :math:`f_i^{\pm}` is the fermionic creation or annihilation operator, or a product + thereof, depending on the particular expression and what Green's function it corresponds to. + + The vector defines the excitaiton manifold probed by the Green's function corresponding to + the expression. + + Args: + orbital: Orbital index. + + Returns: + Excitation vector. """ + pass - t = np.zeros((nmom, self.nphys, self.nphys)) + def get_excitation_bra(self, orbital: int) -> Array: + r"""Obtain the bra vector corresponding to a fermionic operator acting on the ground state. - if left: - get_wavefunction_bra = self.get_wavefunction_ket - get_wavefunction_ket = self.get_wavefunction_bra - apply_hamiltonian = self.apply_hamiltonian_left - else: - get_wavefunction_bra = self.get_wavefunction_bra - get_wavefunction_ket = self.get_wavefunction_ket - apply_hamiltonian = self.apply_hamiltonian + The bra vector is the excitation vector corresponding to the bra state, which may or may not + be the same as the ket state vector. - if store_vectors: - v = [get_wavefunction_bra(i) for i in range(self.nphys)] + Args: + orbital: Orbital index. - for i in range(self.nphys): - u = get_wavefunction_ket(i) + Returns: + Bra excitation vector. - for n in range(nmom): - for j in range(i if self.hermitian else 0, self.nphys): - if not store_vectors: - v = {j: get_wavefunction_bra(j)} + See Also: + :func:`get_excitation_vector`: Function to get the excitation vector when the bra and + ket are the same. + """ + return self.get_excitation_vector(orbital) - t[n, i, j] = np.dot(v[j], u) + def get_excitation_ket(self, orbital: int) -> Array: + r"""Obtain the ket vector corresponding to a fermionic operator acting on the ground state. - if self.hermitian: - t[n, j, i] = t[n, i, j] + The ket vector is the excitation vector corresponding to the ket state, which may or may not + be the same as the bra state vector. - if n != (nmom - 1): - u = apply_hamiltonian(u) + Args: + orbital: Orbital index. - if left: - t = t.transpose(0, 2, 1).conj() - - return t - - def build_gf_chebyshev_moments(self, nmom, store_vectors=True, left=False, scaling=None): - """Build moments of the Green's function using Chebyshev polynomials. - - Parameters - ---------- - nmom : int or tuple of int - Number of moments to compute. - store_vectors : bool, optional - Store all vectors on disk rather than storing them all - ahead of time. With `store_vectors=True`, the memory - overhead of the vectors is O(N) larger. With - `store_vectors=False`, the CPU overhead of the vectors is - O(N) larger. Default value is `True`. - left : bool, optional - Use the left-handed Hamiltonian application instead of the - right-handed one. Default value is `False`. - scaling : tuple of float - Scaling parameters, such that the energy scale of the - Lehmann representation is scaled as - `(energies - scaling[1]) / scaling[0]`. If `None`, the - scaling paramters are computed as - `(max(energies) - min(energies)) / (2.0 - 1e-3)` and - `(max(energies) + min(energies)) / 2.0`, respectively. - - Returns - ------- - t : numpy.ndarray - Chebyshev moments of the Green's function. + Returns: + Ket excitation vector. + + See Also: + :func:`get_excitation_vector`: Function to get the excitation vector when the bra and + ket are the same. """ + return self.get_excitation_vector(orbital) - if scaling is not None: - a, b = scaling - else: - # Calculate the scaling parameters by the range of the - # eigenvalues of the Hamiltonian. These can be approximated - # using the diagonal of the Hamiltonian. A more effective - # method would be to use the Lanczos or Davidson algorithms - # to compute the extremum eigenvalues and pass them in as - # an argument. - diag = self.diagonal() - emin = min(diag) - emax = max(diag) - a = (emax - emin) / (2.0 - 1e-3) - b = (emax + emin) / 2.0 + def get_excitation_vectors(self) -> list[Array]: + """Get the excitation vectors for all orbitals. - t = np.zeros((nmom, self.nphys, self.nphys)) + Returns: + List of excitation vectors for all orbitals. - if left: - get_wavefunction_bra = self.get_wavefunction_ket - get_wavefunction_ket = self.get_wavefunction_bra - apply_hamiltonian = self.apply_hamiltonian_left - else: - get_wavefunction_bra = self.get_wavefunction_bra - get_wavefunction_ket = self.get_wavefunction_ket - apply_hamiltonian = self.apply_hamiltonian + See Also: + :func:`get_excitation_vector`: Function to get the excitation vector for a single + orbital. + """ + return [self.get_excitation_vector(i) for i in range(self.nphys)] - def apply_scaled_hamiltonian(v): - # [(H - b) / a] v = H (v / a) - b (v / a) - v_scaled = v / a - return apply_hamiltonian(v_scaled) - b * v_scaled + def get_excitation_bras(self) -> list[Array]: + """Get the bra excitation vectors for all orbitals. + + Returns: + List of bra excitation vectors for all orbitals. + + See Also: + :func:`get_excitation_bra`: Function to get the bra excitation vector for a single + orbital. + """ + return [self.get_excitation_bra(i) for i in range(self.nphys)] + + def get_excitation_kets(self) -> list[Array]: + """Get the ket excitation vectors for all orbitals. + + Returns: + List of ket excitation vectors for all orbitals. + See Also: + :func:`get_excitation_ket`: Function to get the ket excitation vector for a single + orbital. + """ + return [self.get_excitation_ket(i) for i in range(self.nphys)] + + def _build_gf_moments( + self, + get_bra: Callable[[int], Array], + get_ket: Callable[[int], Array], + apply_hamiltonian_poly: Callable[[Array, Array, int], Array], + nmom: int, + store_vectors: bool = True, + left: bool = False, + reduction: Reduction = Reduction.NONE, + ) -> Array: + """Build the moments of the Green's function.""" + # Precompute bra vectors if needed if store_vectors: - v = [get_wavefunction_bra(i) for i in range(self.nphys)] + bras = list(map(get_bra, range(self.nphys))) + # Loop over ket vectors + moments: dict[tuple[int, int, int], Array] = {} for i in range(self.nphys): - u = get_wavefunction_ket(i) + ket = ket_prev = get_ket(i) + # Loop over moment orders for n in range(nmom): - for j in range(i if self.hermitian else 0, self.nphys): - if not store_vectors: - v = {j: get_wavefunction_bra(j)} - - t[n, i, j] = np.dot(v[j], u) + if Reduction(reduction) == Reduction.NONE: + # Loop over bra vectors + for j in range(i if self.hermitian_downfolded else 0, self.nphys): + bra = bras[j] if store_vectors else get_bra(j) + + # Contract the bra and ket vectors + moments[n, j, i] = bra.conj() @ ket + if self.hermitian_downfolded: + moments[n, i, j] = moments[n, j, i].conj() + + else: + # Contract the bra and ket vectors + bra = bras[i] if store_vectors else get_bra(i) + moments[n, i, i] = bra.conj() @ ket + + # Apply the Hamiltonian to the ket vector + if n != nmom - 1: + ket, ket_prev = apply_hamiltonian_poly(ket, ket_prev, n), ket + + if Reduction(reduction) == Reduction.NONE: + # Convert the moments to a numpy array + moments_array = np.array( + [ + moments[n, i, j] + for n in range(nmom) + for i in range(self.nphys) + for j in range(self.nphys) + ] + ) + moments_array = moments_array.reshape(nmom, self.nphys, self.nphys) + + # Transpose if necessary + if left: + moments_array = moments_array.transpose(0, 2, 1).conj() + + elif Reduction(reduction) == Reduction.DIAG: + # Convert the moments to a numpy array, only keeping the diagonal elements + moments_array = np.array( + [moments[n, i, i] for n in range(nmom) for i in range(self.nphys)] + ) + moments_array = moments_array.reshape(nmom, self.nphys) + + elif Reduction(reduction) == Reduction.TRACE: + # Convert the moments to a numpy array, only keeping the trace + moments_array = np.array( + [sum([moments[n, i, i] for i in range(self.nphys)]) for n in range(nmom)] + ) - if self.hermitian: - t[n, j, i] = t[n, i, j] + else: + Reduction(reduction).raise_invalid_representation() + + return moments_array + + def build_gf_moments( + self, + nmom: int, + store_vectors: bool = True, + left: bool = False, + reduction: Reduction = Reduction.NONE, + ) -> Array: + """Build the moments of the Green's function. + + Args: + nmom: Number of moments to compute. + store_vectors: Whether to store the vectors on disk. Storing the vectors makes the + memory overhead scale worse, but the CPU overhead scales better. + left: Whether to use the left-handed Hamiltonian application. + reduction: Reduction to apply to the moments. + + Returns: + Moments of the Green's function. - if n != (nmom - 1): - if n == 0: - # u_{1} = H u_{0} - u, u_prev = apply_scaled_hamiltonian(u), u - else: - # u_{i} = 2 H u_{i-1} - u_{i-2} - u, u_prev = 2.0 * apply_scaled_hamiltonian(u) - u_prev, u + Notes: + Unlike :func:`dyson.representations.lehmann.Lehmann.moments`, this function takes the + number of moments to compute as an argument, rather than a single order or list of + orders. This is because in this case, the moments are computed recursively. + """ + # Get the appropriate functions + if left: + get_bra = self.get_excitation_ket + get_ket = self.get_excitation_bra + apply_hamiltonian = self.apply_hamiltonian_left + else: + get_bra = self.get_excitation_bra + get_ket = self.get_excitation_ket + apply_hamiltonian = self.apply_hamiltonian_right + + return self._build_gf_moments( + get_bra, + get_ket, + lambda vector, vector_prev, n: apply_hamiltonian(vector), + nmom, + store_vectors=store_vectors, + left=left, + reduction=reduction, + ) + + def build_gf_chebyshev_moments( + self, + nmom: int, + store_vectors: bool = True, + left: bool = False, + scaling: tuple[float, float] | None = None, + reduction: Reduction = Reduction.NONE, + ) -> Array: + """Build the moments of the Green's function using Chebyshev polynomials. + + Args: + nmom: Number of moments to compute. + store_vectors: Whether to store the vectors on disk. Storing the vectors makes the + memory overhead scale worse, but the CPU overhead scales better. + left: Whether to use the left-handed Hamiltonian application. + scaling: Scaling factors to ensure the energy scale of the Lehmann representation is + in ``[-1, 1]``. The scaling is applied as ``(energies - scaling[1]) / scaling[0]``. + If ``None``, the default scaling is computed as + ``(max(energies) - min(energies)) / (2.0 - 1e-3)`` and + ``(max(energies) + min(energies)) / 2.0``, respectively. + reduction: Reduction to apply to the moments. + + Returns: + Chebyshev polynomial moments of the Green's function. + + Notes: + Unlike :func:`dyson.representations.lehmann.Lehmann.chebyshev_moments`, this function + takes the number of moments to compute as an argument, rather than a single order or + list of orders. This is because in this case, the moments are computed recursively. + """ + if scaling is None: + # Approximate the energy scale of the spectrum using the diagonal -- can also use an + # iterative eigensolver to better approximate this + diag = self.diagonal() + scaling = util.get_chebyshev_scaling_parameters(diag.min(), diag.max()) + # Get the appropriate functions if left: - t = t.transpose(1, 2).conj() + get_bra = self.get_excitation_ket + get_ket = self.get_excitation_bra + apply_hamiltonian = self.apply_hamiltonian_left + else: + get_bra = self.get_excitation_bra + get_ket = self.get_excitation_ket + apply_hamiltonian = self.apply_hamiltonian_right - return t + def _apply_hamiltonian_poly(vector: Array, vector_prev: Array, n: int) -> Array: + """Apply the scaled Hamiltonian polynomial to a vector.""" + # [(H - b) / a] v = H (v / a) - b (v / a) + vector_scaled = vector / scaling[0] + result = apply_hamiltonian(vector_scaled) - scaling[1] * vector_scaled + if n == 0: + return result # u_{1} = H u_{0} + return 2.0 * result - vector_prev # u_{n} = 2 H u_{n-1} - u_{n-2} + + return self._build_gf_moments( + get_bra, + get_ket, + _apply_hamiltonian_poly, + nmom, + store_vectors=store_vectors, + left=left, + reduction=reduction, + ) + + @abstractmethod + def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) -> Array: + """Build the self-energy moments. + + Args: + nmom: Number of moments to compute. + reduction: Reduction to apply to the moments. + + Returns: + Moments of the self-energy. + """ + pass - def build_se_moments(self, nmom): - """Build moments of the self-energy. + @property + @abstractmethod + def mol(self) -> Mole: + """Molecule object.""" + pass - Parameters - ---------- - nmom : int or tuple of int - Number of moments to compute. + @property + @abstractmethod + def non_dyson(self) -> bool: + """Whether the expression produces a non-Dyson Green's function.""" + pass - Returns - ------- - t : numpy.ndarray - Moments of the self-energy. - """ + @property + def hermitian(self) -> bool: + """Whether the expression is Hermitian.""" + return self.hermitian_downfolded and self.hermitian_upfolded - raise NotImplementedError + @property + def nphys(self) -> int: + """Number of physical orbitals.""" + return self.mol.nao @property - def nmo(self): - return self.mo_coeff.shape[-1] + @abstractmethod + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + pass @property - def nocc(self): - return np.sum(self.mo_occ > 0) + def nconfig(self) -> int: + """Number of configurations in the non-singles sectors.""" + return self.diagonal().size - self.nsingle @property - def nvir(self): - return np.sum(self.mo_occ == 0) + def shape(self) -> tuple[int, int]: + """Shape of the Hamiltonian matrix.""" + return (self.nconfig + self.nsingle, self.nconfig + self.nsingle) @property - def nalph(self): - return self.nocc + def nocc(self) -> int: + """Number of occupied orbitals.""" + if self.mol.nelectron % 2: + raise NotImplementedError("Open-shell systems are not supported.") + return self.mol.nelectron // 2 @property - def nbeta(self): - return self.nocc + def nvir(self) -> int: + """Number of virtual orbitals.""" + return self.nphys - self.nocc + + +class _ExpressionCollectionMeta(type): + """Metaclass for the ExpressionCollection class.""" + + __wrapped__: bool = False + + def __getattr__(cls, key: str) -> type[BaseExpression]: + """Get an expression by its name.""" + if key in {"hole", "ip", "o", "h"}: + if cls._hole is None: + raise ValueError("Hole expression is not set.") + return cls._hole + elif key in {"particle", "ea", "v", "p"}: + if cls._particle is None: + raise ValueError("Particle expression is not set.") + return cls._particle + elif key in {"central", "dyson"}: + if cls._dyson is None: + raise ValueError("Central (Dyson) expression is not set.") + return cls._dyson + elif key in {"neutral", "ee", "ph"}: + if cls._neutral is None: + raise ValueError("Neutral expression is not set.") + return cls._neutral + else: + raise ValueError(f"Expression '{key}' is not defined in the collection.") + + __getitem__ = __getattr__ @property - def nphys(self): - return self.nmo + def _classes(cls) -> set[type[BaseExpression]]: + """Get all classes in the collection.""" + return { + cls for cls in [cls._hole, cls._particle, cls._dyson, cls._neutral] if cls is not None + } + + def __contains__(cls, key: str) -> bool: + """Check if an expression exists by its name.""" + try: + cls[key] # type: ignore[index] + return True + except ValueError: + return False + + +class ExpressionCollection(metaclass=_ExpressionCollectionMeta): + """Collection of expressions for different parts of the Green's function.""" + + _hole: type[BaseExpression] | None = None + _particle: type[BaseExpression] | None = None + _dyson: type[BaseExpression] | None = None + _neutral: type[BaseExpression] | None = None + _name: str | None = None + + @classmethod + def __getattr__(cls, key: str) -> type[BaseExpression]: + """Get an expression by its name.""" + return getattr(type(cls), key) + + __getitem__ = __getattr__ + + def __contains__(cls, key: str) -> bool: + """Check if an expression exists by its name.""" + return getattr(type(cls), key, None) is not None + + @classmethod + def __repr__(cls) -> str: + """String representation of the collection.""" + return f"ExpressionCollection({cls._name})" if cls._name else "ExpressionCollection" diff --git a/dyson/expressions/fci.py b/dyson/expressions/fci.py index 2cf4792..ca7d203 100644 --- a/dyson/expressions/fci.py +++ b/dyson/expressions/fci.py @@ -1,148 +1,252 @@ +"""Full configuration interaction (FCI) expressions [1]_. + +.. [1] Knowles, P., & Handy, N. (1984). A new determinant-based full configuration + interaction method. Chemical Physics Letters, 111(4–5), 315–321. + https://doi.org/10.1016/0009-2614(84)85513-x """ -FCI expressions. -""" -import numpy as np -from pyscf import ao2mo, fci, lib +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING + +from pyscf import ao2mo, fci + +from dyson.expressions.expression import BaseExpression, ExpressionCollection +from dyson.representations.enums import Reduction + +if TYPE_CHECKING: + from typing import Callable + + from pyscf.gto.mole import Mole + from pyscf.scf.hf import RHF + + from dyson.typing import Array + + +class BaseFCI(BaseExpression): + """Base class for FCI expressions.""" + + hermitian_downfolded = True + hermitian_upfolded = True + + SIGN: int + DELTA_ALPHA: int + DELTA_BETA: int + STATE_FUNC: Callable[[Array, int, tuple[int, int], int], Array] -from dyson import default_log, util -from dyson.expressions import BaseExpression + def __init__( + self, + mol: Mole, + e_fci: Array, + c_fci: Array, + hamiltonian: Array, + diagonal: Array, + chempot: Array | float = 0.0, + ): + """Initialise the expression. + Args: + mol: Molecule object. + e_fci: FCI eigenvalues. + c_fci: FCI eigenvectors. + hamiltonian: Hamiltonian matrix. + diagonal: Diagonal of the FCI Hamiltonian. + chempot: Chemical potential. + """ + self._mol = mol + self._e_fci = e_fci + self._c_fci = c_fci + self._hamiltonian = hamiltonian + self._diagonal = diagonal + self._chempot = chempot + + @classmethod + def from_fci(cls, ci: fci.FCI, h1e: Array, h2e: Array) -> BaseFCI: + """Create an expression from an FCI object. + + Args: + ci: FCI object. + h1e: One-electron Hamiltonian matrix. + h2e: Two-electron Hamiltonian matrix. + + Returns: + Expression object. + """ + if ci.mol is None: + raise ValueError("FCI object must be initialised with a molecule.") + nelec = (ci.mol.nelec[0] + cls.DELTA_ALPHA, ci.mol.nelec[1] + cls.DELTA_BETA) + hamiltonian = ci.absorb_h1e(h1e, h2e, ci.mol.nao, nelec, 0.5) + diagonal = ci.make_hdiag(h1e, h2e, ci.mol.nao, nelec) + return cls( + ci.mol, + ci.eci, + ci.ci, + hamiltonian, + diagonal, + ) + + @classmethod + def from_mf(cls, mf: RHF) -> BaseFCI: + """Create an expression from a mean-field object. + + Args: + mf: Mean-field object. + + Returns: + Expression object. + """ + h1e = mf.mo_coeff.T.conj() @ mf.get_hcore() @ mf.mo_coeff + h2e = ao2mo.kernel(mf._eri, mf.mo_coeff) # pylint: disable=protected-access + ci = fci.direct_spin1.FCI(mf.mol) + ci.verbose = 0 + ci.kernel(h1e, h2e, mf.mol.nao, mf.mol.nelec) + return cls.from_fci(ci, h1e, h2e) -def _fci_constructor(δalph, δbeta, func_sq, sign): - """Construct FCI expressions classes for a given change in the - number of alpha and beta electrons. - """ + def apply_hamiltonian(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector. - @util.inherit_docstrings - class _FCI(BaseExpression): + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. """ - FCI expressions. + nelec = (self.nocc + self.DELTA_ALPHA, self.nocc + self.DELTA_BETA) + result = fci.direct_spin1.contract_2e( + self.hamiltonian, + vector, + self.nphys, + nelec, + self.link_index, + ) + result -= (self.e_fci + self.chempot) * vector + return self.SIGN * result + + def diagonal(self) -> Array: + """Get the diagonal of the Hamiltonian. + + Returns: + Diagonal of the Hamiltonian. """ + return self.SIGN * (self._diagonal - (self.e_fci + self.chempot)) + + def get_excitation_vector(self, orbital: int) -> Array: + r"""Obtain the vector corresponding to a fermionic operator acting on the ground state. + + This vector is a generalisation of - hermitian = True - - def __init__( - self, - *args, - h1e=None, - h2e=None, - e_ci=None, - c_ci=None, - chempot=None, - nelec=None, - **kwargs, - ): - if len(args): - if nelec is not None: - raise ValueError( - "nelec keyword argument only valid when mean-field object is not " "passed." - ) - BaseExpression.__init__(self, *args, **kwargs) - else: - # Allow initialisation without MF object - if h1e is None or h2e is None: - raise ValueError( - "h1e and h2e keyword arguments are required to initialise FCI " - "without mean-field object." - ) - self.log = kwargs.get("log", default_log) - self.mf = None - self._nelec = nelec - self._nmo = h1e.shape[0] - - if c_ci is None: - if h1e is None: - h1e = np.linalg.multi_dot( - ( - self.mo_coeff.T, - self.mf.get_hcore(), - self.mo_coeff, - ) - ) - if h2e is None: - h2e = ao2mo.kernel(self.mf._eri, self.mo_coeff) - - ci = fci.direct_spin1.FCI() - ci.verbose = 0 - e_ci, c_ci = ci.kernel(h1e, h2e, self.nmo, (self.nalph, self.nbeta)) - - assert e_ci is not None - assert c_ci is not None - assert h1e is not None - assert h2e is not None - - self.e_ci = e_ci - self.c_ci = c_ci - self.chempot = e_ci if chempot is None else chempot - - self.link_index = ( - fci.cistring.gen_linkstr_index_trilidx(range(self.nmo), self.nalph + δalph), - fci.cistring.gen_linkstr_index_trilidx(range(self.nmo), self.nbeta + δbeta), - ) - - self.hamiltonian = fci.direct_spin1.absorb_h1e( - h1e, - h2e, - self.nmo, - (self.nalph + δalph, self.nbeta + δbeta), - 0.5, - ) - - self.diag = fci.direct_spin1.make_hdiag( - h1e, - h2e, - self.nmo, - (self.nalph + δalph, self.nbeta + δbeta), - ) - - def diagonal(self): - return sign * self.diag - - def apply_hamiltonian(self, vector): - hvec = fci.direct_spin1.contract_2e( - self.hamiltonian, - vector, - self.nmo, - (self.nalph + δalph, self.nbeta + δbeta), - self.link_index, - ) - hvec -= self.chempot * vector - - return sign * hvec.ravel() - - def get_wavefunction(self, orb): - wfn = func_sq(self.c_ci, self.nmo, (self.nalph, self.nbeta), orb) - return wfn.ravel() - - # Override properties to allow initialisation without a mean-field - - @property - def nmo(self): - if self.mf is None: - return self._nmo - return self.mo_coeff.shape[-1] - - @property - def nocc(self): - if self.mf is None: - return self._nelec // 2 - return np.sum(self.mo_occ > 0) - - @property - def nvir(self): - if self.mf is None: - return self.nmo - self.nocc - return np.sum(self.mo_occ == 0) - - return _FCI - - -FCI_1h = _fci_constructor(-1, 0, fci.addons.des_a, -1) - -FCI_1p = _fci_constructor(1, 0, fci.addons.cre_a, 1) - -FCI = { - "1h": FCI_1h, - "1p": FCI_1p, -} + .. math:: + f_i^{\pm} \left| \Psi_0 \right> + + where :math:`f_i^{\pm}` is the fermionic creation or annihilation operator, or a product + thereof, depending on the particular expression and what Green's function it corresponds to. + + The vector defines the excitaiton manifold probed by the Green's function corresponding to + the expression. + + Args: + orbital: Orbital index. + + Returns: + Excitation vector. + """ + return self.STATE_FUNC( + self.c_fci, + self.nphys, + (self.nocc, self.nocc), + orbital, + ).ravel() + + def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) -> Array: + """Build the self-energy moments. + + Args: + nmom: Number of moments to compute. + reduction: Reduction method to apply to the moments. + + Returns: + Moments of the self-energy. + """ + raise NotImplementedError("Self-energy moments not implemented for FCI.") + + @property + def mol(self) -> Mole: + """Molecule object.""" + return self._mol + + @property + def e_fci(self) -> Array: + """FCI eigenvalues.""" + return self._e_fci + + @property + def c_fci(self) -> Array: + """FCI eigenvectors.""" + return self._c_fci + + @property + def hamiltonian(self) -> Array: + """Hamiltonian matrix.""" + return self._hamiltonian + + @property + def chempot(self) -> Array | float: + """Chemical potential.""" + return self._chempot + + @functools.cached_property + def link_index(self) -> tuple[Array, Array]: + """Index helpers.""" + nelec = (self.nocc + self.DELTA_ALPHA, self.nocc + self.DELTA_BETA) + return ( + fci.cistring.gen_linkstr_index_trilidx(range(self.nphys), nelec[0]), + fci.cistring.gen_linkstr_index_trilidx(range(self.nphys), nelec[1]), + ) + + @property + def non_dyson(self) -> bool: + """Whether the expression produces a non-Dyson Green's function.""" + return False + + @property + def nconfig(self) -> int: + """Number of configurations.""" + link_index = self.link_index + return len(link_index[0]) * len(link_index[1]) - self.nsingle + + +class FCI_1h(BaseFCI): # pylint: disable=invalid-name + """FCI expressions for the hole Green's function.""" + + SIGN = -1 + DELTA_ALPHA = -1 + DELTA_BETA = 0 + STATE_FUNC = staticmethod(fci.addons.des_a) + + @property + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + return self.nocc + + +class FCI_1p(BaseFCI): # pylint: disable=invalid-name + """FCI expressions for the particle Green's function.""" + + SIGN = 1 + DELTA_ALPHA = 1 + DELTA_BETA = 0 + STATE_FUNC = staticmethod(fci.addons.cre_a) + + @property + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + return self.nvir + + +class FCI(ExpressionCollection): + """Collection of FCI expressions for different parts of the Green's function.""" + + _hole = FCI_1h + _particle = FCI_1p + _name = "FCI" diff --git a/dyson/expressions/gw.py b/dyson/expressions/gw.py index c65bf87..e23abf9 100644 --- a/dyson/expressions/gw.py +++ b/dyson/expressions/gw.py @@ -1,155 +1,252 @@ -""" -GW expressions. -""" - -import numpy as np -from pyscf import agf2, ao2mo, lib - -from dyson import util -from dyson.expressions import BaseExpression - - -@util.inherit_docstrings -class GW_Dyson(BaseExpression): - """ - GW expressions without a non-Dyson approximation. - """ - - hermitian = True - polarizability = "drpa" - - def __init__(self, *args, **kwargs): - BaseExpression.__init__(self, *args, **kwargs) - - try: - from momentGW import GW - - self._gw = GW(self.mf) - self._gw.mo_occ = self.mo_occ - self._gw.mo_coeff = self.mo_coeff - self._gw.mo_energy = self.mo_energy - self._gw.polarizability = self.polarizability - except ImportError as e: - raise ImportError("momentGW is required for GW expressions.") - - def get_static_part(self): - static = self._gw.build_se_static(self._gw.ao2mo()) - - return static - - def apply_hamiltonian(self, vector, static=None): - # From Bintrim & Berkelbach - - if static is None: - static = self.get_static_part() - - i = slice(None, self.nocc) - a = slice(self.nocc, self.nmo) - ija = slice(self.nmo, self.nmo + self.nocc * self.nocc * self.nvir) - iab = slice(self.nmo + self.nocc * self.nocc * self.nvir, None) +"""GW approximation expressions [1]_ [2]_ [3]_. - Lpq = self._gw.ao2mo(self.mo_coeff) - Lia = Lpq[:, i, a] - Lai = Lpq[:, a, i] - Lij = Lpq[:, i, i] - Lab = Lpq[:, a, a] +.. [1] Hedin, L. (1965). New Method for Calculating the One-Particle Green’s Function with + Application to the Electron-Gas Problem. Physical Review, 139(3A), A796–A823. + https://doi.org/10.1103/physrev.139.a796 - nocc, nvir = Lia.shape[1:] +.. [2] Aryasetiawan, F., & Gunnarsson, O. (1998). The GW method. Reports on Progress + in Physics, 61(3), 237–312. https://doi.org/10.1088/0034-4885/61/3/002 - vi = vector[i] - va = vector[a] - vija = vector[ija].reshape(nocc, nocc, nvir) - viab = vector[iab].reshape(nocc, nvir, nvir) +.. [3] Zhu, T., & Chan, G. K. (2021). All-Electron Gaussian-Based G0W0 for valence and core + excitation energies of periodic systems. Journal of Chemical Theory and Computation, 17(2), + 727–741. https://doi.org/10.1021/acs.jctc.0c00704 - eija = lib.direct_sum("i+j-a->ija", self.mo_energy[i], self.mo_energy[i], self.mo_energy[a]) - eiab = lib.direct_sum("a+b-i->iab", self.mo_energy[a], self.mo_energy[a], self.mo_energy[i]) - - r = np.zeros_like(vector) - - if self.polarizability == "dtda": - r[i] += lib.einsum("ij,j->i", static[i, i], vi) - r[i] += lib.einsum("ib,b->i", static[i, a], va) - r[i] += lib.einsum("Qik,Qcl,klc->i", Lij, Lai, vija) - r[i] += lib.einsum("Qid,Qkc,kcd->i", Lia, Lia, viab) - - r[a] += lib.einsum("aj,j->a", static[a, i], vi) - r[a] += lib.einsum("ab,b->a", static[a, a], va) - r[a] += lib.einsum("Qak,Qcl,klc->a", Lai, Lai, vija) - r[a] += lib.einsum("Qad,Qkc,kcd->a", Lab, Lia, viab) - - r[ija] += lib.einsum("Qki,Qaj,k->ija", Lij, Lai, vi).ravel() - r[ija] += lib.einsum("Qbi,Qaj,b->ija", Lai, Lai, va).ravel() - r[ija] += lib.einsum("ija,ija->ija", eija, vija).ravel() - r[ija] -= lib.einsum("Qja,Qcl,ilc->ija", Lia, Lai, vija).ravel() - - r[iab] += lib.einsum("Qjb,Qia,j->iab", Lia, Lia, vi).ravel() - r[iab] += lib.einsum("Qcb,Qia,c->iab", Lab, Lia, va).ravel() - r[iab] += lib.einsum("iab,iab->iab", eiab, viab).ravel() - r[iab] += lib.einsum("Qai,Qkc,kcb->iab", Lai, Lia, viab).ravel() - - elif self.polarizability == "drpa": - raise NotImplementedError - - return r - - def diagonal(self, static=None): - # From Bintrim & Berkelbach - - if static is None: - static = self.get_static_part() - - i = slice(None, self.nocc) - a = slice(self.nocc, self.nmo) - ija = slice(self.nmo, self.nmo + self.nocc * self.nocc * self.nvir) - iab = slice(self.nmo + self.nocc * self.nocc * self.nvir, None) - - integrals = self._gw.ao2mo() - Lpq = integrals.Lpx - Lia = integrals.Lia - Lia = Lpq[:, i, a] - Lai = Lpq[:, a, i] - Lij = Lpq[:, i, i] - Lab = Lpq[:, a, a] - - nocc, nvir = Lia.shape[1:] - - eija = lib.direct_sum("i+j-a->ija", self.mo_energy[i], self.mo_energy[i], self.mo_energy[a]) - eiab = lib.direct_sum("a+b-i->iab", self.mo_energy[a], self.mo_energy[a], self.mo_energy[i]) +""" - diag = np.zeros((self.nmo + eija.size + eiab.size,)) +from __future__ import annotations - if self.polarizability == "dtda": - diag[i] += np.diag(static[i, i]) +from typing import TYPE_CHECKING - diag[a] += np.diag(static[a, a]) +from pyscf import gw, lib - diag[ija] += eija.ravel() - diag[ija] -= lib.einsum("Qja,Qaj,ii->ija", Lia, Lai, np.eye(nocc)).ravel() +from dyson import numpy as np +from dyson import util +from dyson.expressions.expression import BaseExpression, ExpressionCollection +from dyson.representations.enums import Reduction - diag[iab] += eiab.ravel() - diag[iab] += lib.einsum("Qai,Qia,bb->iab", Lai, Lia, np.eye(nvir)).ravel() +if TYPE_CHECKING: + from pyscf.gto.mole import Mole + from pyscf.scf.hf import RHF - elif self.polarizability == "drpa": - raise NotImplementedError + from dyson.typing import Array - return diag - def get_wavefunction(self, orb): - nija = self.nocc * self.nocc * self.nvir - nabi = self.nocc * self.nvir * self.nvir +class BaseGW_Dyson(BaseExpression): + """Base class for GW expressions for the Dyson Green's function.""" - r = np.zeros((self.nmo + nija + nabi,)) - r[orb] = 1.0 + hermitian_downfolded = True + hermitian_upfolded = False - return r + def __init__( + self, + mol: Mole, + gw_obj: gw.GW, + eris: Array | None = None, + ) -> None: + """Initialise the expression. + + Args: + mol: Molecule object. + gw_obj: GW object from PySCF. + eris: Density fitted electron repulsion integrals from PySCF. + """ + self._mol = mol + self._gw = gw_obj + self._eris = eris if eris is not None else gw_obj.ao2mo() - def build_se_moments(self, nmom): - integrals = self._gw.ao2mo() - moments = self._gw.build_se_moments(nmom, integrals) + if getattr(self._gw._scf, "xc", "hf") != "hf": + raise NotImplementedError( + "GW expressions currently only support Hartree--Fock mean-field objects." + ) - return moments + @classmethod + def from_mf(cls, mf: RHF) -> BaseGW_Dyson: + """Create an expression from a mean-field object. + Args: + mf: Mean-field object. -GW = { - "Dyson": GW_Dyson, -} + Returns: + Expression object. + """ + return cls.from_gw(gw.GW(mf)) + + @classmethod + def from_gw(cls, gw: gw.GW) -> BaseGW_Dyson: + """Create an expression from a GW object. + + Args: + gw: GW object. + + Returns: + Expression object. + """ + return cls(gw._scf.mol, gw) + + def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) -> Array: + """Build the self-energy moments. + + Args: + nmom: Number of moments to compute. + reduction: Reduction method to apply to the moments. + + Returns: + Moments of the self-energy. + """ + raise NotImplementedError("Self-energy moments not implemented for GW.") + + def get_excitation_vector(self, orbital: int) -> Array: + r"""Obtain the vector corresponding to a fermionic operator acting on the ground state. + + This vector is a generalisation of + + .. math:: + f_i^{\pm} \left| \Psi_0 \right> + + where :math:`f_i^{\pm}` is the fermionic creation or annihilation operator, or a product + thereof, depending on the particular expression and what Green's function it corresponds to. + + The vector defines the excitaiton manifold probed by the Green's function corresponding to + the expression. + + Args: + orbital: Orbital index. + + Returns: + Excitation vector. + """ + return util.unit_vector(self.shape[0], orbital) + + @property + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + return self.nocc + self.nvir + + @property + def nconfig(self) -> int: + """Number of configurations.""" + return self.nocc * self.nocc * self.nvir + self.nvir * self.nvir * self.nocc + + @property + def mol(self) -> Mole: + """Molecule object.""" + return self._mol + + @property + def gw(self) -> gw.GW: + """GW object.""" + return self._gw + + @property + def eris(self) -> Array: + """Density fitted electron repulsion integrals.""" + return self._eris + + @property + def non_dyson(self) -> bool: + """Whether the expression produces a non-Dyson Green's function.""" + return False + + +class TDAGW_Dyson(BaseGW_Dyson): + """GW expressions with Tamm--Dancoff (TDA) approximation for the Dyson Green's function.""" + + def apply_hamiltonian(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector. + + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. + """ + # Get the slices for each sector + o1 = slice(None, self.nocc) + v1 = slice(self.nocc, self.nocc + self.nvir) + o2 = slice(self.nocc + self.nvir, self.nocc + self.nvir + self.nocc * self.nocc * self.nvir) + v2 = slice(self.nocc + self.nvir + self.nocc * self.nocc * self.nvir, None) + + # Get the blocks of the ERIs + Lia = self.eris[:, o1, v1] + Lai = self.eris[:, v1, o1] + Lij = self.eris[:, o1, o1] + Lab = self.eris[:, v1, v1] + + # Get the blocks of the vector + vector_o1 = vector[o1] + vector_v1 = vector[v1] + vector_o2 = vector[o2].reshape(self.nocc, self.nocc, self.nvir) + vector_v2 = vector[v2].reshape(self.nocc, self.nvir, self.nvir) + + # Get the energy denominators + mo_energy = self.gw._scf.mo_energy if self.gw.mo_energy is None else self.gw.mo_energy + e_ija = lib.direct_sum("i+j-a->ija", mo_energy[o1], mo_energy[o1], mo_energy[v1]) + e_iab = lib.direct_sum("a+b-i->iab", mo_energy[v1], mo_energy[v1], mo_energy[o1]) + + # Perform the contractions + r_o1 = mo_energy[o1] * vector_o1 + r_o1 += util.einsum("Qik,Qcl,klc->i", Lij, Lai, vector_o2) * 2 + r_o1 += util.einsum("Qid,Qkc,kcd->i", Lia.conj(), Lia.conj(), vector_v2) * 2 + + r_v1 = mo_energy[v1] * vector_v1 + r_v1 += util.einsum("Qak,Qcl,klc->a", Lai, Lai, vector_o2) * 2 + r_v1 += util.einsum("Qad,Qkc,kcd->a", Lab.conj(), Lia.conj(), vector_v2) * 2 + + r_o2 = util.einsum("Qki,Qaj,k->ija", Lij.conj(), Lai.conj(), vector_o1) + r_o2 += util.einsum("Qbi,Qaj,b->ija", Lai.conj(), Lai.conj(), vector_v1) + r_o2 += util.einsum("ija,ija->ija", e_ija, vector_o2) + r_o2 -= util.einsum("Qja,Qlc,ilc->ija", Lia, Lia, vector_o2) * 2 + + r_v2 = util.einsum("Qjb,Qia,j->iab", Lia, Lia, vector_o1) + r_v2 += util.einsum("Qcb,Qia,c->iab", Lab, Lia, vector_v1) + r_v2 += util.einsum("iab,iab->iab", e_iab, vector_v2) + r_v2 += util.einsum("Qia,Qkc,kcb->iab", Lia, Lia, vector_v2) * 2 + + return np.concatenate([r_o1, r_v1, r_o2.ravel(), r_v2.ravel()]) + + def apply_hamiltonian_left(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector on the left. + + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. + """ + raise NotImplementedError("Left application of Hamiltonian is not implemented for TDA-GW.") + + def diagonal(self) -> Array: + """Get the diagonal of the Hamiltonian. + + Returns: + Diagonal of the Hamiltonian. + """ + # Get the slices for each sector + o1 = slice(None, self.nocc) + v1 = slice(self.nocc, None) + + # Get the blocks of the ERIs + Lia = self.eris[:, o1, v1] + Lai = self.eris[:, v1, o1] + + # Get the energy denominators + mo_energy = self.gw._scf.mo_energy if self.gw.mo_energy is None else self.gw.mo_energy + e_ija = lib.direct_sum("i+j-a->ija", mo_energy[o1], mo_energy[o1], mo_energy[v1]) + e_iab = lib.direct_sum("a+b-i->iab", mo_energy[v1], mo_energy[v1], mo_energy[o1]) + + # Build the diagonal + diag_o1 = mo_energy[o1].copy() + diag_v1 = mo_energy[v1].copy() + diag_o2 = e_ija.ravel() + diag_o2 -= util.einsum("Qja,Qaj,ii->ija", Lia, Lai, np.eye(self.nocc)).ravel() + diag_v2 = e_iab.ravel() + diag_v2 += util.einsum("Qai,Qia,bb->iab", Lai, Lia, np.eye(self.nvir)).ravel() + + return np.concatenate([diag_o1, diag_v1, diag_o2, diag_v2]) + + +class TDAGW(ExpressionCollection): + """Collection of TDAGW expressions for different parts of the Green's function.""" + + _dyson = TDAGW_Dyson + _name = "TDA-GW" diff --git a/dyson/expressions/hamiltonian.py b/dyson/expressions/hamiltonian.py new file mode 100644 index 0000000..a43c91c --- /dev/null +++ b/dyson/expressions/hamiltonian.py @@ -0,0 +1,203 @@ +"""Hamiltonian-driven expressions.""" + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING, cast + +from dyson import numpy as np +from dyson import util +from dyson.expressions.expression import BaseExpression +from dyson.representations.enums import Reduction + +if TYPE_CHECKING: + from pyscf.gto.mole import Mole + from pyscf.scf.hf import RHF + from scipy.sparse import spmatrix as SparseArray + + from dyson.typing import Array + + +class Hamiltonian(BaseExpression): + """Hamiltonian-driven expressions for the Green's function.""" + + def __init__( + self, + hamiltonian: Array | SparseArray, + bra: Array | None = None, + ket: Array | None = None, + ): + """Initialise the expression. + + Args: + hamiltonian: Hamiltonian matrix. + bra: Bra excitation vector. If not passed, a unit vectors are used. See + :meth:`~dyson.expressions.expression.BaseExpression.get_excitation_bra`. + ket: Ket excitation vector. If not passed, ``bra`` is used. See + :meth:`~dyson.expressions.expression.BaseExpression.get_excitation_ket`. + """ + self._hamiltonian = hamiltonian + self._bra = bra + self._ket = ket + + if isinstance(hamiltonian, np.ndarray): + self.hermitian_upfolded = np.allclose(hamiltonian, hamiltonian.T.conj()) + else: + self.hermitian_upfolded = (hamiltonian - hamiltonian.H).nnz == 0 + self.hermitian_downfolded = self.hermitian_upfolded and ket is None + + @classmethod + def from_mf(cls, mf: RHF) -> Hamiltonian: + """Create an expression from a mean-field object. + + Args: + mf: Mean-field object. + + Returns: + Expression object. + """ + raise NotImplementedError("Cannot create Hamiltonian expression from mean-field object.") + + def apply_hamiltonian(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector. + + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. + """ + return self._hamiltonian @ vector + + def apply_hamiltonian_left(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector on the left. + + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. + """ + return vector @ self._hamiltonian + + def diagonal(self) -> Array: + """Get the diagonal of the Hamiltonian. + + Returns: + Diagonal of the Hamiltonian. + """ + return self._hamiltonian.diagonal() + + def build_matrix(self) -> Array: + """Build the Hamiltonian matrix. + + Returns: + Hamiltonian matrix. + """ + if isinstance(self._hamiltonian, np.ndarray): + return self._hamiltonian + else: + size = self.diagonal().size + if size > 2048: + warnings.warn( + "The Hamiltonian matrix is large. This may take a while to compute.", + UserWarning, + 2, + ) + return self._hamiltonian.todense() + + def get_excitation_bra(self, orbital: int) -> Array: + r"""Obtain the bra vector corresponding to a fermionic operator acting on the ground state. + + The bra vector is the excitation vector corresponding to the bra state, which may or may not + be the same as the ket state vector. + + Args: + orbital: Orbital index. + + Returns: + Bra excitation vector. + + See Also: + :func:`get_excitation_vector`: Function to get the excitation vector when the bra and + ket are the same. + """ + if self._bra is None: + return util.unit_vector(self.shape[0], orbital) + return self._bra[orbital] + + def get_excitation_ket(self, orbital: int) -> Array: + r"""Obtain the ket vector corresponding to a fermionic operator acting on the ground state. + + The ket vector is the excitation vector corresponding to the ket state, which may or may not + be the same as the bra state vector. + + Args: + orbital: Orbital index. + + Returns: + Ket excitation vector. + + See Also: + :func:`get_excitation_vector`: Function to get the excitation vector when the bra and + ket are the same. + """ + if self._ket is None: + return self.get_excitation_bra(orbital) + return self._ket[orbital] + + get_excitation_vector = get_excitation_ket + get_excitation_vector.__doc__ = BaseExpression.get_excitation_vector.__doc__ + + def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) -> Array: + """Build the self-energy moments. + + Args: + nmom: Number of moments to compute. + reduction: Reduction method to apply to the moments. + + Returns: + Moments of the self-energy. + """ + raise NotImplementedError("Self-energy moments not implemented for Hamiltonian.") + + @property + def mol(self) -> Mole: + """Molecule object.""" + raise NotImplementedError("Molecule object not available for Hamiltonian expression.") + + @property + def non_dyson(self) -> bool: + """Whether the expression produces a non-Dyson Green's function.""" + return False + + @property + def nphys(self) -> int: + """Number of physical orbitals.""" + return self._bra.shape[0] if self._bra is not None else 1 + + @property + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + raise NotImplementedError("Excitation sectors not implemented for Hamiltonian.") + + @property + def nconfig(self) -> int: + """Number of configurations in the non-singles sectors.""" + raise NotImplementedError("Excitation sectors not implemented for Hamiltonian.") + + @property + def shape(self) -> tuple[int, int]: + """Shape of the Hamiltonian matrix.""" + assert self._hamiltonian.ndim == 2 + return cast(tuple[int, int], self._hamiltonian.shape) + + @property + def nocc(self) -> int: + """Number of occupied orbitals.""" + raise NotImplementedError("Orbital occupancy not defined for Hamiltonian.") + + @property + def nvir(self) -> int: + """Number of virtual orbitals.""" + raise NotImplementedError("Orbital occupancy not defined for Hamiltonian.") diff --git a/dyson/expressions/hf.py b/dyson/expressions/hf.py new file mode 100644 index 0000000..f212ba3 --- /dev/null +++ b/dyson/expressions/hf.py @@ -0,0 +1,241 @@ +"""Hartree--Fock (HF) expressions [1]_. + +.. [1] Slater, J. C. (1928). The self consistent field and the structure of atoms. Physical + Review, 32(3), 339–348. https://doi.org/10.1103/physrev.32.339 +""" + +from __future__ import annotations + +from abc import abstractmethod +from typing import TYPE_CHECKING + +from dyson import numpy as np +from dyson import util +from dyson.expressions.expression import BaseExpression, ExpressionCollection +from dyson.representations.enums import Reduction + +if TYPE_CHECKING: + from pyscf.gto.mole import Mole + from pyscf.scf.hf import RHF + + from dyson.typing import Array + + +class BaseHF(BaseExpression): + """Base class for HF expressions.""" + + hermitian_downfolded = True + hermitian_upfolded = True + + def __init__( + self, + mol: Mole, + mo_energy: Array, + ): + """Initialise the expression. + + Args: + mol: Molecule object. + mo_energy: Molecular orbital energies. + """ + self._mol = mol + self._mo_energy = mo_energy + + @classmethod + def from_mf(cls, mf: RHF) -> BaseHF: + """Create an expression from a mean-field object. + + Args: + mf: Mean-field object. + + Returns: + Expression object. + """ + return cls(mf.mol, mf.mo_energy) + + def apply_hamiltonian(self, vector: Array) -> Array: + """Apply the Hamiltonian to a vector. + + Args: + vector: Vector to apply Hamiltonian to. + + Returns: + Output vector. + """ + return self.diagonal() * vector + + @abstractmethod + def diagonal(self) -> Array: + """Get the diagonal of the Hamiltonian. + + Returns: + Diagonal of the Hamiltonian. + """ + pass + + def build_se_moments(self, nmom: int, reduction: Reduction = Reduction.NONE) -> Array: + """Build the self-energy moments. + + Args: + nmom: Number of moments. + reduction: Reduction method to apply to the moments. + + Returns: + Self-energy moments. + """ + return np.zeros((nmom,) + (self.nphys,) * Reduction(reduction).ndim) + + @property + def mol(self) -> Mole: + """Molecule object.""" + return self._mol + + @property + def mo_energy(self) -> Array: + """Molecular orbital energies.""" + return self._mo_energy + + @property + def non_dyson(self) -> bool: + """Whether the expression produces a non-Dyson Green's function.""" + return True + + @property + def nconfig(self) -> int: + """Number of configurations.""" + return 0 + + +class HF_1h(BaseHF): # pylint: disable=invalid-name + """HF expressions for the hole Green's function.""" + + def diagonal(self) -> Array: + """Get the diagonal of the Hamiltonian. + + Returns: + Diagonal of the Hamiltonian. + """ + return self.mo_energy[: self.nocc] + + def get_excitation_vector(self, orbital: int) -> Array: + r"""Obtain the vector corresponding to a fermionic operator acting on the ground state. + + This vector is a generalisation of + + .. math:: + f_i^{\pm} \left| \Psi_0 \right> + + where :math:`f_i^{\pm}` is the fermionic creation or annihilation operator, or a product + thereof, depending on the particular expression and what Green's function it corresponds to. + + The vector defines the excitaiton manifold probed by the Green's function corresponding to + the expression. + + Args: + orbital: Orbital index. + + Returns: + Excitation vector. + """ + if orbital < self.nocc: + return util.unit_vector(self.shape[0], orbital) + return np.zeros(self.shape[0]) + + @property + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + return self.nocc + + +class HF_1p(BaseHF): # pylint: disable=invalid-name + """HF expressions for the particle Green's function.""" + + def diagonal(self) -> Array: + """Get the diagonal of the Hamiltonian. + + Returns: + Diagonal of the Hamiltonian. + """ + return self.mo_energy[self.nocc :] + + def get_excitation_vector(self, orbital: int) -> Array: + r"""Obtain the vector corresponding to a fermionic operator acting on the ground state. + + This vector is a generalisation of + + .. math:: + f_i^{\pm} \left| \Psi_0 \right> + + where :math:`f_i^{\pm}` is the fermionic creation or annihilation operator, or a product + thereof, depending on the particular expression and what Green's function it corresponds to. + + The vector defines the excitaiton manifold probed by the Green's function corresponding to + the expression. + + Args: + orbital: Orbital index. + + Returns: + Excitation vector. + """ + if orbital >= self.nocc: + return util.unit_vector(self.shape[0], orbital - self.nocc) + return np.zeros(self.shape[0]) + + @property + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + return self.nvir + + +class HF_Dyson(BaseHF): # pylint: disable=invalid-name + """HF expressions for the Dyson Green's function.""" + + def diagonal(self) -> Array: + """Get the diagonal of the Hamiltonian. + + Returns: + Diagonal of the Hamiltonian. + """ + return self.mo_energy + + def get_excitation_vector(self, orbital: int) -> Array: + r"""Obtain the vector corresponding to a fermionic operator acting on the ground state. + + This vector is a generalisation of + + .. math:: + f_i^{\pm} \left| \Psi_0 \right> + + where :math:`f_i^{\pm}` is the fermionic creation or annihilation operator, or a product + thereof, depending on the particular expression and what Green's function it corresponds to. + + The vector defines the excitaiton manifold probed by the Green's function corresponding to + the expression. + + Args: + orbital: Orbital index. + + Returns: + Excitation vector. + """ + return util.unit_vector(self.shape[0], orbital) + + @property + def nsingle(self) -> int: + """Number of configurations in the singles sector.""" + return self.nocc + self.nvir + + @property + def non_dyson(self) -> bool: + """Whether the expression produces a non-Dyson Green's function.""" + return False + + +class HF(ExpressionCollection): + """Collection of HF expressions for different parts of the Green's function.""" + + _hole = HF_1h + _particle = HF_1p + _dyson = HF_Dyson + _name = "HF" diff --git a/dyson/expressions/mp2.py b/dyson/expressions/mp2.py deleted file mode 100644 index f163b1c..0000000 --- a/dyson/expressions/mp2.py +++ /dev/null @@ -1,206 +0,0 @@ -""" -MP2 expressions. -""" - -import numpy as np -from pyscf import agf2, ao2mo, lib - -from dyson import util -from dyson.expressions import BaseExpression - -# TODO only separate for non-Dyson - - -def _mp2_constructor(occ, vir): - """Construct MP2 expressions classes for a given occupied and - virtual mask. These classes use a non-Dyson approximation. - """ - - @util.inherit_docstrings - class _MP2(BaseExpression): - """ - MP2 expressions with non-Dyson approximation. - """ - - hermitian = False - - def __init__(self, *args, **kwargs): - BaseExpression.__init__(self, *args, **kwargs) - - self.ijka = self._integrals_for_hamiltonian() - - def get_static_part(self): - co = self.mo_coeff[:, occ(self)] - cv = self.mo_coeff[:, vir(self)] - eo = self.mo_energy[occ(self)] - ev = self.mo_energy[vir(self)] - - e_iajb = lib.direct_sum("i-a+j-b->iajb", eo, ev, eo, ev) - - iajb = ao2mo.incore.general(self.mf._eri, (co, cv, co, cv), compact=False) - iajb = iajb.reshape([x.shape[-1] for x in (co, cv, co, cv)]) - t2 = iajb / e_iajb - iajb = 2 * iajb - iajb.transpose(0, 3, 2, 1) - - h1 = lib.einsum("iakb,jakb->ij", iajb, t2) * 0.5 - h1 += h1.T - h1 += np.diag(eo) # FIXME or C* F C? - - return h1 - - def _integrals_for_hamiltonian(self): - c = self.mo_coeff[:, occ(self)] - e = self.mo_energy[occ(self)] - p = slice(None, e.size) - a = slice(e.size, None) - - co = self.mo_coeff[:, occ(self)] - cv = self.mo_coeff[:, vir(self)] - - ijka = ao2mo.incore.general(self.mf._eri, (co, co, co, cv), compact=False) - ijka = ijka.reshape([x.shape[-1] for x in (co, co, co, cv)]) - - return ijka - - def apply_hamiltonian(self, vector, static=None): - if static is None: - static = self.get_static_part() - - e = self.mo_energy[occ(self)] - p = slice(None, e.size) - a = slice(e.size, None) - - eo = self.mo_energy[occ(self)] - ev = self.mo_energy[vir(self)] - e_jka = lib.direct_sum("j+k-a->jka", eo, eo, ev) - ijka = self.ijka - - r = np.zeros_like(vector) - r[p] += np.dot(static, vector[p]) - r[p] += lib.einsum("ijka,jka->i", ijka, vector[a].reshape(e_jka.shape)) - r[a] += lib.einsum("ijka,i->jka", ijka, vector[p]).ravel() * 2.0 - r[a] -= lib.einsum("ikja,i->jka", ijka, vector[p]).ravel() - r[a] += vector[a] * e_jka.ravel() - - return r - - def diagonal(self, static=None): - if static is None: - static = self.get_static_part() - - eo = self.mo_energy[occ(self)] - ev = self.mo_energy[vir(self)] - e_ija = lib.direct_sum("i+j-a->ija", eo, eo, ev) - - r = np.concatenate([np.diag(static), e_ija.ravel()]) - - return r - - def get_wavefunction(self, orb): - nocc = np.sum(occ(self)) - nvir = np.sum(vir(self)) - nija = nocc * nocc * nvir - - r = np.zeros((nocc + nija,)) - r[orb] = 1.0 - - return r - - def build_se_moments(self, nmom): - eo = self.mo_energy[occ(self)] - ev = self.mo_energy[vir(self)] - ijka = self.ijka - - t = [] - for n in range(nmom): - tn = 0 - for j in range(eo.size): - vl = ijka[:, j] - vr = 2.0 * ijka[:, j] - ijka[:, :, j] - eka = eo[j] + lib.direct_sum("k-a->ka", eo, ev) - tn += lib.einsum("ika,jka,ka->ij", vl, vr, eka**n) - t.append(tn) - - return np.array(t) - - @property - def nphys(self): - return np.sum(occ(self)) - - return _MP2 - - -@util.inherit_docstrings -class MP2_Dyson(BaseExpression): - """ - MP2 expressions without non-Dyson approximation. - """ - - def __init__(self, *args, **kwargs): - BaseExpression.__init__(self, *args, **kwargs) - - self._agf2 = agf2.ragf2_slow.RAGF2( - self.mf, - mo_energy=self.mo_energy, - mo_coeff=self.mo_coeff, - mo_occ=self.mo_occ, - ) - - def get_static_part(self): - raise NotImplementedError # TODO - - def apply_hamiltonian(self, vector, static=None): - raise NotImplementedError # TODO - - def diagonal(self, static=None): - raise NotImplementedError # TODO - - def get_wavefunction(self, orb): - nija = self.nocc * self.nocc * self.nvir - niab = self.nocc * self.nvir * self.nvir - - r = np.zeros((self.nphys + nija + niab,)) - r[orb] = 1.0 - - return r - - def build_se_moments(self, nmom): - eo = self.mo_energy[: self.nocc] - ev = self.mo_energy[self.nocc :] - c = self.mo_coeff - co = self.mo_coeff[:, : self.nocc] - cv = self.mo_coeff[:, self.nocc :] - - xija = ao2mo.incore.general(self.mf._eri, (c, co, co, cv), compact=False) - xija = xija.reshape([x.shape[-1] for x in (c, co, co, cv)]) - - xabi = ao2mo.incore.general(self.mf._eri, (c, cv, cv, co), compact=False) - xabi = xabi.reshape([x.shape[-1] for x in (c, cv, cv, co)]) - - th = np.zeros((nmom, self.nmo, self.nmo)) - tp = np.zeros((nmom, self.nmo, self.nmo)) - for n in range(nmom): - for i in range(eo.size): - vl = xija[:, i] - vr = 2.0 * xija[:, i] - xija[:, :, i] - eja = eo[i] + lib.direct_sum("j-a->ja", eo, ev) - th[n] += lib.einsum("xja,yja,ja->xy", vl, vr, eja**n) - - for a in range(ev.size): - vl = xabi[:, a] - vr = 2.0 * xabi[:, a] - xabi[:, :, a] - ebi = ev[a] + lib.direct_sum("b-i->bi", ev, eo) - tp[n] += lib.einsum("xbi,ybi,bi->xy", vl, vr, ebi**n) - - return th, tp - - -MP2_1h = _mp2_constructor(lambda self: self.mo_occ > 0, lambda self: self.mo_occ == 0) - -MP2_1p = _mp2_constructor(lambda self: self.mo_occ == 0, lambda self: self.mo_occ > 0) - -MP2 = { - "Dyson": MP2_Dyson, - "1h": MP2_1h, - "1p": MP2_1p, -} diff --git a/dyson/grids/__init__.py b/dyson/grids/__init__.py new file mode 100644 index 0000000..887cf98 --- /dev/null +++ b/dyson/grids/__init__.py @@ -0,0 +1,17 @@ +r"""Grids for Green's functions and self-energies. + +Grids are arrays of points in either the frequency or time domain. + + +Submodules +---------- + +.. autosummary:: + :toctree: + + grid + frequency +""" + +from dyson.grids.frequency import RealFrequencyGrid, GridRF +from dyson.grids.frequency import ImaginaryFrequencyGrid, GridIF diff --git a/dyson/grids/frequency.py b/dyson/grids/frequency.py new file mode 100644 index 0000000..fc13f28 --- /dev/null +++ b/dyson/grids/frequency.py @@ -0,0 +1,318 @@ +"""Frequency grids.""" + +from __future__ import annotations + +from abc import abstractmethod +from typing import TYPE_CHECKING + +import scipy.special + +from dyson import numpy as np +from dyson import util +from dyson.grids.grid import BaseGrid +from dyson.representations.enums import Component, Ordering, Reduction + +if TYPE_CHECKING: + from typing import Any + + from dyson.representations.dynamic import Dynamic + from dyson.representations.lehmann import Lehmann + from dyson.typing import Array + + +class BaseFrequencyGrid(BaseGrid): + """Base class for frequency grids.""" + + def evaluate_lehmann( + self, + lehmann: Lehmann, + reduction: Reduction = Reduction.NONE, + component: Component = Component.FULL, + **kwargs: Any, + ) -> Dynamic[BaseFrequencyGrid]: + r"""Evaluate a Lehmann representation on the grid. + + The imaginary frequency representation is defined as + + .. math:: + \sum_{k} \frac{v_{pk} u_{qk}^*}{i \omega - \epsilon_k}, + + and the real frequency representation is defined as + + .. math:: + \sum_{k} \frac{v_{pk} u_{qk}^*}{\omega - \epsilon_k \pm i \eta}, + + where :math:`\omega` is the frequency grid, :math:`\epsilon_k` are the poles, and the sign + of the broadening factor is determined by the time ordering. + + Args: + lehmann: Lehmann representation to evaluate. + reduction: The reduction of the dynamic representation. + component: The component of the dynamic representation. + kwargs: Additional keyword arguments for the resolvent. + + Returns: + Lehmann representation, realised on the grid. + """ + from dyson.representations.dynamic import Dynamic # noqa: PLC0415 + + left, right = lehmann.unpack_couplings() + resolvent = self.resolvent(lehmann.energies, lehmann.chempot, **kwargs) + reduction = Reduction(reduction) + component = Component(component) + + # Get the input and output indices based on the reduction type + inp = "qk" + out = "wpq" + if reduction == reduction.NONE: + pass + elif reduction == reduction.DIAG: + inp = "pk" + out = "wp" + elif reduction == reduction.TRACE: + inp = "pk" + out = "w" + else: + reduction.raise_invalid_representation() + + # Perform the downfolding operation + array = util.einsum(f"pk,{inp},wk->{out}", right, left.conj(), resolvent) + + # Get the required component + # TODO: Save time by not evaluating the full array when not needed + if component == Component.REAL: + array = array.real + elif component == Component.IMAG: + array = array.imag + + return Dynamic( + self, array, reduction=reduction, component=component, hermitian=lehmann.hermitian + ) + + @property + def domain(self) -> str: + """Get the domain of the grid. + + Returns: + Domain of the grid. + """ + return "frequency" + + @abstractmethod + def resolvent( # noqa: D417 + self, energies: Array, chempot: float | Array, **kwargs: Any + ) -> Array: + """Get the resolvent of the grid. + + Args: + energies: Energies of the poles. + chempot: Chemical potential. + + Returns: + Resolvent of the grid. + """ + pass + + +class RealFrequencyGrid(BaseFrequencyGrid): + """Real frequency grid.""" + + eta: float = 1e-2 + + _options = {"eta"} + + def __init__( # noqa: D417 + self, points: Array, weights: Array | None = None, **kwargs: Any + ) -> None: + """Initialise the grid. + + Args: + points: Points of the grid. + weights: Weights of the grid. + eta: Broadening factor. + """ + self._points = np.asarray(points) + self._weights = np.asarray(weights) if weights is not None else None + self.set_options(**kwargs) + + @property + def reality(self) -> bool: + """Get the reality of the grid. + + Returns: + Reality of the grid. + """ + return True + + @staticmethod + def _resolvent_signs(energies: Array, ordering: Ordering) -> Array: + """Get the signs for the resolvent based on the time ordering.""" + ordering = Ordering(ordering) + signs: Array + if ordering == ordering.ORDERED: + signs = np.where(energies >= 0, 1.0, -1.0) + elif ordering == ordering.ADVANCED: + signs = -np.ones_like(energies) + elif ordering == ordering.RETARDED: + signs = np.ones_like(energies) + else: + ordering.raise_invalid_representation() + return signs + + def resolvent( # noqa: D417 + self, + energies: Array, + chempot: float | Array, + ordering: Ordering = Ordering.ORDERED, + invert: bool = True, + **kwargs: Any, + ) -> Array: + r"""Get the resolvent of the grid. + + For real frequency grids, the resolvent is given by + + .. math:: + R(\omega) = \frac{1}{\omega - E \pm i \eta}, + + where :math:`\eta` is a small broadening factor, and :math:`E` are the pole energies. The + sign of :math:`i \eta` depends on the time ordering of the resolvent. + + Args: + energies: Energies of the poles. + chempot: Chemical potential. + ordering: Time ordering of the resolvent. + invert: Whether to apply the inversion in the resolvent formula. + + Returns: + Resolvent of the grid. + """ + if kwargs: + raise TypeError(f"resolvent() got unexpected keyword argument: {next(iter(kwargs))}") + signs = self._resolvent_signs(energies - chempot, ordering) + grid = np.expand_dims(self.points, axis=tuple(range(1, energies.ndim + 1))) + energies = np.expand_dims(energies, axis=0) + denominator = grid + (signs * 1.0j * self.eta - energies) + return 1.0 / denominator if invert else denominator + + @classmethod + def from_uniform( + cls, start: float, stop: float, num: int, eta: float | None = None + ) -> RealFrequencyGrid: + """Create a uniform real frequency grid. + + Args: + start: Start of the grid. + stop: End of the grid. + num: Number of points in the grid. + eta: Broadening factor. + + Returns: + Uniform real frequency grid. + """ + points = np.linspace(start, stop, num, endpoint=True) + return cls(points, eta=eta) + + +GridRF = RealFrequencyGrid + + +class ImaginaryFrequencyGrid(BaseFrequencyGrid): + """Imaginary frequency grid.""" + + beta: float = 256 + + _options = {"beta"} + + def __init__( # noqa: D417 + self, points: Array, weights: Array | None = None, **kwargs: Any + ) -> None: + """Initialise the grid. + + Args: + points: Points of the grid. + weights: Weights of the grid. + beta: Inverse temperature. + """ + self._points = np.asarray(points) + self._weights = np.asarray(weights) if weights is not None else None + self.set_options(**kwargs) + + @property + def reality(self) -> bool: + """Get the reality of the grid. + + Returns: + Reality of the grid. + """ + return False + + def resolvent( # noqa: D417 + self, + energies: Array, + chempot: float | Array, + invert: bool = True, + **kwargs: Any, + ) -> Array: + r"""Get the resolvent of the grid. + + For imaginary frequency grids, the resolvent is given by + + .. math:: + R(i \omega_n) = \frac{1}{i \omega_n - E}, + + where :math:`E` are the pole energies. + + Args: + energies: Energies of the poles. + chempot: Chemical potential. + invert: Whether to apply the inversion in the resolvent formula. + + Returns: + Resolvent of the grid. + """ + if kwargs: + raise TypeError(f"resolvent() got unexpected keyword argument: {next(iter(kwargs))}") + grid = np.expand_dims(self.points, axis=tuple(range(1, energies.ndim + 1))) + energies = np.expand_dims(energies, axis=0) + denominator = 1.0j * grid - energies + return 1.0 / denominator if invert else denominator + + @classmethod + def from_uniform(cls, num: int, beta: float | None = None) -> ImaginaryFrequencyGrid: + """Create a uniform imaginary frequency grid. + + Args: + num: Number of points in the grid. + beta: Inverse temperature. + + Returns: + Uniform imaginary frequency grid. + """ + if beta is None: + beta = cls.beta + separation = 2.0 * np.pi / beta + start = 0.5 * separation + stop = (num - 0.5) * separation + points = np.linspace(start, stop, num, endpoint=True) + return cls(points, beta=beta) + + @classmethod + def from_legendre( + cls, num: int, diffuse_factor: float = 1.0, beta: float | None = None + ) -> ImaginaryFrequencyGrid: + """Create a Legendre imaginary frequency grid. + + Args: + num: Number of points in the grid. + diffuse_factor: Diffuse factor for the grid. + beta: Inverse temperature. + + Returns: + Legendre imaginary frequency grid. + """ + points, weights = scipy.special.roots_legendre(num) + points = (1 - points) / (diffuse_factor * (1 + points)) + return cls(points, weights=weights, beta=beta) + + +GridIF = ImaginaryFrequencyGrid diff --git a/dyson/grids/grid.py b/dyson/grids/grid.py new file mode 100644 index 0000000..f22cba2 --- /dev/null +++ b/dyson/grids/grid.py @@ -0,0 +1,156 @@ +"""Base class for grids.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from dyson import numpy as np +from dyson.representations.enums import Component, Reduction, RepresentationEnum + +if TYPE_CHECKING: + from typing import Any + + from dyson.representations.dynamic import Dynamic + from dyson.representations.lehmann import Lehmann + from dyson.typing import Array + + +class BaseGrid(ABC): + """Base class for grids.""" + + _options: set[str] = set() + + _points: Array + _weights: Array | None = None + + def __init__( # noqa: D417 + self, points: Array, weights: Array | None = None, **kwargs: Any + ) -> None: + """Initialise the grid. + + Args: + points: Points of the grid. + weights: Weights of the grid. + """ + self._points = np.asarray(points) + self._weights = np.asarray(weights) if weights is not None else None + self.set_options(**kwargs) + + def set_options(self, **kwargs: Any) -> None: + """Set options for the solver. + + Args: + kwargs: Keyword arguments to set as options. + """ + for key, val in kwargs.items(): + if key not in self._options: + raise ValueError(f"Unknown option for {self.__class__.__name__}: {key}") + if isinstance(getattr(self, key), RepresentationEnum): + # Casts string to the appropriate enum type if the default value is an enum + val = getattr(self, key).__class__(val) + setattr(self, key, val) + + @abstractmethod + def evaluate_lehmann( + self, + lehmann: Lehmann, + reduction: Reduction = Reduction.NONE, + component: Component = Component.FULL, + ) -> Dynamic[Any]: + """Evaluate a Lehmann representation on the grid. + + Args: + lehmann: Lehmann representation to evaluate. + reduction: The reduction of the dynamic representation. + component: The component of the dynamic representation. + + Returns: + Lehmann representation, realised on the grid. + """ + pass + + @property + def points(self) -> Array: + """Get the points of the grid. + + Returns: + Points of the grid. + """ + return self._points + + @property + def weights(self) -> Array: + """Get the weights of the grid. + + Returns: + Weights of the grid. + """ + if self._weights is None: + return np.ones_like(self.points) / len(self) + return self._weights + + def __getitem__(self, key: int | slice | list[int] | Array) -> BaseGrid: + """Get a subset of the grid. + + Args: + key: Index or slice to get. + + Returns: + Subset of the grid. + """ + points = self.points[key] + weights = self.weights[key] if self._weights is not None else None + kwargs = {opt: getattr(self, opt) for opt in self._options} + return self.__class__(points, weights=weights, **kwargs) + + def __len__(self) -> int: + """Get the size of the grid. + + Returns: + Size of the grid. + """ + return self.points.shape[0] + + @property + def uniformly_spaced(self) -> bool: + """Check if the grid is uniformly spaced. + + Returns: + True if the grid is uniformly spaced, False otherwise. + """ + if len(self) < 2: + raise ValueError("Grid is too small to compute separation.") + return np.allclose(np.diff(self.points), self.points[1] - self.points[0]) + + @property + def uniformly_weighted(self) -> bool: + """Check if the grid is uniformly weighted. + + Returns: + True if the grid is uniformly weighted, False otherwise. + """ + return np.allclose(self.weights, self.weights[0]) + + @property + def separation(self) -> float: + """Get the separation of the grid. + + Returns: + Separation of the grid. + """ + if not self.uniformly_spaced: + raise ValueError("Grid is not uniformly spaced.") + return np.abs(self.points[1] - self.points[0]) + + @property + @abstractmethod + def domain(self) -> str: + """Get the domain of the grid.""" + pass + + @property + @abstractmethod + def reality(self) -> bool: + """Get the reality of the grid.""" + pass diff --git a/dyson/lehmann.py b/dyson/lehmann.py deleted file mode 100644 index c128a80..0000000 --- a/dyson/lehmann.py +++ /dev/null @@ -1,611 +0,0 @@ -""" -Containers for Lehmann representations. -""" - -import numpy as np -from pyscf import lib - - -class Lehmann: - """ - Lehmann representations. - - Parameters - ---------- - energies : numpy.ndarray - Energies of the poles in the Lehmann representation. - couplings : numpy.ndarray or tuple of numpy.ndarray - Couplings of the poles in the Lehmann representation to a - physical space (i.e. Dyson orbitals in the case of a Green's - function). Can be a tuple of `(left, right)` for a - non-Hermitian representation. - chempot : float, optional - Chemical potential, indicating the position of the Fermi - energy. Default value is `0.0`. - sort : bool, optional - Whether to sort the Lehmann representation by energy. Default - value is `True`. - """ - - def __init__(self, energies, couplings, chempot=0.0, sort=True): - # Input: - self.energies = energies - self.couplings = couplings - self.chempot = chempot - - # Check sanity: - if self.hermitian: - assert self.energies.size == self.couplings.shape[-1] - else: - assert self.energies.size == self.couplings[0].shape[-1] - assert self.energies.size == self.couplings[1].shape[-1] - - # Sort by energies: - if sort: - self.sort_() - - @classmethod - def from_pyscf(cls, lehmann_pyscf): - """ - Construct a Lehmann representation from a PySCF SelfEnergy - or GreensFunction. - - Parameters - ---------- - lehmann_pyscf : pyscf.agf2.aux.AuxSpace - PySCF Lehmann representation. - - Returns - ------- - lehmann : Lehmann - Lehmann representation. - """ - - if isinstance(lehmann_pyscf, cls): - return lehmann_pyscf - else: - return cls( - lehmann_pyscf.energy, - lehmann_pyscf.coupling, - chempot=lehmann_pyscf.chempot, - ) - - def sort_(self): - """ - Sort the Lehmann representation by energy. - """ - - idx = np.argsort(self.energies) - self.energies = self.energies[idx] - if self.hermitian: - self.couplings = self.couplings[:, idx] - else: - self.couplings = (self.couplings[0][:, idx], self.couplings[1][:, idx]) - - def moment(self, order): - """ - Get a spectral moment (or range of moments) of the Lehmann - representation. - - Parameters - ---------- - order : int or iterable of int - Order(s) to calculate. - - Returns - ------- - moment : numpy.ndarray - Spectral moments, if `order` is an `int` then the moment - is a 2D matrix, and if `order` is an `iterable` then it is - a 3D matrix enumerating the orders. - """ - - squeeze = False - if isinstance(order, int): - order = [order] - squeeze = True - order = np.array(order) - - couplings_l, couplings_r = self._unpack_couplings() - - moment = lib.einsum( - "pk,qk,nk->npq", - couplings_l, - couplings_r.conj(), - self.energies[None] ** order[:, None], - optimize=True, - ) - - if squeeze: - moment = moment[0] - - return moment - - def chebyshev_moment(self, order, scaling=None, scale_couplings=False): - """ - Get a spectral Chebyshev moment (or range of moments) of the - Lehmann representation. - - Parameters - ---------- - order : int or iterable of int - Order(s) to calculate - scaling : tuple of float - Scaling parameters, such that the energy scale of the - Lehmann representation is scaled as - `(energies - scaling[1]) / scaling[0]`. If `None`, the - scaling paramters are computed as - `(max(energies) - min(energies)) / (2.0 - 1e-3)` and - `(max(energies) + min(energies)) / 2.0`, respectively. - scale_couplings : bool, optional - Whether to also scale the couplings. Necessary when one - wishes to calculate Chebyshev moments for a self-energy. - Default value is `False`. - - Returns - ------- - chebyshev : numpy.ndarray - Spectral Chebyshev moments, if `order` is an `int` then the - moment is a 2D matrix, and if `order` is an `iterable` then - it is a 3D matrix enumerating the orders. - """ - - if scaling is not None: - a, b = scaling - else: - emin = min(self.energies) - emax = max(self.energies) - a = (emax - emin) / (2.0 - 1e-3) - b = (emax + emin) / 2.0 - - nmoms = set((order,) if isinstance(order, int) else order) - nmom_max = max(nmoms) - - couplings_l, couplings_r = self._unpack_couplings() - energies_scaled = (self.energies - b) / a - if scale_couplings: - couplings_l = couplings_l / a - couplings_r = couplings_r / a - - moments = np.zeros((len(nmoms), self.nphys, self.nphys), dtype=self.dtype) - vecs = (couplings_l, couplings_l * energies_scaled) - - j = 0 - for i in range(2): - if i in nmoms: - moments[i] = np.dot(vecs[i], couplings_r.T.conj()) - j += 1 - - for i in range(2, nmom_max + 1): - vec_next = 2 * energies_scaled * vecs[1] - vecs[0] - vecs = (vecs[1], vec_next) - if i in nmoms: - moments[j] = np.dot(vec_next, couplings_r.T.conj()) - j += 1 - - return moments - - def matrix(self, physical, chempot=False, out=None): - """ - Build a dense matrix consisting of a matrix (i.e. a - Hamiltonian) in the physical space coupling to a series of - energies as defined by the Lehmann representation. - - Parameters - ---------- - physical : numpy.ndarray - Physical space part of the matrix. - chempot : bool or float, optional - Include the chemical potential in the energies. If given - as a `bool`, use `self.chempot`. If a `float` then use - this as the chemical potential. Default value is `False`. - - Returns - ------- - matrix : numpy.ndarray - Dense matrix representation. - """ - - couplings_l, couplings_r = self._unpack_couplings() - - energies = self.energies - if chempot: - energies = energies - chempot - - if out is None: - dtype = np.result_type(couplings_l.dtype, couplings_r.dtype, physical.dtype) - out = np.zeros((self.nphys + self.naux, self.nphys + self.naux), dtype=dtype) - - out[: self.nphys, : self.nphys] = physical - out[: self.nphys, self.nphys :] = couplings_l - out[self.nphys :, : self.nphys] = couplings_r.T.conj() - out[self.nphys :, self.nphys :] = np.diag(energies) - - return out - - def matvec(self, physical, vector, chempot=False, out=None): - """ - Apply the dense matrix representation of the Lehmann - representation to a vector. This is equivalent to - `self.matrix(physical, chempot=chempot) @ vector`. - - Parameters - ---------- - physical : numpy.ndarray - Physical space part of the matrix. - vector : numpy.ndarray - Vector to apply the matrix to. - chempot : bool or float, optional - Include the chemical potential in the energies. If given - as a `bool`, use `self.chempot`. If a `float` then use - this as the chemical potential. Default value is `False`. - - Returns - ------- - result : numpy.ndarray - Result of applying the matrix to the vector. - """ - - couplings_l, couplings_r = self._unpack_couplings() - - energies = self.energies - if chempot: - energies = energies - chempot - - assert vector.shape == (self.nphys + self.naux,) - - if out is None: - dtype = np.result_type( - couplings_l.dtype, - couplings_r.dtype, - physical.dtype, - vector.dtype, - ) - out = np.zeros(self.nphys + self.naux, dtype=dtype) - - out[: self.nphys] += physical @ vector[: self.nphys] - out[: self.nphys] += couplings_l @ vector[self.nphys :] - out[self.nphys :] += couplings_r.T.conj() @ vector[: self.nphys] - out[self.nphys :] += energies * vector[self.nphys :] - - return out - - def diagonalise_matrix(self, physical, chempot=False, out=None): - """ - Diagonalise the dense matrix representation of the Lehmann - representation. - - Parameters - ---------- - physical : numpy.ndarray - Physical space part of the matrix. - chempot : bool or float, optional - Include the chemical potential in the energies. If given - as a `bool`, use `self.chempot`. If a `float` then use - this as the chemical potential. Default value is `False`. - - Returns - ------- - eigenvalues : numpy.ndarray - Eigenvalues of the matrix. - eigenvectors : numpy.ndarray - Eigenvectors of the matrix. - """ - - matrix = self.matrix(physical, chempot=chempot, out=out) - - if self.hermitian: - w, v = np.linalg.eigh(matrix) - else: - w, v = np.linalg.eig(matrix) - - return w, v - - def diagonalise_matrix_with_projection(self, physical, chempot=False, out=None): - """ - Diagonalise the dense matrix representation of the Lehmann - representation, and project the eigenvectors back to the - physical space. - - Parameters - ---------- - physical : numpy.ndarray - Physical space part of the matrix. - chempot : bool or float, optional - Include the chemical potential in the energies. If given - as a `bool`, use `self.chempot`. If a `float` then use - this as the chemical potential. Default value is `False`. - - Returns - ------- - eigenvalues : numpy.ndarray - Eigenvalues of the matrix. - eigenvectors : numpy.ndarray - Eigenvectors of the matrix, projected into the physical - space. - """ - - w, v = self.diagonalise_matrix(physical, chempot=chempot, out=out) - - if self.hermitian: - v = v[: self.nphys] - else: - vl = v[: self.nphys] - vr = np.linalg.inv(v).T.conj()[: self.nphys] - v = (vl, vr) - - return w, v - - def weights(self, occupancy=1): - """ - Get the weights of the residues in the Lehmann representation. - - Parameters - ---------- - occupancy : int or float, optional - Occupancy of the states. Default value is `1`. - - Returns - ------- - weights : numpy.ndarray - Weights of the states. - """ - - couplings_l, couplings_r = self._unpack_couplings() - wt = np.sum(couplings_l * couplings_r.conj(), axis=0) * occupancy - - return wt - - def as_orbitals(self, occupancy=1, mo_coeff=None): - """ - Convert the Lehmann representation to an orbital representation. - - Parameters - ---------- - occupancy : int or float, optional - Occupancy of the states. Default value is `1`. - mo_coeff : numpy.ndarray, optional - Molecular orbital coefficients. If given, the orbitals - will be rotated into the basis of these coefficients. - Default value is `None`. - - Returns - ------- - orb_energy : numpy.ndarray - Orbital energies. - orb_coeff : numpy.ndarray - Orbital coefficients. - orb_occ : numpy.ndarray - Orbital occupancies. - """ - - if not self.hermitian: - raise NotImplementedError - - orb_energy = self.energies - - if mo_coeff is not None: - orb_coeff = np.dot(mo_coeff, self.couplings) - else: - orb_coeff = self.couplings - - orb_occ = np.zeros_like(orb_energy) - orb_occ[orb_energy < self.chempot] = np.abs(self.occupied().weights(occupancy=occupancy)) - - return orb_energy, orb_coeff, orb_occ - - def as_static_potential(self, mo_energy, eta=1e-2): - """ - Convert the Lehmann representation to a static potential, for - example for us in qsGW when the Lehmann representation is of a - self-energy. - - Parameters - ---------- - mo_energy : numpy.ndarray - Molecular orbital energies. - eta : float, optional - Broadening parameter. Default value is `1e-2`. - - Returns - ------- - static_potential : numpy.ndarray - Static potential. - """ - - energies = self.energies + np.sign(self.energies - self.chempot) * 1.0j * eta - denom = mo_energy[:, None] - energies[None, :] - - couplings_l, couplings_r = self._unpack_couplings() - - static_potential = np.einsum("pk,qk,pk->pq", couplings_l, couplings_r, 1.0 / denom).real - static_potential = 0.5 * (static_potential + static_potential.T) - - return static_potential - - def as_perturbed_mo_energy(self): - """ - Return a list akin to the `mo_energy` attribute of a - `pyscf.scf.hf.SCF` object, but with the energies replaced by - those in the Lehmann representation that overlap the most with - each orbital. - - Returns - ------- - mo_energy : numpy.ndarray - Perturbed molecular orbital energies. May not necessarily - be sorted. - """ - - mo_energy = np.zeros((self.nphys,)) - - couplings_l, couplings_r = self._unpack_couplings() - weights = couplings_l * couplings_r.conj() - - for i in range(self.nphys): - mo_energy[i] = self.energies[np.argmax(np.abs(weights[i, :]))] - - return mo_energy - - def on_grid(self, grid, eta=1e-1, ordering="time-ordered", axis="real", trace=False): - """ - Return the Lehmann representation realised on a frequency - grid. - - Parameters - ---------- - grid : numpy.ndarray - Array of frequency points. - eta : float, optional - Broadening parameter. Default value is `1e-1`. - Only relevant for real axis. - ordering : str, optional - Time ordering. Can be one of `{"time-ordered", - "advanced", "retarded"}`. Default value is - `"time-ordered"`. - axis : str, optional - Frequency axis. Can be one of `{"real", "imag"}`. Default - value is `"real"`. - trace : bool, optional - Only return the trace. - - Returns - ------- - f : numpy.ndarray - Lehmann representation realised at each frequency point. - """ - - if ordering == "time-ordered": - signs = np.sign(self.energies - self.chempot) - elif ordering == "advanced": - signs = -np.ones_like(self.energies) - elif ordering == "retarded": - signs = np.ones_like(self.energies) - else: - raise ValueError("ordering = {}".format(ordering)) - - couplings_l, couplings_r = self._unpack_couplings() - - if axis == "real": - prednom = signs * 1.0j * eta - self.energies - denom = 1.0 / lib.direct_sum("w+k->wk", grid, prednom) - del predenom - elif axis == "imag": - denom = 1.0 / lib.direct_sum("w-k->wk", 1j * grid, self.energies) - else: - raise ValueError("axis = {}".format(axis)) - if trace: - f = lib.einsum("pk,pk,wk->w", couplings_l, couplings_r.conj(), denom) - else: - f = lib.einsum("pk,qk,wk->wpq", couplings_l, couplings_r.conj(), denom) - - return f - - @property - def hermitian(self): - """Boolean flag for the Hermiticity.""" - - return not isinstance(self.couplings, tuple) - - def _unpack_couplings(self): - if self.hermitian: - couplings_l = couplings_r = self.couplings - else: - couplings_l, couplings_r = self.couplings - - return couplings_l, couplings_r - - def _mask(self, mask, deep=True): - """Return a part of the Lehmann representation using a mask.""" - - if deep: - energies = self.energies[mask].copy() - if self.hermitian: - couplings = self.couplings[:, mask].copy() - else: - couplings = ( - self.couplings[0][:, mask].copy(), - self.couplings[1][:, mask].copy(), - ) - else: - energies = self.energies[mask] - couplings = self.couplings[:, mask] - - return self.__class__(energies, couplings, chempot=self.chempot) - - def physical(self, deep=True, weight=0.1): - """Return the part of the Lehmann representation with large - weights in the physical space. - """ - - return self._mask(self.weights() > weight, deep=deep) - - def occupied(self, deep=True): - """Return the occupied part of the Lehmann representation.""" - - return self._mask(self.energies < self.chempot, deep=deep) - - def virtual(self, deep=True): - """Return the virtual part of the Lehmann representation.""" - - return self._mask(self.energies >= self.chempot, deep=deep) - - def copy(self, chempot=None, deep=True): - """Return a copy with optionally updated chemical potential.""" - - if chempot is None: - chempot = self.chempot - - if deep: - energies = self.energies.copy() - if self.hermitian: - couplings = self.couplings.copy() - else: - couplings = (self.couplings[0].copy(), self.couplings[1].copy()) - else: - energies = self.energies - couplings = self.couplings - - return self.__class__(energies, couplings, chempot=chempot) - - @property - def nphys(self): - """Number of physical degrees of freedom.""" - return self._unpack_couplings()[0].shape[0] - - @property - def naux(self): - """Number of auxiliary degrees of freedom.""" - return self._unpack_couplings()[0].shape[1] - - def __add__(self, other): - """Combine two Lehmann representations.""" - - if self.nphys != other.nphys: - raise ValueError("Number of physical degrees of freedom do not match.") - - if self.chempot != other.chempot: - raise ValueError("Chemical potentials do not match.") - - energies = np.concatenate((self.energies, other.energies)) - - couplings_a_l, couplings_a_r = self._unpack_couplings() - couplings_b_l, couplings_b_r = other._unpack_couplings() - - if self.hermitian: - couplings = np.concatenate((couplings_a_l, couplings_b_l), axis=1) - else: - couplings = ( - np.concatenate((couplings_a_l, couplings_b_l), axis=1), - np.concatenate((couplings_a_r, couplings_b_r), axis=1), - ) - - return self.__class__(energies, couplings, chempot=self.chempot) - - @property - def dtype(self): - """Data type of the Lehmann representation.""" - if self.hermitian: - return np.result_type(self.energies, self.couplings) - else: - return np.result_type(self.energies, *self.couplings) diff --git a/dyson/plotting.py b/dyson/plotting.py new file mode 100644 index 0000000..4b3135e --- /dev/null +++ b/dyson/plotting.py @@ -0,0 +1,232 @@ +"""Plotting utilities.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import matplotlib.pyplot as plt +import scipy.constants + +from dyson import numpy as np +from dyson.representations.enums import Component, Reduction + +if TYPE_CHECKING: + from typing import Any, Literal + + from matplotlib.axes import Axes + from matplotlib.lines import Line2D + + from dyson.grids.frequency import BaseFrequencyGrid + from dyson.representations.dynamic import Dynamic + from dyson.representations.lehmann import Lehmann + + +theme = { + # Lines + "lines.linewidth": 3.0, + "lines.markersize": 10.0, + "lines.markeredgewidth": 1.0, + "lines.markeredgecolor": "black", + # Font + "font.size": 12, + "font.family": "sans-serif", + "font.weight": "medium", + # Axes + "axes.titlesize": 12, + "axes.labelsize": 12, + "axes.labelweight": "medium", + "axes.facecolor": "whitesmoke", + "axes.linewidth": 1.5, + "axes.unicode_minus": False, + "axes.prop_cycle": plt.cycler( + color=[ + "#1f77b4", + "#ff7f0e", + "#2ca02c", + "#d62728", + "#9467bd", + "#8c564b", + "#e377c2", + "#7f7f7f", + "#bcbd22", + "#17becf", + ] + ), + # Ticks + "xtick.labelsize": 12, + "xtick.major.pad": 7, + "xtick.major.size": 7, + "xtick.major.width": 1.2, + "xtick.minor.size": 4, + "xtick.minor.width": 0.6, + "ytick.labelsize": 12, + "ytick.major.pad": 7, + "ytick.major.size": 7, + "ytick.major.width": 1.2, + "ytick.minor.size": 4, + "ytick.minor.width": 0.6, + # Grid + "grid.linewidth": 1.3, + "grid.alpha": 0.5, + # Legend + "legend.fontsize": 11, + # Figure + "figure.figsize": (8, 6), + "figure.facecolor": "white", + "figure.autolayout": True, + # LaTeX + "pgf.texsystem": "pdflatex", +} + +plt.rcParams.update(theme) + + +def _unit_name(unit: str) -> str: + """Return the name of the unit for SciPy.""" + if unit == "Ha": + return "hartree" + elif unit == "eV": + return "electron volt" + else: + raise ValueError(f"Unknown energy unit: {unit}. Use 'Ha' or 'eV'.") + + +def _convert(energy: float, unit_from: str, unit_to: str) -> float: + """Convert energies between Hartree and eV.""" + if unit_from == unit_to: + return energy + unit_from = _unit_name(unit_from) + unit_to = _unit_name(unit_to) + return energy * scipy.constants.physical_constants[f"{unit_from}-{unit_to} relationship"][0] + + +def plot_lehmann( + lehmann: Lehmann, + ax: Axes | None = None, + energy_unit: Literal["Ha", "eV"] = "eV", + height_by_weight: bool = True, + height_factor: float = 1.0, + fmt: str = "k-", + **kwargs: Any, +) -> list[Line2D]: + """Plot a Lehmann representation as delta functions. + + Args: + lehmann: The Lehmann representation to plot. + ax: The axes to plot on. If ``None``, a new figure and axes are created. + energy_unit: The unit of the energy values. + height_by_weight: If ``True``, the height of each delta function is scaled by its weight. + If ``False``, all delta functions have the same height. + height_factor: A factor to scale the height of the delta functions. + fmt: The format string for the lines. + **kwargs: Additional keyword arguments passed to ``ax.plot``. + + Returns: + A list of Line2D objects representing the plotted delta functions. + """ + if ax is None: + fig, ax = plt.subplots() + lines: list[Line2D] = [] + for i, (energy, weight) in enumerate(zip(lehmann.energies, lehmann.weights())): + energy = _convert(energy, "Ha", energy_unit) + height = weight * height_factor if height_by_weight else height_factor + lines += ax.plot([energy, energy], [0, height], fmt, **kwargs) + return lines + + +def plot_dynamic( + dynamic: Dynamic, + ax: Axes | None = None, + energy_unit: Literal["Ha", "eV"] = "eV", + normalise: bool = False, + height_factor: float = 1.0, + fmt: str = "k-", + **kwargs: Any, +) -> list[Line2D]: + """Plot a dynamic representation as a line plot. + + Args: + dynamic: The dynamic representation to plot. + ax: The axes to plot on. If ``None``, a new figure and axes are created. + energy_unit: The unit of the energy values. + normalise: If ``True``, the representation is normalised to have a maximum value of 1. + height_factor: A factor to scale the height of the line. + fmt: The format string for the lines. + **kwargs: Additional keyword arguments passed to ``ax.plot``. + + Returns: + A list of Line2D objects representing the plotted dynamic. + """ + if ax is None: + fig, ax = plt.subplots() + if dynamic.reduction != Reduction.TRACE: + raise ValueError( + f"Dynamic object reduction must be {Reduction.TRACE.name} to plot as a line plot, but " + f"got {dynamic.reduction.name}. If you intended to plot the trace, use " + 'dynamic.copy(reduction="trace") to create a copy with the trace reduction.' + ) + if dynamic.component == Component.FULL: + raise ValueError( + f"Dynamic object component must be {Component.REAL.name} or {Component.IMAG.name} to " + f"plot as a line plot, but got {dynamic.component.name}. If you intended to plot the " + 'real or imaginary part, use dynamic.copy(component="real") or ' + 'dynamic.copy(component="imag") to create a copy with the desired component.' + ) + grid = _convert(dynamic.grid.points, "Ha", energy_unit) + array = dynamic.array + if normalise: + array = array / np.max(np.abs(array)) + return ax.plot(grid, array * height_factor, fmt, **kwargs) + + +def format_axes_spectral_function( + grid: BaseFrequencyGrid, + ax: Axes | None = None, + energy_unit: Literal["Ha", "eV"] = "eV", + xlabel: str = "Frequency ({})", + ylabel: str = "Spectral function", +) -> None: + """Format the axes for a spectral function plot. + + Args: + grid: The frequency grid used for the spectral function. + ax: The axes to format. If ``None``, the current axes are used. + energy_unit: The unit of the energy values. + xlabel: The label for the x-axis. + ylabel: The label for the y-axis. + """ + if ax is None: + ax = plt.gca() + ax.set_xlabel(xlabel.format(energy_unit)) + ax.set_ylabel(ylabel) + ax.set_yticks([]) + ax.set_xlim( + _convert(grid.points.min(), "Ha", energy_unit), + _convert(grid.points.max(), "Ha", energy_unit), + ) + + +def unknown_pleasures(dynamics: list[Dynamic]) -> Axes: + """Channel your inner Ian Curtis.""" + fig, ax = plt.subplots(figsize=(5, 7), facecolor="black") + norm = max([np.max(np.abs(d.array)) for d in dynamics]) + xmin = min([d.grid.points.min() for d in dynamics]) + xmax = max([d.grid.points.max() for d in dynamics]) + xmin -= (xmax - xmin) * 0.05 # Add some padding + xmax += (xmax - xmin) * 0.05 # Add some padding + ymax = 0.0 + spacing = 0.2 + zorder = 1 + for i, dynamic in list(enumerate(dynamics))[::-1]: + grid = _convert(dynamic.grid.points, "Ha", "eV") + array = dynamic.array / norm + array += i * spacing + array += np.random.uniform(-0.015, 0.015, size=array.shape) # Add some noise + ymax = max(ymax, np.max(array)) + ax.fill_between(grid, i * spacing, array, color="k", zorder=zorder) + ax.plot(grid, array, "-", color="white", linewidth=2.0, zorder=zorder + 1) + zorder += 2 + ax.axis("off") + ax.set_xlim(_convert(xmin, "Ha", "eV"), _convert(xmax, "Ha", "eV")) + ax.set_ylim(-0.1, ymax + spacing) + return ax diff --git a/dyson/printing.py b/dyson/printing.py new file mode 100644 index 0000000..6665fe7 --- /dev/null +++ b/dyson/printing.py @@ -0,0 +1,297 @@ +"""Printing utilities.""" + +from __future__ import annotations + +import importlib +import os +import subprocess +from typing import TYPE_CHECKING + +from rich import box +from rich.console import Console +from rich.errors import LiveError +from rich.progress import Progress +from rich.table import Table +from rich.theme import Theme + +from dyson import __version__ + +if TYPE_CHECKING: + from typing import Any, Literal + + from rich.progress import TaskID + + +theme = Theme( + { + "good": "green", + "okay": "yellow", + "bad": "red", + "output": "cyan", + "input": "bright_magenta", + "method": "bold underline", + "header": "bold", + } +) + +console = Console( + highlight=False, + theme=theme, + log_path=False, + quiet=os.environ.get("DYSON_QUIET", "").lower() in ("1", "true"), +) + +HEADER = r""" _ + | | + __| | _ _ ___ ___ _ __ + / _` || | | |/ __| / _ \ | '_ \ +| (_| || |_| |\__ \| (_) || | | | + \__,_| \__, ||___/ \___/ |_| |_| + __/ | + |___/ %s +""" + + +def init_console() -> None: + """Initialise the console with a header.""" + if globals().get("_DYSON_LOG_INITIALISED", False): + return + + # Print header + header_with_version = "[header]" + HEADER + "[/header]" + header_with_version %= " " * (18 - len(__version__)) + "[input]" + __version__ + "[/input]" + console.print(header_with_version) + + # Print versions of dependencies and ebcc + def get_git_hash(directory: str) -> str: + git_directory = os.path.join(directory, ".git") + cmd = ["git", "--git-dir=%s" % git_directory, "rev-parse", "--short", "HEAD"] + try: + git_hash = subprocess.check_output( + cmd, universal_newlines=True, stderr=subprocess.STDOUT + ).rstrip() + except subprocess.CalledProcessError: + git_hash = "N/A" + return git_hash + + for name in ["numpy", "pyscf", "dyson"]: + module = importlib.import_module(name) + if module.__file__ is None: + git_hash = "N/A" + else: + git_hash = get_git_hash(os.path.join(os.path.dirname(module.__file__), "..")) + console.print(f"{name}:") + console.print(f" > Version: [input]{module.__version__}[/]") + console.print(f" > Git hash: [input]{git_hash}[/]") + + console.print("OMP_NUM_THREADS = [input]%s[/]" % os.environ.get("OMP_NUM_THREADS", "")) + + globals()["_DYSON_LOG_INITIALISED"] = True + + +class Quiet: + """Context manager to disable console output.""" + + def __init__(self, console: Console = console): + """Initialise the object.""" + self._memo: list[bool] = [] + self._console = console + + def __enter__(self) -> None: + """Enter the context manager.""" + self._memo.append(self.console.quiet) + self.console.quiet = True + + def __exit__(self, *args: Any) -> None: + """Exit the context manager.""" + quiet = self._memo.pop() + self.console.quiet = quiet + + def __call__(self) -> None: + """Call the context manager.""" + self.console.quiet = True + + @property + def console(self) -> Console: + """Get the console.""" + return self._console + + +quiet = Quiet(console) + + +def rate_error( + value: float | complex, threshold: float, threshold_okay: float | None = None +) -> Literal["good", "okay", "bad"]: + """Rate the error based on a threshold. + + Args: + value: The value to rate. + threshold: The threshold for the rating. + threshold_okay: Separate threshold for ``"okay"`` rating. Default is 10 times + ``threshold``. + + Returns: + str: The rating, one of "good", "okay", or "bad". + """ + if threshold_okay is None: + threshold_okay = 10 * threshold + if abs(value) < threshold: + return "good" + elif abs(value) < threshold_okay: + return "okay" + else: + return "bad" + + +def format_float( + value: float | complex | None, + precision: int = 10, + scientific: bool = False, + threshold: float | None = None, +) -> str: + """Format a float or complex number to a string with a given precision. + + Args: + value: The value to format. + precision: The number of decimal places to include. + scientific: Whether to use scientific notation for large or small values. + threshold: If provided, the value will be rated based on this threshold. + + Returns: + str: The formatted string. + """ + if isinstance(value, complex): + real = format_float(value.real, precision, scientific, threshold) + if abs(value.imag) < (1e-1**precision): + return real + sign = "+" if value.imag >= 0 else "-" + imag = format_float(abs(value.imag), precision, scientific, threshold) + return f"{real}{sign}{imag}i" + if value is None: + return "N/A" + if value.imag < (1e-1**precision): + value = value.real + out = f"{value:.{precision}g}" if scientific else f"{value:.{precision}f}" + if threshold is not None: + rating = rate_error(value, threshold) + out = f"[{rating}]{out}[/]" + return out + + +class ConvergencePrinter: + """Table for printing convergence information.""" + + def __init__( + self, + quantities: tuple[str, ...], + quantity_errors: tuple[str, ...], + thresholds: tuple[float, ...], + console: Console = console, + cycle_name: str = "Cycle", + ): + """Initialise the object.""" + self._console = console + self._table = Table(box=box.SIMPLE) + self._table.add_column(cycle_name, style="dim", justify="left") + for quantity in quantities: + self._table.add_column(quantity, justify="right") + for quantity_error in quantity_errors: + self._table.add_column(quantity_error, justify="right") + self._thresholds = thresholds + + def add_row( + self, + cycle: int, + quantities: tuple[float | None, ...], + quantity_errors: tuple[float | None, ...], + ) -> None: + """Add a row to the table.""" + self._table.add_row( + str(cycle), + *[format_float(quantity) for quantity in quantities], + *[ + format_float(error, precision=4, scientific=True, threshold=threshold) + for error, threshold in zip(quantity_errors, self._thresholds) + ], + ) + + def print(self) -> None: + """Print the table.""" + self._console.print(self._table) + + @property + def thresholds(self) -> tuple[float, ...]: + """Get the thresholds.""" + return self._thresholds + + +class IterationsPrinter: + """Progress bar for iterations.""" + + def __init__(self, max_cycle: int, console: Console = console, description: str = "Iteration"): + """Initialise the object.""" + self._max_cycle = max_cycle + self._console = console + self._description = description + self._ignore = False + self._progress = Progress(transient=True) + self._task: TaskID | None = None + + def start(self) -> None: + """Start the progress bar.""" + if self.console.quiet: + return + self._ignore = False + try: + self.progress.start() + except LiveError: + # If there is already a live print, don't start a progress bar + self._ignore = True + return + self._task = self.progress.add_task( + f"{self.description} 0 / {self.max_cycle}", total=self.max_cycle + ) + + def update(self, cycle: int) -> None: + """Update the progress bar for the given cycle.""" + if self.console.quiet or self._ignore: + return + if self.task is None: + raise RuntimeError("Progress bar has not been started. Call start() first.") + self.progress.update( + self.task, advance=1, description=f"{self.description} {cycle} / {self.max_cycle}" + ) + + def stop(self) -> None: + """Stop the progress bar.""" + if self.console.quiet or self._ignore: + return + if self.task is None: + raise RuntimeError("Progress bar has not been started. Call start() first.") + self.progress.stop() + + @property + def max_cycle(self) -> int: + """Get the maximum number of cycles.""" + return self._max_cycle + + @property + def console(self) -> Console: + """Get the console.""" + return self._console + + @property + def description(self) -> str: + """Get the description of the progress bar.""" + return self._description + + @property + def progress(self) -> Progress: + """Get the progress bar.""" + return self._progress + + @property + def task(self) -> TaskID | None: + """Get the current task.""" + return self._task diff --git a/dyson/representations/__init__.py b/dyson/representations/__init__.py new file mode 100644 index 0000000..056baf1 --- /dev/null +++ b/dyson/representations/__init__.py @@ -0,0 +1,87 @@ +r"""Representations for Green's functions and self-energies. + +Both the Green's function and self-energy can be represented in the frequency domain according to +their Lehmann representation, for the Green's function + +.. math:: + G_{pq}(\omega) = \sum_{x} \frac{u_{px} u_{qx}^*}{\omega - \varepsilon_x}, + +where poles :math:`\varepsilon_x` couple to the physical states of the system according to the +Dyson orbitals :math:`u_{px}`. For the self-energy, the representation is given by + +.. math:: + \Sigma_{pq}(\omega) = \sum_{k} \frac{v_{pk} v_{qk}^*}{\omega - \epsilon_k}, + +where :math:`v_{px}` are the couplings between auxiliary states and the physical states, and +:math:`\epsilon_k` are the auxiliary state energies. + +These two Lehmann representations can be relataed to each other via the Dyson equation, which can +be written as an eigenvalue problem in the upfolded configuration space as + +.. math:: + \begin{bmatrix} \boldsymbol{\Sigma}(\omega) & \mathbf{v} \\ \mathbf{v}^\dagger & + \boldsymbol{\epsilon} \mathbf{I} \end{bmatrix} \begin{bmatrix} \mathbf{u} \\ \mathbf{w} + \end{bmatrix} = \boldsymbol{\varepsilon} \begin{bmatrix} \mathbf{u} \\ \mathbf{w} \end{bmatrix}. + +The Lehmann representations of either the Green's function or self-energy are contained in +:class:`~dyson.representations.lehmann.Lehmann` objects, which is a simple container for the +energies and couplings, along with a chemical potential. The +:class:`~dyson.representations.spectral.Spectral` representation provides a container for the full +eigenspectrum (including :math:`\mathbf{w}`), and can provide the Lehmann representation of both the +Green's function and self-energy. + +>>> from dyson import util, quiet, FCI, Exact +>>> quiet() # Suppress output +>>> mf = util.get_mean_field("H 0 0 0; H 0 0 1", "6-31g") +>>> fci = FCI.h.from_mf(mf) +>>> solver = Exact.from_expression(fci) +>>> result = solver.kernel() +>>> type(result) + +>>> self_energy = result.get_self_energy() +>>> type(self_energy) + +>>> greens_function = result.get_greens_function() +>>> type(greens_function) + + +Lehmann representations can be realised onto a subclass :class:`~dyson.grids.grid.BaseGrid` to +provide a dynamic representation of the function, which is stored in a +:class:`~dyson.representations.dynamic.Dynamic` object. This dynamic representation has varied +formats, principally depending on the type of grid used, but also according to the so-called +:class:`~dyson.representations.enums.Reduction` and :class:`~dyson.representations.enums.Component` +of the representation. The :class:`~dyson.representations.enums.Reduction` enum encodes the format +of the matrix, i.e. whether it is the full matrix, the diagonal part, or the trace. The +:class:`~dyson.representations.enums.Component` enum encodes the numerical component of the matrix, +i.e. whether it is the real or imaginary part, or the full complex matrix. + +>>> from dyson.grids import GridRF +>>> grid = GridRF.from_uniform(-3.0, 3.0, 256, eta=1e-1) +>>> spectrum = grid.evaluate_lehmann( +... greens_function, ordering="retarded", reduction="trace", component="imag" +... ) +>>> type(spectrum) + + +The various solvers in :mod:`~dyson.solvers` have different representations as their inputs and +outputs. + + +Submodules +---------- + +.. autosummary:: + :toctree: + + representation + lehmann + spectral + dynamic + enums + +""" + +from dyson.representations.enums import Reduction, Component +from dyson.representations.lehmann import Lehmann +from dyson.representations.spectral import Spectral +from dyson.representations.dynamic import Dynamic diff --git a/dyson/representations/dynamic.py b/dyson/representations/dynamic.py new file mode 100644 index 0000000..3bf1328 --- /dev/null +++ b/dyson/representations/dynamic.py @@ -0,0 +1,352 @@ +"""Container for a dynamic representation.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, Generic, TypeVar + +from dyson import numpy as np +from dyson import util +from dyson.grids.grid import BaseGrid +from dyson.representations.enums import Component, Reduction +from dyson.representations.representation import BaseRepresentation + +if TYPE_CHECKING: + from dyson.representations.lehmann import Lehmann + from dyson.typing import Array + +_TGrid = TypeVar("_TGrid", bound=BaseGrid) + + +def _cast_reduction(first: Reduction, second: Reduction) -> Reduction: + """Find the reduction that is compatible with both reductions.""" + values = {Reduction.NONE: 0, Reduction.DIAG: 1, Reduction.TRACE: 2} + if values[first] <= values[second]: + return first + return second + + +def _cast_component(first: Component, second: Component) -> Component: + """Find the component that is compatible with both components.""" + if first == second: + return first + return Component.FULL + + +def _cast_arrays(first: Dynamic[_TGrid], second: Dynamic[_TGrid]) -> tuple[Array, Array]: + """Cast the arrays of two dynamic representations to the same component and reduction.""" + component = _cast_component(first.component, second.component) + reduction = _cast_reduction(first.reduction, second.reduction) + array_first = first.as_dynamic(component=component, reduction=reduction).array + array_second = second.as_dynamic(component=component, reduction=reduction).array + return array_first, array_second + + +def _same_grid(first: Dynamic[_TGrid], second: Dynamic[_TGrid]) -> bool: + """Check if two dynamic representations have the same grid.""" + # TODO: Move to BaseGrid + if not isinstance(second.grid, type(first.grid)): + return False + if len(first.grid) != len(second.grid): + return False + if not all( + getattr(first.grid, attr) == getattr(second.grid, attr) for attr in first.grid._options + ): + return False + if not np.allclose(first.grid.weights, second.grid.weights): + return False + return np.allclose(first.grid.points, second.grid.points) + + +class Dynamic(BaseRepresentation, Generic[_TGrid]): + r"""Dynamic representation. + + The dynamic representation is a set of arrays in some physical space at each point in a time or + frequency grid. This class contains the arrays and the grid information. + """ + + def __init__( + self, + grid: _TGrid, + array: Array, + reduction: Reduction = Reduction.NONE, + component: Component = Component.FULL, + hermitian: bool = False, + ): + """Initialise the object. + + Args: + grid: The grid on which the dynamic representation is defined. + array: The array of values at each point in the grid. + reduction: The reduction of the dynamic representation. + component: The component of the dynamic representation. + hermitian: Whether the array is Hermitian. + """ + self._grid = grid + self._array = array + self._hermitian = hermitian + self._reduction = Reduction(reduction) + self._component = Component(component) + if array.shape[0] != len(grid): + raise ValueError( + f"Array must have the same size as the grid in the first dimension, but got " + f"{array.shape[0]} for grid size {len(grid)}." + ) + if (array.ndim - 1) != self.reduction.ndim: + raise ValueError( + f"Array must be {self.reduction.ndim}D for reduction {self.reduction}, but got " + f"{array.ndim}D." + ) + if int(np.iscomplexobj(array)) + 1 != self.component.ncomp: + raise ValueError( + f"Array must only be complex valued for component {Component.FULL}, but got " + f"{array.dtype} for {self.component}." + ) + + @classmethod + def from_lehmann( + cls, + lehmann: Lehmann, + grid: _TGrid, + reduction: Reduction = Reduction.NONE, + component: Component = Component.FULL, + ) -> Dynamic[_TGrid]: + """Construct a dynamic representation from a Lehmann representation. + + Args: + lehmann: The Lehmann representation to convert. + grid: The grid on which the dynamic representation is defined. + reduction: The reduction of the dynamic representation. + component: The component of the dynamic representation. + + Returns: + A dynamic representation. + """ + return grid.evaluate_lehmann(lehmann, reduction=reduction, component=component) + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return self.array.shape[-1] + + @property + def grid(self) -> _TGrid: + """Get the grid on which the dynamic representation is defined.""" + return self._grid + + @property + def array(self) -> Array: + """Get the array of values at each point in the grid.""" + return self._array + + @property + def reduction(self) -> Reduction: + """Get the reduction of the dynamic representation.""" + return self._reduction + + @property + def component(self) -> Component: + """Get the component of the dynamic representation.""" + return self._component + + @property + def hermitian(self) -> bool: + """Get a boolean indicating if the system is Hermitian.""" + return self._hermitian or self.reduction != Reduction.NONE + + @property + def dtype(self) -> np.dtype: + """Get the data type of the array.""" + return self._array.dtype + + def __repr__(self) -> str: + """Get a string representation of the dynamic representation.""" + return f"Dynamic(grid={self.grid}, shape={self.array.shape}, hermitian={self.hermitian})" + + def copy( + self, + deep: bool = True, + reduction: Reduction | None = None, + component: Component | None = None, + ) -> Dynamic[_TGrid]: + """Return a copy of the dynamic representation. + + Args: + deep: Whether to return a deep copy of the energies and couplings. + component: The component of the dynamic representation. + reduction: The reduction of the dynamic representation. + + Returns: + A new dynamic representation. + """ + grid = self.grid + array = self.array + if reduction is None: + reduction = self.reduction + reduction = Reduction(reduction) + if component is None: + component = self.component + component = Component(component) + + # Copy the array if requested + if deep: + array = array.copy() + + # Adjust the reduction if necessary + if reduction != self.reduction: + if (self.reduction, reduction) == (Reduction.NONE, Reduction.DIAG): + array = np.diagonal(array, axis1=1, axis2=2) + elif (self.reduction, reduction) == (Reduction.NONE, Reduction.TRACE): + array = np.trace(array, axis1=1, axis2=2) + elif (self.reduction, reduction) == (Reduction.DIAG, Reduction.TRACE): + array = np.sum(array, axis=1) + elif (self.reduction, reduction) == (Reduction.DIAG, Reduction.NONE): + array_new = np.zeros((len(grid), self.nphys, self.nphys), dtype=array.dtype) + np.fill_diagonal(array_new, array) + array = array_new + else: + raise ValueError( + f"Cannot convert from {self.reduction} to {reduction} for dynamic " + "representation." + ) + + # Adjust the component if necessary + if component != self.component: + if (self.component, component) == (Component.FULL, Component.REAL): + array = np.real(array) + elif (self.component, component) == (Component.FULL, Component.IMAG): + array = np.imag(array) + elif (self.component, component) == (Component.REAL, Component.FULL): + array = array + 1.0j * np.zeros_like(array) + elif (self.component, component) == (Component.IMAG, Component.FULL): + array = np.zeros_like(array) + 1.0j * array + else: + raise ValueError( + f"Cannot convert from {self.component} to {component} for dynamic " + "representation." + ) + + return self.__class__(grid, array, hermitian=self.hermitian) + + def as_dynamic( + self, component: Component | None = None, reduction: Reduction | None = None + ) -> Dynamic[_TGrid]: + """Return the dynamic representation with the specified component and reduction. + + Args: + component: The component of the dynamic representation. + reduction: The reduction of the dynamic representation. + + Returns: + A new dynamic representation with the specified component and reduction. + """ + return self.copy(deep=False, component=component, reduction=reduction) + + def rotate(self, rotation: Array | tuple[Array, Array]) -> Dynamic[_TGrid]: + """Rotate the dynamic representation. + + Args: + rotation: The rotation matrix to apply to the array. If the matrix has three dimensions, + the first dimension is used to rotate on the left, and the second dimension is used + to rotate on the right. + + Returns: + A new dynamic representation with the rotated array. + """ + left, right = rotation if isinstance(rotation, tuple) else (rotation, rotation) + + if np.iscomplexobj(left) or np.iscomplexobj(right): + array = self.as_dynamic(component=Component.FULL).array + component = Component.FULL + else: + array = self.array + component = self.component + + if self.reduction == Reduction.NONE: + array = util.einsum("wpq,pi,qj->wij", array, left.conj(), right) + elif self.reduction == Reduction.DIAG: + array = util.einsum("wp,pi,pj->wij", array, left.conj(), right) + elif self.reduction == Reduction.TRACE: + raise ValueError("Cannot rotate a dynamic representation with trace reduction.") + + return self.__class__( + self.grid, + array, + component=component, + reduction=Reduction.NONE, + hermitian=self.hermitian, + ) + + def __add__(self, other: Dynamic[_TGrid]) -> Dynamic[_TGrid]: + """Add two dynamic representations.""" + if not isinstance(other, Dynamic): + return NotImplemented + if not _same_grid(self, other): + raise ValueError("Cannot add dynamic representations with different grids.") + return self.__class__( + self.grid, + np.add(*_cast_arrays(self, other)), + component=_cast_component(self.component, other.component), + reduction=_cast_reduction(self.reduction, other.reduction), + hermitian=self.hermitian or other.hermitian, + ) + + def __sub__(self, other: Dynamic[_TGrid]) -> Dynamic[_TGrid]: + """Subtract two dynamic representations.""" + if not isinstance(other, Dynamic): + return NotImplemented + if not _same_grid(self, other): + raise ValueError("Cannot subtract dynamic representations with different grids.") + return self.__class__( + self.grid, + np.subtract(*_cast_arrays(self, other)), + component=_cast_component(self.component, other.component), + reduction=_cast_reduction(self.reduction, other.reduction), + hermitian=self.hermitian or other.hermitian, + ) + + def __mul__(self, other: float | int) -> Dynamic[_TGrid]: + """Multiply the dynamic representation by a scalar.""" + if not isinstance(other, (float, int)): + return NotImplemented + return self.__class__( + self.grid, + self.array * other, + component=self.component, + reduction=self.reduction, + hermitian=self.hermitian, + ) + + __rmul__ = __mul__ + + def __neg__(self) -> Dynamic[_TGrid]: + """Negate the dynamic representation.""" + return -1 * self + + def __array__(self) -> Array: + """Return the dynamic representation as a NumPy array.""" + return self.array + + def __eq__(self, other: object) -> bool: + """Check if two dynamic representations are equal.""" + if not isinstance(other, Dynamic): + return NotImplemented + if other.nphys != self.nphys: + return False + if len(other.grid) != len(self.grid): + return False + if other.hermitian != self.hermitian: + return False + if not _same_grid(self, other): + return False + return np.allclose(other.array, self.array) + + def __hash__(self) -> int: + """Return a hash of the dynamic representation.""" + return hash( + ( + tuple(self.grid.points), + tuple(self.grid.weights), + tuple(self.array.ravel()), + self.hermitian, + ) + ) diff --git a/dyson/representations/enums.py b/dyson/representations/enums.py new file mode 100644 index 0000000..b9dba3a --- /dev/null +++ b/dyson/representations/enums.py @@ -0,0 +1,74 @@ +"""Enumerations for representations.""" + +from __future__ import annotations + +from enum import Enum +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + pass + + +class RepresentationEnum(Enum): + """Base enumeration for representations.""" + + def raise_invalid_representation(self) -> None: + """Raise an error for invalid representation.""" + name = self.__class__.__name__.lower() + valid = [r.name for r in self.__class__] + raise ValueError(f"Invalid {name}: {self.name}. Valid {name}s are: {', '.join(valid)}") + + +class Reduction(RepresentationEnum): + """Enumeration for the reduction of the dynamic representation. + + The valid reductions are: + + * ``none``: No reduction, i.e. the full 2D array. + * ``diag``: Reduction to the diagonal, i.e. a 1D array of diagonal elements. + * ``trace``: Reduction to the trace, i.e. a scalar value. + """ + + NONE = "none" + DIAG = "diag" + TRACE = "trace" + + @property + def ndim(self) -> int: + """Get the number of dimensions of the array for this reduction.""" + return {Reduction.NONE: 2, Reduction.DIAG: 1, Reduction.TRACE: 0}[self] + + +class Component(RepresentationEnum): + """Enumeration for the component of the dynamic representation. + + The valid components are: + + * ``full``: The full (real-valued or complex-valued) representation. + * ``real``: The real part of the representation. + * ``imag``: The imaginary part of the representation, represented as a real-valued array. + """ + + FULL = "full" + REAL = "real" + IMAG = "imag" + + @property + def ncomp(self) -> int: + """Get the number of components for this component type.""" + return 2 if self == Component.FULL else 1 + + +class Ordering(RepresentationEnum): + """Enumeration for the time ordering of the dynamic representation. + + The valid orderings are: + + * ``ordered``: Time-ordered representation. + * ``advanced``: Advanced representation, i.e. affects the past (non-causal). + * ``retarded``: Retarded representation, i.e. affects the future (causal). + """ + + ORDERED = "ordered" + ADVANCED = "advanced" + RETARDED = "retarded" diff --git a/dyson/representations/lehmann.py b/dyson/representations/lehmann.py new file mode 100644 index 0000000..78484bb --- /dev/null +++ b/dyson/representations/lehmann.py @@ -0,0 +1,902 @@ +"""Container for a Lehmann representation.""" + +from __future__ import annotations + +from contextlib import contextmanager +from typing import TYPE_CHECKING, cast + +import scipy.linalg + +from dyson import numpy as np +from dyson import util +from dyson.representations.enums import Reduction +from dyson.representations.representation import BaseRepresentation +from dyson.typing import Array + +if TYPE_CHECKING: + from typing import Iterable, Iterator + + import pyscf.agf2.aux + + +@contextmanager +def shift_energies(lehmann: Lehmann, shift: float) -> Iterator[None]: + """Shift the energies of a Lehmann representation using a context manager. + + Args: + lehmann: The Lehmann representation to shift. + shift: The amount to shift the energies by. + + Yields: + None + """ + original_energies = lehmann.energies + lehmann._energies = original_energies + shift # pylint: disable=protected-access + try: + yield + finally: + lehmann._energies = original_energies # pylint: disable=protected-access + + +class Lehmann(BaseRepresentation): + r"""Lehman representation. + + The Lehmann representation is a set of poles :math:`\epsilon_k` and couplings :math:`v_{pk}` + that can be downfolded into a frequency-dependent function as + + .. math:: + \sum_{k} \frac{v_{pk} u_{qk}^*}{\omega - \epsilon_k}, + + where the couplings are between the poles :math:`k` and the physical space :math:`p` and + :math:`q`, and may be non-Hermitian. The couplings :math:`v` are right-handed vectors, and + :math:`u` are left-handed vectors. + + Note that the order of the couplings is ``(left, right)``, whilst they act in the order + ``(right, left)`` in the numerator. The naming convention is chosen to be consistent with the + eigenvalue decomposition, where :math:`v` may be an eigenvector acting on the right of a matrix, + and :math:`u` is an eigenvector acting on the left of a matrix. + """ + + def __init__( + self, + energies: Array, + couplings: Array, + chempot: float = 0.0, + sort: bool = True, + ): + """Initialise the object. + + Args: + energies: Energies of the poles. + couplings: Couplings of the poles to a physical space. For a non-Hermitian system, they + should be have three dimensions, with the first dimension indexing + ``(left, right)``. + chempot: Chemical potential. + sort: Sort the poles by energy. + """ + self._energies = energies + self._couplings = couplings + self._chempot = chempot + if sort: + self.sort_() + if not self.hermitian: + if couplings.ndim != 3: + raise ValueError( + f"Couplings must be 3D for a non-Hermitian system, but got {couplings.ndim}D." + ) + if couplings.shape[0] != 2: + raise ValueError( + f"Couplings must have shape (2, nphys, naux) for a non-Hermitian system, " + f"but got {couplings.shape}." + ) + + @classmethod + def from_pyscf(cls, auxspace: pyscf.agf2.aux.AuxSpace | Lehmann) -> Lehmann: + """Construct a Lehmann representation from a PySCF auxiliary space. + + Args: + auxspace: The auxiliary space. + + Returns: + A Lehmann representation. + """ + if isinstance(auxspace, Lehmann): + return auxspace + return cls( + energies=auxspace.energy, + couplings=auxspace.coupling, + chempot=auxspace.chempot, + ) + + @classmethod + def from_empty(cls, nphys: int) -> Lehmann: + """Construct an empty Lehmann representation. + + Args: + nphys: The number of physical degrees of freedom. + + Returns: + An empty Lehmann representation. + """ + return cls( + energies=np.zeros((0,)), + couplings=np.zeros((nphys, 0)), + chempot=0.0, + sort=False, + ) + + def sort_(self) -> None: + """Sort the poles by energy. + + Note: + The object is sorted in place. + """ + idx = np.argsort(self.energies) + self._energies = self.energies[idx] + self._couplings = self.couplings[..., idx] + + @property + def energies(self) -> Array: + """Get the energies.""" + return self._energies + + @property + def couplings(self) -> Array: + """Get the couplings.""" + return self._couplings + + @property + def chempot(self) -> float: + """Get the chemical potential.""" + return self._chempot + + @property + def hermitian(self) -> bool: + """Get a boolean indicating if the system is Hermitian.""" + return self.couplings.ndim == 2 + + def unpack_couplings(self) -> tuple[Array, Array]: + """Unpack the couplings. + + Returns: + A tuple of left and right couplings. + """ + if self.hermitian: + return cast(tuple[Array, Array], (self.couplings, self.couplings)) + return cast(tuple[Array, Array], (self.couplings[0], self.couplings[1])) + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return self.unpack_couplings()[0].shape[0] + + @property + def naux(self) -> int: + """Get the number of auxiliary degrees of freedom.""" + return self.unpack_couplings()[0].shape[1] + + @property + def dtype(self) -> np.dtype: + """Get the data type of the couplings.""" + return np.result_type(self.energies, self.couplings) + + def __repr__(self) -> str: + """Return a string representation of the Lehmann representation.""" + return f"Lehmann(nphys={self.nphys}, naux={self.naux}, chempot={self.chempot})" + + def mask(self, mask: Array | slice, deep: bool = True) -> Lehmann: + """Return a part of the Lehmann representation according to a mask. + + Args: + mask: The mask to apply. + deep: Whether to return a deep copy of the energies and couplings. + + Returns: + A new Lehmann representation including only the masked states. + """ + # Mask the energies and couplings + energies = self.energies[mask] + couplings = self.couplings[..., mask] + + # Copy the couplings if requested + if deep: + energies = energies.copy() + couplings = couplings.copy() + + return self.__class__(energies, couplings, chempot=self.chempot, sort=False) + + def physical(self, weight: float = 0.1, deep: bool = True) -> Lehmann: + """Return the physical (large weight) part of the Lehmann representation. + + Args: + weight: The weight to use for the physical part. + deep: Whether to return a deep copy of the energies and couplings. + + Returns: + A new Lehmann representation including only the physical part. + """ + return self.mask(self.weights() > weight, deep=deep) + + def occupied(self, deep: bool = True) -> Lehmann: + """Return the occupied part of the Lehmann representation. + + Args: + deep: Whether to return a deep copy of the energies and couplings. + + Returns: + A new Lehmann representation including only the occupied part. + """ + return self.mask(self.energies < self.chempot, deep=deep) + + def virtual(self, deep: bool = True) -> Lehmann: + """Return the virtual part of the Lehmann representation. + + Args: + deep: Whether to return a deep copy of the energies and couplings. + + Returns: + A new Lehmann representation including only the virtual part. + """ + return self.mask(self.energies >= self.chempot, deep=deep) + + def copy(self, chempot: float | None = None, deep: bool = True) -> Lehmann: + """Return a copy of the Lehmann representation. + + Args: + chempot: The chemical potential to use for the copy. If ``None``, the original + chemical potential is used. + deep: Whether to return a deep copy of the energies and couplings. + + Returns: + A new Lehmann representation. + """ + energies = self.energies + couplings = self.couplings + if chempot is None: + chempot = self.chempot + + # Copy the couplings if requested + if deep: + energies = energies.copy() + couplings = couplings.copy() + + return self.__class__(energies, couplings, chempot=self.chempot, sort=False) + + def rotate_couplings(self, rotation: Array | tuple[Array, Array]) -> Lehmann: + r"""Rotate the couplings and return a new Lehmann representation. + + For rotation matrix :math:`R`, the couplings are rotated as + + .. math:: + \tilde{\mathbf{v}} = R^\dagger \mathbf{v}, \quad + \tilde{\mathbf{u}} = R^\dagger \mathbf{u}, + + where :math:`v` are the right couplings and :math:`u` are the left couplings. + + Args: + rotation: The rotation matrix to apply to the couplings. If the matrix has three + dimensions, the first dimension is used to rotate the left couplings, and the + second dimension is used to rotate the right couplings. + + Returns: + A new Lehmann representation with the couplings rotated into the new basis. + """ + if not isinstance(rotation, tuple) and rotation.ndim == 2: + couplings = util.einsum("...pk,pq->...qk", self.couplings, rotation.conj()) + else: + left, right = self.unpack_couplings() + if isinstance(rotation, tuple) or rotation.ndim == 3: + rot_left, rot_right = rotation + else: + rot_left = rot_right = rotation + couplings = np.array( + [ + rot_left.T.conj() @ left, + rot_right.T.conj() @ right, + ], + ) + return self.__class__( + self.energies, + couplings, + chempot=self.chempot, + sort=False, + ) + + # Methods to calculate moments: + + def moments(self, order: int | Iterable[int], reduction: Reduction = Reduction.NONE) -> Array: + r"""Calculate the moment(s) of the Lehmann representation. + + The moments are defined as + + .. math:: + T_{pq}^{n} = \sum_{k} v_{pk} u_{qk}^* \epsilon_k^n, + + where :math:`T_{pq}^{n}` is the moment of order :math:`n` in the physical space. In terms of + the frequency-dependency, the moments can be written as the integral + + .. math:: + T_{pq}^{n} = \int_{-\infty}^{\infty} d\omega \, \left[ \sum_{k} + \frac{v_{pk} u_{qk}^*}{\omega - \epsilon_k} \right] \, \omega^n, + + where the integral is over the entire real line for central moments. + + Args: + order: The order(s) of the moment(s). + reduction: The reduction to apply to the moments. + + Returns: + The moment(s) of the Lehmann representation. + """ + squeeze = False + if isinstance(order, int): + order = [order] + squeeze = True + orders = np.asarray(order) + + # Get the subscript depending on the reduction + if Reduction(reduction) == Reduction.NONE: + subscript = "pk,qk,nk->npq" + elif Reduction(reduction) == Reduction.DIAG: + subscript = "pk,pk,nk->np" + elif Reduction(reduction) == Reduction.TRACE: + subscript = "pk,pk,nk->n" + else: + Reduction(reduction).raise_invalid_representation() + + # Contract the moments + left, right = self.unpack_couplings() + moments = util.einsum( + subscript, + right, + left.conj(), + self.energies[None] ** orders[:, None], + ) + if squeeze: + moments = moments[0] + + return moments + + moment = moments + + def chebyshev_moments( + self, + order: int | Iterable[int], + scaling: tuple[float, float] | None = None, + scale_couplings: bool = False, + ) -> Array: + r"""Calculate the Chebyshev polynomial moment(s) of the Lehmann representation. + + The Chebyshev moments are defined as + + .. math:: + T_{pq}^{n} = \sum_{k} v_{pk} u_{qk}^* P_n(\epsilon_k), + + where :math:`P_n(x)` is the Chebyshev polynomial of order :math:`n`. + + Args: + order: The order(s) of the moment(s). + scaling: Scaling factors to ensure the energy scale of the Lehmann representation is + in ``[-1, 1]``. The scaling is applied as ``(energies - scaling[1]) / scaling[0]``. + If ``None``, the default scaling is computed as + ``(max(energies) - min(energies)) / (2.0 - 1e-3)`` and + ``(max(energies) + min(energies)) / 2.0``, respectively. + scale_couplings: Scale the couplings as well as the energy spectrum. This is generally + necessary for Chebyshev moments of a self-energy, but not for a Green's function. + + Returns: + The Chebyshev polynomial moment(s) of the Lehmann representation. + """ + if scaling is None: + scaling = util.get_chebyshev_scaling_parameters( + self.energies.min(), self.energies.max() + ) + squeeze = False + if isinstance(order, int): + order = [order] + squeeze = True + max_order = max(order) + orders = set(order) + + # Scale the spectrum + left, right = self.unpack_couplings() + energies = (self.energies - scaling[1]) / scaling[0] + if scale_couplings: + left = left / scaling[0] + right = right / scaling[0] + + # Calculate the Chebyshev moments + moments = np.zeros((len(orders), self.nphys, self.nphys), dtype=self.dtype) + vecs = (right, right * energies[None]) + idx = 0 + if 0 in orders: + moments[idx] = vecs[0] @ left.T.conj() + idx += 1 + if 1 in orders: + moments[idx] = vecs[1] @ left.T.conj() + idx += 1 + for i in range(2, max_order + 1): + vecs = (vecs[1], 2 * energies * vecs[1] - vecs[0]) + if i in orders: + moments[idx] = vecs[1] @ left.T.conj() + idx += 1 + if squeeze: + moments = moments[0] + + return moments + + chebyshev_moment = chebyshev_moments + + # Methods associated with the supermatrix: + + def matrix(self, physical: Array, chempot: bool | float = False) -> Array: + r"""Build the dense supermatrix form of the Lehmann representation. + + The supermatrix is defined as + + .. math:: + \begin{bmatrix} + \mathbf{f} & \mathbf{v} \\ + \mathbf{u}^\dagger & \boldsymbol{\epsilon} \mathbf{I} + \end{bmatrix}, + + where :math:`\mathbf{f}` is the physical space part of the supermatrix, provided as an + argument. + + Args: + physical: The matrix to use for the physical space part of the supermatrix. + chempot: Whether to include the chemical potential in the supermatrix. If `True`, the + chemical potential from :attr:`chempot` is used. If a float is given, that value is + used. + + Returns: + The supermatrix form of the Lehmann representation. + """ + energies = self.energies + left, right = self.unpack_couplings() + if chempot: + if chempot is True: + chempot = self.chempot + energies -= chempot + + # If there are no auxiliary states, return the physical matrix + if self.naux == 0: + return physical + + # Build the supermatrix + matrix = np.block([[physical, right], [left.T.conj(), np.diag(energies)]]) + + return matrix + + def diagonal(self, physical: Array, chempot: bool | float = False) -> Array: + r"""Build the diagonal supermatrix form of the Lehmann representation. + + The diagonal supermatrix is defined as + + .. math:: + \begin{bmatrix} \mathrm{diag}(\mathbf{f}) & \boldsymbol{\epsilon} \end{bmatrix}, + + where :math:`\mathbf{f}` is the physical space part of the supermatrix, provided as an + argument. + + Args: + physical: The matrix to use for the physical space part of the supermatrix. + chempot: Whether to include the chemical potential in the supermatrix. If `True`, the + chemical potential from :attr:`chempot` is used. If a float is given, that value is + used. + + Returns: + The diagonal supermatrix form of the Lehmann representation. + """ + energies = self.energies + if chempot: + if chempot is True: + chempot = self.chempot + energies -= chempot + + # Build the supermatrix diagonal + diagonal = np.concatenate((np.diag(physical), energies)) + + return diagonal + + def matvec(self, physical: Array, vector: Array, chempot: bool | float = False) -> Array: + r"""Apply the supermatrix to a vector. + + The matrix-vector product is defined as + + .. math:: + \begin{bmatrix} + \mathbf{x}_\mathrm{phys} \\ + \mathbf{x}_\mathrm{aux} + \end{bmatrix} + = + \begin{bmatrix} + \mathbf{f} & \mathbf{v} \\ + \mathbf{u}^\dagger & \mathbf{\epsilon} \mathbf{I} + \end{bmatrix} + \begin{bmatrix} + \mathbf{r}_\mathrm{phys} \\ + \mathbf{r}_\mathrm{aux} + \end{bmatrix}, + + where :math:`\mathbf{f}` is the physical space part of the supermatrix, and the input + vector :math:`\mathbf{r}` is spans both the physical and auxiliary spaces. + + Args: + physical: The matrix to use for the physical space part of the supermatrix. + vector: The vector to apply the supermatrix to. + chempot: Whether to include the chemical potential in the supermatrix. If `True`, the + chemical potential from :attr:`chempot` is used. If a float is given, that value is + used. + + Returns: + The result of applying the supermatrix to the vector. + """ + left, right = self.unpack_couplings() + energies = self.energies + if chempot: + if chempot is True: + chempot = self.chempot + energies -= chempot + if vector.shape[0] != (self.nphys + self.naux): + raise ValueError( + f"Vector shape {vector.shape} does not match supermatrix shape " + f"{(self.nphys + self.naux, self.nphys + self.naux)}" + ) + + # Contract the supermatrix + vector_phys, vector_aux = np.split(vector, [self.nphys]) + result_phys = util.einsum("pq,q...->p...", physical, vector_phys) + result_phys += util.einsum("pk,k...->p...", right, vector_aux) + result_aux = util.einsum("pk,p...->k...", left.conj(), vector_phys) + result_aux += util.einsum("k,k...->k...", energies, vector_aux) + result = np.concatenate((result_phys, result_aux), axis=0) + + return result + + def diagonalise_matrix( + self, physical: Array, chempot: bool | float = False, overlap: Array | None = None + ) -> tuple[Array, Array]: + r"""Diagonalise the supermatrix. + + The eigenvalue problem is defined as + + .. math:: + \begin{bmatrix} + \mathbf{f} & \mathbf{v} \\ + \mathbf{u}^\dagger & \mathbf{\epsilon} \mathbf{1} + \end{bmatrix} + \begin{bmatrix} + \mathbf{x}_\mathrm{phys} \\ + \mathbf{x}_\mathrm{aux} + \end{bmatrix} + = + E + \begin{bmatrix} + \mathbf{x}_\mathrm{phys} \\ + \mathbf{x}_\mathrm{aux} + \end{bmatrix}, + + where :math:`\mathbf{f}` is the physical space part of the supermatrix, and the eigenvectors + :math:`\mathbf{x}` span both the physical and auxiliary spaces. + + Args: + physical: The matrix to use for the physical space part of the supermatrix. + chempot: Whether to include the chemical potential in the supermatrix. If ``True``, the + chemical potential from :attr:`chempot` is used. If a float is given, that value is + used. + overlap: The overlap matrix to use for the physical space part of the supermatrix. If + ``None``, the identity matrix is used. + + Returns: + The eigenvalues and eigenvectors of the supermatrix. + + Note: + If a non-identity overlap matrix is provided, this is equivalent to performing a + generalised eigenvalue decomposition of the supermatrix, with the overlap in the + auxiliary space assumed to be the identity. + """ + # Orthogonalise the physical space if overlap is provided + lehmann = self + if overlap is not None: + orth = util.matrix_power(overlap, -0.5, hermitian=False)[0] + unorth = util.matrix_power(overlap, 0.5, hermitian=False)[0] + physical = orth @ physical @ orth + lehmann = lehmann.rotate_couplings(orth if self.hermitian else (orth, orth.T.conj())) + + # Get the chemical potential + if chempot is True: + chempot = self.chempot + else: + chempot = float(chempot) + + # Diagonalise the supermatrix + matrix = lehmann.matrix(physical, chempot=chempot) + if self.hermitian: + eigvals, eigvecs = util.eig(matrix, hermitian=True) + if overlap is not None: + eigvecs = util.rotate_subspace(eigvecs, unorth.T.conj()) + else: + eigvals, eigvecs_tuple = util.eig_lr(matrix, hermitian=False) + if overlap is not None: + left, right = eigvecs_tuple + left = util.rotate_subspace(left, unorth.T.conj()) + right = util.rotate_subspace(right, unorth) + eigvecs_tuple = (left, right) + eigvecs = np.array(eigvecs_tuple) + + return eigvals, eigvecs + + def diagonalise_matrix_with_projection( + self, physical: Array, chempot: bool | float = False, overlap: Array | None = None + ) -> tuple[Array, Array]: + r"""Diagonalise the supermatrix and project the eigenvectors into the physical space. + + The projection of the eigenvectors is + + .. math:: + \mathbf{x}_\mathrm{phys} = \mathbf{P}_\mathrm{phys} \mathbf{x}, + + where :math:`\mathbf{P}_\mathrm{phys}` is the projection operator onto the physical space, + which can be written as + + .. math:: + \mathbf{P}_\mathrm{phys} = \begin{bmatrix} \mathbf{I} & 0 \\ 0 & 0 \end{bmatrix}, + + within the supermatrix block structure of :meth:`matrix`. + + Args: + physical: The matrix to use for the physical space part of the supermatrix. + chempot: Whether to include the chemical potential in the supermatrix. If ``True``, the + chemical potential from :attr:`chempot` is used. If a float is given, that value is + used. + overlap: The overlap matrix to use for the physical space part of the supermatrix. If + ``None``, the identity matrix is used. + + Returns: + The eigenvalues and eigenvectors of the supermatrix, with the eigenvectors projected + into the physical space. + + See Also: + :meth:`diagonalise_matrix` for the full eigenvalue decomposition of the supermatrix. + """ + eigvals, eigvecs = self.diagonalise_matrix(physical, chempot=chempot, overlap=overlap) + eigvecs_projected = eigvecs[..., : self.nphys, :] + return eigvals, eigvecs_projected + + # Methods associated with a quasiparticle representation: + + def weights(self, occupancy: float = 1.0) -> Array: + r"""Get the weights of the residues in the Lehmann representation. + + The weights are defined as + + .. math:: + w_k = \sum_{p} v_{pk} u_{pk}^*, + + where :math:`w_k` is the weight of residue :math:`k`. + + Args: + occupancy: The occupancy of the states. + + Returns: + The weights of each state. + """ + left, right = self.unpack_couplings() + weights = util.einsum("pk,pk->k", right, left.conj()) * occupancy + return weights + + def as_orbitals( + self, occupancy: float = 1.0, mo_coeff: Array | None = None + ) -> tuple[ + Array, + Array, + Array, + ]: + """Convert the Lehmann representation to an orbital representation. + + Args: + occupancy: The occupancy of the states. + mo_coeff: The molecular orbital coefficients. If given, the couplings will have their + physical dimension rotated into the AO basis according to these coefficients. + + Returns: + The energies, coefficients, and occupancies of the states. + + Note: + This representation is intended to be compatible with PySCF's mean-field representation + of molecular orbitals. + """ + if not self.hermitian: + raise NotImplementedError("Cannot convert non-Hermitian system orbitals.") + energies = self.energies + couplings, _ = self.unpack_couplings() + coeffs = couplings if mo_coeff is None else mo_coeff @ couplings + occupancies = np.concatenate( + [ + np.abs(self.occupied().weights(occupancy=occupancy)), + np.zeros(self.virtual().naux), + ] + ) + return energies, coeffs, occupancies + + def as_perturbed_mo_energy(self) -> Array: + r"""Return an array of :math:`N_\mathrm{phys}` pole energies according to best overlap. + + The pole energies are selected as + + .. math:: + \epsilon_p = \epsilon_k \quad \text{where} \quad k = \arg\max_{k} |v_{pk} u_{pk}^*|, + + where :math:`\epsilon_p` is the energy of the physical state :math:`p`, and :math:`k` is the + index of a pole in the Lehmann representation. + + Returns: + The selected energies. + + Note: + The return value of this function is intended to be compatible with + :attr:`pyscf.scf.hf.SCF.mo_energy`, i.e. it represents a reduced quasiparticle picture + consisting of :math:`N_\mathrm{phys}` energies that are picked from the poles of the + Lehmann representation, according to the best overlap with the MO of the same index. + """ + left, right = self.unpack_couplings() + weights = right * left.conj() + energies = [self.energies[np.argmax(np.abs(weights[i]))] for i in range(self.nphys)] + return np.asarray(energies) + + # Methods associated with a static approximation to a self-energy: + + def as_static_potential(self, mo_energy: Array, eta: float = 1e-2) -> Array: + r"""Convert the Lehmann representation to a static potential. + + The static potential is defined as + + .. math:: + V_{pq} = \mathrm{Re}\left[ \sum_{k} \frac{v_{pk} u_{qk}^*}{\epsilon_p - \epsilon_k + \pm i \eta} \right]. + + Args: + mo_energy: The molecular orbital energies. + eta: The broadening parameter. + + Returns: + The static potential. + + Note: + The static potential in this format is common in methods such as quasiparticle + self-consistent :math:`GW` calculations. + """ + left, right = self.unpack_couplings() + energies = self.energies + np.sign(self.energies - self.chempot) * 1.0j * eta + denom = mo_energy[:, None] - energies[None] + + # Calculate the static potential + static = util.einsum("pk,qk,pk->pq", right, left.conj(), 1.0 / denom).real + static = 0.5 * (static + static.T) + + return static + + # Methods for combining Lehmann representations: + + def split_physical(self, nocc: int) -> tuple[Lehmann, Lehmann]: + """Split the physical domain of Lehmann representation into occupied and virtual parts. + + Args: + nocc: The number of occupied states. + + Returns: + The Lehmann representation coupled with the occupied and virtual parts, as separate + Lehmann representations. + + Note: + The Fermi level (value at which the parts are separated) is defined by the chemical + potential :attr:`chempot`. + """ + occ = self.__class__( + self.energies, + self.couplings[..., :nocc, :], + chempot=self.chempot, + sort=False, + ) + vir = self.__class__( + self.energies, + self.couplings[..., nocc:, :], + chempot=self.chempot, + sort=False, + ) + return occ, vir + + def combine_physical(self, other: Lehmann) -> Lehmann: + """Combine the physical domain of two Lehmann representations. + + Args: + other: The other Lehmann representation to combine with. + + Returns: + A new Lehmann representation that is the combination of the two. + + Raises: + ValueError: If the two representations have different chemical potentials. + """ + if not np.isclose(self.chempot, other.chempot): + raise ValueError( + f"Cannot combine Lehmann representations with different chemical potentials: " + f"{self.chempot} and {other.chempot}" + ) + + # Combine the energies and couplings + energies = np.concatenate((self.energies, other.energies), axis=0) + if self.hermitian and other.hermitian: + couplings = scipy.linalg.block_diag(self.couplings, other.couplings) + else: + left_self, right_self = self.unpack_couplings() + left_other, right_other = other.unpack_couplings() + couplings = np.array( + [ + np.concatenate((left_self, left_other), axis=-1), + np.concatenate((right_self, right_other), axis=-1), + ] + ) + + return self.__class__(energies, couplings, chempot=self.chempot, sort=True) + + def concatenate(self, other: Lehmann) -> Lehmann: + """Concatenate two Lehmann representations. + + Args: + other: The other Lehmann representation to concatenate. + + Returns: + A new Lehmann representation that is the concatenation of the two. + + Raises: + ValueError: If the two representations have different physical dimensions or chemical + potentials. + """ + if self.nphys != other.nphys: + raise ValueError( + f"Cannot combine Lehmann representations with different physical dimensions: " + f"{self.nphys} and {other.nphys}" + ) + if not np.isclose(self.chempot, other.chempot): + raise ValueError( + f"Cannot combine Lehmann representations with different chemical potentials: " + f"{self.chempot} and {other.chempot}" + ) + + # Combine the energies and couplings + energies = np.concatenate((self.energies, other.energies)) + if self.hermitian and other.hermitian: + couplings = np.concatenate((self.couplings, other.couplings), axis=-1) + else: + left_self, right_self = self.unpack_couplings() + left_other, right_other = other.unpack_couplings() + couplings = np.array( + [ + np.concatenate((left_self, left_other), axis=-1), + np.concatenate((right_self, right_other), axis=-1), + ] + ) + + return self.__class__(energies, couplings, chempot=self.chempot, sort=False) + + def __eq__(self, other: object) -> bool: + """Check if two spectral representations are equal.""" + if not isinstance(other, Lehmann): + return NotImplemented + if other.nphys != self.nphys: + return False + if other.naux != self.naux: + return False + if other.hermitian != self.hermitian: + return False + if other.chempot != self.chempot: + return False + return np.allclose(other.energies, self.energies) and ( + np.allclose(other.couplings, self.couplings) + ) + + def __hash__(self) -> int: + """Return a hash of the Lehmann representation.""" + return hash((tuple(self.energies), tuple(self.couplings.flatten()), self.chempot)) diff --git a/dyson/representations/representation.py b/dyson/representations/representation.py new file mode 100644 index 0000000..c3da2d3 --- /dev/null +++ b/dyson/representations/representation.py @@ -0,0 +1,25 @@ +"""Base class for representations.""" + +from __future__ import annotations + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +if TYPE_CHECKING: + pass + + +class BaseRepresentation(ABC): + """Base class for representations.""" + + @property + @abstractmethod + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + pass + + @property + @abstractmethod + def hermitian(self) -> bool: + """Get a boolean indicating if the system is Hermitian.""" + pass diff --git a/dyson/representations/spectral.py b/dyson/representations/spectral.py new file mode 100644 index 0000000..a02c2bb --- /dev/null +++ b/dyson/representations/spectral.py @@ -0,0 +1,347 @@ +"""Container for an spectral representation (eigenvalues and eigenvectors) of a matrix.""" + +from __future__ import annotations + +from functools import cached_property +from typing import TYPE_CHECKING + +from dyson import numpy as np +from dyson import util +from dyson.representations.lehmann import Lehmann +from dyson.representations.representation import BaseRepresentation + +if TYPE_CHECKING: + from dyson.typing import Array + + +class Spectral(BaseRepresentation): + r"""Spectral representation matrix with a known number of physical degrees of freedom. + + The eigendecomposition (spectral decomposition) of a matrix consists of the eigenvalues + :math:`\lambda_k` and eigenvectors :math:`v_{pk}` that represent the matrix as + + .. math:: + \sum_{k} \lambda_k v_{pk} u_{qk}^*, + + where the eigenvectors have right-handed components :math:`v` and left-handed components + :math:`u`. + + Note that the order of eigenvectors is ``(left, right)``, whilst they act in the order + ``(right, left)`` in the above equation. The naming convention is chosen to be consistent with + the eigenvalue decomposition, where :math:`v` may be an eigenvector acting on the right of a + matrix, and :math:`u` is an eigenvector acting on the left of a matrix. + """ + + def __init__( + self, + eigvals: Array, + eigvecs: Array, + nphys: int, + sort: bool = False, + chempot: float | None = None, + ): + """Initialise the object. + + Args: + eigvals: Eigenvalues of the matrix. + eigvecs: Eigenvectors of the matrix. + nphys: Number of physical degrees of freedom. + sort: Sort the eigenfunctions by eigenvalue. + chempot: Chemical potential to be used in the Lehmann representations of the self-energy + and Green's function. + """ + self._eigvals = eigvals + self._eigvecs = eigvecs + self._nphys = nphys + self.chempot = chempot + if sort: + self.sort_() + if not self.hermitian: + if eigvecs.ndim != 3: + raise ValueError( + f"Couplings must be 3D for a non-Hermitian system, but got {eigvecs.ndim}D." + ) + if eigvecs.shape[0] != 2: + raise ValueError( + f"Couplings must have shape (2, nphys, naux) for a non-Hermitian system, " + f"but got {eigvecs.shape}." + ) + + @classmethod + def from_matrix( + cls, matrix: Array, nphys: int, hermitian: bool = True, chempot: float | None = None + ) -> Spectral: + """Create a spectrum from a matrix by diagonalising it. + + Args: + matrix: Matrix to diagonalise. + nphys: Number of physical degrees of freedom. + hermitian: Whether the matrix is Hermitian. + chempot: Chemical potential to be used in the Lehmann representations of the self-energy + and Green's function. + + Returns: + Spectrum object. + """ + if hermitian: + eigvals, eigvecs = util.eig(matrix, hermitian=True) + else: + eigvals, (left, right) = util.eig_lr(matrix, hermitian=False) + eigvecs = np.array([left, right]) + return cls(eigvals, eigvecs, nphys, chempot=chempot) + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + ) -> Spectral: + """Create a spectrum from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + + Returns: + Spectrum object. + """ + return cls( + *self_energy.diagonalise_matrix(static, overlap=overlap), + self_energy.nphys, + chempot=self_energy.chempot, + ) + + def sort_(self) -> None: + """Sort the eigenfunctions by eigenvalue. + + Note: + The object is sorted in place. + """ + idx = np.argsort(self.eigvals) + self._eigvals = self.eigvals[idx] + self._eigvecs = self.eigvecs[..., idx] + + def _get_matrix_block(self, slices: tuple[slice, slice]) -> Array: + """Get a block of the matrix. + + Args: + slices: Slices to select. + + Returns: + Block of the matrix. + """ + left, right = util.unpack_vectors(self.eigvecs) + return util.einsum("pk,qk,k->pq", right[slices[0]], left[slices[1]].conj(), self.eigvals) + + def get_static_self_energy(self) -> Array: + """Get the static part of the self-energy. + + Returns: + Static self-energy. + + Note: + The static part of the self-energy is defined as the physical space part of the matrix + from which the spectrum is derived. + """ + return self._get_matrix_block((slice(self.nphys), slice(self.nphys))) + + def get_auxiliaries(self) -> tuple[Array, Array]: + """Get the auxiliary energies and couplings contributing to the dynamic self-energy. + + Returns: + Auxiliary energies and couplings. + + Note: + The auxiliary energies are the eigenvalues of the auxiliary subspace, and the couplings + are the eigenvectors projected back to the auxiliary subspace using the + physical-auxiliary block of the matrix from which the spectrum is derived. + """ + phys = slice(None, self.nphys) + aux = slice(self.nphys, None) + + # Project back to the auxiliary subspace + subspace = self._get_matrix_block((aux, aux)) + + # If there are no auxiliaries, return here + if subspace.size == 0: + energies = np.empty((0)) + couplings = np.empty((self.nphys, 0) if self.hermitian else (2, self.nphys, 0)) + return energies, couplings + + # Diagonalise the subspace to get the energies and basis for the couplings + energies, rotation = util.eig_lr(subspace, hermitian=self.hermitian) + + # Project back to the couplings + couplings_right = self._get_matrix_block((phys, aux)) + if not self.hermitian: + couplings_left = self._get_matrix_block((aux, phys)).T.conj() + + # Rotate the couplings to the auxiliary basis + if self.hermitian: + couplings = couplings_right @ rotation[0] + else: + couplings = np.array([couplings_left @ rotation[0], couplings_right @ rotation[1]]) + + return energies, couplings + + def get_dyson_orbitals(self) -> tuple[Array, Array]: + """Get the Dyson orbitals. + + Returns: + Dyson orbitals. + """ + return self.eigvals, self.eigvecs[..., : self.nphys, :] + + def get_overlap(self) -> Array: + """Get the overlap matrix in the physical space. + + Returns: + Overlap matrix. + + Note: + The overlap matrix is defined as the zeroth moment of the Green's function, and is given + by the inner product of the Dyson orbitals. + """ + _, orbitals = self.get_dyson_orbitals() + left, right = util.unpack_vectors(orbitals) + return util.einsum("pk,qk->pq", right, left.conj()) + + def get_self_energy(self, chempot: float | None = None) -> Lehmann: + """Get the Lehmann representation of the self-energy. + + Args: + chempot: Chemical potential. + + Returns: + Lehmann representation of the self-energy. + """ + if chempot is None: + chempot = self.chempot + if chempot is None: + chempot = 0.0 + return Lehmann(*self.get_auxiliaries(), chempot=chempot) + + def get_greens_function(self, chempot: float | None = None) -> Lehmann: + """Get the Lehmann representation of the Green's function. + + Args: + chempot: Chemical potential. + + Returns: + Lehmann representation of the Green's function. + """ + if chempot is None: + chempot = self.chempot + if chempot is None: + chempot = 0.0 + return Lehmann(*self.get_dyson_orbitals(), chempot=chempot) + + def combine(self, *args: Spectral, chempot: float | None = None) -> Spectral: + """Combine multiple spectral representations. + + Args: + args: Spectral representations to combine. + chempot: Chemical potential to be used in the Lehmann representations of the self-energy + and Green's function. + + Returns: + Combined spectral representation. + """ + # TODO: just concatenate the eigenvectors...? + args = (self, *args) + if len(set(arg.nphys for arg in args)) != 1: + raise ValueError( + "All Spectral objects must have the same number of physical degrees of freedom." + ) + nphys = args[0].nphys + + # Sum the overlap and static self-energy matrices -- double counting is not an issue + # with shared static parts because the overlap matrix accounts for the separation + static = sum([arg.get_static_self_energy() for arg in args], np.zeros((nphys, nphys))) + overlap = sum([arg.get_overlap() for arg in args], np.zeros((nphys, nphys))) + + # Check the chemical potentials + if chempot is None: + if any(arg.chempot is not None for arg in args): + chempots = [arg.chempot for arg in args if arg.chempot is not None] + if not all(np.isclose(chempots[0], part) for part in chempots[1:]): + raise ValueError( + "If not chempot is passed to combine, all chemical potentials must be " + "equal in the inputs." + ) + chempot = chempots[0] + + # Get the auxiliaries + energies = np.zeros((0)) + left = np.zeros((nphys, 0)) + right = np.zeros((nphys, 0)) + for arg in args: + energies_i, couplings_i = arg.get_auxiliaries() + energies = np.concatenate([energies, energies_i]) + if arg.hermitian: + left = np.concatenate([left, couplings_i], axis=1) + else: + left_i, right_i = util.unpack_vectors(couplings_i) + left = np.concatenate([left, left_i], axis=1) + right = np.concatenate([right, right_i], axis=1) + couplings = np.array([left, right]) if not args[0].hermitian else left + + # Solve the eigenvalue problem + self_energy = Lehmann(energies, couplings) + result = Spectral( + *self_energy.diagonalise_matrix(static, overlap=overlap), nphys, chempot=chempot + ) + + return result + + @cached_property + def overlap(self) -> Array: + """Get the overlap matrix (the zeroth moment of the Green's function).""" + orbitals = self.get_dyson_orbitals()[1] + left, right = util.unpack_vectors(orbitals) + return util.einsum("pk,qk->pq", right, left.conj()) + + @property + def eigvals(self) -> Array: + """Get the eigenvalues.""" + return self._eigvals + + @property + def eigvecs(self) -> Array: + """Get the eigenvectors.""" + return self._eigvecs + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return self._nphys + + @property + def neig(self) -> int: + """Get the number of eigenvalues.""" + return self.eigvals.shape[0] + + @property + def hermitian(self) -> bool: + """Check if the spectrum is Hermitian.""" + return self.eigvecs.ndim == 2 + + def __eq__(self, other: object) -> bool: + """Check if two Lehmann representations are equal.""" + if not isinstance(other, Spectral): + return NotImplemented + if other.nphys != self.nphys: + return False + if other.neig != self.neig: + return False + if other.hermitian != self.hermitian: + return False + if other.chempot != self.chempot: + return False + return np.allclose(other.eigvals, self.eigvals) and np.allclose(other.eigvecs, self.eigvecs) + + def __hash__(self) -> int: + """Hash the object.""" + return hash((tuple(self.eigvals), tuple(self.eigvecs.flatten()), self.nphys, self.chempot)) diff --git a/dyson/solvers/__init__.py b/dyson/solvers/__init__.py index 3807718..3b9efb6 100644 --- a/dyson/solvers/__init__.py +++ b/dyson/solvers/__init__.py @@ -1,11 +1,75 @@ -from dyson.solvers.solver import BaseSolver -from dyson.solvers.exact import Exact -from dyson.solvers.davidson import Davidson -from dyson.solvers.downfolded import DiagonalDownfolded, Downfolded -from dyson.solvers.mblse import MBLSE, MixedMBLSE -from dyson.solvers.mblgf import MBLGF, MixedMBLGF -from dyson.solvers.kpmgf import KPMGF -from dyson.solvers.cpgf import CPGF -from dyson.solvers.chempot import AufbauPrinciple, AufbauPrincipleBisect, AuxiliaryShift -from dyson.solvers.density import DensityRelaxation -from dyson.solvers.self_consistent import SelfConsistent +r"""Solvers for solving the Dyson equation. + +Solvers generally provide some method to solving the Dyson equation + +.. math:: + \mathbf{G}(\omega) = \left( \left[ \mathbf{G}^0(\omega) \right]^{-1} - + \boldsymbol{\Sigma}(\omega) \right)^{-1}, + +which can also be written recursively as + +.. math:: + \mathbf{G}(\omega) = \mathbf{G}^0(\omega) + \mathbf{G}^0(\omega) \boldsymbol{\Sigma}(\omega) + \mathbf{G}(\omega), + +and can be expressed as an eigenvalue problem as + +.. math:: + \begin{bmatrix} \boldsymbol{\Sigma}(\omega) & \mathbf{v} \\ \mathbf{v}^\dagger & \mathbf{K} + + \mathbf{C} \end{bmatrix} \mathbf{u} = \omega \mathbf{u}. + +For more details on the equivalence of these representations, see the :mod:`~dyson.representations` +module. + +The :class:`~dyson.solvers.solver.BaseSolver` interface provides the constructors +:func:`~dyson.solvers.solver.BaseSolver.from_expression` and +:func:`~dyson.solvers.solver.BaseSolver.from_self_energy` to create a solver of that type from +either an instance of a subclass of :class:`~dyson.expressions.expression.BaseExpression` or a +self-energy in the form of an instance of :class:`~dyson.representations.lehmann.Lehmann` object, +respectively + +>>> from dyson import util, quiet, CCSD, Exact +>>> quiet() # Suppress output +>>> mf = util.get_mean_field("H 0 0 0; H 0 0 1", "6-31g") +>>> ccsd = CCSD.h.from_mf(mf) +>>> solver = Exact.from_expression(ccsd) + +Solvers can be run by calling the :meth:`~dyson.solvers.solver.BaseSolver.kernel` method, which +in the case of :mod:`~dyson.solvers.static` solvers sets the attribute and returns :attr:`result`, +an instance of :class:`~dyson.representations.spectral.Spectral` + +>>> result = solver.kernel() +>>> type(result) + + +The result can then be used to construct Lehmann representations of the Green's function and +self-energy, details of which can be found in the :mod:`~dyson.representations` module. On the other +hand, solvers in :mod:`~dyson.solvers.dynamic` return an instance of +:class:`~dyson.representations.dynamic.Dynamic`, which contains the dynamic Green's function in the +format requested by the solver arguments. + +A list of available solvers is provided in the documentation of :mod:`dyson`, along with their +expected inputs. + + +Submodules +---------- + +.. autosummary:: + :toctree: + + solver + static + dynamic + +""" + +from dyson.solvers.static.exact import Exact +from dyson.solvers.static.davidson import Davidson +from dyson.solvers.static.downfolded import Downfolded +from dyson.solvers.static.mblse import MBLSE +from dyson.solvers.static.mblgf import MBLGF +from dyson.solvers.static.chempot import AufbauPrinciple, AuxiliaryShift +from dyson.solvers.static.density import DensityRelaxation +from dyson.solvers.dynamic.corrvec import CorrectionVector +from dyson.solvers.dynamic.cpgf import CPGF diff --git a/dyson/solvers/chempot.py b/dyson/solvers/chempot.py deleted file mode 100644 index fd18ab3..0000000 --- a/dyson/solvers/chempot.py +++ /dev/null @@ -1,356 +0,0 @@ -""" -Chemical potential optimisation. -""" - -import numpy as np -import scipy.optimize - -from dyson import NullLogger -from dyson.lehmann import Lehmann -from dyson.solvers import MBLGF, MBLSE, BaseSolver - - -class AufbauPrinciple(BaseSolver): - """ - Fill a series of orbitals according to the Aufbau principle. - - Parameters - ---------- - *args : tuple - Input arguments. Either `(gf, nelec)` or `(fock, se, nelec)` - where `gf` is the Lehmann representation of the Green's - function, `fock` is the Fock matrix, `se` is the Lehmann - representation of the self-energy and `nelec` is the number - of electrons in the physical space. - occupancy : int, optional - Occupancy of each state, i.e. `2` for a restricted reference - and `1` for other references. Default value is `2`. - """ - - # Default parameters: - occupancy = 2 - - def __init__(self, *args, **kwargs): - # Input: - if len(args) == 2: - gf, nelec = args - self.se = None - self.gf = Lehmann.from_pyscf(gf) - self.nelec = nelec - else: - fock, se, nelec = args - self.se = Lehmann.from_pyscf(se) - w, v = self.se.diagonalise_matrix_with_projection(fock) - self.gf = Lehmann(w, v, chempot=se.chempot) - self.nelec = nelec - - # Parameters: - self.occupancy = kwargs.pop("occupancy", self.occupancy) - - # Base class: - super().__init__(self, **kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > occupancy: %s", self.occupancy) - - # Caching: - self.converged = False - self.homo = None - self.lumo = None - self.chempot = None - self.error = None - - def _kernel(self): - energies = self.gf.energies - couplings_l, couplings_r = self.gf._unpack_couplings() - - sum0 = sum1 = 0.0 - for i in range(self.gf.naux): - n = np.dot(couplings_l[:, i], couplings_r[:, i].conj()).real - n *= self.occupancy - sum0, sum1 = sum1, sum1 + n - - self.log.debug("Number of electrons [0:%d] = %.6f", i + 1, sum1) - - if i: - if sum0 <= self.nelec and self.nelec <= sum1: - break - - if abs(sum0 - self.nelec) < abs(sum1 - self.nelec): - homo = i - 1 - error = self.nelec - sum0 - else: - homo = i - error = self.nelec - sum1 - - try: - lumo = homo + 1 - chempot = 0.5 * (energies[homo] + energies[lumo]) - except: - raise ValueError("Failed to find Fermi energy.") - - self.log.info("HOMO = %.6f", energies[homo]) - self.log.info("LUMO = %.6f", energies[lumo]) - self.log.info("Chemical potential = %.6f", chempot) - self.log.info("Error in nelec = %.3g", error) - - self.converged = True - self.homo = energies[homo] - self.lumo = energies[lumo] - self.chempot = chempot - self.error = error - - return chempot, error - - def get_auxiliaries(self): - if self.se is None: - raise ValueError("`AufbauPrinciple` was initialised with a Green's function.") - return self.se.energies, self.se.couplings - - def get_dyson_orbitals(self): - return self.gf.energies, self.gf.couplings - - def get_self_energy(self): - if self.se is None: - raise ValueError("`AufbauPrinciple` was initialised with a Green's function.") - return self.se.copy(chempot=self.chempot, deep=False) - - def get_greens_function(self): - return self.gf.copy(chempot=self.chempot, deep=False) - - -class AufbauPrincipleBisect(AufbauPrinciple): - """ - Fill a series of orbitals according to the Aufbau principle using a bisection algorithim. - - Parameters - ---------- - *args : tuple - Input arguments. Either `(gf, nelec)` or `(fock, se, nelec)` - where `gf` is the Lehmann representation of the Green's - function, `fock` is the Fock matrix, `se` is the Lehmann - representation of the self-energy and `nelec` is the number - of electrons in the physical space. - occupancy : int, optional - Occupancy of each state, i.e. `2` for a restricted reference - and `1` for other references. Default value is `2`. - """ - - def _kernel(self): - energies = self.gf.energies - weights = self.gf.weights() - low, high = 0, self.gf.naux - mid = high // 2 - for iter in range(100): - sum = self.occupancy * weights[:mid].sum() - self.log.debug("Number of electrons [0:%d] = %.6f", iter + 1, sum) - if sum < self.nelec: - low = mid - mid = mid + (high - low) // 2 - else: - high = mid - mid = mid - (high - low) // 2 - - if low == mid or high == mid: - break - - n_low, n_high = self.occupancy * weights[:low].sum(), self.occupancy * weights[:high].sum() - - if abs(n_low - self.nelec) < abs(n_high - self.nelec): - homo = low - 1 - error = self.nelec - n_low - else: - homo = high - 1 - error = self.nelec - n_high - - try: - lumo = homo + 1 - chempot = 0.5 * (energies[homo] + energies[lumo]) - except: - raise ValueError("Failed to find Fermi energy.") - - self.log.info("HOMO LUMO %s %s" % (homo, lumo)) - self.log.info("HOMO = %.6f", energies[homo]) - self.log.info("LUMO = %.6f", energies[lumo]) - self.log.info("Chemical potential = %.6f", chempot) - self.log.info("Error in nelec = %.3g", error) - - self.converged = True - self.homo = energies[homo] - self.lumo = energies[lumo] - self.chempot = chempot - self.error = error - - return chempot, error - - -class AuxiliaryShift(BaseSolver): - """ - Shift the self-energy auxiliaries with respect to the Green's - function, operating on a MBLSE or MBLGF solver. - - Parameters - ---------- - fock : numpy.ndarray - Fock matrix. - se : dyson.lehmann.Lehmann - Lehmann representation of the self-energy. - nelec : int - Number of electrons in the physical space. - occupancy : int, optional - Occupancy of each state, i.e. `2` for a restricted reference - and `1` for other references. Default value is `2`. - max_cycle : float, optional - Maximum number of iterations. Default value is `200`. - conv_tol : float, optional - Threshold for convergence in the number of electrons. Default - value is `1e-8`. - guess : float, optional - Initial guess for the shift. Default value is 0.0. - """ - - # Default parameters: - occupancy = 2 - max_cycle = 200 - conv_tol = 1e-8 - guess = 0.0 - - def __init__(self, fock, se, nelec, **kwargs): - # Input: - self.fock = fock - self.se = Lehmann.from_pyscf(se) - self.nelec = nelec - - # Parameters: - self.occupancy = kwargs.pop("occupancy", self.occupancy) - self.max_cycle = kwargs.pop("max_cycle", self.max_cycle) - self.conv_tol = kwargs.pop("conv_tol", self.conv_tol) - self.guess = kwargs.pop("guess", self.guess) - - # Base class: - super().__init__(self, **kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > occupancy: %s", self.occupancy) - self.log.info(" > max_cycle: %s", self.max_cycle) - self.log.info(" > conv_tol: %s", self.conv_tol) - self.log.info(" > guess: %s", self.guess) - - # Caching: - self.iteration = 0 - self.converged = False - self.shift = 0.0 - self.chempot = None - self.error = None - - def objective(self, x, fock=None, out=None): - """Objective function.""" - - if fock is None: - fock = self.fock - - e, c = self.se.diagonalise_matrix_with_projection(fock, chempot=np.ravel(x)[0], out=out) - gf = Lehmann(e, c) - - aufbau = AufbauPrinciple(gf, self.nelec, occupancy=self.occupancy, log=NullLogger()) - aufbau.conv_tol = self.conv_tol - aufbau.kernel() - - return aufbau.error**2 - - def gradient(self, x, fock=None, out=None): - """Gradient of the objective function.""" - - if fock is None: - fock = self.fock - - e, c = self.se.diagonalise_matrix(fock, chempot=np.ravel(x)[0], out=out) - if self.se.hermitian: - c_phys = c[: self.se.nphys] - gf = Lehmann(e, c_phys) - else: - c_phys = ( - c[: self.se.nphys], - np.linalg.inv(c).T.conj()[: self.se.nphys], - ) - gf = Lehmann(e, c_phys) - - aufbau = AufbauPrinciple(gf, self.nelec, occupancy=self.occupancy, log=NullLogger()) - aufbau.conv_tol = self.conv_tol - gf.chempot, error = aufbau.kernel() - - gf_occ = gf.occupied() - gf_vir = gf.virtual() - - nphys = gf.nphys - nocc = np.sum(gf.energies < gf.chempot) - - h1 = -np.dot(c[gf.nphys :, gf_occ.naux :].conj().T, c[gf.nphys :, : gf_occ.naux]) - z = h1 / (gf_vir.energies[:, None] - gf_occ.energies[None]) - - c_occ = np.dot(gf_vir.couplings, z) - d_rdm1 = np.dot(c_occ, c_occ.T.conj()) * 4.0 - - dif = np.trace(d_rdm1).real * error * self.occupancy - - return error**2, dif - - def callback(self, xk): - self.log.info("Iteration %d: Chemical potential = %.6f", self.iteration, xk) - self.iteration += 1 - - def _kernel(self): - opt = scipy.optimize.minimize( - self.gradient, - x0=self.guess, - method="TNC", - jac=True, - options=dict( - maxfun=self.max_cycle, - ftol=self.conv_tol**2, - xtol=0, - gtol=0, - ), - callback=self.callback, - ) - - shift = -opt.x - se = self.se.copy() - se.energies += shift - - aufbau = AufbauPrinciple( - self.fock, se, self.nelec, occupancy=self.occupancy, log=NullLogger() - ) - aufbau.conv_tol = self.conv_tol - aufbau.kernel() - - self.log.info("Auxiliary shift = %.6f", shift) - self.log.info("Chemical potential = %.6f", aufbau.chempot) - self.log.info("Error in nelec = %.3g", aufbau.error) - self.flag_convergence(opt.success) - - self.converged = opt.success - self.shift = shift - self.chempot = aufbau.chempot - self.error = aufbau.error - - return self.chempot, aufbau.error - - def get_auxiliaries(self): - return self.se.energies, self.se.couplings - - def get_dyson_orbitals(self): - return self.gf.energies, self.gf.couplings - - def get_self_energy(self): - se = self.se.copy(chempot=self.chempot, deep=False) - se.energies = se.energies.copy() + self.shift - return se - - def get_greens_function(self): - se = self.get_self_energy() - w, v = se.diagonalise_matrix_with_projection(self.fock) - gf = Lehmann(w, v, chempot=self.chempot) - return gf diff --git a/dyson/solvers/cpgf.py b/dyson/solvers/cpgf.py deleted file mode 100644 index 1e9a6b7..0000000 --- a/dyson/solvers/cpgf.py +++ /dev/null @@ -1,136 +0,0 @@ -""" -Chebyshev polynomial Green's function method, similar to the KPMGF -and also conserves Chebyshev moments of the Green's function. - -Ref: https://doi.org/10.1103/PhysRevLett.115.106601 -""" - -import numpy as np -import scipy.integrate - -from dyson import util -from dyson.solvers import BaseSolver -from dyson.solvers.kpmgf import as_trace - - -class CPGF(BaseSolver): - """ - Chebyshev polynomial Green's function method. - - Input - ----- - moments : numpy.ndarray - Chebyshev moments of the Green's function. - grid : numpy.ndarray - Real-valued frequency grid to plot the spectral function upon. - scale : tuple of int - Scaling parameters used to scale the spectrum to [-1, 1], - given as `(a, b)` where - - a = (ωmax - ωmin) / (2 - ε) - b = (ωmax + ωmin) / 2 - - where ωmax and ωmin are the maximum and minimum energies in - the spectrum, respectively, and ε is a small number shifting - the spectrum values away from the boundaries. - trace : bool, optional - Whether to compute the trace of the Green's function. If - `False`, the entire Green's function is computed. Default - value is `True`. - include_real : bool, optional - Whether to include the real part of the Green's function in - the computation. Default value is `False`. - - Parameters - ---------- - max_cycle : int, optional - Maximum number of iterations. If `None`, perform as many as - the inputted number of moments permits. Default value is - `None`. - eta : float, optional - Regularisation parameter. Default value is 0.1. - """ - - def __init__(self, moments, grid, scale, **kwargs): - # Input: - self.moments = moments - self.grid = grid - self.scale = scale - - # Parameters - self.max_cycle = kwargs.pop("max_cycle", None) - # self.hermitian = True - self.eta = kwargs.pop("eta", 0.1) - self.trace = kwargs.pop("trace", True) - self.include_real = kwargs.pop("include_real", False) - - max_cycle_limit = len(moments) - 1 - if self.max_cycle is None: - self.max_cycle = max_cycle_limit - if self.max_cycle > max_cycle_limit: - raise ValueError( - "`max_cycle` cannot be more than the number of inputted moments minus one." - ) - - # Base class: - super().__init__(**kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > max_cycle: %s", self.max_cycle) - # self.log.info(" > hermitian: %s", self.hermitian) - self.log.info(" > grid: %s[%d]", type(self.grid), len(self.grid)) - self.log.info(" > scale: %s", scale) - self.log.info(" > eta: %s", self.eta) - self.log.info(" > trace: %s", self.trace) - self.log.info(" > include_real: %s", self.include_real) - - def initialise_recurrence(self): - self.log.info("-" * 21) - self.log.info("{:^4s} {:^16s}".format("Iter", "Integral")) - self.log.info("{:^4s} {:^16s}".format("-" * 4, "-" * 16)) - - def _kernel(self, iteration=None, trace=True): - self.initialise_recurrence() - - if iteration is None: - iteration = self.max_cycle - - filter_type = lambda arr: arr.imag if not self.include_real else arr - - # Get the moments - allow input to already be traced - if self.trace: - moments = as_trace(self.moments[: iteration + 1]).astype(complex) - else: - moments = self.moments[: iteration + 1].astype(complex) - - # Initialise scaled grids - a, b = self.scale - scaled_grid = (self.grid - b) / a - scaled_eta = self.eta / a - z = scaled_grid + 1.0j * scaled_eta - - # Initialise the Green's function - fac = lambda n: -1.0j * (2.0 - int(n == 0)) - num = z - 1.0j * np.sqrt(1.0 - z**2) - den = np.sqrt(1.0 - z**2) - gn = lambda n: fac(n) * num**n / den - gf = np.zeros((len(z), *moments[0].shape), dtype=complex) - - integral = scipy.integrate.simps(as_trace(gf.imag), self.grid) - - for niter in range(iteration + 1): - part = np.einsum("z,...->z...", gn(niter), moments[niter]) - part /= a * np.pi - gf -= part - - if niter in (0, 1, 2, 3, 4, 5, 10, iteration) or niter % 100 == 0: - integral = scipy.integrate.simps(as_trace(gf.imag), self.grid) - self.log.info("%4d %16.8g", niter, integral) - - # Not sure why we need to do this... - gf = -gf.conj() - - self.log.info("-" * 21) - - return filter_type(gf) diff --git a/dyson/solvers/davidson.py b/dyson/solvers/davidson.py deleted file mode 100644 index 6d6a4df..0000000 --- a/dyson/solvers/davidson.py +++ /dev/null @@ -1,217 +0,0 @@ -""" -Davidson eigensolver using the matrix-vector operation on the -upfolded self-energy. - -Interfaces pyscf.lib. -""" - -import warnings - -import numpy as np -from pyscf import lib - -from dyson import util -from dyson.solvers import BaseSolver - -# TODO abs picker - - -def pick_real_eigs(w, v, nroots, env, threshold=1e-3): - """Pick real eigenvalues, sorting by absolute value.""" - - iabs = np.abs(w.imag) - tol = max(threshold, np.sort(iabs)[min(w.size, nroots) - 1]) - idx = np.where(iabs <= tol)[0] - num = np.count_nonzero(iabs[idx] < threshold) - - if num < nroots and w.size >= nroots: - warnings.warn( - "Only %d eigenvalues (out of %3d requested roots) with imaginary part < %4.3g.\n" - % (num, min(w.size, nroots), threshold), - ) - - real_eigenvectors = env.get("dtype") == np.float64 - w, v, idx = lib.linalg_helper._eigs_cmplx2real(w, v, idx, real_eigenvectors=real_eigenvectors) - - mask = np.argsort(np.abs(w)) - w = w[mask] - v = v[:, mask] - - return w, v, 0 - - -class Davidson(BaseSolver): - """ - Davidson eigensolver using the matrix-vector operation on the - upfolded self-energy. - - Input - ----- - matvec : callable - Function returning the result of the dot-product of the - upfolded self-energy with an arbitrary state vector. Input - arguments are `vector`. - diagonal : numpy.ndarray (n,) - Diagonal entries of the upfolded self-energy to precondition - the solver. - - Parameters - ---------- - nroots : int, optional - Number of roots to solve for. Default value is `5`. - picker : callable, optional - Function to pick eigenvalues. Input arguments are `eigvals`, - `eigvecs`, `nroots`, `**env`. Default value is - `pyscf.lib.pick_real_eigs`. - guess : numpy.ndarray, optional - Guess vector. If not `None`, the diagonal is used to construct - a guess based on `diag`. Default value is `None`. - max_cycle : int, optional - Maximum number of iterations. Default value is `50`. - max_space : int, optional - Maximum number of trial vectors to store. Default value is - `12`. - conv_tol : float, optional - Threshold for convergence. Default value is `1e-12`. - hermitian : bool, optional - If `True`, the input matrix is assumed to be hermitian, - otherwise it is assumed to be non-hermitian. Default value - is `False`. - - Returns - ------- - eigvals : numpy.ndarray (nroots,) - Eigenvalues of the matrix, representing the energies of the - Green's function. - eigvecs : numpy.ndarray (n, nroots) - Eigenvectors of the matrix, which provide the Dyson orbitals - once projected into the physical space. - """ - - def __init__(self, matvec, diagonal, **kwargs): - # Input: - self.matvec = matvec - self.diagonal = diagonal - - if lib is None: - raise ImportError("PySCF installation required for %s." % self.__class__) - - # Parameters - self.nroots = kwargs.pop("nroots", 5) - self.picker = kwargs.pop("picker", pick_real_eigs) - self.guess = kwargs.pop("guess", None) - self.max_cycle = kwargs.pop("max_cycle", 50) - self.max_space = kwargs.pop("max_space", 12) - self.conv_tol = kwargs.pop("conv_tol", 1e-8) - self.tol_residual = kwargs.pop("tol_residual", 1e-6) - self.hermitian = kwargs.pop("hermitian", True) - self.nphys = kwargs.pop("nphys", None) - - # Base class: - super().__init__(matvec, diagonal, **kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > nroots: %s", self.nroots) - self.log.info(" > max_cycle: %s", self.max_cycle) - self.log.info(" > max_space: %s", self.max_space) - self.log.info(" > conv_tol: %s", self.conv_tol) - self.log.info(" > hermitian: %s", self.hermitian) - - # Caching: - self.converged = [] - self.eigvals = None - self.eigvecs = None - - def _kernel(self): - # if self.hermitian: - # convs, eigvals, eigvecs = self._kernel_hermitian() - # else: - # convs, eigvals, eigvecs = self._kernel_nonhermitian() - - # Sometimes Hermitian theories may have non-Hermitian matrices, - # i.e. perturbation theories, so always use the non-Hermitian - # solver. - convs, eigvals, eigvecs = self._kernel_nonhermitian() - - self.log.info(util.print_eigenvalues(eigvals, nroots=self.nroots)) - - return eigvals, eigvecs - - def _kernel_hermitian(self): - matvecs = lambda vs: [self.matvec(v) for v in vs] - - guess = self.guess - if guess is None: - args = np.argsort(np.abs(self.diagonal)) - guess = np.zeros((self.nroots, self.diagonal.size)) - for root, idx in enumerate(args[: self.nroots]): - guess[root, idx] = 1.0 - - convs, eigvals, eigvecs = lib.davidson1( - lambda vs: [self.matvec(v) for v in vs], - guess, - self.diagonal, - pick=self.picker, - tol=self.conv_tol, - max_cycle=self.max_cycle, - max_space=self.max_space, - tol_residual=self.tol_residual, - nroots=self.nroots, - verbose=0, - ) - eigvals = np.array(eigvals) - eigvecs = np.array(eigvecs).T - - mask = np.argsort(eigvals) - eigvals = eigvals[mask] - eigvecs = eigvecs[:, mask] - - self.eigvals = eigvals - self.eigvecs = eigvecs - self.converged = convs - - return convs, eigvals, eigvecs - - def _kernel_nonhermitian(self): - matvecs = lambda vs: [self.matvec(v) for v in vs] - - guess = self.guess - if guess is None: - args = np.argsort(np.abs(self.diagonal)) - guess = np.zeros((self.nroots, self.diagonal.size)) - for root, idx in enumerate(args[: self.nroots]): - guess[root, idx] = 1.0 - - convs, eigvals, eigvecs = lib.davidson_nosym1( - lambda vs: [self.matvec(v) for v in vs], - guess, - self.diagonal, - pick=self.picker, - tol=self.conv_tol, - max_cycle=self.max_cycle, - max_space=self.max_space, - nroots=self.nroots, - verbose=0, - ) - eigvals = np.array(eigvals) - eigvecs = np.array(eigvecs).T - - mask = np.argsort(eigvals) - eigvals = eigvals[mask] - eigvecs = eigvecs[:, mask] - - self.eigvals = eigvals - self.eigvecs = eigvecs - self.converged = convs - - return convs, eigvals, eigvecs - - def get_dyson_orbitals(self): - if self.nphys is None: - raise ValueError("`nphys` must be set to use `Exact.get_dyson_orbitals`") - - return super().get_dyson_orbitals() - - def get_auxiliaries(self): - raise ValueError("Cannot determine auxiliaries using `Davidson`.") diff --git a/dyson/solvers/density.py b/dyson/solvers/density.py deleted file mode 100644 index d4f2ca5..0000000 --- a/dyson/solvers/density.py +++ /dev/null @@ -1,226 +0,0 @@ -""" -Relax the density matrix in the presence of a self-energy. -""" - -import numpy as np -from pyscf import lib - -from dyson import NullLogger -from dyson.lehmann import Lehmann -from dyson.solvers import AufbauPrinciple, AuxiliaryShift, BaseSolver - - -class DensityRelaxation(BaseSolver): - """ - Relax the density matrix in the presence of a self-energy. - - Parameters - ---------- - get_fock : callable - Callable that returns the Fock matrix in the MO basis. Takes - a density matrix as input. - se : dyson.lehmann.Lehmann - Lehmann representation of the self-energy. - nelec : int - Number of electrons. - occupancy : int, optional - Occupancy of each state, i.e. `2` for a restricted reference - and `1` for other references. Default value is `2`. - chempot_solver : BaseSolver, optional - Solver for the chemical potential. One of - {`dyson.solvers.AufbauPrinciple`, - `dyson.solvers.AuxiliaryShift`}. Default value is - `dyson.solvers.AuxiliaryShift`. - diis_space : int, optional - Size of the DIIS space. Default value is `8`. - diis_min_space : int, optional - Minimum size of the DIIS space. Default value is `2`. - max_cycle_outer : int, optional - Maximum number of outer iterations. Default value is `20`. - max_cycle_inner : int, optional - Maximum number of inner iterations. Default value is `50`. - conv_tol : float, optional - Threshold for convergence in the change in the density matrix. - Default value is `1e-8`. - """ - - def __init__(self, get_fock, se, nelec, **kwargs): - # Input: - self._get_fock = get_fock - self.se = Lehmann.from_pyscf(se) - self.nelec = nelec - - # Parameters: - self.occupancy = kwargs.pop("occupancy", 2) - self.chempot_solver = kwargs.pop("chempot_solver", AuxiliaryShift) - self.diis_space = kwargs.pop("diis_space", 8) - self.diis_min_space = kwargs.pop("diis_min_space", 2) - self.max_cycle_outer = kwargs.pop("max_cycle_outer", 20) - self.max_cycle_inner = kwargs.pop("max_cycle_inner", 50) - self.conv_tol = kwargs.pop("conv_tol", 1e-8) - - # Base class: - super().__init__(**kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > occupancy: %s", self.occupancy) - self.log.info(" > chempot_solver: %s", self.chempot_solver) - self.log.info(" > diis_space: %s", self.diis_space) - self.log.info(" > diis_min_space: %s", self.diis_min_space) - self.log.info(" > max_cycle_outer: %s", self.max_cycle_outer) - self.log.info(" > max_cycle_inner: %s", self.max_cycle_inner) - self.log.info(" > conv_tol: %s", self.conv_tol) - - # Caching: - self.converged = False - self.se_res = None - self.gf_res = None - - def get_fock(self, rdm1): - """ - Get the Fock matrix in the MO basis. - - Parameters - ---------- - rdm1 : numpy.ndarray - One-particle reduced density matrix in the MO basis. - - Returns - ------- - fock : numpy.ndarray - Fock matrix. - """ - - return self._get_fock(rdm1) - - def optimise_chempot(self, se, fock): - """ - Optimise the chemical potential. - - Parameters - ---------- - se : dyson.lehmann.Lehmann - Lehmann representation of the self-energy. - fock : numpy.ndarray - Fock matrix. - - Returns - ------- - se : dyson.lehmann.Lehmann - Lehmann representation of the self-energy, with the - chemical potential optimised. - error : float - Error in the chemical potential. - """ - - if self.chempot_solver: - solver = self.chempot_solver(fock, se, self.nelec, guess=se.chempot, log=NullLogger()) - solver.kernel() - - se = solver.get_self_energy() - error = solver.error - converged = solver.converged - - else: - error = 0.0 - converged = True - - return se, error, converged - - def _kernel_rhf(self): - """ - Perform the self-consistent field for a restricted reference. - """ - - se = self.se - nocc = self.nelec // self.occupancy - rdm1 = np.zeros((se.nphys, se.nphys)) - rdm1[:nocc, :nocc] = np.eye(nocc) * self.occupancy - rdm1_prev = rdm1.copy() - fock = self.get_fock(rdm1) - - self.log.info("-" * 47) - self.log.info( - "{:^6s} {:^6s} {:^16s} {:^16s}".format( - "Iter", - "DM iter", - "DM error", - "Chempot error", - ) - ) - self.log.info("%6s %6s %16s %16s" % ("-" * 6, "-" * 6, "-" * 16, "-" * 16)) - - for niter_outer in range(1, self.max_cycle_outer + 1): - se, error_chempot, converged_chempot = self.optimise_chempot(se, fock) - - diis = lib.diis.DIIS() - diis.space = self.diis_space - diis.min_space = self.diis_min_space - diis.verbose = 0 - - for niter_inner in range(1, self.max_cycle_inner + 1): - w, v = se.diagonalise_matrix_with_projection(fock) - gf = Lehmann(w, v, chempot=se.chempot) - - aufbau = AufbauPrinciple(gf, self.nelec, log=NullLogger()) - aufbau.kernel() - se.chempot = gf.chempot = aufbau.chempot - - rdm1 = gf.occupied().moment(0) * self.occupancy - fock = self.get_fock(rdm1) - - try: - fock = diis.update(fock, xerr=None) - except np.linalg.LinAlgError: - pass - - error_rdm1 = np.max(np.abs(rdm1 - rdm1_prev)) - self.log.debug("%6d %6d %16.5g", niter_outer, niter_inner, error_rdm1) - if error_rdm1 < self.conv_tol: - break - - rdm1_prev = rdm1.copy() - - self.log.info( - "%6d %6d %16.5g %16.5g", - niter_outer, - niter_inner, - error_rdm1, - error_chempot, - ) - - if error_rdm1 < self.conv_tol and abs(aufbau.error) < self.chempot_solver.conv_tol: - self.converged = True - break - - self.log.info("-" * 47) - - self.flag_convergence(self.converged) - - self.se_res = se - self.gf_res = gf - - return gf, se, self.converged - - def _kernel(self): - """ - Perform the self-consistent field. - """ - - if not isinstance(self.nelec, tuple): - return self._kernel_rhf() - else: - raise NotImplementedError("UHF not implemented.") - - def get_auxiliaries(self): - return self.se_res.energies, self.se_res.couplings - - def get_dyson_orbitals(self): - return self.gf_res.energies, self.gf_res.couplings - - def get_self_energy(self): - return self.se_res - - def get_greens_function(self): - return self.gf_res diff --git a/dyson/solvers/downfolded.py b/dyson/solvers/downfolded.py deleted file mode 100644 index fc1f4e7..0000000 --- a/dyson/solvers/downfolded.py +++ /dev/null @@ -1,195 +0,0 @@ -""" -Downfolded frequency space eigensolver. -""" - -import numpy as np - -from dyson import util -from dyson.solvers import BaseSolver - - -class Downfolded(BaseSolver): - """ - Downfolded frequency space eigensolver, satisfies self-consistency - in C* Σ(ω) C = ω. - - Input - ----- - static : numpy.ndarray - Static part of the matrix (i.e. self-energy). - function : callable - Function returning the matrix (i.e. self-energy) at a given - argument (i.e. frequency). Input arguments are `argument`. - - Parameters - ---------- - guess : float or numpy.ndarray, optional - Initial guess for the argument entering `function`. A single - float uses the same guess for every index in `orbitals`, - whilst a list allows different initial guesses per orbital. - Default value is `0.0`. - target : int or str, optional - Method used to target a particular root. If input is of type - `int`, take the eigenvalue at this index. Otherwise one of - `{"min", "max", "mindif"}`. The first two take the minimnum - and maximum eigenvalues, and `"mindif"` takes the eigenvalue - closest to the guess (and then closest to the previous one at - each subsequent iteration). - max_cycle : int, optional - Maximum number of iterations. Default value is `50`. - conv_tol : float, optional - Threshold for convergence. Default value is `1e-8`. - hermitian : bool, optional - If `True`, the input matrix is assumed to be hermitian, - otherwise it is assumed to be non-hermitian. Default value - is `True`. - - Returns - ------- - eigvals : numpy.ndarray (n,) - Eigenvalues of the matrix, representing the energies of the - Green's function. - eigvecs : numpy.ndarray (n, n) - Eigenvectors of the matrix, which are proportional to the - Dyson orbitals. - """ - - # TODO: Can probably use a newton solver as C* Σ(w) C - w = 0 - - def __init__(self, static, function, **kwargs): - # Input: - self.static = static - self.function = function - - # Parameters: - self.guess = kwargs.pop("guess", 0.0) - self.target = kwargs.pop("target", "mindif") - self.max_cycle = kwargs.pop("max_cycle", 50) - self.conv_tol = kwargs.pop("conv_tol", 1e-8) - self.hermitian = kwargs.pop("hermitian", True) - - # Base class: - super().__init__(function, **kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > guess: %s", self.guess) - self.log.info(" > target: %s", self.target) - self.log.info(" > max_cycle: %s", self.max_cycle) - self.log.info(" > conv_tol: %s", self.conv_tol) - self.log.info(" > hermitian: %s", self.hermitian) - - # Caching: - self.eigvals = None - self.eigvecs = None - - def picker(self, roots): - if isinstance(self.target, int): - root = roots[self.target] - else: - if self.target == "min": - root = np.min(roots) - elif self.target == "max": - root = np.max(roots) - elif self.target == "mindif": - root = roots[np.argmin(np.abs(roots - self.guess))] - else: - raise ValueError("`target = %s`" % self.target) - - return root - - def eig(self, matrix): - if self.hermitian: - return np.linalg.eigh(matrix) - else: - return np.linalg.eig(matrix) - - def _kernel(self): - root = self.guess - root_prev = None - - self.log.info("-" * 38) - self.log.info("%4s %16s %16s", "Iter", "Root", "Delta") - self.log.info("%4s %16s %16s", "-" * 4, "-" * 16, "-" * 16) - - for cycle in range(1, self.max_cycle + 1): - matrix = self.static + self.function(root) - roots = self.eig(matrix)[0] - root = self.picker(roots) - - if cycle > 1: - self.log.info("%4d %16.8f %16.3g", cycle, root, abs(root - root_prev)) - if abs(root - root_prev) < self.conv_tol: - break - else: - self.log.info("%4d %16.8f", cycle, root) - - root_prev = root - - self.log.info("%4s %16s %16s", "-" * 4, "-" * 16, "-" * 16) - - converged = abs(root - root_prev) < self.conv_tol - self.flag_convergence(converged) - - matrix = self.static + self.function(root) - eigvals, eigvecs = self.eig(matrix) - - self.eigvals = eigvals - self.eigvecs = eigvecs - - self.log.info(util.print_eigenvalues(eigvals)) - - return eigvals, eigvecs - - -class DiagonalDownfolded(BaseSolver): - """ - Downfolded frequency space eigensolver, satisfies self-consistency - in C* Σ(ω) C = ω where Σ is diagonal. - - Input - ----- - function : callable - Function returning elements of the matrix (i.e. self-energy) - at a given argument (i.e. frequency). Input arguments are - `argument`, `orbital1`, `orbital2`. - derivative : callable, optional - Function returning elements of the derivative of the matrix - (i.e. self-energy) at a given argument (i.e. frequency), with - the derivative being with respect to the variable of said - argument. Input arguments are the same a `function`. - orbitals : list, optional - Orbital indices to solve for eigenvalues and eigenvectors at. - Default value solves for every orbital the result of function - spans. - - Parameters - ---------- - method : str, optional - Method used to minimise the solution. One of `{"newton"}. - Default value is `"newton"`. - guess : float or numpy.ndarray, optional - Initial guess for the argument entering `function`. A single - float uses the same guess for every index in `orbitals`, - whilst a list allows different initial guesses per orbital. - Default value is `0.0`. - linearised : bool, optional - Linearise the problem using the derivative. If `True` then - `derivative` must be provided, and `diagonal=True`. Default - value is `False`. - diagonal : bool, optional - Apply a diagonal approximation to the input. Default value is - `False`. - - Returns - ------- - eigvals : numpy.ndarray (n,) - Eigenvalues of the matrix, representing the energies of the - Green's function. - eigvecs : numpy.ndarray (n, n) - Eigenvectors of the matrix, which are proportional to the - Dyson orbitals. - """ - - def __init__(self, *args, **kwargs): - raise NotImplementedError # TODO diff --git a/dyson/solvers/dynamic/__init__.py b/dyson/solvers/dynamic/__init__.py new file mode 100644 index 0000000..68ad18f --- /dev/null +++ b/dyson/solvers/dynamic/__init__.py @@ -0,0 +1,12 @@ +r"""Solvers for solving the Dyson equation dynamically. + +Submodules +---------- + +.. autosummary:: + :toctree: + + corrvec + cpgf + +""" diff --git a/dyson/solvers/dynamic/corrvec.py b/dyson/solvers/dynamic/corrvec.py new file mode 100644 index 0000000..19bd6cb --- /dev/null +++ b/dyson/solvers/dynamic/corrvec.py @@ -0,0 +1,327 @@ +"""Correction vector Green's function solver [1]_. + +.. [1] Nocera, A., & Alvarez, G. (2016). Spectral functions with the density matrix + renormalization group: Krylov-space approach for correction vectors. Physical Review. E, 94(5). + https://doi.org/10.1103/physreve.94.053308 +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING, cast + +from scipy.sparse.linalg import LinearOperator, lgmres + +from dyson import console, printing +from dyson import numpy as np +from dyson.grids.frequency import RealFrequencyGrid +from dyson.representations.dynamic import Dynamic +from dyson.representations.enums import Component, Ordering, Reduction +from dyson.solvers.solver import DynamicSolver +from dyson.solvers.static.exact import orthogonalise_self_energy + +if TYPE_CHECKING: + from typing import Any, Callable + + from dyson.expressions.expression import BaseExpression + from dyson.representations.lehmann import Lehmann + from dyson.typing import Array + +# TODO: Can we use DIIS? + + +class CorrectionVector(DynamicSolver): + """Correction vector Green's function solver. + + Args: + matvec: The matrix-vector operation for the self-energy supermatrix. + diagonal: The diagonal of the self-energy supermatrix. + nphys: The number of physical degrees of freedom. + grid: Real frequency grid upon which to evaluate the Green's function. + """ + + reduction: Reduction = Reduction.NONE + component: Component = Component.FULL + ordering: Ordering = Ordering.ORDERED + conv_tol: float = 1e-8 + _options: set[str] = {"reduction", "component", "ordering", "conv_tol"} + + def __init__( # noqa: D417 + self, + matvec: Callable[[Array], Array], + diagonal: Array, + nphys: int, + grid: RealFrequencyGrid, + get_state_bra: Callable[[int], Array] | None = None, + get_state_ket: Callable[[int], Array] | None = None, + **kwargs: Any, + ): + r"""Initialise the solver. + + Args: + matvec: The matrix-vector operation for the self-energy supermatrix. + diagonal: The diagonal of the self-energy supermatrix. + nphys: The number of physical degrees of freedom. + grid: Real frequency grid upon which to evaluate the Green's function. + get_state_bra: Function to get the bra vector corresponding to a fermion operator acting + on the ground state. If ``None``, the state vector is :math:`v_{i} = \delta_{ij}` + for orbital :math:`j`. + get_state_ket: Function to get the ket vector corresponding to a fermion operator acting + on the ground state. If ``None``, the :arg:`get_state_bra` function is used. + component: The component of the dynamic representation to solve for. + reduction: The reduction of the dynamic representation to solve for. + conv_tol: Convergence tolerance for the solver. + ordering: Time ordering of the resolvent. + """ + self._matvec = matvec + self._diagonal = diagonal + self._nphys = nphys + self._grid = grid + self._get_state_bra = get_state_bra + self._get_state_ket = get_state_ket + self.set_options(**kwargs) + + def __post_init__(self) -> None: + """Hook called after :meth:`__init__`.""" + # Check the input + if self.diagonal.ndim != 1: + raise ValueError("diagonal must be a 1D array.") + if not callable(self.matvec): + raise ValueError("matvec must be a callable function.") + + # Print the input information + console.print(f"Matrix shape: [input]{(self.diagonal.size, self.diagonal.size)}[/input]") + console.print(f"Number of physical states: [input]{self.nphys}[/input]") + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> CorrectionVector: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + if "grid" not in kwargs: + raise ValueError("Missing required argument grid.") + static, self_energy, bra, ket = orthogonalise_self_energy( + static, self_energy, overlap=overlap + ) + return cls( + lambda vector: self_energy.matvec(static, vector), + self_energy.diagonal(static), + self_energy.nphys, + kwargs.pop("grid"), + bra.__getitem__, + ket.__getitem__ if ket is not None else None, + **kwargs, + ) + + @classmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> CorrectionVector: + """Create a solver from an expression. + + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + if "grid" not in kwargs: + raise ValueError("Missing required argument grid.") + diagonal = expression.diagonal() + matvec = expression.apply_hamiltonian + return cls( + matvec, + diagonal, + expression.nphys, + kwargs.pop("grid"), + expression.get_excitation_bra, + expression.get_excitation_ket, + **kwargs, + ) + + def matvec_dynamic(self, vector: Array, grid_index: int) -> Array: + r"""Perform the matrix-vector operation for the dynamic self-energy supermatrix. + + .. math:: + \mathbf{x}_\omega = \left(\omega - \mathbf{H} - i\eta\right) \mathbf{r} + + Args: + vector: The vector to operate on. + grid_index: Index of the real frequency grid. + + Returns: + The result of the matrix-vector operation. + """ + grid = cast(RealFrequencyGrid, self.grid[[grid_index]]) + resolvent = grid.resolvent( + np.array(0.0), -self.diagonal, ordering=self.ordering, invert=False + ) + result: Array = vector[None] * resolvent + result -= self.matvec(vector.real)[None] + if np.any(np.abs(vector.imag) > 1e-14): + result -= self.matvec(vector.imag)[None] * 1.0j + return result + + def matdiv_dynamic(self, vector: Array, grid_index: int) -> Array: + r"""Approximately perform a matrix-vector division for the dynamic self-energy supermatrix. + + .. math:: + \mathbf{x}_\omega = \frac{\mathbf{r}}{\omega - \mathbf{H} - i\eta} + + Args: + vector: The vector to operate on. + grid_index: Index of the real frequency grid. + + Returns: + The result of the matrix-vector division. + + Notes: + The inversion is approximated using the diagonal of the matrix. + """ + grid = cast(RealFrequencyGrid, self.grid[[grid_index]]) + resolvent = grid.resolvent(self.diagonal, 0.0, ordering=self.ordering) + result = vector[None] * resolvent[:, None] + result[np.isinf(result)] = np.nan # or 0? + return result + + def get_state_bra(self, orbital: int) -> Array: + """Get the bra vector corresponding to a fermion operator acting on the ground state. + + Args: + orbital: Orbital index. + + Returns: + Bra vector. + """ + if self._get_state_bra is None: + return np.eye(self.nphys, 1, k=orbital).ravel() + return self._get_state_bra(orbital) + + def get_state_ket(self, orbital: int) -> Array: + """Get the ket vector corresponding to a fermion operator acting on the ground state. + + Args: + orbital: Orbital index. + + Returns: + Ket vector. + """ + if self._get_state_ket is None: + return self.get_state_bra(orbital) + return self._get_state_ket(orbital) + + def kernel(self) -> Dynamic[RealFrequencyGrid]: + """Run the solver. + + Returns: + The Green's function on the real frequency grid. + """ + # Get the printing helpers + progress = printing.IterationsPrinter(self.nphys * len(self.grid), description="Frequency") + progress.start() + + # Precompute bra vectors # TODO: Optional + bras = list(map(self.get_state_bra, range(self.nphys))) + + # Loop over ket vectors + shape = (len(self.grid),) + (self.nphys,) * self.reduction.ndim + greens_function = np.zeros(shape, dtype=complex) + failed: set[int] = set() + for i in range(self.nphys): + ket = self.get_state_ket(i) + + # Loop over frequencies + x: Array | None = None + outer_v: list[tuple[Array, Array]] = [] + for w in range(len(self.grid)): + progress.update(i * len(self.grid) + w + 1) + if w in failed: + continue + + linop_shape = (self.diagonal.size, self.diagonal.size) + matvec = LinearOperator( + linop_shape, lambda v: self.matvec_dynamic(v, w), dtype=complex + ) + matdiv = LinearOperator( + linop_shape, lambda v: self.matdiv_dynamic(v, w), dtype=complex + ) + + # Solve the linear system + x, info = lgmres( + matvec, + ket, + x0=x, + M=matdiv, + rtol=0.0, + atol=self.conv_tol, + outer_v=outer_v, + ) + + # Contract the Green's function + if info != 0: + greens_function[w] = np.nan + failed.add(w) + elif self.reduction == Reduction.NONE: + for j in range(self.nphys): + greens_function[w, j, i] = bras[j] @ x + elif self.reduction == Reduction.DIAG: + greens_function[w, i] = bras[i] @ x + elif self.reduction == Reduction.TRACE: + greens_function[w] += bras[i] @ x + else: + self.reduction.raise_invalid_representation() + + # Post-process the Green's function component + # TODO: Can we do this earlier to avoid computing unnecessary components? + if self.component == Component.REAL: + greens_function = greens_function.real + elif self.component == Component.IMAG: + greens_function = greens_function.imag + + progress.stop() + rating = printing.rate_error(len(failed) / len(self.grid), 1e-100, 1e-2) + console.print("") + console.print( + f"Converged [output]{len(self.grid) - len(failed)} of {len(self.grid)}[/output] " + f"frequencies ([{rating}]{1 - len(failed) / len(self.grid):.2%}[/{rating}])." + ) + + return Dynamic( + self.grid, + greens_function, + reduction=self.reduction, + component=self.component, + hermitian=self.get_state_ket is None, + ) + + @property + def matvec(self) -> Callable[[Array], Array]: + """Get the matrix-vector operation for the self-energy supermatrix.""" + return self._matvec + + @property + def diagonal(self) -> Array: + """Get the diagonal of the self-energy supermatrix.""" + return self._diagonal + + @property + def grid(self) -> RealFrequencyGrid: + """Get the real frequency grid.""" + return self._grid + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return self._nphys diff --git a/dyson/solvers/dynamic/cpgf.py b/dyson/solvers/dynamic/cpgf.py new file mode 100644 index 0000000..321bf29 --- /dev/null +++ b/dyson/solvers/dynamic/cpgf.py @@ -0,0 +1,246 @@ +"""Chebyshev polynomial Green's function solver [1]_. + +.. [1] Ferreira, A., & Mucciolo, E. R. (2015). Critical delocalization of Chiral zero + energy modes in graphene. Physical Review Letters, 115(10). + https://doi.org/10.1103/physrevlett.115.106601 +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from dyson import console, printing, util +from dyson import numpy as np +from dyson.representations.dynamic import Dynamic +from dyson.representations.enums import Component, Ordering, Reduction +from dyson.solvers.solver import DynamicSolver + +if TYPE_CHECKING: + from typing import Any + + from dyson.expressions.expression import BaseExpression + from dyson.grids.frequency import RealFrequencyGrid + from dyson.representations.lehmann import Lehmann + from dyson.typing import Array + + +def _infer_max_cycle(moments: Array) -> int: + """Infer the maximum number of cycles from the moments.""" + return moments.shape[0] - 1 + + +class CPGF(DynamicSolver): + """Chebyshev polynomial Green's function solver. + + Args: + moments: Chebyshev moments of the Green's function. + grid: Real frequency grid upon which to evaluate the Green's function. + scaling: Scaling factors to ensure the energy scale of the Lehmann representation is in + ``[-1, 1]``. The scaling is applied as ``(energies - scaling[1]) / scaling[0]``. + """ + + reduction: Reduction = Reduction.NONE + component: Component = Component.FULL + ordering: Ordering = Ordering.ORDERED + _options: set[str] = {"reduction", "component", "ordering"} + + def __init__( # noqa: D417 + self, + moments: Array, + grid: RealFrequencyGrid, + scaling: tuple[float, float], + max_cycle: int | None = None, + **kwargs: Any, + ): + """Initialise the solver. + + Args: + moments: Chebyshev moments of the Green's function. + grid: Real frequency grid upon which to evaluate the Green's function. + scaling: Scaling factors to ensure the energy scale of the Lehmann representation is in + `[-1, 1]`. The scaling is applied as `(energies - scaling[1]) / scaling[0]`. + max_cycle: Maximum number of iterations. + component: The component of the dynamic representation to solve for. + reduction: The reduction of the dynamic representation to solve for. + ordering: Time ordering of the resolvent. + """ + self._moments = moments + self._grid = grid + self._scaling = scaling + self.max_cycle = max_cycle if max_cycle is not None else _infer_max_cycle(moments) + self.set_options(**kwargs) + + def __post_init__(self) -> None: + """Hook called after :meth:`__init__`.""" + # Check the input + if self.moments.ndim != 3 or self.moments.shape[1] != self.moments.shape[2]: + raise ValueError( + "moments must be a 3D array with the second and third dimensions equal." + ) + if _infer_max_cycle(self.moments) < self.max_cycle: + raise ValueError("not enough moments provided for the specified max_cycle.") + if self.ordering == Ordering.ORDERED: + raise NotImplementedError(f"{self.ordering} is not implemented for CPGF.") + + # Print the input information + cond = printing.format_float( + np.linalg.cond(self.moments[0]), threshold=1e10, scientific=True, precision=4 + ) + scaling = ( + printing.format_float(self.scaling[0]), + printing.format_float(self.scaling[1]), + ) + console.print(f"Number of physical states: [input]{self.nphys}[/input]") + console.print(f"Number of moments: [input]{self.moments.shape[0]}[/input]") + console.print(f"Scaling parameters: [input]({scaling[0]}, {scaling[1]})[/input]") + console.print(f"Overlap condition number: {cond}") + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> CPGF: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + if "grid" not in kwargs: + raise ValueError("Missing required argument grid.") + max_cycle = kwargs.pop("max_cycle", 16) + energies, couplings = self_energy.diagonalise_matrix_with_projection( + static, overlap=overlap + ) + if "scaling" not in kwargs: + scaling = util.get_chebyshev_scaling_parameters(energies.min(), energies.max()) + else: + scaling = kwargs.pop("scaling") + greens_function = self_energy.__class__(energies, couplings, chempot=self_energy.chempot) + moments = greens_function.chebyshev_moments(range(max_cycle + 1), scaling=scaling) + return cls(moments, kwargs.pop("grid"), scaling, max_cycle=max_cycle, **kwargs) + + @classmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> CPGF: + """Create a solver from an expression. + + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + if "grid" not in kwargs: + raise ValueError("Missing required argument grid.") + max_cycle = kwargs.pop("max_cycle", 16) + if "scaling" not in kwargs: + diag = expression.diagonal() + emin = np.min(diag) + emax = np.max(diag) + emin -= 0.5 * (emax - emin) + emax += 0.5 * (emax - emin) + scaling = util.get_chebyshev_scaling_parameters(emin, emax) + else: + scaling = kwargs.pop("scaling") + moments = expression.build_gf_chebyshev_moments(max_cycle + 1, scaling=scaling) + return cls(moments, kwargs.pop("grid"), scaling, max_cycle=max_cycle, **kwargs) + + def kernel(self, iteration: int | None = None) -> Dynamic[RealFrequencyGrid]: + """Run the solver. + + Args: + iteration: The iteration number. + + Returns: + The Green's function on the real frequency grid. + """ + if iteration is None: + iteration = self.max_cycle + + # Get the printing helpers + progress = printing.IterationsPrinter(iteration + 1, description="Polynomial order") + progress.start() + + # Get the moments -- allow input to already be traced or diagonal + if self.reduction == Reduction.NONE: + moments = self.moments[: iteration + 1].astype(complex) + elif self.reduction == Reduction.DIAG: + moments = util.as_diagonal(self.moments[: iteration + 1], 1).astype(complex) + elif self.reduction == Reduction.TRACE: + moments = util.as_trace(self.moments[: iteration + 1], 1).astype(complex) + else: + self.reduction.raise_invalid_representation() + if (moments.ndim - 1) != self.reduction.ndim: + raise ValueError( + f"moments must be {self.reduction.ndim + 1}D for reduction {self.reduction}, got " + f"{moments.ndim}D." + ) + + # Scale the grid + scaled_grid = (self.grid.points - self.scaling[1]) / self.scaling[0] + scaled_eta = self.grid.eta / self.scaling[0] + shifted_grid = scaled_grid + 1j * scaled_eta + + # Initialise factors + numerator = shifted_grid - 1j * np.sqrt(1 - shifted_grid**2) + denominator = np.sqrt(1 - shifted_grid**2) + kernel = 1.0 / denominator + + # Iteratively compute the Green's function + greens_function = util.einsum("z,...->z...", kernel, moments[0]) + for cycle in range(1, iteration + 1): + progress.update(cycle) + kernel *= numerator + greens_function += util.einsum("z,...->z...", kernel, moments[cycle]) * 2 + + # Apply factors + greens_function /= self.scaling[0] + greens_function *= -1.0j + if self.ordering == Ordering.ADVANCED: + greens_function = greens_function.conj() + + # Post-process the Green's function component + # TODO: Can we do this earlier to avoid computing unnecessary components? + if self.component == Component.REAL: + greens_function = greens_function.real + elif self.component == Component.IMAG: + greens_function = greens_function.imag + + progress.stop() + + return Dynamic( + self.grid, + greens_function, + reduction=self.reduction, + component=self.component, + hermitian=np.allclose(self.moments, self.moments.transpose(0, 2, 1).conj()), + ) + + @property + def moments(self) -> Array: + """Get the moments of the self-energy.""" + return self._moments + + @property + def grid(self) -> RealFrequencyGrid: + """Get the real frequency grid.""" + return self._grid + + @property + def scaling(self) -> tuple[float, float]: + """Get the scaling factors.""" + return self._scaling + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return self.moments.shape[-1] diff --git a/dyson/solvers/exact.py b/dyson/solvers/exact.py deleted file mode 100644 index 1d148bc..0000000 --- a/dyson/solvers/exact.py +++ /dev/null @@ -1,111 +0,0 @@ -""" -Exact eigensolver on the dense upfolded self-energy. -""" - -import numpy as np -import scipy.linalg - -from dyson import util -from dyson.solvers import BaseSolver - - -class Exact(BaseSolver): - """ - Exact eigensolver on the dense upfolded self-energy. - - Input - ----- - matrix : numpy.ndarray (n, n) - Dense representation of the upfolded self-energy matrix. - - Parameters - ---------- - hermitian : bool, optional - If `True`, the input matrix is assumed to be hermitian, - otherwise it is assumed to be non-hermitian. Default value - is `True`. - overlap : numpy.ndarray, optional - If provided, use as part of a generalised eigenvalue problem. - Default value is `None`. - - Returns - ------- - eigvals : numpy.ndarray (n,) - Eigenvalues of the matrix, representing the energies of the - Green's function. - eigvecs : numpy.ndarray (n, n) - Eigenvectors of the matrix, which provide the Dyson orbitals - once projected into the physical space. - """ - - def __init__(self, matrix, **kwargs): - # Input: - self.matrix = matrix - - # Parameters: - self.hermitian = kwargs.pop("hermitian", True) - self.overlap = kwargs.pop("overlap", None) - self.nphys = kwargs.pop("nphys", None) - - # Base class: - super().__init__(matrix, **kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > hermitian: %s", self.hermitian) - self.log.info(" > overlap: %s", None if not self.generalised else type(self.overlap)) - self.log.info(" > nphys: %s", self.nphys) - - # Caching: - self.eigvals = None - self.eigvecs = None - - def _kernel(self): - if self.hermitian: - eigvals, eigvecs = self._kernel_hermitian() - else: - eigvals, eigvecs = self._kernel_nonhermitian() - - self.eigvals = eigvals - self.eigvecs = eigvecs - - self.log.info(util.print_eigenvalues(eigvals)) - - return eigvals, eigvecs - - def _kernel_hermitian(self): - if self.generalised: - return np.linalg.eigh(self.matrix) - else: - return scipy.linalg.eigh(self.matrix, b=self.overlap) - - def _kernel_nonhermitian(self): - if self.generalised: - return np.linalg.eig(self.matrix) - else: - return scipy.linalg.eig(self.matrix, b=self.overlap) - - def get_dyson_orbitals(self): - if self.nphys is None: - raise ValueError("`nphys` must be set to use `Exact.get_dyson_orbitals`") - - return super().get_dyson_orbitals() - - def get_auxiliaries(self): - if self.nphys is None: - raise ValueError("`nphys` must be set to use `Exact.get_dyson_orbitals`") - - energies = self.matrix[: self.nphys, self.nphys :] - - if self.hermitian: - couplings = self.matrix[: self.nphys, self.nphys :] - else: - couplings_l = self.matrix[: self.nphys, self.nphys :] - couplings_r = self.matrix[self.nphys :, : self.nphys].conj().T - couplings = (couplings_l, couplings_r) - - return energies, couplings - - @property - def generalised(self): - return self.overlap is not None diff --git a/dyson/solvers/kpmgf.py b/dyson/solvers/kpmgf.py deleted file mode 100644 index 4b57236..0000000 --- a/dyson/solvers/kpmgf.py +++ /dev/null @@ -1,216 +0,0 @@ -""" -Kernel polynomial method (moment-conserving Chebyshev eigensolver), -conserving Chebyshev moments of the Green's function. -""" - -import numpy as np -import scipy.integrate - -from dyson import util -from dyson.solvers import BaseSolver - - -def as_trace(arr): - """Return the trace of `arr`, if it has more than one dimension.""" - - if arr.ndim > 1: - arr = np.trace(arr, axis1=-2, axis2=-1) - - return arr - - -class KPMGF(BaseSolver): - """ - Kernel polynomial method. - - Input - ----- - moments : numpy.ndarray - Chebyshev moments of the Green's function. - grid : numpy.ndarray - Real-valued frequency grid to plot the spectral function upon. - scale : tuple of int - Scaling parameters used to scale the spectrum to [-1, 1], - given as `(a, b)` where - - a = (ωmax - ωmin) / (2 - ε) - b = (ωmax + ωmin) / 2 - - where ωmax and ωmin are the maximum and minimum energies in - the spectrum, respectively, and ε is a small number shifting - the spectrum values away from the boundaries. - - Parameters - ---------- - max_cycle : int, optional - Maximum number of iterations. If `None`, perform as many as - the inputted number of moments permits. Default value is - `None`. - kernel_type : str, optional - Kernel to apply to regularise the Chebyshev representation. - Can be one of `{None, "lorentz", "lanczos", "jackson"}, or a - callable whose arguments are the solver object and the - iteration number. Default value is `None`. - trace : bool, optional - Whether to compute the trace of the Green's function. If - `False`, the entire Green's function is computed. Default - value is `True`. - lorentz_parameter : float or callable, optional - Lambda parameter for the Lorentz kernel, a float value which - is then scaled by the number of Chebyshev moments. Default - value is 0.1. - lanczos_order : int - Order parameter for the Lanczos kernel. Default value is 2. - - Returns - ------- - spectral_function : numpy.ndarray - Spectral function expressed on the input grid. - """ - - def __init__(self, moments, grid, scale, **kwargs): - # Input: - self.moments = moments - self.grid = grid - self.scale = scale - - # Parameters - self.max_cycle = kwargs.pop("max_cycle", None) - # self.hermitian = True - self.kernel_type = kwargs.pop("kernel_type", None) - self.trace = kwargs.pop("trace", True) - self.lorentz_parameter = kwargs.pop("lorentz_parameter", 0.1) - self.lanczos_order = kwargs.pop("lanczos_order", 2) - - max_cycle_limit = len(moments) - 1 - if self.max_cycle is None: - self.max_cycle = max_cycle_limit - if self.max_cycle > max_cycle_limit: - raise ValueError( - "`max_cycle` cannot be more than the number of inputted moments minus one." - ) - - # Base class: - super().__init__(moments, **kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > max_cycle: %s", self.max_cycle) - # self.log.info(" > hermitian: %s", self.hermitian) - self.log.info(" > grid: %s[%d]", type(self.grid), len(self.grid)) - self.log.info(" > scale: %s", scale) - self.log.info(" > kernel_type: %s", self.kernel_type) - self.log.info(" > trace: %s", self.trace) - self.log.info(" > lorentz_parameter: %s", self.lorentz_parameter) - self.log.info(" > lanczos_order: %s", self.lanczos_order) - - def get_expansion_coefficients(self, iteration): - """ - Compute the expansion coefficients to modify the moments, - thereby damping the Gibbs oscillations. - """ - - n = iteration - x = np.arange(1, iteration + 1) - - if self.kernel_type is None or self.kernel_type.lower() == "dirichlet": - coefficients = np.ones((n,)) - - elif callable(self.kernel_type): - coefficients = self.kernel_type(n) - - elif self.kernel_type.lower() == "lorentz": - if callable(self.lorentz_parameter): - λ = self.lorentz_parameter(n) - else: - λ = self.lorentz_parameter - coefficients = np.sinh(λ * (1 - x / n)) - coefficients /= np.sinh(λ) - - elif self.kernel_type.lower() == "fejer": - coefficients = 1 - x / (n + 1) - - elif self.kernel_type.lower() == "lanczos": - xp = np.pi * x / n - m = self.lanczos_order - coefficients = (np.sin(xp) / xp) ** m - - elif self.kernel_type.lower() == "jackson": - norm = 1 / (n + 1) - coefficients = (n - x + 1) * np.cos(np.pi * x * norm) - coefficients += np.sin(np.pi * x * norm) / np.tan(np.pi * norm) - coefficients *= norm - - else: - raise ValueError("Invalid self.kernel_type `%s`" % self.kernel_type) - - return coefficients - - def initialise_recurrence(self): - self.log.info("-" * 21) - self.log.info("{:^4s} {:^16s}".format("Iter", "Integral")) - self.log.info("{:^4s} {:^16s}".format("-" * 4, "-" * 16)) - - def _kernel(self, iteration=None): - self.initialise_recurrence() - - if iteration is None: - iteration = self.max_cycle - - # Get the moments - allow input to already be traced - if self.trace: - moments = as_trace(self.moments[: iteration + 1]) - else: - moments = self.moments[: iteration + 1] - - # Initialise scaled grids - a, b = self.scale - scaled_grid = (self.grid - b) / a - grids = (np.ones_like(scaled_grid), scaled_grid) - - # Initialise the polynomial - coefficients = self.get_expansion_coefficients(iteration + 1) - moments = np.einsum("n,n...->n...", coefficients, moments[: iteration + 1]) - polynomial = np.array([moments[0]] * self.grid.size) - - def _get_spectral_function(polynomial): - f = polynomial / np.pi - f /= np.sqrt(1 - scaled_grid**2) - # FIXME should this be here? - # f /= np.pi - f /= np.sqrt(a**2 - (self.grid - b**2)) - return f - - f = _get_spectral_function(as_trace(polynomial)) - integral = scipy.integrate.simps(f, self.grid) - self.log.info("%4d %16.8g", 0, integral) - - for niter in range(1, iteration + 1): - polynomial += np.multiply.outer(grids[-1], moments[niter]) * 2 - grids = (grids[-1], 2 * scaled_grid * grids[-1] - grids[-2]) - - if niter in (1, 2, 3, 4, 5, 10, iteration) or niter % 100 == 0: - f = _get_spectral_function(as_trace(polynomial)) - integral = scipy.integrate.simps(f, self.grid) - self.log.info("%4d %16.8g", niter, integral) - - f = _get_spectral_function(polynomial) - - self.log.info("-" * 21) - - return f - - def _get_spectral_function(self, polynomial): - """ - Get the spectral function corresponding to the current - iteration. - """ - - a, b = self.scale - grid = (self.grid - b) / a - - return f - - @property - def nphys(self): - return self.moments[0].shape[0] diff --git a/dyson/solvers/mblgf.py b/dyson/solvers/mblgf.py deleted file mode 100644 index 4f6a886..0000000 --- a/dyson/solvers/mblgf.py +++ /dev/null @@ -1,992 +0,0 @@ -""" -Moment-conserving block Lanczos eigensolver, conserving moments -of the Green's function. -""" - -import warnings - -import numpy as np -import scipy.linalg - -from dyson import util -from dyson.lehmann import Lehmann -from dyson.solvers import BaseSolver - -# TODO inherit things from MBLSE or vice versa? - - -class RecurrenceCoefficients: - """ - Recurrence coefficients container. - """ - - def __init__(self, shape, hermitian=True, force_orthogonality=True, dtype=np.float64): - self.hermitian = hermitian - self.zero = np.zeros(shape, dtype=dtype) - self.data = {} - - def __getitem__(self, key): - i, j = key - - if i == j == 1: - return np.eye(self.zero.shape[0]) - - if i < 1 or j < 1 or i < j: - # Zero order Lanczos vectors are zero - return self.zero - else: - # Return ∑ Σ^{j-1} Q_{1} C_{i,j} - return self.data[i, j] - - def __setitem__(self, key, val): - i, j = key - - self.data[i, j] = val - - -class MBLGF_Symm(BaseSolver): - """ - Moment-conserving block Lanczos eigensolver, conserving the - moments of the Green's function, for a Hermitian Green's function. - - Input - ----- - moments : numpy.ndarray - Moments of the Green's function. - - Parameters - ---------- - max_cycle : int, optional - Maximum number of iterations. If `None`, perform as many as - the inputted number of moments permits. Default value is - `None`. - - Returns - ------- - eigvals : numpy.ndarray - Eigenvalues of the matrix, representing the energies of the - Green's function. - eigvecs : numpy.ndarray - Eigenvectors of the matrix, which provide the Dyson orbitals - once projected into the physical space. - """ - - def __init__(self, moments, **kwargs): - # Input: - self.moments = moments - - # Parameters: - self.max_cycle = kwargs.pop("max_cycle", None) - self.hermitian = True - - max_cycle_limit = (len(moments) - 2) // 2 - if self.max_cycle is None: - self.max_cycle = max_cycle_limit - if self.max_cycle > max_cycle_limit: - raise ValueError( - "`max_cycle` cannot be more than (M-2)/2, where " - "M is the number of inputted moments." - ) - - # Base class: - super().__init__(moments, **kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > max_cycle: %s", self.max_cycle) - self.log.info(" > hermitian: %s", self.hermitian) - - # Caching: - self._cache = {} - self.coefficients = RecurrenceCoefficients( - self.moments[0].shape, - hermitian=self.hermitian, - dtype=np.result_type(*self.moments), - ) - self.on_diagonal = {} - self.off_diagonal = {} - self.orth = None - self.iteration = None - - @util.cache - def orthogonalised_moment(self, n): - """ - Compute an orthogonalised moment. - """ - - orth = self.orth - if orth is None: - orth = util.matrix_power(self.moments[0], -0.5, hermitian=self.hermitian) - - return np.linalg.multi_dot( - ( - orth, - self.moments[n], - orth, - ) - ) - - def initialise_iteration_table(self): - """ - Print the header for the table summarising the iterations. - """ - - self.log.info("-" * 89) - self.log.info( - "{:^4s} {:^16s} {:^33s} {:^33}".format( - "", - "", - "Norm of matrix", - "Norm of removed space", - ) - ) - self.log.info( - "{:^4s} {:^16s} {:^33s} {:^33}".format( - "Iter", - "Moment error", - "-" * 33, - "-" * 33, - ) - ) - self.log.info( - "%4s %16s %16s %16s %16s %16s", - "", - "", - "On-diagonal", - "Off-diagonal", - "Square root", - "Inv. square root", - ) - self.log.info( - "%4s %16s %16s %16s %16s %16s", - "-" * 4, - "-" * 16, - "-" * 16, - "-" * 16, - "-" * 16, - "-" * 16, - ) - - def _check_moment_error(self, iteration=None): - """ - Check the error in the moments at a given iteration. - """ - - if iteration is None: - iteration = self.iteration - - energies, dyson_orbitals = self.get_dyson_orbitals(iteration=iteration) - - left = dyson_orbitals.copy() - moments_recovered = [] - for n in range(2 * iteration + 2): - moments_recovered.append(np.dot(left, dyson_orbitals.T.conj())) - left = left * energies[None] - - error_moments = sum( - util.scaled_error(a, b) - for a, b in zip(moments_recovered, self.moments[: 2 * iteration + 2]) - ) - - return error_moments - - def initialise_recurrence(self): - """ - Initialise the recurrences - performs the 'zeroth' iteration. - - This iteration is essentially equivalent to solving a generalised - eigenvalue problem on the Fock matrix in the physical space. - """ - - # Initialise the table - self.initialise_iteration_table() - - self.iteration = 0 - - # Calculate the orthogonalisation matrix - self.orth, error_inv_sqrt = util.matrix_power( - self.moments[0], - -0.5, - hermitian=self.hermitian, - return_error=True, - ) - - # Add zero matrix to out-of-bounds off-diagonal to simplify logic - self.off_diagonal[-1] = self.coefficients.zero - - # Zeroth order on-diagonal block is the orthogonalised first - # moment (equal to the orthogonalised static part of the - # matrix corresponding to the solution moments) - self.on_diagonal[0] = self.orthogonalised_moment(1) - - # Check the error in the moments up to this iteration - error_moments = self._check_moment_error() - - # Logging - self.log.info( - "%4d %16.3g %16.3g %16s %16s %16.3g", - 0, - error_moments, - np.linalg.norm(self.on_diagonal[0]), - "", - "", - error_inv_sqrt, - ) - - def recurrence_iteration(self): - """ - Perform an iteration of the recurrence. - """ - - self.iteration += 1 - i = self.iteration - 1 - - if self.iteration > self.max_cycle: - raise ValueError( - "Cannot perform more iterations than permitted " - "by `max_cycle` or (M-2)/2 where M is the number " - "of inputted moments." - ) - - # Find the square of the next off-diagonal block - off_diagonal_squared = self.coefficients.zero.copy() - for j in range(i + 2): - for k in range(i + 1): - off_diagonal_squared += np.linalg.multi_dot( - ( - self.coefficients[i + 1, k + 1].T.conj(), - self.orthogonalised_moment(j + k + 1), - self.coefficients[i + 1, j], - ) - ) - - off_diagonal_squared -= np.dot( - self.on_diagonal[i], - self.on_diagonal[i], - ) - if i: - off_diagonal_squared -= np.dot( - self.off_diagonal[i - 1], - self.off_diagonal[i - 1], - ) - - # Get the next off-diagonal block - self.off_diagonal[i], error_sqrt = util.matrix_power( - off_diagonal_squared, - 0.5, - hermitian=self.hermitian, - return_error=True, - ) - - # Get the inverse of the off-diagonal block - off_diagonal_inv, error_inv_sqrt = util.matrix_power( - off_diagonal_squared, - -0.5, - hermitian=self.hermitian, - return_error=True, - ) - - for j in range(i + 2): - residual = ( - +self.coefficients[i + 1, j] - - np.dot(self.coefficients[i + 1, j + 1], self.on_diagonal[i]) - - np.dot(self.coefficients[i, j + 1], self.off_diagonal[i - 1]) - ) - self.coefficients[i + 2, j + 1] = np.dot(residual, off_diagonal_inv) - - self.on_diagonal[i + 1] = self.coefficients.zero.copy() - for j in range(i + 2): - for k in range(i + 2): - self.on_diagonal[i + 1] += np.linalg.multi_dot( - ( - self.coefficients[i + 2, k + 1].T.conj(), - self.orthogonalised_moment(j + k + 1), - self.coefficients[i + 2, j + 1], - ) - ) - - # Check the error in the moments up to this iteration - error_moments = self._check_moment_error() - - # Logging - self.log.info( - "%4d %16.3g %16.3g %16.3g %16.3g %16.3g", - self.iteration, - error_moments, - np.linalg.norm(self.on_diagonal[i + 1]), - np.linalg.norm(self.off_diagonal[i]), - error_sqrt, - error_inv_sqrt, - ) - - def get_eigenfunctions(self, iteration=None): - """ - Return the eigenfunctions. - """ - - if iteration is None: - iteration = self.iteration - - h_tri = util.build_block_tridiagonal( - [self.on_diagonal[i] for i in range(iteration + 1)], - [self.off_diagonal[i] for i in range(iteration)], - ) - - orth = util.matrix_power( - self.moments[0], - 0.5, - hermitian=self.hermitian, - return_error=False, - ) - - eigvals, eigvecs = np.linalg.eigh(h_tri) - eigvecs[: self.nphys] = np.dot(orth, eigvecs[: self.nphys]) - - return eigvals, eigvecs - - def get_auxiliaries(self, iteration=None): - """ - Return the self-energy auxiliaries. - """ - - if iteration is None: - iteration = self.iteration - - h_tri = util.build_block_tridiagonal( - [self.on_diagonal[i] for i in range(iteration + 1)], - [self.off_diagonal[i] for i in range(iteration)], - ) - - energies, rotated_couplings = np.linalg.eigh(h_tri[self.nphys :, self.nphys :]) - if energies.size: - couplings = np.dot(self.off_diagonal[0].T.conj(), rotated_couplings[: self.nphys]) - else: - couplings = np.zeros((self.nphys, 0), dtype=rotated_couplings.dtype) - - return energies, couplings - - def _kernel(self, iteration=None): - if self.iteration is None: - self.initialise_recurrence() - if iteration is None: - iteration = self.max_cycle - while self.iteration < iteration: - self.recurrence_iteration() - - self.log.info("-" * 89) - self.log.info("Block Lanczos moment recurrence completed to iteration %d.", self.iteration) - - if iteration is None: - iteration = self.max_cycle - - eigvals, eigvecs = self.get_eigenfunctions(iteration=iteration) - - self.log.info(util.print_dyson_orbitals(eigvals, eigvecs, self.nphys)) - - return eigvals, eigvecs - - @property - def static(self): - # Static part of the self-energy is equal to the zeroth order - # moment of the Green's function - return self.moments[0] - - @property - def nphys(self): - return self.moments[0].shape[0] - - -class MBLGF_NoSymm(MBLGF_Symm): - """ - Moment-conserving block Lanczos eigensolver, conserving the - moments of the Green's function, for a non-Hermitian Green's - function. - - Input - ----- - moments : numpy.ndarray - Moments of the Green's function. - - Parameters - ---------- - max_cycle : int, optional - Maximum number of iterations. If `None`, perform as many as - the inputted number of moments permits. Default value is - `None`. - - Returns - ------- - eigvals : numpy.ndarray - Eigenvalues of the matrix, representing the energies of the - Green's function. - eigvecs : tuple of numpy.ndarray - Left- and right-hand eigenvectors of the matrix, which provide - the Dyson orbitals once projected into the physical space. - """ - - def __init__(self, moments, **kwargs): - # Input: - self.moments = moments - - # Parameters: - self.max_cycle = kwargs.pop("max_cycle", None) - self.hermitian = False - - max_cycle_limit = (len(moments) - 2) // 2 - if self.max_cycle is None: - self.max_cycle = max_cycle_limit - if self.max_cycle > max_cycle_limit: - raise ValueError( - "`max_cycle` cannot be more than (M-2)/2, where " - "M is the number of inputted moments." - ) - - # Base class: - BaseSolver.__init__(self, moments, **kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > max_cycle: %s", self.max_cycle) - self.log.info(" > hermitian: %s", self.hermitian) - - # Caching: - self._cache = {} - self.coefficients = [ - RecurrenceCoefficients( - self.moments[0].shape, - hermitian=self.hermitian, - dtype=np.result_type(*self.moments), - ), - RecurrenceCoefficients( - self.moments[0].shape, - hermitian=self.hermitian, - dtype=np.result_type(*self.moments), - ), - ] - self.on_diagonal = {} - self.off_diagonal = [{}, {}] - self.orth = None - self.iteration = None - - def initialise_iteration_table(self): - """ - Print the header for the table sumarising the iterations. - """ - - self.log.info("-" * 106) - self.log.info( - "{:^4s} {:^16s} {:^50s} {:^33}".format( - "", - "", - "Norm of matrix", - "Norm of removed space", - ) - ) - self.log.info( - "{:^4s} {:^16s} {:^50s} {:^33}".format( - "Iter", - "Moment error", - "-" * 50, - "-" * 33, - ) - ) - self.log.info( - "%4s %16s %16s %16s %16s %16s %16s", - "", - "", - "On-diagonal", - "Off-diagonal ↑", - "Off-diagonal ↓", - "Square root", - "Inv. square root", - ) - self.log.info( - "%4s %16s %16s %16s %16s %16s %16s", - "-" * 4, - "-" * 16, - "-" * 16, - "-" * 16, - "-" * 16, - "-" * 16, - "-" * 16, - ) - - def _check_moment_error(self, iteration=None): - """ - Check the error in the moments at a given iteration. - """ - - if iteration is None: - iteration = self.iteration - - energies, (left, right) = self.get_dyson_orbitals(iteration=iteration) - - moments_recovered = [] - for n in range(2 * iteration + 2): - moments_recovered.append(np.dot(left, right.T.conj())) - left = left * energies[None] - - error_moments = util.scaled_error( - np.array(moments_recovered), - self.moments[: 2 * iteration + 2], - ) - - return error_moments - - def initialise_recurrence(self): - """ - Initialise the recurrences - performs the 'zeroth' iteration. - - This iteration is essentially equivalent to solving a generalised - eigenvalue problem on the Fock matrix in the physical space. - """ - - # Initialise the table - self.initialise_iteration_table() - - self.iteration = 0 - - # Calculate the orthogonalisation matrix - self.orth, error_inv_sqrt = util.matrix_power( - self.moments[0], - -0.5, - hermitian=self.hermitian, - return_error=True, - ) - - # Add zero matrix to out-of-bounds off-diagonal to simplify logic - self.off_diagonal[0][-1] = self.coefficients[0].zero - self.off_diagonal[1][-1] = self.coefficients[1].zero - - # Zeroth order on-diagonal block is the orthogonalised first - # moment (equal to the orthogonalised static part of the - # matrix corresponding to the solution moments) - self.on_diagonal[0] = self.orthogonalised_moment(1) - - # Check the error in the moments up to this iteration - error_moments = self._check_moment_error() - - # Logging - self.log.info( - "%4d %16.3g %16.3g %16s %16s %16s %16.3g", - 0, - error_moments, - np.linalg.norm(self.on_diagonal[0]), - "", - "", - "", - error_inv_sqrt, - ) - - def recurrence_iteration(self): - """ - Perform an iteration of the recurrence. - """ - - self.iteration += 1 - i = self.iteration - 1 - - if self.iteration > self.max_cycle: - raise ValueError( - "Cannot perform more iterations than permitted " - "by `max_cycle` or (M-2)/2 where M is the number " - "of inputted moments." - ) - - # Find the square of the next off-diagonal blocks - off_diagonal_squared = [ - self.coefficients[0].zero.astype(complex).copy(), - self.coefficients[1].zero.astype(complex).copy(), - ] - for j in range(i + 2): - for k in range(i + 1): - off_diagonal_squared[0] += np.linalg.multi_dot( - ( - self.coefficients[1][i + 1, k + 1], - self.orthogonalised_moment(j + k + 1), - self.coefficients[0][i + 1, j], - ) - ) - off_diagonal_squared[1] += np.linalg.multi_dot( - ( - self.coefficients[1][i + 1, j], - self.orthogonalised_moment(j + k + 1), - self.coefficients[0][i + 1, k + 1], - ) - ) - - off_diagonal_squared[0] -= np.dot( - self.on_diagonal[i], - self.on_diagonal[i], - ) - off_diagonal_squared[1] -= np.dot( - self.on_diagonal[i], - self.on_diagonal[i], - ) - if i: - off_diagonal_squared[0] -= np.dot( - self.off_diagonal[1][i - 1], - self.off_diagonal[1][i - 1], - ) - off_diagonal_squared[1] -= np.dot( - self.off_diagonal[0][i - 1], - self.off_diagonal[0][i - 1], - ) - - # Get the next off-diagonal blocks - self.off_diagonal[0][i], error_sqrt_upper = util.matrix_power( - off_diagonal_squared[0], - 0.5, - hermitian=self.hermitian, - return_error=True, - ) - self.off_diagonal[1][i], error_sqrt_lower = util.matrix_power( - off_diagonal_squared[1], - 0.5, - hermitian=self.hermitian, - return_error=True, - ) - error_sqrt = np.sqrt(error_sqrt_upper**2 + error_sqrt_lower**2) - - # Get the inverse of the off-diagonal blocks - off_diagonal_inv_upper, error_inv_sqrt_upper = util.matrix_power( - off_diagonal_squared[0], - -0.5, - hermitian=self.hermitian, - return_error=True, - ) - off_diagonal_inv_lower, error_inv_sqrt_lower = util.matrix_power( - off_diagonal_squared[1], - -0.5, - hermitian=self.hermitian, - return_error=True, - ) - error_inv_sqrt = np.sqrt(error_inv_sqrt_upper**2 + error_inv_sqrt_lower**2) - - for j in range(i + 2): - residual = ( - +self.coefficients[0][i + 1, j] - - np.dot(self.coefficients[0][i + 1, j + 1], self.on_diagonal[i]) - - np.dot(self.coefficients[0][i, j + 1], self.off_diagonal[0][i - 1]) - ) - self.coefficients[0][i + 2, j + 1] = np.dot(residual, off_diagonal_inv_lower) - - residual = ( - +self.coefficients[1][i + 1, j] - - np.dot(self.on_diagonal[i], self.coefficients[1][i + 1, j + 1]) - - np.dot(self.off_diagonal[1][i - 1], self.coefficients[1][i, j + 1]) - ) - self.coefficients[1][i + 2, j + 1] = np.dot(off_diagonal_inv_upper, residual) - - self.on_diagonal[i + 1] = self.coefficients[0].zero.astype(complex).copy() - for j in range(i + 2): - for k in range(i + 2): - self.on_diagonal[i + 1] += np.linalg.multi_dot( - ( - self.coefficients[1][i + 2, k + 1], - self.orthogonalised_moment(j + k + 1), - self.coefficients[0][i + 2, j + 1], - ) - ) - - # Check the error in the moments up to this iteration - error_moments = self._check_moment_error() - - # Logging - self.log.info( - "%4d %16.3g %16.3g %16.3g %16.3g %16.3g %16.3g", - self.iteration, - error_moments, - np.linalg.norm(self.on_diagonal[i + 1]), - np.linalg.norm(self.off_diagonal[0][i]), - np.linalg.norm(self.off_diagonal[1][i]), - error_sqrt, - error_inv_sqrt, - ) - if self.iteration == self.max_cycle: - self.log.info("-" * 106) - - def get_eigenfunctions(self, iteration=None): - """ - Return the eigenfunctions. - """ - - if iteration is None: - iteration = self.iteration - - h_tri = util.build_block_tridiagonal( - [self.on_diagonal[i] for i in range(iteration + 1)], - [self.off_diagonal[0][i] for i in range(iteration)], - [self.off_diagonal[1][i] for i in range(iteration)], - ) - - orth = util.matrix_power( - self.moments[0], - 0.5, - hermitian=self.hermitian, - return_error=False, - ) - - eigvals, eigvecs = np.linalg.eig(h_tri) - mask = np.argsort(eigvals.real) - eigvals = eigvals[mask] - eigvecs = eigvecs[:, mask] - - eigvecs_l = eigvecs - eigvecs_r = np.linalg.inv(eigvecs).T.conj() - - eigvecs_l[: self.nphys] = np.dot(orth, eigvecs_l[: self.nphys]) - eigvecs_r[: self.nphys] = np.dot(orth.T.conj(), eigvecs_r[: self.nphys]) - eigvecs = (eigvecs_l, eigvecs_r) - - return eigvals, eigvecs - - def get_auxiliaries(self, iteration=None): - """ - Return the self-energy auxiliaries. - """ - - if iteration is None: - iteration = self.iteration - - h_tri = util.build_block_tridiagonal( - [self.on_diagonal[i] for i in range(iteration + 1)], - [self.off_diagonal[0][i] for i in range(iteration)], - [self.off_diagonal[1][i] for i in range(iteration)], - ) - - energies, rotated_couplings = np.linalg.eig(h_tri[self.nphys :, self.nphys :]) - if energies.size: - couplings_l = np.dot( - self.off_diagonal[0][0].T.conj(), - rotated_couplings[: self.nphys], - ) - couplings_r = np.dot( - self.off_diagonal[1][0].T.conj(), - np.linalg.inv(rotated_couplings).T.conj()[: self.nphys], - ) - else: - couplings_l = np.zeros((self.nphys, 0), dtype=rotated_couplings.dtype) - couplings_r = np.zeros((self.nphys, 0), dtype=rotated_couplings.dtype) - - return energies, (couplings_l, couplings_r) - - -def MBLGF(moments, **kwargs): - """ - Wrapper to construct a solver based on the Hermiticity of the - input, either by the `hermitian` keyword argument or by the - structure of the input matrices. - """ - - if "hermitian" in kwargs: - hermitian = kwargs.pop("hermitian") - else: - hermitian = all(np.allclose(m, m.T.conj()) for m in moments) - - if hermitian: - return MBLGF_Symm(moments, **kwargs) - else: - return MBLGF_NoSymm(moments, **kwargs) - - -class MixedMBLGF: - """ - Mix multiple moment block Lanczos solvers for moments of the - Green's function, overloading the appropriate functions - useful - for example when applying particle and hole separation. Solvers - must correspond to the same physical space (same dimension), but - not necessarily the same physical part. - - Input - ----- - solvers : iterable of MBLGF_Symm or MBLGF_NoSymm - List of solvers to combine. - """ - - def __init__(self, *solvers): - # Input: - assert len(solvers) - self.solvers = solvers - - # Check that the physical spaces are the same: - try: - assert len(set(solver.nphys for solver in self.solvers)) == 1 - - static_parts = [] - for solver in solvers: - static_parts.append(solver.static) - - except AssertionError as e: - raise NotImplementedError( - "Solvers with different numbers of physical degrees of freedom cannot currently be " - "mixed." - ) - - # Caching: - self._static = None - - def initialise_recurrence(self): - for solver in self.solvers: - solver.initialise_recurrence - - def recurrence_iteration(self): - for solver in self.solvers: - solver.recurrence_iteration - - def kernel(self, *args, **kwargs): - for solver in self.solvers: - solver.kernel(*args, **kwargs) - - def get_auxiliaries(self, *args, **kwargs): - energies, orbitals = self.get_dyson_orbitals(*args, **kwargs) - - if isinstance(orbitals, tuple): - # Work with transpose of orbitals: - orbitals_l, orbitals_r = orbitals - orbitals_l = orbitals_l.T.conj() - orbitals_r = orbitals_r.T.conj() - - # Biorthogonalise orbitals: - mat = np.dot(orbitals_l.T.conj(), orbitals_r) - l, r = scipy.linalg.lu(mat, permute_l=True) - orbitals_l = np.dot(orbitals_l, np.linalg.inv(l)) - orbitals_r = np.dot(orbitals_r, np.linalg.inv(r).T.conj()) - - # Find a basis for the null space: - null_space = np.eye(orbitals_l.shape[0]) - np.dot(orbitals_l, orbitals_r.T.conj()) - w, rest_l = np.linalg.eig(null_space) - rest_r = np.linalg.inv(rest_l).T.conj() - rest_r = rest_r[:, np.abs(w) > 0.5] - rest_l = rest_l[:, np.abs(w) > 0.5] - - # Combine vectors: - vectors_l = np.block([orbitals_l, rest_l]) - vectors_r = np.block([orbitals_r, rest_r]) - - # Construct the Hamiltonian: - ham = np.dot(vectors_l.T.conj() * energies[None], vectors_r) - - # Rotate into arrowhead form: - w, v = np.linalg.eig(ham[self.nphys :, self.nphys :]) - v = np.block( - [ - [np.eye(self.nphys), np.zeros((self.nphys, w.size))], - [np.zeros((w.size, self.nphys)), v], - ] - ) - ham = np.linalg.multi_dot((np.linalg.inv(v), ham, v)) - - # Extract auxiliary parameters: - static = ham[: self.nphys, : self.nphys] - energies = np.diag(ham[self.nphys :, self.nphys :]) - couplings = ( - ham[: self.nphys, self.nphys :], - ham[self.nphys :, : self.nphys].T.conj(), - ) - - else: - # Work with transpose of orbitals: - orbitals = orbitals.T.conj() - - # Find a basis for the null space: - null_space = np.eye(orbitals.shape[0]) - np.dot(orbitals, orbitals.T.conj()) - w, rest = np.linalg.eigh(null_space) - rest = rest[:, np.abs(w) > 0.5] - - # Combine vectors: - vectors = np.block([orbitals, rest]) - - # Construct the Hamiltonian: - ham = np.dot(vectors.T.conj() * energies[None], vectors) - - # Rotate into arrowhead form: - w, v = np.linalg.eigh(ham[self.nphys :, self.nphys :]) - v = np.block( - [ - [np.eye(self.nphys), np.zeros((self.nphys, w.size))], - [np.zeros((w.size, self.nphys)), v], - ] - ) - ham = np.linalg.multi_dot((v.T.conj(), ham, v)) - - # Extract auxiliary parameters: - static = ham[: self.nphys, : self.nphys] - energies = np.diag(ham[self.nphys :, self.nphys :]) - couplings = ham[: self.nphys, self.nphys :] - - self._static = static - - return energies, couplings - - def get_eigenfunctions(self, *args, **kwargs): - hermitian = True - eigvals = [] - eigvecs_l = [] - eigvecs_r = [] - - for solver in self.solvers: - eigvals_, eigvecs_ = solver.get_eigenfunctions(*args, **kwargs) - eigvals.append(eigvals_) - - if isinstance(eigvecs_, tuple): - hermitian = False - eigvecs_l.append(eigvecs_[0]) - eigvecs_r.append(eigvecs_[1]) - else: - eigvecs_l.append(eigvecs_) - eigvecs_r.append(eigvecs_) - - eigvals = np.concatenate(eigvals) - - if hermitian: - eigvecs = np.concatenate(eigvecs_l, axis=1) - else: - eigvecs_l = np.concatenate(eigvecs_l, axis=1) - eigvecs_r = np.concatenate(eigvecs_r, axis=1) - eigvecs = (eigvecs_l, eigvecs_r) - - return eigvals, eigvecs - - def get_dyson_orbitals(self, *args, **kwargs): - eigvals, eigvecs = self.get_eigenfunctions(*args, **kwargs) - - if isinstance(eigvecs, tuple): - # The eigvecs are already inverted if in a tuple - eigvecs = (eigvecs[0][: self.nphys], eigvecs[1][: self.nphys]) - else: - eigvecs = eigvecs[: self.nphys] - - return eigvals, eigvecs - - def get_self_energy(self, *args, chempot=0.0, **kwargs): - return Lehmann(*self.get_auxiliaries(*args, **kwargs), chempot=chempot) - - def get_greens_function(self, *args, chempot=0.0, **kwargs): - return Lehmann(*self.get_dyson_orbitals(*args, **kwargs), chempot=chempot) - - def _check_moment_error(self, *args, **kwargs): - error = 0 - for solver in self.solvers: - error += solver._check_moment_error(*args, **kwargs) - return error - - @property - def static(self): - # Static part of the combined system can be determined when - # rotating back into an auxiliary representation using - # self.get_auxiliaries() - if self._static is None: - raise ValueError( - "To determine `MixedMBLGF.static`, one must first call " - "`MixedMBLGF.get_auxiliaries()`." - ) - return self._static - - @property - def nphys(self): - return self.solvers[0].nphys - - @property - def log(self): - return self.solvers[0].log diff --git a/dyson/solvers/mblse.py b/dyson/solvers/mblse.py deleted file mode 100644 index 0387737..0000000 --- a/dyson/solvers/mblse.py +++ /dev/null @@ -1,981 +0,0 @@ -""" -Moment-conserving block Lanczos eigensolver, conserving moments of -the self-energy. -""" - -import warnings - -import numpy as np - -from dyson import util -from dyson.lehmann import Lehmann -from dyson.solvers import BaseSolver - -# TODO improve inheritence - - -class RecurrenceCoefficients: - """ - Recurrence coefficient container. - """ - - def __init__(self, shape, hermitian=True, force_orthogonality=True, dtype=np.float64): - self.hermitian = hermitian - self.force_orthogonality = force_orthogonality - self.zero = np.zeros(shape, dtype=dtype) - self.data = {} - - def __getitem__(self, key): - i, j, n = key - - if i == 0 or j == 0: - # Zeroth order Lanczos vectors are zero - return self.zero - else: - # Return Q_{i}^† H^{n} Q_{j} - if (not self.hermitian) or i >= j: - return self.data[i, j, n] - else: - return self.data[j, i, n].T.conj() - - def __setitem__(self, key, val): - i, j, n = key - - if n == 0 and self.force_orthogonality: - val = np.eye(self.zero.shape[0]) - - if self.hermitian and i == j: - val = 0.5 * util.hermi_sum(val) - - # Set Q_{i}^† H^{n} Q_{j} - if (not self.hermitian) or i >= j: - self.data[i, j, n] = val - else: - self.data[j, i, n] = val.T.conj() - - -class MBLSE_Symm(BaseSolver): - """ - Moment-conserving block Lanczos eingsolver, conserving the - moments of the self-energy, for a Hermitian self-energy. - - Input - ----- - static : numpy.ndarray - Static part of the self-energy. - moments : numpy.ndarray - Moments of the self-energy. - - Parameters - ---------- - max_cycle : int, optional - Maximum number of iterations. If `None`, perform as many as - the inputted number of moments permits. Default value is - `None`. - - Returns - ------- - eigvals : numpy.ndarray - Eigenvalues of the matrix, representing the energies of the - Green's function. - eigvecs : numpy.ndarray - Eigenvectors of the matrix, which provide the Dyson orbitals - once projected into the physical space. - """ - - def __init__(self, static, moments, **kwargs): - # Input: - self.static = static - self.moments = moments - - # Parameters: - self.max_cycle = kwargs.pop("max_cycle", None) - self.hermitian = True - - max_cycle_limit = (len(moments) - 2) // 2 - if self.max_cycle is None: - self.max_cycle = max_cycle_limit - if self.max_cycle > max_cycle_limit: - raise ValueError( - "`max_cycle` cannot be more than (M-2)/2, where " - "M is the number of inputted moments." - ) - - # Base class: - super().__init__(static, moments, **kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > max_cycle: %s", self.max_cycle) - self.log.info(" > hermitian: %s", self.hermitian) - - # Caching: - self._cache = {} - self.coefficients = RecurrenceCoefficients( - static.shape, - hermitian=True, - dtype=np.result_type(self.static, *self.moments), - ) - self.on_diagonal = {} - self.off_diagonal = {} - self.iteration = None - - # @util.cache - # def coefficient_times_off_diagonal(self, i, j, n): - # """ - # Compute Q_{i}^† H^{n} Q_{j} B_{j}^† - # """ - - # return np.dot( - # self.coefficients[i, j, n], - # self.off_diagonal[j].T.conj(), - # ) - - # @util.cache - # def coefficient_times_on_diagonal(self, i, j, n): - # """ - # Compute Q_{i}^† H^{n} Q_{j} A_{j} - # """ - - # return np.dot( - # self.coefficients[i, j, n], - # self.on_diagonal[j], - # ) - - def orthogonalised_moment(self, n): - """ - Compute an orthogonalised moment. - """ - - orth = util.matrix_power(self.moments[0], -0.5, hermitian=True) - - return np.linalg.multi_dot( - ( - orth, - self.moments[n], - orth, - ) - ) - - def _check_moment_error(self, iteration=None): - """ - Check the error in the moments at a given iteration. - """ - - if iteration is None: - iteration = self.iteration - - energies, couplings = self.get_auxiliaries(iteration=iteration) - - left = couplings.copy() - moments_recovered = [] - for n in range(2 * iteration + 2): - moments_recovered.append(np.dot(left, couplings.T.conj())) - left = left * energies[None] - - error_moments = sum( - util.scaled_error(a, b) - for a, b in zip(moments_recovered, self.moments[: 2 * iteration + 2]) - ) - - return error_moments - - def initialise_recurrence(self): - """ - Initialise the recurrences - performs the 'zeroth' iteration. - """ - - self.log.info("-" * 89) - self.log.info( - "{:^4s} {:^16s} {:^33s} {:^33}".format( - "", - "", - "Norm of matrix", - "Norm of removed space", - ) - ) - self.log.info( - "{:^4s} {:^16s} {:^33s} {:^33}".format( - "Iter", - "Moment error", - "-" * 33, - "-" * 33, - ) - ) - self.log.info( - "%4s %16s %16s %16s %16s %16s", - "", - "", - "On-diagonal", - "Off-diagonal", - "Square root", - "Inv. square root", - ) - self.log.info( - "%4s %16s %16s %16s %16s %16s", - "-" * 4, - "-" * 16, - "-" * 16, - "-" * 16, - "-" * 16, - "-" * 16, - ) - - self.iteration = 0 - - # Zeroth order on-diagonal block is the static self-energy - self.on_diagonal[0] = self.static - - # Zeroth order off-diagonal block is the square-root of the - # zeroth order moment - self.off_diagonal[0], error_sqrt = util.matrix_power( - self.moments[0], - 0.5, - hermitian=True, - return_error=True, - ) - - # Populate the other orthogonalised moments - orth, error_inv_sqrt = util.matrix_power( - self.moments[0], - -0.5, - hermitian=True, - return_error=True, - ) - for n in range(2 * self.max_cycle + 2): - # FIXME orth recalculated n+1 times - self.coefficients[1, 1, n] = self.orthogonalised_moment(n) - - # First order on-diagonal block is the orthogonalised first - # order moment - self.on_diagonal[1] = self.coefficients[1, 1, 1] - - # Check the error in the moments up to this iteration - error_moments = self._check_moment_error(iteration=0) - - # Logging - self.log.info( - "%4d %16.3g %16.3g %16.3g %16.3g %16.3g", - 0, - error_moments, - np.linalg.norm(self.on_diagonal[1]), - np.linalg.norm(self.off_diagonal[0]), - error_sqrt, - error_inv_sqrt, - ) - - def recurrence_iteration(self): - """ - Perform an iteration of the recurrence for hermitian systems. - """ - - self.iteration += 1 - i = self.iteration - - if self.iteration > self.max_cycle: - raise ValueError( - "Cannot perform more iterations than permitted " - "by `max_cycle` or (M-2)/2 where M is the number " - "of inputted moments." - ) - - # Find the square of the next off-diagonal block - off_diagonal_squared = ( - +self.coefficients[i, i, 2] - - util.hermi_sum(np.dot(self.coefficients[i, i - 1, 1], self.off_diagonal[i - 1])) - - np.dot(self.coefficients[i, i, 1], self.coefficients[i, i, 1]) - ) - if self.iteration > 1: - off_diagonal_squared += np.dot( - self.off_diagonal[i - 1].T.conj(), - self.off_diagonal[i - 1], - ) - - # Get the next off-diagonal block - self.off_diagonal[i], error_sqrt = util.matrix_power( - off_diagonal_squared, - 0.5, - hermitian=True, - return_error=True, - ) - - # Get the inverse of the off-diagonal block - off_diagonal_inv, error_inv_sqrt = util.matrix_power( - off_diagonal_squared, - -0.5, - hermitian=True, - return_error=True, - ) - - for n in range(2 * (self.max_cycle - self.iteration + 1)): - residual = ( - +self.coefficients[i, i, n + 1] - - np.dot(self.off_diagonal[i - 1].T.conj(), self.coefficients[i - 1, i, n]) - - np.dot(self.on_diagonal[i], self.coefficients[i, i, n]) - ) - self.coefficients[i + 1, i, n] = np.dot(off_diagonal_inv, residual) - - residual = ( - +self.coefficients[i, i, n + 2] - - util.hermi_sum( - np.dot(self.coefficients[i, i - 1, n + 1], self.off_diagonal[i - 1]) - ) - - util.hermi_sum(np.dot(self.coefficients[i, i, n + 1], self.on_diagonal[i])) - + util.hermi_sum( - np.linalg.multi_dot( - ( - self.on_diagonal[i], - self.coefficients[i, i - 1, n], - self.off_diagonal[i - 1], - ) - ) - ) - + np.linalg.multi_dot( - ( - self.off_diagonal[i - 1].T.conj(), - self.coefficients[i - 1, i - 1, n], - self.off_diagonal[i - 1], - ) - ) - + np.linalg.multi_dot( - (self.on_diagonal[i], self.coefficients[i, i, n], self.on_diagonal[i]) - ) - ) - self.coefficients[i + 1, i + 1, n] = np.linalg.multi_dot( - ( - off_diagonal_inv, - residual, - off_diagonal_inv.T.conj(), - ) - ) - - # Extract the next on-diagonal block - self.on_diagonal[i + 1] = self.coefficients[i + 1, i + 1, 1] - - # Check the error in the moments up to this iteration - error_moments = self._check_moment_error() - - # Logging - self.log.info( - "%4d %16.3g %16.3g %16.3g %16.3g %16.3g", - self.iteration, - error_moments, - np.linalg.norm(self.on_diagonal[i + 1]), - np.linalg.norm(self.off_diagonal[i]), - error_sqrt, - error_inv_sqrt, - ) - - def get_auxiliaries(self, iteration=None): - """ - Return the compressed self-energy auxiliaries. - """ - - if iteration is None: - iteration = self.iteration - - h_tri = util.build_block_tridiagonal( - [self.on_diagonal[i] for i in range(iteration + 2)], - [self.off_diagonal[i] for i in range(iteration + 1)], - ) - - energies, rotated_couplings = np.linalg.eigh(h_tri[self.nphys :, self.nphys :]) - couplings = np.dot(self.off_diagonal[0].T.conj(), rotated_couplings[: self.nphys]) - - return energies, couplings - - def get_eigenfunctions(self, iteration=None): - """ - Return the eigenfunctions. - """ - - if iteration is None: - iteration = self.iteration - - energies, couplings = self.get_auxiliaries(iteration=iteration) - h_aux = np.block( - [ - [self.static, couplings], - [couplings.T.conj(), np.diag(energies)], - ] - ) - - eigvals, eigvecs = np.linalg.eigh(h_aux) - - return eigvals, eigvecs - - def _kernel(self, iteration=None): - if self.iteration is None: - self.initialise_recurrence() - if iteration is None: - iteration = self.max_cycle - while self.iteration < iteration: - self.recurrence_iteration() - - self.log.info("-" * 89) - self.log.info("Block Lanczos moment recurrence completed to iteration %d.", self.iteration) - - if iteration is None: - iteration = self.max_cycle - - eigvals, eigvecs = self.get_eigenfunctions(iteration=iteration) - - self.log.info(util.print_dyson_orbitals(eigvals, eigvecs, self.nphys)) - - return eigvals, eigvecs - - @property - def nphys(self): - return self.static.shape[0] - - -class MBLSE_NoSymm(MBLSE_Symm): - """ - Moment-conserving block Lanczos eingsolver, conserving the - moments of the self-energy, for a non-Hermitian self-energy. - - Input - ----- - static : numpy.ndarray - Static part of the self-energy. - moments : numpy.ndarray - Moments of the self-energy. - - Parameters - ---------- - max_cycle : int, optional - Maximum number of iterations. If `None`, perform as many as - the inputted number of moments permits. Default value is - `None`. - - Returns - ------- - eigvals : numpy.ndarray - Eigenvalues of the matrix, representing the energies of the - Green's function. - eigvecs : tuple of numpy.ndarray - Left- and right-hand eigenvectors of the matrix, which provide - the Dyson orbitals once projected into the physical space. - """ - - def __init__(self, static, moments, **kwargs): - # Input: - self.static = static - self.moments = moments - - # Parameters: - self.max_cycle = kwargs.pop("max_cycle", None) - self.hermitian = False - - max_cycle_limit = (len(moments) - 2) // 2 - if self.max_cycle is None: - self.max_cycle = max_cycle_limit - if self.max_cycle > max_cycle_limit: - raise ValueError( - "`max_cycle` cannot be more than (M-2)/2, where " - "M is the number of inputted moments." - ) - - # Base class: - BaseSolver.__init__(self, static, moments, **kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > max_cycle: %s", self.max_cycle) - self.log.info(" > hermitian: %s", self.hermitian) - - # Caching: - self._cache = {} - self.coefficients = RecurrenceCoefficients( - static.shape, - hermitian=False, - dtype=np.result_type(self.static, *self.moments), - ) - self.on_diagonal = {} - self.off_diagonal = {} - self.iteration = None - - def orthogonalised_moment(self, n): - """ - Compute an orthogonalised moment. - """ - - orth = util.matrix_power( - self.moments[0], - -0.5, - hermitian=False, - ) - - return np.linalg.multi_dot( - ( - orth, - self.moments[n], - orth, - ) - ) - - def _check_moment_error(self, iteration=None): - """ - Check the error in the moments at a given iteration. - """ - - if iteration is None: - iteration = self.iteration - - energies, couplings = self.get_auxiliaries(iteration=iteration) - - left = couplings[0].copy() - right = couplings[1] - moments_recovered = [] - for n in range(2 * iteration + 2): - moments_recovered.append(np.dot(left, right.T.conj())) - left = left * energies[None] - - error_moments = sum( - util.scaled_error(a, b) - for a, b in zip(moments_recovered, self.moments[: 2 * iteration + 2]) - ) - - return error_moments - - def initialise_recurrence(self): - """ - Initialise the recurrences - performs the 'zeroth' iteration. - """ - - self.log.info("-" * 106) - self.log.info( - "{:^4s} {:^16s} {:^50s} {:^33}".format( - "", - "", - "Norm of matrix", - "Norm of removed space", - ) - ) - self.log.info( - "{:^4s} {:^16s} {:^50s} {:^33}".format( - "Iter", - "Moment error", - "-" * 50, - "-" * 33, - ) - ) - self.log.info( - "%4s %16s %16s %16s %16s %16s %16s", - "", - "", - "On-diagonal", - "Off-diagonal ↑", - "Off-diagonal ↓", - "Square root", - "Inv. square root", - ) - self.log.info( - "%4s %16s %16s %16s %16s %16s %16s", - "-" * 4, - "-" * 16, - "-" * 16, - "-" * 16, - "-" * 16, - "-" * 16, - "-" * 16, - ) - - self.iteration = 0 - - # Zeroth order on-diagonal block is the static self-energy - self.on_diagonal[0] = self.static - - # Zeroth order off-diagonal blocks are the square-root of the - # zeroth order moment - self.off_diagonal[0], error_sqrt = util.matrix_power( - self.moments[0], - 0.5, - hermitian=False, - return_error=True, - ) - assert np.allclose(np.dot(self.off_diagonal[0], self.off_diagonal[0]), self.moments[0]) - - # Populate the other orthogonalised moments - orth, error_inv_sqrt = util.matrix_power( - self.moments[0], - -0.5, - hermitian=False, - return_error=True, - ) - for n in range(2 * self.max_cycle + 2): - # FIXME orth recalculated n+1 times - self.coefficients[1, 1, n] = self.orthogonalised_moment(n) - - # First order on-diagonal block is the orthogonalised first - # order moment - self.on_diagonal[1] = self.coefficients[1, 1, 1] - - # Check the error in the moments up to this iteration - error_moments = self._check_moment_error(iteration=0) - - # Logging - self.log.info( - "%4d %16.3g %16.3g %16.3g %16.3g %16.3g %16.3g", - 0, - error_moments, - np.linalg.norm(self.on_diagonal[1]), - np.linalg.norm(self.off_diagonal[0]), - np.linalg.norm(self.off_diagonal[0]), - error_sqrt, - error_inv_sqrt, - ) - - def recurrence_iteration(self): - """ - Perform an iteration of the recurrence for hermitian systems. - """ - - self.iteration += 1 - i = self.iteration - - if self.iteration > self.max_cycle: - raise ValueError( - "Cannot perform more iterations than permitted " - "by `max_cycle` or (M-2)/2 where M is the number " - "of inputted moments." - ) - - # Find the square of the next off-diagonal block - off_diagonal_squared = ( - +self.coefficients[i, i, 2] - - np.dot(self.coefficients[i, i, 1], self.coefficients[i, i, 1]) - - np.dot(self.coefficients[i, i - 1, 1], self.off_diagonal[i - 1]) - - np.dot(self.off_diagonal[i - 1], self.coefficients[i, i - 1, 1]) - ) - if self.iteration > 1: - off_diagonal_squared += np.dot( - self.off_diagonal[i - 1], - self.off_diagonal[i - 1], - ) - - # Get the next off-diagonal blocks - self.off_diagonal[i], error_sqrt = util.matrix_power( - off_diagonal_squared, - 0.5, - hermitian=False, - return_error=True, - ) - - # Get the inverse of the off-diagonal blocks - off_diagonal_inv, error_inv_sqrt = util.matrix_power( - off_diagonal_squared, - -0.5, - hermitian=False, - return_error=True, - ) - - for n in range(2 * (self.max_cycle - self.iteration + 1)): - residual = ( - +self.coefficients[i, i, n + 1] - - np.dot(self.off_diagonal[i - 1], self.coefficients[i - 1, i, n]) - - np.dot(self.on_diagonal[i], self.coefficients[i, i, n]) - ) - self.coefficients[i + 1, i, n] = np.dot(off_diagonal_inv, residual) - - residual = ( - +self.coefficients[i, i, n + 1] - - np.dot(self.coefficients[i, i - 1, n], self.off_diagonal[i - 1]) - - np.dot(self.coefficients[i, i, n], self.on_diagonal[i]) - ) - self.coefficients[i, i + 1, n] = np.dot(residual, off_diagonal_inv) - - residual = ( - +self.coefficients[i, i, n + 2] - - np.dot(self.coefficients[i, i - 1, n + 1], self.off_diagonal[i - 1]) - - np.dot(self.coefficients[i, i, n + 1], self.on_diagonal[i]) - - np.dot(self.off_diagonal[i - 1], self.coefficients[i - 1, i, n + 1]) - + np.linalg.multi_dot( - ( - self.off_diagonal[i - 1], - self.coefficients[i - 1, i - 1, n], - self.off_diagonal[i - 1], - ) - ) - + np.linalg.multi_dot( - ( - self.off_diagonal[i - 1], - self.coefficients[i - 1, i, n], - self.on_diagonal[i], - ) - ) - - np.dot(self.on_diagonal[i], self.coefficients[i, i, n + 1]) - + np.linalg.multi_dot( - ( - self.on_diagonal[i], - self.coefficients[i, i - 1, n], - self.off_diagonal[i - 1], - ) - ) - + np.linalg.multi_dot( - ( - self.on_diagonal[i], - self.coefficients[i, i, n], - self.on_diagonal[i], - ) - ) - ) - self.coefficients[i + 1, i + 1, n] = np.linalg.multi_dot( - ( - off_diagonal_inv, - residual, - off_diagonal_inv, - ) - ) - - # Extract the next on-diagonal block - self.on_diagonal[i + 1] = self.coefficients[i + 1, i + 1, 1] - - # Check the error in the moments up to this iteration - error_moments = self._check_moment_error() - - # Logging - self.log.info( - "%4d %16.3g %16.3g %16.3g %16.3g %16.3g %16.3g", - self.iteration, - error_moments, - np.linalg.norm(self.on_diagonal[i + 1]), - np.linalg.norm(self.off_diagonal[i]), - np.linalg.norm(self.off_diagonal[i]), - error_sqrt, - error_inv_sqrt, - ) - if self.iteration == self.max_cycle: - self.log.info("-" * 106) - - def get_auxiliaries(self, iteration=None): - """ - Return the compressed self-energy auxiliaries. - """ - - if iteration is None: - iteration = self.iteration - - h_tri = util.build_block_tridiagonal( - [self.on_diagonal[i] for i in range(iteration + 2)], - [self.off_diagonal[i] for i in range(iteration + 1)], - [self.off_diagonal[i] for i in range(iteration + 1)], - ) - - energies, rotated_couplings = np.linalg.eig(h_tri[self.nphys :, self.nphys :]) - couplings_left = np.dot( - self.off_diagonal[0], - rotated_couplings[: self.nphys], - ) - couplings_right = np.dot( - self.off_diagonal[0].T.conj(), - np.linalg.inv(rotated_couplings).T.conj()[: self.nphys], - ) - - return energies, (couplings_left, couplings_right) - - def get_eigenfunctions(self, iteration=None): - """ - Return the eigenfunctions. - """ - - if iteration is None: - iteration = self.iteration - - energies, (couplings_left, couplings_right) = self.get_auxiliaries(iteration=iteration) - h_aux = np.block( - [ - [self.static, couplings_left], - [couplings_right.T.conj(), np.diag(energies)], - ] - ) - - eigvals, eigvecs = np.linalg.eig(h_aux) - - return eigvals, eigvecs - - def _kernel(self, iteration=None): - if self.iteration is None: - self.initialise_recurrence() - if iteration is None: - iteration = self.max_cycle - while self.iteration < iteration: - self.recurrence_iteration() - - self.log.info("Block Lanczos moment recurrence completed to iteration %d.", self.iteration) - - if iteration is None: - iteration = self.max_cycle - - eigvals, eigvecs = self.get_eigenfunctions(iteration=iteration) - - self.log.info(util.print_dyson_orbitals(eigvals, eigvecs, self.nphys)) - - return eigvals, eigvecs - - @property - def nphys(self): - return self.static.shape[0] - - -def MBLSE(static, moments, **kwargs): - """ - Wrapper to construct a solver based on the Hermiticity of the - input, either by the `hermitian` keyword argument or by the - structure of the input matrices. - """ - - if "hermitian" in kwargs: - hermitian = kwargs.pop("hermitian") - else: - hermitian = all(np.allclose(m, m.T.conj()) for m in [static, *moments]) - - if hermitian: - return MBLSE_Symm(static, moments, **kwargs) - else: - return MBLSE_NoSymm(static, moments, **kwargs) - - -class MixedMBL: - """ - Deprecated class - immediately raises a `NotImplementedError`. - """ - - def __init__(self, *args, **kwargs): - raise NotImplementedError( - "`MixedMBL` is deprecated in favour of " - "`dyson.mblse.MixedMBLSE` and `dyson.mblgf.MixedMBLGF." - ) - - -class MixedMBLSE: - """ - Mix multiple moment block Lanczos solvers for moments of the - self-energy, overloading the appropriate functions - useful for - example when applying particle and hole separation. Solvers must - correspond to the same physical space (same dimension, and same - static part). - - Input - ----- - solvers : iterable of MBLSE_Symm - List of solvers to combine. - """ - - def __init__(self, *solvers): - # Input: - assert len(solvers) - self.solvers = solvers - - # Check that the physical spaces are the same: - try: - assert len(set(solver.nphys for solver in self.solvers)) == 1 - - static_parts = [] - for solver in solvers: - static_parts.append(solver.static) - - except AssertionError as e: - raise NotImplementedError( - "Solvers with different physical degrees of freedom cannot currently be mixed." - ) - - def initialise_recurrence(self): - for solver in self.solvers: - solver.initialise_recurrence - - def recurrence_iteration(self): - for solver in self.solvers: - solver.recurrence_iteration - - def kernel(self, *args, **kwargs): - for solver in self.solvers: - solver.kernel(*args, **kwargs) - - def get_auxiliaries(self, *args, **kwargs): - hermitian = True - energies = [] - couplings_l = [] - couplings_r = [] - - for solver in self.solvers: - energies_, couplings_ = solver.get_auxiliaries(*args, **kwargs) - energies.append(energies_) - - if isinstance(couplings_, tuple): - hermitian = False - couplings_l.append(couplings_[0]) - couplings_r.append(couplings_[1]) - else: - couplings_l.append(couplings_) - couplings_r.append(couplings_) - - energies = np.concatenate(energies) - - if hermitian: - couplings = np.concatenate(couplings_l, axis=1) - else: - couplings_l = np.concatenate(couplings_l, axis=1) - couplings_r = np.concatenate(couplings_r, axis=1) - couplings = (couplings_l, couplings_r) - - return energies, couplings - - def get_eigenfunctions(self, *args, **kwargs): - energies, couplings = self.get_auxiliaries(*args, **kwargs) - - if isinstance(couplings, tuple): - couplings_l, couplings_r = couplings - else: - couplings_l = couplings_r = couplings - - h_aux = np.block( - [ - [self.static, couplings_l], - [couplings_r.T.conj(), np.diag(energies)], - ] - ) - - if isinstance(couplings, tuple): - eigvals, eigvecs = np.linalg.eig(h_aux) - else: - eigvals, eigvecs = np.linalg.eigh(h_aux) - - return eigvals, eigvecs - - def get_dyson_orbitals(self, *args, **kwargs): - eigvals, eigvecs = self.get_eigenfunctions(*args, **kwargs) - - if any( - not solver.hermitian for solver in self.solvers - ): # FIXME make more rigorous throughout - eigvecs = (eigvecs, np.linalg.inv(eigvecs).T.conj()) - eigvecs = (eigvecs[0][: self.nphys], eigvecs[1][: self.nphys]) - else: - eigvecs = eigvecs[: self.nphys] - - return eigvals, eigvecs - - def get_self_energy(self, *args, chempot=0.0, **kwargs): - return Lehmann(*self.get_auxiliaries(*args, **kwargs), chempot=chempot) - - def get_greens_function(self, *args, chempot=0.0, **kwargs): - return Lehmann(*self.get_dyson_orbitals(*args, **kwargs), chempot=chempot) - - def _check_moment_error(self, *args, **kwargs): - error = 0 - for solver in self.solvers: - error += solver._check_moment_error(*args, **kwargs) - return error - - @property - def nphys(self): - return self.solvers[0].nphys - - @property - def static(self): - return self.solvers[0].static - - @property - def log(self): - return self.solvers[0].log diff --git a/dyson/solvers/self_consistent.py b/dyson/solvers/self_consistent.py deleted file mode 100644 index c8ac71d..0000000 --- a/dyson/solvers/self_consistent.py +++ /dev/null @@ -1,223 +0,0 @@ -""" -Self-consistent solution to the Dyson equation. -""" - -import numpy as np - -from dyson import Lehmann, NullLogger -from dyson.solvers import AufbauPrinciple, AuxiliaryShift, BaseSolver, DensityRelaxation - - -class SelfConsistent(BaseSolver): - """ - Self-consistent solution to the Dyson equation. - - Parameters - ---------- - get_se : callable - Callable that returns the self-energy. Takes a Green's - function in the format of a `Lehmann` object as input, which - provides the basis in which the self-energy is to be - constructed. - get_fock : callable - Callable that returns the Fock matrix. Takes a density matrix - in the MO basis as input. Default value is `None`. - gf_init : dyson.Lehmann - Initial guess for the Green's function. - nelec : int, optional - Number of electrons. If not provided, the number is inferred - from the initial guess for the Green's function. Default - value is `None`. - occupancy : int, optional - Occupancy of each state, i.e. `2` for a restricted reference - and `1` for other references. Default value is `2`. - relax_solver : dyson.solvers.BaseSolver, optional - Solver for relaxing the density matrix or chemical potential. - Must be one of {`None`, `dyson.solvers.AufbauPrinciple`, - `dyson.solvers.AuxiliaryShift`, - `dyson.solvers.DensityRelaxation`}. If provided, the - `get_fock` argument must be provided. Default value is - `None`. - max_cycle : int, optional - Maximum number of iterations. Default value is `50`. - conv_tol : float, optional - Convergence threshold in the first moment of the Green's - function. Default value is `1e-8`. - """ - - def __init__(self, get_se, get_fock, gf_init, **kwargs): - # Input: - self._get_se = get_se - self._get_fock = get_fock - self.gf_init = gf_init - - # Parameters: - self._nelec = kwargs.pop("nelec", None) - self.occupancy = kwargs.pop("occupancy", 2) - self.relax_solver = kwargs.pop("relax_solver", None) - self.max_cycle = kwargs.pop("max_cycle", 50) - self.conv_tol = kwargs.pop("conv_tol", 1e-8) - - # Base class: - super().__init__(**kwargs) - - # Logging: - self.log.info("Options:") - self.log.info(" > nelec: %s", self.nelec) - self.log.info(" > occupancy: %s", self.occupancy) - self.log.info(" > relax_solver: %s", self.relax_solver) - self.log.info(" > max_cycle: %s", self.max_cycle) - self.log.info(" > conv_tol: %s", self.conv_tol) - - # Caching: - self.converged = False - self.se_res = None - self.gf_res = None - - def get_se(self, gf, se_prev=None): - """ - Update the self-energy using a particular Green's function. - - Parameters - ---------- - gf : dyson.Lehmann - Green's function. - se_prev : dyson.Lehmann, optional - Previous self-energy. Default value is `None`. - - Returns - ------- - se : dyson.Lehmann - Self-energy. - """ - - return self._get_se(gf) - - def get_fock(self, rdm1): - """ - Update the Fock matrix using a particular density matrix. - - Parameters - ---------- - rdm1 : numpy.ndarray - Density matrix in the MO basis. - - Returns - ------- - fock : numpy.ndarray - Fock matrix. - """ - - return self._get_fock(rdm1) - - def _kernel(self): - """ - Perform the self-consistent solution of the Dyson equation. - """ - - gf = self.gf_init - gf_prev = gf - se = self.get_se(gf) - se_prev = None - gap = gf.physical().virtual().energies[0] - gf.physical().occupied().energies[-1] - - self.log.info("-" * 58) - self.log.info( - "{:^6s} {:^12s} {:^12s} {:^12s} {:^12s}".format( - "Iter", - "Gap", - "Gap error", - "Nelec error", - "Chempot", - ) - ) - self.log.info( - "%6s %12s %12s %12s %12s", - "-" * 6, - "-" * 12, - "-" * 12, - "-" * 12, - "-" * 12, - ) - - for i in range(1, self.max_cycle + 1): - gf_prev = gf - gap_prev = gap - - if self.relax_solver: - if self.relax_solver is DensityRelaxation: - fock = self.get_fock - else: - rdm1 = gf.occupied().moment(0) * self.occupancy - fock = self.get_fock(rdm1) - - solver = self.relax_solver(fock, se, self.nelec, log=NullLogger()) - solver.kernel() - - gf = solver.get_greens_function() - se = solver.get_self_energy() - - else: - rdm1 = gf.occupied().moment(0) * self.occupancy - fock = self.get_fock(rdm1) - - w, v = se.diagonalise_matrix_with_projection(fock) - gf = Lehmann(w, v, chempot=se.chempot) - - se_prev = se.copy() - se = self.get_se(gf, se_prev=se_prev) - - gap_prev = gap - ip = -gf.physical().occupied().energies[-1] - ea = gf.physical().virtual().energies[0] - gap = ip + ea - - n_error = abs(np.trace(gf.occupied().moment(0)) * self.occupancy - self.nelec) - gap_error = abs(gap - gap_prev) - - self.log.info( - "{:6d} {:12.8f} {:12.6g} {:12.6g} {:12.6f}".format( - i, - gap, - gap_error, - n_error, - gf.chempot, - ) - ) - - if gap_error < self.conv_tol: - self.converged = True - break - - self.log.info("-" * 58) - - self.flag_convergence(self.converged) - - self.se_res = se - self.gf_res = gf - - return gf, se, self.converged - - @property - def nelec(self): - """ - Number of electrons. - """ - - if self._nelec is None: - rdm1 = self.gf_init.occupied().moment(0) * self.occupancy - self._nelec = int(np.round(np.trace(rdm1))) - - return self._nelec - - def get_auxiliaries(self): - return self.se_res.energies, self.se_res.couplings - - def get_dyson_orbitals(self): - return self.gf_res.energies, self.gf_res.couplings - - def get_self_energy(self): - return self.se_res - - def get_greens_function(self): - return self.gf_res diff --git a/dyson/solvers/solver.py b/dyson/solvers/solver.py index 6c01223..d5df977 100644 --- a/dyson/solvers/solver.py +++ b/dyson/solvers/solver.py @@ -1,94 +1,185 @@ -""" -Solver base class. -""" +"""Base class for Dyson equation solvers.""" -import numpy as np +from __future__ import annotations -from dyson import Lehmann, default_log, init_logging +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING +from rich import box +from rich.table import Table -class BaseSolver: - """ - Base class for all solvers. - """ +from dyson import console, printing +from dyson.representations.enums import RepresentationEnum +from dyson.representations.lehmann import Lehmann +from dyson.typing import Array - def __init__(self, *args, **kwargs): - self.log = kwargs.pop("log", default_log) - if self.log is None: - self.log = default_log - init_logging(self.log) - self.log.info("") - self.log.info("%s", self.__class__.__name__) - self.log.info("%s", "*" * len(self.__class__.__name__)) +if TYPE_CHECKING: + from typing import Any - # Check all the arguments have now been consumed: - if len(kwargs): - for key, val in kwargs.items(): - self.log.warn("Argument `%s` invalid" % key) + from dyson.expressions.expression import BaseExpression + from dyson.representations.dynamic import Dynamic + from dyson.representations.spectral import Spectral - def kernel(self, *args, **kwargs): - """ - Driver function. Classes inheriting the `BaseSolver` should - implement `_kernel`, which is called by this function. If - the solver has a `_cache`, this function clears it. - """ - out = self._kernel(*args, **kwargs) +class BaseSolver(ABC): + """Base class for Dyson equation solvers.""" - # Clear the cache if it is used: - if hasattr(self, "_cache"): - self._cache.clear() + _options: set[str] = set() - return out + def __init_subclass__(cls, *args: Any, **kwargs: Any) -> None: + """Initialise a subclass of :class:`BaseSolver`.""" - def flag_convergence(self, converged): - """Preset logging for convergence message.""" + def wrap_init(init: Any) -> Any: + """Wrapper to call __post_init__ after __init__.""" - if converged: - self.log.info("Successfully converged.") - else: - self.log.info("Failed to converge.") + def wrapped_init(self: BaseSolver, *args: Any, **kwargs: Any) -> None: + init(self, *args, **kwargs) + if init.__name__ == "__init__": + self.__log_init__() + self.__post_init__() - def get_auxiliaries(self, *args, **kwargs): - """ - Return the auxiliary energies and couplings. - """ + return wrapped_init - raise NotImplementedError + def wrap_kernel(kernel: Any) -> Any: + """Wrapper to call __post_kernel__ after kernel.""" - def get_dyson_orbitals(self, *args, **kwargs): - """ - Return the Dyson orbitals and their energies. - """ + def wrapped_kernel(self: BaseSolver, *args: Any, **kwargs: Any) -> Any: + result = kernel(self, *args, **kwargs) + if kernel.__name__ == "kernel": + self.__post_kernel__() + return result - eigvals, eigvecs = self.get_eigenfunctions(*args, **kwargs) + return wrapped_kernel - if self.hermitian: - eigvecs = eigvecs[: self.nphys] - elif isinstance(eigvecs, tuple): - eigvecs = (eigvecs[0][: self.nphys], eigvecs[1][: self.nphys]) - else: - eigvecs = (eigvecs[: self.nphys], np.linalg.inv(eigvecs).T.conj()[: self.nphys]) + cls.__init__ = wrap_init(cls.__init__) # type: ignore[method-assign] + cls.kernel = wrap_kernel(cls.kernel) # type: ignore[method-assign] - return eigvals, eigvecs + def __log_init__(self) -> None: + """Hook called after :meth:`__init__` for logging purposes.""" + printing.init_console() + console.print("") - def get_eigenfunctions(self, *args, **kwargs): - """ - Return the eigenvalues and eigenfunctions. - """ + # Print the solver name + console.print(f"[method]{self.__class__.__name__}[/method]") + + # Print the options table + table = Table(box=box.SIMPLE) + table.add_column("Option") + table.add_column("Value", style="input") + for key in sorted(self._options): + if not hasattr(self, key): + raise ValueError(f"Option {key} not set in {self.__class__.__name__}") + value = getattr(self, key) + if hasattr(value, "__name__"): + name = value.__name__ + else: + name = str(value) + table.add_row(key, name) + console.print(table) - return self.eigvals, self.eigvecs + def __post_init__(self) -> None: + """Hook called after :meth:`__init__`.""" + pass - def get_self_energy(self, *args, chempot=0.0, **kwargs): + def __post_kernel__(self) -> None: + """Hook called after :meth:`kernel`.""" + pass + + def set_options(self, **kwargs: Any) -> None: + """Set options for the solver. + + Args: + kwargs: Keyword arguments to set as options. """ - Get the self-energy in the format of `pyscf.agf2`. + for key, val in kwargs.items(): + if key not in self._options: + raise ValueError(f"Unknown option for {self.__class__.__name__}: {key}") + if isinstance(getattr(self, key), RepresentationEnum): + # Casts string to the appropriate enum type if the default value is an enum + val = getattr(self, key).__class__(val) + setattr(self, key, val) + + @abstractmethod + def kernel(self) -> Any: + """Run the solver.""" + pass + + @classmethod + @abstractmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> BaseSolver: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + + Notes: + This method will extract the appropriate quantities or functions from the self-energy + to instantiate the solver. In some cases, additional keyword arguments may required. """ + pass + + @classmethod + @abstractmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> BaseSolver: + """Create a solver from an expression. - return Lehmann(*self.get_auxiliaries(*args, **kwargs), chempot=chempot) + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. - def get_greens_function(self, *args, chempot=0.0, **kwargs): + Returns: + Solver instance. + + Notes: + This method will extract the appropriate quantities or functions from the expression + to instantiate the solver. In some cases, additional keyword arguments may required. """ - Get the Green's function in the format of `pyscf.agf2`. + pass + + @property + @abstractmethod + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + pass + + +class StaticSolver(BaseSolver): + """Base class for static Dyson equation solvers.""" + + _options: set[str] = set() + + result: Spectral | None = None + + @abstractmethod + def kernel(self) -> Spectral: + """Run the solver. + + Returns: + The eigenvalues and eigenvectors of the self-energy supermatrix. """ + pass + - return Lehmann(*self.get_dyson_orbitals(*args, **kwargs), chempot=chempot) +class DynamicSolver(BaseSolver): + """Base class for dynamic Dyson equation solvers.""" + + @abstractmethod + def kernel(self) -> Dynamic[Any]: + """Run the solver. + + Returns: + Dynamic Green's function resulting from the Dyson equation. + """ + pass diff --git a/dyson/solvers/static/__init__.py b/dyson/solvers/static/__init__.py new file mode 100644 index 0000000..26713a5 --- /dev/null +++ b/dyson/solvers/static/__init__.py @@ -0,0 +1,17 @@ +r"""Solvers for solving the Dyson equation statically. + +Submodules +---------- + +.. autosummary:: + :toctree: + + exact + davidson + downfolded + mblse + mblgf + chempot + density + +""" diff --git a/dyson/solvers/static/_mbl.py b/dyson/solvers/static/_mbl.py new file mode 100644 index 0000000..a675c1a --- /dev/null +++ b/dyson/solvers/static/_mbl.py @@ -0,0 +1,257 @@ +"""Common functionality for moment block Lanczos solvers.""" + +from __future__ import annotations + +import functools +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from dyson import console, printing, util +from dyson import numpy as np +from dyson.solvers.solver import StaticSolver + +if TYPE_CHECKING: + from dyson.representations.spectral import Spectral + from dyson.typing import Array + +# TODO: reimplement caching + + +class BaseRecursionCoefficients(ABC): + """Base class for recursion coefficients for the moment block Lanczos algorithms. + + Args: + nphys: Number of physical degrees of freedom. + """ + + NDIM: int = 2 + + def __init__( + self, + nphys: int, + hermitian: bool = True, + force_orthogonality: bool = True, + ): + """Initialise the recursion coefficients.""" + self._nphys = nphys + self._zero = np.zeros((nphys,) * self.NDIM) + self._data: dict[tuple[int, ...], Array] = {} + self.hermitian = hermitian + self.force_orthogonality = force_orthogonality + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return self._nphys + + @property + def dtype(self) -> str: + """Get the data type of the recursion coefficients.""" + if any([np.iscomplexobj(v) for v in self._data.values()]): + return "complex128" + return "float64" + + @abstractmethod + def __getitem__(self, key: tuple[int, ...]) -> Array: + """Get the recursion coefficients for the given key. + + Args: + key: The key for the recursion coefficients. + + Returns: + The recursion coefficients. + """ + pass + + @abstractmethod + def __setitem__(self, key: tuple[int, ...], value: Array) -> None: + """Set the recursion coefficients for the given key. + + Args: + key: The key for the recursion coefficients. + value: The recursion coefficients. + """ + pass + + +class BaseMBL(StaticSolver): + """Base class for moment block Lanczos solvers.""" + + Coefficients: type[BaseRecursionCoefficients] + + _moments: Array + + max_cycle: int + hermitian: bool = True + force_orthogonality: bool = True + calculate_errors: bool = True + _options: set[str] = {"max_cycle", "hermitian", "force_orthogonality", "calculate_errors"} + + def __post_kernel__(self) -> None: + """Hook called after :meth:`kernel`.""" + assert self.result is not None + emin = printing.format_float(self.result.eigvals.min()) + emax = printing.format_float(self.result.eigvals.max()) + console.print( + f"Found [output]{self.result.neig}[/output] roots between [output]{emin}[/output] and " + f"[output]{emax}[/output]." + ) + if self.calculate_errors: + error = printing.format_float( + self.moment_error(iteration=self.max_cycle), + threshold=1e-10, + scientific=True, + precision=4, + ) + console.print(f"Error in the moments: {error}") + + @abstractmethod + def solve(self, iteration: int | None = None) -> Spectral: + """Solve the eigenvalue problem at a given iteration. + + Args: + iteration: The iteration to get the results for. + + Returns: + The :cls:`Spectral` object. + """ + pass + + def kernel(self) -> Spectral: + """Run the solver. + + Returns: + The eigenvalues and eigenvectors of the self-energy supermatrix. + """ + # Get the table + table = printing.ConvergencePrinter( + (), ("Error in moments", "Error in sqrt", "Error in inv. sqrt"), (1e-10, 1e-10, 1e-10) + ) + progress = printing.IterationsPrinter(self.max_cycle) + progress.start() + + # Run the solver + for iteration in range(self.max_cycle + 1): # TODO: check + error_sqrt, error_inv_sqrt, error_moments = self.recurrence_iteration(iteration) + if not self.calculate_errors: + error_sqrt = error_inv_sqrt = error_moments = np.nan + table.add_row(iteration, (), (error_moments, error_sqrt, error_inv_sqrt)) + progress.update(iteration) + + progress.stop() + table.print() + + # Diagonalise the compressed self-energy + self.result = self.solve(iteration=self.max_cycle) + + return self.result + + @functools.cached_property + def orthogonalisation_metric(self) -> Array: + """Get the orthogonalisation metric.""" + return util.matrix_power(self.moments[0], -0.5, hermitian=self.hermitian)[0] + + @functools.cached_property + def orthogonalisation_metric_inv(self) -> Array: + """Get the inverse of the orthogonalisation metric.""" + return util.matrix_power(self.moments[0], 0.5, hermitian=self.hermitian)[0] + + @functools.lru_cache(maxsize=64) + def orthogonalised_moment(self, order: int) -> Array: + """Compute an orthogonalised moment. + + Args: + order: The order of the moment. + + Returns: + The orthogonalised moment. + """ + return self.orthogonalisation_metric @ self.moments[order] @ self.orthogonalisation_metric + + @abstractmethod + def reconstruct_moments(self, iteration: int) -> Array: + """Reconstruct the moments. + + Args: + iteration: The iteration number. + + Returns: + The reconstructed moments. + """ + pass + + def moment_error(self, iteration: int | None = None): + """Get the moment error at a given iteration. + + Args: + iteration: The iteration to check. + """ + if iteration is None: + iteration = self.max_cycle + + # Construct the recovered moments + moments = self.reconstruct_moments(iteration) + + # Get the error + error = sum( + util.scaled_error(predicted, actual) + for predicted, actual in zip(moments, self.moments[: 2 * iteration + 2]) + ) + + return error + + @abstractmethod + def initialise_recurrence(self) -> tuple[float | None, float | None, float | None]: + """Initialise the recurrence (zeroth iteration). + + Returns: + If :attr:`calculate_errors`, the error metrics in the square root of the off-diagonal + block, the inverse square root of the off-diagonal block, and the error in the + recovered moments. If not, all three are ``None``. + """ + pass + + @abstractmethod + def _recurrence_iteration_hermitian( + self, iteration: int + ) -> tuple[float | None, float | None, float | None]: + """Perform an iteration of the recurrence for a Hermitian self-energy.""" + pass + + @abstractmethod + def _recurrence_iteration_non_hermitian( + self, iteration: int + ) -> tuple[float | None, float | None, float | None]: + """Perform an iteration of the recurrence for a non-Hermitian self-energy.""" + pass + + def recurrence_iteration( + self, iteration: int + ) -> tuple[float | None, float | None, float | None]: + """Perform an iteration of the recurrence. + + Args: + iteration: The iteration to perform. + + Returns: + If :attr:`calculate_errors`, the error metrics in the square root of the off-diagonal + block, the inverse square root of the off-diagonal block, and the error in the + recovered moments. If not, all three are ``None``. + """ + if iteration == 0: + return self.initialise_recurrence() + if iteration > self.max_cycle: + raise ValueError(f"Iteration {iteration} exceeds max_cycle {self.max_cycle}.") + if self.hermitian: + return self._recurrence_iteration_hermitian(iteration) + return self._recurrence_iteration_non_hermitian(iteration) + + @property + def moments(self) -> Array: + """Get the moments of the self-energy.""" + return self._moments + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return self.moments.shape[-1] diff --git a/dyson/solvers/static/chempot.py b/dyson/solvers/static/chempot.py new file mode 100644 index 0000000..52a8694 --- /dev/null +++ b/dyson/solvers/static/chempot.py @@ -0,0 +1,636 @@ +"""Chemical potential optimising solvers.""" + +from __future__ import annotations + +import functools +import warnings +from typing import TYPE_CHECKING + +import scipy.optimize + +from dyson import console, printing, util +from dyson import numpy as np +from dyson.representations.lehmann import Lehmann, shift_energies +from dyson.solvers.solver import StaticSolver +from dyson.solvers.static.exact import Exact + +if TYPE_CHECKING: + from typing import Any, Literal + + from dyson.expressions.expression import BaseExpression + from dyson.representations.spectral import Spectral + from dyson.typing import Array + + +def _warn_or_raise_if_negative_weight( + weight: float | Array, hermitian: bool, tol: float = 1e-6 +) -> None: + """Warn or raise an error for negative weights. + + Args: + weight: Weight to check. + hermitian: Whether the Green's function is hermitian. + tol: Tolerance for the weight to be considered negative. + + Raises: + ValueError: If the weight is negative and the Green's function is hermitian. + UserWarning: If the weight is negative and the Green's function is not hermitian. + """ + if not isinstance(weight, np.ndarray): + weight = np.array(weight) + if np.any(weight < -tol): + if hermitian: + raise ValueError( + f"Negative number of electrons in state: {weight:.6f}. This should be " + "impossible for a hermitian Green's function." + ) + else: + warnings.warn( + f"Negative number of electrons in state: {weight:.6f}. This is possible for " + "a non-hermitian Green's function, but may be problematic for finding the " + "chemical potential. Consider using the global method.", + UserWarning, + ) + + +def search_aufbau_global( + greens_function: Lehmann, nelec: int, occupancy: float = 2.0 +) -> tuple[float, float]: + """Search for a chemical potential in a Green's function using a global minimisation. + + Args: + greens_function: Green's function. + nelec: Target number of electrons. + occupancy: Occupancy of each state, typically 2 for a restricted reference and 1 + otherwise. + + Returns: + The chemical potential and the error in the number of electrons. + """ + energies = greens_function.energies + weights = greens_function.weights(occupancy=occupancy) + cumweights = np.cumsum(weights) + + # Find the global minimum + i = np.argmin(np.abs(cumweights - nelec)) + error = cumweights[i] - nelec + homo = i + lumo = i + 1 + + # Find the chemical potential + if homo == -1: + chempot = energies[lumo].real - 1e-4 + elif lumo == energies.size: + chempot = energies[homo].real + 1e-4 + else: + chempot = 0.5 * (energies[homo] + energies[lumo]).real + + return chempot, error + + +def search_aufbau_direct( + greens_function: Lehmann, nelec: int, occupancy: float = 2.0 +) -> tuple[float, float]: + """Search for a chemical potential in a Green's function using the Aufbau principle. + + Args: + greens_function: Green's function. + nelec: Target number of electrons. + occupancy: Occupancy of each state, typically 2 for a restricted reference and 1 + otherwise. + + Returns: + The chemical potential and the error in the number of electrons. + """ + energies = greens_function.energies + left, right = greens_function.unpack_couplings() + + # Find the two states bounding the chemical potential + sum_i = sum_j = 0.0 + for i in range(greens_function.naux): + number = np.vdot(left[:, i], right[:, i]).real * occupancy + _warn_or_raise_if_negative_weight(number, greens_function.hermitian) + sum_i, sum_j = sum_j, sum_j + number + if sum_i < nelec <= sum_j: + break + + # Find the best HOMO + if abs(sum_i - nelec) < abs(sum_j - nelec): + homo = i - 1 + error = nelec - sum_i + else: + homo = i + error = nelec - sum_j + lumo = homo + 1 + + # Find the chemical potential + if homo == -1: + chempot = energies[lumo].real - 1e-4 + elif lumo == energies.size: + chempot = energies[homo].real + 1e-4 + else: + chempot = 0.5 * (energies[homo] + energies[lumo]).real + + return chempot, error + + +def search_aufbau_bisect( + greens_function: Lehmann, nelec: int, occupancy: float = 2.0, max_cycle: int = 1000 +) -> tuple[float, float]: + """Search for a chemical potential in a Green's function using Aufbau principle and bisection. + + Args: + greens_function: Green's function. + nelec: Target number of electrons. + occupancy: Occupancy of each state, typically 2 for a restricted reference and 1 + otherwise. + max_cycle: Maximum number of iterations. + + Returns: + The chemical potential and the error in the number of electrons. + """ + energies = greens_function.energies + weights = greens_function.weights(occupancy=occupancy) + cumweights = np.cumsum(weights) + + # Find the two states bounding the chemical potential + low, mid, high = 0, greens_function.naux // 2, greens_function.naux + for cycle in range(1, max_cycle + 1): + number = cumweights[mid] + if number < nelec: + low = mid + mid = mid + (high - low) // 2 + else: + high = mid + mid = mid - (high - low) // 2 + if mid in {low, high}: + break + else: + raise ValueError("Failed to converge bisection") + sum_i = cumweights[low] + sum_j = cumweights[high] + + # Find the best HOMO + if abs(sum_i - nelec) < abs(sum_j - nelec): + homo = low + error = nelec - sum_i + else: + homo = high + error = nelec - sum_j + lumo = homo + 1 + + # Find the chemical potential + if homo == -1: + chempot = energies[lumo].real - 1e-4 + elif lumo == energies.size: + chempot = energies[homo].real + 1e-4 + else: + chempot = 0.5 * (energies[homo] + energies[lumo]).real + + return chempot, error + + +class ChemicalPotentialSolver(StaticSolver): + """Base class for a solver for a self-energy that optimises the chemical potential. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + nelec: Target number of electrons. + """ + + _static: Array + _self_energy: Lehmann + _nelec: int + _overlap: Array | None + + error: float | None = None + chempot: float | None = None + converged: bool | None = None + + def __post_init__(self) -> None: + """Hook called after :meth:`__init__`.""" + # Check the input + if self.static.ndim != 2 or self.static.shape[0] != self.static.shape[1]: + raise ValueError("static must be a square matrix.") + if self.self_energy.nphys != self.static.shape[0]: + raise ValueError( + "self_energy must have the same number of physical degrees of freedom as static." + ) + if self.overlap is not None and ( + self.overlap.ndim != 2 or self.overlap.shape[0] != self.overlap.shape[1] + ): + raise ValueError("overlap must be a square matrix or None.") + if self.overlap is not None and self.overlap.shape != self.static.shape: + raise ValueError("overlap must have the same shape as static.") + + # Print the input information + console.print(f"Number of physical states: [input]{self.nphys}[/input]") + console.print(f"Number of auxiliary states: [input]{self.self_energy.naux}[/input]") + console.print(f"Target number of electrons: [input]{self.nelec}[/input]") + if self.overlap is not None: + cond = printing.format_float( + np.linalg.cond(self.overlap), threshold=1e10, scientific=True, precision=4 + ) + console.print(f"Overlap condition number: {cond}") + + @property + def static(self) -> Array: + """Get the static part of the self-energy.""" + return self._static + + @property + def self_energy(self) -> Lehmann: + """Get the self-energy.""" + return self._self_energy + + @property + def nelec(self) -> int: + """Get the target number of electrons.""" + return self._nelec + + @property + def overlap(self) -> Array | None: + """Get the overlap matrix, if available.""" + return self._overlap + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return self._self_energy.nphys + + +class AufbauPrinciple(ChemicalPotentialSolver): + """Solve a self-energy and assign a chemical potential based on the Aufbau principle. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + nelec: Target number of electrons. + """ + + occupancy: float = 2.0 + solver: type[Exact] = Exact + method: Literal["direct", "bisect", "global"] = "global" + _options: set[str] = {"occupancy", "solver", "method"} + + def __init__( # noqa: D417 + self, + static: Array, + self_energy: Lehmann, + nelec: int, + overlap: Array | None = None, + **kwargs: Any, + ): + """Initialise the solver. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + nelec: Target number of electrons. + overlap: Overlap matrix for the physical space. + occupancy: Occupancy of each state, typically 2 for a restricted reference and 1 + otherwise. + solver: Solver to use for the self-energy. + method: Method to use for the chemical potential search. + """ + self._static = static + self._self_energy = self_energy + self._nelec = nelec + self._overlap = overlap + self.set_options(**kwargs) + + def __post_kernel__(self) -> None: + """Hook called after :meth:`kernel`.""" + assert self.result is not None + assert self.chempot is not None + assert self.error is not None + emin = printing.format_float(self.result.eigvals.min()) + emax = printing.format_float(self.result.eigvals.max()) + console.print("") + console.print( + f"Found [output]{self.result.neig}[/output] roots between [output]{emin}[/output] and " + f"[output]{emax}[/output]." + ) + cpt = printing.format_float(self.chempot) + err = printing.format_float(self.error, threshold=1e-3, precision=4, scientific=True) + console.print(f"Chemical potential: [output]{cpt}[/output]") + console.print(f"Error in number of electrons: [output]{err}[/output]") + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> AufbauPrinciple: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + + Notes: + To initialise this solver from a self-energy, the `nelec` keyword argument must be + provided. + """ + if "nelec" not in kwargs: + raise ValueError("Missing required argument nelec.") + kwargs = kwargs.copy() + nelec = kwargs.pop("nelec") + return cls(static, self_energy, nelec, overlap=overlap, **kwargs) + + @classmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> AufbauPrinciple: + """Create a solver from an expression. + + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + raise NotImplementedError( + "Cannot instantiate AufbauPrinciple from expression, use from_self_energy instead." + ) + + def kernel(self) -> Spectral: + """Run the solver. + + Returns: + The eigenvalues and eigenvectors of the self-energy supermatrix. + """ + # Solve the self-energy + with printing.quiet: + solver = self.solver.from_self_energy( + self.static, self.self_energy, overlap=self.overlap + ) + result = solver.kernel() + greens_function = result.get_greens_function() + + # Get the chemical potential and error + if self.method == "direct": + chempot, error = search_aufbau_direct(greens_function, self.nelec, self.occupancy) + elif self.method == "bisect": + chempot, error = search_aufbau_bisect(greens_function, self.nelec, self.occupancy) + elif self.method == "global": + chempot, error = search_aufbau_global(greens_function, self.nelec, self.occupancy) + else: + raise ValueError(f"Unknown method: {self.method}") + result.chempot = chempot + + # Set the results + self.result = result + self.chempot = chempot + self.error = error + self.converged = True + + return result + + +class AuxiliaryShift(ChemicalPotentialSolver): + """Shift the self-energy auxiliaries to best assign a chemical potential. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + nelec: Target number of electrons. + + Notes: + Convergence is met when either of the thresholds `conv_tol` or `conv_tol_grad` are met, + rather than both, due to constraints of the :meth:`scipy.optimize.minimize` method. + """ + + occupancy: float = 2.0 + solver: type[AufbauPrinciple] = AufbauPrinciple + max_cycle: int = 200 + conv_tol: float = 1e-8 + conv_tol_grad: float = 0.0 + guess: float = 0.0 + _options: set[str] = {"occupancy", "solver", "max_cycle", "conv_tol", "conv_tol_grad", "guess"} + + shift: float | None = None + + def __init__( # noqa: D417 + self, + static: Array, + self_energy: Lehmann, + nelec: int, + overlap: Array | None = None, + **kwargs: Any, + ): + """Initialise the solver. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + nelec: Target number of electrons. + overlap: Overlap matrix for the physical space. + occupancy: Occupancy of each state, typically 2 for a restricted reference and 1 + otherwise. + solver: Solver to use for the self-energy and chemical potential search. + max_cycle: Maximum number of iterations. + conv_tol: Convergence tolerance for the number of electrons. + conv_tol_grad: Convergence tolerance for the gradient of the objective function. + guess: Initial guess for the chemical potential. + """ + self._static = static + self._self_energy = self_energy + self._nelec = nelec + self._overlap = overlap + self.set_options(**kwargs) + + def __post_kernel__(self) -> None: + """Hook called after :meth:`kernel`.""" + assert self.result is not None + assert self.chempot is not None + assert self.error is not None + assert self.shift is not None + emin = printing.format_float(self.result.eigvals.min()) + emax = printing.format_float(self.result.eigvals.max()) + console.print( + f"Found [output]{self.result.neig}[/output] roots between [output]{emin}[/output] and " + f"[output]{emax}[/output]." + ) + cpt = printing.format_float(self.chempot) + err = printing.format_float(self.error, threshold=1e-3, precision=4, scientific=True) + shift = printing.format_float(self.shift, precision=4, scientific=True) + console.print(f"Chemical potential: [output]{cpt}[/output]") + console.print(f"Auxiliary shift: [output]{shift}[/output]") + console.print(f"Error in number of electrons: [output]{err}[/output]") + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> AuxiliaryShift: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + + Notes: + To initialise this solver from a self-energy, the `nelec` keyword argument must be + provided. + """ + if "nelec" not in kwargs: + raise ValueError("Missing required argument nelec.") + nelec = kwargs.pop("nelec") + return cls(static, self_energy, nelec, overlap=overlap, **kwargs) + + @classmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> AuxiliaryShift: + """Create a solver from an expression. + + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + raise NotImplementedError( + "Cannot instantiate AuxiliaryShift from expression, use from_self_energy instead." + ) + + def objective(self, shift: float) -> float: + """Objective function for the chemical potential search. + + Args: + shift: Shift to apply to the self-energy. + + Returns: + The error in the number of electrons. + """ + with printing.quiet: + with shift_energies(self.self_energy, np.ravel(shift)[0]): + solver = self.solver.from_self_energy( + self.static, self.self_energy, nelec=self.nelec, overlap=self.overlap + ) + solver.kernel() + assert solver.error is not None + return solver.error**2 + + @functools.lru_cache(maxsize=16) + def gradient(self, shift: float) -> tuple[float, Array]: + """Gradient of the objective function. + + Args: + shift: Shift to apply to the self-energy. + + Returns: + The error in the number of electrons, and the gradient of the error. + """ + with printing.quiet: + with shift_energies(self.self_energy, np.ravel(shift)[0]): + solver = self.solver.from_self_energy( + self.static, self.self_energy, nelec=self.nelec, overlap=self.overlap + ) + solver.kernel() + assert solver.error is not None + assert solver.result is not None + eigvals = solver.result.eigvals + left, right = util.unpack_vectors(solver.result.eigvecs) + nphys = self.nphys + nocc = np.count_nonzero(eigvals < solver.chempot) + + h1 = -left[nphys:, nocc:].T.conj() @ right[nphys:, :nocc] + z = h1 / (eigvals[nocc:, None] - eigvals[None, :nocc]) + pert_coeff_occ_left = left[:nphys, nocc:] @ z + pert_coeff_occ_right = right[:nphys, nocc:] @ z + pert_rdm1 = pert_coeff_occ_left @ pert_coeff_occ_right.T.conj() * 4.0 # occupancy? + grad = np.trace(pert_rdm1).real * solver.error * self.occupancy + + return solver.error**2, grad + + def _minimize(self) -> scipy.optimize.OptimizeResult: + """Minimise the objective function. + + Returns: + The :class:`scipy.optimize.OptimizeResult` object from the minimizer. + """ + # Get the table and callback function + table = printing.ConvergencePrinter( + ("Shift",), ("Error", "Gradient"), (self.conv_tol, self.conv_tol_grad) + ) + progress = printing.IterationsPrinter(self.max_cycle) + progress.start() + cycle = 1 + + def _callback(xk: Array) -> None: + """Callback function for the minimizer.""" + nonlocal cycle + error, grad = self.gradient(np.ravel(xk)[0]) + error = np.sqrt(error) + table.add_row(cycle, (np.ravel(xk)[0],), (error, np.ravel(grad)[0])) + progress.update(cycle) + cycle += 1 + + with util.catch_warnings(np.exceptions.ComplexWarning): + opt = scipy.optimize.minimize( + lambda x: self.gradient(np.ravel(x)[0]), + x0=self.guess, + method="TNC", + jac=True, + options=dict( + maxfun=self.max_cycle, + xtol=0.0, + ftol=self.conv_tol**2, + gtol=self.conv_tol_grad, + ), + callback=_callback, + ) + + progress.stop() + table.print() + + return opt + + def kernel(self) -> Spectral: + """Run the solver. + + Returns: + The eigenvalues and eigenvectors of the self-energy supermatrix. + """ + # Minimize the objective function + opt = self._minimize() + + # Get the shifted self-energy + self_energy = Lehmann( + self.self_energy.energies + opt.x, + self.self_energy.couplings, + chempot=self.self_energy.chempot, + sort=False, + ) + + # Solve the self-energy + with printing.quiet: + solver = self.solver.from_self_energy( + self.static, self_energy, nelec=self.nelec, overlap=self.overlap + ) + result = solver.kernel() + + # Set the results + self.result = result + self.chempot = solver.chempot + self.error = solver.error + self.converged = opt.success + self.shift = np.ravel(opt.x)[0] + + return result diff --git a/dyson/solvers/static/davidson.py b/dyson/solvers/static/davidson.py new file mode 100644 index 0000000..929a746 --- /dev/null +++ b/dyson/solvers/static/davidson.py @@ -0,0 +1,335 @@ +"""Davidson algorithm [1]_ [2]_. + +.. [1] Davidson, E. R. (1975). The iterative calculation of a few of the lowest + eigenvalues and corresponding eigenvectors of large real-symmetric matrices. Journal of + Computational Physics, 17(1), 87–94. https://doi.org/10.1016/0021-9991(75)90065-0 + +.. [2] Morgan, R. B. (1990). Davidson’s method and preconditioning for generalized + eigenvalue problems. Journal of Computational Physics, 89(1), 241–245. + https://doi.org/10.1016/0021-9991(90)90124-j +""" + +from __future__ import annotations + +import warnings +from typing import TYPE_CHECKING + +from pyscf import lib + +from dyson import console, printing, util +from dyson import numpy as np +from dyson.representations.lehmann import Lehmann +from dyson.representations.spectral import Spectral +from dyson.solvers.solver import StaticSolver +from dyson.solvers.static.exact import orthogonalise_self_energy, project_eigenvectors + +if TYPE_CHECKING: + from typing import Any, Callable + + from dyson.expressions.expression import BaseExpression + from dyson.typing import Array + + +def _pick_real_eigenvalues( + eigvals: Array, + eigvecs: Array, + nroots: int, + env: dict[str, Any], + threshold: float = 1e-3, +) -> tuple[Array, Array, Array]: + """Pick real eigenvalues.""" + iabs = np.abs(eigvals.imag) + tol = max(threshold, np.sort(iabs)[min(eigvals.size, nroots) - 1]) + real_idx = np.where(iabs <= tol)[0] + + # Check we have enough real eigenvalues + num = np.count_nonzero(iabs[real_idx] < threshold) + if num < nroots and eigvals.size >= nroots: + warnings.warn( + f"Only {num} of requested {nroots} real eigenvalues found with threshold {tol:.2e}.", + RuntimeWarning, + stacklevel=2, + ) + + # Sort the eigenvalues + idx = real_idx[np.argsort(np.abs(eigvals[real_idx]))] + eigvals = eigvals[idx] + eigvecs = eigvecs[:, idx] + + # Make the eigenvalues real + real_system = issubclass(env.get("dtype", np.float64).type, (complex, np.complexfloating)) + if real_system: + degen_idx = np.where(eigvals.imag != 0)[0] + if degen_idx.size > 0: + eigvecs[:, degen_idx[1::2]] = eigvecs[:, degen_idx[1::2]].imag + eigvecs = eigvecs.real + + return eigvals, eigvecs, idx + + +class Davidson(StaticSolver): + """Davidson algorithm for diagonalisation of the supermatrix form of the self-energy. + + Args: + matvec: The matrix-vector operation for the self-energy supermatrix. + diagonal: The diagonal of the self-energy supermatrix. + bra: The bra excitation vector mapping the supermatrix to the physical space. + ket: The ket excitation vector mapping the supermatrix to the physical space. + """ + + hermitian: bool = True + nroots: int = 1 + max_cycle: int = 100 + max_space: int = 16 + conv_tol: float = 1e-8 + conv_tol_residual: float = 1e-5 + _options: set[str] = { + "hermitian", + "nroots", + "max_cycle", + "max_space", + "conv_tol", + "conv_tol_residual", + } + + converged: Array | None = None + + def __init__( # noqa: D417 + self, + matvec: Callable[[Array], Array], + diagonal: Array, + bra: Array, + ket: Array | None = None, + **kwargs: Any, + ): + """Initialise the solver. + + Args: + matvec: The matrix-vector operation for the self-energy supermatrix. + diagonal: The diagonal of the self-energy supermatrix. + bra: The bra excitation vector mapping the supermatrix to the physical space. + ket: The ket excitation vector mapping the supermatrix to the physical space. If `None`, + use the same vectors as ``bra``. + hermitian: Whether the matrix is hermitian. + nroots: Number of roots to find. + max_cycle: Maximum number of iterations. + max_space: Maximum size of the subspace. + conv_tol: Convergence tolerance for the eigenvalues. + conv_tol_residual: Convergence tolerance for the residual. + """ + self._matvec = matvec + self._diagonal = diagonal + self._bra = bra + self._ket = ket if ket is not None else bra + self.set_options(**kwargs) + + def __post_init__(self) -> None: + """Hook called after :meth:`__init__`.""" + # Check the input + if self.diagonal.ndim != 1: + raise ValueError("diagonal must be a 1D array.") + if self.bra.ndim != 2 or self.bra.shape[1] != self.diagonal.size: + raise ValueError("bra must be a 2D array with the same number of columns as diagonal.") + if self.ket is not None and (self.ket.ndim != 2 or self.ket.shape[1] != self.diagonal.size): + raise ValueError("ket must be a 2D array with the same number of columns as diagonal.") + if self.ket is not None and self.ket.shape[0] != self.bra.shape[0]: + raise ValueError("ket must have the same number of rows as bra.") + if not callable(self.matvec): + raise ValueError("matvec must be a callable function.") + + # Print the input information + console.print(f"Matrix shape: [input]{(self.diagonal.size, self.diagonal.size)}[/input]") + console.print(f"Number of physical states: [input]{self.nphys}[/input]") + + def __post_kernel__(self) -> None: + """Hook called after :meth:`kernel`.""" + assert self.result is not None + assert self.converged is not None + emin = printing.format_float(self.result.eigvals.min()) + emax = printing.format_float(self.result.eigvals.max()) + console.print( + f"Found [output]{self.result.neig}[/output] roots between [output]{emin}[/output] and " + f"[output]{emax}[/output]." + ) + rating = "good" if np.all(self.converged) else "okay" if np.any(self.converged) else "bad" + console.print( + f"Converged [{rating}]{np.sum(self.converged)} of {self.nroots}[/{rating}] roots." + ) + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> Davidson: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + static, self_energy, bra, ket = orthogonalise_self_energy( + static, self_energy, overlap=overlap + ) + return cls( + lambda vector: self_energy.matvec(static, vector), + self_energy.diagonal(static), + bra, + ket, + hermitian=self_energy.hermitian, + **kwargs, + ) + + @classmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> Davidson: + """Create a solver from an expression. + + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + diagonal = expression.diagonal() + matvec = expression.apply_hamiltonian + bra = np.array(expression.get_excitation_bras()) + ket = ( + np.array(expression.get_excitation_kets()) + if not expression.hermitian_upfolded + else None + ) + return cls( + matvec, + diagonal, + bra, + ket, + hermitian=expression.hermitian_upfolded, + **kwargs, + ) + + def get_guesses(self) -> list[Array]: + """Get the initial guesses for the eigenvectors. + + Returns: + Initial guesses for the eigenvectors. + """ + args = np.argsort(np.abs(self.diagonal)).astype(int) + dtype = " Spectral: + """Run the solver. + + Returns: + The eigenvalues and eigenvectors of the self-energy supermatrix. + """ + # Get the table and callback function + table = printing.ConvergencePrinter( + ("Smallest root",), ("Change", "Residual"), (self.conv_tol, self.conv_tol_residual) + ) + progress = printing.IterationsPrinter(self.max_cycle) + progress.start() + + def _callback(env: dict[str, Any]) -> None: + """Callback function for the Davidson algorithm.""" + root = env["e"][np.argmin(np.abs(env["e"]))] + table.add_row( + env["icyc"] + 1, (root,), (np.max(np.abs(env["de"])), np.max(env["dx_norm"])) + ) + progress.update(env["icyc"] + 1) + del env + + # Call the Davidson function + if self.hermitian: + converged, eigvals, eigvecs = lib.linalg_helper.davidson1( + lambda vectors: [self.matvec(vector) for vector in vectors], + self.get_guesses(), + self.diagonal, + pick=_pick_real_eigenvalues, + tol=self.conv_tol, + tol_residual=self.conv_tol_residual, + max_cycle=self.max_cycle, + max_space=self.max_space, + nroots=self.nroots, + callback=_callback, + verbose=0, + ) + + eigvals = np.array(eigvals) + eigvecs = np.array(eigvecs).T + + else: + with util.catch_warnings(UserWarning): + converged, eigvals, left, right = lib.linalg_helper.davidson_nosym1( + lambda vectors: [self.matvec(vector) for vector in vectors], + self.get_guesses(), + self.diagonal, + pick=_pick_real_eigenvalues, + tol=self.conv_tol, + tol_residual=self.conv_tol_residual, + max_cycle=self.max_cycle, + max_space=self.max_space, + nroots=self.nroots, + left=True, + callback=_callback, + verbose=0, + ) + + eigvals = np.array(eigvals) + left = np.array(left).T + right = np.array(right).T + eigvecs = np.array([left, right]) + + # TODO: How to print the final iteration? + progress.stop() + table.print() + + # Sort the eigenvalues + mask = np.argsort(eigvals) + eigvals = eigvals[mask] + eigvecs = eigvecs[..., mask] + converged = converged[mask] + + # Get the full map onto physical + auxiliary and rotate the eigenvectors + eigvecs = project_eigenvectors(eigvecs, self.bra, self.ket if not self.hermitian else None) + + # Store the results + self.result = Spectral(eigvals, eigvecs, self.nphys) + self.converged = converged + + return self.result + + @property + def matvec(self) -> Callable[[Array], Array]: + """Get the matrix-vector operation for the self-energy supermatrix.""" + return self._matvec + + @property + def diagonal(self) -> Array: + """Get the diagonal of the self-energy supermatrix.""" + return self._diagonal + + @property + def bra(self) -> Array: + """Get the bra excitation vector mapping the supermatrix to the physical space.""" + return self._bra + + @property + def ket(self) -> Array: + """Get the ket excitation vector mapping the supermatrix to the physical space.""" + if self._ket is None: + return self._bra + return self._ket + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return self.bra.shape[0] diff --git a/dyson/solvers/static/density.py b/dyson/solvers/static/density.py new file mode 100644 index 0000000..ce4e22b --- /dev/null +++ b/dyson/solvers/static/density.py @@ -0,0 +1,379 @@ +"""Density matrix relaxing solver.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from pyscf import lib + +from dyson import console, printing +from dyson import numpy as np +from dyson.representations.lehmann import Lehmann +from dyson.solvers.solver import StaticSolver +from dyson.solvers.static.chempot import AufbauPrinciple, AuxiliaryShift + +if TYPE_CHECKING: + from typing import Any, Protocol + + from pyscf import scf + + from dyson.expressions.expression import BaseExpression + from dyson.representations.spectral import Spectral + from dyson.typing import Array + + class StaticFunction(Protocol): + """Protocol for a function that computes the static self-energy.""" + + def __call__( + self, + rdm1: Array, + rdm1_prev: Array | None = None, + static_prev: Array | None = None, + ) -> Array: + """Compute the static self-energy for a given density matrix. + + Args: + rdm1: Density matrix. + rdm1_prev: Previous density matrix. Used for direct build. + static_prev: Previous Fock matrix. Used for direct build. + + Returns: + Static self-energy. + """ + ... + + +def get_fock_matrix_function(mf: scf.hf.RHF) -> StaticFunction: + """Get a function to compute the Fock matrix for a given density matrix. + + Args: + mf: Mean-field object. + + Returns: + Function to compute the Fock matrix. + """ + h1e = mf.get_hcore() + s1e = mf.get_ovlp() + + def get_fock( + rdm1: Array, rdm1_prev: Array | None = None, static_prev: Array | None = None + ) -> Array: + """Compute the Fock matrix for a given density matrix. + + Args: + rdm1: Density matrix. + rdm1_prev: Previous density matrix. Used for direct build. + static_prev: Previous Fock matrix. Used for direct build. + + Returns: + Fock matrix. + """ + # Transform to AO basis + rdm1 = mf.mo_coeff @ rdm1 @ mf.mo_coeff.T.conj() + if (rdm1_prev is None) != (static_prev is None): + raise ValueError( + "Both rdm1_prev and static_prev must be None or both must be provided." + ) + if rdm1_prev is not None and static_prev is not None: + rdm1_prev = mf.mo_coeff @ rdm1_prev @ mf.mo_coeff.T.conj() + static_prev = mf.mo_coeff @ static_prev @ mf.mo_coeff.T.conj() + + # Compute the new Fock matrix + veff_last = static_prev - h1e if static_prev is not None else None + veff = mf.get_veff(dm=rdm1, dm_last=rdm1_prev, vhf_last=veff_last) + fock = mf.get_fock(h1e=h1e, s1e=s1e, vhf=veff, dm=rdm1) + + # Transform back to MO basis + fock = mf.mo_coeff.T.conj() @ fock @ mf.mo_coeff + + return fock + + return get_fock + + +class DensityRelaxation(StaticSolver): + """Solve a self-energy and relax the density matrix in the presence of the auxiliaries. + + Args: + get_static: Function to get the static self-energy (including Fock contributions) for a + given density matrix. + self_energy: Self-energy. + nelec: Target number of electrons. + """ + + occupancy: float = 2.0 + solver_outer: type[AuxiliaryShift] = AuxiliaryShift + solver_inner: type[AufbauPrinciple] = AufbauPrinciple + diis_min_space: int = 2 + diis_max_space: int = 8 + max_cycle_outer: int = 20 + max_cycle_inner: int = 50 + conv_tol: float = 1e-8 + favour_rdm: bool = True + _options: set[str] = { + "occupancy", + "solver_outer", + "solver_inner", + "diis_min_space", + "diis_max_space", + "max_cycle_outer", + "max_cycle_inner", + "conv_tol", + "favour_rdm", + } + + converged: bool | None = None + + def __init__( # noqa: D417 + self, + get_static: StaticFunction, + self_energy: Lehmann, + nelec: int, + overlap: Array | None = None, + **kwargs: Any, + ): + """Initialise the solver. + + Args: + get_static: Function to get the static self-energy (including Fock contributions) for a + given density matrix. + self_energy: Self-energy. + nelec: Target number of electrons. + overlap: Overlap matrix for the physical space. + occupancy: Occupancy of each state, typically 2 for a restricted reference and 1 + otherwise. + solver_outer: Solver to use for the self-energy and chemical potential search in the + outer loop. + solver_inner: Solver to use for the self-energy and chemical potential search in the + inner loop. + diis_min_space: Minimum size of the DIIS space. + diis_max_space: Maximum size of the DIIS space. + max_cycle_outer: Maximum number of outer iterations. + max_cycle_inner: Maximum number of inner iterations. + conv_tol: Convergence tolerance in the density matrix. + favour_rdm: Whether to favour the density matrix over the number of electrons in the + non-commuting solutions. + """ + self._get_static = get_static + self._self_energy = self_energy + self._nelec = nelec + self._overlap = overlap + self.set_options(**kwargs) + + def __post_init__(self) -> None: + """Hook called after :meth:`__init__`.""" + # Check the input + if not callable(self.get_static): + raise TypeError("get_static must be a callable function.") + if self.overlap is not None and ( + self.overlap.ndim != 2 or self.overlap.shape[0] != self.overlap.shape[1] + ): + raise ValueError("overlap must be a square matrix or None.") + if self.overlap is not None and self.overlap.shape[0] != self.self_energy.nphys: + raise ValueError( + "overlap must have the same number of physical states as the self-energy." + ) + + # Print the input information + console.print(f"Number of physical states: [input]{self.nphys}[/input]") + console.print(f"Number of auxiliary states: [input]{self.self_energy.naux}[/input]") + console.print(f"Target number of electrons: [input]{self.nelec}[/input]") + if self.overlap is not None: + cond = printing.format_float( + np.linalg.cond(self.overlap), threshold=1e10, scientific=True, precision=4 + ) + console.print(f"Overlap condition number: {cond}") + + def __post_kernel__(self) -> None: + """Hook called after :meth:`kernel`.""" + assert self.result is not None + emin = printing.format_float(self.result.eigvals.min()) + emax = printing.format_float(self.result.eigvals.max()) + console.print( + f"Found [output]{self.result.neig}[/output] roots between [output]{emin}[/output] and " + f"[output]{emax}[/output]." + ) + nelec = np.trace(self.result.get_greens_function().occupied().moment(0)) * self.occupancy + if self.result.chempot is not None: + cpt = printing.format_float(self.result.chempot) + console.print(f"Chemical potential: [output]{cpt}[/output]") + err = printing.format_float( + self.nelec - nelec, threshold=1e-3, precision=4, scientific=True + ) + console.print(f"Error in number of electrons: [output]{err}[/output]") + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> DensityRelaxation: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + + Notes: + To initialise this solver from a self-energy, the ``nelec`` and ``get_static`` keyword + arguments must be provided. + """ + if "nelec" not in kwargs: + raise ValueError("Missing required argument nelec.") + if "get_static" not in kwargs: + raise ValueError("Missing required argument get_static.") + kwargs = kwargs.copy() + nelec = kwargs.pop("nelec") + get_static = kwargs.pop("get_static") + return cls(get_static, self_energy, nelec, overlap=overlap, **kwargs) + + @classmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> DensityRelaxation: + """Create a solver from an expression. + + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + raise NotImplementedError( + "Cannot instantiate DensityRelaxation from expression, use from_self_energy instead." + ) + + def kernel(self) -> Spectral: + """Run the solver. + + Returns: + The eigenvalues and eigenvectors of the self-energy supermatrix. + """ + # Get the table + table = printing.ConvergencePrinter( + ("Shift",), + ( + "Error", + "Gradient", + "Change in RDM", + ), + ( + self.solver_outer.conv_tol, + self.solver_outer.conv_tol_grad, + self.conv_tol, + ), + ) + progress = printing.IterationsPrinter(self.max_cycle_outer) + progress.start() + + # Get the initial parameters + self_energy = self.self_energy + nocc = self.nelec // self.occupancy + rdm1 = np.diag(np.arange(self.nphys) < nocc).astype(self_energy.dtype) * self.occupancy + static = self.get_static(rdm1) + + converged = False + for cycle_outer in range(1, self.max_cycle_outer + 1): + if self.favour_rdm: + # Solve the self-energy + with printing.quiet: + solver_outer = self.solver_outer.from_self_energy( + static, self_energy, nelec=self.nelec, overlap=self.overlap + ) + result = solver_outer.kernel() + self_energy = result.get_self_energy() + + # Initialise DIIS for the inner loop + diis = lib.diis.DIIS() + diis.space = self.diis_min_space + diis.max_space = self.diis_max_space + diis.incore = True + diis.verbose = 0 + + for cycle_inner in range(1, self.max_cycle_inner + 1): + # Solve the self-energy + with printing.quiet: + solver_inner = self.solver_inner.from_self_energy( + static, self_energy, nelec=self.nelec, overlap=self.overlap + ) + result = solver_inner.kernel() + self_energy = result.get_self_energy() + + # Get the density matrix + greens_function = result.get_greens_function() + rdm1_prev = rdm1.copy() + rdm1 = greens_function.occupied().moment(0) * self.occupancy + + # Update the static self-energy + static = self.get_static(rdm1, rdm1_prev=rdm1_prev, static_prev=static) + try: + if not self_energy.hermitian and not np.iscomplexobj(static): + # Avoid casting errors if non-Hermitian self-energy starts as real and + # becomes complex during the iterations... probably more efficient to + # subclass DIIS to handle this. + static = static.astype(np.complex128) + static = diis.update(static, xerr=None) + except np.linalg.LinAlgError: + pass + + # Check for convergence + error = np.max(np.abs(rdm1 - rdm1_prev)) + if error < self.conv_tol: + break + + if not self.favour_rdm: + # Solve the self-energy + with printing.quiet: + solver_outer = self.solver_outer.from_self_energy( + static, self_energy, nelec=self.nelec, overlap=self.overlap + ) + result = solver_outer.kernel() + self_energy = result.get_self_energy() + + # Check for convergence + converged = bool(error < self.conv_tol and solver_outer.converged) + grad = np.ravel(solver_outer.gradient(solver_outer.shift)[1])[0] + table.add_row(cycle_outer, (solver_outer.shift,), (solver_outer.error, grad, error)) + progress.update(cycle_outer) + if converged: + break + + progress.stop() + table.print() + + # Set the results + self.converged = converged + self.result = result + + return result + + @property + def get_static(self) -> StaticFunction: + """Get the static self-energy function.""" + return self._get_static + + @property + def self_energy(self) -> Lehmann: + """Get the self-energy.""" + return self._self_energy + + @property + def nelec(self) -> int: + """Get the target number of electrons.""" + return self._nelec + + @property + def overlap(self) -> Array | None: + """Get the overlap matrix for the physical space.""" + return self._overlap + + @property + def nphys(self) -> int: + """Get the number of physical states.""" + return self.self_energy.nphys diff --git a/dyson/solvers/static/downfolded.py b/dyson/solvers/static/downfolded.py new file mode 100644 index 0000000..bf47223 --- /dev/null +++ b/dyson/solvers/static/downfolded.py @@ -0,0 +1,223 @@ +"""Downfolded frequency-space diagonalisation.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from dyson import console, printing, util +from dyson import numpy as np +from dyson.grids.frequency import RealFrequencyGrid +from dyson.representations.enums import Ordering +from dyson.representations.lehmann import Lehmann +from dyson.representations.spectral import Spectral +from dyson.solvers.solver import StaticSolver + +if TYPE_CHECKING: + from typing import Any, Callable + + from dyson.expressions.expression import BaseExpression + from dyson.typing import Array + +# TODO: Use Newton solver as C* Σ(ω) C - ω = 0 +# TODO: Diagonal version + + +class Downfolded(StaticSolver): + r"""Downfolded frequency-space diagonalisation. + + Self-consistently satisfies the eigenvalue problem + + .. math:: + \Sigma(\omega) C = \omega C + + where :math:`\Sigma(\omega)` is the downfolded self-energy. + + Args: + static: The static part of the self-energy. + function: The function to return the downfolded self-energy at a given frequency, the only + argument. + """ + + guess: float = 0.0 + max_cycle: int = 100 + conv_tol: float = 1e-8 + hermitian: bool = True + _options: set[str] = {"guess", "max_cycle", "conv_tol", "hermitian"} + + converged: bool | None = None + + def __init__( # noqa: D417 + self, + static: Array, + function: Callable[[float], Array], + overlap: Array | None = None, + **kwargs: Any, + ): + """Initialise the solver. + + Args: + static: The static part of the self-energy. + function: The function to return the downfolded self-energy at a given frequency, the + only argument. + overlap: Overlap matrix for the physical space. + guess: Initial guess for the eigenvalue. + max_cycle: Maximum number of iterations. + conv_tol: Convergence tolerance for the eigenvalue. + hermitian: Whether the matrix is hermitian. + """ + self._static = static + self._function = function + self._overlap = overlap + self.set_options(**kwargs) + + def __post_init__(self) -> None: + """Hook called after :meth:`__init__`.""" + # Check the input + if self.static.ndim != 2 or self.static.shape[0] != self.static.shape[1]: + raise ValueError("static must be a square matrix.") + if not callable(self.function): + raise ValueError("function must be a callable that takes a single float argument.") + if self.overlap is not None and ( + self.overlap.ndim != 2 or self.overlap.shape[0] != self.overlap.shape[1] + ): + raise ValueError("overlap must be a square matrix or None.") + if self.overlap is not None and self.overlap.shape != self.static.shape: + raise ValueError("overlap must have the same shape as static.") + + # Print the input information + console.print(f"Matrix shape: [input]{self.static.shape}[/input]") + console.print(f"Number of physical states: [input]{self.nphys}[/input]") + if self.overlap is not None: + cond = printing.format_float( + np.linalg.cond(self.overlap), threshold=1e10, scientific=True, precision=4 + ) + console.print(f"Overlap condition number: {cond}") + + def __post_kernel__(self) -> None: + """Hook called after :meth:`kernel`.""" + assert self.result is not None + emin = printing.format_float(self.result.eigvals.min()) + emax = printing.format_float(self.result.eigvals.max()) + console.print( + f"Found [output]{self.result.neig}[/output] roots between [output]{emin}[/output] and " + f"[output]{emax}[/output]." + ) + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> Downfolded: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + kwargs = kwargs.copy() + eta = kwargs.pop("eta", 1e-3) + + def _function(freq: float) -> Array: + """Evaluate the self-energy at the frequency.""" + grid = RealFrequencyGrid(np.array([freq]), eta=eta) + return grid.evaluate_lehmann(self_energy, ordering=Ordering.ORDERED).array[0] + + return cls( + static, + _function, + overlap=overlap, + hermitian=self_energy.hermitian, + **kwargs, + ) + + @classmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> Downfolded: + """Create a solver from an expression. + + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + raise NotImplementedError( + "Cannot instantiate Downfolded solver from an expression, use from_self_energy instead." + ) + + def kernel(self) -> Spectral: + """Run the solver. + + Returns: + The eigenvalues and eigenvectors of the self-energy supermatrix. + """ + # Get the table + table = printing.ConvergencePrinter(("Best root",), ("Change",), (self.conv_tol,)) + progress = printing.IterationsPrinter(self.max_cycle) + progress.start() + + # Initialise the guess + root = self.guess + root_prev = 0.0 + + converged = False + for cycle in range(1, self.max_cycle + 1): + # Update the root + matrix = self.static + self.function(root) + roots, _ = util.eig(matrix, overlap=self.overlap, hermitian=self.hermitian) + root_prev = root + root = roots[np.argmin(np.abs(roots - self.guess))] + + # Check for convergence + converged = np.abs(root - root_prev) < self.conv_tol + table.add_row(cycle, (root,), (root - root_prev,)) + progress.update(cycle) + if converged: + break + + progress.stop() + table.print() + + # Get final eigenvalues and eigenvectors + matrix = self.static + self.function(root) + if self.hermitian: + eigvals, eigvecs = util.eig(matrix, hermitian=self.hermitian, overlap=self.overlap) + else: + eigvals, eigvecs_tuple = util.eig_lr( + matrix, hermitian=self.hermitian, overlap=self.overlap + ) + eigvecs = np.array(eigvecs_tuple) + + # Store the results + self.result = Spectral(eigvals, eigvecs, self.nphys) + self.converged = converged + + return self.result + + @property + def static(self) -> Array: + """Get the static part of the self-energy.""" + return self._static + + @property + def function(self) -> Callable[[float], Array]: + """Get the function to return the downfolded self-energy at a given frequency.""" + return self._function + + @property + def overlap(self) -> Array | None: + """Get the overlap matrix for the physical space.""" + return self._overlap + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return self._static.shape[0] diff --git a/dyson/solvers/static/exact.py b/dyson/solvers/static/exact.py new file mode 100644 index 0000000..4111644 --- /dev/null +++ b/dyson/solvers/static/exact.py @@ -0,0 +1,281 @@ +"""Exact diagonalisation.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +from dyson import console, printing, util +from dyson import numpy as np +from dyson.representations.lehmann import Lehmann +from dyson.representations.spectral import Spectral +from dyson.solvers.solver import StaticSolver + +if TYPE_CHECKING: + from typing import Any + + from dyson.expressions.expression import BaseExpression + from dyson.typing import Array + + +def project_eigenvectors( + eigvecs: Array, + bra: Array, + ket: Array | None = None, +) -> Array: + """Project eigenvectors onto the physical plus auxiliary space. + + Args: + eigvecs: Eigenvectors to be projected. + bra: Bra state vector mapping the supermatrix to the physical space. + ket: Ket state vector mapping the supermatrix to the physical space. If ``None``, use the + same vectors as ``bra``. + + Returns: + Projected eigenvectors. + + Notes: + The physical space is defined by the ``bra`` and ``ket`` vectors, while the auxiliary part + is defined by the null space of the projector formed by the outer product of these vectors. + """ + hermitian = ket is None + nphys = bra.shape[0] + if not hermitian and eigvecs.ndim == 2: + raise ValueError( + "bra and ket both passed implying a non-hermitian system, but eigvecs is 2D." + ) + if ket is None: + ket = bra + + # Find a basis for the null space of the bra and ket vectors + projector = ket.T @ bra.conj() + vectors = util.null_space_basis(projector, hermitian=hermitian) + + # If the system is hermitian, the rotation is trivial + if hermitian: + rotation = np.concatenate([bra.T, vectors[0]], axis=1) + return rotation.T.conj() @ eigvecs + + # If the system is not hermitian, we need to ensure biorthonormality + overlap = ket.conj() @ bra.T + orth, orth_error = util.matrix_power(overlap, -0.5, hermitian=hermitian, return_error=True) + unorth, unorth_error = util.matrix_power(overlap, 0.5, hermitian=hermitian, return_error=True) + + # Work in an orthonormal physical basis + bra = bra.T @ orth + ket = ket.T @ orth.T.conj() + + # Biorthonormalise the physical plus auxiliary vectors + left = np.concatenate([ket, vectors[0]], axis=1) + right = np.concatenate([bra, vectors[1]], axis=1) + left, right = util.biorthonormalise(left, right) + + # Return the physical vectors to the original basis + left[:, :nphys] = left[:, :nphys] @ unorth.T.conj() + right[:, :nphys] = right[:, :nphys] @ unorth + + # Rotate the eigenvectors + eigvecs = np.array([left.T.conj() @ eigvecs[0], right.T.conj() @ eigvecs[1]]) + + return eigvecs + + +def orthogonalise_self_energy( + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, +) -> tuple[Array, Lehmann, Array, Array | None]: + """Orthogonalise a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. If ``None``, assume identity. + + Returns: + The static part of the self-energy and the self-energy itself, projected into an orthogonal + basis, along with the ``bra`` and ``ket`` vectors to map the supermatrix to the original + physical space. + + Notes: + The ``bra`` and ``ket`` vectors essentially transform the orthogonalised self-energy from + the orthogonal basis to the original basis. The main use of this function is to generate a + self-energy and corresponding ``bra`` and ``ket`` vectors that can reproduce a Green's + function with a non-identity zeroth moment (overlap). + """ + size = self_energy.nphys + self_energy.naux + hermitian = self_energy.hermitian + bra = np.array([util.unit_vector(size, i) for i in range(self_energy.nphys)]) + ket = bra if not hermitian else None + + if overlap is not None: + if hermitian: + orth = util.matrix_power(overlap, 0.5, hermitian=hermitian)[0] + unorth = util.matrix_power(overlap, -0.5, hermitian=hermitian)[0] + bra = util.rotate_subspace(bra, orth) + ket = util.rotate_subspace(ket, orth.T.conj()) if ket is not None else None + static = unorth @ static @ unorth + self_energy = self_energy.rotate_couplings(unorth) + else: + bra = util.rotate_subspace(bra, overlap) + orth = util.matrix_power(overlap, -1, hermitian=hermitian)[0] + eye = np.eye(self_energy.nphys) + static = orth @ static + self_energy = self_energy.rotate_couplings((eye, orth.T.conj())) + + return static, self_energy, bra, ket + + +class Exact(StaticSolver): + """Exact diagonalisation of the supermatrix form of the self-energy. + + Args: + matrix: The self-energy supermatrix. + bra: The bra excitation vector mapping the supermatrix to the physical space. + ket: The ket excitation vector mapping the supermatrix to the physical space. + """ + + hermitian: bool = True + _options: set[str] = {"hermitian"} + + def __init__( # noqa: D417 + self, + matrix: Array, + bra: Array, + ket: Array | None = None, + **kwargs: Any, + ): + """Initialise the solver. + + Args: + matrix: The self-energy supermatrix. + bra: The bra excitation vector mapping the supermatrix to the physical space. + ket: The ket excitation vector mapping the supermatrix to the physical space. If `None`, + use the same vectors as `bra`. + hermitian: Whether the matrix is hermitian. + """ + self._matrix = matrix + self._bra = bra + self._ket = ket + self.set_options(**kwargs) + + def __post_init__(self) -> None: + """Hook called after :meth:`__init__`.""" + # Check the input + if self.matrix.ndim != 2 or self.matrix.shape[0] != self.matrix.shape[1]: + raise ValueError("matrix must be a square matrix.") + if self.bra.ndim != 2 or self.bra.shape[1] != self.matrix.shape[0]: + raise ValueError("bra must be a 2D array with the same number of columns as matrix.") + if self.ket is not None and ( + self.ket.ndim != 2 or self.ket.shape[1] != self.matrix.shape[0] + ): + raise ValueError("ket must be a 2D array with the same number of columns as matrix.") + if self.ket is not None and self.ket.shape[0] != self.bra.shape[0]: + raise ValueError("ket must have the same number of rows as bra.") + + # Print the input information + console.print(f"Matrix shape: [input]{self.matrix.shape}[/input]") + console.print(f"Number of physical states: [input]{self.nphys}[/input]") + + def __post_kernel__(self) -> None: + """Hook called after :meth:`kernel`.""" + assert self.result is not None + emin = printing.format_float(self.result.eigvals[np.argmin(np.abs(self.result.eigvals))]) + emax = printing.format_float(self.result.eigvals[np.argmax(np.abs(self.result.eigvals))]) + console.print("") + console.print( + f"Found [output]{self.result.neig}[/output] roots between [output]{emin}[/output] and " + f"[output]{emax}[/output]." + ) + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> Exact: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + static, self_energy, bra, ket = orthogonalise_self_energy( + static, self_energy, overlap=overlap + ) + return cls( + self_energy.matrix(static), + bra, + ket, + hermitian=self_energy.hermitian, + **kwargs, + ) + + @classmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> Exact: + """Create a solver from an expression. + + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + matrix = expression.build_matrix() + bra = np.array(expression.get_excitation_bras()) + ket = ( + np.array(expression.get_excitation_kets()) + if not expression.hermitian_upfolded + else None + ) + return cls(matrix, bra, ket, hermitian=expression.hermitian_upfolded, **kwargs) + + def kernel(self) -> Spectral: + """Run the solver. + + Returns: + The eigenvalues and eigenvectors of the self-energy supermatrix. + """ + # Get the raw eigenvalues and eigenvectors + if self.hermitian: + eigvals, eigvecs = util.eig(self.matrix, hermitian=self.hermitian) + else: + eigvals, (left, right) = util.eig_lr(self.matrix, hermitian=self.hermitian) + eigvecs = np.array([left, right]) + + # Get the full map onto physical + auxiliary and rotate the eigenvectors + eigvecs = project_eigenvectors(eigvecs, self.bra, self.ket if not self.hermitian else None) + + # Store the result + self.result = Spectral(eigvals, eigvecs, self.nphys) + + return self.result + + @property + def matrix(self) -> Array: + """Get the self-energy supermatrix.""" + return self._matrix + + @property + def bra(self) -> Array: + """Get the bra state vector mapping the supermatrix to the physical space.""" + return self._bra + + @property + def ket(self) -> Array: + """Get the ket state vector mapping the supermatrix to the physical space.""" + if self._ket is None: + return self._bra + return self._ket + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return self.bra.shape[0] diff --git a/dyson/solvers/static/mblgf.py b/dyson/solvers/static/mblgf.py new file mode 100644 index 0000000..2030cc1 --- /dev/null +++ b/dyson/solvers/static/mblgf.py @@ -0,0 +1,459 @@ +"""Moment block Lanczos for moments of the Green's function [1]_. + +.. [1] Backhouse, O. J., & Booth, G. H. (2022). Constructing “Full-Frequency” spectra + via moment constraints for coupled cluster Green’s functions. Journal of Chemical Theory and + Computation, 18(11), 6622–6636. https://doi.org/10.1021/acs.jctc.2c00670 +""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING + +from dyson import console, printing, util +from dyson import numpy as np +from dyson.representations.spectral import Spectral +from dyson.solvers.static._mbl import BaseMBL, BaseRecursionCoefficients + +if TYPE_CHECKING: + from typing import Any + + from dyson.expressions.expression import BaseExpression + from dyson.representations.lehmann import Lehmann + from dyson.typing import Array + + +class RecursionCoefficients(BaseRecursionCoefficients): + """Recursion coefficients for the moment block Lanczos algorithm for the Green's function. + + Args: + nphys: Number of physical degrees of freedom. + """ + + def __getitem__(self, key: tuple[int, ...]) -> Array: + """Get the recursion coefficients for the given key. + + Args: + key: The key for the recursion coefficients. + + Returns: + The recursion coefficients. + """ + i, j = key + if i == j == 1: + return np.eye(self.nphys) + if i < 1 or j < 1 or i < j: + return self._zero + return self._data[i, j] + + def __setitem__(self, key: tuple[int, ...], value: Array) -> None: + """Set the recursion coefficients for the given key. + + Args: + key: The key for the recursion coefficients. + value: The recursion coefficients. + """ + self._data[key] = value + + +def _infer_max_cycle(moments: Array) -> int: + """Infer the maximum number of cycles from the moments.""" + return (moments.shape[0] - 2) // 2 + + +class MBLGF(BaseMBL): + """Moment block Lanczos for moments of the Green's function. + + Args: + moments: Moments of the Green's function. + """ + + Coefficients = RecursionCoefficients + + def __init__( # noqa: D417 + self, + moments: Array, + **kwargs: Any, + ) -> None: + """Initialise the solver. + + Args: + moments: Moments of the Green's function. + max_cycle: Maximum number of cycles. + hermitian: Whether the Green's function is hermitian. + force_orthogonality: Whether to force orthogonality of the recursion coefficients. + calculate_errors: Whether to calculate errors. + """ + self._moments = moments + self.max_cycle = kwargs["max_cycle"] if "max_cycle" in kwargs else _infer_max_cycle(moments) + self.set_options(**kwargs) + + if self.hermitian: + self._coefficients = ( + self.Coefficients( + self.nphys, + hermitian=self.hermitian, + force_orthogonality=self.force_orthogonality, + ), + ) * 2 + else: + self._coefficients = ( + self.Coefficients( + self.nphys, + hermitian=self.hermitian, + force_orthogonality=self.force_orthogonality, + ), + self.Coefficients( + self.nphys, + hermitian=self.hermitian, + force_orthogonality=self.force_orthogonality, + ), + ) + self._on_diagonal: dict[int, Array] = {} + self._off_diagonal_upper: dict[int, Array] = {} + self._off_diagonal_lower: dict[int, Array] = {} + + def __post_init__(self) -> None: + """Hook called after :meth:`__init__`.""" + # Check the input + if self.moments.ndim != 3 or self.moments.shape[1] != self.moments.shape[2]: + raise ValueError( + "moments must be a 3D array with the second and third dimensions equal." + ) + if _infer_max_cycle(self.moments) < self.max_cycle: + raise ValueError("not enough moments provided for the specified max_cycle.") + + # Print the input information + cond = printing.format_float( + np.linalg.cond(self.moments[0]), threshold=1e10, scientific=True, precision=4 + ) + console.print(f"Number of physical states: [input]{self.nphys}[/input]") + console.print(f"Number of moments: [input]{self.moments.shape[0]}[/input]") + console.print(f"Overlap condition number: {cond}") + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> MBLGF: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + max_cycle = kwargs.get("max_cycle", 0) + energies, couplings = self_energy.diagonalise_matrix_with_projection( + static, overlap=overlap + ) + greens_function = self_energy.__class__(energies, couplings, chempot=self_energy.chempot) + moments = greens_function.moments(range(2 * max_cycle + 2)) + return cls(moments, hermitian=greens_function.hermitian, **kwargs) + + @classmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> MBLGF: + """Create a solver from an expression. + + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + moments = expression.build_gf_moments(2 * kwargs.get("max_cycle", 0) + 2) + return cls(moments, hermitian=expression.hermitian_downfolded, **kwargs) + + def reconstruct_moments(self, iteration: int) -> Array: + """Reconstruct the moments. + + Args: + iteration: The iteration number. + + Returns: + The reconstructed moments. + """ + greens_function = self.solve(iteration=iteration).get_greens_function() + return greens_function.moments(range(2 * iteration + 2)) + + @functools.lru_cache(maxsize=None) + def _rotated_moment(self, i: int, j: int, k: int, jk: int) -> Array: # noqa: D417 + """Compute an orthogonalised moment rotated by given coefficients. + + Equivalent to the expression + + .. code-block:: python + coefficients[0][i, k].T.conj() @ moments[jk] @ coefficients[0][i, j] + + for Hermitian Green's functions, or + + .. code-block:: python + coefficients[1][i, k] @ moments[jk] @ coefficients[0][i, j] + + for non-Hermitian Green's functions. + + Args: + i, j, k, jk: Indices for the coefficients and moments, as defined above. + + Returns: + The orthogonalised moment rotated by the coefficients. + """ + moment = self.orthogonalised_moment(jk) + if self.hermitian: + return self.coefficients[0][i, k].T.conj() @ moment @ self.coefficients[0][i, j] + else: + return self.coefficients[1][i, k] @ moment @ self.coefficients[0][i, j] + + def initialise_recurrence(self) -> tuple[float | None, float | None, float | None]: + """Initialise the recurrence (zeroth iteration). + + Returns: + If :attr:`calculate_errors`, the error metrics in the square root of the off-diagonal + block, the inverse square root of the off-diagonal block, and the error in the + recovered moments. If not, all three are ``None``. + """ + # Get the inverse square-root error + error_inv_sqrt: float | None = None + if self.calculate_errors: + _, error_inv_sqrt = util.matrix_power( + self.moments[0], -0.5, hermitian=self.hermitian, return_error=True + ) + + # Initialise the blocks + dtype = np.result_type(self.coefficients[0].dtype, self.coefficients[1].dtype) + self.off_diagonal_upper[-1] = np.zeros((self.nphys, self.nphys), dtype=dtype) + self.off_diagonal_lower[-1] = np.zeros((self.nphys, self.nphys), dtype=dtype) + self.on_diagonal[0] = self.orthogonalised_moment(1) + error_sqrt = 0.0 + + # Get the error in the moments + error_moments: float | None = None + if self.calculate_errors: + error_moments = self.moment_error(iteration=0) + + return error_sqrt, error_inv_sqrt, error_moments + + def _recurrence_iteration_hermitian( + self, iteration: int + ) -> tuple[float | None, float | None, float | None]: + """Perform an iteration of the recurrence for a Hermitian Green's function.""" + i = iteration - 1 + coefficients = self.coefficients[0] + on_diagonal = self.on_diagonal + off_diagonal = self.off_diagonal_upper + dtype = np.result_type( + coefficients.dtype, + *[self.orthogonalised_moment(k).dtype for k in range(2 * i + 3)], + on_diagonal[i].dtype, + off_diagonal[i - 1].dtype if i else np.float64, + ) + + # Find the squre of the off-diagonal block + off_diagonal_squared = np.zeros((self.nphys, self.nphys), dtype=dtype) + for j in range(i + 2): + for k in range(i + 1): + off_diagonal_squared += self._rotated_moment(i + 1, j, k + 1, j + k + 1) + off_diagonal_squared -= on_diagonal[i] @ on_diagonal[i] + if i: + off_diagonal_squared -= off_diagonal[i - 1] @ off_diagonal[i - 1] + + # Get the off-diagonal block + off_diagonal[i], error_sqrt = util.matrix_power( + off_diagonal_squared, 0.5, hermitian=self.hermitian, return_error=self.calculate_errors + ) + + # Invert the off-diagonal block + off_diagonal_inv, error_inv_sqrt = util.matrix_power( + off_diagonal_squared, -0.5, hermitian=self.hermitian, return_error=self.calculate_errors + ) + + # Update the dtype + dtype = np.result_type(dtype, off_diagonal_inv.dtype) + + for j in range(i + 2): + # Horizontal recursion + residual = coefficients[i + 1, j].astype(dtype, copy=True) + residual -= coefficients[i + 1, j + 1] @ on_diagonal[i] + residual -= coefficients[i, j + 1] @ off_diagonal[i - 1] + coefficients[i + 2, j + 1] = residual @ off_diagonal_inv + + # Calculate the on-diagonal block + on_diagonal[i + 1] = np.zeros((self.nphys, self.nphys), dtype=dtype) + for j in range(i + 2): + for k in range(i + 2): + on_diagonal[i + 1] += self._rotated_moment(i + 2, j + 1, k + 1, j + k + 1) + + # Get the error in the moments + error_moments: float | None = None + if self.calculate_errors: + error_moments = self.moment_error(iteration=iteration) + + return error_sqrt, error_inv_sqrt, error_moments + + def _recurrence_iteration_non_hermitian( + self, iteration: int + ) -> tuple[float | None, float | None, float | None]: + """Perform an iteration of the recurrence for a non-Hermitian Green's function.""" + i = iteration - 1 + coefficients = self.coefficients + on_diagonal = self.on_diagonal + off_diagonal_upper = self.off_diagonal_upper + off_diagonal_lower = self.off_diagonal_lower + dtype = np.result_type( + coefficients[0].dtype, + coefficients[1].dtype, + *[self.orthogonalised_moment(k).dtype for k in range(2 * i + 3)], + on_diagonal[i].dtype, + off_diagonal_upper[i - 1].dtype if i else np.float64, + ) + + # Find the square of the off-diagonal blocks + off_diagonal_upper_squared = np.zeros((self.nphys, self.nphys), dtype=dtype) + off_diagonal_lower_squared = np.zeros((self.nphys, self.nphys), dtype=dtype) + for j in range(i + 2): + for k in range(i + 1): + off_diagonal_upper_squared += self._rotated_moment(i + 1, j, k + 1, j + k + 1) + off_diagonal_lower_squared += self._rotated_moment(i + 1, k + 1, j, j + k + 1) + off_diagonal_upper_squared -= on_diagonal[i] @ on_diagonal[i] + off_diagonal_lower_squared -= on_diagonal[i] @ on_diagonal[i] + if i: + off_diagonal_upper_squared -= off_diagonal_lower[i - 1] @ off_diagonal_lower[i - 1] + off_diagonal_lower_squared -= off_diagonal_upper[i - 1] @ off_diagonal_upper[i - 1] + + # Get the off-diagonal blocks + off_diagonal_upper[i], error_sqrt_upper = util.matrix_power( + off_diagonal_upper_squared, + 0.5, + hermitian=self.hermitian, + return_error=self.calculate_errors, + ) + off_diagonal_lower[i], error_sqrt_lower = util.matrix_power( + off_diagonal_lower_squared, + 0.5, + hermitian=self.hermitian, + return_error=self.calculate_errors, + ) + error_sqrt: float | None = None + if self.calculate_errors: + assert error_sqrt_upper is not None and error_sqrt_lower is not None + error_sqrt = np.sqrt(error_sqrt_upper**2 + error_sqrt_lower**2) + + # Invert the off-diagonal blocks + off_diagonal_upper_inv, error_inv_sqrt_upper = util.matrix_power( + off_diagonal_upper_squared, + -0.5, + hermitian=self.hermitian, + return_error=self.calculate_errors, + ) + off_diagonal_lower_inv, error_inv_sqrt_lower = util.matrix_power( + off_diagonal_lower_squared, + -0.5, + hermitian=self.hermitian, + return_error=self.calculate_errors, + ) + error_inv_sqrt: float | None = None + if self.calculate_errors: + assert error_inv_sqrt_upper is not None and error_inv_sqrt_lower is not None + error_inv_sqrt = np.sqrt(error_inv_sqrt_upper**2 + error_inv_sqrt_lower**2) + + # Update the dtype + dtype = np.result_type(dtype, off_diagonal_upper_inv.dtype, off_diagonal_lower_inv.dtype) + + for j in range(i + 2): + # Horizontal recursion + residual = coefficients[0][i + 1, j].astype(dtype, copy=True) + residual -= coefficients[0][i + 1, j + 1] @ on_diagonal[i] + residual -= coefficients[0][i, j + 1] @ off_diagonal_upper[i - 1] + coefficients[0][i + 2, j + 1] = residual @ off_diagonal_lower_inv + + # Vertical recursion + residual = coefficients[1][i + 1, j].astype(dtype, copy=True) + residual -= on_diagonal[i] @ coefficients[1][i + 1, j + 1] + residual -= off_diagonal_lower[i - 1] @ coefficients[1][i, j + 1] + coefficients[1][i + 2, j + 1] = off_diagonal_upper_inv @ residual + + # Calculate the on-diagonal block + on_diagonal[i + 1] = np.zeros((self.nphys, self.nphys), dtype=dtype) + for j in range(i + 2): + for k in range(i + 2): + on_diagonal[i + 1] += self._rotated_moment(i + 2, j + 1, k + 1, j + k + 1) + + # Get the error in the moments + error_moments: float | None = None + if self.calculate_errors: + error_moments = self.moment_error(iteration=iteration) + + return error_sqrt, error_inv_sqrt, error_moments + + def solve(self, iteration: int | None = None) -> Spectral: + """Solve the eigenvalue problem at a given iteration. + + Args: + iteration: The iteration to get the results for. + + Returns: + The :class:`~dyson.representations.spectral.Spectral` object. + """ + if iteration is None: + iteration = self.max_cycle + + # Check if we're just returning the result + if iteration == self.max_cycle and self.result is not None: + return self.result + + # Diagonalise the block tridiagonal Hamiltonian + on_diag = [self.on_diagonal[i] for i in range(iteration + 1)] + off_diag_upper = [self.off_diagonal_upper[i] for i in range(iteration)] + off_diag_lower = ( + [self.off_diagonal_lower[i] for i in range(iteration)] if not self.hermitian else None + ) + hamiltonian = util.build_block_tridiagonal(on_diag, off_diag_upper, off_diag_lower) + + # Allow Hermitian solution even for non-Hermitian solver if the Hamiltonian is Hermitian + if self.hermitian: + eigvals, eigvecs = util.eig(hamiltonian, hermitian=self.hermitian) + else: + eigvals, eigvecs_tuple = util.eig_lr(hamiltonian, hermitian=self.hermitian) + eigvecs = np.array(eigvecs_tuple) + + # Unorthogonalise the eigenvectors + metric_inv = self.orthogonalisation_metric_inv + if self.hermitian: + eigvecs[: self.nphys] = metric_inv @ eigvecs[: self.nphys] + else: + eigvecs[:, : self.nphys] = np.array( + [ + metric_inv.T.conj() @ eigvecs[0, : self.nphys], + metric_inv @ eigvecs[1, : self.nphys], + ], + ) + + return Spectral(eigvals, eigvecs, self.nphys) + + @property + def coefficients(self) -> tuple[BaseRecursionCoefficients, BaseRecursionCoefficients]: + """Get the recursion coefficients.""" + return self._coefficients + + @property + def on_diagonal(self) -> dict[int, Array]: + """Get the on-diagonal blocks of the self-energy.""" + return self._on_diagonal + + @property + def off_diagonal_upper(self) -> dict[int, Array]: + """Get the upper off-diagonal blocks of the self-energy.""" + return self._off_diagonal_upper + + @property + def off_diagonal_lower(self) -> dict[int, Array]: + """Get the lower off-diagonal blocks of the self-energy.""" + return self._off_diagonal_lower diff --git a/dyson/solvers/static/mblse.py b/dyson/solvers/static/mblse.py new file mode 100644 index 0000000..e2108c8 --- /dev/null +++ b/dyson/solvers/static/mblse.py @@ -0,0 +1,647 @@ +"""Moment block Lanczos for moments of the self-energy [1]_. + +.. [1] Backhouse, O. J., Santana-Bonilla, A., & Booth, G. H. (2021). Scalable and + Predictive Spectra of Correlated Molecules with Moment Truncated Iterated Perturbation Theory. + The Journal of Physical Chemistry Letters, 12(31), 7650–7658. + https://doi.org/10.1021/acs.jpclett.1c02383 +""" + +from __future__ import annotations + +import functools +from typing import TYPE_CHECKING + +from dyson import console, printing, util +from dyson import numpy as np +from dyson.representations.enums import Reduction +from dyson.representations.lehmann import Lehmann +from dyson.representations.spectral import Spectral +from dyson.solvers.static._mbl import BaseMBL, BaseRecursionCoefficients + +if TYPE_CHECKING: + from typing import Any, TypeVar + + from dyson.expressions.expression import BaseExpression + from dyson.typing import Array + + T = TypeVar("T", bound="BaseMBL") + + +class RecursionCoefficients(BaseRecursionCoefficients): + """Recursion coefficients for the moment block Lanczos algorithm for the self-energy. + + Args: + nphys: Number of physical degrees of freedom. + """ + + def __getitem__(self, key: tuple[int, ...]) -> Array: + """Get the recursion coefficients for the given key. + + Args: + key: The key for the recursion coefficients. + + Returns: + The recursion coefficients. + """ + i, j, order = key + if i == 0 or j == 0: + return self._zero + if i < j and self.hermitian: + return self._data[j, i, order].T.conj() + return self._data[i, j, order] + + def __setitem__(self, key: tuple[int, ...], value: Array) -> None: + """Set the recursion coefficients for the given key. + + Args: + key: The key for the recursion coefficients. + value: The recursion coefficients. + """ + i, j, order = key + if order == 0 and self.force_orthogonality: + value = np.eye(self.nphys) + if self.hermitian and i == j: + value = 0.5 * util.hermi_sum(value) + if i < j and self.hermitian: + self._data[j, i, order] = value.T.conj() + else: + self._data[i, j, order] = value + + +class ScalarRecursionCoefficients(BaseRecursionCoefficients): + """Scalar recursion coefficients for the moment block Lanczos algorithm for the self-energy. + + Args: + nphys: Number of physical degrees of freedom. + """ + + NDIM = 0 + + def __getitem__(self, key: tuple[int, ...]) -> Array: + """Get the recursion coefficients for the given key. + + Args: + key: The key for the recursion coefficients. + + Returns: + The recursion coefficients. + """ + i, j, order = key + if i == 0 or j == 0: + return 0.0 # type: ignore[return-value] + if i < j: + i, j = j, i + return self._data[i, j, order] + + def __setitem__(self, key: tuple[int, ...], value: Array) -> None: + """Set the recursion coefficients for the given key. + + Args: + key: The key for the recursion coefficients. + value: The recursion coefficients. + """ + i, j, order = key + if order == 0 and self.force_orthogonality: + value = 1.0 # type: ignore[assignment] + if i < j: + i, j = j, i + self._data[i, j, order] = np.asarray(value).item() + + +def _infer_max_cycle(moments: Array) -> int: + """Infer the maximum number of cycles from the moments.""" + return (moments.shape[0] - 2) // 2 + + +class MBLSE(BaseMBL): + """Moment block Lanczos for moments of the self-energy. + + Args: + static: Static part of the self-energy. + moments: Moments of the self-energy. + """ + + Coefficients = RecursionCoefficients + + def __init__( # noqa: D417 + self, + static: Array, + moments: Array, + overlap: Array | None = None, + **kwargs: Any, + ) -> None: + """Initialise the solver. + + Args: + static: Static part of the self-energy. + moments: Moments of the self-energy. + overlap: Overlap matrix for the physical space. + max_cycle: Maximum number of cycles. + hermitian: Whether the self-energy is hermitian. + force_orthogonality: Whether to force orthogonality of the recursion coefficients. + calculate_errors: Whether to calculate errors. + """ + self._static = static + self._moments = moments + self._overlap = overlap + self.max_cycle = kwargs["max_cycle"] if "max_cycle" in kwargs else _infer_max_cycle(moments) + self.set_options(**kwargs) + + self._coefficients = self.Coefficients( + self.nphys, + hermitian=self.hermitian, + force_orthogonality=self.force_orthogonality, + ) + self._on_diagonal: dict[int, Array] = {} + self._off_diagonal: dict[int, Array] = {} + + def __post_init__(self) -> None: + """Hook called after :meth:`__init__`.""" + # Check the input + if self.static.ndim != 2 or self.static.shape[0] != self.static.shape[1]: + raise ValueError("static must be a square matrix.") + if self.moments.ndim != 3 or self.moments.shape[1] != self.moments.shape[2]: + raise ValueError( + "moments must be a 3D array with the second and third dimensions equal." + ) + if self.moments.shape[1] != self.static.shape[0]: + raise ValueError( + "moments must have the same shape as static in the last two dimensions." + ) + if _infer_max_cycle(self.moments) < self.max_cycle: + raise ValueError("not enough moments provided for the specified max_cycle.") + + # Print the input information + console.print(f"Number of physical states: [input]{self.nphys}[/input]") + console.print(f"Number of moments: [input]{self.moments.shape[0]}[/input]") + if self.overlap is not None: + cond = printing.format_float( + np.linalg.cond(self.overlap), threshold=1e10, scientific=True, precision=4 + ) + console.print(f"Overlap condition number: {cond}") + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> MBLSE: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + max_cycle = kwargs.get("max_cycle", 0) + moments = self_energy.moments(range(2 * max_cycle + 2)) + return cls(static, moments, hermitian=self_energy.hermitian, overlap=overlap, **kwargs) + + @classmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> MBLSE: + """Create a solver from an expression. + + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + overlap, static = expression.build_gf_moments(2) + moments = expression.build_se_moments(2 * kwargs.get("max_cycle", 0) + 2) + return cls( + static, + moments, + overlap=overlap, + hermitian=expression.hermitian_downfolded, + **kwargs, + ) + + def reconstruct_moments(self, iteration: int) -> Array: + """Reconstruct the moments. + + Args: + iteration: The iteration number. + + Returns: + The reconstructed moments. + """ + self_energy = self.solve(iteration=iteration).get_self_energy() + return self_energy.moments(range(2 * iteration)) + + def initialise_recurrence(self) -> tuple[float | None, float | None, float | None]: + """Initialise the recurrence (zeroth iteration). + + Returns: + If :attr:`calculate_errors`, the error metrics in the square root of the off-diagonal + block, the inverse square root of the off-diagonal block, and the error in the + recovered moments. If not, all three are ``None``. + """ + # Get the inverse square-root error + error_inv_sqrt: float | None = None + if self.calculate_errors: + _, error_inv_sqrt = util.matrix_power( + self.moments[0], -0.5, hermitian=self.hermitian, return_error=True + ) + + # Initialise the coefficients + for n in range(2 * self.max_cycle + 2): + self.coefficients[1, 1, n] = self.orthogonalised_moment(n) + + # Initialise the blocks + self.off_diagonal[0], error_sqrt = util.matrix_power( + self.moments[0], 0.5, hermitian=self.hermitian, return_error=self.calculate_errors + ) + self.on_diagonal[0] = self.static + self.on_diagonal[1] = self.coefficients[1, 1, 1] + + # Get the error in the moments + error_moments: float | None = None + if self.calculate_errors: + error_moments = self.moment_error(iteration=0) + + return error_sqrt, error_inv_sqrt, error_moments + + def _recurrence_iteration_hermitian( + self, iteration: int + ) -> tuple[float | None, float | None, float | None]: + """Perform an iteration of the recurrence for a Hermitian self-energy.""" + i = iteration + coefficients = self.coefficients + on_diagonal = self.on_diagonal + off_diagonal = self.off_diagonal + dtype = np.result_type(coefficients.dtype, off_diagonal[i - 1].dtype).char + + # Find the squre of the off-diagonal block + off_diagonal_squared = coefficients[i, i, 2].astype(dtype, copy=True) + off_diagonal_squared -= util.hermi_sum(coefficients[i, i - 1, 1] @ off_diagonal[i - 1]) + off_diagonal_squared -= coefficients[i, i, 1] @ coefficients[i, i, 1] + if iteration > 1: + off_diagonal_squared += off_diagonal[i - 1].T.conj() @ off_diagonal[i - 1] + + # Get the off-diagonal block + off_diagonal[i], error_sqrt = util.matrix_power( + off_diagonal_squared, 0.5, hermitian=self.hermitian, return_error=self.calculate_errors + ) + + # Invert the off-diagonal block + off_diagonal_inv, error_inv_sqrt = util.matrix_power( + off_diagonal_squared, + -0.5, + hermitian=self.hermitian, + return_error=self.calculate_errors, + ) + + # Update the dtype + dtype = np.result_type(dtype, off_diagonal_inv.dtype).char + + for n in range(2 * (self.max_cycle - iteration + 1)): + # Horizontal recursion + residual = coefficients[i, i, n + 1].astype(dtype, copy=True) + residual -= off_diagonal[i - 1].T.conj() @ coefficients[i - 1, i, n] + residual -= on_diagonal[i] @ coefficients[i, i, n] + coefficients[i + 1, i, n] = off_diagonal_inv @ residual + + # Diagonal recursion + residual = coefficients[i, i, n + 2].astype(dtype, copy=True) + residual -= util.hermi_sum(coefficients[i, i - 1, n + 1] @ off_diagonal[i - 1]) + residual -= util.hermi_sum(coefficients[i, i, n + 1] @ on_diagonal[i]) + residual += util.hermi_sum( + on_diagonal[i] @ coefficients[i, i - 1, n] @ off_diagonal[i - 1] + ) + residual += ( + off_diagonal[i - 1].T.conj() @ coefficients[i - 1, i - 1, n] @ off_diagonal[i - 1] + ) + residual += on_diagonal[i] @ coefficients[i, i, n] @ on_diagonal[i] + coefficients[i + 1, i + 1, n] = off_diagonal_inv @ residual @ off_diagonal_inv.T.conj() + + # Extract the on-diagonal block + on_diagonal[i + 1] = coefficients[i + 1, i + 1, 1].copy() + + # Get the error in the moments + error_moments: float | None = None + if self.calculate_errors: + error_moments = self.moment_error(iteration=iteration) + + return error_sqrt, error_inv_sqrt, error_moments + + def _recurrence_iteration_non_hermitian( + self, iteration: int + ) -> tuple[float | None, float | None, float | None]: + """Perform an iteration of the recurrence for a non-Hermitian self-energy.""" + i = iteration + coefficients = self.coefficients + on_diagonal = self.on_diagonal + off_diagonal = self.off_diagonal + dtype = np.result_type(coefficients.dtype, off_diagonal[i - 1].dtype).char + + # Find the squre of the off-diagonal block + off_diagonal_squared = coefficients[i, i, 2].astype(dtype, copy=True) + off_diagonal_squared -= coefficients[i, i, 1] @ coefficients[i, i, 1] + off_diagonal_squared -= coefficients[i, i - 1, 1] @ off_diagonal[i - 1] + off_diagonal_squared -= off_diagonal[i - 1] @ coefficients[i, i - 1, 1] + if iteration > 1: + off_diagonal_squared += off_diagonal[i - 1] @ off_diagonal[i - 1] + + # Get the off-diagonal block + off_diagonal[i], error_sqrt = util.matrix_power( + off_diagonal_squared, 0.5, hermitian=self.hermitian, return_error=self.calculate_errors + ) + + # Invert the off-diagonal block + off_diagonal_inv, error_inv_sqrt = util.matrix_power( + off_diagonal_squared, -0.5, hermitian=self.hermitian, return_error=self.calculate_errors + ) + + # Update the dtype + dtype = np.result_type(dtype, off_diagonal_inv.dtype).char + + for n in range(2 * (self.max_cycle - iteration + 1)): + # Horizontal recursion + residual = coefficients[i, i, n + 1].astype(dtype, copy=True) + residual -= off_diagonal[i - 1] @ coefficients[i - 1, i, n] + residual -= on_diagonal[i] @ coefficients[i, i, n] + coefficients[i + 1, i, n] = off_diagonal_inv @ residual + + # Vertical recursion + residual = coefficients[i, i, n + 1].astype(dtype, copy=True) + residual -= coefficients[i, i - 1, n] @ off_diagonal[i - 1] + residual -= coefficients[i, i, n] @ on_diagonal[i] + coefficients[i, i + 1, n] = residual @ off_diagonal_inv + + # Diagonal recursion + residual = coefficients[i, i, n + 2].astype(dtype, copy=True) + residual -= coefficients[i, i - 1, n + 1] @ off_diagonal[i - 1] + residual -= coefficients[i, i, n + 1] @ on_diagonal[i] + residual -= off_diagonal[i - 1] @ coefficients[i - 1, i, n + 1] + residual += off_diagonal[i - 1] @ coefficients[i - 1, i - 1, n] @ off_diagonal[i - 1] + residual += off_diagonal[i - 1] @ coefficients[i - 1, i, n] @ on_diagonal[i] + residual -= on_diagonal[i] @ coefficients[i, i, n + 1] + residual += on_diagonal[i] @ coefficients[i, i - 1, n] @ off_diagonal[i - 1] + residual += on_diagonal[i] @ coefficients[i, i, n] @ on_diagonal[i] + coefficients[i + 1, i + 1, n] = off_diagonal_inv @ residual @ off_diagonal_inv + + # Extract the on-diagonal block + on_diagonal[i + 1] = coefficients[i + 1, i + 1, 1].copy() + + # Get the error in the moments + error_moments: float | None = None + if self.calculate_errors: + error_moments = self.moment_error(iteration=iteration) + + return error_sqrt, error_inv_sqrt, error_moments + + def solve(self, iteration: int | None = None) -> Spectral: + """Solve the eigenvalue problem at a given iteration. + + Args: + iteration: The iteration to get the results for. + + Returns: + The :cls:`Spectral` object. + """ + # TODO inherit + if iteration is None: + iteration = self.max_cycle + + # Check if we're just returning the result + if iteration == self.max_cycle and self.result is not None: + return self.result + + # Get the supermatrix + on_diag = [self.on_diagonal[i] for i in range(iteration + 2)] + off_diag_upper = [self.off_diagonal[i] for i in range(iteration + 1)] + off_diag_lower = ( + [self.off_diagonal[i] for i in range(iteration + 1)] if not self.hermitian else None + ) + hamiltonian = util.build_block_tridiagonal(on_diag, off_diag_upper, off_diag_lower) + + # Diagonalise the subspace + subspace = hamiltonian[self.nphys :, self.nphys :] + energies, rotated = util.eig_lr(subspace, hermitian=self.hermitian) + if self.hermitian: + couplings = np.atleast_2d(self.off_diagonal[0]) @ rotated[0][: self.nphys] + else: + couplings = np.array( + [ + np.atleast_2d(self.off_diagonal[0]).T.conj() @ rotated[0][: self.nphys], + np.atleast_2d(self.off_diagonal[0]) @ rotated[1][: self.nphys], + ] + ) + + return Spectral.from_self_energy( + np.atleast_2d(self.static), + Lehmann(energies, couplings), + overlap=np.atleast_2d(self.overlap) if self.overlap is not None else None, + ) + + @property + def static(self) -> Array: + """Get the static part of the self-energy.""" + return self._static + + @property + def overlap(self) -> Array | None: + """Get the overlap matrix for the physical space.""" + return self._overlap + + @property + def coefficients(self) -> BaseRecursionCoefficients: + """Get the recursion coefficients.""" + return self._coefficients + + @property + def on_diagonal(self) -> dict[int, Array]: + """Get the on-diagonal blocks of the self-energy.""" + return self._on_diagonal + + @property + def off_diagonal(self) -> dict[int, Array]: + """Get the off-diagonal blocks of the self-energy.""" + return self._off_diagonal + + +class MLSE(MBLSE): + """Moment Lanczos for scalar moments of the self-energy. + + This is a specialisation of :class:`MBLSE` for scalar moments. + + Args: + static: Static part of the self-energy. + moments: Moments of the self-energy. + """ + + Coefficients = ScalarRecursionCoefficients # type: ignore[assignment] + + def __post_init__(self) -> None: + """Hook called after :meth:`__init__`.""" + # Check the input + if np.size(self.static.ndim) != 1: + raise ValueError("static must be scalar.") + if self.moments.ndim != 1: + raise ValueError("moments must be a 1D array of scalar elements for each order.") + if _infer_max_cycle(self.moments) < self.max_cycle: + raise ValueError("not enough moments provided for the specified max_cycle.") + + # Print the input information + console.print("Number of physical states: [input]1[/input]") + console.print(f"Number of moments: [input]{self.moments.shape[0]}[/input]") + + @classmethod + def from_self_energy( + cls, + static: Array, + self_energy: Lehmann, + overlap: Array | None = None, + **kwargs: Any, + ) -> MBLSE: + """Create a solver from a self-energy. + + Args: + static: Static part of the self-energy. + self_energy: Self-energy. + overlap: Overlap matrix for the physical space. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + raise NotImplementedError("Cannot instantiate MLSE from a self-energy.") + + @classmethod + def from_expression(cls, expression: BaseExpression, **kwargs: Any) -> MBLSE: + """Create a solver from an expression. + + Args: + expression: Expression to be solved. + kwargs: Additional keyword arguments for the solver. + + Returns: + Solver instance. + """ + raise NotImplementedError("Cannot instantiate MLSE from an expression.") + + @property + def orthogonalisation_metric(self) -> Array: + """Get the orthogonalisation metric.""" + return self.moments[0] ** -0.5 + + @property + def orthogonalisation_metric_inv(self) -> Array: + """Get the inverse of the orthogonalisation metric.""" + return self.moments[0] ** 0.5 + + @functools.lru_cache(maxsize=64) + def orthogonalised_moment(self, order: int) -> Array: + """Compute an orthogonalised moment. + + Args: + order: The order of the moment. + + Returns: + The orthogonalised moment. + """ + return self.moments[order] / self.moments[0] + + def reconstruct_moments(self, iteration: int) -> Array: + """Reconstruct the moments. + + Args: + iteration: The iteration number. + + Returns: + The reconstructed moments. + """ + self_energy = self.solve(iteration=iteration).get_self_energy() + return self_energy.moments(range(2 * iteration), reduction=Reduction.TRACE) + + def initialise_recurrence(self) -> tuple[float | None, float | None, float | None]: + """Initialise the recurrence (zeroth iteration). + + Returns: + If :attr:`calculate_errors`, the error metrics in the square root of the off-diagonal + block, the inverse square root of the off-diagonal block, and the error in the + recovered moments. If not, all three are ``None``. + """ + # Initialise the coefficients + for n in range(2 * self.max_cycle + 2): + self.coefficients[1, 1, n] = self.orthogonalised_moment(n) + + # Initialise the blocks + self.off_diagonal[0] = self.moments[0] ** 0.5 + self.on_diagonal[0] = self.static + self.on_diagonal[1] = self.coefficients[1, 1, 1] + + # Get the error in the moments + error_moments: float | None = None + if self.calculate_errors: + error_moments = self.moment_error(iteration=0) + + return 0.0, 0.0, error_moments + + def _recurrence_iteration_hermitian( + self, iteration: int + ) -> tuple[float | None, float | None, float | None]: + """Perform an iteration of the recurrence for a Hermitian self-energy.""" + i = iteration + coefficients = self.coefficients + on_diagonal = self.on_diagonal + off_diagonal = self.off_diagonal + + # Find the squre of the off-diagonal block + off_diagonal_squared = coefficients[i, i, 2] + off_diagonal_squared -= coefficients[i, i - 1, 1] * off_diagonal[i - 1] * 2.0 + off_diagonal_squared -= coefficients[i, i, 1] * coefficients[i, i, 1] + if iteration > 1: + off_diagonal_squared += off_diagonal[i - 1].conj() * off_diagonal[i - 1] + + # Get the off-diagonal block + off_diagonal[i] = off_diagonal_squared**0.5 + + # Invert the off-diagonal block + off_diagonal_inv = off_diagonal_squared**-0.5 + + for n in range(2 * (self.max_cycle - iteration + 1)): + # Horizontal recursion + residual = coefficients[i, i, n + 1] + residual -= off_diagonal[i - 1].conj() * coefficients[i - 1, i, n] + residual -= on_diagonal[i] * coefficients[i, i, n] + coefficients[i + 1, i, n] = off_diagonal_inv * residual + + # Diagonal recursion + residual = coefficients[i, i, n + 2] + residual -= coefficients[i, i - 1, n + 1] * off_diagonal[i - 1] * 2.0 + residual -= coefficients[i, i, n + 1] * on_diagonal[i] * 2.0 + residual += on_diagonal[i] * coefficients[i, i - 1, n] * off_diagonal[i - 1] * 2.0 + residual += ( + off_diagonal[i - 1].conj() * coefficients[i - 1, i - 1, n] * off_diagonal[i - 1] + ) + residual += on_diagonal[i] * coefficients[i, i, n] * on_diagonal[i] + coefficients[i + 1, i + 1, n] = off_diagonal_inv * residual * off_diagonal_inv.conj() + + # Extract the on-diagonal block + on_diagonal[i + 1] = coefficients[i + 1, i + 1, 1] + + # Get the error in the moments + error_moments: float | None = None + if self.calculate_errors: + error_moments = self.moment_error(iteration=iteration) + + return 0.0, 0.0, error_moments + + _recurrence_iteration_non_hermitian = _recurrence_iteration_hermitian + _recurrence_iteration_non_hermitian.__doc__ = ( + BaseMBL._recurrence_iteration_non_hermitian.__doc__ + ) + + @property + def nphys(self) -> int: + """Get the number of physical degrees of freedom.""" + return 1 diff --git a/dyson/typing.py b/dyson/typing.py new file mode 100644 index 0000000..8cf9534 --- /dev/null +++ b/dyson/typing.py @@ -0,0 +1,7 @@ +"""Typing.""" + +from __future__ import annotations + +from dyson import numpy + +Array = numpy.ndarray diff --git a/dyson/util/__init__.py b/dyson/util/__init__.py index 65c2a45..d968eef 100644 --- a/dyson/util/__init__.py +++ b/dyson/util/__init__.py @@ -1,6 +1,43 @@ -from dyson.util.misc import * -from dyson.util.spectra import * -from dyson.util.logging import * -from dyson.util.linalg import * -from dyson.util.moments import * -from dyson.util.energy import * +"""Utility functions. + +Submodules +---------- + +.. autosummary:: + :toctree: + + linalg + moments + energy + misc + +""" + +from dyson.util.misc import catch_warnings, cache_by_id, get_mean_field +from dyson.util.linalg import ( + einsum, + orthonormalise, + biorthonormalise, + biorthonormalise_with_overlap, + eig, + eig_lr, + matrix_power, + hermi_sum, + scaled_error, + as_diagonal, + as_trace, + unit_vector, + null_space_basis, + concatenate_paired_vectors, + unpack_vectors, + block_diag, + set_subspace, + rotate_subspace, +) +from dyson.util.moments import ( + se_moments_to_gf_moments, + gf_moments_to_se_moments, + build_block_tridiagonal, + get_chebyshev_scaling_parameters, +) +from dyson.util.energy import gf_moments_galitskii_migdal diff --git a/dyson/util/energy.py b/dyson/util/energy.py index 42fa685..086597e 100644 --- a/dyson/util/energy.py +++ b/dyson/util/energy.py @@ -1,34 +1,29 @@ -""" -Energy functionals. -""" +"""Energy functionals.""" -import numpy as np +from __future__ import annotations +from typing import TYPE_CHECKING -def greens_function_galitskii_migdal(gf_moments_hole, hcore, factor=1.0): - """ - Compute the energy using the Galitskii-Migdal formula in terms of - the hole Green's function moments and the core Hamiltonian. - - Parameters - ---------- - gf_moments : numpy.ndarray (m, n, n) - Moments of the hole Green's function. Only the first two - (n=0 and n=1) are required. - hcore : numpy.ndarray (n, n) - Core Hamiltonian. - factor : float, optional - Factor to scale energy. For UHF and GHF calculations, this - should likely be 0.5, for RHF it is 1.0. Default value is - `1.0`. - - Returns - ------- - e_gm : float - Galitskii-Migdal energy. - """ +from dyson import numpy as np +from dyson import util + +if TYPE_CHECKING: + from dyson.typing import Array - e_gm = np.einsum("pq,qp->", gf_moments_hole[0], hcore) - e_gm += np.trace(gf_moments_hole[1]) - return e_gm +def gf_moments_galitskii_migdal(gf_moments_hole: Array, hcore: Array, factor: float = 1.0) -> float: + """Compute the Galitskii--Migdal energy in terms of the moments of the hole Green's function. + + Args: + gf_moments_hole: Moments of the hole Green's function. Only the first two (zeroth and first + moments) are required. + hcore: Core Hamiltonian. + factor: Factor to scale energy. For UHF and GHF calculations, this should likely be 0.5, + for RHF it is 1.0. + + Returns: + Galitskii--Migdal energy. + """ + e_gm = util.einsum("pq,qp->", gf_moments_hole[0], hcore) + e_gm += np.trace(gf_moments_hole[1]) + return e_gm * factor diff --git a/dyson/util/linalg.py b/dyson/util/linalg.py index e8fcdf0..7f8d9c6 100644 --- a/dyson/util/linalg.py +++ b/dyson/util/linalg.py @@ -1,158 +1,560 @@ +"""Linear algebra.""" + +from __future__ import annotations + +import functools +import warnings +from typing import TYPE_CHECKING, cast + +import scipy.linalg + +from dyson import numpy as np +from dyson.util import cache_by_id + +if TYPE_CHECKING: + from typing import Literal + + from dyson.typing import Array + +einsum = functools.partial(np.einsum, optimize=True) + +"""Flag to avoid using :func:`scipy.linalg.eig` and :func:`scipy.linalg.eigh`. + +On some platforms, mixing :mod:`numpy` and :mod:`scipy` eigenvalue solvers can lead to performance +issues, likely from repeating warm-up overhead from conflicting BLAS and/or LAPACK libraries. """ -Linear algebra utilities. -""" +AVOID_SCIPY_EIG: bool = True + +"""Default biorthonormalisation method.""" +BIORTH_METHOD: Literal["lu", "eig", "eig-balanced"] = "lu" + + +def is_orthonormal(vectors_left: Array, vectors_right: Array | None = None) -> bool: + """Check if a set of vectors is orthonormal. + + Args: + vectors_left: The left set of vectors to be checked. + vectors_right: The right set of vectors to be checked. If ``None``, use the left vectors. + + Returns: + A boolean array indicating whether each vector is orthonormal. + """ + if vectors_right is None: + vectors_right = vectors_left + if vectors_left.ndim == 1: + vectors_left = vectors_left[:, None] + if vectors_right.ndim == 1: + vectors_right = vectors_right[:, None] + overlap = einsum("ij,ik->jk", vectors_left.conj(), vectors_right) + return np.allclose(overlap, np.eye(overlap.shape[0]), atol=1e-10, rtol=0.0) + + +@cache_by_id +def orthonormalise(vectors: Array, transpose: bool = False) -> Array: + """Orthonormalise a set of vectors. + + Args: + vectors: The set of vectors to be orthonormalised. + transpose: Whether to transpose the vectors before and after orthonormalisation. + + Returns: + The orthonormalised set of vectors. + """ + if transpose: + vectors = vectors.T.conj() + + overlap = vectors.T.conj() @ vectors + orth = matrix_power(overlap, -0.5, hermitian=False) + vectors = vectors @ orth.T.conj() + + if transpose: + vectors = vectors.T.conj() + + return vectors + + +def biorthonormalise_with_overlap( + left: Array, + right: Array, + overlap: Array, + method: Literal["eig", "eig-balanced", "lu"] = BIORTH_METHOD, +) -> tuple[Array, Array]: + """Biorthonormalise two sets of vectors with a given overlap matrix. + + Args: + left: The left set of vectors. + right: The right set of vectors. + overlap: The overlap matrix to be used for biorthonormalisation. + method: The method to use for biorthonormalisation. See :func:`biorthonormalise` for + available methods. + + Returns: + The biorthonormalised left and right sets of vectors. -import numpy as np + See Also: + :func:`biorthonormalise` for details on the available methods. + """ + if method == "eig": + orth, error = matrix_power(overlap, -1, hermitian=False, return_error=True) + right = right @ orth + elif method == "eig-balanced": + orth, error = matrix_power(overlap, -0.5, hermitian=False, return_error=True) + left = left @ orth.T.conj() + right = right @ orth + elif method == "lu": + l, u = scipy.linalg.lu(overlap, permute_l=True) + try: + left = left @ np.linalg.inv(l).T.conj() + right = right @ np.linalg.inv(u) + except np.linalg.LinAlgError as e: + warnings.warn( + f"Inverse of LU decomposition failed with error: {e}. " + "Falling back to eigenvalue decomposition.", + UserWarning, + ) + return biorthonormalise_with_overlap(left, right, overlap, method="eig-balanced") + else: + raise ValueError(f"Unknown biorthonormalisation method: {method}") + + return left, right + + +@cache_by_id +def biorthonormalise( + left: Array, + right: Array, + transpose: bool = False, + method: Literal["eig", "eig-balanced", "lu"] = BIORTH_METHOD, +) -> tuple[Array, Array]: + """Biorthonormalise two sets of vectors. + + Args: + left: The left set of vectors. + right: The right set of vectors. + transpose: Whether to transpose the vectors before and after biorthonormalisation. + method: The method to use for biorthonormalisation. The ``"eig"`` method uses the + eigenvalue decomposition, the ``"eig-balanced"`` method uses the same decomposition but + applies a balanced transformation to the left- and right-hand vectors, and the ``"lu"`` + method uses the LU decomposition. + + Returns: + The biorthonormalised left and right sets of vectors. + + See Also: + :func:`biorthonormalise_with_overlap` for a more general method that allows for a custom + overlap matrix. + """ + if transpose: + left = left.T.conj() + right = right.T.conj() + + overlap = left.T.conj() @ right + left, right = biorthonormalise_with_overlap(left, right, overlap, method=method) + + if transpose: + left = left.T.conj() + right = right.T.conj() + + return left, right + + +def _sort_eigvals(eigvals: Array, eigvecs: Array, threshold: float = 1e-11) -> tuple[Array, Array]: + """Sort eigenvalues and eigenvectors. + Args: + eigvals: The eigenvalues to be sorted. + eigvecs: The eigenvectors to be sorted. + threshold: Threshold for rounding the eigenvalues to avoid numerical noise. -def matrix_power(m, power, hermitian=True, threshold=1e-10, return_error=False): + Returns: + The sorted eigenvalues and eigenvectors. + + Note: + The indirect sort attempts to sort the eigenvalues such that complex conjugate pairs are + ordered correctly, regardless of any numerical noise in the real part. This is done by + first ordering based on the rounded real and imaginary parts of the eigenvalues, and then + sorting by the true real and imaginary parts. """ - Compute the power of the matrix `m` via the eigenvalue - decomposition. + decimals = round(-np.log10(threshold)) + real_approx = np.round(eigvals.real, decimals=decimals) + imag_approx = np.round(eigvals.imag, decimals=decimals) + idx = np.lexsort((eigvals.imag, eigvals.real, imag_approx, real_approx)) + eigvals = eigvals[idx] + eigvecs = eigvecs[:, idx] + return eigvals, eigvecs + + +@cache_by_id +def eig(matrix: Array, hermitian: bool = True, overlap: Array | None = None) -> tuple[Array, Array]: + """Compute the eigenvalues and eigenvectors of a matrix. + + Args: + matrix: The matrix to be diagonalised. + hermitian: Whether the matrix is hermitian. + overlap: An optional overlap matrix to be used for the eigenvalue decomposition. + + Returns: + The eigenvalues and eigenvectors of the matrix. + """ + # Find the eigenvalues and eigenvectors + if AVOID_SCIPY_EIG and overlap is not None: + matrix = np.linalg.solve(overlap, matrix) + if AVOID_SCIPY_EIG and hermitian: + eigvals, eigvecs = np.linalg.eigh(matrix) + elif AVOID_SCIPY_EIG: + eigvals, eigvecs = np.linalg.eig(matrix) + elif hermitian: + eigvals, eigvecs = scipy.linalg.eigh(matrix, b=overlap) + else: + eigvals, eigvecs = scipy.linalg.eig(matrix, b=overlap) - Parameters - ---------- - m : numpy.ndarray (n, n) - The matrix to be raised to a power. - power : float - The power to which the matrix is to be raised. - hermitian : bool, optional - Whether the matrix is hermitian. Default value is `True`. - threshold : float, optional - Threshold for removing singularities. Default value is - `1e-10`. - return_error : bool, optional - Whether to return the error in the power. Default value is - `False`. + # See if we can remove the imaginary part of the eigenvalues + if not hermitian and np.all(eigvals.imag == 0.0): + eigvals = eigvals.real - Returns - ------- - m_pow : numpy.ndarray (n, n) - The matrix raised to the power. - error : float, optional - The error in the power. Only returned if `return_error` is - `True`. + return _sort_eigvals(eigvals, eigvecs) + + +@cache_by_id +def eig_lr( + matrix: Array, hermitian: bool = True, overlap: Array | None = None +) -> tuple[Array, tuple[Array, Array]]: + """Compute the eigenvalues and biorthogonal left- and right-hand eigenvectors of a matrix. + + Args: + matrix: The matrix to be diagonalised. + hermitian: Whether the matrix is hermitian. + overlap: An optional overlap matrix to be used for the eigenvalue decomposition. + + Returns: + The eigenvalues and biorthogonal left- and right-hand eigenvectors of the matrix. """ + # Find the eigenvalues and eigenvectors + eigvals_left: Array | None = None + if AVOID_SCIPY_EIG and hermitian: + if overlap is not None: + matrix = np.linalg.solve(overlap, matrix) + eigvals, eigvecs_right = _sort_eigvals(*np.linalg.eigh(matrix)) + eigvecs_left = eigvecs_right + elif AVOID_SCIPY_EIG: + matrix_right = matrix + matrix_left = matrix.T.conj() + if overlap is not None: + matrix_right = np.linalg.solve(overlap, matrix_right) + matrix_left = np.linalg.solve(overlap.T.conj(), matrix_left) + eigvals, eigvecs_right = _sort_eigvals(*np.linalg.eig(matrix_right)) + eigvals_left, eigvecs_left = np.linalg.eig(matrix_left) + eigvals_left, eigvecs_left = _sort_eigvals(eigvals_left.conj(), eigvecs_left) + elif hermitian: + eigvals, eigvecs_right = _sort_eigvals(*scipy.linalg.eigh(matrix, b=overlap)) + eigvecs_left = eigvecs_right + else: + eigvals_raw, eigvecs_left, eigvecs_right = scipy.linalg.eig( + matrix, + left=True, + right=True, + b=overlap, + ) + eigvals, eigvecs_right = _sort_eigvals(eigvals_raw, eigvecs_right) + eigvals, eigvecs_left = _sort_eigvals(eigvals_raw, eigvecs_left) + if not hermitian: + eigvecs_left, eigvecs_right = biorthonormalise(eigvecs_left, eigvecs_right) + + # See if we can remove the imaginary part of the eigenvalues + if not hermitian and np.all(eigvals.imag == 0.0): + eigvals = eigvals.real + + return eigvals, (eigvecs_left, eigvecs_right) + +@cache_by_id +def null_space_basis( + matrix: Array, threshold: float = 1e-11, hermitian: bool | None = None +) -> tuple[Array, Array]: + r"""Find a basis for the null space of a matrix. + + Args: + matrix: The matrix for which to find the null space. + threshold: Threshold for removing vectors to obtain the null space. + hermitian: Whether the matrix is hermitian. If ``None``, infer from the matrix. + + Returns: + The basis for the null space. + + Note: + The full vector space may not be biorthonormal. + """ + if hermitian is None: + hermitian = np.allclose(matrix, matrix.T.conj()) + + # Find the null space + null = np.eye(matrix.shape[1]) - matrix + + # Diagonalise the null space to find the basis + weights, (left, right) = eig_lr(null, hermitian=hermitian) + mask = (1 - np.abs(weights)) < threshold + left = left[:, mask] + right = right[:, mask] + + return (left, right) if hermitian else (left, left) + + +@cache_by_id +def matrix_power( + matrix: Array, + power: int | float, + hermitian: bool = True, + threshold: float = 1e-10, + return_error: bool = False, + ord: int | float = np.inf, +) -> tuple[Array, float | None]: + """Compute the power of a matrix via the eigenvalue decomposition. + + Args: + matrix: The matrix to be exponentiated. + power: The power to which the matrix is to be raised. + hermitian: Whether the matrix is hermitian. + threshold: Threshold for removing singularities. + return_error: Whether to return the error in the power. + ord: The order of the norm to be used for the error. + + Returns: + The matrix raised to the power, and the error if requested. + """ + # TODO: Check if scipy.linalg.fractional_matrix_power is better + + # Get the eigenvalues and eigenvectors -- don't need to be biorthogonal, avoid recursive calls + eigvals, right = eig(matrix, hermitian=hermitian) if hermitian: - # assert np.allclose(m, m.T.conj()) - eigvals, eigvecs = np.linalg.eigh(m) + left = right else: - eigvals, eigvecs = np.linalg.eig(m) + left = np.linalg.inv(right).T.conj() + # Get the mask for removing singularities if power < 0: - # Remove singularities mask = np.abs(eigvals) > threshold else: mask = np.ones_like(eigvals, dtype=bool) - if hermitian and not np.iscomplexobj(m): + # Get the mask for removing negative eigenvalues + if hermitian and not np.iscomplexobj(matrix): if np.abs(power) < 1: - mask = np.logical_and(mask, eigvals > 0) - eigvecs_right = eigvecs.T.conj() - elif hermitian and np.iscomplexobj(m): - power = power + 0.0j - eigvecs_right = eigvecs.T.conj() + mask &= eigvals > 0 else: - power = power + 0.0j - eigvecs_right = np.linalg.inv(eigvecs) + if np.any(eigvals < 0): + power: complex = power + 0.0j # type: ignore[no-redef] - left = eigvecs[:, mask] * eigvals[mask][None] ** power - right = eigvecs_right[mask] - m_pow = np.dot(left, right) + # Contract the eigenvalues and eigenvectors + matrix_power: Array = (right[:, mask] * eigvals[mask][None] ** power) @ left[:, mask].T.conj() + # Get the error if requested + error: float | None = None if return_error: - left = eigvecs[:, ~mask] * eigvals[~mask][None] - right = eigvecs_right[~mask] - m_res = np.dot(left, right) - error = np.linalg.norm(np.linalg.norm(m_res)) - return m_pow, error + null = (right[:, ~mask] * eigvals[~mask][None]) @ left[:, ~mask].T.conj() + if null.size == 0: + error = 0.0 + else: + error = cast(float, np.linalg.norm(null, ord=ord)) + + # See if we can remove the imaginary part of the matrix power + if np.iscomplexobj(matrix_power) and np.all(np.isclose(matrix_power.imag, 0.0)): + matrix_power = matrix_power.real + + return matrix_power, error + + +def hermi_sum(matrix: Array) -> Array: + """Return the sum of a matrix with its Hermitian conjugate. + + Args: + matrix: The matrix to be summed with its hermitian conjugate. + + Returns: + The sum of the matrix with its hermitian conjugate. + """ + return matrix + matrix.T.conj() + + +def scaled_error(matrix1: Array, matrix2: Array, ord: int | float = np.inf) -> float: + """Return the scaled error between two matrices. + + Args: + matrix1: The first matrix. + matrix2: The second matrix. + ord: The order of the norm to be used for the error. + + Returns: + The scaled error between the two matrices. + """ + matrix1 = np.atleast_1d(matrix1 / max(np.max(np.abs(matrix1)), 1)) + matrix2 = np.atleast_1d(matrix2 / max(np.max(np.abs(matrix2)), 1)) + return cast(float, np.linalg.norm(matrix1 - matrix2, ord=ord)) + + +def as_trace(matrix: Array, ndim: int, axis1: int = -2, axis2: int = -1) -> Array: + """Return the trace of a matrix, unless it has been passed as a trace. + + Args: + matrix: The matrix to be traced. + ndim: The number of dimensions of the matrix before the trace. + axis1: The first axis of the trace. + axis2: The second axis of the trace. + + Returns: + The trace of the matrix. + """ + if matrix.ndim == ndim: + return matrix + elif matrix.ndim > ndim: + return np.trace(matrix, axis1=axis1, axis2=axis2) else: - return m_pow + raise ValueError(f"Matrix has invalid shape {matrix.shape} for trace.") + + +def as_diagonal(matrix: Array, ndim: int) -> Array: + """Return the diagonal of a matrix, unless it has been passed as a diagonal. + Args: + matrix: The matrix to be diagonalised. + ndim: The number of dimensions of the matrix before the diagonal. -def hermi_sum(m): + Returns: + The diagonal of the matrix. """ - Return m + m^† + if matrix.ndim == ndim: + return matrix + elif matrix.ndim > ndim: + return np.diagonal(matrix, axis1=-2, axis2=-1) + else: + raise ValueError(f"Matrix has invalid shape {matrix.shape} for diagonal.") - Parameters - ---------- - m : numpy.ndarray (n, n) - The matrix to be summed with its hermitian conjugate. - Returns - ------- - m_sum : numpy.ndarray (n, n) - The sum of the matrix with its hermitian conjugate. +def unit_vector(size: int, index: int, dtype: str = "float64") -> Array: + """Return a unit vector of size ``size`` with a 1 at index ``index``. + + Args: + size: The size of the vector. + index: The index of the vector. + dtype: The data type of the vector. + + Returns: + The unit vector. """ + return np.eye(1, size, k=index, dtype=dtype).ravel() + + +def concatenate_paired_vectors(vectors: list[Array], size: int) -> Array: + r"""Concatenate vectors that are partitioned into two spaces, the first of which is common. + + Args: + vectors: The vectors to be concatenated. + size: The size of the first space. - return m + m.T.conj() + Returns: + The concatenated vectors. + Note: + The concatenation is -def scaled_error(a, b): + .. math:: + \begin{pmatrix} + p_1 & p_2 & \cdots & p_n \\ + a_1 & & & \\ + & a_2 & & \\ + & & \ddots & \\ + & & & a_n \\ + \end{pmatrix} + = + \begin{pmatrix} p_1 \\ a_1 \end{pmatrix} + + \begin{pmatrix} p_2 \\ a_2 \end{pmatrix} + + \cdots + + \begin{pmatrix} p_n \\ a_n \end{pmatrix} + + where :math:`p_i` are the vectors in the first space and :math:`a_i` are the vectors in the + second space. + + This is useful for combining couplings between a common physical space and a set of + auxiliary degrees of freedom. """ - Return the scaled error between two matrices. + space1 = slice(0, size) + space2 = slice(size, None) + vectors1 = np.concatenate([vector[space1] for vector in vectors], axis=1) + vectors2 = block_diag(*[vector[space2] for vector in vectors]) + return np.concatenate([vectors1, vectors2], axis=0) - Parameters - ---------- - a : numpy.ndarray (n, n) - The first matrix. - b : numpy.ndarray (n, n) - The second matrix. - Returns - ------- - error : float - The scaled error between the two matrices. +def unpack_vectors(vector: Array) -> tuple[Array, Array]: + """Unpack a block vector in the :mod:`dyson` convention. + + Args: + vector: The vector to be unpacked. The vector should either be a 2D array ``(n, m)`` or a 3D + array ``(2, n, m)``. The latter case is non-Hermitian. + + Returns: + Left- and right-hand vectors. """ + if vector.ndim == 2: + return vector, vector + elif vector.ndim == 3: + return vector[0], vector[1] + raise ValueError( + f"Vector has invalid shape {vector.shape} for unpacking. Must be 2D or 3D array." + ) - a = a / max(np.max(np.abs(a)), 1) - b = b / max(np.max(np.abs(b)), 1) - return np.linalg.norm(a - b) +def block_diag(*arrays: Array) -> Array: + """Return a block diagonal matrix from a list of arrays. + Args: + arrays: The arrays to be combined into a block diagonal matrix. -def remove_unphysical(eigvecs, nphys, eigvals=None, tol=1e-8): + Returns: + The block diagonal matrix. """ - Remove eigenvectors with a small physical component. + if not all(array.ndim == 2 for array in arrays): + raise ValueError("All arrays must be 2D.") + rows = [array.shape[0] for array in arrays] + cols = [array.shape[1] for array in arrays] + arrays_full = [[np.zeros((row, col)) for col in cols] for row in rows] + for i, array in enumerate(arrays): + arrays_full[i][i] = array + return np.block(arrays_full) + + +def set_subspace(vectors: Array, subspace: Array) -> Array: + """Set the subspace of a set of vectors. - Parameters - ---------- - eigvecs : numpy.ndarray or tuple of numpy.ndarray - Eigenvectors. If a tuple, the first element is the left - eigenvectors and the second element is the right - eigenvectors. - nphys : int - Number of physical orbitals. - eigvals : numpy.ndarray, optional - Eigenvalues. Default value is `None`. - tol : float, optional - Threshold for removing eigenvectors. Default value is - `1e-8`. + Args: + vectors: The vectors to be set. + subspace: The subspace to be applied to the vectors. - Returns - ------- - eigvals : numpy.ndarray, optional - Eigenvalues. Only returned if `eigvals` is not `None`. - eigvecs : numpy.ndarray or tuple of numpy.ndarray - Eigenvectors. If a tuple, the first element is the left - eigenvectors and the second element is the right - eigenvectors. + Returns: + The vectors with the subspace applied. + + Note: + This operation is equivalent to applying ``vectors[: n] = subspace`` where ``n`` is the size + of both dimensions in the subspace. """ + size = subspace.shape[0] + return np.concatenate([subspace, vectors[size:]], axis=0) - if isinstance(eigvecs, tuple): - eigvecs_l, eigvecs_r = eigvecs - else: - eigvecs_l = eigvecs_r = eigvecs - mask = np.abs(np.sum(eigvecs_l[:nphys] * eigvecs_r.conj()[:nphys], axis=0)) > tol +def rotate_subspace(vectors: Array, rotation: Array) -> Array: + """Rotate the subspace of a set of vectors. - if isinstance(eigvecs, tuple): - eigvecs_out = (eigvecs_l[:, mask], eigvecs_r[:, mask]) - else: - eigvecs_out = eigvecs[:, mask] + Args: + vectors: The vectors to be rotated. + rotation: The rotation matrix to be applied to the vectors. - if eigvals is not None: - return eigvals[mask], eigvecs_out - else: - return eigvecs_out + Returns: + The rotated vectors. + + Note: + This operation is equivalent to applying ``vectors[: n] = rotation @ vectors[: n]`` where + ``n`` is the size of both dimensions in the rotation matrix. + """ + if rotation.shape[0] != rotation.shape[1]: + raise ValueError(f"Rotation matrix must be square, got shape {rotation.shape}.") + size = rotation.shape[0] + subspace = rotation @ vectors[:size] + return set_subspace(vectors, subspace) diff --git a/dyson/util/logging.py b/dyson/util/logging.py deleted file mode 100644 index b552d0a..0000000 --- a/dyson/util/logging.py +++ /dev/null @@ -1,166 +0,0 @@ -""" -Logging utilities. -""" - -import warnings - -import numpy as np - - -def format_value(val, prec=8): - """ - Format a float or complex value. - - Parameters - ---------- - val : float or complex - The value to be formatted. - prec : int, optional - The number of decimal places to use. Default value is `8`. - - Returns - ------- - out : str - String representation of the value. - """ - - if not np.iscomplexobj(val): - return "%.*f" % (prec, val) - else: - op = "+" if val.imag >= 0 else "-" - return "%.*f%s%.*fj" % (prec, val.real, op, prec, np.abs(val.imag)) - - -def print_eigenvalues(eigvals, nroots=5, abs_sort=True, header=True): - """ - Return a string summarising some eigenvalues. - - Parameters - ---------- - eigvals : numpy.ndarray - The eigenvalues. - nroots : int, optional - The number of eigenvalues to print. Default value is `5`. - abs_sort : bool, optional - Whether to sort the eigenvalues by absolute value. Default - value is `True`. - header : bool, optional - Whether to print a header. Default value is `True`. - - Returns - ------- - out : str - String summarising the eigenvalues. - """ - - lines = ["-" * 30] - if header: - lines += ["{:^30s}".format("Eigenvalue summary")] - lines += ["-" * 30] - lines += [ - "%4s %25s" % ("Root", "Value"), - "%4s %25s" % ("-" * 4, "-" * 25), - ] - - inds = np.argsort(np.abs(eigvals.real)) if abs_sort else np.argsort(eigvals.real) - for i in inds[: min(nroots, len(eigvals))]: - lines.append("%4d %25s" % (i, format_value(eigvals[i]))) - - if nroots < len(eigvals): - lines.append(" ...") - - lines.append("-" * 30) - - return "\n".join(lines) - - -def print_dyson_orbitals( - eigvals, eigvecs, nphys, nroots=5, abs_sort=True, phys_threshold=1e-8, header=True -): - """ - Returns a string summarising the projection of some eigenfunctions - into the physical space, resulting in Dyson orbitals. - - Parameters - ---------- - eigvals : numpy.ndarray - The eigenvalues. - eigvecs : numpy.ndarray - The eigenvectors. - nphys : int - The number of physical orbitals. - nroots : int, optional - The number of eigenvalues to print. Default value is `5`. - abs_sort : bool, optional - Whether to sort the eigenvalues by absolute value. Default - value is `True`. - phys_threshold : float, optional - The threshold for the projection of the eigenvectors into the - physical space. Default value is `1e-8`. - header : bool, optional - Whether to print a header. Default value is `True`. - - Returns - ------- - out : str - String summarising the dyson obritals. - """ - - lines = ["-" * 98] - if header: - lines += ["{:^98s}".format("Dyson orbital summary")] - lines += ["-" * 98] - lines += [ - "{:>4s} {:^25s} {:^33s} {:^33s}".format("", "", "Weight", ""), - "{:>4s} {:^25s} {:^33s} {:^33s}".format( - "Orb", - "Energy", - "-" * 33, - "Dominant physical contributions", - ), - "{:>4s} {:^25s} {:>16s} {:>16s} {:>16s}".format("", "", "Physical", "Auxiliary", ""), - "{:>4s} {:^25s} {:>16s} {:>16s} {:>16s}".format( - "-" * 4, - "-" * 25, - "-" * 16, - "-" * 16, - "-" * 33, - ), - ] - - if isinstance(eigvecs, tuple): - eigvecs_l, eigvecs_r = eigvecs - else: - eigvecs_l = eigvecs_r = eigvecs - - mask = np.sum(np.abs(eigvecs_l * eigvecs_r.conj()), axis=0) > phys_threshold - inds = np.arange(eigvals.size)[mask] - inds = inds[ - np.argsort(np.abs(eigvals[inds].real)) if abs_sort else np.argsort(eigvals[inds].real) - ] - for i in inds[: min(nroots, len(eigvals))]: - v = np.abs(eigvecs_l[:, i] * eigvecs_r[:, i].conj()) - phys = np.sum(v[:nphys]) - aux = np.sum(v[nphys:]) - chars = [] - for j in np.argsort(v[:nphys]): - if v[j] > 0.2: - chars.append("%d (%.2f)" % (j, v[j])) - chars = ", ".join(chars) - lines.append( - "%4d %25s %16.3g %16.3g %33s" - % ( - i, - format_value(eigvals[i]), - phys, - aux, - chars, - ) - ) - - if nroots < len(inds): - lines.append(" ...") - - lines += ["-" * 98] - - return "\n".join(lines) diff --git a/dyson/util/misc.py b/dyson/util/misc.py index 95a5640..a3c8782 100644 --- a/dyson/util/misc.py +++ b/dyson/util/misc.py @@ -1,38 +1,96 @@ -""" -Miscellaneous utilities. -""" +"""Miscellaneous utility functions.""" -import inspect +from __future__ import annotations -import numpy as np +import functools +import warnings +import weakref +from contextlib import contextmanager +from typing import TYPE_CHECKING +from pyscf import gto, scf -def cache(function): +if TYPE_CHECKING: + from typing import Any, Callable, Iterator + from warnings import WarningMessage + + +@contextmanager +def catch_warnings(warning_type: type[Warning] = Warning) -> Iterator[list[WarningMessage]]: + """Context manager to catch warnings. + + Returns: + A list of caught warnings. """ - Caches return values according to positional and keyword arguments - in the `_cache` property of an object. + # Remove any user filters + user_filters = warnings.filters[:] + warnings.simplefilter("always", warning_type) + + # Catch warnings + with warnings.catch_warnings(record=True) as caught_warnings: + yield caught_warnings + + # Restore user filters + warnings.filters[:] = user_filters # type: ignore[index] + + +def cache_by_id(func: Callable) -> Callable: + """Decorator to cache function results based on the ``id`` of the arguments. + + Args: + func: The function to cache. + + Returns: + A wrapper function that caches results based on the id of the arguments. """ + cache: dict[tuple[tuple[int, ...], tuple[tuple[str, int], ...]], Any] = {} + watchers: dict[tuple[tuple[int, ...], tuple[tuple[str, int], ...]], list[weakref.ref]] = {} - def wrapper(obj, *args, **kwargs): - if (function.__name__, args, tuple(kwargs.items())) in obj._cache: - return obj._cache[function.__name__, args, tuple(kwargs.items())] - else: - out = function(obj, *args, **kwargs) - obj._cache[function.__name__, args, tuple(kwargs.items())] = out - return out + def _remove(key: tuple[tuple[int, ...], tuple[tuple[str, int], ...]]) -> None: + """Remove an entry from the cache.""" + cache.pop(key, None) + watchers.pop(key, None) + + @functools.wraps(func) + def wrapper(*args: Any, **kwargs: Any) -> Any: + """Cache results based on the id of the arguments.""" + key_args = tuple(id(arg) for arg in args) + key_kwargs = tuple(sorted((k, id(v)) for k, v in kwargs.items())) + key = (key_args, key_kwargs) + if key in cache: + return cache[key] + + result = func(*args, **kwargs) + cache[key] = result + + refs: list[weakref.ref] = [] + for obj in [*args, *kwargs.values()]: + try: + refs.append(weakref.ref(obj, lambda _ref, k=key: _remove(k))) # type: ignore[misc] + except TypeError: + continue + if refs: + watchers[key] = refs + + return result return wrapper -def inherit_docstrings(cls): - """ - Inherit docstring from superclass. - """ +def get_mean_field(atom: str, basis: str, charge: int = 0, spin: int = 0) -> scf.RHF: + """Get a mean-field object for a given system. - for name, func in inspect.getmembers(cls, inspect.isfunction): - if not func.__doc__: - for parent in cls.__mro__[1:]: - if hasattr(parent, name): - func.__doc__ = getattr(parent, name).__doc__ + Intended as a convenience function for examples. - return cls + Args: + atom: The atomic symbol of the system. + basis: The basis set to use. + charge: The total charge of the system. + spin: The total spin of the system. + + Returns: + A mean-field object for the system. + """ + mol = gto.M(atom=atom, basis=basis, charge=charge, spin=spin, verbose=0) + mf = scf.RHF(mol).run() + return mf diff --git a/dyson/util/moments.py b/dyson/util/moments.py index c2e0fd7..3e91a71 100644 --- a/dyson/util/moments.py +++ b/dyson/util/moments.py @@ -1,83 +1,129 @@ -""" -Moment utilities. -""" +"""Moment utilities.""" -import numpy as np +from __future__ import annotations +import warnings +from typing import TYPE_CHECKING -def se_moments_to_gf_moments(se_static, se_moments): - """ - Convert moments of the self-energy to those of the Green's - function. The first m moments of the self-energy, along with - the static part, are sufficient to define the first m+2 moments - of the Green's function. See Eqns 2.103-105 of Backhouse's thesis. - - Parameters - ---------- - se_static : numpy.ndarray (n, n) - Static part of the self-energy. - se_moments : numpy.ndarray (m, n, n) - Moments of the self-energy. - - Returns - ------- - gf_moments : numpy.ndarray (m+2, n, n) +from dyson import numpy as np +from dyson.util.linalg import einsum, matrix_power + +if TYPE_CHECKING: + from dyson.typing import Array + + +def se_moments_to_gf_moments( + static: Array, se_moments: Array, overlap: Array | None = None, check_error: bool = True +) -> Array: + """Convert moments of the self-energy to those of the Green's function. + + Args: + static: Static part of the self-energy. + se_moments: Moments of the self-energy. + overlap: The overlap matrix (zeroth moment of the Green's function). If ``None``, the zeroth + moment of the Green's function is assumed to be the identity matrix. + check_error: Whether to check the errors in the orthogonalisation of the moments. + + Returns: Moments of the Green's function. - """ + Notes: + The first :math:`m` moments of the self-energy, along with the static part, are sufficient + to define the first :math:`m+2` moments of the Green's function. + """ nmom, nphys, _ = se_moments.shape - gf_moments = np.zeros((nmom + 2, nphys, nphys), dtype=se_moments.dtype) + # Orthogonalise the moments + if overlap is not None: + hermitian = np.allclose(overlap, overlap.T.conj()) + orth, error_orth = matrix_power( + overlap, -0.5, hermitian=hermitian, return_error=check_error + ) + unorth, error_unorth = matrix_power( + overlap, 0.5, hermitian=hermitian, return_error=check_error + ) + if check_error: + assert error_orth is not None and error_unorth is not None + error = max(error_orth, error_unorth) + if error > 1e-10: + warnings.warn( + "Space contributing non-zero weight to the zeroth moments " + f"({max(error_orth, error_unorth)}) was removed during moment conversion.", + UserWarning, + 2, + ) + static = orth @ static @ orth + se_moments = einsum("npq,ip,qj->nij", se_moments, orth, orth) + + # Get the powers of the static part + powers = [np.eye(static.shape[-1], dtype=static.dtype), static] + for i in range(2, nmom + 2): + powers.append(powers[i - 1] @ static) + gf_moments = np.zeros( + (nmom + 2, nphys, nphys), dtype=np.result_type(se_moments.dtype, powers[0].dtype) + ) + + # Perform the recursion for i in range(nmom + 2): - gf_moments[i] += np.linalg.matrix_power(se_static, i) + gf_moments[i] += powers[i] for n in range(i - 1): for m in range(i - n - 1): k = i - n - m - 2 - gf_moments[i] += np.linalg.multi_dot( - ( - np.linalg.matrix_power(se_static, n), - se_moments[m], - gf_moments[k], - ) - ) + gf_moments[i] += powers[n] @ se_moments[m] @ gf_moments[k] + + # Unorthogonalise the moments + if overlap is not None: + gf_moments = einsum("npq,ip,qj->nij", gf_moments, unorth, unorth) return gf_moments -def gf_moments_to_se_moments(gf_moments): - """ - Convert moments of the Green's function to those of the - self-energy. The first m+2 moments of the Green's function - are sufficient to define the first m moments of the self-energy, - along with the static part. See Eqns 2.103-105 of Backhouse's - thesis. - - Parameters - ---------- - gf_moments : numpy.ndarray (m+2, n, n) - Moments of the Green's function. +def gf_moments_to_se_moments(gf_moments: Array, check_error: bool = True) -> tuple[Array, Array]: + """Convert moments of the Green's function to those of the self-energy. - Returns - ------- - se_static : numpy.ndarray (n, n) - Static part of the self-energy. - se_moments : numpy.ndarray (m, n, n) - Moments of the self-energy. - """ + Args: + gf_moments: Moments of the Green's function. + check_error: Whether to check the errors in the orthogonalisation of the moments. - nmom, nphys, _ = gf_moments.shape + Returns: + static: Static part of the self-energy. + moments: Moments of the self-energy. + Notes: + The first :math:`m+2` moments of the Green's function are sufficient to define the first + :math:`m` moments of the self-energy, along with the static part. + """ + nmom, nphys, _ = gf_moments.shape if nmom < 2: raise ValueError( - "At least 2 moments of the Green's function are required to " - "find those of the self-energy." + "Need at least 2 moments of the Green's function to compute those of the self-energy." ) - if not np.allclose(gf_moments[0], np.eye(nphys)): - raise ValueError("The first moment of the Green's function must be the identity.") + # Orthogonalise the moments + ident = np.allclose(gf_moments[0], np.eye(nphys)) + if not ident: + hermitian = np.allclose(gf_moments[0], gf_moments[0].T.conj()) + orth, error_orth = matrix_power( + gf_moments[0], -0.5, hermitian=hermitian, return_error=check_error + ) + unorth, error_unorth = matrix_power( + gf_moments[0], 0.5, hermitian=hermitian, return_error=check_error + ) + if check_error: + assert error_orth is not None and error_unorth is not None + error = max(error_orth, error_unorth) + if error > 1e-10: + warnings.warn( + "Space contributing non-zero weight to the zeroth moments " + f"({max(error_orth, error_unorth)}) was removed during moment conversion.", + UserWarning, + 2, + ) + gf_moments = einsum("npq,ip,qj->nij", gf_moments, orth, orth) - se_moments = np.zeros((nmom - 2, nphys, nphys), dtype=gf_moments.dtype) + # Get the static part and the moments of the self-energy se_static = gf_moments[1] + se_moments = np.zeros((nmom - 2, nphys, nphys), dtype=gf_moments.dtype) # Invert the recurrence relations: # @@ -88,145 +134,90 @@ def gf_moments_to_se_moments(gf_moments): # with the constraint that m != n. This case is F^{0} \Sigma_{n} G_{0} # which is equal to the desired LHS. + # Get the powers of the static part + powers = [np.eye(nphys, dtype=gf_moments.dtype), se_static] + for i in range(2, nmom): + powers.append(powers[i - 1] @ se_static) + + # Perform the recursion for i in range(nmom - 2): - se_moments[i] = gf_moments[i + 2].copy() - se_moments[i] -= np.linalg.matrix_power(se_static, i + 2) + se_moments[i] = gf_moments[i + 2] - powers[i + 2] for l in range(i + 1): for m in range(i + 1 - l): k = i - l - m if m != i: - se_moments[i] -= np.linalg.multi_dot( - ( - np.linalg.matrix_power(se_static, l), - se_moments[m], - gf_moments[k], - ) - ) + se_moments[i] -= powers[l] @ se_moments[m] @ gf_moments[k] + + # Unorthogonalise the moments + if not ident: + se_static = unorth @ se_static @ unorth + se_moments = einsum("npq,ip,qj->nij", se_moments, unorth, unorth) return se_static, se_moments -def build_block_tridiagonal(on_diagonal, off_diagonal_upper, off_diagonal_lower=None): - """ - Build a block tridiagonal matrix. - - Parameters - ---------- - on_diagonal : numpy.ndarray (m+1, n, n) - On-diagonal blocks. - off_diagonal_upper : numpy.ndarray (m, n, n) - Off-diagonal blocks for the upper half of the matrix. - off_diagonal_lower : numpy.ndarray (m, n, n), optional - Off-diagonal blocks for the lower half of the matrix. If - `None`, use the transpose of `off_diagonal_upper`. - """ +def build_block_tridiagonal( + on_diagonal: list[Array], + off_diagonal_upper: list[Array], + off_diagonal_lower: list[Array] | None = None, +) -> Array: + """Build a block tridiagonal matrix. - zero = np.zeros_like(on_diagonal[0]) - - if off_diagonal_lower is None: - off_diagonal_lower = [m.T.conj() for m in off_diagonal_upper] - - m = np.block( - [ - [ - ( - on_diagonal[i] - if i == j - else ( - off_diagonal_upper[j] - if j == i - 1 - else off_diagonal_lower[i] if i == j - 1 else zero - ) - ) - for j in range(len(on_diagonal)) - ] - for i in range(len(on_diagonal)) - ] - ) - - return m + Args: + on_diagonal: On-diagonal blocks. + off_diagonal_upper: Off-diagonal blocks for the upper half of the matrix. + off_diagonal_lower: Off-diagonal blocks for the lower half of the matrix. If + ``None``, use the transpose of ``off_diagonal_upper``. + Returns: + A block tridiagonal matrix with the given blocks. -def matvec_to_greens_function(matvec, nmom, bra, ket=None): + Notes: + The number of on-diagonal blocks should be one greater than the number of off-diagonal + blocks. """ - Build a set of moments using the matrix-vector product for a - given Hamiltonian and a bra and ket vector. - - Parameters - ---------- - matvec : callable - Matrix-vector product function, takes a vector as input. - nmom : int - Number of moments to compute. - bra : numpy.ndarray (n, m) - Bra vector. - ket : numpy.ndarray (n, m), optional - Ket vector, if `None` then use the bra. - """ - - nphys, nconf = bra.shape - moments = np.zeros((nmom, nphys, nphys)) - - if ket is None: - ket = bra - ket = ket.copy() - - for n in range(nmom): - moments[n] = np.dot(bra, ket.T.conj()) - if n != (nmom - 1): - for i in range(nphys): - ket[i] = matvec(ket[i]) + on_diagonal = [np.atleast_2d(matrix) for matrix in on_diagonal] + off_diagonal_upper = [np.atleast_2d(matrix) for matrix in off_diagonal_upper] + if off_diagonal_lower is not None: + off_diagonal_lower = [np.atleast_2d(matrix) for matrix in off_diagonal_lower] + if len(on_diagonal) == 0: + return np.zeros((0, 0)) + zero = np.zeros_like(on_diagonal[0]) + if off_diagonal_lower is None: + off_diagonal_lower = [matrix.T.conj() for matrix in off_diagonal_upper] + + def _block(i: int, j: int) -> Array: + """Return the block at position (i, j).""" + if i == j: + return on_diagonal[i] + elif j == i - 1: + return off_diagonal_upper[j] + elif i == j - 1: + return off_diagonal_lower[i] + return zero + + # Construct the block tridiagonal matrix + matrix = np.block( + [[_block(i, j) for j in range(len(on_diagonal))] for i in range(len(on_diagonal))] + ) - return moments + return matrix -matvec_to_greens_function_monomial = matvec_to_greens_function +def get_chebyshev_scaling_parameters( + min_value: float, max_value: float, epsilon: float = 1e-3 +) -> tuple[float, float]: + """Get the Chebyshev scaling parameters. + Args: + min_value: Minimum value of the range. + max_value: Maximum value of the range. + epsilon: Small value to avoid division by zero. -def matvec_to_greens_function_chebyshev(matvec, nmom, scale_factors, bra, ket=None): - """ - Build a set of Chebyshev moments using the matrix-vector product - for a given Hamiltonian and a bra and ket vector. - - Parameters - ---------- - matvec : callable - Matrix-vector product function, takes a vector as input. - nmom : int - Number of moments to compute. - scale_factors : tuple of int - Factors to scale the Hamiltonian as `(H - b) / a`, in order - to keep the spectrum within [-1, 1]. These are typically - defined as - `a = (emax - emin) / (2 - eps)` - `b = (emax + emin) / 2` - where `emin` and `emax` are the minimum and maximum eigenvalues - of H, and `eps` is a small number. - bra : numpy.ndarray (n, m) - Bra vector. - ket : numpy.ndarray (n, m), optional - Ket vector, if `None` then use the bra. + Returns: + A tuple containing the scaling factor and the shift. """ - - nphys, nconf = bra.shape - moments = np.zeros((nmom, nphys, nphys)) - a, b = scale_factors - - if ket is None: - ket = bra - - ket0 = ket.copy() - ket1 = np.zeros_like(ket0) - for i in range(nphys): - ket1[i] = (matvec(ket0[i]) - b * ket0[i]) / a - - moments[0] = np.dot(bra, ket0.T.conj()) - - for n in range(1, nmom): - moments[n] = np.dot(bra, ket1.T.conj()) - if n != (nmom - 1): - for i in range(nphys): - ket2i = 2.0 * (matvec(ket1[i]) - b * ket1[i]) / a - ket0[i] - ket0[i], ket1[i] = ket1[i], ket2i - - return moments + return ( + (max_value - min_value) / (2.0 - epsilon), + (max_value + min_value) / 2.0, + ) diff --git a/dyson/util/spectra.py b/dyson/util/spectra.py deleted file mode 100644 index 8330b4d..0000000 --- a/dyson/util/spectra.py +++ /dev/null @@ -1,54 +0,0 @@ -""" -Spectral function utilities. -""" - -import numpy as np -from pyscf import lib - - -def build_spectral_function(energy, coupling, grid, eta=1e-1, trace=True, imag=True): - """ - Build a spectral function. - - Parameters - ---------- - energy : numpy.ndarray - Energies of the states. - coupling : numpy.ndarray or tuple of numpy.ndarray - Coupling of the states to the spectral function. If a tuple - is given, the first element is the left coupling and the - second element is the right coupling. - grid : numpy.ndarray - Grid on which to evaluate the spectral function. - eta : float, optional - Broadening parameter. Default value is `1e-1`. - trace : bool, optional - Whether to trace over the spectral function before returning. - Default value is `True`. - imag : bool, optional - Whether to return only the imaginary part of the spectral - function. Default value is `True`. - - Returns - ------- - sf : numpy.ndarray - Spectral function. - """ - - if isinstance(coupling, tuple): - coupling_l, coupling_r = coupling - else: - coupling_l = coupling_r = coupling - - if not trace: - subscript = "pk,qk,wk->wpq" - else: - subscript = "pk,pk,wk->w" - - denom = 1.0 / (grid[:, None] - energy[None] + 1.0j * eta) - sf = -lib.einsum(subscript, coupling_l, coupling_r.conj(), denom) / np.pi - - if imag: - sf = sf.imag - - return sf diff --git a/examples/00-mblse.py b/examples/00-mblse.py deleted file mode 100644 index 1322ded..0000000 --- a/examples/00-mblse.py +++ /dev/null @@ -1,37 +0,0 @@ -""" -Example of the moment block Lanczos recursion for moments of the -self-energy (MBLSE) solver. -""" - -import numpy as np -import matplotlib.pyplot as plt -from pyscf import gto, scf, agf2, lib -from dyson import MBLSE, util - -niter = 1 -grid = np.linspace(-40, 20, 1024) - -# Define a self-energy using PySCF -mol = gto.M(atom="O 0 0 0; O 0 0 1", basis="6-31g", verbose=0) -mf = scf.RHF(mol).run() -se_static = np.diag(mf.mo_energy) -se = agf2.AGF2(mf, nmom=(None, None)).build_se() -se_moms = se.moment(range(2*niter+2)) - -# Use the solver to get the spectral function -solver = MBLSE(se_static, se_moms) -solver.kernel() -e, v = solver.get_dyson_orbitals() -sf = util.build_spectral_function(e, v, grid, eta=1.0) - -# Get a reference spectral function for comparison -gf = se.get_greens_function(se_static) -sf_ref = util.build_spectral_function(gf.energy, gf.coupling, grid, eta=1.0) - -# Plot the results -plt.plot(grid, sf_ref, "C0-", label="Reference") -plt.plot(grid, sf, "C1-", label="MBLSE") -plt.legend() -plt.xlabel("Frequency (Ha)") -plt.ylabel("Spectral function") -plt.show() diff --git a/examples/01-mblgf.py b/examples/01-mblgf.py deleted file mode 100644 index 112e0a2..0000000 --- a/examples/01-mblgf.py +++ /dev/null @@ -1,37 +0,0 @@ -""" -Example of the moment block Lanczos recursion for moments of the -Green's function (MBLGF) solver. -""" - -import numpy as np -import matplotlib.pyplot as plt -from pyscf import gto, scf, agf2, lib -from dyson import MBLGF, util - -niter = 1 -grid = np.linspace(-40, 20, 1024) - -# Define a Green's function using PySCF -mol = gto.M(atom="O 0 0 0; O 0 0 1", basis="6-31g", verbose=0) -mf = scf.RHF(mol).run() -se_static = np.diag(mf.mo_energy) -se = agf2.AGF2(mf, nmom=(None, None)).build_se() -gf = se.get_greens_function(se_static) -gf_moms = gf.moment(range(2*niter+2)) - -# Use the solver to get the spectral function -solver = MBLGF(gf_moms) -solver.kernel() -e, v = solver.get_dyson_orbitals() -sf = util.build_spectral_function(e, v, grid, eta=1.0) - -# Get a reference spectral function for comparison -sf_ref = util.build_spectral_function(gf.energy, gf.coupling, grid, eta=1.0) - -# Plot the results -plt.plot(grid, sf_ref, "C0-", label="Reference") -plt.plot(grid, sf, "C1-", label="MBLGF") -plt.legend() -plt.xlabel("Frequency (Ha)") -plt.ylabel("Spectral function") -plt.show() diff --git a/examples/02-kpmgf.py b/examples/02-kpmgf.py deleted file mode 100644 index a8e6b7c..0000000 --- a/examples/02-kpmgf.py +++ /dev/null @@ -1,55 +0,0 @@ -""" -Example of the Green's function moment kernel polynomial method -(KMPGF) solver, leveraging a Chebyshev moment representation. -""" - -import numpy as np -import matplotlib.pyplot as plt -from pyscf import gto, scf, agf2, lib -from dyson import KPMGF, util - -ncheb = 10 # Number of Chebyshev moments -kernel_type = "lorentz" # Kernel method - -# Define a self-energy using PySCF -mol = gto.M(atom="O 0 0 0; O 0 0 1", basis="6-31g", verbose=0) -mf = scf.RHF(mol).run() -se_static = np.diag(mf.mo_energy) -se = agf2.AGF2(mf, nmom=(None, None)).build_se() - -# Found the bounds of the self-energy - in practice this should be -# done using a lower-scaling solver. -gf = se.get_greens_function(se_static) -emin = gf.energy.min() -emax = gf.energy.max() -grid = np.linspace(emin, emax, 1024) - -# Scale the energies of the Green's function -a = (emax - emin) / (2.0 - 1e-2) -b = (emax + emin) / 2.0 -energy_scaled = (gf.energy - b) / a - -# Compute the Chebyshev moments -c = np.zeros((ncheb, mol.nao, energy_scaled.size)) -c[0] = gf.coupling -c[1] = gf.coupling * energy_scaled -for i in range(2, ncheb): - c[i] = 2.0 * c[i-1] * energy_scaled - c[i-2] -moments = lib.einsum("qx,npx->npq", gf.coupling, c) - -# Use the solver to get the spectral function -solver = KPMGF(moments, grid, (a, b), kernel_type=kernel_type) -sf = solver.kernel() - -# Get a reference spectral function for comparison -sf_ref = util.build_spectral_function(gf.energy, gf.coupling, grid, eta=1.0) - -# Plot the results -plt.plot(grid, sf_ref, "C0-", label="Reference") -ylim = plt.ylim() -plt.plot(grid, sf, "C1-", label="KPMGF") -plt.legend() -plt.xlabel("Frequency (Ha)") -plt.ylabel("Spectral function") -plt.ylim(ylim) -plt.show() diff --git a/examples/03-cpgf.py b/examples/03-cpgf.py deleted file mode 100644 index 8dd439c..0000000 --- a/examples/03-cpgf.py +++ /dev/null @@ -1,55 +0,0 @@ -""" -Example of the Chebyshev polynomial Green's function method -(CPGF) solver, which is similar to KPMGF but more accurately -produces the correctly normalised Lorentzian spectral function. -""" - -import numpy as np -import matplotlib.pyplot as plt -from pyscf import gto, scf, agf2, lib -from dyson import CPGF, util - -ncheb = 50 # Number of Chebyshev moments - -# Define a self-energy using PySCF -mol = gto.M(atom="O 0 0 0; O 0 0 1", basis="6-31g", verbose=0) -mf = scf.RHF(mol).run() -se_static = np.diag(mf.mo_energy) -se = agf2.AGF2(mf, nmom=(None, None)).build_se() - -# Found the bounds of the self-energy - in practice this should be -# done using a lower-scaling solver. -gf = se.get_greens_function(se_static) -emin = gf.energy.min() -emax = gf.energy.max() -grid = np.linspace(emin, emax, 1024) - -# Scale the energies of the Green's function -a = (emax - emin) / (2.0 - 1e-2) -b = (emax + emin) / 2.0 -energy_scaled = (gf.energy - b) / a - -# Compute the Chebyshev moments -c = np.zeros((ncheb, mol.nao, energy_scaled.size)) -c[0] = gf.coupling -c[1] = gf.coupling * energy_scaled -for i in range(2, ncheb): - c[i] = 2.0 * c[i-1] * energy_scaled - c[i-2] -moments = lib.einsum("qx,npx->npq", gf.coupling, c) - -# Use the solver to get the spectral function -solver = CPGF(moments, grid, (a, b), eta=1.0) -sf = solver.kernel() - -# Get a reference spectral function for comparison -sf_ref = util.build_spectral_function(gf.energy, gf.coupling, grid, eta=1.0) - -# Plot the results -plt.plot(grid, sf_ref, "C0-", label="Reference") -ylim = plt.ylim() -plt.plot(grid, sf, "C1-", label="CPGF") -plt.legend() -plt.xlabel("Frequency (Ha)") -plt.ylabel("Spectral function") -plt.ylim(ylim) -plt.show() diff --git a/examples/10-ph_separation.py b/examples/10-ph_separation.py deleted file mode 100644 index 68def42..0000000 --- a/examples/10-ph_separation.py +++ /dev/null @@ -1,47 +0,0 @@ -""" -Example of mixing MBL solvers. -""" - -import numpy as np -import matplotlib.pyplot as plt -from pyscf import gto, scf, agf2, lib -from dyson import MBLSE, MixedMBLSE, util - -niter_occ = 1 -niter_vir = 2 -grid = np.linspace(-40, 20, 1024) - -# Define a self-energy using PySCF -mol = gto.M(atom="O 0 0 0; O 0 0 1", basis="6-31g", verbose=0) -mf = scf.RHF(mol).run() -se_static = np.diag(mf.mo_energy) -se = agf2.AGF2(mf, nmom=(None, None)).build_se() -se_occ = se.get_occupied() -se_vir = se.get_virtual() - -# Apply a solver to the occupied sector -solver_occ = MBLSE(se_static, se_occ.moment(range(2*niter_occ+2))) -solver_occ.kernel() - -# Apply a solver to the virtual sector -solver_vir = MBLSE(se_static, se_vir.moment(range(2*niter_vir+2))) -solver_vir.kernel() - -# Mix the solvers -mix = MixedMBLSE(solver_occ, solver_vir) - -# Use the mixed solver to get the spectral function -e, v = mix.get_dyson_orbitals() -sf = util.build_spectral_function(e, v, grid, eta=1.0) - -# Get a reference spectral function for comparison -gf = se.get_greens_function(se_static) -sf_ref = util.build_spectral_function(gf.energy, gf.coupling, grid, eta=1.0) - -# Plot the results -plt.plot(grid, sf_ref, "C0-", label="Reference") -plt.plot(grid, sf, "C1-", label="MBLSE") -plt.legend() -plt.xlabel("Frequency (Ha)") -plt.ylabel("Spectral function") -plt.show() diff --git a/examples/11-aufbau.py b/examples/11-aufbau.py deleted file mode 100644 index f0ef8a5..0000000 --- a/examples/11-aufbau.py +++ /dev/null @@ -1,21 +0,0 @@ -""" -Example of applying Aufbau principle. -""" - -import numpy as np -from dyson import Lehmann, AufbauPrinciple - -# Define some energies -e = np.arange(10).astype(float) - -# Put them into a Lehmann representation for a Green's function -c = np.eye(e.size) -gf = Lehmann(e, c) - -# Define the number of electrons filling them, at double filling (i.e. RHF) -nelec = 6 - -# Use the AufbauPrinciple class to get the HOMO and LUMO, and therefore the -# chemical potential -solver = AufbauPrinciple(gf, nelec, occupancy=2) -solver.kernel() diff --git a/examples/12-auxiliary_shift.py b/examples/12-auxiliary_shift.py deleted file mode 100644 index 8399907..0000000 --- a/examples/12-auxiliary_shift.py +++ /dev/null @@ -1,37 +0,0 @@ -""" -Example of applying auxiliary shifts to satisfy the number of -electrons when one has auxiliaries, and the Aufbau principle alone -cannot satisfy the number of electrons. -""" - -import numpy as np -from dyson import Lehmann, MBLSE, AufbauPrinciple, AuxiliaryShift - -np.random.seed(1) - -# Define a Fock matrix -n = 10 -fock = np.diag(np.random.random(n)) - -# Define a self-energy -moms = np.random.random((6, n, n)) -moms = moms + moms.transpose(0, 2, 1) -mblse = MBLSE(fock, moms) -mblse.kernel() -se = mblse.get_self_energy() - -# Define the number of electrons filling them, at double filling (i.e. RHF) -nelec = 6 - -# Use the AufbauPrinciple class to get the chemical potential - this -# won't be satisfied exactly -w, v = se.diagonalise_matrix(fock) -gf = Lehmann(w, v[:n]) -solver = AufbauPrinciple(gf, nelec, occupancy=2) -solver.kernel() - -# Use the AuxiliaryShift class to get the chemical potential more -# accurately, by shifting the self-energy poles with respect to those -# of the Green's function -solver = AuxiliaryShift(fock, se, nelec, occupancy=2) -solver.kernel() diff --git a/examples/13-density_relaxation.py b/examples/13-density_relaxation.py deleted file mode 100644 index fc13465..0000000 --- a/examples/13-density_relaxation.py +++ /dev/null @@ -1,30 +0,0 @@ -""" -Example of optimising the density matrix such that it is self-consistent -with the Fock matrix, in the presence of some self-energy, with -intermediate chemical potential optimisation. -""" - -import numpy as np -from pyscf import gto, scf, agf2 -from dyson import DensityRelaxation - -# Define a self-energy using PySCF -mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="aug-cc-pvdz", verbose=0) -mf = scf.RHF(mol).run() -fock = np.diag(mf.mo_energy) -gf2 = agf2.AGF2(mf, nmom=(None, 0)) -se = gf2.build_se() - -# Define a function to obtain the Fock matrix in the MO basis -def get_fock(dm): - dm_ao = np.linalg.multi_dot((mf.mo_coeff, dm, mf.mo_coeff.T)) - fock = mf.get_fock(dm=dm_ao) - return np.linalg.multi_dot((mf.mo_coeff.T, fock, mf.mo_coeff)) - -# Use the DensityRelaxation class to relax the density matrix such -# that it is self-consistent with the Fock matrix, and the number of -# electrons is correct -solver = DensityRelaxation(get_fock, se, mol.nelectron) -solver.conv_tol = 1e-10 -solver.max_cycle_inner = 30 -solver.kernel() diff --git a/examples/20-kernel_types.py b/examples/20-kernel_types.py deleted file mode 100644 index fcd3d67..0000000 --- a/examples/20-kernel_types.py +++ /dev/null @@ -1,62 +0,0 @@ -""" -Comparison of different kernel types in the kernel polynomial method. -""" - -import numpy as np -import matplotlib.pyplot as plt -from pyscf import gto, scf, agf2, lib -from dyson import KPMGF, util - -ncheb = 50 # Number of Chebyshev moments -kernel_types = [ - "dirichlet", - "lorentz", - "lanczos", - "jackson", - "fejer", -] # Kernel methods - -# Define a self-energy using PySCF -mol = gto.M(atom="O 0 0 0; O 0 0 1", basis="6-31g", verbose=0) -mf = scf.RHF(mol).run() -se_static = np.diag(mf.mo_energy) -se = agf2.AGF2(mf, nmom=(None, None)).build_se() - -# Found the bounds of the self-energy - in practice this should be -# done using a lower-scaling solver. -gf = se.get_greens_function(se_static) -emin = gf.energy.min() -emax = gf.energy.max() -grid = np.linspace(emin, emax, 1024) - -# Scale the energies of the Green's function -a = (emax - emin) / (2.0 - 1e-2) -b = (emax + emin) / 2.0 -energy_scaled = (gf.energy - b) / a - -# Compute the Chebyshev moments -c = np.zeros((ncheb, mol.nao, energy_scaled.size)) -c[0] = gf.coupling -c[1] = gf.coupling * energy_scaled -for i in range(2, ncheb): - c[i] = 2.0 * c[i-1] * energy_scaled - c[i-2] -moments = lib.einsum("qx,npx->npq", gf.coupling, c) - -# Get a reference spectral function for comparison -sf_ref = util.build_spectral_function(gf.energy, gf.coupling, grid, eta=1.0) - -plt.figure() -plt.plot(grid, sf_ref, "C0-", label="Reference") - -for i, kernel_type in enumerate(kernel_types): - # Use the solver to get the spectral function - solver = KPMGF(moments, grid, (a, b), kernel_type=kernel_type) - sf = solver.kernel() - - # Plot the results - plt.plot(grid, sf, "C%d-"%(i+1), label=kernel_type.capitalize()) - -plt.legend() -plt.xlabel("Frequency (Ha)") -plt.ylabel("Spectral function") -plt.show() diff --git a/examples/30-agf2.py b/examples/30-agf2.py deleted file mode 100644 index 8b91bbd..0000000 --- a/examples/30-agf2.py +++ /dev/null @@ -1,66 +0,0 @@ -""" -Example performing an AGF2 calculation, using the `SelfConsistent` -and `DensityRelaxation` solvers, along with the `MP2` expressions, -leverage the plug-and-play callback. -""" - -import numpy as np -from pyscf import gto, scf, lib -from dyson import Lehmann, NullLogger, MBLSE, MixedMBLSE, DensityRelaxation, SelfConsistent -from dyson.expressions import MP2 - -nmom = 2 - -# Define a system using PySCF -mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto3g", verbose=0) -mf = scf.RHF(mol).run() - -# Define a function to calculate the Fock matrix in the MO basis -def get_fock(rdm1_mo): - rdm1_ao = np.linalg.multi_dot((mf.mo_coeff, rdm1_mo, mf.mo_coeff.T)) - fock_ao = mf.get_fock(dm=rdm1_ao) - fock_mo = np.linalg.multi_dot((mf.mo_coeff.T, fock_ao, mf.mo_coeff)) - return fock_mo - -# Define a function to calculate the self-energy - also uses DIIS -# to extrapolate those moments. Note that the `diis` object would -# need to be cleared between calculations in the same script. -diis = lib.diis.DIIS() -def get_se(gf, se_prev=None): - mo_energy, mo_coeff, mo_occ = gf.as_orbitals(mo_coeff=mf.mo_coeff, occupancy=2) - fock = get_fock(gf.occupied().moment(0) * 2) - - mp2 = MP2["Dyson"](mf, mo_energy=mo_energy, mo_coeff=mo_coeff, mo_occ=mo_occ) - th, tp = mp2.build_se_moments(nmom) - th = lib.einsum("...ij,pi,qj->...pq", th, gf.couplings, gf.couplings) - tp = lib.einsum("...ij,pi,qj->...pq", tp, gf.couplings, gf.couplings) - th, tp = diis.update(np.array([th, tp]), xerr=None) - - solverh = MBLSE(fock, th, log=NullLogger()) - solverp = MBLSE(fock, tp, log=NullLogger()) - solver = MixedMBLSE(solverh, solverp) - solver.kernel() - - return solver.get_self_energy() - -# Define the initial Green's function -gf = Lehmann(mf.mo_energy, np.eye(mf.mo_energy.size)) - -# Run the solver -solver = SelfConsistent( - get_se, - get_fock, - gf, - relax_solver=DensityRelaxation, - conv_tol=1e-10, -) -solver.kernel() -solver.log.info("IP: %.8f", -solver.get_greens_function().occupied().energies[-1]) -solver.log.info("EA: %.8f", solver.get_greens_function().virtual().energies[0]) - -# Compare to PySCF -print("\nPySCF:") -from pyscf import agf2 -gf2 = agf2.AGF2(mf) -gf2.verbose = 3 -gf2.kernel() diff --git a/examples/31-fci.py b/examples/31-fci.py deleted file mode 100644 index dd1d670..0000000 --- a/examples/31-fci.py +++ /dev/null @@ -1,31 +0,0 @@ -""" -Example performing an FCI calculation for the IP, using the `Davidson` -solver and the `FCI` expressions. -""" - -import numpy as np -from pyscf import gto, scf, lib -from dyson import Lehmann, NullLogger, Davidson, MBLGF -from dyson.expressions import FCI - - -# Define a system using PySCF -mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto3g", verbose=0) -mf = scf.RHF(mol).run() - -# Get the expressions -fci = FCI["1h"](mf) -diag = fci.diagonal() -matvec = fci.apply_hamiltonian - -# Run the Davidson algorithm -solver = Davidson(matvec, diag, nroots=5, nphys=fci.nocc) -solver.conv_tol = 1e-10 -solver.kernel() -solver.log.info("IP: %.8f", -solver.get_greens_function().occupied().energies[-1]) - -# Use MBLGF -moments = fci.build_gf_moments(4) -solver = MBLGF(moments) -solver.kernel() -solver.log.info("IP: %.8f", -solver.get_greens_function().occupied().energies[-1]) diff --git a/examples/32-momgfccsd.py b/examples/32-momgfccsd.py deleted file mode 100644 index cde22e8..0000000 --- a/examples/32-momgfccsd.py +++ /dev/null @@ -1,59 +0,0 @@ -""" -Example performaing a MomGF-CCSD calculation. -""" - -import numpy as np -import matplotlib.pyplot as plt -from pyscf import gto, scf, cc -from dyson import MBLGF, NullLogger, util -from dyson.expressions import CCSD - -nmom = 4 - -# Define a system using PySCF -mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto3g", verbose=0) -mf = scf.RHF(mol).run() - -# Run a CCSD calculation -ccsd = cc.CCSD(mf) -ccsd.kernel() -ccsd.solve_lambda() - -# Find the moments -expr = CCSD["1h"](mf, t1=ccsd.t1, t2=ccsd.t2, l1=ccsd.l1, l2=ccsd.l2) -th = expr.build_gf_moments(nmom) -expr = CCSD["1p"](mf, t1=ccsd.t1, t2=ccsd.t2, l1=ccsd.l1, l2=ccsd.l2) -tp = expr.build_gf_moments(nmom) - -# Solve for the Green's function -solverh = MBLGF(th, hermitian=False) -solverh.kernel() -gfh = solverh.get_greens_function() -solverp = MBLGF(tp, hermitian=False) -solverp.kernel() -gfp = solverp.get_greens_function() -gf = gfh + gfp - -# Get the spectrum -grid = np.linspace(-5, 5, 1024) -eta = 1e-1 -sf = util.build_spectral_function(gf.energies, gf.couplings, grid, eta=eta) - -# If PySCF version is new enough, plot a reference -try: - momgfcc = cc.momgfccsd.MomGFCCSD(ccsd, ((nmom-2)//2, (nmom-2)//2)) - eh, vh, ep, vp = momgfcc.kernel() - e = np.concatenate((eh, ep), axis=0) - v = np.concatenate((vh[0], vp[0]), axis=1) - u = np.concatenate((vh[1], vp[1]), axis=1) - sf_ref = util.build_spectral_function(e, (v, u), grid, eta=eta) - plt.plot(grid, sf_ref, "C1--", label="MomGF-CCSD (PySCF)", zorder=10) -except AttributeError: - pass - -# Plot the results -plt.plot(grid, sf, "C0-", label="MomGF-CCSD") -plt.legend() -plt.xlabel("Frequency (Ha)") -plt.ylabel("Spectral function") -plt.show() diff --git a/examples/33-adc2.py b/examples/33-adc2.py deleted file mode 100644 index e3cda91..0000000 --- a/examples/33-adc2.py +++ /dev/null @@ -1,33 +0,0 @@ -""" -Example performing an ADC(2) calculation, using the `Davidson` -solver and the `MP2` expressions. -""" - -import numpy as np -from pyscf import gto, scf, lib -from dyson import Lehmann, NullLogger, Davidson -from dyson.expressions import MP2 - - -# Define a system using PySCF -mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto3g", verbose=0) -mf = scf.RHF(mol).run() - -# Get the expressions -mp2 = MP2["1h"](mf) -static = mp2.get_static_part() -diag = mp2.diagonal(static=static) -matvec = lambda v: mp2.apply_hamiltonian(v, static=static) - -# Run the Davidson algorithm -solver = Davidson(matvec, diag, nroots=5, nphys=mp2.nocc) -solver.conv_tol = 1e-10 -solver.kernel() -solver.log.info("IP: %.8f", -solver.get_greens_function().occupied().energies[-1]) - -# Compare to PySCF -print("\nPySCF:") -from pyscf import adc -adc2 = adc.ADC(mf) -adc2.verbose = 4 -adc2.kernel(nroots=5) diff --git a/examples/34-eom_ccsd.py b/examples/34-eom_ccsd.py deleted file mode 100644 index 345851d..0000000 --- a/examples/34-eom_ccsd.py +++ /dev/null @@ -1,34 +0,0 @@ -""" -Example performing an EOM-CCSD calculation, using the `Davidson` -solver and the `CCSD` expressions. -""" - -import numpy as np -from pyscf import gto, scf, lib -from dyson import Lehmann, NullLogger, Davidson -from dyson.expressions import CCSD - - -# Define a system using PySCF -mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto3g", verbose=0) -mf = scf.RHF(mol).run() - -# Get the expressions -ccsd = CCSD["1h"](mf) -diag = ccsd.diagonal() -matvec = ccsd.apply_hamiltonian - -# Run the Davidson algorithm -solver = Davidson(matvec, diag, nroots=5, nphys=ccsd.nocc) -solver.conv_tol = 1e-10 -solver.kernel() -solver.log.info("IP: %.8f", -solver.get_greens_function().occupied().energies[-1]) - -# Compare to PySCF -print("\nPySCF:") -from pyscf import cc -ccsd = cc.CCSD(mf) -ccsd.verbose = 4 -ccsd.kernel() -ccsd.ipccsd(nroots=5) - diff --git a/examples/36-moment_gw.py b/examples/36-moment_gw.py deleted file mode 100644 index 4fd683f..0000000 --- a/examples/36-moment_gw.py +++ /dev/null @@ -1,41 +0,0 @@ -""" -Example performing a momentGW calculation. -""" - -import numpy as np -from pyscf import gto, dft, cc -from dyson import MBLSE, MixedMBLSE, NullLogger, util -from dyson.expressions import GW - -nmom_max = 5 - -# Define a system using PySCF -mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto3g", verbose=0) -mf = dft.RKS(mol) -mf = mf.density_fit() -mf.xc = "hf" -mf.kernel() - -# Find the moments -gw = GW["Dyson"](mf) -static = gw.get_static_part() -th, tp = gw.build_se_moments(nmom_max+1) - -# Solve for the Green's function -solverh = MBLSE(static, th) -solverp = MBLSE(static, tp) -solver = MixedMBLSE(solverh, solverp) -solver.kernel() -gf = solver.get_greens_function() -gf = gf.physical() -solver.log.info("IP: %.8f", -gf.occupied().energies[-1]) -solver.log.info("EA: %.8f", gf.virtual().energies[0]) - -# Compare to momentGW -import momentGW -gw_ref = momentGW.GW(mf) -conv, gf_ref, se_ref = gw_ref.kernel(nmom_max) -gf_ref.remove_uncoupled(tol=0.1) -solver.log.info("") -solver.log.info("IP (ref): %.8f", -gf_ref.get_occupied().energy[-1]) -solver.log.info("EA (ref): %.8f", gf_ref.get_virtual().energy[0]) diff --git a/examples/38-fci_static_self_energy.py b/examples/38-fci_static_self_energy.py deleted file mode 100644 index 29d3696..0000000 --- a/examples/38-fci_static_self_energy.py +++ /dev/null @@ -1,44 +0,0 @@ -""" -Example showing the relationship between the FCI static self-energy -and the Green's function moments. -""" - -import numpy as np -from pyscf import gto, scf, lib -from dyson import Lehmann, NullLogger, MBLGF, MixedMBLGF -from dyson.expressions import FCI - - -# Define a system using PySCF -mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto3g", verbose=0) -mf = scf.RHF(mol).run() - -# Get the expressions -fci_1h = FCI["1h"](mf) -fci_1p = FCI["1p"](mf) - -# Use MBLGF -th = fci_1h.build_gf_moments(4) -tp = fci_1p.build_gf_moments(4) -solver_h = MBLGF(th) -solver_h.kernel() -solver_p = MBLGF(tp) -solver_p.kernel() -solver = MixedMBLGF(solver_h, solver_p) - -# Get the Green's function -gf = solver.get_greens_function() - -# Back-transform to the self-energy and use the sum of the first order -# moments as the static self-energy -se = solver.get_self_energy() -se_static = th[1] + tp[1] -gf_recov = Lehmann(*se.diagonalise_matrix_with_projection(se_static)) -assert np.allclose(gf.moment(range(4)), gf_recov.moment(range(4))) - -# This is equivalent to the Fock matrix evaluated at the FCI density -dm = th[0] * 2.0 -dm = np.linalg.multi_dot((mf.mo_coeff, dm, mf.mo_coeff.T)) -f = mf.get_fock(dm=dm) -f = np.linalg.multi_dot((mf.mo_coeff.T, f, mf.mo_coeff)) -assert np.allclose(f, se_static) diff --git a/examples/particle-hole-separation.py b/examples/particle-hole-separation.py new file mode 100644 index 0000000..41eef27 --- /dev/null +++ b/examples/particle-hole-separation.py @@ -0,0 +1,59 @@ +"""Example of particle-hole separated calculations.""" + +import matplotlib.pyplot as plt +import numpy +from pyscf import gto, scf + +from dyson import ADC2, MBLGF, Spectral +from dyson.grids import GridRF +from dyson.plotting import format_axes_spectral_function, plot_dynamic + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto-3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Use FCI one-hole and one-particle expressions for the Hamiltonian +exp_h = ADC2.h.from_mf(mf) +exp_p = ADC2.p.from_mf(mf) + +# Use MBLGF to solve the Hamiltonian for each case separately +solver_h = MBLGF.from_expression(exp_h, max_cycle=1) +solver_h.kernel() +solver_p = MBLGF.from_expression(exp_p, max_cycle=1) +solver_p.kernel() + +# Combine the results -- this function operators by projecting the result back into a self-energy +# and combining the two self-energies, before diagonalising the combined self-energy to get a new +# result spectrum. This may have unwanted consequences for some methodology, so use with care. +result = Spectral.combine(solver_h.result, solver_p.result) + +# Get the spectral functions +grid = GridRF.from_uniform(-3.0, 3.0, 1024, eta=0.05) +spectrum_h = (1 / numpy.pi) * grid.evaluate_lehmann( + solver_h.result.get_greens_function(), + ordering="advanced", + reduction="trace", + component="imag", +) +spectrum_p = (1 / numpy.pi) * grid.evaluate_lehmann( + solver_p.result.get_greens_function(), + ordering="advanced", + reduction="trace", + component="imag", +) +spectrum_combined = (1 / numpy.pi) * grid.evaluate_lehmann( + result.get_greens_function(), + ordering="advanced", + reduction="trace", + component="imag", +) + +# Plot the spectra +fig, ax = plt.subplots() +plot_dynamic(spectrum_combined, fmt="k-", label="Combined Spectrum", energy_unit="eV", ax=ax) +plot_dynamic(spectrum_h, fmt="C0--", label="Hole Spectrum", energy_unit="eV", ax=ax) +plot_dynamic(spectrum_p, fmt="C1--", label="Particle Spectrum", energy_unit="eV", ax=ax) +format_axes_spectral_function(grid, ax=ax, energy_unit="eV") +plt.legend() +plt.show() diff --git a/examples/solver-aufbau.py b/examples/solver-aufbau.py new file mode 100644 index 0000000..b81896e --- /dev/null +++ b/examples/solver-aufbau.py @@ -0,0 +1,53 @@ +"""Example of the Aufbau principle solver. + +This solver applies another solver before filling the resulting solution according to the Aufbau +principle. This attaches a chemical potential to the solution, and the resulting self-energy and +Green's function. +""" + +from pyscf import gto, scf + +from dyson import MBLGF, TDAGW, AufbauPrinciple, Exact +from dyson.solvers.static.chempot import search_aufbau_global + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto-3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Use a TDA-GW Dyson expression for the Hamiltonian +exp = TDAGW.dyson.from_mf(mf) + +# Use the exact solver to get the self-energy for demonstration purposes +exact = Exact.from_expression(exp) +exact.kernel() +static = exact.result.get_static_self_energy() +self_energy = exact.result.get_self_energy() +overlap = exact.result.get_overlap() + +# Solve the Hamiltonian using the Aufbau solver, initialisation via either: + +# 1) Create the solver from a self-energy +solver = AufbauPrinciple.from_self_energy(static, self_energy, overlap=overlap, nelec=mol.nelectron) +solver.kernel() + +# 2) Create the solver directly from the self-energy +solver = AufbauPrinciple( + static, + self_energy, + overlap=overlap, + nelec=mol.nelectron, +) +solver.kernel() + +# By default, this is solving the input self-energy using the Exact solver. To use another solver, +# e.g. MBLSE, as the base solver, you can specify it as the `solver` argument +solver = AufbauPrinciple.from_self_energy( + static, self_energy, overlap=overlap, nelec=mol.nelectron, solver=MBLGF +) +solver.kernel() + +# If you don't want to solve the self-energy at all and just want to find a chemical potential for +# an existing solution, you can pass the Green's function directly to the search functions +greens_function = solver.result.get_greens_function() +chempot, error = search_aufbau_global(greens_function, mol.nelectron) diff --git a/examples/solver-corrvec.py b/examples/solver-corrvec.py new file mode 100644 index 0000000..03272ce --- /dev/null +++ b/examples/solver-corrvec.py @@ -0,0 +1,59 @@ +"""Example of the correction vector solver. + +This solver is a dynamic iterative solver to build the downfolded frequency-dependent Green's +function. It uses a GMRES algorithm under the hood to iteratively improve the correction vector +and contract to the Green's function. +""" + +import numpy +from pyscf import gto, scf + +from dyson import FCI, CorrectionVector, Exact +from dyson.grids import RealFrequencyGrid + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto-3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Use an FCI one-hole expression for the Hamiltonian +exp = FCI.hole.from_mf(mf) + +# Initialise a real frequency grid for the correction vector solver +grid = RealFrequencyGrid.from_uniform(-3.0, 3.0, 128, eta=1e-2) + +# Use the exact solver to get the self-energy for demonstration purposes +exact = Exact.from_expression(exp) +exact.kernel() +static = exact.result.get_static_self_energy() +self_energy = exact.result.get_self_energy() +overlap = exact.result.get_overlap() + +# Solve the Hamiltonian using the correction vector solver, initialisation via either: + +# 1) Create the solver from the expression +solver = CorrectionVector.from_expression(exp, grid=grid, ordering="ordered") +gf = solver.kernel() + +# 2) Create the solver from a self-energy +solver = CorrectionVector.from_self_energy( + static, self_energy, overlap=overlap, grid=grid, ordering="ordered" +) +gf = solver.kernel() + +# 3) Create the solver directly from the matrix and excitation vectors +solver = CorrectionVector( + exp.apply_hamiltonian, + exp.diagonal(), + exp.nphys, + grid, + exp.get_excitation_bra, + exp.get_excitation_ket, + ordering="ordered", +) +gf = solver.kernel() + +# Compare to that of the Exact solver, by downfolding the Green's function corresponding to the +# exact result onto the same grid +gf_exact = grid.evaluate_lehmann(exact.result.get_greens_function(), ordering="ordered") +print("Correction vector error:", numpy.max(numpy.abs(gf - gf_exact))) diff --git a/examples/solver-cpgf.py b/examples/solver-cpgf.py new file mode 100644 index 0000000..5aa7fba --- /dev/null +++ b/examples/solver-cpgf.py @@ -0,0 +1,74 @@ +"""Example of the Chebyshev polynomial Green's function solver. + +This solver uses Chebyshev polynomials to evaluate the Green's function on a real frequency grid. It +is systematically improvable by increasing the order of the Chebyshev polynomial expansion. It is +related to MBLGF, however, it does not offer a static result, rather a dynamic Green's function. +""" + +import numpy +from pyscf import gto, scf + +from dyson import CPGF, FCI, Exact, util +from dyson.grids import RealFrequencyGrid + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto-3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Use an FCI one-hole expression for the Hamiltonian +exp = FCI.hole.from_mf(mf) + +# Initialise a real frequency grid for the correction vector solver +grid = RealFrequencyGrid.from_uniform(-3.0, 0.0, 128, eta=1e-2) + +# Use the exact solver to get the self-energy for demonstration purposes +exact = Exact.from_expression(exp) +exact.kernel() +static = exact.result.get_static_self_energy() +self_energy = exact.result.get_self_energy() +overlap = exact.result.get_overlap() + +# CPGF requires a pair of scaling parameters, which are used to scale the spectrum onto the range +# [-1, 1]. The scaling parameters can be obtained from the minimum and maximum eigenvalues of the +# Hamiltonian if they are known a priori, or they can be approximated from the diagonal of the +# expression. If the approximation is used, it is recommended to use an additional factor to avoid +# cases where the minimum and maximum of the diagonal are not representative of the spectrum. +energies, _ = self_energy.diagonalise_matrix(static, overlap=overlap) +scaling = util.get_chebyshev_scaling_parameters(energies.min(), energies.max()) + +# Solve the Hamiltonian using the CPGF solver, initialisation via either: + +# 1) Create the solver from the expression +max_cycle = 1024 +solver = CPGF.from_expression( + exp, grid=grid, max_cycle=max_cycle, scaling=scaling, ordering="advanced" +) +gf = solver.kernel() + +# 2) Create the solver from a self-energy +solver = CPGF.from_self_energy( + static, + self_energy, + overlap=overlap, + grid=grid, + max_cycle=max_cycle, + scaling=scaling, + ordering="advanced", +) +gf = solver.kernel() + +# 3) Create the solver directly from the matrix and excitation vectors +solver = CPGF( + exp.build_gf_chebyshev_moments(max_cycle + 1, scaling=scaling), + grid, + max_cycle=max_cycle, + scaling=scaling, + ordering="advanced", +) +gf = solver.kernel() + +# Compare to that of the Exact solver, by downfolding the Green's function corresponding to the +# exact result onto the same grid +gf_exact = grid.evaluate_lehmann(exact.result.get_greens_function(), ordering="advanced") +print("Correction vector error:", numpy.max(numpy.abs(gf - gf_exact))) diff --git a/examples/solver-davidson.py b/examples/solver-davidson.py new file mode 100644 index 0000000..ad0b3b3 --- /dev/null +++ b/examples/solver-davidson.py @@ -0,0 +1,47 @@ +"""Example of the Davidson eigenvalue solver. + +This solver is a traditional iterative Jacobi--Davidson eigenvalue solver to find the eigenvalues +and eigenvectors of a Hamiltonian corresponding to the lowest-lying states. As such, it does not +fully calculate the Green's function, but targets the energy region close to the Fermi level. +""" + +import numpy +from pyscf import gto, scf + +from dyson import FCI, Davidson, Exact + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto-3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Use an FCI one-hole expression for the Hamiltonian +exp = FCI.hole.from_mf(mf) + +# Use the exact solver to get the self-energy for demonstration purposes +exact = Exact.from_expression(exp) +exact.kernel() +static = exact.result.get_static_self_energy() +self_energy = exact.result.get_self_energy() +overlap = exact.result.get_overlap() + +# Solve the Hamiltonian using the Davidson solver, initialisation via either: + +# 1) Create the solver from the expression +solver = Davidson.from_expression(exp, nroots=5) +solver.kernel() + +# 2) Create the solver from a self-energy +solver = Davidson.from_self_energy(static, self_energy, overlap=overlap, nroots=5) +solver.kernel() + +# 3) Create the solver directly from the matrix and excitation vectors +solver = Davidson( + exp.apply_hamiltonian, + exp.diagonal(), + numpy.asarray(exp.get_excitation_bras()), + numpy.asarray(exp.get_excitation_kets()), + hermitian=exp.hermitian_upfolded, + nroots=5, +) +solver.kernel() diff --git a/examples/solver-density.py b/examples/solver-density.py new file mode 100644 index 0000000..f8a009e --- /dev/null +++ b/examples/solver-density.py @@ -0,0 +1,92 @@ +"""Example of the density relaxation solver. + +This solver relaxes the density matrix of a system in the presence of a self-energy. Between +iterations, the self-energy is shifted in order to allow the Aufbau principle to assign a +chemical potential that best matches the particle number of the system. The resulting Green's +function is a minimum with respect to the self-consistent field. + +The solvers require a function to evaluate the static part of the self-energy (i.e. the Fock matrix) +for a given density matrix. We provide a convenience function to get this from a PySCF RHF object. + +Note that for some Hamiltonians, the relaxation of the density and the shifting of the self-energy +may not commute, i.e. their solutions cannot be obtain simultaneously. In this case, one solution +is favoured according to a parameter. +""" + +from pyscf import gto, scf + +from dyson import MBLSE, TDAGW, AufbauPrinciple, AuxiliaryShift, DensityRelaxation, Exact +from dyson.solvers.static.density import get_fock_matrix_function + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto-3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Use a TDA-GW Dyson expression for the Hamiltonian +exp = TDAGW.dyson.from_mf(mf) + +# Use the exact solver to get the self-energy for demonstration purposes +exact = Exact.from_expression(exp) +exact.kernel() +static = exact.result.get_static_self_energy() +self_energy = exact.result.get_self_energy() +overlap = exact.result.get_overlap() + +# Solve the Hamiltonian using the density relaxation solver, initialisation via either: + +# 1) Create the solver from a self-energy +solver = DensityRelaxation.from_self_energy( + static, + self_energy, + overlap=overlap, + get_static=get_fock_matrix_function(mf), + nelec=mol.nelectron, +) +solver.kernel() + +# 2) Create the solver directly from the self-energy +solver = DensityRelaxation( + get_fock_matrix_function(mf), + self_energy, + overlap=overlap, + nelec=mol.nelectron, +) +solver.kernel() + +# Like the auxiliary shift solver, we can customise the solvers + + +class MyAufbauPrinciple(AufbauPrinciple): # noqa: D101 + solver = MBLSE + + +class MyAuxiliaryShift(AuxiliaryShift): # noqa: D101 + solver = MyAufbauPrinciple + + +solver = DensityRelaxation.from_self_energy( + static, + self_energy, + overlap=overlap, + get_static=get_fock_matrix_function(mf), + nelec=mol.nelectron, + solver_outer=MyAuxiliaryShift, + solver_inner=MyAufbauPrinciple, +) +solver.kernel() + +# By default, the non-commutative solutions favour the self-consistency in the density matrix, +# rather than the particle number. To favour the particle number, we can pass an additional +# parameter +solver = DensityRelaxation.from_self_energy( + static, + self_energy, + overlap=overlap, + get_static=get_fock_matrix_function(mf), + nelec=mol.nelectron, + solver_outer=MyAuxiliaryShift, + solver_inner=MyAufbauPrinciple, + favour_rdm=False, # Favour the particle number over the density matrix +) +solver.kernel() diff --git a/examples/solver-downfolded.py b/examples/solver-downfolded.py new file mode 100644 index 0000000..56ee804 --- /dev/null +++ b/examples/solver-downfolded.py @@ -0,0 +1,51 @@ +r"""Example of the downfolded eigenvalue solver. + +This solver finds the eigenvalues of the frequency-dependent downfolded self-energy matrix in a +self-consistent manner. It converges on the pole of the Green's function closest to the initial +guess, but does not account for a fully featured Green's function. The eigenvalue problem also +maintains a dependency on the broadening parameter :math:`\eta`. +""" + +import numpy +from pyscf import gto, scf + +from dyson import FCI, Downfolded, Exact +from dyson.grids import GridRF + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto-3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Use an FCI one-hole expression for the Hamiltonian +exp = FCI.hole.from_mf(mf) + +# Use the exact solver to get the self-energy for demonstration purposes +exact = Exact.from_expression(exp) +exact.kernel() +static = exact.result.get_static_self_energy() +self_energy = exact.result.get_self_energy() +overlap = exact.result.get_overlap() + +# Solve the Hamiltonian using the Downfolded solver, initialisation via either: + +# 1) Create the solver from a self-energy +solver = Downfolded.from_self_energy(static, self_energy, overlap=overlap, eta=1e-2) +solver.kernel() + +# 2) Create the solver directly from the generating function + + +def _function(freq: float) -> numpy.ndarray: + """Evaluate the self-energy at the frequency.""" + grid = GridRF(numpy.array([freq]), eta=1e-2) + return grid.evaluate_lehmann(self_energy, ordering="ordered").array[0] + + +solver = Downfolded( + static, + _function, + overlap=overlap, + hermitian=exp.hermitian_downfolded, +) +solver.kernel() diff --git a/examples/solver-exact.py b/examples/solver-exact.py new file mode 100644 index 0000000..86cbeef --- /dev/null +++ b/examples/solver-exact.py @@ -0,0 +1,41 @@ +"""Example of the exact diagonalisation solver. + +This solver is a non-scalable solver that exactly diagonalises the dense Hamiltonian as a +demonstration for small systems. When constructing from an expression, it constructs the Hamiltonian +matrix using repeated applications of the matrix-vector product to unit vectors, which is also slow. +""" + +import numpy +from pyscf import gto, scf + +from dyson import FCI, Exact + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto-3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Use an FCI one-hole expression for the Hamiltonian +exp = FCI.hole.from_mf(mf) + +# Solve the Hamiltonian using the Exact solver, initialisation via either: + +# 1) Create the solver from the expression +solver = Exact.from_expression(exp) +solver.kernel() + +# 2) Create the solver from a self-energy +static = solver.result.get_static_self_energy() +self_energy = solver.result.get_self_energy() +overlap = solver.result.get_overlap() +solver = Exact.from_self_energy(static, self_energy, overlap=overlap) +solver.kernel() + +# 3) Create the solver directly from the matrix and excitation vectors +solver = Exact( + exp.build_matrix(), + numpy.asarray(exp.get_excitation_bras()), + numpy.asarray(exp.get_excitation_kets()), + hermitian=exp.hermitian_upfolded, +) +solver.kernel() diff --git a/examples/solver-mblgf.py b/examples/solver-mblgf.py new file mode 100644 index 0000000..798ce04 --- /dev/null +++ b/examples/solver-mblgf.py @@ -0,0 +1,44 @@ +"""Example of the MBLGF solver. + +This solver diagonalises the self-energy via conservation of the spectral moments of the resulting +Green's function, using recursion relations of those moments. The resulting Green's function is +approximate, and is systematically improved by increasing the number of moments (maximum algorithm +cycle) used in the calculation. +""" + +from pyscf import gto, scf + +from dyson import ADC2, MBLGF, Exact + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto-3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Use an ADC(2) one-hole expression for the Hamiltonian +exp = ADC2.hole.from_mf(mf) + +# Use the exact solver to get the self-energy for demonstration purposes +exact = Exact.from_expression(exp) +exact.kernel() +static = exact.result.get_static_self_energy() +self_energy = exact.result.get_self_energy() +overlap = exact.result.get_overlap() + +# Solve the Hamiltonian using the MBLGF solver, initialisation via either: + +# 1) Create the solver from the expression +solver = MBLGF.from_expression(exp, max_cycle=1) +solver.kernel() + +# 2) Create the solver from a self-energy +solver = MBLGF.from_self_energy(static, self_energy, overlap=overlap, max_cycle=1) +solver.kernel() + +# 3) Create the solver directly from the moments +max_cycle = 1 +solver = MBLGF( + solver.result.get_greens_function().moments(range(2 * max_cycle + 2)), + hermitian=exp.hermitian_downfolded, + max_cycle=max_cycle, +) diff --git a/examples/solver-mblse.py b/examples/solver-mblse.py new file mode 100644 index 0000000..9788d6d --- /dev/null +++ b/examples/solver-mblse.py @@ -0,0 +1,46 @@ +"""Example of the MBLSE solver. + +This solver diagonalises the self-energy via conservation of its spectral moments, using recursion +relations of those moments. The resulting Green's function is approximate, and is systematically +improved by increasing the number of moments (maximum algorithm cycle) used in the calculation. +""" + +from pyscf import gto, scf + +from dyson import ADC2, MBLSE, Exact + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto-3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Use an ADC(2) one-hole expression for the Hamiltonian +exp = ADC2.hole.from_mf(mf) + +# Use the exact solver to get the self-energy for demonstration purposes +exact = Exact.from_expression(exp) +exact.kernel() +static = exact.result.get_static_self_energy() +self_energy = exact.result.get_self_energy() +overlap = exact.result.get_overlap() + +# Solve the Hamiltonian using the MBLSE solver, initialisation via either: + +# 1) Create the solver from the expression +solver = MBLSE.from_expression(exp, max_cycle=1) +solver.kernel() + +# 2) Create the solver from a self-energy +solver = MBLSE.from_self_energy(static, self_energy, overlap=overlap, max_cycle=1) +solver.kernel() + +# 3) Create the solver directly from the moments +max_cycle = 1 +solver = MBLSE( + static, + self_energy.moments(range(2 * max_cycle + 2)), + overlap=overlap, + hermitian=exp.hermitian_downfolded, + max_cycle=max_cycle, +) +solver.kernel() diff --git a/examples/solver-shift.py b/examples/solver-shift.py new file mode 100644 index 0000000..0a93525 --- /dev/null +++ b/examples/solver-shift.py @@ -0,0 +1,54 @@ +"""Example of the auxiliary shift solver. + +This solver applies another solver with a variable shift in the energies of the self-energy. This +shift is optimised to allow an Aufbau principle to arrive at the best possible solution with respect +to the particle number. This modifies the self-energy and attaches a chemical potential to the +solution, and the resulting self-energy and Green's function. +""" + +from pyscf import gto, scf + +from dyson import MBLSE, TDAGW, AufbauPrinciple, AuxiliaryShift, Exact + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; H 0 0 1.64", basis="sto-3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Use a TDA-GW Dyson expression for the Hamiltonian +exp = TDAGW.dyson.from_mf(mf) + +# Use the exact solver to get the self-energy for demonstration purposes +exact = Exact.from_expression(exp) +exact.kernel() +static = exact.result.get_static_self_energy() +self_energy = exact.result.get_self_energy() +overlap = exact.result.get_overlap() + +# Solve the Hamiltonian using the auxiliary shift solver, initialisation via either: + +# 1) Create the solver from a self-energy +solver = AuxiliaryShift.from_self_energy(static, self_energy, overlap=overlap, nelec=mol.nelectron) +solver.kernel() + +# 2) Create the solver directly from the self-energy +solver = AuxiliaryShift( + static, + self_energy, + overlap=overlap, + nelec=mol.nelectron, +) +solver.kernel() + +# By default, this is solving the input self-energy using the default Aufbau solver. To use another +# solver, e.g. one that uses MBLSE for the base solver, you can specify it as the `solver` argument + + +class MyAufbauPrincple(AufbauPrinciple): # noqa: D101 + solver = MBLSE + + +solver = AuxiliaryShift.from_self_energy( + static, self_energy, overlap=overlap, nelec=mol.nelectron, solver=MyAufbauPrincple +) +solver.kernel() diff --git a/examples/spectra.py b/examples/spectra.py new file mode 100644 index 0000000..c91c959 --- /dev/null +++ b/examples/spectra.py @@ -0,0 +1,76 @@ +"""Comparison of spectra from different solvers.""" + +import matplotlib.pyplot as plt +import numpy +from pyscf import gto, scf + +from dyson.expressions import ADC2 +from dyson.grids import GridRF +from dyson.plotting import format_axes_spectral_function, plot_dynamic +from dyson.solvers import CPGF, MBLGF, MBLSE, CorrectionVector, Downfolded, Exact + +# Get a molecule and mean-field from PySCF +mol = gto.M(atom="Li 0 0 0; Li 0 0 1.64", basis="sto3g", verbose=0) +mf = scf.RHF(mol) +mf.kernel() + +# Define a grid for the spectra +grid = GridRF.from_uniform(-3.0, 3.0, 256, eta=1e-1) + +# Get a complete self-energy (identity overlap) to solve for demonstration purposes +exp_h = ADC2.h.from_mf(mf) +exp_p = ADC2.p.from_mf(mf) +exact_h = Exact.from_expression(exp_h) +exact_h.kernel() +exact_p = Exact.from_expression(exp_p) +exact_p.kernel() +result = exact_h.result.combine(exact_p.result) +static = result.get_static_self_energy() +self_energy = result.get_self_energy() + +# Solve the self-energy using each static solver -- since ADC(2) is non-Dyson, we can just add +# the Green's function rather than using the spectral combination utility +spectra = {} +for key, solver_cls, kwargs in [ + ("Exact", Exact, dict()), + ("Downfolded", Downfolded, dict()), + ("MBLSE(1)", MBLSE, dict(max_cycle=1)), + ("MBLGF(1)", MBLGF, dict(max_cycle=1)), +]: + solver = solver_cls.from_self_energy(static, self_energy, **kwargs) + solver.kernel() + gf = solver.result.get_greens_function() + spectra[key] = (1 / numpy.pi) * ( + grid.evaluate_lehmann(gf, ordering="advanced", reduction="trace", component="imag") + ) + +# Solve the self-energy using each dynamic solver +for key, solver_cls, kwargs in [ + ("CorrectionVector", CorrectionVector, dict()), + ("CPGF(256)", CPGF, dict(max_cycle=256)), +]: + solver = solver_cls.from_self_energy( + static, + self_energy, + grid=grid, + ordering="advanced", + reduction="trace", + component="imag", + **kwargs, + ) + gf = solver.kernel() + spectra[key] = (1 / numpy.pi) * gf + +# Plot the spectra +fig, ax = plt.subplots() +for i, (key, spectrum) in enumerate(spectra.items()): + plot_dynamic( + spectrum, + fmt=f"C{i}", + label=key, + energy_unit="eV", + ax=ax, + ) +format_axes_spectral_function(grid, ax=ax, energy_unit="eV") +plt.legend() +plt.show() diff --git a/examples/unknown-pleasures.py b/examples/unknown-pleasures.py new file mode 100644 index 0000000..239b7f0 --- /dev/null +++ b/examples/unknown-pleasures.py @@ -0,0 +1,35 @@ +"""Plot spectra in the style of the cover of Joy Division's 'Unknown Pleasures' album.""" + +import matplotlib.pyplot as plt +import numpy +from pyscf import gto, scf + +from dyson import ADC2, MBLGF, Lehmann, quiet +from dyson.grids import GridRF +from dyson.plotting import unknown_pleasures + +# Suppress output +quiet() + +# Define a grid for the spectra +grid = GridRF.from_uniform(-5.0, 7.0, 128, eta=0.25) + +# Generate random bond distances for a pair of nitrogen atoms +spectra = [] +for _ in range(64): + bond_distance = numpy.random.uniform(0.8, 2.5) + mol = gto.M(atom=f"N 0 0 0; N 0 0 {bond_distance}", basis="cc-pvdz", verbose=0) + mf = scf.RHF(mol).run() + adc2_h = ADC2.h.from_mf(mf) + adc2_p = ADC2.p.from_mf(mf) + + # Solve the ADC(2) Green's function for the hole and particle sectors + gf_h = MBLGF.from_expression(adc2_h, max_cycle=4).kernel().get_greens_function() + gf_p = MBLGF.from_expression(adc2_p, max_cycle=4).kernel().get_greens_function() + gf = Lehmann.concatenate(gf_h, gf_p) + sf = grid.evaluate_lehmann(gf, ordering="advanced", reduction="trace", component="imag") + spectra.append(sf) + +# Plot the spectra in the style of 'Unknown Pleasures' +unknown_pleasures(spectra) +plt.show() diff --git a/pyproject.toml b/pyproject.toml index 49db775..d5b864e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,5 +1,6 @@ [project] name = "dyson" +version = "1.0.0" description = "Dyson equation solvers for electron propagator methods" keywords = [ "quantum", "chemistry", @@ -10,9 +11,9 @@ keywords = [ "greens", "function", ] readme = "README.md" -requires-python = ">=3.8" +requires-python = ">=3.10" classifiers = [ - "Development Status :: 2 - Pre-Alpha", + "Development Status :: 4 - Beta", "Intended Audience :: Science/Research", "Intended Audience :: Developers", "Topic :: Scientific/Engineering", @@ -27,36 +28,77 @@ classifiers = [ ] dependencies = [ "numpy>=1.19.0", + "scipy>=1.5.0", "pyscf>=2.0.0", + "rich>=11.0.0", ] -dynamic = [ - "version", -] - -[tools.setuptools.dynamic] -version = {attr = "dyson.__version__"} - -[build-system] -requires = [ - "setuptools>=46.1.0", -] -build-backend = "setuptools.build_meta" [project.optional-dependencies] dev = [ - "black>=22.6.0", - "isort>=5.10.1", - "coverage[toml]", - "pytest", - "pytest-cov", + "ruff>=0.10.0", + "mypy>=1.5.0", + "coverage[toml]>=5.5.0", + "pytest>=6.2.4", + "pytest-cov>=4.0.0", + "matplotlib>=3.4.0", + "sphinx>=7.0", + "sphinx-book-theme>=1.0", + "sphinx-mdinclude>=0.5", ] -[tool.black] +[tool.ruff] line-length = 100 -target-version = [ - "py311", +target-version = "py312" +include = ["pyproject.toml", "dyson/**/*.py", "tests/**/*.py", "examples/**/*.py"] + +[tool.ruff.format] +quote-style = "double" +indent-style = "space" +line-ending = "lf" + +[tool.ruff.lint] +select = [ + "E", # pycodestyle + "F", # pyflakes + "I", # isort + "D", # pydocstyle + "PL", # pylint +] +ignore = [ + "E722", # bare-except + "E731", # lambda-assignment + "E741", # ambiguous-variable-name + "PLR0911", # too-many-return-statements + "PLR0912", # too-many-branches + "PLR0913", # too-many-arguments + "PLR0915", # too-many-statements + "PLR2004", # magic-value-comparison + "PLR5501", # collapsible-else-if + "PLW2901", # redefined-loop-name +] + +[tool.ruff.lint.per-file-ignores] +"dyson/__init__.py" = [ + "D205", # missing-blank-line-after-summary + "D212", # multi-line-summary-first-line + "D415", # missing-terminal-punctuation +] +"dyson/**/__init__.py" = [ + "I001", # unsorted-imports + "F401", # unused-import ] -include = "dyson" +"tests/**/*.py" = [ + "D103", # undocumented-public-function +] + +[tool.ruff.lint.pydocstyle] +convention = "google" + +[tool.ruff.lint.isort] +known-first-party = ["dyson"] + +[tool.mypy] +python_version = "3.12" exclude = """ /( | __pycache__ @@ -64,17 +106,13 @@ exclude = """ )/ """ -[tool.isort] -atomic = true -profile = "black" -line_length = 100 -src_paths = [ - "dyson", -] -skip_glob = [ - "*__pycache__*", - "*__init__*", -] +[[tool.mypy.overrides]] +module = "scipy.*" +ignore_missing_imports = true + +[[tool.mypy.overrides]] +module = "pyscf.*" +ignore_missing_imports = true [tool.coverage.run] branch = true @@ -96,12 +134,11 @@ exclude_lines = [ directory = "cov_html" [tool.pytest.ini_options] -addopts = "-m 'not slow'" +filterwarnings = [ + "ignore::DeprecationWarning", +] testpaths = [ + "dyson", "tests", ] -markers = [ - "slow", - "reference", - "regression", -] +addopts = "--doctest-modules --cov=dyson" diff --git a/tests/__init__.py b/tests/__init__.py index e69de29..7a57fb0 100644 --- a/tests/__init__.py +++ b/tests/__init__.py @@ -0,0 +1 @@ +"""Test module for :mod:`~dyson`.""" diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 0000000..5fbd1ad --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,190 @@ +"""Configuration for :mod:`pytest`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest +from pyscf import gto, scf + +from dyson import numpy as np +from dyson.expressions import ADC2, CCSD, FCI, HF, TDAGW, ADC2x +from dyson.representations.lehmann import Lehmann +from dyson.representations.spectral import Spectral +from dyson.solvers import Exact + +if TYPE_CHECKING: + from typing import Callable, Hashable + + from dyson.expressions.expression import BaseExpression, ExpressionCollection + from dyson.typing import Array + + ExactGetter = Callable[[scf.hf.RHF, type[BaseExpression]], Exact] + + +MOL_CACHE = { + "h2-631g": gto.M( + atom="H 0 0 0; H 0 0 1.4", + basis="6-31g", + verbose=0, + ), + "h2o-sto3g": gto.M( + atom="O 0 0 0; H 0.758602 0.504284 0; H 0.758602 -0.504284 0", + basis="sto-3g", + verbose=0, + ), + "he-ccpvdz": gto.M( + atom="He 0 0 0", + basis="cc-pvdz", + verbose=0, + ), +} + +MF_CACHE = { + "h2-631g": scf.RHF(MOL_CACHE["h2-631g"]).run(conv_tol=1e-12), + "h2o-sto3g": scf.RHF(MOL_CACHE["h2o-sto3g"]).run(conv_tol=1e-12), + "he-ccpvdz": scf.RHF(MOL_CACHE["he-ccpvdz"]).run(conv_tol=1e-12), +} + +for key, mf in MF_CACHE.items(): + mo = mf.stability()[0] + dm = mf.make_rdm1(mo, mf.mo_occ) + mf = mf.run(dm) + +METHODS = [HF, CCSD, FCI, ADC2, ADC2x, TDAGW] +METHOD_NAMES = ["HF", "CCSD", "FCI", "ADC2", "ADC2x", "TDAGW"] + + +def pytest_generate_tests(metafunc): # type: ignore + if "mf" in metafunc.fixturenames: + metafunc.parametrize("mf", MF_CACHE.values(), ids=MF_CACHE.keys()) + if "expression_cls" in metafunc.fixturenames: + expressions = [] + ids = [] + for method, name in zip(METHODS, METHOD_NAMES): + for expression in method._classes: + expressions.append(expression) + ids.append(expression.__name__) + metafunc.parametrize("expression_cls", expressions, ids=ids) + if "expression_method" in metafunc.fixturenames: + expressions = [] + ids = [] + for method, name in zip(METHODS, METHOD_NAMES): + expressions.append(method) + ids.append(name) + metafunc.parametrize("expression_method", expressions, ids=ids) + + +class Helper: + """Helper class for tests.""" + + @staticmethod + def are_equal_arrays(moment1: Array, moment2: Array, tol: float = 1e-8) -> bool: + """Check if two arrays are equal to within a threshold.""" + print( + f"Error in {object.__repr__(moment1)} and {object.__repr__(moment2)}: " + f"{np.max(np.abs(moment1 - moment2))}" + ) + return np.allclose(moment1, moment2, atol=tol) + + @staticmethod + def have_equal_moments( + lehmann1: Lehmann | Array, lehmann2: Lehmann | Array, num: int, tol: float = 1e-8 + ) -> bool: + """Check if two :class:`Lehmann` objects have equal moments to within a threshold.""" + moments1 = lehmann1.moments(range(num)) if isinstance(lehmann1, Lehmann) else lehmann1 + moments2 = lehmann2.moments(range(num)) if isinstance(lehmann2, Lehmann) else lehmann2 + checks: list[bool] = [] + for i, (m1, m2) in enumerate(zip(moments1, moments2)): + errors = np.abs(m1 - m2) + errors_scaled = errors / np.maximum(np.max(np.abs(m1)), 1.0) + checks.append(bool(np.all(errors_scaled < tol))) + print( + f"Error in moment {i} of {object.__repr__(lehmann1)} and " + f"{object.__repr__(lehmann2)}: {np.max(errors_scaled)} ({np.max(errors)})" + ) + return all(checks) + + @staticmethod + def recovers_greens_function( + static: Array, + self_energy: Lehmann, + greens_function: Lehmann, + num: int = 2, + tol: float = 1e-8, + ) -> bool: + """Check if a self-energy recovers the Green's function to within a threshold.""" + overlap = greens_function.moment(0) + greens_function_other = Lehmann( + *self_energy.diagonalise_matrix_with_projection(static, overlap=overlap) + ) + return Helper.have_equal_moments(greens_function, greens_function_other, num, tol=tol) + + @staticmethod + def has_orthonormal_couplings(greens_function: Lehmann, tol: float = 1e-8) -> bool: + """Check if the Green's function Dyson orbitals are orthonormal to within a threshold.""" + return Helper.are_equal_arrays( + greens_function.moment(0), np.eye(greens_function.nphys), tol=tol + ) + + +@pytest.fixture(scope="session") +def helper() -> Helper: + """Fixture for the :class:`Helper` class.""" + return Helper() + + +_EXACT_CACHE: dict[Hashable, Exact] = {} + + +def get_exact(mf: scf.hf.RHF, expression_cls: type[BaseExpression]) -> Exact: + """Get the exact solver for a given mean-field object and expression.""" + key = (mf.__class__, mf.mol.dumps(), expression_cls) + if key not in _EXACT_CACHE: + expression = expression_cls.from_mf(mf) + exact = Exact.from_expression(expression) + exact.kernel() + _EXACT_CACHE[key] = exact + + exact = _EXACT_CACHE[key] + assert exact.result is not None + + return exact + + +@pytest.fixture(scope="session") +def exact_cache() -> ExactGetter: + """Fixture for a getter function for cached :class:`Exact` classes.""" + return get_exact + + +def _get_central_result( + helper: Helper, + mf: scf.hf.RHF, + expression_method: ExpressionCollection, + exact_cache: ExactGetter, + allow_hermitian: bool = True, +) -> Spectral: + """Get the central result for the given mean-field method.""" + if "dyson" in expression_method: + expression = expression_method.dyson.from_mf(mf) + if expression.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + if not expression.hermitian and not allow_hermitian: + pytest.skip("Skipping test for non-Hermitian Hamiltonian with negative weights") + exact = exact_cache(mf, expression_method.dyson) + assert exact.result is not None + return exact.result + + # Combine hole and particle results + expression_h = expression_method.h.from_mf(mf) + expression_p = expression_method.p.from_mf(mf) + if expression_h.nconfig > 1024 or expression_p.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + if not expression_h.hermitian and not allow_hermitian: + pytest.skip("Skipping test for non-Hermitian Hamiltonian with negative weights") + exact_h = exact_cache(mf, expression_method.h) + exact_p = exact_cache(mf, expression_method.p) + assert exact_h.result is not None + assert exact_p.result is not None + return Spectral.combine(exact_h.result, exact_p.result) diff --git a/tests/expressions/test_ccsd.py b/tests/expressions/test_ccsd.py deleted file mode 100644 index af5574c..0000000 --- a/tests/expressions/test_ccsd.py +++ /dev/null @@ -1,105 +0,0 @@ -""" -Tests for CCSD. -""" - -import unittest -import pytest - -from pyscf import gto, scf, cc, lib -from pyscf.cc.momgfccsd import MomGFCCSD -import numpy as np -import scipy.linalg - -from dyson import util, Lehmann, NullLogger -from dyson import MBLGF, Davidson -from dyson.expressions import CCSD - - -@pytest.mark.regression -class CCSD_Tests(unittest.TestCase): - """ - Test the `CCSD` expressions. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="cc-pvdz", verbose=0) - mf = scf.RHF(mol).run() - ccsd = cc.CCSD(mf).run() - ccsd.solve_lambda() - cls.mf, cls.ccsd = mf, ccsd - - @classmethod - def tearDownClass(cls): - del cls.mf, cls.ccsd - - def test_ip_ccsd(self): - mf = self.mf - - ccsd = CCSD["1h"](mf) - diag = ccsd.diagonal() - matvec = ccsd.apply_hamiltonian - - solver = Davidson(matvec, diag, nroots=5, nphys=ccsd.nocc, log=NullLogger()) - solver.conv_tol = 1e-10 - solver.kernel() - ip1 = -solver.get_greens_function().energies[-3:][::-1] - - ip2 = self.ccsd.ipccsd(nroots=5)[0] - - self.assertAlmostEqual(ip1[0], ip2[0], 7) - self.assertAlmostEqual(ip1[1], ip2[1], 7) - self.assertAlmostEqual(ip1[2], ip2[2], 7) - - def test_ea_ccsd(self): - mf = self.mf - - ccsd = CCSD["1p"](mf) - diag = ccsd.diagonal() - matvec = ccsd.apply_hamiltonian - - solver = Davidson(matvec, diag, nroots=5, nphys=ccsd.nocc, log=NullLogger()) - solver.conv_tol = 1e-10 - solver.kernel() - ea1 = solver.get_greens_function().energies[:3] - - ea2 = self.ccsd.eaccsd(nroots=5)[0] - - self.assertAlmostEqual(ea1[0], ea2[0], 6) - self.assertAlmostEqual(ea1[1], ea2[1], 6) - self.assertAlmostEqual(ea1[2], ea2[2], 6) - - def test_momgfccsd(self): - mf = self.mf - ccsd = self.ccsd - nmom = 6 - - expr = CCSD["1h"](mf, t1=ccsd.t1, t2=ccsd.t2, l1=ccsd.l1, l2=ccsd.l2) - th = expr.build_gf_moments(nmom) - expr = CCSD["1p"](mf, t1=ccsd.t1, t2=ccsd.t2, l1=ccsd.l1, l2=ccsd.l2) - tp = expr.build_gf_moments(nmom) - - solverh = MBLGF(th, hermitian=False, log=NullLogger()) - solverh.kernel() - gfh = solverh.get_greens_function() - solverp = MBLGF(tp, hermitian=False, log=NullLogger()) - solverp.kernel() - gfp = solverp.get_greens_function() - gf = gfh + gfp - - grid = np.linspace(-5, 5, 1024) - eta = 1e-1 - sf = util.build_spectral_function(gf.energies, gf.couplings, grid, eta=eta) - - momgfcc = MomGFCCSD(ccsd, ((nmom-2)//2, (nmom-2)//2)) - eh, vh, ep, vp = momgfcc.kernel() - e = np.concatenate((eh, ep), axis=0) - v = np.concatenate((vh[0], vp[0]), axis=1) - u = np.concatenate((vh[1], vp[1]), axis=1) - sf_ref = util.build_spectral_function(e, (v, u), grid, eta=eta) - - np.testing.assert_allclose(sf, sf_ref, atol=1e-5) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/expressions/test_gw.py b/tests/expressions/test_gw.py deleted file mode 100644 index be53f1f..0000000 --- a/tests/expressions/test_gw.py +++ /dev/null @@ -1,70 +0,0 @@ -""" -Tests for GW. -""" - -import unittest -import pytest - -from pyscf import gto, dft, adc, agf2, lib -import numpy as np -import scipy.linalg - -try: - import momentGW -except ImportError: - momentGW = None - -from dyson import util, Lehmann, NullLogger -from dyson import MBLSE, MixedMBLSE, Davidson -from dyson.expressions import GW - - -@pytest.mark.regression -@pytest.mark.skipif(momentGW is None, reason="Moment GW tests require momentGW") -class GW_Tests(unittest.TestCase): - """ - Test the `GW` expressions. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="6-31g", verbose=0) - mf = dft.RKS(mol, xc="hf").density_fit().run() - cls.mf = mf - - @classmethod - def tearDownClass(cls): - del cls.mf - - def test_moment_gw(self): - gw = GW["Dyson"](self.mf) - static = gw.get_static_part() - th, tp = gw.build_se_moments(9) - - solverh = MBLSE(static, th, log=NullLogger()) - solverp = MBLSE(static, tp, log=NullLogger()) - solver = MixedMBLSE(solverh, solverp) - solver.kernel() - - gf = solver.get_greens_function() - gf = gf.physical() - - import momentGW - gw_ref = momentGW.GW(self.mf) - _, gf_ref, se_ref, _ = gw_ref.kernel(9) - gf_ref = gf_ref.physical(weight=0.1) - - np.testing.assert_allclose(gf_ref.moment(0), gf.moment(0), rtol=1e10, atol=1e-10) - np.testing.assert_allclose(gf_ref.moment(1), gf.moment(1), rtol=1e10, atol=1e-10) - - def test_tda_gw(self): - gw = GW["Dyson"](self.mf) - gw.polarizability = "dtda" - static = gw.get_static_part() - matvec = lambda v: gw.apply_hamiltonian(v, static=static) - diag = gw.diagonal(static=static) - # TODO - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/expressions/test_mp2.py b/tests/expressions/test_mp2.py deleted file mode 100644 index e7d659f..0000000 --- a/tests/expressions/test_mp2.py +++ /dev/null @@ -1,155 +0,0 @@ -""" -Tests for MP2. -""" - -import unittest -import pytest - -from pyscf import gto, scf, adc, agf2, lib -import numpy as np -import scipy.linalg - -from dyson import util, Lehmann, NullLogger -from dyson import MBLSE, MixedMBLSE, DensityRelaxation, SelfConsistent, Davidson -from dyson.expressions import MP2 - - -@pytest.mark.regression -class MP2_Tests(unittest.TestCase): - """ - Test the `MP2` expressions. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="cc-pvdz", verbose=0) - mf = scf.RHF(mol).run() - gf2 = agf2.AGF2(mf) - gf2.conv_tol_rdm1 = gf2.conv_tol_nelec = gf2.conv_tol = 1e-10 - gf2.kernel() - cls.mf, cls.gf2 = mf, gf2 - - @classmethod - def tearDownClass(cls): - del cls.mf, cls.gf2 - - def test_agf2(self): - # Tests AGF2 implementation using `dyson`, as done in - # `examples/30-agf2.py`. - - mf = self.mf - - def get_fock(rdm1_mo): - rdm1_ao = np.linalg.multi_dot((mf.mo_coeff, rdm1_mo, mf.mo_coeff.T)) - fock_ao = mf.get_fock(dm=rdm1_ao) - fock_mo = np.linalg.multi_dot((mf.mo_coeff.T, fock_ao, mf.mo_coeff)) - return fock_mo - - diis = lib.diis.DIIS() - - def get_se(gf, se_prev=None): - mo_energy, mo_coeff, mo_occ = gf.as_orbitals(mo_coeff=mf.mo_coeff, occupancy=2) - fock = get_fock(gf.occupied().moment(0) * 2) - - mp2 = MP2["Dyson"](mf, mo_energy=mo_energy, mo_coeff=mo_coeff, mo_occ=mo_occ) - th, tp = mp2.build_se_moments(2) - th = lib.einsum("...ij,pi,qj->...pq", th, gf.couplings, gf.couplings) - tp = lib.einsum("...ij,pi,qj->...pq", tp, gf.couplings, gf.couplings) - th, tp = diis.update(np.array([th, tp]), xerr=None) - - solverh = MBLSE(fock, th, log=NullLogger()) - solverp = MBLSE(fock, tp, log=NullLogger()) - solver = MixedMBLSE(solverh, solverp) - solver.kernel() - - return solver.get_self_energy() - - gf = Lehmann(mf.mo_energy, np.eye(mf.mo_energy.size)) - - solver = SelfConsistent( - get_se, - get_fock, - gf, - relax_solver=DensityRelaxation, - conv_tol=1e-10, - log=NullLogger(), - ) - solver.kernel() - - ip1 = -solver.get_greens_function().occupied().energies[-1] - ea1 = solver.get_greens_function().virtual().energies[0] - - ip2 = -self.gf2.gf.get_occupied().energy[-1] - ea2 = self.gf2.gf.get_virtual().energy[0] - - self.assertAlmostEqual(ip1, ip2, 8) - self.assertAlmostEqual(ea1, ea2, 8) - - def test_ip_adc2(self): - mf = self.mf - - mp2 = MP2["1h"](mf) - static = mp2.get_static_part() - diag = mp2.diagonal(static=static) - matvec = lambda v: mp2.apply_hamiltonian(v, static=static) - - solver = Davidson(matvec, diag, nroots=5, nphys=mp2.nocc, log=NullLogger()) - solver.conv_tol = 1e-10 - solver.kernel() - ip1 = -solver.get_greens_function().energies[-3:][::-1] - - adc2 = adc.ADC(mf) - ip2 = adc2.kernel(nroots=5)[0] - - self.assertAlmostEqual(ip1[0], ip2[0], 8) - self.assertAlmostEqual(ip1[1], ip2[1], 8) - self.assertAlmostEqual(ip1[2], ip2[2], 8) - - def test_ea_adc2(self): - mf = self.mf - - mp2 = MP2["1p"](mf) - static = mp2.get_static_part() - diag = mp2.diagonal(static=static) - matvec = lambda v: mp2.apply_hamiltonian(v, static=static) - - solver = Davidson(matvec, diag, nroots=5, nphys=mp2.nocc, log=NullLogger()) - solver.conv_tol = 1e-10 - solver.kernel() - ea1 = solver.get_greens_function().energies[:3] - - adc2 = adc.ADC(mf) - adc2.method_type = "ea" - ea2 = adc2.kernel(nroots=5)[0] - - self.assertAlmostEqual(ea1[0], ea2[0], 8) - self.assertAlmostEqual(ea1[1], ea2[1], 8) - self.assertAlmostEqual(ea1[2], ea2[2], 8) - - def test_gf_moments(self): - mf = self.mf - - mp2 = MP2["1h"](mf) - diag = mp2.diagonal() - matrix = np.array([mp2.apply_hamiltonian(v) for v in np.eye(diag.size)]) - w, v = np.linalg.eig(matrix) - v = (v, np.linalg.inv(v).T) - v = (v[0][:mp2.nocc], v[1][:mp2.nocc]) - gf = Lehmann(w, v) - - t1 = gf.moment(range(6)) - t2 = mp2.build_gf_moments(6) - np.testing.assert_allclose(t1, t2, atol=1e-8) - - a = (np.max(w) - np.min(w)) / (2.0 - 1e-3) - b = (np.max(w) + np.min(w)) / 2.0 - t1 = gf.chebyshev_moment(range(20), scaling=(a, b)) - t2 = mp2.build_gf_chebyshev_moments(20, scaling=(a, b)) - for i in range(20): - assert np.allclose(t1[i], t2[i]), i - np.testing.assert_allclose(t1, t2, atol=1e-8) - - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/solvers/test_aufbau.py b/tests/solvers/test_aufbau.py deleted file mode 100644 index 22b43f1..0000000 --- a/tests/solvers/test_aufbau.py +++ /dev/null @@ -1,83 +0,0 @@ -""" -Tests for AufbauPrinciple. -""" - -import unittest -import pytest - -from pyscf import gto, scf, agf2 -import numpy as np - -from dyson import util, AufbauPrinciple, AufbauPrincipleBisect, NullLogger, MBLGF -from dyson.lehmann import Lehmann - - -@pytest.mark.regression -class AufbauPrinciple_Tests(unittest.TestCase): - """ - Test the `AufbauPrinciple` solver. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="6-31g", verbose=0) - mf = scf.RHF(mol).run() - f = np.diag(mf.mo_energy) - se = agf2.AGF2(mf, nmom=(None, None)).build_se().get_virtual() - cls.mf, cls.mol = mf, mol - cls.f, cls.se = f, se - - @classmethod - def tearDownClass(cls): - del cls.mf, cls.mol, cls.f, cls.se - - def test_hf(self): - gf = Lehmann(self.mf.mo_energy, np.eye(self.mf.mo_energy.size)) - - solver = AufbauPrinciple(gf, self.mol.nelectron, log=NullLogger()) - solver.kernel() - - self.assertTrue(solver.converged) - self.assertAlmostEqual(solver.error, 0.0, 7) - self.assertAlmostEqual(solver.homo, self.mf.mo_energy[self.mf.mo_occ > 0].max(), 7) - self.assertAlmostEqual(solver.lumo, self.mf.mo_energy[self.mf.mo_occ == 0].min(), 7) - - def test_agf2(self): - f = self.f - e = self.se.energy - v = self.se.coupling - h = np.block([[f, v], [v.T, np.diag(e)]]) - w, v = np.linalg.eigh(h) - v = v[:f.shape[0]] - gf = Lehmann(w, v) - - solver = AufbauPrinciple(gf, self.mol.nelectron, log=NullLogger()) - solver.kernel() - - self.assertTrue(solver.converged) - self.assertAlmostEqual(solver.error, 0.017171058925, 7) - - -@pytest.mark.regression -class AufbauPrincipleBisect_Tests(unittest.TestCase): - def test_wrt_AufbauPrinciple(self): - for i in range(10): - n = 100 - moms = np.random.random((16, n, n)) - moms = moms + moms.transpose(0, 2, 1) - mblgf = MBLGF(moms) - mblgf.kernel() - gf = mblgf.get_greens_function() - nelec = 25 - - solver = AufbauPrinciple(gf, nelec, occupancy=2, log=NullLogger()) - solver.kernel() - - solver_bisect = AufbauPrincipleBisect(gf, nelec, occupancy=2, log=NullLogger()) - solver_bisect.kernel() - - assert np.allclose(solver.chempot, solver_bisect.chempot) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/solvers/test_auxiliary_shift.py b/tests/solvers/test_auxiliary_shift.py deleted file mode 100644 index 322fd67..0000000 --- a/tests/solvers/test_auxiliary_shift.py +++ /dev/null @@ -1,44 +0,0 @@ -""" -Tests for AuxiliaryShift. -""" - -import unittest -import pytest - -from pyscf import gto, scf, agf2 -import numpy as np - -from dyson import util, AuxiliaryShift, NullLogger -from dyson.lehmann import Lehmann - - -@pytest.mark.regression -class AuxiliaryShift_Tests(unittest.TestCase): - """ - Test the `AuxiliaryShift` solver. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="cc-pvdz", verbose=0) - mf = scf.RHF(mol).run() - f = np.diag(mf.mo_energy) - se = agf2.AGF2(mf, nmom=(None, None)).build_se().get_virtual() - cls.mf, cls.mol = mf, mol - cls.f, cls.se = f, se - - @classmethod - def tearDownClass(cls): - del cls.mf, cls.mol, cls.f, cls.se - - def test_agf2(self): - solver = AuxiliaryShift(self.f, self.se, self.mol.nelectron, log=NullLogger()) - solver.conv_tol = 1e-6 - solver.kernel() - - self.assertTrue(solver.converged) - self.assertAlmostEqual(solver.error, 0.0, 5) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/solvers/test_davidson.py b/tests/solvers/test_davidson.py deleted file mode 100644 index 35c52cf..0000000 --- a/tests/solvers/test_davidson.py +++ /dev/null @@ -1,85 +0,0 @@ -""" -Tests for the Davidson solver. -""" - -import unittest -import pytest - -from pyscf import gto, scf, agf2 -from pyscf.lib.linalg_helper import pick_real_eigs -import numpy as np -import scipy.linalg - -from dyson import NullLogger, Davidson - - -@pytest.mark.regression -class Davidson_Tests(unittest.TestCase): - """ - Test for the `Davidson` solver. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="6-31g", verbose=0) - mf = scf.RHF(mol).run() - f = np.diag(mf.mo_energy[mf.mo_occ > 0]) - se = agf2.AGF2(mf, nmom=(None, None)).build_se().get_virtual() - e = se.energy - v = se.coupling[mf.mo_occ > 0] - h = np.block([[f, v], [v.T, np.diag(e)]]) - cls.e, cls.v, cls.f, cls.h = e, v, f, h - cls.w0, cls.v0 = np.linalg.eigh(h) - - @classmethod - def tearDownClass(cls): - del cls.e, cls.v, cls.f, cls.h - del cls.w0, cls.v0 - - def test_hermitian(self): - m = lambda v: np.dot(self.h, v) - d = np.diag(self.h) - solver = Davidson(m, d, picker=pick_real_eigs, log=NullLogger()) - w, v = solver.kernel() - self.assertAlmostEqual(w[0], self.w0[0], 8) - self.assertAlmostEqual(w[1], self.w0[1], 8) - - def test_hermitian_guess(self): - m = lambda v: np.dot(self.h, v) - d = np.diag(self.h) - guess = np.zeros((1, d.size)) - guess[0, np.argmin(d)] = 1 - solver = Davidson(m, d, picker=pick_real_eigs, guess=guess, nroots=1, log=NullLogger()) - w, v = solver.kernel() - self.assertAlmostEqual(w[0], self.w0[0], 8) - - def test_nonhermitian(self): - pert = (np.random.random(self.v.shape) - 0.5) / 100 - h = np.block([[self.f, self.v+pert], [self.v.T, np.diag(self.e)]]) - m = lambda v: np.dot(h, v) - d = np.diag(h) - solver = Davidson(m, d, picker=pick_real_eigs, hermitian=False, log=NullLogger()) - w, v = solver.kernel() - w0, v0 = np.linalg.eig(h) - w0 = w0[np.argsort(w0.real)] - self.assertAlmostEqual(w[0], w0[0], 8) - self.assertAlmostEqual(w[1], w0[1], 8) - - def test_nonhermitian_guess(self): - pert = (np.random.random(self.v.shape) - 0.5) / 100 - h = np.block([[self.f, self.v+pert], [self.v.T, np.diag(self.e)]]) - m = lambda v: np.dot(h, v) - d = np.diag(h) - guess = np.zeros((1, d.size)) - guess[0, np.argmin(d)] = 1 - solver = Davidson(m, d, picker=pick_real_eigs, hermitian=False, guess=guess, log=NullLogger()) - w, v = solver.kernel() - w0, v0 = np.linalg.eig(h) - w0 = w0[np.argsort(w0.real)] - self.assertAlmostEqual(w[0], w0[0], 8) - self.assertAlmostEqual(w[1], w0[1], 8) - - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/solvers/test_density.py b/tests/solvers/test_density.py deleted file mode 100644 index afb3300..0000000 --- a/tests/solvers/test_density.py +++ /dev/null @@ -1,54 +0,0 @@ -""" -Tests for DensityRelaxation. -""" - -import unittest -import pytest - -from pyscf import gto, scf, agf2, lib -import numpy as np - -from dyson import util, DensityRelaxation, NullLogger -from dyson.lehmann import Lehmann - - -@pytest.mark.regression -class DensityRelaxation_Tests(unittest.TestCase): - """ - Test the `DensityRelaxation` solver. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="6-31g", verbose=0) - mf = scf.RHF(mol).run() - f = np.diag(mf.mo_energy) - se = agf2.AGF2(mf, nmom=(None, 2)).build_se() - cls.mf, cls.mol = mf, mol - cls.f, cls.se = f, se - - @classmethod - def tearDownClass(cls): - del cls.mf, cls.mol, cls.f, cls.se - - def test_agf2(self): - def get_fock(rdm1): - rdm1_ao = np.linalg.multi_dot((self.mf.mo_coeff, rdm1, self.mf.mo_coeff.T)) - fock_ao = self.mf.get_fock(dm=rdm1_ao) - fock = np.linalg.multi_dot((self.mf.mo_coeff.T, fock_ao, self.mf.mo_coeff)) - return fock - - solver = DensityRelaxation(get_fock, self.se, self.mol.nelectron, log=NullLogger()) - solver.conv_tol = 1e-12 - solver.chempot_solver.conv_tol = 1e-8 - solver.kernel() - - rdm1 = solver.gf_res.occupied().moment(0) * 2 - fock = get_fock(rdm1) - - self.assertTrue(solver.converged) - self.assertAlmostEqual(lib.fp(np.linalg.eigvalsh(fock)), -3.7636639657, 7) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/solvers/test_downfolded.py b/tests/solvers/test_downfolded.py deleted file mode 100644 index 85e3885..0000000 --- a/tests/solvers/test_downfolded.py +++ /dev/null @@ -1,75 +0,0 @@ -""" -Tests for the self-consistent solvers. -""" - -import unittest -import pytest - -from pyscf import gto, scf, agf2 -import numpy as np -import scipy.linalg - -from dyson import NullLogger, Downfolded - - -@pytest.mark.regression -class Downfolded_Tests(unittest.TestCase): - """ - Test the `Downfolded` solver. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="6-31g", verbose=0) - mf = scf.RHF(mol).run() - f = np.diag(mf.mo_energy) - se = agf2.AGF2(mf, nmom=(None, None)).build_se().get_virtual() - e = se.energy - v = se.coupling - h = np.block([[f, v], [v.T, np.diag(e)]]) - cls.e, cls.v, cls.f, cls.h = e, v, f, h - cls.w0, cls.v0 = np.linalg.eigh(h) - - @classmethod - def tearDownClass(cls): - del cls.e, cls.v, cls.f, cls.h - del cls.w0, cls.v0 - - def test_orbital_target(self): - m = lambda w: np.einsum("pk,qk,k->pq", self.v, self.v, 1/(w-self.e)) - solver = Downfolded(self.f, m, log=NullLogger()) - - solver.target = 0 - w, v = solver.kernel() - self.assertAlmostEqual(w[0], self.w0[0], 8) - - solver.target = 1 - w, v = solver.kernel() - self.assertAlmostEqual(w[1], self.w0[1], 8) - - solver.target = 2 - w, v = solver.kernel() - self.assertAlmostEqual(w[2], self.w0[2], 8) - - solver.target = 3 - w, v = solver.kernel() - self.assertAlmostEqual(w[3], self.w0[3], 8) - - def test_min_target(self): - m = lambda w: np.einsum("pk,qk,k->pq", self.v, self.v, 1/(w-self.e)) - solver = Downfolded(self.f, m, log=NullLogger()) - solver.target = "min" - w, v = solver.kernel() - self.assertAlmostEqual(w[0], self.w0[0], 8) - - def test_mindif_target(self): - m = lambda w: np.einsum("pk,qk,k->pq", self.v, self.v, 1/(w-self.e)) - solver = Downfolded(self.f, m, log=NullLogger()) - solver.target = "mindif" - solver.guess = self.w0[3] + 0.01 - w, v = solver.kernel() - self.assertAlmostEqual(w[3], self.w0[3], 8) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/solvers/test_exact.py b/tests/solvers/test_exact.py deleted file mode 100644 index f02d462..0000000 --- a/tests/solvers/test_exact.py +++ /dev/null @@ -1,96 +0,0 @@ -""" -Tests for the exact solver. -""" - -import unittest -import pytest - -import numpy as np -import scipy.linalg - -from dyson import NullLogger, Exact - - -@pytest.mark.regression -class Exact_Tests(unittest.TestCase): - """ - Test the `Exact` solver. - """ - - @classmethod - def setUpClass(cls): - pass - - @classmethod - def tearDownClass(cls): - pass - - def test_real_hermitian(self): - m = np.random.random((100, 100)) - m = 0.5 * (m + m.T.conj()) - solver = Exact(m, log=NullLogger()) - w, v = solver.kernel() - m0 = np.dot(v * w[None], v.T.conj()) - np.testing.assert_almost_equal(m, m0) - - def test_complex_hermitian(self): - m = np.random.random((100, 100)) + np.random.random((100, 100)) + 1.0j - m = 0.5 * (m + m.T.conj()) - solver = Exact(m, log=NullLogger()) - w, v = solver.kernel() - m0 = np.dot(v * w[None], v.T.conj()) - np.testing.assert_almost_equal(m, m0) - - def test_real_nonhermitian(self): - m = np.random.random((100, 100)) - solver = Exact(m, hermitian=False, log=NullLogger()) - w, v = solver.kernel() - m0 = np.dot(v * w[None], np.linalg.inv(v)) - np.testing.assert_almost_equal(m, m0) - - def test_complex_nonhermitian(self): - m = np.random.random((100, 100)) + np.random.random((100, 100)) + 1.0j - solver = Exact(m, hermitian=False, log=NullLogger()) - w, v = solver.kernel() - m0 = np.dot(v * w[None], np.linalg.inv(v)) - np.testing.assert_almost_equal(m, m0) - - def test_real_hermitian_generalised(self): - m = np.random.random((100, 100)) - m = 0.5 * (m + m.T.conj()) - s = np.random.random((100, 100)) - s = 0.5 * (s + s.T.conj()) - solver = Exact(m, overlap=s, log=NullLogger()) - w, v = solver.kernel() - m0 = np.dot(v * w[None], v.T.conj()) - np.testing.assert_almost_equal(m, m0) - - def test_complex_hermitian_generalised(self): - m = np.random.random((100, 100)) + np.random.random((100, 100)) + 1.0j - m = 0.5 * (m + m.T.conj()) - s = np.random.random((100, 100)) - s = 0.5 * (s + s.T.conj()) - solver = Exact(m, overlap=s, log=NullLogger()) - w, v = solver.kernel() - m0 = np.dot(v * w[None], v.T.conj()) - np.testing.assert_almost_equal(m, m0) - - def test_real_nonhermitian_generalised(self): - m = np.random.random((100, 100)) - s = np.random.random((100, 100)) - solver = Exact(m, overlap=s, hermitian=False, log=NullLogger()) - w, v = solver.kernel() - m0 = np.dot(v * w[None], np.linalg.inv(v)) - np.testing.assert_almost_equal(m, m0) - - def test_complex_nonhermitian_generalised(self): - m = np.random.random((100, 100)) + np.random.random((100, 100)) + 1.0j - s = np.random.random((100, 100)) - solver = Exact(m, overlap=s, hermitian=False, log=NullLogger()) - w, v = solver.kernel() - m0 = np.dot(v * w[None], np.linalg.inv(v)) - np.testing.assert_almost_equal(m, m0) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/solvers/test_mblgf.py b/tests/solvers/test_mblgf.py deleted file mode 100644 index b59cca7..0000000 --- a/tests/solvers/test_mblgf.py +++ /dev/null @@ -1,101 +0,0 @@ -""" -Tests for MBLGF. -""" - -import unittest -import pytest - -from pyscf import gto, scf, agf2 -import numpy as np -import scipy.linalg - -from dyson import util, MBLGF, NullLogger - - -@pytest.mark.regression -class MBLGF_Tests(unittest.TestCase): - """ - Test the `MBLGF` solver. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="6-31g", verbose=0) - mf = scf.RHF(mol).run() - f = np.diag(mf.mo_energy[mf.mo_occ == 0]) - se = agf2.AGF2(mf, nmom=(None, None)).build_se().get_virtual() - se.coupling = se.coupling[mf.mo_occ == 0] - cls.f, cls.se = f, se - - @classmethod - def tearDownClass(cls): - del cls.f, cls.se - - def test_hermitian(self): - f = self.f - e = self.se.energy - v = self.se.coupling - h = np.block([[f, v], [v.T, np.diag(e)]]) - w0, v0 = np.linalg.eigh(h) - nmo = self.se.nphys - t = np.einsum("pk,qk,nk->npq", v0[:nmo], v0[:nmo], w0[None]**np.arange(16)[:, None]) - w0, v0 = util.remove_unphysical(v0, nmo, eigvals=w0, tol=1e-2) - - solver = MBLGF(t, max_cycle=2, log=NullLogger()) - w, v = solver.kernel() - w, v = util.remove_unphysical(v, nmo, eigvals=w, tol=1e-2) - error = solver._check_moment_error() - self.assertAlmostEqual(error, 0.0, 10) - self.assertAlmostEqual(w[0], w0[0], 1) - - solver = MBLGF(t, max_cycle=3, log=NullLogger()) - w, v = solver.kernel() - w, v = util.remove_unphysical(v, nmo, eigvals=w, tol=1e-2) - error = solver._check_moment_error() - self.assertAlmostEqual(error, 0.0, 10) - self.assertAlmostEqual(w[0], w0[0], 2) - - solver = MBLGF(t, max_cycle=4, log=NullLogger()) - w, v = solver.kernel() - w, v = util.remove_unphysical(v, nmo, eigvals=w, tol=1e-2) - error = solver._check_moment_error() - self.assertAlmostEqual(error, 0.0, 10) - self.assertAlmostEqual(w[0], w0[0], 3) - - def test_nonhermitian(self): - f = self.f - e = self.se.energy - v = self.se.coupling - pert = (np.ones(v.shape) - 0.5) / 200 - h = np.block([[f, v+pert], [v.T, np.diag(e)]]) - w0, v0 = np.linalg.eig(h) - mask = np.argsort(w0.real) - w0, v0 = w0[mask], v0[:, mask] - v0i = np.linalg.inv(v0).T - nmo = self.se.nphys - t = np.einsum("pk,qk,nk->npq", v0[:nmo], v0i[:nmo], w0[None]**np.arange(16)[:, None]) - w0, v0 = util.remove_unphysical(v0, nmo, eigvals=w0, tol=1e-2) - - solver = MBLGF(t, max_cycle=2, log=NullLogger()) - w, v = solver.kernel() - w, v = util.remove_unphysical(v, nmo, eigvals=w, tol=1e-2) - error = solver._check_moment_error() - self.assertAlmostEqual(error, 0.0, 10) - - # FIXME: these tests are incredibly flaky - - #solver = MBLGF(t, max_cycle=3, log=NullLogger()) - #w, v = solver.kernel() - #w, v = util.remove_unphysical(v, nmo, eigvals=w, tol=1e-2) - #error = solver._check_moment_error() - #self.assertAlmostEqual(error, 0.0, 10) - - #solver = MBLGF(t, max_cycle=4, log=NullLogger()) - #w, v = solver.kernel() - #w, v = util.remove_unphysical(v, nmo, eigvals=w, tol=1e-2) - #error = solver._check_moment_error() - #self.assertAlmostEqual(error, 0.0, 10) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/solvers/test_mblse.py b/tests/solvers/test_mblse.py deleted file mode 100644 index c0cdf76..0000000 --- a/tests/solvers/test_mblse.py +++ /dev/null @@ -1,83 +0,0 @@ -""" -Tests for MBLSE. -""" - -import unittest -import pytest - -from pyscf import gto, scf, agf2 -import numpy as np -import scipy.linalg - -from dyson import util, MBLSE, NullLogger - - -@pytest.mark.regression -class MBLSE_Tests(unittest.TestCase): - """ - Test the `MBLSE` solver. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="6-31g", verbose=0) - mf = scf.RHF(mol).run() - f = np.diag(mf.mo_energy) - se = agf2.AGF2(mf, nmom=(None, None)).build_se().get_virtual() - cls.f, cls.se = f, se - - @classmethod - def tearDownClass(cls): - del cls.f, cls.se - - def test_hermitian(self): - f = self.f - e = self.se.energy - v = self.se.coupling - h = np.block([[f, v], [v.T, np.diag(e)]]) - t = np.einsum("pk,qk,nk->npq", v, v, e[None]**np.arange(16)[:, None]) - w0, v0 = np.linalg.eigh(h) - - solver = MBLSE(f, t, max_cycle=0, log=NullLogger()) - w, v = solver.kernel() - error = solver._check_moment_error() - self.assertAlmostEqual(error, 0.0, 10) - self.assertAlmostEqual(w[0], w0[0], 3) - - solver = MBLSE(f, t, max_cycle=1, log=NullLogger()) - w, v = solver.kernel() - error = solver._check_moment_error() - self.assertAlmostEqual(error, 0.0, 10) - self.assertAlmostEqual(w[0], w0[0], 5) - - solver = MBLSE(f, t, max_cycle=3, log=NullLogger()) - w, v = solver.kernel() - error = solver._check_moment_error() - self.assertAlmostEqual(error, 0.0, 10) - self.assertAlmostEqual(w[0], w0[0], 7) - - def test_nonhermitian(self): - f = self.f - e = self.se.energy - v = self.se.coupling - pert = (np.ones(v.shape) - 0.5) / 100 - h = np.block([[f, v+pert], [v.T, np.diag(e)]]) - t = np.einsum("pk,qk,nk->npq", v+pert, v, e[None]**np.arange(16)[:, None]) - w0, v0 = np.linalg.eigh(h) - - solver = MBLSE(f, t, max_cycle=0, log=NullLogger()) - w, v = solver.kernel() - error = solver._check_moment_error() - self.assertAlmostEqual(error, 0.0, 10) - self.assertAlmostEqual(w[0], w0[0], 2) - - solver = MBLSE(f, t, max_cycle=1, log=NullLogger()) - w, v = solver.kernel() - error = solver._check_moment_error() - self.assertAlmostEqual(error, 0.0, 10) - self.assertAlmostEqual(w[0], w0[0], 3) - - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_chempot.py b/tests/test_chempot.py new file mode 100644 index 0000000..cf0aeaa --- /dev/null +++ b/tests/test_chempot.py @@ -0,0 +1,98 @@ +"""Tests for :module:`~dyson.results.static.chempot`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +from dyson.solvers import AufbauPrinciple, AuxiliaryShift + +from .conftest import _get_central_result + +if TYPE_CHECKING: + from pyscf import scf + + from dyson.expressions.expression import ExpressionCollection + + from .conftest import ExactGetter, Helper + + +@pytest.mark.parametrize("method", ["direct", "bisect", "global"]) +def test_aufbau_vs_exact_solver( + helper: Helper, + mf: scf.hf.RHF, + expression_method: ExpressionCollection, + exact_cache: ExactGetter, + method: str, +) -> None: + """Test AufbauPrinciple compared to the exact solver.""" + result_exact = _get_central_result( + helper, mf, expression_method, exact_cache, allow_hermitian=method == "global" + ) + + # Solve the Hamiltonian with AufbauPrinciple + with pytest.raises(ValueError): + # Needs nelec + aufbau = AufbauPrinciple.from_self_energy( + result_exact.get_static_self_energy(), + result_exact.get_self_energy(), + method=method, + ) + aufbau = AufbauPrinciple.from_self_energy( + result_exact.get_static_self_energy(), + result_exact.get_self_energy(), + nelec=mf.mol.nelectron, + method=method, + ) + aufbau.kernel() + assert aufbau.result is not None + + # Get the Green's function and number of electrons + greens_function = aufbau.result.get_greens_function() + nelec: int = np.sum(greens_function.occupied().weights(2.0)) + + # Find the best number of electrons + best = np.min( + [ + np.abs(mf.mol.nelectron - np.sum(greens_function.mask(slice(i)).weights(2.0))) + for i in range(greens_function.naux) + ] + ) + + assert np.isclose(np.abs(mf.mol.nelectron - nelec), best) + + +def test_shift_vs_exact_solver( + helper: Helper, + mf: scf.hf.RHF, + expression_method: ExpressionCollection, + exact_cache: ExactGetter, +) -> None: + """Test AuxiliaryShift compared to the exact solver.""" + result_exact = _get_central_result( + helper, mf, expression_method, exact_cache, allow_hermitian=True + ) + + # Solve the Hamiltonian with AuxiliaryShift + with pytest.raises(ValueError): + # Needs nelec + solver = AuxiliaryShift.from_self_energy( + result_exact.get_static_self_energy(), + result_exact.get_self_energy(), + ) + solver = AuxiliaryShift.from_self_energy( + result_exact.get_static_self_energy(), + result_exact.get_self_energy(), + nelec=mf.mol.nelectron, + conv_tol=1e-9, + ) + solver.kernel() + assert solver.result is not None + + # Get the Green's function and number of electrons + greens_function = solver.result.get_greens_function() + nelec: int = np.sum(greens_function.occupied().weights(2.0)) + + assert np.abs(mf.mol.nelectron - nelec) < 1e-7 diff --git a/tests/test_corrvec.py b/tests/test_corrvec.py new file mode 100644 index 0000000..a39171a --- /dev/null +++ b/tests/test_corrvec.py @@ -0,0 +1,51 @@ +"""Tests for :module:`~dyson.solvers.dynamic.corrvec`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +from dyson.grids import RealFrequencyGrid +from dyson.solvers import CorrectionVector + +if TYPE_CHECKING: + from pyscf import scf + + from dyson.expressions.expression import BaseExpression + + from .conftest import ExactGetter, Helper + + +def test_vs_exact_solver( + helper: Helper, + mf: scf.hf.RHF, + expression_cls: type[BaseExpression], + exact_cache: ExactGetter, +) -> None: + """Test correction vector compared to the exact solver.""" + expression = expression_cls.from_mf(mf) + if expression.nconfig > 1024: # TODO: Make larger for CI runs? + pytest.skip("Skipping test for large Hamiltonian") + if expression.nsingle == (expression.nocc + expression.nvir): + pytest.skip("Skipping test for central Hamiltonian") + grid = RealFrequencyGrid.from_uniform(-2, 2, 16, 0.1) + + # Solve the Hamiltonian exactly + exact = exact_cache(mf, expression_cls) + assert exact.result is not None + gf_exact = grid.evaluate_lehmann(exact.result.get_greens_function()) + + # Solve the Hamiltonian with CorrectionVector + corrvec = CorrectionVector( + expression.apply_hamiltonian, + expression.diagonal(), + expression.nphys, + grid, + expression.get_excitation_bra, + expression.get_excitation_ket, + ) + gf = corrvec.kernel() + + assert np.allclose(gf, gf_exact) diff --git a/tests/test_cpgf.py b/tests/test_cpgf.py new file mode 100644 index 0000000..0f0ba11 --- /dev/null +++ b/tests/test_cpgf.py @@ -0,0 +1,54 @@ +"""Tests for :module:`~dyson.solvers.dynamic.cpgf`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +from dyson.expressions.hf import BaseHF +from dyson.grids import RealFrequencyGrid +from dyson.solvers import CPGF + +if TYPE_CHECKING: + from pyscf import scf + + from dyson.expressions.expression import BaseExpression + + from .conftest import ExactGetter, Helper + + +def test_vs_exact_solver( + helper: Helper, + mf: scf.hf.RHF, + expression_cls: type[BaseExpression], + exact_cache: ExactGetter, +) -> None: + """Test CPGF compared to the exact solver.""" + expression = expression_cls.from_mf(mf) + if expression.nconfig > 1024: # TODO: Make larger for CI runs? + pytest.skip("Skipping test for large Hamiltonian") + if expression.nsingle == (expression.nocc + expression.nvir): + pytest.skip("Skipping test for central Hamiltonian") + if isinstance(expression, BaseHF): + pytest.skip("Skipping test for HF Hamiltonian") + grid = RealFrequencyGrid.from_uniform(-2, 2, 16, 0.1) + + # Solve the Hamiltonian exactly + exact = exact_cache(mf, expression_cls) + assert exact.result is not None + gf_exact = grid.evaluate_lehmann(exact.result.get_greens_function(), ordering="advanced") + + # Solve the Hamiltonian with CorrectionVector + cpgf = CPGF.from_self_energy( + exact.result.get_static_self_energy(), + exact.result.get_self_energy(), + overlap=exact.result.get_overlap(), + grid=grid, + max_cycle=2048, # Converge fully for all systems + ordering="advanced", + ) + gf = cpgf.kernel() + + assert np.allclose(gf, gf_exact) diff --git a/tests/test_davidson.py b/tests/test_davidson.py new file mode 100644 index 0000000..1dc347c --- /dev/null +++ b/tests/test_davidson.py @@ -0,0 +1,180 @@ +"""Tests for :module:`~dyson.results.static.davidson`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +from dyson.representations.lehmann import Lehmann +from dyson.representations.spectral import Spectral +from dyson.solvers import Davidson + +if TYPE_CHECKING: + from pyscf import scf + + from dyson.expressions.expression import BaseExpression, ExpressionCollection + from dyson.typing import Array + + from .conftest import ExactGetter, Helper + + +def test_vs_exact_solver( + helper: Helper, + mf: scf.hf.RHF, + expression_cls: type[BaseExpression], + exact_cache: ExactGetter, +) -> None: + """Test Davidson compared to the exact solver.""" + expression = expression_cls.from_mf(mf) + if expression.nconfig > 1024: # TODO: Make larger for CI runs? + pytest.skip("Skipping test for large Hamiltonian") + if expression.nsingle == (expression.nocc + expression.nvir): + pytest.skip("Skipping test for central Hamiltonian") + bra: Array = np.array(expression.get_excitation_bras()) + ket: Array = np.array(expression.get_excitation_kets()) + + # Solve the Hamiltonian exactly + exact = exact_cache(mf, expression_cls) + assert exact.result is not None + + # Solve the Hamiltonian with Davidson + davidson = Davidson( + expression.apply_hamiltonian, + expression.diagonal(), + bra, + ket, + nroots=expression.nsingle + expression.nconfig, # Get all the roots + hermitian=expression.hermitian_upfolded, + ) + davidson.kernel() + assert davidson.result is not None + + assert davidson.matvec == expression.apply_hamiltonian + assert np.all(davidson.diagonal == expression.diagonal()) + assert davidson.nphys == expression.nphys + assert exact.matrix.shape == (davidson.nroots, davidson.nroots) + + # Get the self-energy and Green's function from the Davidson solver + static = davidson.result.get_static_self_energy() + self_energy = davidson.result.get_self_energy() + greens_function = davidson.result.get_greens_function() + + # Get the self-energy and Green's function from the exact solver + static_exact = exact.result.get_static_self_energy() + self_energy_exact = exact.result.get_self_energy() + greens_function_exact = exact.result.get_greens_function() + + if expression.hermitian_upfolded: + # Left-handed eigenvectors not converged for non-Hermitian Davidson # TODO + assert helper.are_equal_arrays(static, static_exact) + assert helper.have_equal_moments(self_energy, self_energy_exact, 4) + assert helper.have_equal_moments(greens_function, greens_function_exact, 4) + + +def test_vs_exact_solver_central( + helper: Helper, + mf: scf.hf.RHF, + expression_method: ExpressionCollection, + exact_cache: ExactGetter, +) -> None: + """Test the exact solver for central moments.""" + # Get the quantities required from the expressions + if "o" not in expression_method or "p" not in expression_method: + pytest.skip("Skipping test for Dyson only expression") + expression_h = expression_method.h.from_mf(mf) + expression_p = expression_method.p.from_mf(mf) + if expression_h.nconfig > 1024 or expression_p.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + diagonal = [expression_h.diagonal(), expression_p.diagonal()] + bra = ( + np.array(expression_h.get_excitation_bras()), + np.array(expression_p.get_excitation_bras()), + ) + ket = ( + np.array(expression_h.get_excitation_kets()), + np.array(expression_p.get_excitation_kets()), + ) + + # Solve the Hamiltonian exactly + exact_h = exact_cache(mf, expression_method.h) + exact_p = exact_cache(mf, expression_method.p) + assert exact_h.result is not None + assert exact_p.result is not None + + # Solve the Hamiltonian with Davidson + davidson_h = Davidson( + expression_h.apply_hamiltonian, + diagonal[0], + bra[0], + ket[0], + nroots=expression_h.nsingle + expression_h.nconfig, # Get all the roots + hermitian=expression_h.hermitian_upfolded, + conv_tol=1e-11, + conv_tol_residual=1e-8, + ) + davidson_h.kernel() + davidson_p = Davidson( + expression_p.apply_hamiltonian, + diagonal[1], + bra[1], + ket[1], + nroots=expression_p.nsingle + expression_p.nconfig, # Get all the roots + hermitian=expression_p.hermitian_upfolded, + conv_tol=1e-11, + conv_tol_residual=1e-8, + ) + davidson_p.kernel() + assert davidson_h.result is not None + assert davidson_p.result is not None + + # Get the self-energy and Green's function from the Davidson solver + static = davidson_h.result.get_static_self_energy() + davidson_p.result.get_static_self_energy() + self_energy = Lehmann.concatenate( + davidson_h.result.get_self_energy(), davidson_p.result.get_self_energy() + ) + greens_function = Lehmann.concatenate( + davidson_h.result.get_greens_function(), davidson_p.result.get_greens_function() + ) + + # Get the self-energy and Green's function from the exact solvers + static_exact = exact_h.result.get_static_self_energy() + exact_p.result.get_static_self_energy() + self_energy_exact = Lehmann.concatenate( + exact_h.result.get_self_energy(), exact_p.result.get_self_energy() + ) + greens_function_exact = Lehmann.concatenate( + exact_h.result.get_greens_function(), exact_p.result.get_greens_function() + ) + + if expression_h.hermitian_upfolded and expression_p.hermitian_upfolded: + # Left-handed eigenvectors not converged for non-Hermitian Davidson # TODO + assert helper.are_equal_arrays(static, static_exact) + assert helper.have_equal_moments(self_energy, self_energy_exact, 2) + + # Use the component-wise solvers + result_exact = Spectral.combine(exact_h.result, exact_p.result) + result_davidson = Spectral.combine(davidson_h.result, davidson_p.result) + + # Get the self-energy and Green's function from the Davidson solver + static = result_davidson.get_static_self_energy() + self_energy = result_davidson.get_self_energy() + greens_function = result_davidson.get_greens_function() + + # Get the self-energy and Green's function from the exact solver + static_exact = result_exact.get_static_self_energy() + self_energy_exact = result_exact.get_self_energy() + greens_function_exact = result_exact.get_greens_function() + + if expression_h.hermitian_upfolded and expression_p.hermitian_upfolded: + # Left-handed eigenvectors not converged for non-Hermitian Davidson # TODO + assert helper.are_equal_arrays(static, static_exact) + assert helper.have_equal_moments(self_energy, self_energy_exact, 2) + assert helper.are_equal_arrays(greens_function.moment(1), static) + assert helper.are_equal_arrays(greens_function_exact.moment(1), static_exact) + assert helper.recovers_greens_function(static, self_energy, greens_function) + assert helper.recovers_greens_function( + static_exact, self_energy_exact, greens_function_exact + ) + assert helper.has_orthonormal_couplings(greens_function) + assert helper.has_orthonormal_couplings(greens_function_exact) diff --git a/tests/test_density.py b/tests/test_density.py new file mode 100644 index 0000000..b87682d --- /dev/null +++ b/tests/test_density.py @@ -0,0 +1,59 @@ +"""Tests for :module:`~dyson.results.static.density`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +from dyson.representations.spectral import Spectral +from dyson.solvers import DensityRelaxation +from dyson.solvers.static.density import get_fock_matrix_function + +if TYPE_CHECKING: + from pyscf import scf + + from dyson.expressions.expression import ExpressionCollection + + from .conftest import ExactGetter, Helper + + +def test_vs_exact_solver( + helper: Helper, + mf: scf.hf.RHF, + expression_method: ExpressionCollection, + exact_cache: ExactGetter, +) -> None: + """Test DensityRelaxation compared to the exact solver.""" + if "h" not in expression_method or "p" not in expression_method: + pytest.skip("Skipping test for Dyson only expression") + expression_h = expression_method.h.from_mf(mf) + expression_p = expression_method.p.from_mf(mf) + if expression_h.nconfig > 1024 or expression_p.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + + # Solve the Hamiltonian exactly + exact_h = exact_cache(mf, expression_method.h) + exact_p = exact_cache(mf, expression_method.p) + assert exact_h.result is not None + assert exact_p.result is not None + result_exact = Spectral.combine(exact_h.result, exact_p.result) + + # Solve the Hamiltonian with DensityRelaxation + get_fock = get_fock_matrix_function(mf) + solver = DensityRelaxation.from_self_energy( + result_exact.get_static_self_energy(), + result_exact.get_self_energy(), + nelec=mf.mol.nelectron, + get_static=get_fock, + ) + solver.kernel() + assert solver.result is not None + + # Get the Green's function + greens_function = solver.result.get_greens_function() + rdm1 = greens_function.occupied().moment(0) * 2.0 + + assert solver.converged + assert np.isclose(np.trace(rdm1), mf.mol.nelectron, atol=1e-2) diff --git a/tests/test_downfolded.py b/tests/test_downfolded.py new file mode 100644 index 0000000..32be3d7 --- /dev/null +++ b/tests/test_downfolded.py @@ -0,0 +1,60 @@ +"""Tests for :module:`~dyson.results.static.downfolded`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import numpy as np +import pytest + +from dyson.representations.spectral import Spectral +from dyson.solvers import Downfolded + +if TYPE_CHECKING: + from pyscf import scf + + from dyson.expressions.expression import ExpressionCollection + + from .conftest import ExactGetter, Helper + + +def test_vs_exact_solver( + helper: Helper, + mf: scf.hf.RHF, + expression_method: ExpressionCollection, + exact_cache: ExactGetter, +) -> None: + """Test Downfolded compared to the exact solver.""" + # Get the quantities required from the expressions + if "h" not in expression_method or "p" not in expression_method: + pytest.skip("Skipping test for Dyson only expression") + expression_h = expression_method.h.from_mf(mf) + expression_p = expression_method.p.from_mf(mf) + if expression_h.nconfig > 1024 or expression_p.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + + # Solve the Hamiltonian exactly + exact_h = exact_cache(mf, expression_method.h) + exact_p = exact_cache(mf, expression_method.p) + assert exact_h.result is not None + assert exact_p.result is not None + result_exact = Spectral.combine(exact_h.result, exact_p.result) + + # Solve the Hamiltonian with Downfolded + downfolded = Downfolded.from_self_energy( + result_exact.get_static_self_energy(), + result_exact.get_self_energy(), + eta=1e-9, + conv_tol=1e-10, + ) + downfolded.kernel() + assert downfolded.result is not None + + # Get the targetted energies + guess = downfolded.guess + energy_downfolded = downfolded.result.eigvals[ + np.argmin(np.abs(downfolded.result.eigvals - guess)) + ] + energy_exact = result_exact.eigvals[np.argmin(np.abs(result_exact.eigvals - energy_downfolded))] + + assert np.abs(energy_exact - energy_downfolded) < 1e-8 diff --git a/tests/test_exact.py b/tests/test_exact.py new file mode 100644 index 0000000..e3d4993 --- /dev/null +++ b/tests/test_exact.py @@ -0,0 +1,106 @@ +"""Tests for :mod:`~dyson.solvers.static.exact`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from dyson.representations.spectral import Spectral +from dyson.solvers import Exact + +if TYPE_CHECKING: + from pyscf import scf + + from dyson.expressions.expression import BaseExpression, ExpressionCollection + + from .conftest import ExactGetter, Helper + + +def test_exact_solver( + helper: Helper, + mf: scf.hf.RHF, + expression_cls: type[BaseExpression], + exact_cache: ExactGetter, +) -> None: + """Test the exact solver.""" + # Get the quantities required from the expression + expression = expression_cls.from_mf(mf) + if expression.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + gf_moments = expression.build_gf_moments(4) + + # Solve the Hamiltonian + solver = exact_cache(mf, expression_cls) + + assert solver.result is not None + assert solver.nphys == expression.nphys + assert solver.hermitian == expression.hermitian_upfolded + + # Get the self-energy and Green's function from the solver + static = solver.result.get_static_self_energy() + self_energy = solver.result.get_self_energy() + greens_function = solver.result.get_greens_function() + + assert self_energy.nphys == expression.nphys + assert greens_function.nphys == expression.nphys + assert helper.are_equal_arrays(static, gf_moments[1]) + assert helper.have_equal_moments(greens_function, gf_moments, 4) + + # Recover the Green's function from the recovered self-energy + overlap = greens_function.moment(0) + solver = Exact.from_self_energy(static, self_energy, overlap=overlap) + solver.kernel() + assert solver.result is not None + static_other = solver.result.get_static_self_energy() + self_energy_other = solver.result.get_self_energy() + greens_function_other = solver.result.get_greens_function() + + assert helper.are_equal_arrays(static, static_other) + assert helper.have_equal_moments(self_energy, self_energy_other, 4) + assert helper.have_equal_moments(greens_function, greens_function_other, 4) + + +def test_vs_exact_solver_central( + helper: Helper, + mf: scf.hf.RHF, + expression_method: ExpressionCollection, + exact_cache: ExactGetter, +) -> None: + """Test the exact solver for central moments.""" + # Get the quantities required from the expressions + if "h" not in expression_method or "p" not in expression_method: + pytest.skip("Skipping test for Dyson only expression") + expression_h = expression_method.h.from_mf(mf) + expression_p = expression_method.p.from_mf(mf) + if expression_h.nconfig > 1024 or expression_p.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + + # Solve the Hamiltonian exactly + exact_h = exact_cache(mf, expression_method.h) + exact_p = exact_cache(mf, expression_method.p) + assert exact_h.result is not None + assert exact_p.result is not None + result_ph = Spectral.combine(exact_h.result, exact_p.result) + + # Recover the hole self-energy and Green's function + static = exact_h.result.get_static_self_energy() + self_energy = exact_h.result.get_self_energy() + greens_function = exact_h.result.get_greens_function() + + assert helper.recovers_greens_function(static, self_energy, greens_function, 4) + + # Recover the particle self-energy and Green's function + static = exact_p.result.get_static_self_energy() + self_energy = exact_p.result.get_self_energy() + greens_function = exact_p.result.get_greens_function() + + assert helper.recovers_greens_function(static, self_energy, greens_function, 4) + + # Recover the self-energy and Green's function + static = result_ph.get_static_self_energy() + self_energy = result_ph.get_self_energy() + greens_function = result_ph.get_greens_function() + + assert helper.has_orthonormal_couplings(greens_function) + assert helper.recovers_greens_function(static, self_energy, greens_function, 4) diff --git a/tests/test_expressions.py b/tests/test_expressions.py new file mode 100644 index 0000000..ec521fd --- /dev/null +++ b/tests/test_expressions.py @@ -0,0 +1,309 @@ +"""Tests for :class:`~dyson.expressions`.""" + +from __future__ import annotations + +import itertools +from typing import TYPE_CHECKING + +import numpy as np +import pyscf +import pytest + +from dyson import util +from dyson.expressions import ADC2, CCSD, FCI, HF, TDAGW, ADC2x +from dyson.solvers import Davidson, Exact + +if TYPE_CHECKING: + from pyscf import scf + + from dyson.expressions.expression import BaseExpression + + from .conftest import ExactGetter, Helper + + +def test_init(mf: scf.hf.RHF, expression_cls: type[BaseExpression]) -> None: + """Test the instantiation of the expression from a mean-field object.""" + expression = expression_cls.from_mf(mf) + assert expression.mol is mf.mol + assert expression.nphys == mf.mol.nao + assert expression.nocc == mf.mol.nelectron // 2 + assert expression.nvir == mf.mol.nao - mf.mol.nelectron // 2 + + +def test_hamiltonian(mf: scf.hf.RHF, expression_cls: type[BaseExpression]) -> None: + """Test the Hamiltonian of the expression.""" + expression = expression_cls.from_mf(mf) + if expression.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + diagonal = expression.diagonal() + hamiltonian = expression.build_matrix() + + if expression_cls in ADC2._classes: + # ADC(2)-x diagonal is set to ADC(2) diagonal in PySCF for better Davidson convergence + assert np.allclose(np.diag(hamiltonian), diagonal) + assert hamiltonian.shape == expression.shape + assert (expression.nconfig + expression.nsingle) == diagonal.size + + vector = np.random.random(expression.nconfig + expression.nsingle) + hv = expression.apply_hamiltonian_right(vector) + try: + vh = expression.apply_hamiltonian_left(vector) + except NotImplementedError: + vh = None + + assert np.allclose(hv, hamiltonian @ vector) + if vh is not None: + assert np.allclose(vh, vector @ hamiltonian) + + +def test_gf_moments(mf: scf.hf.RHF, expression_cls: type[BaseExpression]) -> None: + """Test the Green's function moments of the expression.""" + # Get the quantities required from the expression + expression = expression_cls.from_mf(mf) + if expression.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + hamiltonian = expression.build_matrix() + + # Construct the moments + moments = np.zeros((2, expression.nphys, expression.nphys)) + for i, j in itertools.product(range(expression.nphys), repeat=2): + bra = expression.get_excitation_bra(j) + ket = expression.get_excitation_ket(i) + moments[0, j, i] += bra.conj() @ ket + moments[1, j, i] += bra.conj() @ hamiltonian @ ket + + # Compare the moments to the reference + ref = expression.build_gf_moments(2) + + assert np.allclose(ref[0], moments[0]) + assert np.allclose(ref[1], moments[1]) + + +def test_static( + helper: Helper, + mf: scf.hf.RHF, + expression_cls: type[BaseExpression], + exact_cache: ExactGetter, +) -> None: + """Test the static self-energy of the expression.""" + # Get the quantities required from the expression + expression = expression_cls.from_mf(mf) + if expression.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + gf_moments = expression.build_gf_moments(2) + + # Get the static self-energy + exact = exact_cache(mf, expression_cls) + assert exact.result is not None + greens_function = exact.result.get_greens_function() + static = exact.result.get_static_self_energy() + + assert helper.have_equal_moments(gf_moments, greens_function, 2) + assert np.allclose(static, gf_moments[1]) + + +def test_hf(mf: scf.hf.RHF) -> None: + """Test the HF expression.""" + hf_h = HF.h.from_mf(mf) + hf_p = HF.p.from_mf(mf) + hf_dyson = HF["dyson"].from_mf(mf) + gf_h_moments = hf_h.build_gf_moments(2) + gf_p_moments = hf_p.build_gf_moments(2) + gf_dyson_moments = hf_dyson.build_gf_moments(2) + + # Get the energy from the hole moments + h1e = np.einsum("pq,pi,qj->ij", mf.get_hcore(), mf.mo_coeff, mf.mo_coeff) + energy = util.gf_moments_galitskii_migdal(gf_h_moments, h1e, factor=1.0) + + assert np.abs(energy - mf.energy_elec()[0]) < 1e-8 + + # Get the Fock matrix Fock matrix from the moments + fock_ref = np.einsum("pq,pi,qj->ij", mf.get_fock(), mf.mo_coeff, mf.mo_coeff) + fock = gf_h_moments[1] + gf_p_moments[1] + + assert np.allclose(fock, fock_ref) + assert np.allclose(gf_dyson_moments[1], fock) + + # Get the Green's function from the Exact solver + exact_h = Exact.from_expression(hf_h) + exact_h.kernel() + exact_p = Exact.from_expression(hf_p) + exact_p.kernel() + assert exact_h.result is not None + assert exact_p.result is not None + result = exact_h.result.combine(exact_p.result) + + assert np.allclose(result.get_greens_function().as_perturbed_mo_energy(), mf.mo_energy) + + +def test_ccsd(mf: scf.hf.RHF) -> None: + """Test the CCSD expression.""" + ccsd_h = CCSD.h.from_mf(mf) + ccsd_p = CCSD.p.from_mf(mf) + pyscf_ccsd = pyscf.cc.CCSD(mf) + pyscf_ccsd.run(conv_tol=1e-10, conv_tol_normt=1e-8) + gf_moments_h = ccsd_h.build_gf_moments(2) + gf_moments_p = ccsd_p.build_gf_moments(2) + + # Get the energy from the hole moments + h1e = np.einsum("pq,pi,qj->ij", mf.get_hcore(), mf.mo_coeff, mf.mo_coeff) + energy = util.gf_moments_galitskii_migdal(gf_moments_h, h1e, factor=1.0) + energy_ref = pyscf_ccsd.e_tot - mf.mol.energy_nuc() + + correct_energy = np.abs(energy - energy_ref) < 1e-8 + if mf.mol.nelectron > 2: + # Galitskii--Migdal should not capture the energy for CCSD for >2 electrons + with pytest.raises(AssertionError): + assert correct_energy + else: + assert correct_energy + + # Get the Green's function from the Davidson solver + davidson = Davidson.from_expression(ccsd_h, nroots=3) + davidson.kernel() + ip_ref, _ = pyscf_ccsd.ipccsd(nroots=3) + + assert davidson.result is not None + assert np.allclose(davidson.result.eigvals[0], -ip_ref[-1]) + + # Check the RDM + rdm1 = gf_moments_h[0].copy() + rdm1 += rdm1.T.conj() + rdm1_ref = pyscf_ccsd.make_rdm1(with_mf=True) + + assert np.allclose(rdm1, rdm1_ref) + + # Check the zeroth moments add to identity + gf_moment_0 = gf_moments_h[0] + gf_moments_p[0] + + assert np.allclose(gf_moment_0, np.eye(mf.mol.nao)) + + +def test_fci(mf: scf.hf.RHF) -> None: + """Test the FCI expression.""" + fci_h = FCI.h.from_mf(mf) + fci_p = FCI.p.from_mf(mf) + pyscf_fci = pyscf.fci.FCI(mf) + gf_moments_h = fci_h.build_gf_moments(2) + gf_moments_p = fci_p.build_gf_moments(2) + + # Get the energy from the hole moments + h1e = np.einsum("pq,pi,qj->ij", mf.get_hcore(), mf.mo_coeff, mf.mo_coeff) + energy = util.gf_moments_galitskii_migdal(gf_moments_h, h1e, factor=1.0) + energy_ref = pyscf_fci.kernel()[0] - mf.mol.energy_nuc() + + assert np.abs(energy - energy_ref) < 1e-8 + + # Check the RDM + rdm1 = fci_h.build_gf_moments(1)[0] * 2 + rdm1_ref = pyscf_fci.make_rdm1(pyscf_fci.ci, mf.mol.nao, mf.mol.nelectron) + + assert np.allclose(rdm1, rdm1_ref) + + # Check the zeroth moments add to identity + gf_moment_0 = gf_moments_h[0] + gf_moments_p[0] + + assert np.allclose(gf_moment_0, np.eye(mf.mol.nao)) + + if mf.mol.nelectron <= 2: + # CCSD should match FCI for <=2 electrons for the hole moments + ccsd = CCSD.h.from_mf(mf) + gf_moments_ccsd = ccsd.build_gf_moments(2) + + assert np.allclose(gf_moments_ccsd[0], gf_moments_h[0]) + assert np.allclose(gf_moments_ccsd[1], gf_moments_h[1]) + + +def test_adc2(mf: scf.hf.RHF) -> None: + """Test the ADC(2) expression.""" + adc_h = ADC2.h.from_mf(mf) + adc_p = ADC2.p.from_mf(mf) + pyscf_adc = pyscf.adc.ADC(mf) + gf_moments_h = adc_h.build_gf_moments(2) + gf_moments_p = adc_p.build_gf_moments(2) + + # Get the energy from the hole moments + h1e = np.einsum("pq,pi,qj->ij", mf.get_hcore(), mf.mo_coeff, mf.mo_coeff) + energy = util.gf_moments_galitskii_migdal(gf_moments_h, h1e, factor=1.0) + energy_ref = mf.energy_elec()[0] + pyscf_adc.kernel_gs()[0] + + assert np.abs(energy - energy_ref) < 1e-8 + + # Get the Green's function from the Davidson solver + davidson = Davidson.from_expression(adc_h, nroots=3) + davidson.kernel() + ip_ref, _, _, _ = pyscf_adc.kernel(nroots=3) + + assert davidson.result is not None + assert np.allclose(davidson.result.eigvals[0], -ip_ref[-1]) + + # Check the RDM + rdm1 = adc_h.build_gf_moments(1)[0] * 2 + rdm1_ref = np.diag(mf.mo_occ) # No correlated ground state! + + assert np.allclose(rdm1, rdm1_ref) + + # Check the zeroth moments add to identity + gf_moment_0 = gf_moments_h[0] + gf_moments_p[0] + + assert np.allclose(gf_moment_0, np.eye(mf.mol.nao)) + + +def test_adc2x(mf: scf.hf.RHF) -> None: + """Test the ADC(2)-x expression.""" + adc_h = ADC2x.h.from_mf(mf) + adc_p = ADC2x.p.from_mf(mf) + pyscf_adc = pyscf.adc.ADC(mf) + pyscf_adc.method = "adc(2)-x" + gf_moments_h = adc_h.build_gf_moments(2) + gf_moments_p = adc_p.build_gf_moments(2) + + # Get the energy from the hole moments + h1e = np.einsum("pq,pi,qj->ij", mf.get_hcore(), mf.mo_coeff, mf.mo_coeff) + energy = util.gf_moments_galitskii_migdal(gf_moments_h, h1e, factor=1.0) + energy_ref = mf.energy_elec()[0] + pyscf_adc.kernel_gs()[0] + + assert np.abs(energy - energy_ref) < 1e-8 + + # Get the Green's function from the Davidson solver + davidson = Davidson.from_expression(adc_h, nroots=3) + davidson.kernel() + ip_ref, _, _, _ = pyscf_adc.kernel(nroots=3) + + assert davidson.result is not None + assert np.allclose(davidson.result.eigvals[0], -ip_ref[-1]) + + # Check the RDM + rdm1 = adc_h.build_gf_moments(1)[0] * 2 + rdm1_ref = np.diag(mf.mo_occ) # No correlated ground state! + + assert np.allclose(rdm1, rdm1_ref) + + # Check the zeroth moments add to identity + gf_moment_0 = gf_moments_h[0] + gf_moments_p[0] + + assert np.allclose(gf_moment_0, np.eye(mf.mol.nao)) + + +def test_tdagw(mf: scf.hf.RHF, exact_cache: ExactGetter) -> None: + """Test the TDAGW expression.""" + tdagw = TDAGW["dyson"].from_mf(mf) + dft = mf.to_rks() + dft.xc = "hf" + + td = pyscf.tdscf.dTDA(dft) + td.nstates = np.sum(mf.mo_occ > 0) * np.sum(mf.mo_occ == 0) + td.kernel() + td.xy = np.array([(x, np.zeros_like(x)) for x, y in td.xy]) + gw_obj = pyscf.gw.GW(dft, tdmf=td, freq_int="exact") + gw_obj.kernel() + + # Get the IPs and EAs from the Exact solver + solver = exact_cache(mf, TDAGW["dyson"]) + assert solver.result is not None + gf = solver.result.get_greens_function() + mo_energy = gf.as_perturbed_mo_energy() + + # No diagonal approximation in TDAGW so large error + assert np.abs(mo_energy[tdagw.nocc - 1] - gw_obj.mo_energy[tdagw.nocc - 1]) < 1e-3 + assert np.abs(mo_energy[tdagw.nocc] - gw_obj.mo_energy[tdagw.nocc]) < 1e-3 diff --git a/tests/test_lehmann.py b/tests/test_lehmann.py deleted file mode 100644 index 1c7a5cb..0000000 --- a/tests/test_lehmann.py +++ /dev/null @@ -1,146 +0,0 @@ -""" -Tests for Lehmann representations. -""" - -import unittest -import pytest - -import numpy as np -from pyscf import lib - -from dyson import Lehmann - - -@pytest.mark.regression -class Lehmann_Hermitian_Tests(unittest.TestCase): - """ - Tests for the `Lehmann` class. - """ - - @classmethod - def setUpClass(cls): - e = np.cos(np.arange(100)) - c = np.sin(np.arange(1000)).reshape(10, 100) - cls.aux = Lehmann(e, c) - - @classmethod - def tearDownClass(cls): - del cls.aux - - def test_moments(self): - t = self.aux.moment([0, 1]) - self.assertAlmostEqual(lib.fp(t), -38.97393642159078, 10) - t = self.aux.moment(1) - self.assertAlmostEqual(lib.fp(t), 0.09373842776339, 10) - - def test_chebyshev_moments(self): - t = self.aux.chebyshev_moment(range(10)) - self.assertAlmostEqual(lib.fp(t), -59.24704483050994, 10) - t = self.aux.chebyshev_moment(5) - self.assertAlmostEqual(lib.fp(t), -0.89044258131632, 10) - - def test_matrix(self): - phys = np.diag(np.cos(np.arange(self.aux.nphys))) - mat = self.aux.matrix(phys) - self.assertAlmostEqual(lib.fp(mat), -1.7176781717484837, 10) - mat = self.aux.matrix(phys, chempot=0.1, out=mat) - self.assertAlmostEqual(lib.fp(mat), -1.6486800825995238, 10) - - def test_matvec(self): - phys = np.diag(np.cos(np.arange(self.aux.nphys))) - v = np.cos(np.arange(1e3, 1e3+self.aux.nphys+self.aux.naux)) - mat = self.aux.matrix(phys) - u1 = np.dot(mat, v) - u2 = self.aux.matvec(phys, v) - np.testing.assert_allclose(u1, u2, atol=1e-10) - - def test_diagonalise_matrix(self): - phys = np.diag(np.cos(np.arange(self.aux.nphys))) - e, c = self.aux.diagonalise_matrix(phys) - self.assertAlmostEqual(lib.fp(e), -27.601125782799805, 10) - self.assertAlmostEqual(lib.fp(c), 5.734873329655418, 10) - - def test_diagonalise_matrix_with_projection(self): - phys = np.diag(np.cos(np.arange(self.aux.nphys))) - e, c = self.aux.diagonalise_matrix_with_projection(phys) - self.assertAlmostEqual(lib.fp(e), -27.601125782799805, 10) - self.assertAlmostEqual(lib.fp(c), 0.893365900726610, 10) - - def test_weights(self): - w = self.aux.weights() - self.assertAlmostEqual(lib.fp(w), -3.958396736529412, 10) - - def test_as_orbitals(self): - mo_energy, mo_coeff, mo_occ = self.aux.as_orbitals() - self.assertAlmostEqual(lib.fp(mo_energy), -1.3761236354579696, 10) - self.assertAlmostEqual(lib.fp(mo_coeff), 16.769969516330693, 10) - self.assertAlmostEqual(lib.fp(mo_occ), -0.8204576360667741, 10) - mo_energy, mo_coeff, mo_occ = self.aux.as_orbitals(mo_coeff=np.eye(self.aux.nphys)) - self.assertAlmostEqual(lib.fp(mo_energy), -1.3761236354579696, 10) - self.assertAlmostEqual(lib.fp(mo_coeff), 16.769969516330693, 10) - self.assertAlmostEqual(lib.fp(mo_occ), -0.8204576360667741, 10) - - def test_as_static_potential(self): - mo_energy = np.cos(np.arange(self.aux.nphys)) - v = self.aux.as_static_potential(mo_energy) - self.assertAlmostEqual(lib.fp(v), 92.45915312825181, 10) - - def test_as_perturbed_mo_energy(self): - mo_energy = self.aux.as_perturbed_mo_energy() - self.assertAlmostEqual(lib.fp(mo_energy), -1.3838965817318036, 10) - - -@pytest.mark.regression -class Lehmann_NonHermitian_Tests(unittest.TestCase): - """ - Tests for the `Lehmann` class without hermiticity. - """ - - @classmethod - def setUpClass(cls): - e = np.cos(np.arange(100)) - c = ( - np.sin(np.arange(1000)).reshape(10, 100), - np.cos(np.arange(1000, 2000)).reshape(10, 100), - ) - cls.aux = Lehmann(e, c) - - @classmethod - def tearDownClass(cls): - del cls.aux - - def test_moments(self): - t = self.aux.moment([0, 1]) - self.assertAlmostEqual(lib.fp(t), 5.500348836608749, 10) - t = self.aux.moment(1) - self.assertAlmostEqual(lib.fp(t), 0.106194238305493, 10) - - def test_chebyshev_moments(self): - t = self.aux.chebyshev_moment(range(10)) - self.assertAlmostEqual(lib.fp(t), 22.449864768273073, 10) - t = self.aux.chebyshev_moment(5) - self.assertAlmostEqual(lib.fp(t), 0.481350207154633, 10) - - def test_matrix(self): - phys = np.diag(np.cos(np.arange(self.aux.nphys))) - mat = self.aux.matrix(phys) - self.assertAlmostEqual(lib.fp(mat), 0.3486818284611074, 10) - mat = self.aux.matrix(phys, chempot=0.1, out=mat) - self.assertAlmostEqual(lib.fp(mat), 0.4176799176100642, 10) - - def test_weights(self): - w = self.aux.weights() - self.assertAlmostEqual(lib.fp(w), 3.360596091200564, 10) - - def test_as_static_potential(self): - mo_energy = np.cos(np.arange(self.aux.nphys)) - v = self.aux.as_static_potential(mo_energy) - self.assertAlmostEqual(lib.fp(v), 16.51344106239947, 10) - - def test_as_perturbed_mo_energy(self): - mo_energy = self.aux.as_perturbed_mo_energy() - self.assertAlmostEqual(lib.fp(mo_energy), -1.4195771834933937, 10) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/test_mblgf.py b/tests/test_mblgf.py new file mode 100644 index 0000000..ad96789 --- /dev/null +++ b/tests/test_mblgf.py @@ -0,0 +1,133 @@ +"""Tests for :mod:`~dyson.solvers.static.mblgf`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from dyson import util +from dyson.representations.spectral import Spectral +from dyson.solvers import MBLGF + +if TYPE_CHECKING: + from pyscf import scf + + from dyson.expressions.expression import ExpressionCollection + + from .conftest import ExactGetter, Helper + + +@pytest.mark.parametrize("max_cycle", [0, 1, 2, 3]) +def test_central_moments( + helper: Helper, + mf: scf.hf.RHF, + expression_method: ExpressionCollection, + max_cycle: int, +) -> None: + """Test the recovery of the exact central moments from the MBLGF solver.""" + # Get the quantities required from the expression + if "h" not in expression_method or "p" not in expression_method: + pytest.skip("Skipping test for Dyson only expression") + expression_h = expression_method.h.from_mf(mf) + expression_p = expression_method.p.from_mf(mf) + nmom_gf = max_cycle * 2 + 2 + nmom_se = nmom_gf - 2 + gf_moments = expression_h.build_gf_moments(nmom_gf) + expression_p.build_gf_moments(nmom_gf) + se_static, se_moments = util.gf_moments_to_se_moments(gf_moments) + + # Run the MBLGF solver + solver = MBLGF(gf_moments, hermitian=expression_h.hermitian_downfolded) + solver.kernel() + assert solver.result is not None + + # Recover the Green's function and self-energy + static = solver.result.get_static_self_energy() + self_energy = solver.result.get_self_energy() + greens_function = solver.result.get_greens_function() + + if expression_h.hermitian_downfolded: + assert helper.have_equal_moments(greens_function, gf_moments, nmom_gf) + assert helper.have_equal_moments(static, se_static, nmom_se) + assert helper.have_equal_moments(self_energy, se_moments, nmom_se) + else: + # A little more numerical error in some non-Hermitian cases + assert helper.have_equal_moments(greens_function, gf_moments, nmom_gf, tol=1e-7) + assert helper.have_equal_moments(static, se_static, nmom_se, tol=1e-7) + assert helper.have_equal_moments(self_energy, se_moments, nmom_se, 1e-7) + + +@pytest.mark.parametrize("max_cycle", [0, 1, 2, 3]) +def test_vs_exact_solver_central( + helper: Helper, + request: pytest.FixtureRequest, + mf: scf.hf.RHF, + expression_method: ExpressionCollection, + exact_cache: ExactGetter, + max_cycle: int, +) -> None: + """Test the MBLGF solver for central moments.""" + # Get the quantities required from the expressions + if "h" not in expression_method or "p" not in expression_method: + pytest.skip("Skipping test for Dyson only expression") + expression_h = expression_method.h.from_mf(mf) + expression_p = expression_method.p.from_mf(mf) + if expression_h.nconfig > 1024 or expression_p.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + if request.node.name in ( + "test_vs_exact_solver_central[lih-631g-CCSD-3]", + "test_vs_exact_solver_central[h2o-sto3g-CCSD-2]", + "test_vs_exact_solver_central[h2o-sto3g-CCSD-3]", + ): + pytest.skip("Numerical error in this test case is too high.") + nmom_gf = max_cycle * 2 + 2 + + # Solve the Hamiltonian exactly + exact_h = exact_cache(mf, expression_method.h) + exact_p = exact_cache(mf, expression_method.p) + assert exact_h.result is not None + assert exact_p.result is not None + result_exact_ph = Spectral.combine(exact_h.result, exact_p.result) + + # Get the self-energy and Green's function from the exact solver + static_exact = result_exact_ph.get_static_self_energy() + self_energy_exact = result_exact_ph.get_self_energy() + greens_function_exact = result_exact_ph.get_greens_function() + gf_h_moments_exact = exact_h.result.get_greens_function().moments(range(nmom_gf)) + gf_p_moments_exact = exact_p.result.get_greens_function().moments(range(nmom_gf)) + + # Solve the Hamiltonian with MBLGF + mblgf_h = MBLGF(gf_h_moments_exact, hermitian=expression_h.hermitian_downfolded) + mblgf_h.kernel() + mblgf_p = MBLGF(gf_p_moments_exact, hermitian=expression_p.hermitian_downfolded) + mblgf_p.kernel() + assert mblgf_h.result is not None + assert mblgf_p.result is not None + result_ph = Spectral.combine(mblgf_h.result, mblgf_p.result) + + assert helper.have_equal_moments( + mblgf_h.result.get_self_energy(), exact_h.result.get_self_energy(), nmom_gf - 2 + ) + assert helper.have_equal_moments( + mblgf_p.result.get_self_energy(), exact_p.result.get_self_energy(), nmom_gf - 2 + ) + + # Recover the hole Green's function from the MBLGF solver + greens_function = mblgf_h.result.get_greens_function() + + assert helper.have_equal_moments(greens_function, gf_h_moments_exact, nmom_gf) + + # Recover the particle Green's function from the MBLGF solver + greens_function = mblgf_p.result.get_greens_function() + + assert helper.have_equal_moments(greens_function, gf_p_moments_exact, nmom_gf, tol=1e-7) + + # Recover the self-energy and Green's function from the recovered MBLGF solver + static = result_ph.get_static_self_energy() + self_energy = result_ph.get_self_energy() + greens_function = result_ph.get_greens_function() + + assert helper.are_equal_arrays(static, static_exact) + assert helper.have_equal_moments(self_energy, self_energy_exact, nmom_gf - 2) + assert helper.have_equal_moments(greens_function, greens_function_exact, nmom_gf) + assert helper.recovers_greens_function(static, self_energy, greens_function, 4) diff --git a/tests/test_mblse.py b/tests/test_mblse.py new file mode 100644 index 0000000..30f4ca8 --- /dev/null +++ b/tests/test_mblse.py @@ -0,0 +1,119 @@ +"""Tests for :mod:`~dyson.solvers.static.mblse`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from dyson import util +from dyson.representations.spectral import Spectral +from dyson.solvers import MBLSE + +if TYPE_CHECKING: + from pyscf import scf + + from dyson.expressions.expression import ExpressionCollection + + from .conftest import ExactGetter, Helper + + +@pytest.mark.parametrize("max_cycle", [0, 1, 2, 3]) +def test_central_moments( + helper: Helper, + mf: scf.hf.RHF, + expression_method: ExpressionCollection, + max_cycle: int, +) -> None: + """Test the recovery of the exact central moments from the MBLSE solver.""" + # Get the quantities required from the expression + if "h" not in expression_method or "p" not in expression_method: + pytest.skip("Skipping test for Dyson only expression") + expression_h = expression_method.h.from_mf(mf) + expression_p = expression_method.p.from_mf(mf) + nmom_gf = max_cycle * 2 + 4 + nmom_se = nmom_gf - 2 + gf_moments = expression_h.build_gf_moments(nmom_gf) + expression_p.build_gf_moments(nmom_gf) + static, se_moments = util.gf_moments_to_se_moments(gf_moments) + + # Check if we need a non-Hermitian solver + hermitian = expression_h.hermitian_downfolded and expression_p.hermitian_downfolded + + # Run the MBLSE solver + solver = MBLSE(static, se_moments, hermitian=hermitian) + solver.kernel() + assert solver.result is not None + + # Recover the moments + static_recovered = solver.result.get_static_self_energy() + self_energy = solver.result.get_self_energy() + + assert helper.are_equal_arrays(static, static_recovered) + assert helper.have_equal_moments(se_moments, self_energy, nmom_se) + + +@pytest.mark.parametrize("max_cycle", [0, 1, 2, 3]) +def test_vs_exact_solver_central( + helper: Helper, + mf: scf.hf.RHF, + expression_method: ExpressionCollection, + exact_cache: ExactGetter, + max_cycle: int, +) -> None: + # Get the quantities required from the expressions + if "h" not in expression_method or "p" not in expression_method: + pytest.skip("Skipping test for Dyson only expression") + expression_h = expression_method.h.from_mf(mf) + expression_p = expression_method.p.from_mf(mf) + if expression_h.nconfig > 1024 or expression_p.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + nmom_se = max_cycle * 2 + 2 + + # Check if we need a non-Hermitian solver + hermitian = expression_h.hermitian_downfolded and expression_p.hermitian_downfolded + + # Solve the Hamiltonian exactly + exact_h = exact_cache(mf, expression_method.h) + exact_p = exact_cache(mf, expression_method.p) + assert exact_h.result is not None + assert exact_p.result is not None + result_exact_ph = Spectral.combine(exact_h.result, exact_p.result) + + # Get the self-energy and Green's function from the exact solver + static_exact = result_exact_ph.get_static_self_energy() + self_energy_exact = result_exact_ph.get_self_energy() + greens_function_exact = result_exact_ph.get_greens_function() + static_h_exact = exact_h.result.get_static_self_energy() + static_p_exact = exact_p.result.get_static_self_energy() + se_h_moments_exact = exact_h.result.get_self_energy().moments(range(nmom_se)) + se_p_moments_exact = exact_p.result.get_self_energy().moments(range(nmom_se)) + overlap_h = exact_h.result.get_overlap() + overlap_p = exact_p.result.get_overlap() + + # Solve the Hamiltonian with MBLSE + mblse_h = MBLSE( + static_h_exact, + se_h_moments_exact, + overlap=overlap_h, + hermitian=hermitian, + ) + result_h = mblse_h.kernel() + mblse_p = MBLSE( + static_p_exact, + se_p_moments_exact, + overlap=overlap_p, + hermitian=hermitian, + ) + result_p = mblse_p.kernel() + result_ph = Spectral.combine(result_h, result_p) + + # Recover the self-energy and Green's function from the MBLSE solver + static = result_ph.get_static_self_energy() + self_energy = result_ph.get_self_energy() + greens_function = result_ph.get_greens_function() + + assert helper.are_equal_arrays(static, static_exact) + assert helper.have_equal_moments(self_energy, se_h_moments_exact + se_p_moments_exact, nmom_se) + assert helper.have_equal_moments(self_energy, self_energy_exact, nmom_se) + assert helper.recovers_greens_function(static, self_energy, greens_function, 4) + assert helper.have_equal_moments(greens_function, greens_function_exact, nmom_se) diff --git a/tests/test_util.py b/tests/test_util.py new file mode 100644 index 0000000..d59ffd0 --- /dev/null +++ b/tests/test_util.py @@ -0,0 +1,61 @@ +"""Tests for :module:`~dyson.util`.""" + +from __future__ import annotations + +from typing import TYPE_CHECKING + +import pytest + +from dyson import util + +if TYPE_CHECKING: + from pyscf import scf + + from dyson.expressions.expression import BaseExpression + + from .conftest import ExactGetter, Helper + + +def test_moments_conversion( + helper: Helper, + mf: scf.hf.RHF, + expression_cls: type[BaseExpression], + exact_cache: ExactGetter, +) -> None: + """Test the conversion of moments between self-energy and Green's function.""" + # Get the quantities required from the expression + expression = expression_cls.from_mf(mf) + if expression.nconfig > 1024: + pytest.skip("Skipping test for large Hamiltonian") + + # Solve the Hamiltonian + solver = exact_cache(mf, expression_cls) + + assert solver.result is not None + assert solver.nphys == expression.nphys + assert solver.hermitian == expression.hermitian_upfolded + + # Get the self-energy and Green's function from the solver + static = solver.result.get_static_self_energy() + self_energy = solver.result.get_self_energy() + greens_function = solver.result.get_greens_function() + + assert self_energy.nphys == expression.nphys + assert greens_function.nphys == expression.nphys + assert helper.recovers_greens_function(static, self_energy, greens_function) + + # Get the moments from the self-energy and Green's function + se_moments = self_energy.moments(range(4)) + gf_moments = greens_function.moments(range(6)) + + # Recover the self-energy from the Green's function moments + static_other, se_moments_other = util.gf_moments_to_se_moments(gf_moments) + gf_moments_other = util.se_moments_to_gf_moments(static, se_moments, overlap=gf_moments[0]) + + assert helper.are_equal_arrays(static, static_other) + if expression.hermitian_upfolded: + assert helper.have_equal_moments(se_moments, se_moments_other, 4) + assert helper.have_equal_moments(gf_moments, gf_moments_other, 6) + else: + assert helper.have_equal_moments(se_moments, se_moments_other, 4, tol=5e-7) + assert helper.have_equal_moments(gf_moments, gf_moments_other, 6, tol=5e-7) diff --git a/tests/util/test_energy.py b/tests/util/test_energy.py deleted file mode 100644 index dddf101..0000000 --- a/tests/util/test_energy.py +++ /dev/null @@ -1,51 +0,0 @@ -""" -Tests for energy functionals. -""" - -import unittest -import pytest - -from pyscf import gto, scf, agf2, lib -import numpy as np - -from dyson import util - - -@pytest.mark.regression -class Energy_Tests(unittest.TestCase): - """ - Tests for the `util.energy` module. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="6-31g", verbose=0) - mf = scf.RHF(mol) - mf.conv_tol = 1e-14 - mf.kernel() - h = np.linalg.multi_dot((mf.mo_coeff.T, mf.get_hcore(), mf.mo_coeff)) - f = np.diag(mf.mo_energy) - gf2 = agf2.AGF2(mf, nmom=(None, None)) - se = gf2.build_se() - se.coupling[:gf2.nocc, se.energy > se.chempot] = 0 - se.coupling[gf2.nocc:, se.energy < se.chempot] = 0 - gf = se.get_greens_function(f) - cls.gf2, cls.mf, cls.h, cls.f, cls.se, cls.gf = gf2, mf, h, f, se, gf - - @classmethod - def tearDownClass(cls): - del cls.gf2, cls.mf, cls.h, cls.f, cls.se, cls.gf - - def test_greens_function_galitskii_migdal(self): - moments = self.gf.get_occupied().moment(range(2)) - e_gm = util.greens_function_galitskii_migdal(moments, self.h) - e_gm += self.mf.mol.energy_nuc() - - e_ref = self.gf2.energy_1body(self.gf2.ao2mo(), self.gf) - e_ref += self.gf2.energy_2body(self.gf, self.se) - - self.assertAlmostEqual(e_gm, e_ref, 8) - - -if __name__ == "__main__": - unittest.main() diff --git a/tests/util/test_moments.py b/tests/util/test_moments.py deleted file mode 100644 index b805421..0000000 --- a/tests/util/test_moments.py +++ /dev/null @@ -1,79 +0,0 @@ -""" -Tests for moment utilities. -""" - -import unittest -import pytest - -from pyscf import gto, scf, agf2, lib -import numpy as np - -from dyson import util - - -@pytest.mark.regression -class Moments_Tests(unittest.TestCase): - """ - Test for the `util.moments` module. - """ - - @classmethod - def setUpClass(cls): - mol = gto.M(atom="H 0 0 0; Li 0 0 1.64", basis="6-31g", verbose=0) - mf = scf.RHF(mol).run() - f = np.diag(mf.mo_energy[mf.mo_occ > 0]) - se = agf2.AGF2(mf, nmom=(None, None)).build_se().get_virtual() - se.coupling = se.coupling[mf.mo_occ > 0] - gf = se.get_greens_function(f) - cls.f, cls.se, cls.gf, = f, se, gf - - @classmethod - def tearDownClass(cls): - del cls.f, cls.se, cls.gf - - def test_se_moments_to_gf_moments(self): - t_se = self.se.moment(range(10)) - t_gf = self.gf.moment(range(12)) - t_gf_recov = util.se_moments_to_gf_moments(self.f, t_se) - for i, (a, b) in enumerate(zip(t_gf, t_gf_recov)): - self.assertAlmostEqual(util.scaled_error(a, b), 0, 10) - - def test_gf_moments_to_se_moments(self): - t_gf = self.gf.moment(range(12)) - t_se = self.se.moment(range(10)) - static_recov, t_se_recov = util.gf_moments_to_se_moments(t_gf) - self.assertAlmostEqual(util.scaled_error(static_recov, self.f), 0, 10) - for i, (a, b) in enumerate(zip(t_se, t_se_recov)): - self.assertAlmostEqual(util.scaled_error(a, b), 0, 10) - - def test_matvec_to_greens_function(self): - h = np.block([[self.f, self.se.coupling], [self.se.coupling.T, np.diag(self.se.energy)]]) - matvec = lambda v: np.dot(h, v) - bra = np.eye(self.se.nphys, self.se.nphys + self.se.naux) - t_gf = self.gf.moment(range(10)) - t_gf_matvec = util.matvec_to_greens_function(matvec, 10, bra) - for i, (a, b) in enumerate(zip(t_gf, t_gf_matvec)): - self.assertAlmostEqual(util.scaled_error(a, b), 0, 10) - - def test_matvec_to_greens_function_chebyshev(self): - emin = self.gf.energy.min() - emax = self.gf.energy.max() - a = (emax - emin) / (2.0 - 1e-2) - b = (emax + emin) / 2.0 - energy_scaled = (self.gf.energy - b) / a - c = np.zeros((100, self.gf.nphys, energy_scaled.size)) - c[0] = self.gf.coupling - c[1] = self.gf.coupling * energy_scaled - for i in range(2, 100): - c[i] = 2.0 * c[i-1] * energy_scaled - c[i-2] - t_gf = lib.einsum("qx,npx->npq", self.gf.coupling, c) - h = np.block([[self.f, self.se.coupling], [self.se.coupling.T, np.diag(self.se.energy)]]) - matvec = lambda v: np.dot(h, v) - bra = np.eye(self.se.nphys, self.se.nphys + self.se.naux) - t_gf_matvec = util.matvec_to_greens_function_chebyshev(matvec, 100, (a, b), bra) - for i, (a, b) in enumerate(zip(t_gf, t_gf_matvec)): - self.assertAlmostEqual(util.scaled_error(a, b), 0, 10) - - -if __name__ == "__main__": - unittest.main()