<commit_before>#  Author: <NAME>
"""
This and other `_slow` modules implement the time-dependent Kohn-Sham 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 by proxying to pyscf density response routines with a O(N^4) complexity scaling. As a result,
regular `numpy.linalg.eig` can be used to retrieve TDKS roots in a reliable fashion without any issues related to the
Davidson procedure. Several variants of TDKS are available:

 * `pyscf.tdscf.rks_slow`: the molecular implementation;
 * `pyscf.pbc.tdscf.rks_slow`: PBC (periodic boundary condition) implementation for RKS objects of `pyscf.pbc.scf`
   modules;
 * (this module)`pyscf.pbc.tdscf.krks_slow_supercell`: PBC implementation for KRKS objects of `pyscf.pbc.scf` modules.
   Works with an arbitrary number of k-points but has a overhead due to an effective construction of a supercell.
 * `pyscf.pbc.tdscf.krks_slow_gamma`: A Gamma-point calculation resembling the original `pyscf.pbc.tdscf.krks`
   module. Despite its name, it accepts KRKS objects with an arbitrary number of k-points but finds only few TDKS roots
   corresponding to collective oscillations without momentum transfer;
 * `pyscf.pbc.tdscf.krks_slow`: PBC implementation for KRKS objects of `pyscf.pbc.scf` modules. Works with
   an arbitrary number of k-points and employs k-point conservation (diagonalizes matrix blocks separately).
"""

# Convention for these modules:
# * PhysERI, PhysERI4, PhysERI8 are proxy classes for computing the full TDDFT matrix
# * vector_to_amplitudes reshapes and normalizes the solution
# * TDRKS provides a container

from pyscf.tdscf.common_slow import TDProxyMatrixBlocks, PeriodicMFMixin
from pyscf.tdscf import rks_slow
from pyscf.pbc.tdscf import krhf_slow_supercell, KTDDFT
from pyscf.lib import einsum, cartesian_prod, norm, logger
from pyscf.pbc.tools.pbc import super_cell

import numpy
import scipy
from scipy import sparse

from itertools import product


def minus_k(model, threshold=None, degeneracy_threshold=None):
    """
    Retrieves an array of indexes of negative k.
    Args:
        model (KSCF): an arbitrary mean-field pbc model;
        threshold (float): a threshold for determining the negative;
        degeneracy_threshold (float): a threshold for assuming degeneracy;

    Returns:
        A list of integers with indexes of the corresponding k-points.
    """
    if threshold is None:
        threshold = 1e-8
    if degeneracy_threshold is None:
        degeneracy_threshold = 1e-8
    kpts = model.cell.get_scaled_kpts(model.kpts)
    result = []
    for id_k, k in enumerate(kpts):
        delta = norm(((kpts + k[numpy.newaxis, :]) - .5) % 1 - .5, axis=-1)
        i = numpy.argmin(delta)
        if delta[i] > threshold:
            raise RuntimeError("Could not find a negative k-point for k={} (ID: {:d}, best difference: {:.3e})".format(
                repr(k), id_k, delta[i]
            ))
        if abs(model.mo_energy[id_k] - model.mo_energy[i]).max() > degeneracy_threshold:
            raise RuntimeError("Non-symmetric band structure (time-reversal) at k={} (ID: {:d}) and k={} (ID: {:d}). "
                               "This prevents composing real-valued orbitals".format(
                                    repr(k), id_k, repr(kpts[i]), i,
                                ))
        result.append(i)
    return result


def assert_scf_converged(model, threshold=1e-7):
    """
    Tests if scf is converged.
    Args:
        model (KSCF): a model to test;
        threshold (float): threshold for eigenvalue comparison;

    Returns:
        True if model is converged and False otherwise.
    """
    ovlp = model.get_ovlp()
    fock = model.get_fock(dm=model.make_rdm1())
    # dm = model.make_rdm1()
    # fock = model.get_fock(
    #     model.get_hcore(model.cell),
    #     ovlp,
    #     model.get_veff(model.cell, dm),
    #     dm,
    # )
    for k, (m, e, v, o) in enumerate(zip(fock, model.mo_energy, model.mo_coeff, ovlp)):
        delta = norm(numpy.dot(m, v) - e[numpy.newaxis, :] * numpy.dot(o, v), axis=0)
        nabove = (delta > threshold).sum()

        eye = reduce(numpy.dot, (v.conj().T, o, v))
        delta_o = abs(eye - numpy.eye(eye.shape[0])).max()
        if nabove > 0:
            raise AssertionError("{:d} vectors at k={:d} are not converged, max difference: {:.3e}; orthogonality error: {:.3e}".format(
                nabove, k, max(delta), delta_o,
            ))


