Source code for pyscf.tdscf.common_slow

#  Author: Artem Pulkin
"""
This and other `_slow` modules implement the time-dependent procedure. The primary performance drawback is
that, unlike other 'fast' routines with an implicit construction of the eigenvalue problem, these modules construct
TDHF matrices explicitly. As a result, regular `numpy.linalg.eig` can be used to retrieve TDHF roots in a reliable
fashion without any issues related to the Davidson procedure.

This is a helper module defining basic interfaces.
"""

import sys
from pyscf.lib import logger

from pyscf.pbc.tools import get_kconserv

import numpy
from scipy.linalg import solve

from itertools import count, groupby

if sys.version_info >= (3,):
    unicode = str

[docs] def msize(m): """ Checks whether the matrix is square and returns its size. Args: m (numpy.ndarray): the matrix to measure; Returns: An integer with the size. """ s = m.shape[0] if m.shape != (s, s): raise ValueError("Do not recognize the shape (must be a square matrix): {}".format(m.shape)) return s
[docs] def full2ab(full, tolerance=1e-12): """ Transforms a full TD matrix into A and B parts. Args: full (numpy.ndarray): the full TD matrix; tolerance (float): a tolerance for checking whether the full matrix is in the ABBA-form; Returns: A and B submatrices. """ s = msize(full) if s % 2 != 0: raise ValueError("Not an even matrix size: {:d}".format(s)) s2 = s // 2 a, b = full[:s2, :s2], full[:s2, s2:] b_, a_ = full[s2:, :s2].conj(), full[s2:, s2:].conj() delta = max(abs(a + a_).max(), abs(b + b_).max()) if delta > tolerance: raise ValueError("The full matrix is not in the ABBA-form, delta: {:.3e}".format(delta)) return full[:s2, :s2], full[:s2, s2:]
[docs] def ab2full(a, b): """ Transforms A and B TD matrices into a full matrix. Args: a (numpy.ndarray): TD A-matrix; b (numpy.ndarray): TD B-matrix; Returns: The full TD matrix. """ sa = msize(a) sb = msize(b) if sa != sb: raise ValueError("Input matrix dimensions do not match: {:d} vs {:d}".format(sa, sb)) return numpy.block([[a, b], [-b.conj(), -a.conj()]])
[docs] def ab2mkk(a, b, tolerance=1e-12): """ Transforms A and B TD matrices into MK and K matrices. Args: a (numpy.ndarray): TD A-matrix; b (numpy.ndarray): TD B-matrix; tolerance (float): a tolerance for checking whether the input matrices are real; Returns: MK and K submatrices. """ if max(abs(a.imag).max(), abs(b.imag).max()) > tolerance: raise ValueError("A- and/or B-matrixes are complex-valued: no transform is possible") a, b = a.real, b.real tdhf_k, tdhf_m = a - b, a + b tdhf_mk = tdhf_m.dot(tdhf_k) return tdhf_mk, tdhf_k
[docs] def mkk2ab(mk, k): """ Transforms MK and M TD matrices into A and B matrices. Args: mk (numpy.ndarray): TD MK-matrix; k (numpy.ndarray): TD K-matrix; Returns: A and B submatrices. """ if numpy.iscomplexobj(mk) or numpy.iscomplexobj(k): raise ValueError("MK- and/or K-matrixes are complex-valued: no transform is possible") m = solve(k.T, mk.T).T a = 0.5 * (m + k) b = 0.5 * (m - k) return a, b
[docs] def full2mkk(full): """ Transforms a full TD matrix into MK and K parts. Args: full (numpy.ndarray): the full TD matrix; Returns: MK and K submatrices. """ return ab2mkk(*full2ab(full))
[docs] def mkk2full(mk, k): """ Transforms MK and M TD matrices into a full TD matrix. Args: mk (numpy.ndarray): TD MK-matrix; k (numpy.ndarray): TD K-matrix; Returns: The full TD matrix. """ return ab2full(*mkk2ab(mk, k))
[docs] class TDMatrixBlocks:
[docs] def tdhf_primary_form(self, *args, **kwargs): """ A primary form of TDHF matrixes. Returns: Output type: "full", "ab", or "mk" and the corresponding matrix(es). """ raise NotImplementedError
@staticmethod def __check_primary_form__(m): if not isinstance(m, tuple): raise ValueError("The value returned by `tdhf_primary_form` is not a tuple") if len(m) < 1: raise ValueError("Empty tuple returned by `tdhf_primary_form`") if not isinstance(m[0], (str, unicode)): raise ValueError("The first item returned by `tdhf_primary_form` must be a string") forms = dict(ab=3, mk=3, full=2) if m[0] in forms: if len(m) != forms[m[0]]: raise ValueError("The {} form returned by `tdhf_primary_form` must contain {:d} values".format( m[0].upper(), forms[m[0]], )) else: raise ValueError("Unknown form specification returned by `tdhf_primary_form`: {}".format(m[0]))
[docs] def tdhf_ab_form(self, *args, **kwargs): """ The A-B form of the TD problem. Returns: A and B TD matrices. """ m = self.tdhf_primary_form(*args, **kwargs) self.__check_primary_form__(m) if m[0] == "ab": return m[1:] elif m[0] == "full": return full2ab(m[1]) elif m[0] == "mk": return mkk2ab(*m[1:])
[docs] def tdhf_full_form(self, *args, **kwargs): """ The full form of the TD problem. Returns: The full TD matrix. """ m = self.tdhf_primary_form(*args, **kwargs) self.__check_primary_form__(m) if m[0] == "ab": return ab2full(*m[1:]) elif m[0] == "full": return m[1] elif m[0] == "mk": return mkk2full(*m[1:])
[docs] def tdhf_mk_form(self, *args, **kwargs): """ The MK form of the TD problem. Returns: MK and K TD matrixes. """ m = self.tdhf_primary_form(*args, **kwargs) self.__check_primary_form__(m) if m[0] == "ab": return ab2mkk(*m[1:]) elif m[0] == "full": return full2mkk(m[1]) elif m[0] == "mk": return m[1:]
[docs] def mknj2i(item): """ Transforms "mknj" notation into tensor index order for the ERI. Args: item (str): an arbitrary transpose of "mknj" letters; Returns: 4 indexes. """ notation = "mknj" notation = dict(zip(notation, range(len(notation)))) return tuple(notation[i] for i in item)
[docs] class TDERIMatrixBlocks(TDMatrixBlocks): symmetries = [ ((0, 1, 2, 3), False), ] def __init__(self): """ This a prototype class for TD calculations based on ERI (TD-HF). It handles integral blocks and the diagonal part, see Eq. 7.5 of RevModPhys.36.844. """ # Caching self.__eri__ = {} def __get_mo_energies__(self, *args, **kwargs): """This routine collects occupied and virtual MO energies.""" raise NotImplementedError def __calc_block__(self, item, *args): raise NotImplementedError
[docs] def tdhf_diag(self, *args): """ Retrieves the diagonal block. Args: ``*args``: args passed to `__get_mo_energies__`; Returns: The diagonal block. """ e_occ, e_virt = self.__get_mo_energies__(*args) diag = (- e_occ[:, numpy.newaxis] + e_virt[numpy.newaxis, :]).reshape(-1) return numpy.diag(diag).reshape((len(e_occ) * len(e_virt), len(e_occ) * len(e_virt)))
[docs] def eri_ov(self, item, *args): """ Retrieves ERI block using 'ov' notation. Args: item (str): a 4-character string of 'o' and 'v' letters; ``*args``: other args passed to `__calc_block__`; Returns: The corresponding block of ERI (4-tensor, phys notation). """ if len(item) != 4 or not isinstance(item, str) or not set(item).issubset('ov'): raise ValueError("Unknown item: {}".format(repr(item))) args = (tuple(item), ) + args if args in self.__eri__: return self.__eri__[args] result = self.__calc_block__(*args) for permutation, conjugation in self.symmetries: permuted_args = tuple( tuple(arg[_i] for _i in permutation) for arg in args ) if conjugation: self.__eri__[permuted_args] = result.transpose(*permutation).conj() else: self.__eri__[permuted_args] = result.transpose(*permutation) return result
[docs] def eri_mknj(self, item, *args): """ Retrieves ERI block using 'mknj' notation. Args: item (str): a 4-character string of 'mknj' letters; ``*args``: other arguments passed to `get_block_ov_notation`; Returns: The corresponding block of ERI (matrix with paired dimensions). """ if len(item) != 4 or not isinstance(item, str) or set(item) != set('mknj'): raise ValueError("Unknown item: {}".format(repr(item))) item = mknj2i(item) n_ov = ''.join('o' if i % 2 == 0 else 'v' for i in item) args = tuple( tuple(arg[i] for i in item) for arg in args ) result = self.eri_ov(n_ov, *args).transpose(*numpy.argsort(item)) i, j, k, l = result.shape result = result.reshape((i * j, k * l)) return result
def __getitem__(self, item): if isinstance(item, str): spec, args = item, tuple() else: spec, args = item[0], item[1:] if set(spec) == set("mknj"): return self.eri_mknj(spec, *args) elif set(spec).issubset("ov"): return self.eri_ov(spec, *args) else: raise ValueError("Unknown item: {}".format(repr(item)))
[docs] def tdhf_primary_form(self, *args, **kwargs): """ A primary form of TDHF matrixes (AB). Returns: Output type: "ab", and the corresponding matrixes. """ d = self.tdhf_diag(*args, **kwargs) a = d + 2 * self["knmj"] - self["knjm"] b = 2 * self["kjmn"] - self["kjnm"] return "ab", a, b
[docs] class TDProxyMatrixBlocks(TDMatrixBlocks): def __init__(self, model): """ This a prototype class for TD calculations based on proxying pyscf classes such as TDDFT. It is a work-around class. It accepts a `pyscf.tdscf.*` class and uses its matvec to construct a full-sized TD matrix. Args: model: a pyscf base model to extract TD matrix from; """ super(TDProxyMatrixBlocks, self).__init__() self.proxy_model = model self.proxy_vind, self.proxy_diag = self.proxy_model.gen_vind(self.proxy_model._scf) self.proxy_vind = VindTracker(self.proxy_vind)
[docs] def tdhf_primary_form(self, *args, **kwargs): raise NotImplementedError
[docs] def format_frozen_mol(frozen, nmo): """ Formats the argument into a mask array of bools where False values correspond to frozen molecular orbitals. Args: frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals; nmo (int): the total number of molecular orbitals; Returns: The mask array. """ space = numpy.ones(nmo, dtype=bool) if frozen is None: pass elif isinstance(frozen, int): space[:frozen] = False elif isinstance(frozen, (tuple, list, numpy.ndarray)): space[frozen] = False else: raise ValueError("Cannot recognize the 'frozen' argument: expected None, int or Iterable") return space
[docs] class MolecularMFMixin: def __init__(self, model, frozen=None): """ A mixin to support custom slices of mean-field attributes: `mo_coeff`, `mo_energy`, ... Molecular version. Also supports single k-point inputs. Args: model: the base model; frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals; """ self.__is_k__ = False if "kpts" in dir(model): self.__is_k__ = True if len(model.kpts) != 1: raise ValueError("Only a single k-point supported, found: model.kpts = {}".format(model.kpts)) self.model = model self.space = format_frozen_mol(frozen, len(self.squeeze(model.mo_energy)))
[docs] def squeeze(self, x): """Squeezes quantities in the case of a PBC model.""" return x[0] if self.__is_k__ else x
@property def mo_coeff(self): """MO coefficients.""" return self.squeeze(self.model.mo_coeff)[:, self.space] @property def mo_energy(self): """MO energies.""" return self.squeeze(self.model.mo_energy)[self.space] @property def mo_occ(self): """MO occupation numbers.""" return self.squeeze(self.model.mo_occ)[self.space] @property def nocc(self): """The number of occupied orbitals.""" return int(self.squeeze(self.model.mo_occ)[self.space].sum() // 2) @property def nmo(self): """The total number of molecular orbitals.""" return self.space.sum() @property def mo_coeff_full(self): """MO coefficients.""" return self.squeeze(self.model.mo_coeff) @property def nocc_full(self): """The true (including frozen degrees of freedom) number of occupied orbitals.""" return int(self.squeeze(self.model.mo_occ).sum() // 2) @property def nmo_full(self): """The true (including frozen degrees of freedom) total number of molecular orbitals.""" return len(self.space)
[docs] def format_frozen_k(frozen, nmo, nk): """ Formats the argument into a mask array of bools where False values correspond to frozen orbitals for each k-point. Args: frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals for all k-points or multiple lists of frozen orbitals for each k-point; nmo (int): the total number of molecular orbitals; nk (int): the total number of k-points; Returns: The mask array. """ space = numpy.ones((nk, nmo), dtype=bool) if frozen is None: pass elif isinstance(frozen, int): space[:, :frozen] = False elif isinstance(frozen, (tuple, list, numpy.ndarray)): if len(frozen) > 0: if isinstance(frozen[0], int): space[:, frozen] = False else: for i in range(nk): space[i, frozen[i]] = False else: raise ValueError("Cannot recognize the 'frozen' argument: expected None, int or Iterable") return space
[docs] def k_nocc(model): """ Retrieves occupation numbers. Args: model (RHF): the model; Returns: Numbers of occupied orbitals in the model. """ return tuple(int(i.sum() // 2) for i in model.mo_occ)
[docs] def k_nmo(model): """ Retrieves number of AOs per k-point. Args: model (RHF): the model; Returns: Numbers of AOs in the model. """ return tuple(i.shape[1] for i in model.mo_coeff)
[docs] class PeriodicMFMixin: def __init__(self, model, frozen=None): """ A mixin to support custom slices of mean-field attributes: `mo_coeff`, `mo_energy`, ... PBC version. Args: model: the base model; frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals; """ self.model = model self.space = format_frozen_k(frozen, len(model.mo_energy[0]), len(model.kpts)) self.kconserv = get_kconserv(self.model.cell, self.model.kpts).swapaxes(1, 2) @property def mo_coeff(self): """MO coefficients.""" return tuple(i[:, j] for i, j in zip(self.model.mo_coeff, self.space)) @property def mo_energy(self): """MO energies.""" return tuple(i[j] for i, j in zip(self.model.mo_energy, self.space)) @property def mo_occ(self): """MO occupation numbers.""" return tuple(i[j] for i, j in zip(self.model.mo_occ, self.space)) @property def nocc(self): """The number of occupied orbitals.""" return k_nocc(self) @property def nmo(self): """The total number of molecular orbitals.""" return k_nmo(self) @property def mo_coeff_full(self): """MO coefficients.""" return self.model.mo_coeff @property def nocc_full(self): """The true (including frozen degrees of freedom) number of occupied orbitals.""" return k_nocc(self.model) @property def nmo_full(self): """The true (including frozen degrees of freedom) total number of molecular orbitals.""" return k_nmo(self.model)
[docs] class VindTracker: def __init__(self, vind): """ Tracks calls to `vind` (a matrix-vector multiplication density response routine). Args: vind (Callable): a matvec product routine; """ self.vind = vind self.args = self.results = self.errors = None self.reset()
[docs] def reset(self): """ Resets statistics. """ self.args = [] self.results = [] self.errors = []
def __call__(self, v): if not isinstance(v, numpy.ndarray): raise ValueError("The input is not an array") self.args.append(v.shape) try: r = self.vind(v) except Exception as e: self.results.append(None) self.errors.append(e) raise r = numpy.array(r) self.results.append(r.shape) self.errors.append(None) return r def __iter__(self): yield from zip(self.args, self.results, self.errors) @property def ncalls(self): return len(self.args) @property def msize(self): for i in self.results: if i is not None: return i[1] return None @property def elements_total(self): return self.msize ** 2 @property def elements_calc(self): return sum(map(lambda i: i[0] * i[1] if i is not None else 0, self.results)) @property def ratio(self): return 1.0 * self.elements_calc / self.elements_total
[docs] def text_stats(self): return "--------------------\nVind call statistics\n--------------------\n" \ " calls: {total_calls:d}\n" \ " elements total: {total_elems:d} ({size})\n" \ " elements calculated: {total_calc:d}\n" \ " ratio: {ratio:.3f}".format( total_calls=self.ncalls, total_elems=self.elements_total, size="x".join((str(self.msize),) * 2), total_calc=self.elements_calc, ratio=self.ratio, )
[docs] def eig(m, driver=None, nroots=None, half=True): """ Eigenvalue problem solver. Args: m (numpy.ndarray): the matrix to diagonalize; driver (str): one of the drivers; nroots (int): the number of roots ot calculate (ignored for `driver` == 'eig'); half (bool): if True, implies spectrum symmetry and takes only a half of eigenvalues; Returns: """ if driver is None: driver = 'eig' if driver == 'eig': vals, vecs = numpy.linalg.eig(m) order = numpy.argsort(vals) vals, vecs = vals[order], vecs[:, order] if half: vals, vecs = vals[len(vals) // 2:], vecs[:, vecs.shape[1] // 2:] vecs = vecs[:, ] vals, vecs = vals[:nroots], vecs[:, :nroots] else: raise ValueError("Unknown driver: {}".format(driver)) return vals, vecs
[docs] def kernel(eri, driver=None, fast=True, nroots=None, **kwargs): """ Calculates eigenstates and eigenvalues of the TDHF problem. Args: eri (TDDFTMatrixBlocks): ERI; driver (str): one of the eigenvalue problem drivers; fast (bool): whether to run diagonalization on smaller matrixes; nroots (int): the number of roots to calculate; ``**kwargs``: arguments to `eri.tdhf_matrix`; Returns: Positive eigenvalues and eigenvectors. """ if not isinstance(eri, TDMatrixBlocks): raise ValueError("The argument must be ERI object") if fast: logger.debug1(eri.model, "Preparing TDHF matrix (fast) ...") tdhf_mk, tdhf_k = eri.tdhf_mk_form(**kwargs) logger.debug1(eri.model, "Diagonalizing a {} matrix with {} ...".format( 'x'.join(map(str, tdhf_mk.shape)), "'{}'".format(driver) if driver is not None else "a default method", )) vals, vecs_x = eig(tdhf_mk, driver=driver, nroots=nroots, half=False) vals = vals ** .5 vecs_y = (1. / vals)[numpy.newaxis, :] * tdhf_k.dot(vecs_x) vecs_u, vecs_v = vecs_y + vecs_x, vecs_y - vecs_x return vals, numpy.concatenate((vecs_u, vecs_v), axis=0) else: logger.debug1(eri.model, "Preparing TDHF matrix ...") m = eri.tdhf_full_form(**kwargs) logger.debug1(eri.model, "Diagonalizing a {} matrix with {} ...".format( 'x'.join(map(str, m.shape)), "'{}'".format(driver) if driver is not None else "a default method", )) return eig(m, driver=driver, nroots=nroots)
[docs] class TDBase: v2a = None def __init__(self, mf, frozen=None): """ Performs TD calculation. Roots and eigenvectors are stored in `self.e`, `self.xy`. Args: mf: the mean-field model; frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals; """ self._scf = mf self.driver = None self.nroots = None self.eri = None self.xy = None self.e = None self.frozen = frozen self.fast = not numpy.iscomplexobj(numpy.asanyarray(mf.mo_coeff)) def __kernel__(self, **kwargs): """Silent implementation of kernel which does not change attributes.""" if self.eri is None: self.eri = self.ao2mo() e, xy = kernel( self.eri, driver=self.driver, nroots=self.nroots, fast=self.fast, **kwargs ) xy = self.vector_to_amplitudes(xy) return e, xy
[docs] def kernel(self): """ Calculates eigenstates and eigenvalues of the TDHF problem. Returns: Positive eigenvalues and eigenvectors. """ self.e, self.xy = self.__kernel__() return self.e, self.xy
[docs] def ao2mo(self): """ Picks ERI: either 4-fold or 8-fold symmetric. Returns: A suitable ERI. """ raise NotImplementedError
[docs] def vector_to_amplitudes(self, vectors): """ Transforms (reshapes) and normalizes vectors into amplitudes. Args: vectors (numpy.ndarray): raw eigenvectors to transform; Returns: Amplitudes with the following shape: (# of roots, 2 (x or y), # of occupied orbitals, # of virtual orbitals). """ return self.v2a(vectors, self.eri.nocc, self.eri.nmo)
[docs] def format_mask(x): """ Formats a mask into a readable string. Args: x (ndarray): an array with the mask; Returns: A readable string with the mask. """ x = numpy.asanyarray(x) if len(x) == 0: return "(empty)" if x.dtype == bool: x = numpy.argwhere(x)[:, 0] grps = tuple(list(g) for _, g in groupby(x, lambda n, c=count(): n-next(c))) return ",".join("{:d}-{:d}".format(i[0], i[-1]) if len(i) > 1 else "{:d}".format(i[0]) for i in grps)