def sparse_transform(m, *args):
    """
    Performs a sparse transform of a dense tensor.
    Args:
        m (ndarray): a tensor to transform;
        *args: alternating indexes and bases to transform into;

    Returns:
        The transformed tensor.
    """
    result = m
    for i, (index, basis) in enumerate(zip(args[::2], args[1::2])):

        if len(basis.shape) != 2:
            raise ValueError("Transform {:d} is not a matrix: shape = {}".format(
                i, repr(basis.shape)
            ))
        if result.shape[index] != basis.shape[0]:
            raise ValueError("Dimension mismatch of transform {:d}: m.shape[{:d}] = {:d} != basis.shape[0] = {:d}".format(
                i, index, result.shape[index], basis.shape[0],
            ))
        if "getcol" not in dir(basis):
            raise ValueError("No 'getcol' in the transform matrix {:d}: not a CSC sparse matrix?")

        result_shape = result.shape[:index] + (basis.shape[1],) + result.shape[index + 1:]
        new_result = numpy.zeros(result_shape, numpy.common_type(*(
            args[1::2] + (m,)
        )))
        for b2 in range(basis.shape[1]):
            slice_b2 = (slice(None),) * index + (b2,)
            col = basis.getcol(b2)
            assert len(col.indptr) == 2
            assert col.shape[1] == 1
            for b1, v in zip(col.indices, col.data):
                slice_b1 = (slice(None),) * index + (b1,)
                new_result[slice_b2] += v * result[slice_b1]
        result = new_result

    return result


def k2s(model, grid_spec, mf_constructor, threshold=None, degeneracy_threshold=None, imaginary_threshold=None):
    """
    Converts K to supercell with real orbitals.
    Args:
        model (KSCF): an arbitrary mean-field pbc model;
        grid_spec (Iterable): integer dimensions of the k-grid in KSCF;
        mf_constructor (Callable): a function constructing the mean-field object;
        threshold (float): a threshold for determining the negative k-point index;
        degeneracy_threshold (float): a threshold for assuming degeneracy when composing real-valued orbitals;
        imaginary_threshold (float): a threshold for asserting real-valued supercell orbitals;

    Returns:
        The same class where the Cell object was replaced by the supercell and all fields were adjusted accordingly.
    """
    if imaginary_threshold is None:
        imaginary_threshold = 1e-7

    mk = minus_k(model, threshold=threshold, degeneracy_threshold=degeneracy_threshold)
    nk = len(model.kpts)
    t_vecs = cartesian_prod(tuple(numpy.arange(i) for i in grid_spec))
    kpts_frac = model.cell.get_scaled_kpts(model.kpts)

    result = mf_constructor(super_cell(model.cell, grid_spec))
    result_ovlp = result.get_ovlp()[0]

    moe = numpy.concatenate(model.mo_energy)
    moo = numpy.concatenate(model.mo_occ)

    # Complex-valued wf in a supercell
    moc = []
    for mo_coeff, k in zip(model.mo_coeff, kpts_frac):
        psi = (
            mo_coeff[numpy.newaxis, ...] * numpy.exp(2.j * numpy.pi * t_vecs.dot(k))[:, numpy.newaxis, numpy.newaxis]
        ).reshape(-1, mo_coeff.shape[1])
        norms = einsum("ai,ab,bi->i", psi.conj(), result_ovlp, psi) ** .5
        psi /= norms[numpy.newaxis, :]
        moc.append(psi)
    moc = numpy.concatenate(moc, axis=1)

    rotation_matrix = sparse.csc_matrix(moc.shape, dtype=moc.dtype)
    inv_rotation_matrix = sparse.csc_matrix(moc.shape, dtype=moc.dtype)
    nvecs = (0,) + tuple(i.shape[1] for i in model.mo_coeff)
    nvecs = numpy.cumsum(nvecs)<commit_msg>Switch to a generic sparse matrix implementation<commit_after>#  Author: <NAME>
"""
This and other `_slow` modules implement the time-dependent Kohn-Sham 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 by proxying to pyscf density response routines with a O(N^4) complexity scaling. As a result,
regular `numpy.linalg.eig` can be used to retrieve TDKS roots in a reliable fashion without any issues related to the
Davidson procedure. Several variants of TDKS are available:

 * `pyscf.tdscf.rks_slow`: the molecular implementation;
 * `pyscf.pbc.tdscf.rks_slow`: PBC (periodic boundary condition) implementation for RKS objects of `pyscf.pbc.scf`
   modules;
 * (this module)`pyscf.pbc.tdscf.krks_slow_supercell`: PBC implementation for KRKS objects of `pyscf.pbc.scf` modules.
   Works with an arbitrary number of k-points but has a overhead due to an effective construction of a supercell.
 * `pyscf.pbc.tdscf.krks_slow_gamma`: A Gamma-point calculation resembling the original `pyscf.pbc.tdscf.krks`
   module. Despite its name, it accepts KRKS objects with an arbitrary number of k-points but finds only few TDKS roots
   corresponding to collective oscillations without momentum transfer;
 * `pyscf.pbc.tdscf.krks_slow`: PBC implementation for KRKS objects of `pyscf.pbc.scf` modules. Works with
   an arbitrary number of k-points and employs k-point conservation (diagonalizes matrix blocks separately).
"""

# Convention for these modules:
# * PhysERI, PhysERI4, PhysERI8 are proxy classes for computing the full TDDFT matrix
# * vector_to_amplitudes reshapes and normalizes the solution
# * TDRKS provides a container

from pyscf.tdscf.common_slow import TDProxyMatrixBlocks, PeriodicMFMixin
from pyscf.tdscf import rks_slow
from pyscf.pbc.tdscf import krhf_slow_supercell, KTDDFT
from pyscf.lib import einsum, cartesian_prod, norm, logger
from pyscf.pbc.tools.pbc import super_cell

import numpy
import scipy
from scipy import sparse

from itertools import product


def minus_k(model, threshold=None, degeneracy_threshold=None):
    """
    Retrieves an array of indexes of negative k.
    Args:
        model (KSCF): an arbitrary mean-field pbc model;
        threshold (float): a threshold for determining the negative;
        degeneracy_threshold (float): a threshold for assuming degeneracy;

    Returns:
        A list of integers with indexes of the corresponding k-points.
    """
    if threshold is None:
        threshold = 1e-8
    if degeneracy_threshold is None:
        degeneracy_threshold = 1e-8
    kpts = model.cell.get_scaled_kpts(model.kpts)
    result = []
    for id_k, k in enumerate(kpts):
        delta = norm(((kpts + k[numpy.newaxis, :]) - .5) % 1 - .5, axis=-1)
        i = numpy.argmin(delta)
        if delta[i] > threshold:
            raise RuntimeError("Could not find a negative k-point for k={} (ID: {:d}, best difference: {:.3e})".format(
                repr(k), id_k, delta[i]
            ))
        if abs(model.mo_energy[id_k] - model.mo_energy[i]).max() > degeneracy_threshold:
            raise RuntimeError("Non-symmetric band structure (time-reversal) at k={} (ID: {:d}) and k={} (ID: {:d}). "
                               "This prevents composing real-valued orbitals".format(
                                    repr(k), id_k, repr(kpts[i]), i,
                                ))
        result.append(i)
    return result


def assert_scf_converged(model, threshold=1e-7):
    """
    Tests if scf is converged.
    Args:
        model (KSCF): a model to test;
        threshold (float): threshold for eigenvalue comparison;

    Returns:
        True if model is converged and False otherwise.
    """
    ovlp = model.get_ovlp()
    fock = model.get_fock(dm=model.make_rdm1())
    # dm = model.make_rdm1()
    # fock = model.get_fock(
    #     model.get_hcore(model.cell),
    #     ovlp,
    #     model.get_veff(model.cell, dm),
    #     dm,
    # )
    for k, (m, e, v, o) in enumerate(zip(fock, model.mo_energy, model.mo_coeff, ovlp)):
        delta = norm(numpy.dot(m, v) - e[numpy.newaxis, :] * numpy.dot(o, v), axis=0)
        nabove = (delta > threshold).sum()

        eye = reduce(numpy.dot, (v.conj().T, o, v))
        delta_o = abs(eye - numpy.eye(eye.shape[0])).max()
        if nabove > 0:
            raise AssertionError("{:d} vectors at k={:d} are not converged, max difference: {:.3e}; orthogonality error: {:.3e}".format(
                nabove, k, max(delta), delta_o,
            ))


def sparse_transform(m, *args):
    """
    Performs a sparse transform of a dense tensor.
    Args:
        m (ndarray): a tensor to transform;
        *args: alternating indexes and bases to transform into;

    Returns:
        The transformed tensor.
    """
    result = m
    for i, (index, basis) in enumerate(zip(args[::2], args[1::2])):

        if len(basis.shape) != 2:
            raise ValueError("Transform {:d} is not a matrix: shape = {}".format(
                i, repr(basis.shape)
            ))
        if result.shape[index] != basis.shape[0]:
            raise ValueError("Dimension mismatch of transform {:d}: m.shape[{:d}] = {:d} != basis.shape[0] = {:d}".format(
                i, index, result.shape[index], basis.shape[0],
            ))
        if "getcol" not in dir(basis):
            raise ValueError("No 'getcol' in the transform matrix {:d}: not a CSC sparse matrix?")

        result_shape = result.shape[:index] + (basis.shape[1],) + result.shape[index + 1:]
        new_result = numpy.zeros(result_shape, numpy.common_type(*(
            args[1::2] + (m,)
        )))
        for b2 in range(basis.shape[1]):
            slice_b2 = (slice(None),) * index + (b2,)
            col = basis.getcol(b2)
            for b1 in col.nonzero()[0]:
                slice_b1 = (slice(None),) * index + (b1,)
                new_result[slice_b2] += col[b1, 0] * result[slice_b1]
        result = new_result

    return result


def k2s(model, grid_spec, mf_constructor, threshold=None, degeneracy_threshold=None, imaginary_threshold=None):
    """
    Converts K to supercell with real orbitals.
    Args:
        model (KSCF): an arbitrary mean-field pbc model;
        grid_spec (Iterable): integer dimensions of the k-grid in KSCF;
        mf_constructor (Callable): a function constructing the mean-field object;
        threshold (float): a threshold for determining the negative k-point index;
        degeneracy_threshold (float): a threshold for assuming degeneracy when composing real-valued orbitals;
        imaginary_threshold (float): a threshold for asserting real-valued supercell orbitals;

    Returns:
        The same class where the Cell object was replaced by the supercell and all fields were adjusted accordingly.
    """
    if imaginary_threshold is None:
        imaginary_threshold = 1e-7

    mk = minus_k(model, threshold=threshold, degeneracy_threshold=degeneracy_threshold)
    nk = len(model.kpts)
    t_vecs = cartesian_prod(tuple(numpy.arange(i) for i in grid_spec))
    kpts_frac = model.cell.get_scaled_kpts(model.kpts)

    result = mf_constructor(super_cell(model.cell, grid_spec))
    result_ovlp = result.get_ovlp()[0]

    moe = numpy.concatenate(model.mo_energy)
    moo = numpy.concatenate(model.mo_occ)

    # Complex-valued wf in a supercell
    moc = []
    for mo_coeff, k in zip(model.mo_coeff, kpts_frac):
        psi = (
            mo_coeff[numpy.newaxis, ...] * numpy.exp(2.j * numpy.pi * t_vecs.dot(k))[:, numpy.newaxis, numpy.newaxis]
        ).reshape(-1, mo_coeff.shape[1])
        norms = einsum("ai,ab,bi->i", psi.conj(), result_ovlp, psi) ** .5
        psi /= norms[numpy.newaxis, :]
        moc.append(psi)
    moc = numpy.concatenate(moc, axis=1)

    rotation_matrix = sparse.dok_matrix(moc.shape, dtype=moc.dtype)
    inv_rotation_matrix = sparse.dok_matrix(moc.shape, dtype=moc.dtype)
    nvecs = (0,) + tuple(i.shape[1] for i in model.mo_coeff)
    nvecs = numpy.cumsum(nvecs)
{'fraction_non_alphanumeric': 0.11074074074074074, 'fraction_numerical': 0.006913580246913581, 'mean_word_length': 2.9066795273691826, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 15, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4724665', 'n_tokens_mistral': 5357, 'n_tokens_neox': 5079, 'n_words': 1516}