 kwargs['atom']

        if 'gs' in kwargs:
            self.gs = kwargs['gs']

        # Set-up pseudopotential if it exists
        # This must be done before build() because it affects
        # tot_electrons() via the call to .atom_charge()

        self.ecp, self.pseudo = classify_ecp_pseudo(self, self.ecp, self.pseudo)
        if self.pseudo is not None:
            _atom = self.format_atom(self.atom)
            # Unless explicitly input, ECP should not be assigned to ghost atoms
            uniq_atoms = set([a[0] for a in _atom if not is_ghost_atom(a[0])])
            if isinstance(self.pseudo, (str, unicode)):
                # specify global pseudo for whole molecule
                _pseudo = dict([(a, str(self.pseudo)) for a in uniq_atoms])
            elif 'default' in self.pseudo:
                default_pseudo = self.pseudo['default']
                _pseudo = dict(((a, default_pseudo) for a in uniq_atoms))
                _pseudo.update(self.pseudo)
                del(_pseudo['default'])
            else:
                _pseudo = self.pseudo
            self._pseudo = self.format_pseudo(_pseudo)

        # Do regular Mole.build
        _built = self._built
        mole.Mole.build(self, False, parse_arg, *args, **kwargs)

        exp_min = np.array([self.bas_exp(ib).min() for ib in range(self.nbas)])
        if self.exp_to_discard is None:
            if np.any(exp_min < 0.1):
                sys.stderr.write('''WARNING!
  Very diffused basis functions are found in the basis set. They may lead to severe
  linear dependence and numerical instability.  You can set  cell.exp_to_discard=0.1
  to remove the diffused Gaussians whose exponents are less than 0.1.\n\n''')
        elif np.any(exp_min < self.exp_to_discard):
            # Discard functions of small exponents in basis
            _basis = {}
            for symb, basis_now in self._basis.items():
                basis_add = []
                for b in basis_now:
                    l = b[0]
                    if isinstance(b[1], int):
                        kappa = b[1]
                        b_coeff = np.array(b[2:])
                    else:
                        kappa = 0
                        b_coeff = np.array(b[1:])
                    es = b_coeff[:,0]
                    if np.any(es < self.exp_to_discard):
                        b_coeff = b_coeff[es>=self.exp_to_discard]
# contraction coefficients may be completely zero after removing one primitive
# basis. Removing the zero-coefficient basis.
                        b_coeff = b_coeff[:,np.all(b_coeff!=0, axis=0)]
                        if b_coeff.size > 0:
                            if kappa == 0:
                                basis_add.append([l] + b_coeff.tolist())
                            else:
                                basis_add.append([l,kappa] + b_coeff.tolist())
                    else:
                        basis_add.append(b)
                _basis[symb] = basis_add
            self._basis = _basis

            steep_shls = []
            nprim_drop = 0
            nctr_drop = 0
            _env = self._env.copy()
            for ib in range(len(self._bas)):
                l = self.bas_angular(ib)
                nprim = self.bas_nprim(ib)
                nc = self.bas_nctr(ib)
                es = self.bas_exp(ib)
                ptr = self._bas[ib,mole.PTR_COEFF]
                cs = self._env[ptr:ptr+nprim*nc].reshape(nc,nprim).T

                if np.any(es < self.exp_to_discard):
                    cs = cs[es>=self.exp_to_discard]
                    es = es[es>=self.exp_to_discard]
                    nprim_old, nc_old = nprim, nc

# contraction coefficients may be completely zero after removing one primitive
# basis. Removing the zero-coefficient basis.
                    cs = cs[:,np.all(cs!=0, axis=0)]
                    nprim, nc = cs.shape
                    self._bas[ib,mole.NPRIM_OF] = nprim
                    self._bas[ib,mole.NCTR_OF] = nc

                    nprim_drop = nprim_old - nprim + nprim_drop
                    nctr_drop = nc_old - nc + nctr_drop
                    if cs.size > 0:
                        pe = self._bas[ib,mole.PTR_EXP]
                        _env[pe:pe+nprim] = es
                        cs = mole._nomalize_contracted_ao(l, es, cs)
                        _env[ptr:ptr+nprim*nc] = cs.T.reshape(-1)
                if nprim > 0:
                    steep_shls.append(ib)
            self._env = _env
            self._bas = np.asarray(self._bas[steep_shls], order='C')
            logger.info(self, 'Discarded %d diffused primitive functions, '
                        '%d contracted functions', nprim_drop, nctr_drop)
            #logger.debug1(self, 'Old shells %s', steep_shls)

        # The rest initialization requires lattice parameters.  If .a is not
        # set, pass the rest initialization.
        if self.a is None:
            if dump_input and not _built and self.verbose > logger.NOTE:
                self.dump_input()
            return self

        if self.rcut is None or self._rcut_from_build:
            self._rcut = max([self.bas_rcut(ib, self.precision)
                              for ib in range(self.nbas)] + [0])
            self._rcut_from_build = True

        _a = self.lattice_vectors()
        if np.linalg.det(_a) < 0:
            sys.stderr.write('''WARNING!
  Lattice are not in right-handed coordinate system. This can cause wrong value for some integrals.
  It's recommended to resort the lattice vectors to\na = %s\n\n''' % _a[[0,2,1]])

        if self.dimension == 2 and self.low_dim_ft_type != 'inf_vacuum':
            # check vacuum size. See Fig 1 of PRB, 73, 2015119
            #Lz_guess = self.rcut*(1+np.sqrt(2))
            Lz_guess = self.rcut * 2
            if np.linalg.norm(_a[2]) < 0.7 * Lz_guess:
                sys.stderr.write('''WARNING!
  Size of vacuum may not be enough. The recommended vacuum size is %s AA (%s Bohr)\n\n'''
                                 % (Lz_guess*param.BOHR, Lz_guess))

        if self.mesh is None or self._mesh_from_build:
            if self.ke_cutoff is None:
                ke_cutoff = estimate_ke_cutoff(self, self.precision)
            else:
                ke_cutoff = self.ke_cutoff
            self._mesh = pbctools.cutoff_to_mesh(_a, ke_cutoff)

            if (self.dimension < 2 or
                (self.dimension == 2 and self.low_dim_ft_type == 'inf_vacuum')):
                self._mesh[self.dimension:] = _mesh_inf_vaccum(self)
            self._mesh_from_build = True

            # Set minimal mesh grids to handle the case mesh==0. since Madelung
            # constant may be computed even if the unit cell has 0 atoms. In this
            # system, cell.mesh was initialized to 0.
            self._mesh[self._mesh == 0] = 30

        if dump_input and not _built and self.verbose > logger.NOTE:
            self.dump_input()
            logger.info(self, 'lattice vectors  a1 [%.9f, %.9f, %.9f]', *_a[0])
            logger.info(self, '                 a2 [%.9f, %.9f, %.9f]', *_a[1])
            logger.info(self, '                 a3 [%.9f, %.9f, %.9f]', *_a[2])
            logger.info(self, 'dimension = %s', self.dimension)
            logger.info(self, 'low_dim_ft_type = %s', self.low_dim_ft_type)
            logger.info(self, 'Cell volume = %g', self.vol)
            if self.exp_to_discard is not None:
                logger.info(self, 'exp_to_discard = %s', self.exp_to_discard)
            logger.info(self, 'rcut = %s (nimgs = %s)', self.rcut, self.nimgs)
            logger.info(self, 'lattice sum = %d cells', len(self.get_lattice_Ls()))
            logger.info(self, 'precision = %g', self.precision)
            logger.info(self, 'pseudo = %s', self.pseudo)
            if ke_cutoff is not None:
                logger.info(self, 'ke_cutoff = %s', ke_cutoff)
                logger.info(self, '    = %s mesh (%d PWs)',
                            self.mesh, np.prod(self.mesh))
            else:
                logger.info(self, 'mesh = %s (%d PWs)',
                            self.mesh, np.prod(self.mesh))
                Ecut = pbctools.mesh_to_cutoff(self.lattice_vectors(), self.mesh)
                logger.info(self, '    = ke_cutoff %s', Ecut)
        return self
    kernel = build

    @property
    def h(self):
        return np.asarray(self.a).T
    @h.setter
    def h(self, x):
        warnings.warn('cell.h is deprecated.  It is replaced by the '
                      '(row-based) lattice vectors cell.a:  cell.a = cell.h.T\n')
        if isinstance(x, (str, unicode)):
            x = x.replace(';',' ').replace(',',' ').replace('\n',' ')
            self.a = np.asarray([float(z) for z in x.split()]).reshape(3,3).T
        else:
            self.a = np.asarray(x).T

    @property
    def _h(self):
        return self.lattice_vectors().T

    @property
    def vol(self):
        return abs(np.linalg.det(self.lattice_vectors()))

    @property
    def Gv(self):
        return self.get_Gv()

    @lib.with_doc(format_pseudo.__doc__)
    def format_pseudo(self, pseudo_tab):
        return format_pseudo(pseudo_tab)

    @lib.with_doc(format_basis.__doc__)
    def format_basis(self, basis_tab):
        return format_basis(basis_tab)

    @property
    def gs(self):
        return [n//2 for n in self.mesh]
    @gs.setter
    def gs(self, x):
        warnings.warn('cell.gs is deprecated.  It is replaced by cell.mesh,'
                      'the number of PWs (=2*gs+1) along each direction.')
        self.mesh = [2*n+1 for n in x]

    @property
    def drop_exponent(self):
        return self.exp_to_discard
    @drop_exponent.setter
    def drop_exponent(self, x):
        self.exp_to_discard = x

    @property
    def nimgs(self):
        return self.get_bounding_sphere(self.rcut)
    @nimgs.setter
    def nimgs(self, x):
        b = self.reciprocal_vectors(norm_to=1)
        heights_inv = lib.norm(b, axis=1)
        self.rcut = max(np.asarray(x) / heights_inv)

        if self.nbas == 0:
            rcut_guess = _estimate_rcut(.05, 0, 1, 1e-8)
        else:
            rcut_guess = max([self.bas_rcut(ib, self.precision)
                              for ib in range(self.nbas)])
        if self.rcut > rcut_guess*1.5:
            msg = ('.nimgs is a deprecated attribute.  It is replaced by .rcut '
                   'attribute for lattic sum cutoff radius.  The given nimgs '
                   '%s is far over the estimated cutoff radius %s. ' %
                   (x, rcut_guess))
            warnings.warn(msg)

    def make_ecp_env(self, _atm, _ecp, pre_env=[]):
        if _ecp and self._pseudo:
            conflicts = set(self._pseudo.keys()).intersection(set(_ecp.keys()))
            if conflicts:
                raise RuntimeError('Pseudo potential for atoms %s are defined '
                                   'in both .ecp and .pseudo.' % list(conflicts))

        _ecpbas, _env = np.zeros((0,8)), pre_env
        if _ecp:
            _atm, _ecpbas, _env = mole.make_ecp_env(self, _atm, _ecp, _env)
        if self._pseudo:
            _atm, _, _env = make_pseudo_env(self, _atm, self._pseudo, _env)
        return _atm, _ecpbas, _env

    def lattice_vectors(self):
        '''Convert the primitive lattice vectors.

        Return 3x3 array in which each row represents one direction of the
        lattice vectors (unit in Bohr)
        '''
        if isinstance(self.a, (str, unicode)):
            a = self.a.replace(';',' ').replace(',',' ').replace('\n',' ')
            a = np.asarray([float(x) for x in a.split()]).reshape(3,3)
        else:
            a = np.asarray(self.a, dtype=np.double)
        if isinstance(self.unit, (str, unicode)):
            if self.unit.startswith(('B','b','au','AU')):
                return a
            else:
                return a/param.BOHR
        else:
            return a/self.unit

    def reciprocal_vectors(self, norm_to=2*np.pi):
        r'''
        .. math::

            \begin{align}
            \mathbf{b_1} &= 2\pi \frac{\mathbf{a_2} \times \mathbf{a_3}}{\mathbf{a_1} \cdot (\mathbf{a_2} \times \mathbf{a_3})} \\
            \mathbf{b_2} &= 2\pi \frac{\mathbf{a_3} \times \mathbf{a_1}}{\mathbf{a_2} \cdot (\mathbf{a_3} \times \mathbf{a_1})} \\
            \mathbf{b_3} &= 2\pi \frac{\mathbf{a_1} \times \mathbf{a_2}}{\mathbf{a_3} \cdot (\mathbf{a_1} \times \mathbf{a_2})}
            \end{align}

        '''  # noqa: E501
        a = self.lattice_vectors()
        if self.dimension == 1:
            assert(abs(np.dot(a[0], a[1])) < 1e-9 and
                   abs(np.dot(a[0], a[2])) < 1e-9 and
                   abs(np.dot(a[1], a[2])) < 1e-9)
        elif self.dimension == 2:
            assert(abs(np.dot(a[0], a[2])) < 1e-9 and
                   abs(np.dot(a[1], a[2])) < 1e-9)
        b = np.linalg.inv(a.T)
        return norm_to * b

    def get_abs_kpts(self, scaled_kpts):
        '''Get absolute k-points (in 1/Bohr), given "scaled" k-points in
        fractions of lattice vectors.

        Args:
            scaled_kpts : (nkpts, 3) ndarray of floats

        Returns:
            abs_kpts : (nkpts, 3) ndarray of floats
        '''
        return np.dot(scaled_kpts, self.reciprocal_vectors())

    def get_scaled_kpts(self, abs_kpts):
        '''Get scaled k-points, given absolute k-points in 1/Bohr.

        Args:
            abs_kpts : (nkpts, 3) ndarray of floats

        Returns:
            scaled_kpts : (nkpts, 3) ndarray of floats
        '''
        return 1./(2*np.pi)*np.dot(abs_kpts, self.lattice_vectors().T)

    make_kpts = get_kpts = make_kpts

    def copy(self):
        return copy(self)

    pack = pack

    @classmethod
    @lib.with_doc(unpack.__doc__)
    def unpack(cls, moldic):
        return unpack(moldic)

    @lib.with_doc(unpack.__doc__)
    def unpack_(self, moldic):
        self.__dict__.update(moldic)
        return self

    dumps = dumps

    @classmethod
    @lib.with_doc(loads.__doc__)
    def loads(cls, molstr):
        return loads(molstr)

    @lib.with_doc(unpack.__doc__)
    def loads_(self, molstr):
        self.__dict__.update(loads(molstr).__dict__)
        return self

    bas_rcut = bas_rcut

    get_lattice_Ls = pbctools.get_lattice_Ls

    get_nimgs = get_nimgs

    get_ewald_params = get_ewald_params

    get_bounding_sphere = get_bounding_sphere

    get_Gv = get_Gv
    get_Gv_weights = get_Gv_weights

    get_SI = get_SI

    ewald = ewald
    energy_nuc = ewald

    gen_uniform_grids = get_uniform_grids = get_uniform_grids

    __add__ = conc_cell

    def pbc_intor(self, intor, comp=None, hermi=0, kpts=None, kpt=None,
                  shls_slice=None, **kwargs):
        r'''One-electron integrals with PBC.

        .. math::

            \sum_T \int \mu(r) * [intor] * \nu (r-T) dr

        See also Mole.intor
        '''
        if not self._built:
            logger.warn(self, 'Warning: intor envs of %s not initialized.', self)
            # FIXME: Whether to check _built and call build?  ._bas and .basis
            # may not be consistent. calling .build() may leads to wrong intor env.
            #self.build(False, False)
        return intor_cross(intor, self, self, comp, hermi, kpts, kpt,
                           shls_slice, **kwargs)

    @lib.with_doc(pbc_eval_gto.__doc__)
    def pbc_eval_gto(self, eval_name, coords, comp=None, kpts=None, kpt=None,
                     shls_slice=None, non0tab=None, ao_loc=None, out=None):
        return pbc_eval_gto(self, eval_name, coords, comp, kpts, kpt,
                            shls_slice, non0tab, ao_loc, out)
    pbc_eval_ao = pbc_eval_gto

    @lib.with_doc(pbc_eval_gto.__doc__)
    def eval_gto(self, eval_name, coords, comp=None, kpts=None, kpt=None,
                 shls_slice=None, non0tab=None, ao_loc=None, out=None):
        if eval_name[:3] == 'PBC':
            return self.pbc_eval_gto(eval_name, coords, comp, kpts, kpt,
                                     shls_slice, non0tab, ao_loc, out)
        else:
            return mole.eval_gto(self, eval_name, coords, comp,
                                 shls_slice, non0tab, ao_loc, out)
    eval_ao = eval_gto

    def from_ase(self, ase_atom):
        '''Update cell based on given ase atom object

        Examples:

        >>> from ase.lattice import bulk
        >>> cell.from_ase(bulk('C', 'diamond', a=LATTICE_CONST))
        '''
        from pyscf.pbc.tools import pyscf_ase
        self.a = ase_atom.cell
        self.atom = pyscf_ase.ase_atoms_to_pyscf(ase_atom)
        return self

    def to_mol(self):
        '''Return a Mole object using the same atoms and basis functions as
        the Cell object.
        '''
        #FIXME: should cell be converted to mole object?  If cell is converted
        # and a mole object is returned, many attributes (e.g. the GTH basis,
        # gth-PP) will not be recognized by mole.build function.
        mol = self.view(mole.Mole)
        delattr(mol, 'a')
        delattr(mol, '_mesh')
        if mol.symmetry:
            mol._build_symmetry()
        return mol

    def has_ecp(self):
        '''Whether pseudo potential is used in the system.'''
        return self.pseudo or self._pseudo or (len(self._ecpbas) > 0)

    def ao2mo(self, mo_coeffs, intor='int2e', erifile=None, dataname='eri_mo',
              **kwargs):
        raise NotImplementedError

del(INTEGRAL_PRECISION, WRAP_AROUND, WITH_GAMMA, EXP_DELIMITER)


# pyscf/pbc/gto/ecp.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Short range part of ECP under PBC
'''

from functools import reduce
import copy
import numpy
from pyscf import lib
from pyscf.pbc import gto
from pyscf.gto import AS_ECPBAS_OFFSET, AS_NECPBAS


def ecp_int(cell, kpts=None):
    from pyscf.pbc.df import incore
    if kpts is None:
        kpts_lst = numpy.zeros((1,3))
    else:
        kpts_lst = numpy.reshape(kpts, (-1,3))

    cell, contr_coeff = gto.cell._split_basis(cell)
    lib.logger.debug1(cell, 'nao %d -> nao %d', *(contr_coeff.shape))

    ecpcell = copy.copy(cell)
    # append a fictitious s function to mimic the auxiliary index in pbc.incore.
    # ptr2last_env_idx to force PBCnr3c_fill_* function to copy the entire "env"
    ptr2last_env_idx = len(cell._env) - 1
    ecpbas = numpy.vstack([[0, 0, 1, 1, 0, ptr2last_env_idx, 0, 0],
                           cell._ecpbas]).astype(numpy.int32)
    ecpcell._bas = ecpbas
    # In pbc.incore _ecpbas is appended to two sets of cell._bas and the
    # fictitious s function.
    cell._env[AS_ECPBAS_OFFSET] = cell.nbas * 2 + 1
    cell._env[AS_NECPBAS] = len(cell._ecpbas)
    # shls_slice of auxiliary index (0,1) corresponds to the fictitious s function
    shls_slice = (0, cell.nbas, 0, cell.nbas, 0, 1)

    kptij_lst = numpy.hstack((kpts_lst,kpts_lst)).reshape(-1,2,3)
    buf = incore.aux_e2(cell, ecpcell, 'ECPscalar', aosym='s2',
                        kptij_lst=kptij_lst, shls_slice=shls_slice)
    buf = buf.reshape(len(kpts_lst),-1)
    mat = []
    for k, kpt in enumerate(kpts_lst):
        v = lib.unpack_tril(buf[k], lib.HERMITIAN)
        if abs(kpt).sum() < 1e-9:  # gamma_point:
            v = v.real
        mat.append(reduce(numpy.dot, (contr_coeff.T, v, contr_coeff)))
    if kpts is None or numpy.shape(kpts) == (3,):
        mat = mat[0]
    return mat



# pyscf/pbc/gto/eval_gto.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import ctypes
import numpy
from pyscf import lib
from pyscf.gto import moleintor
from pyscf.gto.eval_gto import _get_intor_and_comp, BLKSIZE
from pyscf.pbc.gto import _pbcintor
from pyscf import __config__

EXTRA_PREC = getattr(__config__, 'pbc_gto_eval_gto_extra_precision', 1e-2)

libpbc = _pbcintor.libpbc

def eval_gto(cell, eval_name, coords, comp=None, kpts=None, kpt=None,
             shls_slice=None, non0tab=None, ao_loc=None, out=None):
    r'''Evaluate PBC-AO function value on the given grids,

    Args:
        eval_name : str

            ==========================  =======================
            Function                    Expression
            ==========================  =======================
            "GTOval_sph"                \sum_T exp(ik*T) |AO>
            "GTOval_ip_sph"             nabla \sum_T exp(ik*T) |AO>
            "GTOval_cart"               \sum_T exp(ik*T) |AO>
            "GTOval_ip_cart"            nabla \sum_T exp(ik*T) |AO>
            ==========================  =======================

        atm : int32 ndarray
            libcint integral function argument
        bas : int32 ndarray
            libcint integral function argument
        env : float64 ndarray
            libcint integral function argument

        coords : 2D array, shape (N,3)
            The coordinates of the grids.

    Kwargs:
        shls_slice : 2-element list
            (shl_start, shl_end).
            If given, only part of AOs (shl_start <= shell_id < shl_end) are
            evaluated.  By default, all shells defined in cell will be evaluated.
        non0tab : 2D bool array
            mask array to indicate whether the AO values are zero.  The mask
            array can be obtained by calling :func:`dft.gen_grid.make_mask`
        out : ndarray
            If provided, results are written into this array.

    Returns:
        A list of 2D (or 3D) arrays to hold the AO values on grids.  Each
        element of the list corresponds to a k-point and it has the shape
        (N,nao) Or shape (\*,N,nao).

    Examples:

    >>> cell = pbc.gto.M(a=numpy.eye(3)*4, atom='He 1 1 1', basis='6-31g')
    >>> coords = cell.get_uniform_grids([20,20,20])
    >>> kpts = cell.make_kpts([3,3,3])
    >>> ao_value = cell.pbc_eval_gto("GTOval_sph", coords, kpts)
    >>> len(ao_value)
    27
    >>> ao_value[0].shape
    (100, 2)
    >>> ao_value = cell.pbc_eval_gto("GTOval_ig_sph", coords, kpts, comp=3)
    >>> print(ao_value.shape)
    >>> len(ao_value)
    27
    >>> ao_value[0].shape
    (3, 100, 2)
    '''
    if eval_name[:3] == 'PBC':  # PBCGTOval_xxx
        eval_name, comp = _get_intor_and_comp(cell, eval_name[3:], comp)
    else:
        eval_name, comp = _get_intor_and_comp(cell, eval_name, comp)
    eval_name = 'PBC' + eval_name

    atm = numpy.asarray(cell._atm, dtype=numpy.int32, order='C')
    bas = numpy.asarray(cell._bas, dtype=numpy.int32, order='C')
    env = numpy.asarray(cell._env, dtype=numpy.double, order='C')
    natm = atm.shape[0]
    nbas = bas.shape[0]
    if kpts is None:
        if kpt is not None:
            kpts_lst = numpy.reshape(kpt, (1,3))
        else:
            kpts_lst = numpy.zeros((1,3))
    else:
        kpts_lst = numpy.reshape(kpts, (-1,3))
    nkpts = len(kpts_lst)
    ngrids = len(coords)

    if non0tab is None:
        non0tab = numpy.empty(((ngrids+BLKSIZE-1)//BLKSIZE, nbas),
                              dtype=numpy.uint8)
# non0tab stores the number of images to be summed in real space.
# Initializing it to 255 means all images are summed
        non0tab[:] = 0xff

    if ao_loc is None:
        ao_loc = moleintor.make_loc(bas, eval_name)
    if shls_slice is None:
        shls_slice = (0, nbas)
    sh0, sh1 = shls_slice
    nao = ao_loc[sh1] - ao_loc[sh0]

    out = numpy.empty((nkpts,comp,nao,ngrids), dtype=numpy.complex128)
    coords = numpy.asarray(coords, order='F')

    # For atoms near the boundary of the cell, it is necessary (even in low-
    # dimensional systems) to include lattice translations in all 3 dimensions.
    if cell.dimension < 2 or cell.low_dim_ft_type == 'inf_vacuum':
        Ls = cell.get_lattice_Ls(dimension=cell.dimension)
    else:
        Ls = cell.get_lattice_Ls(dimension=3)
    Ls = Ls[numpy.argsort(lib.norm(Ls, axis=1))]
    expLk = numpy.exp(1j * numpy.asarray(numpy.dot(Ls, kpts_lst.T), order='C'))
    rcut = _estimate_rcut(cell)

    drv = getattr(libpbc, eval_name)
    drv(ctypes.c_int(ngrids),
        (ctypes.c_int*2)(*shls_slice), ao_loc.ctypes.data_as(ctypes.c_void_p),
        Ls.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(len(Ls)),
        expLk.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nkpts),
        out.ctypes.data_as(ctypes.c_void_p),
        coords.ctypes.data_as(ctypes.c_void_p),
        rcut.ctypes.data_as(ctypes.c_void_p),
        non0tab.ctypes.data_as(ctypes.c_void_p),
        atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(natm),
        bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(nbas),
        env.ctypes.data_as(ctypes.c_void_p))

    ao_kpts = []
    for k, kpt in enumerate(kpts_lst):
        v = out[k]
        if abs(kpt).sum() < 1e-9:
            v = numpy.asarray(v.real, order='C')
        v = v.transpose(0,2,1)
        if comp == 1:
            v = v[0]
        ao_kpts.append(v)

    if kpts is None or numpy.shape(kpts) == (3,):  # A single k-point
        ao_kpts = ao_kpts[0]
    return ao_kpts

pbc_eval_gto = eval_gto

def _estimate_rcut(cell):
    '''Cutoff raidus, above which each shell decays to a value less than the
    required precsion'''
    log_prec = numpy.log(cell.precision * EXTRA_PREC)
    rcut = []
    for ib in range(cell.nbas):
        l = cell.bas_angular(ib)
        es = cell.bas_exp(ib)
        cs = abs(cell.bas_ctr_coeff(ib)).max(axis=1)
        r = 5.
        r = (((l+2)*numpy.log(r)+numpy.log(cs) - log_prec) / es)**.5
        r = (((l+2)*numpy.log(r)+numpy.log(cs) - log_prec) / es)**.5
        rcut.append(r.max())
    return numpy.array(rcut)


if __name__ == '__main__':
    from pyscf.pbc import gto
    cell = gto.M(a=numpy.eye(3)*4, atom='He 1 1 1', basis=[[2,(1,.5),(.5,.5)]])
    coords = cell.get_uniform_grids([10]*3)
    ao_value = eval_gto(cell, "GTOval_sph", coords, kpts=cell.make_kpts([3]*3))
    print(lib.finger(numpy.asarray(ao_value)) - (-0.27594803231989179+0.0064644591759109114j))

    cell = gto.M(a=numpy.eye(3)*4, atom='He 1 1 1', basis=[[2,(1,.5),(.5,.5)]])
    coords = cell.get_uniform_grids([10]*3)
    ao_value = eval_gto(cell, "GTOval_ip_cart", coords, kpts=cell.make_kpts([3]*3))
    print(lib.finger(numpy.asarray(ao_value)) - (0.38051517609460028+0.062526488684770759j))



# pyscf/pbc/gw/__init__.py

from .kgw_slow import GW as KRGW


# pyscf/pbc/gw/gw_slow.py

"""
This module implements the G0W0 approximation on top of `pyscf.tdscf.rhf_slow` TDHF implementation. Unlike `gw.py`, all
integrals are stored in memory. Several variants of GW are available:

 * `pyscf.gw_slow`: the molecular implementation;
 * (this module) `pyscf.pbc.gw.gw_slow`: single-kpoint PBC (periodic boundary condition) implementation;
 * `pyscf.pbc.gw.kgw_slow_supercell`: a supercell approach to PBC implementation with multiple k-points. Runs the
   molecular code for a model with several k-points for the cost of discarding momentum conservation and using dense
   instead of sparse matrixes;
 * `pyscf.pbc.gw.kgw_slow`: a PBC implementation with multiple k-points;
"""

# Convention for these modules:
# * IMDS contains routines for intermediates
# * kernel finds GW roots
# * GW provides a container

# This module and kgw_slow_supercell are simply aliases of the molecular code
from pyscf.gw import gw_slow


IMDS = gw_slow.IMDS
kernel = gw_slow.kernel
GW = gw_slow.GW


# pyscf/pbc/gw/kgw_slow.py

# flake8: noqa
"""
This module implements the G0W0 approximation on top of `pyscf.tdscf.rhf_slow` TDHF implementation. Unlike `gw.py`, all
integrals are stored in memory. Several variants of GW are available:

 * `pyscf.gw_slow`: the molecular implementation;
 * `pyscf.pbc.gw.gw_slow`: single-kpoint PBC (periodic boundary condition) implementation;
 * `pyscf.pbc.gw.kgw_slow_supercell`: a supercell approach to PBC implementation with multiple k-points. Runs the
   molecular code for a model with several k-points for the cost of discarding momentum conservation and using dense
   instead of sparse matrixes;
 * (this module) `pyscf.pbc.gw.kgw_slow`: a PBC implementation with multiple k-points;
"""

from pyscf.gw import gw_slow
from pyscf.pbc.gw import kgw_slow_supercell
from pyscf.lib import einsum, direct_sum
from pyscf.pbc.tdscf.krhf_slow import get_block_k_ix

import numpy


# Convention for these modules:
# * IMDS contains routines for intermediates
# * kernel finds GW roots
# * GW provides a container


class IMDS(kgw_slow_supercell.IMDS):
    def __init__(self, td, eri=None):
        """
        GW intermediates (k-version).
        Args:
        Args:
            td: a container with TD solution;
            eri: a container with electron repulsion integrals;

        """
        gw_slow.AbstractIMDS.__init__(self, td, eri=eri)

        self.nk = len(self.td._scf.mo_energy)

        # MF
        self.nocc = self.eri.nocc
        self.o = tuple(e[:nocc] for e, nocc in zip(self.eri.mo_energy, self.eri.nocc))
        self.v = tuple(e[nocc:] for e, nocc in zip(self.eri.mo_energy, self.eri.nocc))

        # TD
        self.td_xy = self.td.xy
        self.td_e = self.td.e

        self.tdm = self.construct_tdm()

    def eri_ov(self, item):
        item, k1, k2, k3, k4 = item
        return self.eri.eri_ov(item, (k1, k2, k3, k4)) / self.nk

    __getitem__ = eri_ov

    def construct_tdm(self):

        # Indexes of td_x:
        # - k_transfer
        # - k
        # - o: k
        # - v: fw[k]

        # Indexes of td_y:
        # - k_transfer
        # - k
        # - o: k
        # - v: bw[k]

        # Original code:
        # tdm_oo = einsum('vxia,ipaq->vxpq', td_xy, self["oovo"])
        # tdm_ov = einsum('vxia,ipaq->vxpq', td_xy, self["oovv"])
        # tdm_vv = einsum('vxia,ipaq->vxpq', td_xy, self["ovvv"])

        # ERI k:
        # ki, kp, ka=?, kq

        # Now fw[kp] = kq, bw[kq] = kp -> bw[ki] = ka, fw[ka] = ki
        # for x amplitudes, the transfer is bw[0] such that ov -> k, bw[k]
        # for y amplitudes, the transfer is also fw[0] such that ov -> k, bw[k]

        result = []

        for k_transfer in range(self.nk):
            xy_k = self.td_xy[k_transfer]

            fw, bw, _, _ = get_block_k_ix(self.eri, k_transfer)
            result.append([[], []])

            for xy_kx, ix_fw, ix_bw, storage in (
                (xy_k[:, 0], fw, bw, result[k_transfer][0]),  # X
                (xy_k[:, 1], bw, fw, result[k_transfer][1]),  # Y
            ):

                for kp in range(self.nk):

                    tdm_oo = tdm_ov = tdm_vv = 0

                    tdm_vo = 0

                    for ki in range(self.nk):

                        x = xy_kx[:, ki] * 2
                        tdm_oo = tdm_oo + einsum('via,ipaq->vpq', x, self["oovo", ki, kp, ix_fw[ki], ix_bw[kp]])
                        tdm_ov = tdm_ov + einsum('via,ipaq->vpq', x, self["oovv", ki, kp, ix_fw[ki], ix_bw[kp]])
                        tdm_vv = tdm_vv + einsum('via,ipaq->vpq', x, self["ovvv", ki, kp, ix_fw[ki], ix_bw[kp]])
                        tdm_vo = tdm_vo + einsum('via,ipaq->vpq', x, self["ovvo", ki, kp, ix_fw[ki], ix_bw[kp]])

                    tdm = numpy.concatenate(
                        (
                            numpy.concatenate((tdm_oo, tdm_ov), axis=2),
                            numpy.concatenate((tdm_vo, tdm_vv), axis=2)
                        ),
                        axis=1,
                    )

                    storage.append(tdm)
        # The output is the following:
        # for each kp, kq pair two 3-tensors are given
        # The last two indexes in each tensor correspond to kp, kq
        # Given fw, bw, _, _ = get_block_k_ix(self.eri, (kp, kq)),
        # The first index of the two tensors will correspond to tdhf.e[bw[0]], tdhf.e[fw[0]] correspondingly
        return numpy.array(result)

    def get_sigma_element(self, omega, p, eta, vir_sgn=1):
        k, p = p

        # Molecular implementation
        # ------------------------
        # tdm = self.tdm.sum(axis=1)
        # evi = direct_sum('v-i->vi', self.td_e, self.o)
        # eva = direct_sum('v+a->va', self.td_e, self.v)
        # sigma = numpy.sum(tdm[:, :self.nocc, p] ** 2 / (omega + evi - 1j * eta))
        # sigma += numpy.sum(tdm[:, self.nocc:, p] ** 2 / (omega - eva + vir_sgn * 1j * eta))
        # return sigma

        sigma = 0
        for k_transfer, tdm_k in enumerate(self.tdm):
            fw, bw, _, _ = get_block_k_ix(self.eri, k_transfer)

            same = fw == bw
            different = numpy.logical_not(same)

            terms = []

            if same.sum() > 0:
                terms.append((tdm_k[0] + tdm_k[1], fw[same], fw[same]))

            if different.sum() > 0:
                terms.append((tdm_k[0], bw[different], fw[different]))
                terms.append((tdm_k[1], fw[different], bw[different]))

            for tdm_kx, ix_fw, ix_bw in terms:
                k1, k2 = ix_bw[k], k
                evi = direct_sum('v-i->vi', self.td_e[k_transfer], self.o[k1])
                eva = direct_sum('v+a->va', self.td_e[k_transfer], self.v[k1])
                sigma += numpy.sum(tdm_kx[k1][:, :self.nocc[k1], p] ** 2 / (omega + evi - 1j * eta))
                sigma += numpy.sum(tdm_kx[k1][:, self.nocc[k1]:, p] ** 2 / (omega - eva + vir_sgn * 1j * eta))
        return sigma


kernel = gw_slow.kernel


class GW(gw_slow.GW):
    base_imds = IMDS


# pyscf/pbc/gw/kgw_slow_supercell.py

"""
This module implements the G0W0 approximation on top of `pyscf.tdscf.rhf_slow` TDHF implementation. Unlike `gw.py`, all
integrals are stored in memory. Several variants of GW are available:

 * `pyscf.gw_slow`: the molecular implementation;
 * `pyscf.pbc.gw.gw_slow`: single-kpoint PBC (periodic boundary condition) implementation;
 * (this module) `pyscf.pbc.gw.kgw_slow_supercell`: a supercell approach to PBC implementation with multiple k-points.
   Runs the molecular code for a model with several k-points for the cost of discarding momentum conservation and using
   dense instead of sparse matrixes;
 * `pyscf.pbc.gw.kgw_slow`: a PBC implementation with multiple k-points;
"""

from pyscf.gw import gw_slow
from pyscf.lib import einsum

import numpy


# Convention for these modules:
# * IMDS contains routines for intermediates
# * kernel finds GW roots
# * GW provides a container


def corrected_moe(eri, k, p):
    """
    Calculates the corrected orbital energy.
    Args:
        eri (PhysERI): a container with electron repulsion integrals;
        k (int): the k-point index;
        p (int): orbital;

    Returns:
        The corrected orbital energy.
    """
    moe = eri.mo_energy[k][p]
    moc = eri.mo_coeff[k][:, p]
    nk = len(eri.mo_energy)
    vk = 0
    for k2 in range(nk):
        vk -= eri.ao2mo_k((
            moc[:, numpy.newaxis],
            eri.mo_coeff_full[k2][:, :eri.nocc_full[k2]],
            eri.mo_coeff_full[k2][:, :eri.nocc_full[k2]],
            moc[:, numpy.newaxis],
        ), (k, k2, k2, k)).squeeze().trace().real
    vk /= nk
    mf = eri.model
    v_mf = mf.get_veff()[k] - mf.get_j()[k]
    v_mf = einsum("i,ij,j", moc.conj(), v_mf, moc).real
    return moe + vk - v_mf


class IMDS(gw_slow.IMDS):
    orb_dims = 2

    def __init__(self, td, eri=None):
        """
        GW intermediates (k-version/supercell).
        Args:
            td: a container with TD solution;
            eri: a container with electron repulsion integrals;
        """
        gw_slow.AbstractIMDS.__init__(self, td, eri=eri)

        self.nk = len(self.td._scf.mo_energy)

        # MF
        self.nocc = sum(self.eri.nocc)
        self.o = numpy.concatenate(tuple(e[:nocc] for e, nocc in zip(self.eri.mo_energy, self.eri.nocc)))
        self.v = numpy.concatenate(tuple(e[nocc:] for e, nocc in zip(self.eri.mo_energy, self.eri.nocc)))

        # TD
        nroots, _, k1, k2, o, v = self.td.xy.shape
        self.td_xy = self.td.xy.transpose(0, 1, 2, 4, 3, 5).reshape(nroots, 2, k1*o, k2*v)
        self.td_e = self.td.e

        self.tdm = self.construct_tdm()

    def eri_ov(self, item):
        result = []
        k = numpy.arange(self.nk)
        for k1 in k:
            result.append([])
            for k2 in k:
                result[-1].append([])
                for k3 in k:
                    result[-1][-1].append([])
                    for k4 in k:
                        result[-1][-1][-1].append(self.eri.eri_ov(item, (k1, k2, k3, k4)))

        r = numpy.block(result)
        return r / len(k)

    __getitem__ = eri_ov

    def __plain_index__(self, p, spec=True):
        k, kp = p
        if kp < self.eri.nocc[k]:
            x = sum(self.eri.nocc[:k]) + kp
            if spec:
                return "o", x
            else:
                return x
        else:
            kp -= self.eri.nocc[k]
            x = sum(self.eri.nmo[:k]) - sum(self.eri.nocc[:k]) + kp
            if spec:
                return "v", x
            else:
                return x + self.nocc

    def get_rhs(self, p, components=False):
        k, kp = p
        # return self.eri.mo_energy[k][kp]
        return corrected_moe(self.eri, k, kp)

    def get_sigma_element(self, omega, p, eta, vir_sgn=1):
        return super(IMDS, self).get_sigma_element(omega, self.__plain_index__(p, spec=False), eta, vir_sgn=vir_sgn)

    def initial_guess(self, p):
        k, kp = p
        return self.eri.mo_energy[k][kp]

    @property
    def entire_space(self):
        assert all(i == self.eri.nmo[0] for i in self.eri.nmo)
        return [numpy.arange(self.nk), numpy.arange(self.eri.nmo[0])]


kernel = gw_slow.kernel


class GW(gw_slow.GW):
    base_imds = IMDS


# pyscf/pbc/gw/krgw_ac.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Tianyu Zhu <zhutianyu1991@gmail.com>
#

'''
PBC spin-restricted G0W0-AC QP eigenvalues with k-point sampling
This implementation has N^4 scaling, and is faster than GW-CD (N^4)
and analytic GW (N^6) methods.
GW-AC is recommended for valence states only, and is inaccuarate for core states.

Method:
    See T. Zhu and G.K.-L. Chan, arxiv:2007.03148 (2020) for details
    Compute Sigma on imaginary frequency with density fitting,
    then analytically continued to real frequency.
    Gaussian density fitting must be used (FFTDF and MDF are not supported).
'''

from functools import reduce
import numpy
import numpy as np
import h5py
from scipy.optimize import newton, least_squares

from pyscf import lib
from pyscf.lib import logger
from pyscf.ao2mo import _ao2mo
from pyscf.ao2mo.incore import _conc_mos
from pyscf.pbc import df, dft, scf
from pyscf.pbc.mp.kmp2 import get_nocc, get_nmo, get_frozen_mask
from pyscf import __config__

einsum = lib.einsum

def kernel(gw, mo_energy, mo_coeff, orbs=None,
           kptlist=None, nw=None, verbose=logger.NOTE):
    '''GW-corrected quasiparticle orbital energies

    Returns:
        A list :  converged, mo_energy, mo_coeff
    '''
    mf = gw._scf
    if gw.frozen is None:
        frozen = 0
    else:
        frozen = gw.frozen
    assert (frozen == 0)

    if orbs is None:
        orbs = range(gw.nmo)
    if kptlist is None:
        kptlist = range(gw.nkpts)
    nkpts = gw.nkpts
    nklist = len(kptlist)

    # v_xc
    dm = np.array(mf.make_rdm1())
    v_mf = np.array(mf.get_veff()) - np.array(mf.get_j(dm_kpts=dm))
    for k in range(nkpts):
        v_mf[k] = reduce(numpy.dot, (mo_coeff[k].T.conj(), v_mf[k], mo_coeff[k]))

    nocc = gw.nocc
    nmo = gw.nmo

    # v_hf from DFT/HF density
    if gw.fc:
        exxdiv = 'ewald'
    else:
        exxdiv = None
    rhf = scf.KRHF(gw.mol, gw.kpts, exxdiv=exxdiv)
    rhf.with_df = gw.with_df
    if getattr(gw.with_df, '_cderi', None) is None:
        raise RuntimeError('Found incompatible integral scheme %s.'
                           'KGWAC can be only used with GDF integrals' %
                           gw.with_df.__class__)

    vk = rhf.get_veff(gw.mol,dm_kpts=dm) - rhf.get_j(gw.mol,dm_kpts=dm)
    for k in range(nkpts):
        vk[k] = reduce(numpy.dot, (mo_coeff[k].T.conj(), vk[k], mo_coeff[k]))

    # Grids for integration on imaginary axis
    freqs,wts = _get_scaled_legendre_roots(nw)

    # Compute self-energy on imaginary axis i*[0,iw_cutoff]
    sigmaI, omega = get_sigma_diag(gw, orbs, kptlist, freqs, wts, iw_cutoff=5.)

    # Analytic continuation
    coeff = []
    if gw.ac == 'twopole':
        for k in range(nklist):
            coeff.append(AC_twopole_diag(sigmaI[k], omega, orbs, nocc))
    elif gw.ac == 'pade':
        for k in range(nklist):
            coeff_tmp, omega_fit = AC_pade_thiele_diag(sigmaI[k], omega)
            coeff.append(coeff_tmp)
    coeff = np.array(coeff)

    conv = True
    # This code does not support metals
    homo = -99.
    lumo = 99.
    for k in range(nkpts):
        if homo < mf.mo_energy[k][nocc-1]:
            homo = mf.mo_energy[k][nocc-1]
        if lumo > mf.mo_energy[k][nocc]:
            lumo = mf.mo_energy[k][nocc]
    ef = (homo+lumo)/2.

    mo_energy = np.zeros_like(np.array(mf.mo_energy))
    for k in range(nklist):
        kn = kptlist[k]
        for p in orbs:
            if gw.linearized:
                # linearized G0W0
                de = 1e-6
                ep = mf.mo_energy[kn][p]
                #TODO: analytic sigma derivative
                if gw.ac == 'twopole':
                    sigmaR = two_pole(ep-ef, coeff[k,:,p-orbs[0]]).real
                    dsigma = two_pole(ep-ef+de, coeff[k,:,p-orbs[0]]).real - sigmaR.real
                elif gw.ac == 'pade':
                    sigmaR = pade_thiele(ep-ef, omega_fit[p-orbs[0]], coeff[k,:,p-orbs[0]]).real
                    dsigma = pade_thiele(ep-ef+de, omega_fit[p-orbs[0]], coeff[k,:,p-orbs[0]]).real - sigmaR.real
                zn = 1.0/(1.0-dsigma/de)
                e = ep + zn*(sigmaR.real + vk[kn,p,p].real - v_mf[kn,p,p].real)
                mo_energy[kn,p] = e
            else:
                # self-consistently solve QP equation
                def quasiparticle(omega):
                    if gw.ac == 'twopole':
                        sigmaR = two_pole(omega-ef, coeff[k,:,p-orbs[0]]).real
                    elif gw.ac == 'pade':
                        sigmaR = pade_thiele(omega-ef, omega_fit[p-orbs[0]], coeff[k,:,p-orbs[0]]).real
                    return omega - mf.mo_energy[kn][p] - (sigmaR.real + vk[kn,p,p].real - v_mf[kn,p,p].real)
                try:
                    e = newton(quasiparticle, mf.mo_energy[kn][p], tol=1e-6, maxiter=100)
                    mo_energy[kn,p] = e
                except RuntimeError:
                    conv = False
    mo_coeff = mf.mo_coeff

    if gw.verbose >= logger.DEBUG:
        numpy.set_printoptions(threshold=nmo)
        for k in range(nkpts):
            logger.debug(gw, '  GW mo_energy @ k%d =\n%s', k,mo_energy[k])
        numpy.set_printoptions(threshold=1000)

    return conv, mo_energy, mo_coeff

def get_rho_response(gw, omega, mo_energy, Lpq, kL, kidx):
    '''
    Compute density response function in auxiliary basis at freq iw
    '''
    nkpts, naux, nmo, nmo = Lpq.shape
    nocc = gw.nocc
    kpts = gw.kpts
    kscaled = gw.mol.get_scaled_kpts(kpts)
    kscaled -= kscaled[0]

    # Compute Pi for kL
    Pi = np.zeros((naux,naux),dtype=np.complex128)
    for i, kpti in enumerate(kpts):
        # Find ka that conserves with ki and kL (-ki+ka+kL=G)
        a = kidx[i]
        eia = mo_energy[i,:nocc,None] - mo_energy[a,None,nocc:]
        eia = eia/(omega**2+eia*eia)
        Pia = einsum('Pia,ia->Pia',Lpq[i][:,:nocc,nocc:],eia)
        # Response from both spin-up and spin-down density
        Pi += 4./nkpts * einsum('Pia,Qia->PQ',Pia,Lpq[i][:,:nocc,nocc:].conj())
    return Pi

def get_sigma_diag(gw, orbs, kptlist, freqs, wts, iw_cutoff=None, max_memory=8000):
    '''
    Compute GW correlation self-energy (diagonal elements)
    in MO basis on imaginary axis
    '''
    mo_energy = np.array(gw._scf.mo_energy)
    mo_coeff = np.array(gw._scf.mo_coeff)
    nocc = gw.nocc
    nmo = gw.nmo
    nkpts = gw.nkpts
    kpts = gw.kpts
    nklist = len(kptlist)
    nw = len(freqs)
    norbs = len(orbs)
    mydf = gw.with_df

    # possible kpts shift center
    kscaled = gw.mol.get_scaled_kpts(kpts)
    kscaled -= kscaled[0]

    # This code does not support metals
    homo = -99.
    lumo = 99.
    for k in range(nkpts):
        if homo < mo_energy[k][nocc-1]:
            homo = mo_energy[k][nocc-1]
        if lumo > mo_energy[k][nocc]:
            lumo = mo_energy[k][nocc]
    if (lumo-homo)<1e-3:
        logger.warn(gw, 'This GW-AC code is not supporting metals!')
    ef = (homo+lumo)/2.

    # Integration on numerical grids
    if iw_cutoff is not None:
        nw_sigma = sum(iw < iw_cutoff for iw in freqs) + 1
    else:
        nw_sigma = nw + 1

    # Compute occ for -iw and vir for iw separately
    # to avoid branch cuts in analytic continuation
    omega_occ = np.zeros((nw_sigma), dtype=np.complex128)
    omega_vir = np.zeros((nw_sigma), dtype=np.complex128)
    omega_occ[1:] = -1j*freqs[:(nw_sigma-1)]
    omega_vir[1:] = 1j*freqs[:(nw_sigma-1)]
    orbs_occ = [i for i in orbs if i < nocc]
    norbs_occ = len(orbs_occ)

    emo_occ = np.zeros((nkpts,nmo,nw_sigma),dtype=np.complex128)
    emo_vir = np.zeros((nkpts,nmo,nw_sigma),dtype=np.complex128)
    for k in range(nkpts):
        emo_occ[k] = omega_occ[None,:] + ef - mo_energy[k][:,None]
        emo_vir[k] = omega_vir[None,:] + ef - mo_energy[k][:,None]

    sigma = np.zeros((nklist,norbs,nw_sigma),dtype=np.complex128)
    omega = np.zeros((norbs,nw_sigma),dtype=np.complex128)
    for p in range(norbs):
        orbp = orbs[p]
        if orbp < nocc:
            omega[p] = omega_occ.copy()
        else:
            omega[p] = omega_vir.copy()

    if gw.fc:
        # Set up q mesh for q->0 finite size correction
        q_pts = np.array([1e-3,0,0]).reshape(1,3)
        q_abs = gw.mol.get_abs_kpts(q_pts)

        # Get qij = 1/sqrt(Omega) * < psi_{ik} | e^{iqr} | psi_{ak-q} > at q: (nkpts, nocc, nvir)
        qij = get_qij(gw, q_abs[0], mo_coeff)

    for kL in range(nkpts):
        # Lij: (ki, L, i, j) for looping every kL
        Lij = []
        # kidx: save kj that conserves with kL and ki (-ki+kj+kL=G)
        # kidx_r: save ki that conserves with kL and kj (-ki+kj+kL=G)
        kidx = np.zeros((nkpts),dtype=np.int64)
        kidx_r = np.zeros((nkpts),dtype=np.int64)
        for i, kpti in enumerate(kpts):
            for j, kptj in enumerate(kpts):
                # Find (ki,kj) that satisfies momentum conservation with kL
                kconserv = -kscaled[i] + kscaled[j] + kscaled[kL]
                is_kconserv = np.linalg.norm(np.round(kconserv) - kconserv) < 1e-12
                if is_kconserv:
                    kidx[i] = j
                    kidx_r[j] = i
                    logger.debug(gw, "Read Lpq (kL: %s / %s, ki: %s, kj: %s)"%(kL+1, nkpts, i, j))
                    Lij_out = None
                    # Read (L|pq) and ao2mo transform to (L|ij)
                    Lpq = []
                    for LpqR, LpqI, sign \
                            in mydf.sr_loop([kpti, kptj], max_memory=0.1*gw._scf.max_memory, compact=False):
                        Lpq.append(LpqR+LpqI*1.0j)
                    # support uneqaul naux on different k points
                    Lpq = np.vstack(Lpq).reshape(-1,nmo**2)
                    tao = []
                    ao_loc = None
                    moij, ijslice = _conc_mos(mo_coeff[i], mo_coeff[j])[2:]
                    Lij_out = _ao2mo.r_e2(Lpq, moij, ijslice, tao, ao_loc, out=Lij_out)
                    Lij.append(Lij_out.reshape(-1,nmo,nmo))
        Lij = np.asarray(Lij)
        naux = Lij.shape[1]

        if kL == 0:
            for w in range(nw):
                # body dielectric matrix eps_body
                Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx)
                eps_body_inv = np.linalg.inv(np.eye(naux)-Pi)

                if gw.fc:
                    # head dielectric matrix eps_00
                    Pi_00 = get_rho_response_head(gw, freqs[w], mo_energy, qij)
                    eps_00 = 1. - 4. * np.pi/np.linalg.norm(q_abs[0])**2 * Pi_00

                    # wings dielectric matrix eps_P0
                    Pi_P0 = get_rho_response_wing(gw, freqs[w], mo_energy, Lij, qij)
                    eps_P0 = -np.sqrt(4.*np.pi) / np.linalg.norm(q_abs[0]) * Pi_P0

                    # inverse dielectric matrix
                    eps_inv_00 = 1./(eps_00 - np.dot(np.dot(eps_P0.conj(),eps_body_inv),eps_P0))
                    eps_inv_P0 = -eps_inv_00 * np.dot(eps_body_inv, eps_P0)

                    # head correction
                    Del_00 = 2./np.pi * (6.*np.pi**2/gw.mol.vol/nkpts)**(1./3.) * (eps_inv_00 - 1.)

                eps_inv_PQ = eps_body_inv
                g0_occ = wts[w] * emo_occ / (emo_occ**2+freqs[w]**2)
                g0_vir = wts[w] * emo_vir / (emo_vir**2+freqs[w]**2)

                for k in range(nklist):
                    kn = kptlist[k]
                    # Find km that conserves with kn and kL (-km+kn+kL=G)
                    km = kidx_r[kn]
                    Qmn = einsum('Pmn,PQ->Qmn',Lij[km][:,:,orbs].conj(),eps_inv_PQ-np.eye(naux))
                    Wmn = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn,Lij[km][:,:,orbs])
                    sigma[k][:norbs_occ] += -einsum('mn,mw->nw',Wmn[:,:norbs_occ],g0_occ[km])/np.pi
                    sigma[k][norbs_occ:] += -einsum('mn,mw->nw',Wmn[:,norbs_occ:],g0_vir[km])/np.pi

                    if gw.fc:
                        # apply head correction
                        assert(kn == km)
                        sigma[k][:norbs_occ] += -Del_00 * g0_occ[kn][orbs][:norbs_occ] /np.pi
                        sigma[k][norbs_occ:] += -Del_00 * g0_vir[kn][orbs][norbs_occ:] /np.pi

                        # apply wing correction
                        Wn_P0 = einsum('Pnm,P->nm',Lij[kn],eps_inv_P0).diagonal()
                        Wn_P0 = Wn_P0.real * 2.
                        Del_P0 = np.sqrt(gw.mol.vol/4./np.pi**3) * (6.*np.pi**2/gw.mol.vol/nkpts)**(2./3.) * Wn_P0[orbs]
                        sigma[k][:norbs_occ] += -einsum('n,nw->nw', Del_P0[:norbs_occ],
                                                        g0_occ[kn][orbs][:norbs_occ]) /np.pi
                        sigma[k][norbs_occ:] += -einsum('n,nw->nw', Del_P0[norbs_occ:],
                                                        g0_vir[kn][orbs][norbs_occ:]) /np.pi
        else:
            for w in range(nw):
                Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx)
                Pi_inv = np.linalg.inv(np.eye(naux)-Pi)-np.eye(naux)
                g0_occ = wts[w] * emo_occ / (emo_occ**2+freqs[w]**2)
                g0_vir = wts[w] * emo_vir / (emo_vir**2+freqs[w]**2)
                for k in range(nklist):
                    kn = kptlist[k]
                    # Find km that conserves with kn and kL (-km+kn+kL=G)
                    km = kidx_r[kn]
                    Qmn = einsum('Pmn,PQ->Qmn',Lij[km][:,:,orbs].conj(),Pi_inv)
                    Wmn = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn,Lij[km][:,:,orbs])
                    sigma[k][:norbs_occ] += -einsum('mn,mw->nw',Wmn[:,:norbs_occ],g0_occ[km])/np.pi
                    sigma[k][norbs_occ:] += -einsum('mn,mw->nw',Wmn[:,norbs_occ:],g0_vir[km])/np.pi

    return sigma, omega

def get_rho_response_head(gw, omega, mo_energy, qij):
    '''
    Compute head (G=0, G'=0) density response function in auxiliary basis at freq iw
    '''
    nkpts, nocc, nvir = qij.shape
    nocc = gw.nocc
    kpts = gw.kpts

    # Compute Pi head
    Pi_00 = 0j
    for i, kpti in enumerate(kpts):
        eia = mo_energy[i,:nocc,None] - mo_energy[i,None,nocc:]
        eia = eia/(omega**2+eia*eia)
        Pi_00 += 4./nkpts * einsum('ia,ia->',eia,qij[i].conj()*qij[i])
    return Pi_00

def get_rho_response_wing(gw, omega, mo_energy, Lpq, qij):
    '''
    Compute wing (G=P, G'=0) density response function in auxiliary basis at freq iw
    '''
    nkpts, naux, nmo, nmo = Lpq.shape
    nocc = gw.nocc
    kpts = gw.kpts

    # Compute Pi wing
    Pi = np.zeros(naux,dtype=np.complex128)
    for i, kpti in enumerate(kpts):
        eia = mo_energy[i,:nocc,None] - mo_energy[i,None,nocc:]
        eia = eia/(omega**2+eia*eia)
        eia_q = eia * qij[i].conj()
        Pi += 4./nkpts * einsum('Pia,ia->P',Lpq[i][:,:nocc,nocc:],eia_q)
    return Pi

def get_qij(gw, q, mo_coeff, uniform_grids=False):
    '''
    Compute qij = 1/Omega * |< psi_{ik} | e^{iqr} | psi_{ak-q} >|^2 at q: (nkpts, nocc, nvir)
    through kp perturbtation theory
    Ref: Phys. Rev. B 83, 245122 (2011)
    '''
    nocc = gw.nocc
    nmo = gw.nmo
    nvir = nmo - nocc
    kpts = gw.kpts
    nkpts = len(kpts)
    cell = gw.mol
    mo_energy = gw._scf.mo_energy

    if uniform_grids:
        mydf = df.FFTDF(cell, kpts=kpts)
        coords = cell.gen_uniform_grids(mydf.mesh)
    else:
        coords, weights = dft.gen_grid.get_becke_grids(cell,level=5)
    ngrid = len(coords)

    qij = np.zeros((nkpts,nocc,nvir),dtype=np.complex128)
    for i, kpti in enumerate(kpts):
        ao_p = dft.numint.eval_ao(cell, coords, kpt=kpti, deriv=1)
        ao = ao_p[0]
        ao_grad = ao_p[1:4]
        if uniform_grids:
            ao_ao_grad = einsum('mg,xgn->xmn',ao.T.conj(),ao_grad) * cell.vol / ngrid
        else:
            ao_ao_grad = einsum('g,mg,xgn->xmn',weights,ao.T.conj(),ao_grad)
        q_ao_ao_grad = -1j * einsum('x,xmn->mn',q,ao_ao_grad)
        q_mo_mo_grad = np.dot(np.dot(mo_coeff[i][:,:nocc].T.conj(), q_ao_ao_grad), mo_coeff[i][:,nocc:])
        enm = 1./(mo_energy[i][nocc:,None] - mo_energy[i][None,:nocc])
        dens = enm.T * q_mo_mo_grad
        qij[i] = dens / np.sqrt(cell.vol)

    return qij

def _get_scaled_legendre_roots(nw):
    """
    Scale nw Legendre roots, which lie in the
    interval [-1, 1], so that they lie in [0, inf)
    Ref: www.cond-mat.de/events/correl19/manuscripts/ren.pdf

    Returns:
        freqs : 1D ndarray
        wts : 1D ndarray
    """
    freqs, wts = np.polynomial.legendre.leggauss(nw)
    x0 = 0.5
    freqs_new = x0*(1.+freqs)/(1.-freqs)
    wts = wts*2.*x0/(1.-freqs)**2
    return freqs_new, wts

def _get_clenshaw_curtis_roots(nw):
    """
    Clenshaw-Curtis qaudrature on [0,inf)
    Ref: J. Chem. Phys. 132, 234114 (2010)
    Returns:
        freqs : 1D ndarray
        wts : 1D ndarray
    """
    freqs = np.zeros(nw)
    wts = np.zeros(nw)
    a = 0.2
    for w in range(nw):
        t = (w+1.0)/nw * np.pi/2.
        freqs[w] = a / np.tan(t)
        if w != nw-1:
            wts[w] = a*np.pi/2./nw/(np.sin(t)**2)
        else:
            wts[w] = a*np.pi/4./nw/(np.sin(t)**2)
    return freqs[::-1], wts[::-1]

def two_pole_fit(coeff, omega, sigma):
    cf = coeff[:5] + 1j*coeff[5:]
    f = cf[0] + cf[1]/(omega+cf[3]) + cf[2]/(omega+cf[4]) - sigma
    f[0] = f[0]/0.01
    return np.array([f.real,f.imag]).reshape(-1)

def two_pole(freqs, coeff):
    cf = coeff[:5] + 1j*coeff[5:]
    return cf[0] + cf[1]/(freqs+cf[3]) + cf[2]/(freqs+cf[4])

def AC_twopole_diag(sigma, omega, orbs, nocc):
    """
    Analytic continuation to real axis using a two-pole model
    Returns:
        coeff: 2D array (ncoeff, norbs)
    """
    norbs, nw = sigma.shape
    coeff = np.zeros((10,norbs))
    for p in range(norbs):
        if orbs[p] < nocc:
            x0 = np.array([0, 1, 1, 1, -1, 0, 0, 0, -1.0, -0.5])
        else:
            x0 = np.array([0, 1, 1, 1, -1, 0, 0, 0, 1.0, 0.5])
        #TODO: analytic gradient
        xopt = least_squares(two_pole_fit, x0, jac='3-point', method='trf', xtol=1e-10,
                             gtol = 1e-10, max_nfev=1000, verbose=0, args=(omega[p], sigma[p]))
        if xopt.success is False:
            print('WARN: 2P-Fit Orb %d not converged, cost function %e'%(p,xopt.cost))
        coeff[:,p] = xopt.x.copy()
    return coeff

def thiele(fn,zn):
    nfit = len(zn)
    g = np.zeros((nfit,nfit),dtype=np.complex128)
    g[:,0] = fn.copy()
    for i in range(1,nfit):
        g[i:,i] = (g[i-1,i-1]-g[i:,i-1])/((zn[i:]-zn[i-1])*g[i:,i-1])
    a = g.diagonal()
    return a

def pade_thiele(freqs,zn,coeff):
    nfit = len(coeff)
    X = coeff[-1]*(freqs-zn[-2])
    for i in range(nfit-1):
        idx = nfit-i-1
        X = coeff[idx]*(freqs-zn[idx-1])/(1.+X)
    X = coeff[0]/(1.+X)
    return X

def AC_pade_thiele_diag(sigma, omega):
    """
    Analytic continuation to real axis using a Pade approximation
    from Thiele's reciprocal difference method
    Reference: J. Low Temp. Phys. 29, 179 (1977)
    Returns:
        coeff: 2D array (ncoeff, norbs)
        omega: 2D array (norbs, npade)
    """
    idx = range(1,40,6)
    sigma1 = sigma[:,idx].copy()
    sigma2 = sigma[:,(idx[-1]+4)::4].copy()
    sigma = np.hstack((sigma1,sigma2))
    omega1 = omega[:,idx].copy()
    omega2 = omega[:,(idx[-1]+4)::4].copy()
    omega = np.hstack((omega1,omega2))
    norbs, nw = sigma.shape
    npade = nw // 2
    coeff = np.zeros((npade*2,norbs),dtype=np.complex128)
    for p in range(norbs):
        coeff[:,p] = thiele(sigma[p,:npade*2], omega[p,:npade*2])

    return coeff, omega[:,:npade*2]

class KRGWAC(lib.StreamObject):

    linearized = getattr(__config__, 'gw_gw_GW_linearized', False)
    # Analytic continuation: pade or twopole
    ac = getattr(__config__, 'gw_gw_GW_ac', 'pade')
    # Whether applying finite size corrections
    fc = getattr(__config__, 'gw_gw_GW_fc', True)

    def __init__(self, mf, frozen=0):
        self.mol = mf.mol
        self._scf = mf
        self.verbose = self.mol.verbose
        self.stdout = self.mol.stdout
        self.max_memory = mf.max_memory

        #TODO: implement frozen orbs
        if frozen > 0:
            raise NotImplementedError
        self.frozen = frozen

        # DF-KGW must use GDF integrals
        if getattr(mf, 'with_df', None):
            self.with_df = mf.with_df
        else:
            raise NotImplementedError
        self._keys.update(['with_df'])

##################################################
# don't modify the following attributes, they are not input options
        self._nocc = None
        self._nmo = None
        self.kpts = mf.kpts
        self.nkpts = len(self.kpts)
        # self.mo_energy: GW quasiparticle energy, not scf mo_energy
        self.mo_energy = None
        self.mo_coeff = mf.mo_coeff
        self.mo_occ = mf.mo_occ
        self.sigma = None

        keys = set(('linearized','ac','fc'))
        self._keys = set(self.__dict__.keys()).union(keys)

    def dump_flags(self):
        log = logger.Logger(self.stdout, self.verbose)
        log.info('')
        log.info('******** %s ********', self.__class__)
        log.info('method = %s', self.__class__.__name__)
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        log.info('GW nocc = %d, nvir = %d, nkpts = %d', nocc, nvir, nkpts)
        if self.frozen is not None:
            log.info('frozen orbitals %s', str(self.frozen))
        logger.info(self, 'use perturbative linearized QP eqn = %s', self.linearized)
        logger.info(self, 'analytic continuation method = %s', self.ac)
        logger.info(self, 'GW finite size corrections = %s', self.fc)
        return self

    @property
    def nocc(self):
        return self.get_nocc()
    @nocc.setter
    def nocc(self, n):
        self._nocc = n

    @property
    def nmo(self):
        return self.get_nmo()
    @nmo.setter
    def nmo(self, n):
        self._nmo = n

    get_nocc = get_nocc
    get_nmo = get_nmo
    get_frozen_mask = get_frozen_mask

    def kernel(self, mo_energy=None, mo_coeff=None, orbs=None, kptlist=None, nw=100):
        """
        Input:
            kptlist: self-energy k-points
            orbs: self-energy orbs
            nw: grid number
        Output:
            mo_energy: GW quasiparticle energy
        """
        if mo_coeff is None:
            mo_coeff = np.array(self._scf.mo_coeff)
        if mo_energy is None:
            mo_energy = np.array(self._scf.mo_energy)

        nmo = self.nmo
        naux = self.with_df.get_naoaux()
        nkpts = self.nkpts
        mem_incore = (2*nkpts*nmo**2*naux) * 16/1e6
        mem_now = lib.current_memory()[0]
        if (mem_incore + mem_now > 0.99*self.max_memory):
            logger.warn(self, 'Memory may not be enough!')
            raise NotImplementedError

        cput0 = (logger.process_clock(), logger.perf_counter())
        self.dump_flags()
        self.converged, self.mo_energy, self.mo_coeff = \
                kernel(self, mo_energy, mo_coeff, orbs=orbs,
                       kptlist=kptlist, nw=nw, verbose=self.verbose)

        logger.warn(self, 'GW QP energies may not be sorted from min to max')
        logger.timer(self, 'GW', *cput0)
        return self.mo_energy

if __name__ == '__main__':
    from pyscf.pbc import gto
    from pyscf.pbc.lib import chkfile
    import os
    # This test takes a few minutes
    cell = gto.Cell()
    cell.build(unit = 'angstrom',
               a = '''
               0.000000     1.783500     1.783500
               1.783500     0.000000     1.783500
               1.783500     1.783500     0.000000
               ''',
               atom = 'C 1.337625 1.337625 1.337625; C 2.229375 2.229375 2.229375',
               dimension = 3,
               max_memory = 8000,
               verbose = 4,
               pseudo = 'gth-pade',
               basis='gth-szv',
               precision=1e-10)

    kpts = cell.make_kpts([3,1,1],scaled_center=[0,0,0])
    gdf = df.GDF(cell, kpts)
    gdf_fname = 'gdf_ints_311.h5'
    gdf._cderi_to_save = gdf_fname
    if not os.path.isfile(gdf_fname):
        gdf.build()

    chkfname = 'diamond_311.chk'
    if os.path.isfile(chkfname):
        kmf = dft.KRKS(cell, kpts)
        kmf.xc = 'pbe'
        kmf.with_df = gdf
        kmf.with_df._cderi = gdf_fname
        data = chkfile.load(chkfname, 'scf')
        kmf.__dict__.update(data)
    else:
        kmf = dft.KRKS(cell, kpts)
        kmf.xc = 'pbe'
        kmf.with_df = gdf
        kmf.with_df._cderi = gdf_fname
        kmf.conv_tol = 1e-12
        kmf.chkfile = chkfname
        kmf.kernel()

    gw = KRGWAC(kmf)
    gw.linearized = False
    gw.ac = 'pade'
    # without finite size corrections
    gw.fc = False
    nocc = gw.nocc
    gw.kernel(kptlist=[0,1,2],orbs=range(0,nocc+3))
    print(gw.mo_energy)
    assert((abs(gw.mo_energy[0][nocc-1]-0.62045797))<1e-5)
    assert((abs(gw.mo_energy[0][nocc]-0.96574324))<1e-5)
    assert((abs(gw.mo_energy[1][nocc-1]-0.52639137))<1e-5)
    assert((abs(gw.mo_energy[1][nocc]-1.07513258))<1e-5)

    # with finite size corrections
    gw.fc = True
    gw.kernel(kptlist=[0,1,2],orbs=range(0,nocc+3))
    print(gw.mo_energy)
    assert((abs(gw.mo_energy[0][nocc-1]-0.54277092))<1e-5)
    assert((abs(gw.mo_energy[0][nocc]-0.80148537))<1e-5)
    assert((abs(gw.mo_energy[1][nocc-1]-0.45073793))<1e-5)
    assert((abs(gw.mo_energy[1][nocc]-0.92910108))<1e-5)


# pyscf/pbc/gw/krgw_cd.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Tianyu Zhu <zhutianyu1991@gmail.com>
#

'''
PBC spin-restricted G0W0-CD QP eigenvalues with k-point sampling
This implementation has the same scaling (N^4) as GW-AC, more robust but slower.
GW-CD is particularly recommended for accurate core and high-energy states.

Method:
    See T. Zhu and G.K.-L. Chan, arxiv:2007.03148 (2020) for details
    Compute Sigma directly on real axis with density fitting
    through a contour deformation method
'''

from functools import reduce
import numpy
import numpy as np
import h5py
from scipy.optimize import newton, least_squares

from pyscf import lib
from pyscf.lib import logger
from pyscf.ao2mo import _ao2mo
from pyscf.ao2mo.incore import _conc_mos
from pyscf.pbc import df, dft, scf
from pyscf.pbc.mp.kmp2 import get_nocc, get_nmo, get_frozen_mask
from pyscf import __config__

einsum = lib.einsum

def kernel(gw, mo_energy, mo_coeff, orbs=None,
           kptlist=None, nw=None, verbose=logger.NOTE):
    '''GW-corrected quasiparticle orbital energies

    Returns:
        A list :  converged, mo_energy, mo_coeff
    '''
    mf = gw._scf
    if gw.frozen is None:
        frozen = 0
    else:
        frozen = gw.frozen
    assert (frozen == 0)

    if orbs is None:
        orbs = range(gw.nmo)
    if kptlist is None:
        kptlist = range(gw.nkpts)
    nkpts = gw.nkpts
    nklist = len(kptlist)

    # v_xc
    dm = np.array(mf.make_rdm1())
    v_mf = np.array(mf.get_veff()) - np.array(mf.get_j(dm_kpts=dm))
    for k in range(nkpts):
        v_mf[k] = reduce(numpy.dot, (mo_coeff[k].T.conj(), v_mf[k], mo_coeff[k]))

    nocc = gw.nocc
    nmo = gw.nmo

    # v_hf from DFT/HF density
    if gw.fc:
        exxdiv = 'ewald'
    else:
        exxdiv = None
    rhf = scf.KRHF(gw.mol, gw.kpts, exxdiv=exxdiv)
    rhf.with_df = gw.with_df
    if getattr(gw.with_df, '_cderi', None) is None:
        raise RuntimeError('Found incompatible integral scheme %s.'
                           'KGWCD can be only used with GDF integrals' %
                           gw.with_df.__class__)

    vk = rhf.get_veff(gw.mol,dm_kpts=dm) - rhf.get_j(gw.mol,dm_kpts=dm)
    for k in range(nkpts):
        vk[k] = reduce(numpy.dot, (mo_coeff[k].T.conj(), vk[k], mo_coeff[k]))

    # Grids for integration on imaginary axis
    freqs,wts = _get_scaled_legendre_roots(nw)

    logger.debug(gw, "Computing the imaginary part")
    Wmn, Del_00, Del_P0, qij, q_abs = get_WmnI_diag(gw, orbs, kptlist, freqs)

    conv = True
    mo_energy = np.zeros_like(np.array(mf.mo_energy))
    for k in range(nklist):
        kn = kptlist[k]
        for p in orbs:
            if p < nocc:
                delta = -2e-2
            else:
                delta = 2e-2
            if gw.linearized:
                # FIXME
                logger.warn(gw,'linearization with CD leads to wrong quasiparticle energy')
                raise NotImplementedError
            else:
                # self-consistently solve QP equation
                def quasiparticle(omega):
                    if gw.fc:
                        sigmaR = get_sigma_diag(gw, omega, kn, p, Wmn[:,k,:,p-orbs[0],:],
                                                Del_00, Del_P0[k,p-orbs[0],:], freqs, wts, qij, q_abs).real
                    else:
                        sigmaR = get_sigma_diag(gw, omega, kn, p, Wmn[:,k,:,p-orbs[0],:],
                                                Del_00, Del_P0, freqs, wts, qij, q_abs).real
                    return omega - mf.mo_energy[kn][p] - (sigmaR.real + vk[kn,p,p].real - v_mf[kn,p,p].real)
                try:
                    e = newton(quasiparticle, mf.mo_energy[kn][p]+delta, tol=1e-6, maxiter=50)
                    logger.debug(gw, "Computing poles for QP (k: %s, orb: %s)"%(kn,p))
                    mo_energy[kn,p] = e
                except RuntimeError:
                    conv = False
    mo_coeff = mf.mo_coeff

    if gw.verbose >= logger.DEBUG:
        numpy.set_printoptions(threshold=nmo)
        for k in range(nkpts):
            logger.debug(gw, '  GW mo_energy @ k%d =\n%s', k,mo_energy[k])
        numpy.set_printoptions(threshold=1000)

    return conv, mo_energy, mo_coeff

def get_sigma_diag(gw, ep, kp, p, Wmn, Del_00, Del_P0, freqs, wts, qij, q_abs):
    '''
    Compute self-energy on real axis using contour deformation
    '''
    nocc = gw.nocc
    nkpts = gw.nkpts
    # This code does not support metals
    homo = -99.
    lumo = 99.
    for k in range(nkpts):
        if homo < gw._scf.mo_energy[k][nocc-1]:
            homo = gw._scf.mo_energy[k][nocc-1]
        if lumo > gw._scf.mo_energy[k][nocc]:
            lumo = gw._scf.mo_energy[k][nocc]
    ef = (homo+lumo)/2.

    nmo = gw.nmo
    sign = np.zeros((nkpts,nmo),dtype=np.int64)
    for k in range(nkpts):
        sign[k] = np.sign(ef-gw._scf.mo_energy[k])
    sigmaI = get_sigmaI_diag(gw, ep, kp, p, Wmn, Del_00, Del_P0, sign, freqs, wts)
    sigmaR = get_sigmaR_diag(gw, ep, kp, p, ef, freqs, qij, q_abs)
    return sigmaI + sigmaR

def get_rho_response(gw, omega, mo_energy, Lpq, kL, kidx):
    '''
    Compute density response function in auxiliary basis at freq iw
    '''
    nkpts, naux, nmo, nmo = Lpq.shape
    nocc = gw.nocc
    kpts = gw.kpts
    kscaled = gw.mol.get_scaled_kpts(kpts)
    kscaled -= kscaled[0]

    # Compute Pi for kL
    Pi = np.zeros((naux,naux),dtype=np.complex128)
    for i, kpti in enumerate(kpts):
        # Find ka that conserves with ki and kL (-ki+ka+kL=G)
        a = kidx[i]
        eia = mo_energy[i,:nocc,None] - mo_energy[a,None,nocc:]
        eia = eia/(omega**2+eia*eia)
        Pia = einsum('Pia,ia->Pia',Lpq[i][:,:nocc,nocc:],eia)
        # Response from both spin-up and spin-down density
        Pi += 4./nkpts * einsum('Pia,Qia->PQ',Pia,Lpq[i][:,:nocc,nocc:].conj())
    return Pi

def get_WmnI_diag(gw, orbs, kptlist, freqs, max_memory=8000):
    '''
    Compute GW correlation self-energy (diagonal elements)
    in MO basis on imaginary axis
    '''
    mo_energy = np.array(gw._scf.mo_energy)
    mo_coeff = np.array(gw._scf.mo_coeff)
    nmo = gw.nmo
    nkpts = gw.nkpts
    kpts = gw.kpts
    nklist = len(kptlist)
    nw = len(freqs)
    norbs = len(orbs)
    mydf = gw.with_df

    # possible kpts shift center
    kscaled = gw.mol.get_scaled_kpts(kpts)
    kscaled -= kscaled[0]

    Del_00, Del_P0, qij, q_abs = None, None, None, None
    if gw.fc:
        # Set up q mesh for q->0 finite size correction
        q_pts = np.array([1e-3,0,0]).reshape(1,3)
        q_abs = gw.mol.get_abs_kpts(q_pts)

        # Get qij = 1/sqrt(Omega) * < psi_{ik} | e^{iqr} | psi_{ak-q} > at q: (nkpts, nocc, nvir)
        qij = get_qij(gw, q_abs[0], mo_coeff)

    Wmn = np.zeros((nkpts,nklist,nmo,norbs,nw),dtype=np.complex128)
    if gw.fc:
        Del_P0 = np.zeros((nklist,norbs,nw),dtype=np.complex128)
        Del_00 = np.zeros(nw,dtype=np.complex128)
    for kL in range(nkpts):
        # Lij: (ki, L, i, j) for looping every kL
        Lij = []
        # kidx: save kj that conserves with kL and ki (-ki+kj+kL=G)
        # kidx_r: save ki that conserves with kL and kj (-ki+kj+kL=G)
        kidx = np.zeros((nkpts),dtype=np.int64)
        kidx_r = np.zeros((nkpts),dtype=np.int64)
        for i, kpti in enumerate(kpts):
            for j, kptj in enumerate(kpts):
                # Find (ki,kj) that satisfies momentum conservation with kL
                kconserv = -kscaled[i] + kscaled[j] + kscaled[kL]
                is_kconserv = np.linalg.norm(np.round(kconserv) - kconserv) < 1e-12
                if is_kconserv:
                    kidx[i] = j
                    kidx_r[j] = i
                    logger.debug(gw, "Read Lpq (kL: %s / %s, ki: %s, kj: %s)"%(kL+1, nkpts, i, j))
                    Lij_out = None
                    # Read (L|pq) and ao2mo transform to (L|ij)
                    Lpq = []
                    for LpqR, LpqI, sign \
                            in mydf.sr_loop([kpti, kptj], max_memory=0.1*gw._scf.max_memory, compact=False):
                        Lpq.append(LpqR+LpqI*1.0j)
                    # support uneqaul naux on different k points
                    Lpq = np.vstack(Lpq).reshape(-1,nmo**2)
                    tao = []
                    ao_loc = None
                    moij, ijslice = _conc_mos(mo_coeff[i], mo_coeff[j])[2:]
                    Lij_out = _ao2mo.r_e2(Lpq, moij, ijslice, tao, ao_loc, out=Lij_out)
                    Lij.append(Lij_out.reshape(-1,nmo,nmo))
        Lij = np.asarray(Lij)
        naux = Lij.shape[1]

        if kL == 0:
            for w in range(nw):
                # body dielectric matrix eps_body
                Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx)
                eps_body_inv = np.linalg.inv(np.eye(naux)-Pi)

                if gw.fc:
                    # head dielectric matrix eps_00
                    Pi_00 = get_rho_response_head(gw, freqs[w], mo_energy, qij)
                    eps_00 = 1. - 4. * np.pi/np.linalg.norm(q_abs[0])**2 * Pi_00

                    # wings dielectric matrix eps_P0
                    Pi_P0 = get_rho_response_wing(gw, freqs[w], mo_energy, Lij, qij)
                    eps_P0 = -np.sqrt(4.*np.pi) / np.linalg.norm(q_abs[0]) * Pi_P0

                    # inverse dielectric matrix
                    eps_inv_00 = 1./(eps_00 - np.dot(np.dot(eps_P0.conj(),eps_body_inv),eps_P0))
                    eps_inv_P0 = -eps_inv_00 * np.dot(eps_body_inv, eps_P0)

                    # head correction
                    Del_00[w] = 2./np.pi * (6.*np.pi**2/gw.mol.vol/nkpts)**(1./3.) * (eps_inv_00 - 1.)
                    wings_const = np.sqrt(gw.mol.vol/4./np.pi**3) * (6.*np.pi**2/gw.mol.vol/nkpts)**(2./3.)

                eps_inv_PQ = eps_body_inv

                for k in range(nklist):
                    kn = kptlist[k]
                    # Find km that conserves with kn and kL (-km+kn+kL=G)
                    km = kidx_r[kn]
                    Qmn = einsum('Pmn,PQ->Qmn',Lij[km][:,:,orbs].conj(),eps_inv_PQ-np.eye(naux))
                    Wmn[km,k,:,:,w] = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn,Lij[km][:,:,orbs])

                    if gw.fc:
                        # compute wing correction
                        Wn_P0 = einsum('Pnm,P->nm',Lij[kn],eps_inv_P0).diagonal()
                        Wn_P0 = Wn_P0.real * 2.
                        Del_P0[k,:,w] = wings_const * Wn_P0[orbs]
        else:
            for w in range(nw):
                Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx)
                Pi_inv = np.linalg.inv(np.eye(naux)-Pi)-np.eye(naux)
                for k in range(nklist):
                    kn = kptlist[k]
                    # Find km that conserves with kn and kL (-km+kn+kL=G)
                    km = kidx_r[kn]
                    Qmn = einsum('Pmn,PQ->Qmn',Lij[km][:,:,orbs].conj(),Pi_inv)
                    Wmn[km,k,:,:,w] = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn,Lij[km][:,:,orbs])

    return Wmn, Del_00, Del_P0, qij, q_abs

def get_sigmaI_diag(gw, omega, kp, p, Wmn, Del_00, Del_P0, sign, freqs, wts):
    '''
    Compute self-energy by integrating on imaginary axis
    '''
    mo_energy = gw._scf.mo_energy
    nkpts = gw.nkpts
    sigma = 0j
    for k in range(nkpts):
        emo = omega - 1j*gw.eta*sign[k] - mo_energy[k]
        g0 = wts[None,:]*emo[:,None] / ((emo**2)[:,None]+(freqs**2)[None,:])
        sigma += -einsum('mw,mw',g0,Wmn[k])/np.pi
        if gw.fc and k == kp:
            sigma += -einsum('w,w->',Del_00,g0[p])/np.pi
            sigma += -einsum('w,w->',Del_P0,g0[p])/np.pi

    return sigma

def get_rho_response_R(gw, omega, mo_energy, Lpq, kL, kidx):
    '''
    Compute density response function in auxiliary basis at freq iw
    '''
    nkpts, naux, nmo, nmo = Lpq.shape
    nocc = gw.nocc
    kpts = gw.kpts
    kscaled = gw.mol.get_scaled_kpts(kpts)
    kscaled -= kscaled[0]

    # Compute Pi for kL
    Pi = np.zeros((naux,naux),dtype=np.complex128)
    for i, kpti in enumerate(kpts):
        # Find ka that conserves with ki and kL (-ki+ka+kL=G)
        a = kidx[i]
        eia = mo_energy[i,:nocc,None] - mo_energy[a,None,nocc:]
        eia = 1./(omega+eia+2j*gw.eta) + 1./(-omega+eia)
        Pia = einsum('Pia,ia->Pia',Lpq[i][:,:nocc,nocc:],eia)
        # Response from both spin-up and spin-down density
        Pi += 2./nkpts * einsum('Pia,Qia->PQ',Pia,Lpq[i][:,:nocc,nocc:].conj())
    return Pi

def get_sigmaR_diag(gw, omega, kn, orbp, ef, freqs, qij, q_abs):
    '''
    Compute self-energy for poles inside coutour
    (more and more expensive away from Fermi surface)
    '''
    mo_energy = np.array(gw._scf.mo_energy)
    mo_coeff = np.array(gw._scf.mo_coeff)
    nmo = gw.nmo
    nkpts = gw.nkpts
    kpts = gw.kpts
    mydf = gw.with_df

    # possible kpts shift center
    kscaled = gw.mol.get_scaled_kpts(kpts)
    kscaled -= kscaled[0]

    idx = []
    for k in range(nkpts):
        if omega > ef:
            fm = 1.0
            idx.append(np.where((mo_energy[k]<omega) & (mo_energy[k]>ef))[0])
        else:
            fm = -1.0
            idx.append(np.where((mo_energy[k]>omega) & (mo_energy[k]<ef))[0])

    sigmaR = 0j
    for kL in range(nkpts):
        # Lij: (ki, L, i, j) for looping every kL
        Lij = []
        # kidx: save kj that conserves with kL and ki (-ki+kj+kL=G)
        # kidx_r: save ki that conserves with kL and kj (-ki+kj+kL=G)
        kidx = np.zeros((nkpts),dtype=np.int64)
        kidx_r = np.zeros((nkpts),dtype=np.int64)
        for i, kpti in enumerate(kpts):
            for j, kptj in enumerate(kpts):
                # Find (ki,kj) that satisfies momentum conservation with kL
                kconserv = -kscaled[i] + kscaled[j] + kscaled[kL]
                is_kconserv = np.linalg.norm(np.round(kconserv) - kconserv) < 1e-12
                if is_kconserv:
                    kidx[i] = j
                    kidx_r[j] = i
        km = kidx_r[kn]
        if len(idx[km]) > 0:
            for i, kpti in enumerate(kpts):
                for j, kptj in enumerate(kpts):
                    # Find (ki,kj) that satisfies momentum conservation with kL
                    kconserv = -kscaled[i] + kscaled[j] + kscaled[kL]
                    is_kconserv = np.linalg.norm(np.round(kconserv) - kconserv) < 1e-12
                    if is_kconserv:
                        kidx[i] = j
                        kidx_r[j] = i
                        #logger.debug(gw, "Read Lpq (kL: %s / %s, ki: %s, kj: %s)"%(kL+1, nkpts, i, j))
                        Lij_out = None
                        # Read (L|pq) and ao2mo transform to (L|ij)
                        Lpq = []
                        for LpqR, LpqI, sign \
                                in mydf.sr_loop([kpti, kptj], max_memory=0.1*gw._scf.max_memory, compact=False):
                            Lpq.append(LpqR+LpqI*1.0j)
                        # support uneqaul naux on different k points
                        Lpq = np.vstack(Lpq).reshape(-1,nmo**2)
                        tao = []
                        ao_loc = None
                        moij, ijslice = _conc_mos(mo_coeff[i], mo_coeff[j])[2:]
                        Lij_out = _ao2mo.r_e2(Lpq, moij, ijslice, tao, ao_loc, out=Lij_out)
                        Lij.append(Lij_out.reshape(-1,nmo,nmo))
            Lij = np.asarray(Lij)
            naux = Lij.shape[1]

            if kL == 0:
                km = kidx_r[kn]
                if len(idx[km]) > 0:
                    for m in idx[km]:
                        em = mo_energy[km][m] - omega

                        # body dielectric matrix eps_body
                        Pi = get_rho_response_R(gw, abs(em), mo_energy, Lij, kL, kidx)
                        eps_body_inv = np.linalg.inv(np.eye(naux)-Pi)

                        if gw.fc and m == orbp:
                            # head dielectric matrix eps_00
                            Pi_00 = get_rho_response_head_R(gw, abs(em), mo_energy, qij)
                            eps_00 = 1. - 4. * np.pi/np.linalg.norm(q_abs[0])**2 * Pi_00

                            # wings dielectric matrix eps_P0
                            Pi_P0 = get_rho_response_wing_R(gw, abs(em), mo_energy, Lij, qij)
                            eps_P0 = -np.sqrt(4.*np.pi) / np.linalg.norm(q_abs[0]) * Pi_P0

                            # inverse dielectric matrix
                            eps_inv_00 = 1./(eps_00 - np.dot(np.dot(eps_P0.conj(),eps_body_inv),eps_P0))
                            eps_inv_P0 = -eps_inv_00 * np.dot(eps_body_inv, eps_P0)

                        eps_inv_PQ = eps_body_inv

                        # body
                        Qmn = einsum('P,PQ->Q',Lij[km][:,m,orbp].conj(),eps_inv_PQ-np.eye(naux))
                        Wmn = 1./nkpts * einsum('Q,Q->',Qmn,Lij[km][:,m,orbp])
                        sigmaR += fm * Wmn

                        if gw.fc and m == orbp:
                            # head correction
                            Del_00 = 2./np.pi * (6.*np.pi**2/gw.mol.vol/nkpts)**(1./3.) * (eps_inv_00 - 1.)
                            sigmaR += fm * Del_00

                            # wings correction
                            wings_const = np.sqrt(gw.mol.vol/4./np.pi**3) * (6.*np.pi**2/gw.mol.vol/nkpts)**(2./3.)
                            Wn_P0 = einsum('P,P->',Lij[kn][:,m,orbp].conj(),eps_inv_P0)
                            Wn_P0 = Wn_P0.real * 2.
                            sigmaR += fm * wings_const * Wn_P0
            else:
                km = kidx_r[kn]
                if len(idx[km]) > 0:
                    for m in idx[km]:
                        em = mo_energy[km][m] - omega
                        Pi = get_rho_response_R(gw, abs(em), mo_energy, Lij, kL, kidx)
                        Pi_inv = np.linalg.inv(np.eye(naux)-Pi)-np.eye(naux)
                        Qmn = einsum('P,PQ->Q',Lij[km][:,m,orbp].conj(),Pi_inv)
                        Wmn = 1./nkpts * einsum('Q,Q->',Qmn,Lij[km][:,m,orbp])
                        sigmaR += fm * Wmn

    return sigmaR

def get_rho_response_head_R(gw, omega, mo_energy, qij):
    '''
    Compute head (G=0, G'=0) density response function in auxiliary basis at freq w
    '''
    nkpts, nocc, nvir = qij.shape
    nocc = gw.nocc
    kpts = gw.kpts

    # Compute Pi head
    Pi_00 = 0j
    for i, kpti in enumerate(kpts):
        eia = mo_energy[i,:nocc,None] - mo_energy[i,None,nocc:]
        eia = 1./(omega+eia+2j*gw.eta) + 1./(-omega+eia)
        Pi_00 += 2./nkpts * einsum('ia,ia->',eia,qij[i].conj()*qij[i])
    return Pi_00

def get_rho_response_wing_R(gw, omega, mo_energy, Lpq, qij):
    '''
    Compute density response function in auxiliary basis at freq iw
    '''
    nkpts, naux, nmo, nmo = Lpq.shape
    nocc = gw.nocc
    kpts = gw.kpts

    # Compute Pi for kL
    Pi = np.zeros(naux,dtype=np.complex128)
    for i, kpti in enumerate(kpts):
        eia = mo_energy[i,:nocc,None] - mo_energy[i,None,nocc:]
        eia = 1./(omega+eia+2j*gw.eta) + 1./(-omega+eia)
        eia_q = eia * qij[i].conj()
        Pi += 2./nkpts * einsum('Pia,ia->P',Lpq[i][:,:nocc,nocc:],eia_q)
    return Pi

def get_rho_response_head(gw, omega, mo_energy, qij):
    '''
    Compute head (G=0, G'=0) density response function in auxiliary basis at freq iw
    '''
    nkpts, nocc, nvir = qij.shape
    nocc = gw.nocc
    kpts = gw.kpts

    # Compute Pi head
    Pi_00 = 0j
    for i, kpti in enumerate(kpts):
        eia = mo_energy[i,:nocc,None] - mo_energy[i,None,nocc:]
        eia = eia/(omega**2+eia*eia)
        Pi_00 += 4./nkpts * einsum('ia,ia->',eia,qij[i].conj()*qij[i])
    return Pi_00

def get_rho_response_wing(gw, omega, mo_energy, Lpq, qij):
    '''
    Compute wing (G=P, G'=0) density response function in auxiliary basis at freq iw
    '''
    nkpts, naux, nmo, nmo = Lpq.shape
    nocc = gw.nocc
    kpts = gw.kpts

    # Compute Pi wing
    Pi = np.zeros(naux,dtype=np.complex128)
    for i, kpti in enumerate(kpts):
        eia = mo_energy[i,:nocc,None] - mo_energy[i,None,nocc:]
        eia = eia/(omega**2+eia*eia)
        eia_q = eia * qij[i].conj()
        Pi += 4./nkpts * einsum('Pia,ia->P',Lpq[i][:,:nocc,nocc:],eia_q)
    return Pi

def get_qij(gw, q, mo_coeff, uniform_grids=False):
    '''
    Compute qij = 1/Omega * |< psi_{ik} | e^{iqr} | psi_{ak-q} >|^2 at q: (nkpts, nocc, nvir)
    '''
    nocc = gw.nocc
    nmo = gw.nmo
    nvir = nmo - nocc
    kpts = gw.kpts
    nkpts = len(kpts)
    cell = gw.mol
    mo_energy = gw._scf.mo_energy

    if uniform_grids:
        mydf = df.FFTDF(cell, kpts=kpts)
        coords = cell.gen_uniform_grids(mydf.mesh)
    else:
        coords, weights = dft.gen_grid.get_becke_grids(cell,level=5)
    ngrid = len(coords)

    qij = np.zeros((nkpts,nocc,nvir),dtype=np.complex128)
    for i, kpti in enumerate(kpts):
        ao_p = dft.numint.eval_ao(cell, coords, kpt=kpti, deriv=1)
        ao = ao_p[0]
        ao_grad = ao_p[1:4]
        if uniform_grids:
            ao_ao_grad = einsum('mg,xgn->xmn',ao.T.conj(),ao_grad) * cell.vol / ngrid
        else:
            ao_ao_grad = einsum('g,mg,xgn->xmn',weights,ao.T.conj(),ao_grad)
        q_ao_ao_grad = -1j * einsum('x,xmn->mn',q,ao_ao_grad)
        q_mo_mo_grad = np.dot(np.dot(mo_coeff[i][:,:nocc].T.conj(), q_ao_ao_grad), mo_coeff[i][:,nocc:])
        enm = 1./(mo_energy[i][nocc:,None] - mo_energy[i][None,:nocc])
        dens = enm.T * q_mo_mo_grad
        qij[i] = dens / np.sqrt(cell.vol)

    return qij

def _get_scaled_legendre_roots(nw):
    """
    Scale nw Legendre roots, which lie in the
    interval [-1, 1], so that they lie in [0, inf)
    Ref: www.cond-mat.de/events/correl19/manuscripts/ren.pdf

    Returns:
        freqs : 1D ndarray
        wts : 1D ndarray
    """
    freqs, wts = np.polynomial.legendre.leggauss(nw)
    x0 = 0.5
    freqs_new = x0*(1.+freqs)/(1.-freqs)
    wts = wts*2.*x0/(1.-freqs)**2
    return freqs_new, wts

def _get_clenshaw_curtis_roots(nw):
    """
    Clenshaw-Curtis qaudrature on [0,inf)
    Ref: J. Chem. Phys. 132, 234114 (2010)
    Returns:
        freqs : 1D ndarray
        wts : 1D ndarray
    """
    freqs = np.zeros(nw)
    wts = np.zeros(nw)
    a = 0.2
    for w in range(nw):
        t = (w+1.0)/nw * np.pi/2.
        freqs[w] = a / np.tan(t)
        if w != nw-1:
            wts[w] = a*np.pi/2./nw/(np.sin(t)**2)
        else:
            wts[w] = a*np.pi/4./nw/(np.sin(t)**2)
    return freqs[::-1], wts[::-1]


class KRGWCD(lib.StreamObject):

    linearized = getattr(__config__, 'gw_gw_GW_linearized', False)
    eta = getattr(__config__, 'gw_gw_GW_eta', 1e-3)
    fc = getattr(__config__, 'gw_gw_GW_fc', True)

    def __init__(self, mf, frozen=0):
        self.mol = mf.mol
        self._scf = mf
        self.verbose = self.mol.verbose
        self.stdout = self.mol.stdout
        self.max_memory = mf.max_memory

        #TODO: implement frozen orbs
        if frozen > 0:
            raise NotImplementedError
        self.frozen = frozen

        # DF-KGW must use GDF integrals
        if getattr(mf, 'with_df', None):
            self.with_df = mf.with_df
        else:
            raise NotImplementedError
        self._keys.update(['with_df'])

##################################################
# don't modify the following attributes, they are not input options
        self._nocc = None
        self._nmo = None
        self.kpts = mf.kpts
        self.nkpts = len(self.kpts)
        # self.mo_energy: GW quasiparticle energy, not scf mo_energy
        self.mo_energy = None
        self.mo_coeff = mf.mo_coeff
        self.mo_occ = mf.mo_occ
        self.sigma = None

        keys = set(('linearized','eta','fc'))
        self._keys = set(self.__dict__.keys()).union(keys)

    def dump_flags(self):
        log = logger.Logger(self.stdout, self.verbose)
        log.info('')
        log.info('******** %s ********', self.__class__)
        log.info('method = %s', self.__class__.__name__)
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        log.info('GW nocc = %d, nvir = %d, nkpts = %d', nocc, nvir, nkpts)
        if self.frozen is not None:
            log.info('frozen orbitals %s', str(self.frozen))
        logger.info(self, 'use perturbative linearized QP eqn = %s', self.linearized)
        logger.info(self, 'GW finite size corrections = %s', self.fc)
        return self

    @property
    def nocc(self):
        return self.get_nocc()
    @nocc.setter
    def nocc(self, n):
        self._nocc = n

    @property
    def nmo(self):
        return self.get_nmo()
    @nmo.setter
    def nmo(self, n):
        self._nmo = n

    get_nocc = get_nocc
    get_nmo = get_nmo
    get_frozen_mask = get_frozen_mask

    def kernel(self, mo_energy=None, mo_coeff=None, orbs=None, kptlist=None, nw=100):
        """
        Input:
            kptlist: self-energy k-points
            orbs: self-energy orbs
            nw: grid number
        Output:
            mo_energy: GW quasiparticle energy
        """
        if mo_coeff is None:
            mo_coeff = np.array(self._scf.mo_coeff)
        if mo_energy is None:
            mo_energy = np.array(self._scf.mo_energy)

        nmo = self.nmo
        naux = self.with_df.get_naoaux()
        nkpts = self.nkpts
        mem_incore = (2*nkpts*nmo**2*naux) * 16/1e6
        mem_now = lib.current_memory()[0]
        if (mem_incore + mem_now > 0.99*self.max_memory):
            logger.warn(self, 'Memory may not be enough!')
            raise NotImplementedError

        cput0 = (logger.process_clock(), logger.perf_counter())
        self.dump_flags()
        self.converged, self.mo_energy, self.mo_coeff = \
                kernel(self, mo_energy, mo_coeff, orbs=orbs,
                       kptlist=kptlist, nw=nw, verbose=self.verbose)

        logger.warn(self, 'GW QP energies may not be sorted from min to max')
        logger.timer(self, 'GW', *cput0)
        return self.mo_energy

if __name__ == '__main__':
    from pyscf.pbc import gto
    from pyscf.pbc.lib import chkfile
    import os
    # This test takes a few minutes
    cell = gto.Cell()
    cell.build(unit = 'angstrom',
               a = '''
               0.000000     1.783500     1.783500
               1.783500     0.000000     1.783500
               1.783500     1.783500     0.000000
               ''',
               atom = 'C 1.337625 1.337625 1.337625; C 2.229375 2.229375 2.229375',
               dimension = 3,
               max_memory = 8000,
               verbose = 4,
               pseudo = 'gth-pade',
               basis='gth-szv',
               precision=1e-10)

    kpts = cell.make_kpts([3,1,1],scaled_center=[0,0,0])
    gdf = df.GDF(cell, kpts)
    gdf_fname = 'gdf_ints_311.h5'
    gdf._cderi_to_save = gdf_fname
    if not os.path.isfile(gdf_fname):
        gdf.build()

    chkfname = 'diamond_311.chk'
    if os.path.isfile(chkfname):
        kmf = dft.KRKS(cell, kpts)
        kmf.xc = 'pbe'
        kmf.with_df = gdf
        kmf.with_df._cderi = gdf_fname
        data = chkfile.load(chkfname, 'scf')
        kmf.__dict__.update(data)
    else:
        kmf = dft.KRKS(cell, kpts)
        kmf.xc = 'pbe'
        kmf.with_df = gdf
        kmf.with_df._cderi = gdf_fname
        kmf.conv_tol = 1e-12
        kmf.chkfile = chkfname
        kmf.kernel()

    gw = KRGWCD(kmf)
    gw.linearized = False
    # without finite size corrections
    gw.fc = False
    nocc = gw.nocc
    gw.kernel(kptlist=[0,1,2],orbs=range(0,nocc+3))
    print(gw.mo_energy)
    assert((abs(gw.mo_energy[0][nocc-1]-0.62045796))<1e-5)
    assert((abs(gw.mo_energy[0][nocc]-0.96574426))<1e-5)
    assert((abs(gw.mo_energy[1][nocc-1]-0.52639129))<1e-5)
    assert((abs(gw.mo_energy[1][nocc]-1.07442235))<1e-5)

    # with finite size corrections
    gw.fc = True
    gw.kernel(kptlist=[0,1,2],orbs=range(0,nocc+3))
    print(gw.mo_energy)
    assert((abs(gw.mo_energy[0][nocc-1]-0.5427707))<1e-5)
    assert((abs(gw.mo_energy[0][nocc]-0.80148557))<1e-5)
    assert((abs(gw.mo_energy[1][nocc-1]-0.45073751))<1e-5)
    assert((abs(gw.mo_energy[1][nocc]-0.92910117))<1e-5)


# pyscf/pbc/gw/kugw_ac.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Tianyu Zhu <zhutianyu1991@gmail.com>
#

'''
PBC spin-unrestricted G0W0-AC QP eigenvalues with k-point sampling
GW-AC is recommended for valence states only, and is inaccuarate for core states.

Method:
    See T. Zhu and G.K.-L. Chan, arxiv:2007.03148 (2020) for details
    Compute Sigma on imaginary frequency with density fitting,
    then analytically continued to real frequency.
    Gaussian density fitting must be used (FFTDF and MDF are not supported).
'''

from functools import reduce
import numpy
import numpy as np
import h5py
from scipy.optimize import newton, least_squares

from pyscf import lib
from pyscf.lib import logger
from pyscf.ao2mo import _ao2mo
from pyscf.ao2mo.incore import _conc_mos
from pyscf.pbc import df, dft, scf
from pyscf.pbc.cc.kccsd_uhf import get_nocc, get_nmo, get_frozen_mask
from pyscf import __config__

einsum = lib.einsum

def kernel(gw, mo_energy, mo_coeff, orbs=None,
           kptlist=None, nw=None, verbose=logger.NOTE):
    '''
    GW-corrected quasiparticle orbital energies
    Returns:
        A list :  converged, mo_energy, mo_coeff
    '''
    mf = gw._scf
    if gw.frozen is None:
        frozen = 0
    else:
        frozen = gw.frozen
    assert (frozen == 0)

    nmoa, nmob = gw.nmo
    nocca, noccb = gw.nocc

    if orbs is None:
        orbs = range(nmoa)
    if kptlist is None:
        kptlist = range(gw.nkpts)

    nkpts = gw.nkpts
    nklist = len(kptlist)

    # v_xc
    dm = np.array(mf.make_rdm1())
    v_mf = np.array(mf.get_veff())
    vj = np.array(mf.get_j(dm_kpts=dm))
    v_mf[0] = v_mf[0] - (vj[0] + vj[1])
    v_mf[1] = v_mf[1] - (vj[0] + vj[1])
    for s in range(2):
        for k in range(nkpts):
            v_mf[s,k] = reduce(numpy.dot, (mo_coeff[s,k].T.conj(), v_mf[s,k], mo_coeff[s,k]))

    # v_hf from DFT/HF density
    if gw.fc:
        exxdiv = 'ewald'
    else:
        exxdiv = None
    uhf = scf.KUHF(gw.mol, gw.kpts, exxdiv=exxdiv)
    uhf.with_df = gw.with_df
    uhf.with_df._cderi = gw.with_df._cderi
    vk = uhf.get_veff(gw.mol,dm_kpts=dm)
    vj = uhf.get_j(gw.mol,dm_kpts=dm)
    vk[0] = vk[0] - (vj[0] + vj[1])
    vk[1] = vk[1] - (vj[0] + vj[1])
    for s in range(2):
        for k in range(nkpts):
            vk[s,k] = reduce(numpy.dot, (mo_coeff[s,k].T.conj(), vk[s,k], mo_coeff[s,k]))

    # Grids for integration on imaginary axis
    freqs,wts = _get_scaled_legendre_roots(nw)

    # Compute self-energy on imaginary axis i*[0,iw_cutoff]
    sigmaI, omega = get_sigma_diag(gw, orbs, kptlist, freqs, wts, iw_cutoff=5.)

    # Analytic continuation
    coeff_a = []
    coeff_b = []
    if gw.ac == 'twopole':
        for k in range(nklist):
            coeff_a.append(AC_twopole_diag(sigmaI[0,k], omega[0], orbs, nocca))
            coeff_b.append(AC_twopole_diag(sigmaI[1,k], omega[1], orbs, noccb))
    elif gw.ac == 'pade':
        for k in range(nklist):
            coeff_a_tmp, omega_fit_a = AC_pade_thiele_diag(sigmaI[0,k], omega[0])
            coeff_b_tmp, omega_fit_b = AC_pade_thiele_diag(sigmaI[1,k], omega[1])
            coeff_a.append(coeff_a_tmp)
            coeff_b.append(coeff_b_tmp)
        omega_fit = np.asarray((omega_fit_a, omega_fit_b))
    coeff = np.asarray((coeff_a, coeff_b))

    conv = True
    # This code does not support metals
    homo = -99.
    lumo = 99.
    mo_energy = np.asarray(mf.mo_energy)
    for k in range(nkpts):
        if homo < max(mo_energy[0,k][nocca-1],mo_energy[1,k][noccb-1]):
            homo = max(mo_energy[0,k][nocca-1],mo_energy[1,k][noccb-1])
        if lumo > min(mo_energy[0,k][nocca],mo_energy[1,k][noccb]):
            lumo = min(mo_energy[0,k][nocca],mo_energy[1,k][noccb])
    ef = (homo+lumo)/2.

    mo_energy = np.zeros_like(np.array(mf.mo_energy))
    for s in range(2):
        for k in range(nklist):
            kn = kptlist[k]
            for p in orbs:
                if gw.linearized:
                    # linearized G0W0
                    de = 1e-6
                    ep = mf.mo_energy[s][kn][p]
                    #TODO: analytic sigma derivative
                    if gw.ac == 'twopole':
                        sigmaR = two_pole(ep-ef, coeff[s,k,:,p-orbs[0]]).real
                        dsigma = two_pole(ep-ef+de, coeff[s,k,:,p-orbs[0]]).real - sigmaR.real
                    elif gw.ac == 'pade':
                        sigmaR = pade_thiele(ep-ef, omega_fit[s,p-orbs[0]], coeff[s,k,:,p-orbs[0]]).real
                        dsigma = pade_thiele(ep-ef+de, omega_fit[s,p-orbs[0]],
                                             coeff[s,k,:,p-orbs[0]]).real - sigmaR.real
                    zn = 1.0/(1.0-dsigma/de)
                    e = ep + zn*(sigmaR.real + vk[s,kn,p,p].real - v_mf[s,kn,p,p].real)
                    mo_energy[s,kn,p] = e
                else:
                    # self-consistently solve QP equation
                    def quasiparticle(omega):
                        if gw.ac == 'twopole':
                            sigmaR = two_pole(omega-ef, coeff[s,k,:,p-orbs[0]]).real
                        elif gw.ac == 'pade':
                            sigmaR = pade_thiele(omega-ef, omega_fit[s,p-orbs[0]], coeff[s,k,:,p-orbs[0]]).real
                        return omega - mf.mo_energy[s][kn][p] - (sigmaR.real + vk[s,kn,p,p].real - v_mf[s,kn,p,p].real)
                    try:
                        e = newton(quasiparticle, mf.mo_energy[s][kn][p], tol=1e-6, maxiter=100)
                        mo_energy[s,kn,p] = e
                    except RuntimeError:
                        conv = False
    mo_coeff = mf.mo_coeff

    if gw.verbose >= logger.DEBUG:
        numpy.set_printoptions(threshold=nmoa)
        for k in range(nkpts):
            logger.debug(gw, '  GW mo_energy spin-up @ k%d =\n%s', k,mo_energy[0,k])
        for k in range(nkpts):
            logger.debug(gw, '  GW mo_energy spin-down @ k%d =\n%s', k,mo_energy[1,k])
        numpy.set_printoptions(threshold=1000)

    return conv, mo_energy, mo_coeff

def get_rho_response(gw, omega, mo_energy, Lpq, kL, kidx):
    '''
    Compute density response function in auxiliary basis at freq iw
    '''
    spin, nkpts, naux, nmo, nmo = Lpq.shape
    nocca, noccb = gw.nocc
    kpts = gw.kpts
    kscaled = gw.mol.get_scaled_kpts(kpts)
    kscaled -= kscaled[0]

    # Compute Pi for kL
    Pi = np.zeros((naux,naux),dtype=np.complex128)
    for i, kpti in enumerate(kpts):
        # Find ka that conserves with ki and kL (-ki+ka+kL=G)
        a = kidx[i]
        eia_a = mo_energy[0,i,:nocca,None] - mo_energy[0,a,None,nocca:]
        eia_b = mo_energy[1,i,:noccb,None] - mo_energy[1,a,None,noccb:]
        eia_a = eia_a/(omega**2+eia_a*eia_a)
        eia_b = eia_b/(omega**2+eia_b*eia_b)
        Pia_a = einsum('Pia,ia->Pia',Lpq[0,i][:,:nocca,nocca:],eia_a)
        Pia_b = einsum('Pia,ia->Pia',Lpq[1,i][:,:noccb,noccb:],eia_b)
        # Response from both spin-up and spin-down density
        Pi += 2./nkpts * (einsum('Pia,Qia->PQ',Pia_a,Lpq[0,i][:,:nocca,nocca:].conj()) +
                          einsum('Pia,Qia->PQ',Pia_b,Lpq[1,i][:,:noccb,noccb:].conj()))
    return Pi

def get_sigma_diag(gw, orbs, kptlist, freqs, wts, iw_cutoff=None, max_memory=8000):
    '''
    Compute GW correlation self-energy (diagonal elements) in MO basis
    on imaginary axis
    '''
    mo_energy = np.array(gw._scf.mo_energy)
    mo_coeff = np.array(gw._scf.mo_coeff)
    nocca, noccb = gw.nocc
    nmoa, nmob = gw.nmo
    nkpts = gw.nkpts
    kpts = gw.kpts
    nklist = len(kptlist)
    nw = len(freqs)
    norbs = len(orbs)
    mydf = gw.with_df

    # possible kpts shift
    kscaled = gw.mol.get_scaled_kpts(kpts)
    kscaled -= kscaled[0]

    # This code does not support metals
    homo = -99.
    lumo = 99.
    for k in range(nkpts):
        if homo < max(mo_energy[0,k][nocca-1],mo_energy[1,k][noccb-1]):
            homo = max(mo_energy[0,k][nocca-1],mo_energy[1,k][noccb-1])
        if lumo > min(mo_energy[0,k][nocca],mo_energy[1,k][noccb]):
            lumo = min(mo_energy[0,k][nocca],mo_energy[1,k][noccb])
    if (lumo-homo)<1e-3:
        logger.warn(gw, 'Current KUGW is not supporting metals!')
    ef = (homo+lumo)/2.

    # Integration on numerical grids
    if iw_cutoff is not None:
        nw_sigma = sum(iw < iw_cutoff for iw in freqs) + 1
    else:
        nw_sigma = nw + 1

    # Compute occ for -iw and vir for iw separately
    # to avoid branch cuts in analytic continuation
    omega_occ = np.zeros((nw_sigma),dtype=np.complex128)
    omega_vir = np.zeros((nw_sigma),dtype=np.complex128)
    omega_occ[1:] = -1j*freqs[:(nw_sigma-1)]
    omega_vir[1:] = 1j*freqs[:(nw_sigma-1)]
    orbs_occ_a = [i for i in orbs if i < nocca]
    orbs_occ_b = [i for i in orbs if i < noccb]
    norbs_occ_a = len(orbs_occ_a)
    norbs_occ_b = len(orbs_occ_b)

    emo_occ_a = np.zeros((nkpts,nmoa,nw_sigma),dtype=np.complex128)
    emo_occ_b = np.zeros((nkpts,nmob,nw_sigma),dtype=np.complex128)
    emo_vir_a = np.zeros((nkpts,nmoa,nw_sigma),dtype=np.complex128)
    emo_vir_b = np.zeros((nkpts,nmob,nw_sigma),dtype=np.complex128)
    for k in range(nkpts):
        emo_occ_a[k] = omega_occ[None,:] + ef - mo_energy[0,k][:,None]
        emo_occ_b[k] = omega_occ[None,:] + ef - mo_energy[1,k][:,None]
        emo_vir_a[k] = omega_vir[None,:] + ef - mo_energy[0,k][:,None]
        emo_vir_b[k] = omega_vir[None,:] + ef - mo_energy[1,k][:,None]

    sigma = np.zeros((2,nklist,norbs,nw_sigma),dtype=np.complex128)
    omega = np.zeros((2,norbs,nw_sigma),dtype=np.complex128)
    for s in range(2):
        for p in range(norbs):
            orbp = orbs[p]
            if orbp < gw.nocc[s]:
                omega[s,p] = omega_occ.copy()
            else:
                omega[s,p] = omega_vir.copy()

    if gw.fc:
        # Set up q mesh for q->0 finite size correction
        q_pts = np.array([1e-3,0,0]).reshape(1,3)
        q_abs = gw.mol.get_abs_kpts(q_pts)

        # Get qij = 1/sqrt(Omega) * < psi_{ik} | e^{iqr} | psi_{ak-q} > at q: (nkpts, nocc, nvir)
        qij = get_qij(gw, q_abs[0], mo_coeff)

    for kL in range(nkpts):
        # Lij: (2, ki, L, i, j) for looping every kL
        #Lij = np.zeros((2,nkpts,naux,nmoa,nmoa),dtype=np.complex128)
        Lij = []
        # kidx: save kj that conserves with kL and ki (-ki+kj+kL=G)
        # kidx_r: save ki that conserves with kL and kj (-ki+kj+kL=G)
        kidx = np.zeros((nkpts),dtype=np.int64)
        kidx_r = np.zeros((nkpts),dtype=np.int64)
        for i, kpti in enumerate(kpts):
            for j, kptj in enumerate(kpts):
                # Find (ki,kj) that satisfies momentum conservation with kL
                kconserv = -kscaled[i] + kscaled[j] + kscaled[kL]
                is_kconserv = np.linalg.norm(np.round(kconserv) - kconserv) < 1e-12
                if is_kconserv:
                    kidx[i] = j
                    kidx_r[j] = i
                    logger.debug(gw, "Read Lpq (kL: %s / %s, ki: %s, kj: %s)"%(kL+1, nkpts, i, j))
                    Lij_out_a = None
                    Lij_out_b = None
                    # Read (L|pq) and ao2mo transform to (L|ij)
                    Lpq = []
                    for LpqR, LpqI, sign \
                            in mydf.sr_loop([kpti, kptj], max_memory=0.1*gw._scf.max_memory, compact=False):
                        Lpq.append(LpqR+LpqI*1.0j)
                    Lpq = np.vstack(Lpq).reshape(-1,nmoa**2)
                    moija, ijslicea = _conc_mos(mo_coeff[0,i], mo_coeff[0,j])[2:]
                    moijb, ijsliceb = _conc_mos(mo_coeff[1,i], mo_coeff[1,j])[2:]
                    tao = []
                    ao_loc = None
                    Lij_out_a = _ao2mo.r_e2(Lpq, moija, ijslicea, tao, ao_loc, out=Lij_out_a)
                    tao = []
                    ao_loc = None
                    Lij_out_b = _ao2mo.r_e2(Lpq, moijb, ijsliceb, tao, ao_loc, out=Lij_out_b)
                    Lij.append(np.asarray((Lij_out_a.reshape(-1,nmoa,nmoa),Lij_out_b.reshape(-1,nmob,nmob))))

        Lij = np.asarray(Lij)
        Lij = Lij.transpose(1,0,2,3,4)
        naux = Lij.shape[2]

        if kL == 0:
            for w in range(nw):
                # body dielectric matrix eps_body
                Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx)
                eps_body_inv = np.linalg.inv(np.eye(naux)-Pi)

                if gw.fc:
                    # head dielectric matrix eps_00
                    Pi_00 = get_rho_response_head(gw, freqs[w], mo_energy, qij)
                    eps_00 = 1. - 4. * np.pi/np.linalg.norm(q_abs[0])**2 * Pi_00

                    # wings dielectric matrix eps_P0
                    Pi_P0 = get_rho_response_wing(gw, freqs[w], mo_energy, Lij, qij)
                    eps_P0 = -np.sqrt(4.*np.pi) / np.linalg.norm(q_abs[0]) * Pi_P0

                    # inverse dielectric matrix
                    eps_inv_00 = 1./(eps_00 - np.dot(np.dot(eps_P0.conj(),eps_body_inv),eps_P0))
                    eps_inv_P0 = -eps_inv_00 * np.dot(eps_body_inv, eps_P0)

                    # head correction
                    Del_00 = 2./np.pi * (6.*np.pi**2/gw.mol.vol/nkpts)**(1./3.) * (eps_inv_00 - 1.)

                eps_inv_PQ = eps_body_inv
                g0_occ_a = wts[w] * emo_occ_a / (emo_occ_a**2+freqs[w]**2)
                g0_occ_b = wts[w] * emo_occ_b / (emo_occ_b**2+freqs[w]**2)
                g0_vir_a = wts[w] * emo_vir_a / (emo_vir_a**2+freqs[w]**2)
                g0_vir_b = wts[w] * emo_vir_b / (emo_vir_b**2+freqs[w]**2)
                for k in range(nklist):
                    kn = kptlist[k]
                    # Find km that conserves with kn and kL (-km+kn+kL=G)
                    km = kidx_r[kn]
                    Qmn_a = einsum('Pmn,PQ->Qmn',Lij[0,km][:,:,orbs].conj(),eps_inv_PQ-np.eye(naux))
                    Qmn_b = einsum('Pmn,PQ->Qmn',Lij[1,km][:,:,orbs].conj(),eps_inv_PQ-np.eye(naux))
                    Wmn_a = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn_a,Lij[0,km][:,:,orbs])
                    Wmn_b = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn_b,Lij[1,km][:,:,orbs])

                    sigma[0,k][:norbs_occ_a] += -einsum('mn,mw->nw',Wmn_a[:,:norbs_occ_a],g0_occ_a[km])/np.pi
                    sigma[1,k][:norbs_occ_b] += -einsum('mn,mw->nw',Wmn_b[:,:norbs_occ_b],g0_occ_b[km])/np.pi
                    sigma[0,k][norbs_occ_a:] += -einsum('mn,mw->nw',Wmn_a[:,norbs_occ_a:],g0_vir_a[km])/np.pi
                    sigma[1,k][norbs_occ_b:] += -einsum('mn,mw->nw',Wmn_b[:,norbs_occ_b:],g0_vir_b[km])/np.pi

                    if gw.fc:
                        # apply head correction
                        assert(kn == km)
                        sigma[0,k][:norbs_occ_a] += -Del_00 * g0_occ_a[kn][orbs][:norbs_occ_a] /np.pi
                        sigma[0,k][norbs_occ_a:] += -Del_00 * g0_vir_a[kn][orbs][norbs_occ_a:] /np.pi
                        sigma[1,k][:norbs_occ_b] += -Del_00 * g0_occ_b[kn][orbs][:norbs_occ_b] /np.pi
                        sigma[1,k][norbs_occ_b:] += -Del_00 * g0_vir_b[kn][orbs][norbs_occ_b:] /np.pi

                        # apply wing correction
                        Wn_P0_a = einsum('Pnm,P->nm',Lij[0,kn],eps_inv_P0).diagonal()
                        Wn_P0_b = einsum('Pnm,P->nm',Lij[1,kn],eps_inv_P0).diagonal()
                        Wn_P0_a = Wn_P0_a.real * 2.
                        Wn_P0_b = Wn_P0_b.real * 2.
                        Del_P0_a = np.sqrt(gw.mol.vol/4./np.pi**3) * (6.*np.pi**2/gw.mol.vol/nkpts)**(2./3.) * Wn_P0_a[orbs]  # noqa: E501
                        Del_P0_b = np.sqrt(gw.mol.vol/4./np.pi**3) * (6.*np.pi**2/gw.mol.vol/nkpts)**(2./3.) * Wn_P0_b[orbs]  # noqa: E501
                        sigma[0,k][:norbs_occ_a] += -einsum('n,nw->nw',Del_P0_a[:norbs_occ_a],g0_occ_a[kn][orbs][:norbs_occ_a]) /np.pi  # noqa: E501
                        sigma[0,k][norbs_occ_a:] += -einsum('n,nw->nw',Del_P0_a[norbs_occ_a:],g0_vir_a[kn][orbs][norbs_occ_a:]) /np.pi  # noqa: E501
                        sigma[1,k][:norbs_occ_b] += -einsum('n,nw->nw',Del_P0_b[:norbs_occ_b],g0_occ_b[kn][orbs][:norbs_occ_b]) /np.pi  # noqa: E501
                        sigma[1,k][norbs_occ_b:] += -einsum('n,nw->nw',Del_P0_b[norbs_occ_b:],g0_vir_b[kn][orbs][norbs_occ_b:]) /np.pi  # noqa: E501
        else:
            for w in range(nw):
                Pi = get_rho_response(gw, freqs[w], mo_energy, Lij, kL, kidx)
                Pi_inv = np.linalg.inv(np.eye(naux)-Pi)-np.eye(naux)
                g0_occ_a = wts[w] * emo_occ_a / (emo_occ_a**2+freqs[w]**2)
                g0_occ_b = wts[w] * emo_occ_b / (emo_occ_b**2+freqs[w]**2)
                g0_vir_a = wts[w] * emo_vir_a / (emo_vir_a**2+freqs[w]**2)
                g0_vir_b = wts[w] * emo_vir_b / (emo_vir_b**2+freqs[w]**2)
                for k in range(nklist):
                    kn = kptlist[k]
                    # Find km that conserves with kn and kL (-km+kn+kL=G)
                    km = kidx_r[kn]
                    Qmn_a = einsum('Pmn,PQ->Qmn',Lij[0,km][:,:,orbs].conj(),Pi_inv)
                    Qmn_b = einsum('Pmn,PQ->Qmn',Lij[1,km][:,:,orbs].conj(),Pi_inv)
                    Wmn_a = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn_a,Lij[0,km][:,:,orbs])
                    Wmn_b = 1./nkpts * einsum('Qmn,Qmn->mn',Qmn_b,Lij[1,km][:,:,orbs])

                    sigma[0,k][:norbs_occ_a] += -einsum('mn,mw->nw',Wmn_a[:,:norbs_occ_a],g0_occ_a[km])/np.pi
                    sigma[1,k][:norbs_occ_b] += -einsum('mn,mw->nw',Wmn_b[:,:norbs_occ_b],g0_occ_b[km])/np.pi
                    sigma[0,k][norbs_occ_a:] += -einsum('mn,mw->nw',Wmn_a[:,norbs_occ_a:],g0_vir_a[km])/np.pi
                    sigma[1,k][norbs_occ_b:] += -einsum('mn,mw->nw',Wmn_b[:,norbs_occ_b:],g0_vir_b[km])/np.pi

    return sigma, omega

def get_rho_response_head(gw, omega, mo_energy, qij):
    '''
    Compute head (G=0, G'=0) density response function in auxiliary basis at freq iw
    '''
    qij_a, qij_b = qij
    nocca, noccb = gw.nocc
    kpts = gw.kpts
    nkpts = len(kpts)

    # Compute Pi head
    Pi_00 = 0j
    for i, kpti in enumerate(kpts):
        eia_a = mo_energy[0,i,:nocca,None] - mo_energy[0,i,None,nocca:]
        eia_b = mo_energy[1,i,:noccb,None] - mo_energy[1,i,None,noccb:]
        eia_a = eia_a/(omega**2+eia_a*eia_a)
        eia_b = eia_b/(omega**2+eia_b*eia_b)
        Pi_00 += 2./nkpts * (einsum('ia,ia->',eia_a,qij_a[i].conj()*qij_a[i]) +
                             einsum('ia,ia->',eia_b,qij_b[i].conj()*qij_b[i]))
    return Pi_00

def get_rho_response_wing(gw, omega, mo_energy, Lpq, qij):
    '''
    Compute wing (G=P, G'=0) density response function in auxiliary basis at freq iw
    '''
    qij_a, qij_b = qij
    spin, nkpts, naux, nmo, nmo = Lpq.shape
    nocca, noccb = gw.nocc
    kpts = gw.kpts
    nkpts = len(kpts)

    # Compute Pi wing
    Pi = np.zeros(naux,dtype=np.complex128)
    for i, kpti in enumerate(kpts):
        eia_a = mo_energy[0,i,:nocca,None] - mo_energy[0,i,None,nocca:]
        eia_b = mo_energy[1,i,:noccb,None] - mo_energy[1,i,None,noccb:]
        eia_a = eia_a/(omega**2+eia_a*eia_a)
        eia_b = eia_b/(omega**2+eia_b*eia_b)
        eia_q_a = eia_a * qij_a[i].conj()
        eia_q_b = eia_b * qij_b[i].conj()
        Pi += 2./nkpts * (einsum('Pia,ia->P',Lpq[0,i][:,:nocca,nocca:],eia_q_a) +
                          einsum('Pia,ia->P',Lpq[1,i][:,:noccb,noccb:],eia_q_b))
    return Pi

def get_qij(gw, q, mo_coeff, uniform_grids=False):
    '''
    Compute qij = 1/Omega * |< psi_{ik} | e^{iqr} | psi_{ak-q} >|^2 at q: (nkpts, nocc, nvir)
    through kp perturbtation theory
    Ref: Phys. Rev. B 83, 245122 (2011)
    '''
    nocca, noccb = gw.nocc
    nmoa, nmob = gw.nmo
    nvira = nmoa - nocca
    nvirb = nmob - noccb
    kpts = gw.kpts
    nkpts = len(kpts)
    cell = gw.mol
    mo_energy = np.asarray(gw._scf.mo_energy)

    if uniform_grids:
        mydf = df.FFTDF(cell, kpts=kpts)
        coords = cell.gen_uniform_grids(mydf.mesh)
    else:
        coords, weights = dft.gen_grid.get_becke_grids(cell,level=4)
    ngrid = len(coords)

    qij_a = np.zeros((nkpts,nocca,nvira),dtype=np.complex128)
    qij_b = np.zeros((nkpts,noccb,nvirb),dtype=np.complex128)
    for i, kpti in enumerate(kpts):
        ao_p = dft.numint.eval_ao(cell, coords, kpt=kpti, deriv=1)
        ao = ao_p[0]
        ao_grad = ao_p[1:4]
        if uniform_grids:
            ao_ao_grad = einsum('mg,xgn->xmn',ao.T.conj(),ao_grad) * cell.vol / ngrid
        else:
            ao_ao_grad = einsum('g,mg,xgn->xmn',weights,ao.T.conj(),ao_grad)
        q_ao_ao_grad = -1j * einsum('x,xmn->mn',q,ao_ao_grad)
        q_mo_mo_grad_a = np.dot(np.dot(mo_coeff[0,i][:,:nocca].T.conj(), q_ao_ao_grad), mo_coeff[0,i][:,nocca:])
        q_mo_mo_grad_b = np.dot(np.dot(mo_coeff[1,i][:,:noccb].T.conj(), q_ao_ao_grad), mo_coeff[1,i][:,noccb:])
        enm_a = 1./(mo_energy[0,i][nocca:,None] - mo_energy[0,i][None,:nocca])
        enm_b = 1./(mo_energy[1,i][noccb:,None] - mo_energy[1,i][None,:noccb])
        dens_a = enm_a.T * q_mo_mo_grad_a
        dens_b = enm_b.T * q_mo_mo_grad_b
        qij_a[i] = dens_a / np.sqrt(cell.vol)
        qij_b[i] = dens_b / np.sqrt(cell.vol)

    return (qij_a, qij_b)

def _get_scaled_legendre_roots(nw):
    """
    Scale nw Legendre roots, which lie in the
    interval [-1, 1], so that they lie in [0, inf)
    Ref: www.cond-mat.de/events/correl19/manuscripts/ren.pdf

    Returns:
        freqs : 1D ndarray
        wts : 1D ndarray
    """
    freqs, wts = np.polynomial.legendre.leggauss(nw)
    x0 = 0.5
    freqs_new = x0*(1.+freqs)/(1.-freqs)
    wts = wts*2.*x0/(1.-freqs)**2
    return freqs_new, wts

def _get_clenshaw_curtis_roots(nw):
    """
    Clenshaw-Curtis qaudrature on [0,inf)
    Ref: J. Chem. Phys. 132, 234114 (2010)
    Returns:
        freqs : 1D ndarray
        wts : 1D ndarray
    """
    freqs = np.zeros(nw)
    wts = np.zeros(nw)
    a = 0.2
    for w in range(nw):
        t = (w+1.0)/nw * np.pi/2.
        freqs[w] = a / np.tan(t)
        if w != nw-1:
            wts[w] = a*np.pi/2./nw/(np.sin(t)**2)
        else:
            wts[w] = a*np.pi/4./nw/(np.sin(t)**2)
    return freqs[::-1], wts[::-1]

def two_pole_fit(coeff, omega, sigma):
    cf = coeff[:5] + 1j*coeff[5:]
    f = cf[0] + cf[1]/(omega+cf[3]) + cf[2]/(omega+cf[4]) - sigma
    f[0] = f[0]/0.01
    return np.array([f.real,f.imag]).reshape(-1)

def two_pole(freqs, coeff):
    cf = coeff[:5] + 1j*coeff[5:]
    return cf[0] + cf[1]/(freqs+cf[3]) + cf[2]/(freqs+cf[4])

def AC_twopole_diag(sigma, omega, orbs, nocc):
    """
    Analytic continuation to real axis using a two-pole model
    Returns:
        coeff: 2D array (ncoeff, norbs)
    """
    norbs, nw = sigma.shape
    coeff = np.zeros((10,norbs))
    for p in range(norbs):
        if orbs[p] < nocc:
            x0 = np.array([0, 1, 1, 1, -1, 0, 0, 0, -1.0, -0.5])
        else:
            x0 = np.array([0, 1, 1, 1, -1, 0, 0, 0, 1.0, 0.5])
        #TODO: analytic gradient
        xopt = least_squares(two_pole_fit, x0, jac='3-point', method='trf', xtol=1e-10,
                             gtol = 1e-10, max_nfev=2000, verbose=0, args=(omega[p], sigma[p]))
        if xopt.success is False:
            print('WARN: 2P-Fit Orb %d not converged, cost function %e'%(p,xopt.cost))
        coeff[:,p] = xopt.x.copy()
    return coeff

def thiele(fn,zn):
    nfit = len(zn)
    g = np.zeros((nfit,nfit),dtype=np.complex128)
    g[:,0] = fn.copy()
    for i in range(1,nfit):
        g[i:,i] = (g[i-1,i-1]-g[i:,i-1])/((zn[i:]-zn[i-1])*g[i:,i-1])
    a = g.diagonal()
    return a

def pade_thiele(freqs,zn,coeff):
    nfit = len(coeff)
    X = coeff[-1]*(freqs-zn[-2])
    for i in range(nfit-1):
        idx = nfit-i-1
        X = coeff[idx]*(freqs-zn[idx-1])/(1.+X)
    X = coeff[0]/(1.+X)
    return X

def AC_pade_thiele_diag(sigma, omega):
    """
    Analytic continuation to real axis using a Pade approximation
    from Thiele's reciprocal difference method
    Reference: J. Low Temp. Phys. 29, 179 (1977)
    Returns:
        coeff: 2D array (ncoeff, norbs)
        omega: 2D array (norbs, npade)
    """
    idx = range(1,40,6)
    sigma1 = sigma[:,idx].copy()
    sigma2 = sigma[:,(idx[-1]+4)::4].copy()
    sigma = np.hstack((sigma1,sigma2))
    omega1 = omega[:,idx].copy()
    omega2 = omega[:,(idx[-1]+4)::4].copy()
    omega = np.hstack((omega1,omega2))
    norbs, nw = sigma.shape
    npade = nw // 2
    coeff = np.zeros((npade*2,norbs),dtype=np.complex128)
    for p in range(norbs):
        coeff[:,p] = thiele(sigma[p,:npade*2], omega[p,:npade*2])

    return coeff, omega[:,:npade*2]


class KUGWAC(lib.StreamObject):

    linearized = getattr(__config__, 'gw_gw_GW_linearized', False)
    # Analytic continuation: pade or twopole
    ac = getattr(__config__, 'gw_gw_GW_ac', 'pade')
    # Whether applying finite size corrections
    fc = getattr(__config__, 'gw_gw_GW_fc', True)

    def __init__(self, mf, frozen=0):
        self.mol = mf.mol
        self._scf = mf
        self.verbose = self.mol.verbose
        self.stdout = self.mol.stdout
        self.max_memory = mf.max_memory

        #TODO: implement frozen orbs
        if frozen > 0:
            raise NotImplementedError
        self.frozen = frozen

        # DF-KGW must use GDF integrals
        if getattr(mf, 'with_df', None):
            self.with_df = mf.with_df
        else:
            raise NotImplementedError
        self._keys.update(['with_df'])

##################################################
# don't modify the following attributes, they are not input options
        self._nocc = None
        self._nmo = None
        self.kpts = mf.kpts
        self.nkpts = len(self.kpts)
        # self.mo_energy: GW quasiparticle energy, not scf mo_energy
        self.mo_energy = None
        self.mo_coeff = mf.mo_coeff
        self.mo_occ = mf.mo_occ
        self.sigma = None

        keys = set(('linearized','ac','fc'))
        self._keys = set(self.__dict__.keys()).union(keys)

    def dump_flags(self):
        log = logger.Logger(self.stdout, self.verbose)
        log.info('')
        log.info('******** %s ********', self.__class__)
        log.info('method = %s', self.__class__.__name__)
        nocca, noccb = self.nocc
        nmoa, nmob = self.nmo
        nvira = nmoa - nocca
        nvirb = nmob - noccb
        nkpts = self.nkpts
        log.info('GW (nocca, noccb) = (%d, %d), (nvira, nvirb) = (%d, %d), nkpts = %d',
                 nocca, noccb, nvira, nvirb, nkpts)
        if self.frozen is not None:
            log.info('frozen orbitals %s', str(self.frozen))
        logger.info(self, 'use perturbative linearized QP eqn = %s', self.linearized)
        logger.info(self, 'analytic continuation method = %s', self.ac)
        logger.info(self, 'GW finite size corrections = %s', self.fc)
        return self

    @property
    def nocc(self):
        return self.get_nocc()
    @nocc.setter
    def nocc(self, n):
        self._nocc = n

    @property
    def nmo(self):
        return self.get_nmo()
    @nmo.setter
    def nmo(self, n):
        self._nmo = n

    get_nocc = get_nocc
    get_nmo = get_nmo
    get_frozen_mask = get_frozen_mask

    def kernel(self, mo_energy=None, mo_coeff=None, orbs=None, kptlist=None, nw=100):
        """
        Input:
            kptlist: self-energy k-points
            orbs: self-energy orbs
            nw: grid number
        Output:
            mo_energy: GW quasiparticle energy
        """
        if mo_coeff is None:
            mo_coeff = np.array(self._scf.mo_coeff)
        if mo_energy is None:
            mo_energy = np.array(self._scf.mo_energy)

        nmoa, nmob = self.nmo
        naux = self.with_df.get_naoaux()
        nkpts = self.nkpts
        mem_incore = (3*nkpts*nmoa**2*naux) * 16/1e6
        mem_now = lib.current_memory()[0]
        if (mem_incore + mem_now > 0.99*self.max_memory):
            logger.warn(self, 'Memory may not be enough!')
            raise NotImplementedError

        cput0 = (logger.process_clock(), logger.perf_counter())
        self.dump_flags()
        self.converged, self.mo_energy, self.mo_coeff = \
                kernel(self, mo_energy, mo_coeff, orbs=orbs,
                       kptlist=kptlist, nw=nw, verbose=self.verbose)

        logger.warn(self, 'GW QP energies may not be sorted from min to max')
        logger.timer(self, 'GW', *cput0)
        return self.mo_energy

if __name__ == '__main__':
    from pyscf.pbc import gto
    from pyscf.pbc.lib import chkfile
    import os
    cell = gto.Cell()
    cell.build(
        unit = 'B',
        a = [[ 0.,          6.74027466,  6.74027466],
             [ 6.74027466,  0.,          6.74027466],
             [ 6.74027466,  6.74027466,  0.        ]],
        atom = '''H 0 0 0
                  H 1.68506866 1.68506866 1.68506866
                  H 3.37013733 3.37013733 3.37013733''',
        basis = 'gth-dzvp',
        pseudo = 'gth-pade',
        verbose = 5,
        charge = 0,
        spin = 1)

    cell.spin = cell.spin * 3
    kpts = cell.make_kpts([3,1,1],scaled_center=[0,0,0])
    gdf = df.GDF(cell, kpts)
    gdf_fname = 'h3_ints_311.h5'
    gdf._cderi_to_save = gdf_fname
    if not os.path.isfile(gdf_fname):
        gdf.build()

    chkfname = 'h_311.chk'
    if os.path.isfile(chkfname):
        kmf = scf.KUHF(cell, kpts, exxdiv=None)
        kmf.with_df = gdf
        kmf.with_df._cderi = gdf_fname
        data = chkfile.load(chkfname, 'scf')
        kmf.__dict__.update(data)
    else:
        kmf = scf.KUHF(cell, kpts, exxdiv=None)
        kmf.with_df = gdf
        kmf.with_df._cderi = gdf_fname
        kmf.conv_tol = 1e-12
        kmf.chkfile = chkfname
        kmf.kernel()

    gw = KUGWAC(kmf)
    gw.linearized = False
    gw.ac = 'pade'
    gw.fc = False
    nocca, noccb = gw.nocc
    gw.kernel(kptlist=[0,1,2],orbs=range(0,nocca+3))
    print(gw.mo_energy)
    assert((abs(gw.mo_energy[0][0][nocca-1]--0.28012813))<1e-5)
    assert((abs(gw.mo_energy[0][0][nocca]-0.13748876))<1e-5)
    assert((abs(gw.mo_energy[0][1][nocca-1]--0.29515851))<1e-5)
    assert((abs(gw.mo_energy[0][1][nocca]-0.14128011))<1e-5)
    assert((abs(gw.mo_energy[1][0][noccb-1]--0.33991721))<1e-5)
    assert((abs(gw.mo_energy[1][0][noccb]-0.10578847))<1e-5)
    assert((abs(gw.mo_energy[1][1][noccb-1]--0.33547973))<1e-5)
    assert((abs(gw.mo_energy[1][1][noccb]-0.08053408))<1e-5)

    gw.fc = True
    nocca, noccb = gw.nocc
    gw.kernel(kptlist=[0,1,2],orbs=range(0,nocca+3))
    print(gw.mo_energy)
    assert((abs(gw.mo_energy[0][0][nocca-1]--0.40244058))<1e-5)
    assert((abs(gw.mo_energy[0][0][nocca]-0.13618348))<1e-5)
    assert((abs(gw.mo_energy[0][1][nocca-1]--0.41743063))<1e-5)
    assert((abs(gw.mo_energy[0][1][nocca]-0.13997427))<1e-5)
    assert((abs(gw.mo_energy[1][0][noccb-1]--0.46133481))<1e-5)
    assert((abs(gw.mo_energy[1][0][noccb]-0.1044926))<1e-5)
    assert((abs(gw.mo_energy[1][1][noccb-1]--0.4568894))<1e-5)
    assert((abs(gw.mo_energy[1][1][noccb]-0.07922511))<1e-5)


# pyscf/pbc/lib/__init__.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from . import kpts_helper


# pyscf/pbc/lib/arnoldi.py

# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
Extension to scipy.linalg module developed for PBC branch.
'''

import numpy as np
import scipy.linalg

def davidson_nosymm(matvec,size,nroots,Adiag=None):
    '''Davidson diagonalization method to solve A c = E c
    when A is not Hermitian.
    '''

    # We don't pass args
    def matvec_args(vec, args):
        return matvec(vec)

    nroots = min(nroots,size)
    #if Adiag == None:
    #   Adiag = matvec(numpy.ones(size))

    # Currently not used:
    x = np.ones((size,1))
    P = np.ones((size,1))

    arnold = Arnoldi(matvec_args, x, P, nroots=nroots)
    return arnold.solve()


VERBOSE = False

class Arnoldi:
    def __init__(self,matr_multiply,xStart,inPreCon,nroots=1,tol=1e-6):
        self.matrMultiply = matr_multiply
        self.size = xStart.shape[0]
        self.nEigen = min(nroots, self.size)
        self.maxM = min(30, self.size)
        self.maxOuterLoop = 10
        self.tol = tol

        #
        #  Creating initial guess and preconditioner
        #
        self.x0 = xStart.real.copy()

        self.iteration = 0
        self.totalIter = 0
        self.converged = False
        self.preCon = inPreCon.copy()
        #
        #  Allocating other vectors
        #
        self.allocateVecs()

    def solve(self):
        while self.converged == 0:
            if self.totalIter == 0:
                self.guessInitial()
            for i in range(self.maxM):
                if self.deflated == 1:
                    self.currentSize = self.nEigen

                if self.deflated == 0 and self.totalIter > 0:
                    self.hMult()
                    self.push_Av()
                    self.constructSubspace()

                self.solveSubspace()
                self.constructSol()
                self.computeResidual()
                self.checkConvergence()
                self.deflated = 0
                if self.converged:
                    break

                self.updateVecs()
                self.checkDeflate()
                self.constructDeflatedSub()

                self.totalIter += 1
                self.currentSize += 1
        print("")
        print("Converged in %3d cycles" % self.totalIter)
        self.constructAllSolV()
        return self.outeigs, self.outevecs

    def allocateVecs(self):
        self.subH = np.zeros( shape=(self.maxM,self.maxM), dtype=complex )
        self.sol = np.zeros( shape=(self.maxM), dtype=complex )
        self.dgks = np.zeros( shape=(self.maxM), dtype=complex )
        self.nConv = np.zeros( shape=(self.maxM), dtype=int )
        self.eigs = np.zeros( shape=(self.maxM), dtype=complex )
        self.evecs = np.zeros( shape=(self.maxM,self.maxM), dtype=complex )
        self.oldeigs = np.zeros( shape=(self.maxM), dtype=complex )
        self.deigs = np.zeros( shape=(self.maxM), dtype=complex )
        self.outeigs = np.zeros( shape=(self.nEigen), dtype=complex )
        self.outevecs = np.zeros( shape=(self.size,self.nEigen), dtype=complex)
        self.currentSize = 0

        self.Ax = np.zeros( shape=(self.size), dtype=complex )
        self.res = np.zeros( shape=(self.size), dtype=complex )
        self.vlist = np.zeros( shape=(self.maxM,self.size), dtype=complex )
        self.cv = np.zeros( shape = (self.size), dtype = complex )
        self.cAv = np.zeros( shape = (self.size), dtype = complex )
        self.Avlist = np.zeros( shape=(self.maxM,self.size), dtype=complex )
        self.dres = 999.9
        self.resnorm = 999.9
        self.cvEig = 0.1
        self.ciEig = 0
        self.deflated = 0

    def guessInitial(self):
        nrm = np.linalg.norm(self.x0)
        self.x0 *= 1./nrm
        self.currentSize = self.nEigen
        for i in range(self.currentSize):
            self.vlist[i] *= 0.0
            self.vlist[i,i] = 1.0 + 0.0*1j
            self.vlist[i] /= np.linalg.norm(self.vlist[i])
        for i in range(self.currentSize):
            self.cv = self.vlist[i].copy()
            self.hMult()
            self.Avlist[i] = self.cAv.copy()
        self.constructSubspace()

    def hMult(self):
        args = 0
        self.cAv = self.matrMultiply(self.cv.reshape(self.size),args)

    def push_Av(self):
        self.Avlist[self.currentSize-1] = self.cAv.reshape(self.size)

    def constructSubspace(self):
        if self.totalIter == 0 or self.deflated == 1: # construct the full block of v^*Av
            for i in range(self.currentSize):
                for j in range(self.currentSize):
                    val = np.vdot(self.vlist[i],self.Avlist[j])
                    self.subH[i,j] = val
        else:
            for j in range(self.currentSize):
                if j <= (self.currentSize-1):
                    val = np.vdot(self.vlist[j],self.Avlist[self.currentSize-1])
                    self.subH[j,self.currentSize-1] = val
                if j < (self.currentSize-1):
                    val = np.vdot(self.vlist[self.currentSize-1],self.Avlist[j])
                    self.subH[self.currentSize-1,j] = val

    def solveSubspace(self):
        w, v = scipy.linalg.eig(self.subH[:self.currentSize,:self.currentSize])
        idx = w.real.argsort()
        #imag_norm = np.linalg.norm(w.imag)
        #if imag_norm > 1e-12:
        #    print " *************************************************** "
        #    print " WARNING  IMAGINARY EIGENVALUE OF NORM %.15g " % (imag_norm)
        #    print " *************************************************** "
        #print "Imaginary norm eigenvectors = ", np.linalg.norm(v.imag)
        #print "Imaginary norm eigenvalue   = ", np.linalg.norm(w.imag)
        v = v[:,idx]
        w = w[idx].real
        self.sol[:self.currentSize] = v[:,self.ciEig]
        self.evecs[:self.currentSize,:self.currentSize] = v
        self.eigs[:self.currentSize] = w[:self.currentSize]
        self.outeigs[:self.nEigen] = w[:self.nEigen]
        self.cvEig = self.eigs[self.ciEig]

    def constructAllSolV(self):
        for i in range(self.nEigen):
            self.sol[:] = self.evecs[:,i]
            self.cv = np.dot(self.vlist[:self.currentSize].transpose(),self.sol[:self.currentSize])
            self.outevecs[:,i] = self.cv

    def constructSol(self):
        self.constructSolV()
        self.constructSolAv()

    def constructSolV(self):
        self.cv = np.dot(self.vlist[:self.currentSize].transpose(),self.sol[:self.currentSize])

    def constructSolAv(self):
        self.cAv = np.dot(self.Avlist[:self.currentSize].transpose(),self.sol[:self.currentSize])

    def computeResidual(self):
        self.res = self.cAv - self.cvEig * self.cv
        self.dres = np.vdot(self.res,self.res)**0.5
        #
        # gram-schmidt for residual vector
        #
        for i in range(self.currentSize):
            self.dgks[i] = np.vdot( self.vlist[i], self.res )
            self.res -= self.dgks[i]*self.vlist[i]
        #
        # second gram-schmidt to make them really orthogonal
        #
        for i in range(self.currentSize):
            self.dgks[i] = np.vdot( self.vlist[i], self.res )
            self.res -= self.dgks[i]*self.vlist[i]
        self.resnorm = np.linalg.norm(self.res)
        self.res /= self.resnorm

        orthog = 0.0
        for i in range(self.currentSize):
            orthog += np.vdot(self.res,self.vlist[i])**2.0
        orthog = orthog ** 0.5
        if not self.deflated:
            if VERBOSE:
                print("%3d %20.14f %20.14f  %10.4g" % (self.ciEig, self.cvEig.real, self.resnorm.real, orthog.real))
        #else:
        #    print "%3d %20.14f %20.14f %20.14f (deflated)" % (self.ciEig, self.cvEig,
        #                                                      self.resnorm, orthog)

        self.iteration += 1

    def updateVecs(self):
        self.vlist[self.currentSize] = self.res.copy()
        self.cv = self.vlist[self.currentSize]

    def checkConvergence(self):
        if self.resnorm < self.tol:
            if VERBOSE:
                print("Eigenvalue %3d converged! (res = %.15g)" % (self.ciEig, self.resnorm))
            self.ciEig += 1
        if self.ciEig == self.nEigen:
            self.converged = True
        if self.resnorm < self.tol and not self.converged:
            if VERBOSE:
                print("")
                print("")
                print("%-3s %-20s %-20s %-8s" % ("#", "  Eigenvalue", "  Res. Norm.", "  Ortho. (should be ~0)"))

    def gramSchmidtCurrentVec(self,northo):
        for i in range(northo):
            self.dgks[i] = np.vdot( self.vlist[i], self.cv )
            self.cv -= self.dgks[i]*self.vlist[i] #/ np.vdot(self.vlist[i],self.vlist[i])
        self.cv /= np.linalg.norm(self.cv)


    def checkDeflate(self):
        if self.currentSize == self.maxM-1:
            self.deflated = 1
            #print "deflating..."
            for i in range(self.nEigen):
                self.sol[:self.currentSize] = self.evecs[:self.currentSize,i]
                # Finds the "best" eigenvector for this eigenvalue
                self.constructSolV()
                # Puts this guess in self.Avlist rather than self.vlist for now...
                # since this would mess up self.constructSolV()'s solution
                self.Avlist[i] = self.cv.copy()
            for i in range(self.nEigen):
                # This is actually the "best" eigenvector v, not A*v (see above)
                self.cv = self.Avlist[i].copy()
                self.gramSchmidtCurrentVec(i)
                self.vlist[i] = self.cv.copy()

            for i in range(self.nEigen):
                # This is actually the "best" eigenvector v, not A*v (see above)
                self.cv = self.vlist[i].copy()
                # Use current vector cv to create cAv
                self.hMult()
                self.Avlist[i] = self.cAv.copy()

    def constructDeflatedSub(self):
        if self.deflated == 1:
            self.currentSize = self.nEigen
            self.constructSubspace()


# pyscf/pbc/lib/chkfile.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

import h5py
import pyscf.pbc.gto
import pyscf.lib.chkfile
from pyscf.lib.chkfile import load_chkfile_key, load  # noqa
from pyscf.lib.chkfile import dump_chkfile_key, dump, save  # noqa

def load_cell(chkfile):
    '''Load Cell object from chkfile.

    Args:
        chkfile : str
            Name of chkfile.

    Returns:
        An (initialized/built) Cell object

    Examples:

    >>> from pyscf.pbc import gto, scf
    >>> cell = gto.Cell()
    >>> cell.build(atom='He 0 0 0')
    >>> scf.chkfile.save_cell(cell, 'He.chk')
    >>> scf.chkfile.load_cell('He.chk')
    <pyscf.pbc.gto.cell.Cell object at 0x7fdcd94d7f50>
    '''
    with h5py.File(chkfile, 'r') as fh5:
        try:
            cell = pyscf.pbc.gto.loads(fh5['mol'][()])
        except Exception:
            from numpy import array  # noqa
            celldic = eval(fh5['mol'][()])
            cell = pyscf.pbc.gto.cell.unpack(celldic)
            cell.build(False, False)
    return cell

dump_cell = save_cell = pyscf.lib.chkfile.save_mol



# pyscf/pbc/lib/kpts_helper.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors: Qiming Sun <osirpt.sun@gmail.com>
#          James D. McClain
#          Timothy Berkelbach <tim.berkelbach@gmail.com>
#

import itertools
from collections import OrderedDict
from numbers import Number
import numpy as np
from pyscf import lib
from pyscf import __config__

KPT_DIFF_TOL = getattr(__config__, 'pbc_lib_kpts_helper_kpt_diff_tol', 1e-6)


def is_zero(kpt):
    return abs(np.asarray(kpt)).sum() < KPT_DIFF_TOL
gamma_point = is_zero

def member(kpt, kpts):
    kpts = np.reshape(kpts, (len(kpts),kpt.size))
    dk = np.einsum('ki->k', abs(kpts-kpt.ravel()))
    return np.where(dk < KPT_DIFF_TOL)[0]

def unique(kpts):
    kpts = np.asarray(kpts)
    try:
        digits = int(-np.log10(KPT_DIFF_TOL))
        uniq_index, uniq_inverse = np.unique(
            kpts.round(digits), return_index=True, return_inverse=True, axis=0)[1:3]
        idx = uniq_index.argsort()
        rank = idx.argsort()
        return kpts[uniq_index[idx]], uniq_index[idx], rank[uniq_inverse]
    except TypeError:
        # Old numpy does not support unique of 2D array
        nkpts = len(kpts)
        uniq_kpts = []
        uniq_index = []
        uniq_inverse = np.zeros(nkpts, dtype=int)
        seen = np.zeros(nkpts, dtype=bool)
        n = 0
        for i, kpt in enumerate(kpts):
            if not seen[i]:
                uniq_kpts.append(kpt)
                uniq_index.append(i)
                idx = abs(kpt-kpts).sum(axis=1) < KPT_DIFF_TOL
                uniq_inverse[idx] = n
                seen[idx] = True
                n += 1
        return np.asarray(uniq_kpts), np.asarray(uniq_index), uniq_inverse

def loop_kkk(nkpts):
    range_nkpts = range(nkpts)
    return itertools.product(range_nkpts, range_nkpts, range_nkpts)

def conj_mapping(cell, kpts):
    '''Find the mapping index: -kpts = kpts[index]'''
    scaled_kpts = cell.get_scaled_kpts(kpts).round(5)
    scaled_kpts = np.modf(scaled_kpts)[0]
    scaled_kpts[scaled_kpts > .5] -= 1
    scaled_kpts[scaled_kpts <= -.5] += 1
    uniq_kpts, index, inverse = unique(scaled_kpts)
    if len(index) != len(inverse):
        raise KPointSymmetryError('duplicated k points')

    minus_kpts = -scaled_kpts
    minus_kpts[minus_kpts == -.5] = .5
    uniq_m_kpts, m_index, m_inverse = unique(minus_kpts)
    lex_order = np.lexsort(uniq_kpts.T)
    m_lex_order = np.lexsort(uniq_m_kpts.T)
    if abs(uniq_kpts[lex_order] - uniq_m_kpts[m_lex_order]).max() > KPT_DIFF_TOL:
        raise KPointSymmetryError('k points not symmetric')

    index[lex_order] = m_lex_order
    return index

def get_kconserv(cell, kpts):
    r'''Get the momentum conservation array for a set of k-points.

    Given k-point indices (k, l, m) the array kconserv[k,l,m] returns
    the index n that satifies momentum conservation,

        (k(k) - k(l) + k(m) - k(n)) \dot a = 2n\pi

    This is used for symmetry e.g. integrals of the form
        [\phi*[k](1) \phi[l](1) | \phi*[m](2) \phi[n](2)]
    are zero unless n satisfies the above.
    '''
    nkpts = kpts.shape[0]
    a = cell.lattice_vectors() / (2*np.pi)

    kconserv = np.zeros((nkpts,nkpts,nkpts), dtype=int)
    kvKLM = kpts[:,None,None,:] - kpts[:,None,:] + kpts
    for N, kvN in enumerate(kpts):
        kvKLMN = np.einsum('wx,klmx->wklm', a, kvKLM - kvN)
        # check whether (1/(2pi) k_{KLMN} dot a) is an integer
        kvKLMN_int = np.rint(kvKLMN)
        mask = np.einsum('wklm->klm', abs(kvKLMN - kvKLMN_int)) < 1e-9
        kconserv[mask] = N
    return kconserv


def check_kpt_antiperm_symmetry(array, idx1, idx2, tolerance=1e-8):
    '''Checks antipermutational symmetry for k-point array.

    Checks whether an array with k-point symmetry has antipermutational symmetry
    with respect to switching the particle indices `idx1`, `idx2`. The particle
    indices switches both the orbital index and k-point index associated with
    the two indices.

    Note:
        One common reason for not obeying antipermutational symmetry in a calculation
        involving FFTs is that the grid to perform the FFT may be too coarse.  This
        symmetry is present in operators in spin-orbital form and 'spin-free'
        operators.

    array (:obj:`ndarray`): array to test permutational symmetry, where for
        an n-particle array, the first (2n-1) array elements are kpoint indices
        while the final 2n array elements are orbital indices.
    idx1 (int): first index
    idx2 (int): second index

    Examples:
        For a 3-particle array, such as the T3 amplitude
            t3[ki, kj, kk, ka, kb, i, j, a, b, c],
        setting `idx1 = 0` and `idx2 = 1` would switch the orbital indices i, j as well
        as the kpoint indices ki, kj.

        >>> nkpts, nocc, nvir = 3, 4, 5
        >>> t2 = numpy.random.random_sample((nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir))
        >>> t2 = t2 + t2.transpose(1,0,2,4,3,5,6)
        >>> check_kpt_antiperm_symmetry(t2, 0, 1)
        True
    '''
    # Checking to make sure bounds of idx1 and idx2 are O.K.
    assert(idx1 >= 0 and idx2 >= 0 and 'indices to swap must be non-negative!')

    array_shape_len = len(array.shape)
    nparticles = (array_shape_len + 1) / 4
    assert(idx1 < (2 * nparticles - 1) and idx2 < (2 * nparticles - 1) and
           'This function does not support the swapping of the last k-point index '
           '(This k-point is implicitly not indexed due to conservation of momentum '
           'between k-points.).')

    if (nparticles > 3):
        raise NotImplementedError('Currently set up for only up to 3 particle '
                                  'arrays. Input array has %d particles.')

    kpt_idx1 = idx1
    kpt_idx2 = idx2

    # Start of the orbital index, located after k-point indices
    orb_idx1 = (2 * nparticles - 1) + idx1
    orb_idx2 = (2 * nparticles - 1) + idx2

    # Sign of permutation
    # sign = (-1)**(abs(idx1 - idx2) + 1)
    out_array_indices = np.arange(array_shape_len)

    out_array_indices[kpt_idx1], out_array_indices[kpt_idx2] = \
            out_array_indices[kpt_idx2], out_array_indices[kpt_idx1]
    out_array_indices[orb_idx1], out_array_indices[orb_idx2] = \
            out_array_indices[orb_idx2], out_array_indices[orb_idx1]
    antisymmetric = (np.linalg.norm(array + array.transpose(out_array_indices)) <
                     tolerance)
    return antisymmetric


def get_kconserv3(cell, kpts, kijkab):
    r'''Get the momentum conservation array for a set of k-points.

    This function is similar to get_kconserv, but instead finds the 'kc'
    that satisfies momentum conservation for 5 k-points,

        (ki + kj + kk - ka - kb - kc) dot a = 2n\pi

    where these kpoints are stored in kijkab[ki, kj, kk, ka, kb].
    '''
    a = cell.lattice_vectors() / (2*np.pi)

    kpts_i, kpts_j, kpts_k, kpts_a, kpts_b = \
            [kpts[x].reshape(-1,3) for x in kijkab]
    shape = [np.size(x) for x in kijkab]
    kconserv = np.zeros(shape, dtype=int)

    kv_kab = kpts_k[:,None,None,:] - kpts_a[:,None,:] - kpts_b
    for i, kpti in enumerate(kpts_i):
        for j, kptj in enumerate(kpts_j):
            kv_ijkab = kv_kab + kpti + kptj
            for c, kptc in enumerate(kpts):
                s = np.einsum('kabx,wx->kabw', kv_ijkab - kptc, a)
                s_int = np.rint(s)
                mask = np.einsum('kabw->kab', abs(s - s_int)) < 1e-9
                kconserv[i,j,mask] = c

    new_shape = [shape[i] for i, x in enumerate(kijkab)
                 if not isinstance(x, (int,np.int))]
    kconserv = kconserv.reshape(new_shape)
    return kconserv


class VectorComposer(object):
    def __init__(self, dtype):
        """
        Composes vectors.
        Args:
            dtype (type): array data type;
        """
        self.__dtype__ = dtype
        self.__transactions__ = []
        self.__total_size__ = 0
        self.__data__ = None

    def put(self, a):
        """
        Puts array into vector.
        Args:
            a (ndarray): array to put;
        """
        if a.dtype != self.__dtype__:
            raise ValueError("dtype mismatch: passed %s vs expected %s" % (a.dtype, self.dtype))
        self.__transactions__.append(a)
        self.__total_size__ += a.size

    def flush(self):
        """
        Composes the vector.
        Returns:
            The composed vector.
        """
        if self.__data__ is None:
            self.__data__ = result = np.empty(self.__total_size__, dtype=self.__dtype__)
            offset = 0
        else:
            offset = self.__data__.size
            self.__data__ = result = np.empty(self.__total_size__ + self.__data__.size, dtype=self.__dtype__)

        for i in self.__transactions__:
            s = i.size
            result[offset:offset + s] = i.reshape(-1)
            offset += s
        self.__transactions__ = []

        return result


class VectorSplitter(object):
    def __init__(self, vector):
        """
        Splits vectors into pieces.
        Args:
            vector (ndarray): vector to split;
        """
        self.__data__ = vector
        self.__offset__ = 0

    def get(self, destination, slc=None):
        """
        Retrieves the next array.
        Args:
            destination: the shape of the destination array or the destination array itself;
            slc: an optional slice;

        Returns:
            The array.
        """
        if isinstance(destination, Number):
            destination = np.zeros((destination,), dtype=self.__data__.dtype)
        elif isinstance(destination, tuple):
            destination = np.zeros(destination, dtype=self.__data__.dtype)
        elif isinstance(destination, np.ndarray):
            pass
        else:
            raise ValueError("Unknown destination: %s" % str(destination))

        if slc is None:
            take_size = np.prod(destination.shape)
            take_shape = destination.shape
        else:
            slc = np.ix_(*slc)
            take_size = destination[slc].size
            take_shape = destination[slc].shape

        avail = self.__data__.size - self.__offset__
        if take_size > avail:
            raise ValueError("Insufficient # of elements: required %d %s, found %d" % (take_size, take_shape, avail))

        if slc is None:
            destination[:] = self.__data__[self.__offset__:self.__offset__ + take_size].reshape(take_shape)
        else:
            destination[slc] = self.__data__[self.__offset__:self.__offset__ + take_size].reshape(take_shape)

        self.__offset__ += take_size
        return destination

    def truncate(self):
        """
        Truncates the data vector.
        """
        self.__data__ = self.__data__[self.__offset__:].copy()
        self.__offset__ = 0


class KptsHelper(lib.StreamObject):
    def __init__(self, cell, kpts):
        '''Helper class for handling k-points in correlated calculations.

        Attributes:
            kconserv : (nkpts,nkpts,nkpts) ndarray
                The index of the fourth momentum-conserving k-point, given
                indices of three k-points
            symm_map : OrderedDict of list of (3,) tuples
                Keys are (3,) tuples of symmetry-unique k-point indices and
                values are lists of (3,) tuples, enumerating all
                symmetry-related k-point indices for ERI generation
        '''
        self.kconserv = get_kconserv(cell, kpts)
        nkpts = len(kpts)
        temp = range(0,nkpts)
        kptlist = lib.cartesian_prod((temp,temp,temp))
        completed = np.zeros((nkpts,nkpts,nkpts), dtype=bool)

        self._operation = np.zeros((nkpts,nkpts,nkpts), dtype=int)
        self.symm_map = OrderedDict()

        for kpt in kptlist:
            kpt = tuple(kpt)
            kp,kq,kr = kpt
            if not completed[kp,kq,kr]:
                self.symm_map[kpt] = list()
                ks = self.kconserv[kp,kq,kr]

                completed[kp,kq,kr] = True
                self._operation[kp,kq,kr] = 0
                self.symm_map[kpt].append((kp,kq,kr))

                completed[kr,ks,kp] = True
                self._operation[kr,ks,kp] = 1 #.transpose(2,3,0,1)
                self.symm_map[kpt].append((kr,ks,kp))

                completed[kq,kp,ks] = True
                self._operation[kq,kp,ks] = 2 #np.conj(.transpose(1,0,3,2))
                self.symm_map[kpt].append((kq,kp,ks))

                completed[ks,kr,kq] = True
                self._operation[ks,kr,kq] = 3 #np.conj(.transpose(3,2,1,0))
                self.symm_map[kpt].append((ks,kr,kq))


    def transform_symm(self, eri_kpt, kp, kq, kr):
        '''Return the symmetry-related ERI at any set of k-points.

        Args:
            eri_kpt : (nmo,nmo,nmo,nmo) ndarray
                An in-cell ERI calculated with a set of symmetry-unique k-points.
            kp, kq, kr : int
                The indices of the k-points at which the ERI is desired.
        '''
        operation = self._operation[kp,kq,kr]
        if operation == 0:
            return eri_kpt
        if operation == 1:
            return eri_kpt.transpose(2,3,0,1)
        if operation == 2:
            return np.conj(eri_kpt.transpose(1,0,3,2))
        if operation == 3:
            return np.conj(eri_kpt.transpose(3,2,1,0))

class KPointSymmetryError(RuntimeError):
    pass


# pyscf/pbc/lib/linalg_helper.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Timothy Berkelbach <tim.berkelbach@gmail.com>
#

'''
Extension to scipy.linalg module developed for PBC branch.
'''

from functools import reduce

import numpy as np
import scipy.linalg

from pyscf.lib import logger

method = 'davidson'
#method = 'arnoldi'

VERBOSE = False

try:
    # Temporary fix for davidson when using mpicc
    from mpi4py import MPI
    FOUND_MPI4PY = True
    MPI_RANK = MPI.COMM_WORLD.Get_rank()
    MPI_COMM = MPI.COMM_WORLD
except (ImportError, OSError):
    FOUND_MPI4PY = False

def eigs(matvec, size, nroots, x0=None, Adiag=None, guess=False, verbose=logger.INFO):
    '''Davidson diagonalization method to solve A c = E c
    when A is not Hermitian.
    '''

    # We don't pass args
    def matvec_args(vec, args=None):
        return matvec(vec)

    nroots = min(nroots, size)

    if method == 'davidson':
        if guess:
            conv, e, c, niter = davidson_guess(matvec, size, nroots, Adiag)
        else:
            conv, e, c, niter = davidson(matvec, size, nroots, x0, Adiag, verbose)
        return conv, e, c
    elif method == 'arnoldi':
        # Currently not used:
        x = np.ones((size,1))
        P = np.ones((size,1))
        arnold = Arnoldi(matvec_args, x, P, nroots=nroots)
        e,c = arnold.solve()
        print("Arnoldi converged in %d iterations"%(arnold.totalIter))
        return e,c
    else:
        david = DavidsonZL()
        david.ndim = size
        david.neig = nroots
        david.diag = matvec(np.ones(size))
        david.matvec = matvec

        return david.solve_iter()


def davidson(mult_by_A, N, neig, x0=None, Adiag=None, verbose=logger.INFO):
    """Diagonalize a matrix via non-symmetric Davidson algorithm.

    mult_by_A() is a function which takes a vector of length N
        and returns a vector of length N.
    neig is the number of eigenvalues requested
    """

    if isinstance(verbose, logger.Logger):
        log = verbose
    else:
        import sys
        log = logger.Logger(sys.stdout, verbose)

    cput1 = (logger.process_clock(), logger.perf_counter())

    Mmin = min(neig,N)
    Mmax = min(N,2000)
    tol = 1e-6

    #Adiagcheck = np.zeros(N,np.complex128)
    #for i in range(N):
    #    test = np.zeros(N,np.complex128)
    #    test[i] = 1.0
    #    Adiagcheck[i] = mult_by_A(test)[i]
    #print "Analytical Adiag == numerical Adiag?", np.allclose(Adiag,Adiagcheck)

    if Adiag is None:
        Adiag = np.zeros(N,np.complex128)
        for i in range(N):
            test = np.zeros(N,np.complex128)
            test[i] = 1.0
            Adiag[i] = mult_by_A(test)[i]

    xi = np.zeros(N,np.complex128)

    lamda_k_old = 0
    lamda_k = 0
    target = 0
    conv = False
    if x0 is not None:
        assert x0.shape == (N, Mmin)
        b = x0.copy()

        Ab = np.zeros((N,Mmin),np.complex128)
        for m in range(Mmin):
            Ab[:,m] = mult_by_A(b[:,m])

    for istep,M in enumerate(range(Mmin,Mmax+1)):
        if M == Mmin:
            # Set of M unit vectors from lowest Adiag (NxM)
            b = np.zeros((N,M))
            idx = Adiag.argsort()
            for m,i in zip(range(M),idx):
                b[i,m] = 1.0
            ## Add random noise and orthogonalize
            #for m in range(M):
            #    b[:,m] += 0.01*np.random.random(N)
            #    b[:,m] /= np.linalg.norm(b[:,m])
            #    b,R = np.linalg.qr(b)

            Ab = np.zeros((N,M),np.complex128)
            for m in range(M):
                Ab[:,m] = mult_by_A(b[:,m])
        else:
            Ab = np.column_stack( (Ab,mult_by_A(b[:,M-1])) )

        Atilde = np.dot(b.conj().transpose(),Ab)
        lamda, alpha = diagonalize_asymm(Atilde)
        lamda_k_old, lamda_k = lamda_k, lamda[target]
        alpha_k = alpha[:,target]

        if M == Mmax:
            break

        q = np.dot( Ab-lamda_k*b, alpha_k )
        log.info('davidson istep = %d  root = %d  E = %.15g  dE = %.9g  residual = %.6g',
                 istep, target, lamda_k.real, (lamda_k - lamda_k_old).real, np.linalg.norm(q))
        cput1 = log.timer('davidson iter', *cput1)
        if np.linalg.norm(q) < tol:
            if target == neig-1:
                conv = True
                break
            else:
                target += 1
        #for i in range(N):
            #eps = 0.
            #if np.allclose(lamda_k,Adiag[i]):
            #    eps = 1e-10
            #xi[i] = q[i]/(lamda_k-Adiag[i]+eps)
        eps = 1e-10
        xi = q/(lamda_k-Adiag+eps)

        # orthonormalize xi wrt b
        bxi,R = np.linalg.qr(np.column_stack((b,xi)))
        if FOUND_MPI4PY:  # Ensure all processes search in same direction
            bxi = MPI_COMM.bcast(bxi)

        # append orthonormalized xi to b
        b = np.column_stack((b,bxi[:,-1]))

    #if M > Mmin and M == Mmax:
    #    print("WARNING: Davidson algorithm reached max basis size "
    #          "M = %d without converging."%(M))

    # Express alpha in original basis
    evecs = np.dot(b,alpha) # b is N x M, alpha is M x M
    return conv, lamda[:neig], evecs[:,:neig], istep


def davidson_guess(mult_by_A,N,neig,Adiag=None):
    """Diagonalize a matrix via non-symmetric Davidson algorithm.

    mult_by_A() is a function which takes a vector of length N
        and returns a vector of length N.
    neig is the number of eigenvalues requested

    Targets the first nroots closest in amplitude to the first
    nroots of the basis. This is useful for the QP peaks in IP/EA.
    """
    Mmax = min(N,2000)
    tol = 1e-6

    #Adiagcheck = np.zeros(N,np.complex128)
    #for i in range(N):
    #    test = np.zeros(N,np.complex128)
    #    test[i] = 1.0
    #    Adiagcheck[i] = mult_by_A(test)[i]
    #print "Analytical Adiag == numerical Adiag?", np.allclose(Adiag,Adiagcheck)

    if Adiag is None:
        Adiag = np.zeros(N,np.complex128)
        for i in range(N):
            test = np.zeros(N,np.complex128)
            test[i] = 1.0
            Adiag[i] = mult_by_A(test)[i]

    xi = np.zeros(N,np.complex128)

    evals = np.zeros(neig,np.complex128)
    evecs = np.zeros((N,neig),np.complex128)

    Mtot = 0
    for guess in range(neig):
        print("Working on guess =", guess+1, "/", neig)
        for M in range(1,Mmax+1):
            if M == 1:
                # Unit vector 'target' as the guess
                b = np.zeros((N,1))
                b[guess,0] = 1.0
                Ab = np.zeros((N,1),np.complex128)
                Ab[:,0] = mult_by_A(b[:,0])
            else:
                Ab = np.column_stack( (Ab,mult_by_A(b[:,M-1])) )

            Atilde = np.dot(b.conj().transpose(),Ab)
            lamda, alpha = diagonalize_asymm(Atilde)

            overlap_guess_j_max = -99
            target = 0
            for j, overlap_guess_j in enumerate(alpha[0,:]):
                if overlap_guess_j > overlap_guess_j_max:
                    overlap_guess_j_max = overlap_guess_j
                    target = j

            lamda_k = lamda[target]
            alpha_k = alpha[:,target]

            if M == Mmax:
                print(" -- M reached Mmax")
                break

            q = np.dot( Ab-lamda_k*b, alpha_k )
            if np.linalg.norm(q) < tol:
                evals[guess] = lamda_k
                evecs[:,guess] = np.dot(b,alpha_k)
                Mtot += M
                print(" -- Converged in", M, "iterations")
                break

            for i in range(N):
                eps = 0.
                if np.allclose(lamda_k,Adiag[i]):
                    eps = 1e-8
                xi[i] = q[i]/(lamda_k-Adiag[i]+eps)

            # orthonormalize xi wrt b
            bxi,R = np.linalg.qr(np.column_stack((b,xi)))
            # append orthonormalized xi to b
            b = np.column_stack((b,bxi[:,-1]))

    if M > 1 and M == Mmax:
        print("WARNING: Davidson algorithm reached max basis size "
              "M = %d without converging."%(M))

    return evals, evecs, Mtot


def diagonalize_asymm(H):
    """
    Diagonalize a real, *asymmetric* matrix and return sorted results.

    Return the eigenvalues and eigenvectors (column matrix)
    sorted from lowest to highest eigenvalue.
    """
    E,C = np.linalg.eig(H)
    #if np.allclose(E.imag, 0*E.imag):
    #    E = np.real(E)
    #else:
    #    print "WARNING: Eigenvalues are complex, will be returned as such."

    idx = E.real.argsort()
    E = E[idx]
    C = C[:,idx]

    return E,C


class Arnoldi(object):
    def __init__(self,matr_multiply,xStart,inPreCon,nroots=1,tol=1e-10):
        self.matrMultiply = matr_multiply
        self.size = xStart.shape[0]
        self.nEigen = min(nroots, self.size)
        self.maxM = min(30, self.size)
        self.maxOuterLoop = 10
        self.tol = tol

        #
        #  Creating initial guess and preconditioner
        #
        self.x0 = xStart.real.copy()

        self.iteration = 0
        self.totalIter = 0
        self.converged = False
        self.preCon = inPreCon.copy()
        #
        #  Allocating other vectors
        #
        self.allocateVecs()

    def solve(self):
        while self.converged == 0:
            if self.totalIter == 0:
                self.guessInitial()
            for i in range(self.maxM):
                if self.deflated == 1:
                    self.currentSize = self.nEigen

                if self.deflated == 0 and self.totalIter > 0:
                    self.hMult()
                    self.push_Av()
                    self.constructSubspace()

                self.solveSubspace()
                self.constructSol()
                self.computeResidual()
                self.checkConvergence()
                self.deflated = 0
                if self.converged:
                    break

                self.updateVecs()
                self.checkDeflate()
                self.constructDeflatedSub()

                self.totalIter += 1
                self.currentSize += 1
        if VERBOSE:
            print("\nConverged in %3d cycles" % self.totalIter)
        self.constructAllSolV()
        return self.outeigs, self.outevecs

    def allocateVecs(self):
        self.subH = np.zeros( shape=(self.maxM,self.maxM), dtype=complex )
        self.sol = np.zeros( shape=(self.maxM), dtype=complex )
        self.dgks = np.zeros( shape=(self.maxM), dtype=complex )
        self.nConv = np.zeros( shape=(self.maxM), dtype=int )
        self.eigs = np.zeros( shape=(self.maxM), dtype=complex )
        self.evecs = np.zeros( shape=(self.maxM,self.maxM), dtype=complex )
        self.oldeigs = np.zeros( shape=(self.maxM), dtype=complex )
        self.deigs = np.zeros( shape=(self.maxM), dtype=complex )
        self.outeigs = np.zeros( shape=(self.nEigen), dtype=complex )
        self.outevecs = np.zeros( shape=(self.size,self.nEigen), dtype=complex)
        self.currentSize = 0

        self.Ax = np.zeros( shape=(self.size), dtype=complex )
        self.res = np.zeros( shape=(self.size), dtype=complex )
        self.vlist = np.zeros( shape=(self.maxM,self.size), dtype=complex )
        self.cv = np.zeros( shape = (self.size), dtype = complex )
        self.cAv = np.zeros( shape = (self.size), dtype = complex )
        self.Avlist = np.zeros( shape=(self.maxM,self.size), dtype=complex )
        self.dres = 999.9
        self.resnorm = 999.9
        self.cvEig = 0.1
        self.ciEig = 0
        self.deflated = 0

    def guessInitial(self):
        nrm = np.linalg.norm(self.x0)
        self.x0 *= 1./nrm
        self.currentSize = self.nEigen
        for i in range(self.currentSize):
            self.vlist[i] *= 0.0
            self.vlist[i,:] += np.random.rand(self.size)
            self.vlist[i,i] /= np.linalg.norm(self.vlist[i,i])
            self.vlist[i,i] *= 12.
            for j in range(i):
                fac = np.vdot( self.vlist[j,:], self.vlist[i,:] )
                self.vlist[i,:] -= fac * self.vlist[j,:]
            self.vlist[i,:] /= np.linalg.norm(self.vlist[i,:])
        for i in range(self.currentSize):
            self.cv = self.vlist[i].copy()
            self.hMult()
            self.Avlist[i] = self.cAv.copy()
        self.constructSubspace()

    def hMult(self):
        args = 0
        self.cAv = self.matrMultiply(self.cv.reshape(self.size),args)

    def push_Av(self):
        self.Avlist[self.currentSize-1] = self.cAv.reshape(self.size)

    def constructSubspace(self):
        if self.totalIter == 0 or self.deflated == 1: # construct the full block of v^*Av
            for i in range(self.currentSize):
                for j in range(self.currentSize):
                    val = np.vdot(self.vlist[i],self.Avlist[j])
                    self.subH[i,j] = val
        else:
            for j in range(self.currentSize):
                if j <= (self.currentSize-1):
                    val = np.vdot(self.vlist[j,:],self.Avlist[self.currentSize-1,:])
                    self.subH[j,self.currentSize-1] = val
                if j < (self.currentSize-1):
                    val = np.vdot(self.vlist[self.currentSize-1,:],self.Avlist[j,:])
                    self.subH[self.currentSize-1,j] = val

    def solveSubspace(self):
        w, v = scipy.linalg.eig(self.subH[:self.currentSize,:self.currentSize])
        idx = w.real.argsort()
        v = v[:,idx]
        w = w[idx].real
#
        imag_norm = np.linalg.norm(w.imag)
        if imag_norm > 1e-12:
            print(" *************************************************** ")
            print(" WARNING  IMAGINARY EIGENVALUE OF NORM %.15g " % (imag_norm))
            print(" *************************************************** ")
        #print "Imaginary norm eigenvectors = ", np.linalg.norm(v.imag)
        #print "Imaginary norm eigenvalue   = ", np.linalg.norm(w.imag)
        #print "eigenvalues = ", w[:min(self.currentSize,7)]
#
        self.sol[:self.currentSize] = v[:,self.ciEig]
        self.evecs[:self.currentSize,:self.currentSize] = v
        self.eigs[:self.currentSize] = w[:self.currentSize]
        self.outeigs[:self.nEigen] = w[:self.nEigen]
        self.cvEig = self.eigs[self.ciEig]

    def constructAllSolV(self):
        for i in range(self.nEigen):
            self.sol[:] = self.evecs[:,i]
            self.cv = np.dot(self.vlist[:self.currentSize].transpose(),self.sol[:self.currentSize])
            self.outevecs[:,i] = self.cv

    def constructSol(self):
        self.constructSolV()
        self.constructSolAv()

    def constructSolV(self):
        self.cv = np.dot(self.vlist[:self.currentSize].transpose(),self.sol[:self.currentSize])

    def constructSolAv(self):
        self.cAv = np.dot(self.Avlist[:self.currentSize].transpose(),self.sol[:self.currentSize])

    def computeResidual(self):
        import sys
        self.res = self.cAv - self.cvEig * self.cv
        self.dres = np.vdot(self.res,self.res)**0.5
        #
        # gram-schmidt for residual vector
        #
        for i in range(self.currentSize):
            self.dgks[i] = np.vdot( self.vlist[i,:], self.res )
            self.res -= self.dgks[i]*self.vlist[i,:]
        #
        # second gram-schmidt to make them really orthogonal
        #
        for i in range(self.currentSize):
            self.dgks[i] = np.vdot( self.vlist[i,:], self.res )
            self.res -= self.dgks[i]*self.vlist[i,:]
        self.resnorm = np.linalg.norm(self.res)
        self.res /= self.resnorm

        orthog = 0.0
        for i in range(self.currentSize):
            orthog += np.vdot(self.res,self.vlist[i])**2.0
            if orthog > 1e-8:
                sys.exit( "Exiting davidson procedure ... orthog = %24.16f" % orthog )
        orthog = orthog ** 0.5
        if not self.deflated:
            if VERBOSE:
                print("%3d %20.14f %20.14f  %10.4g" % (self.ciEig, self.cvEig.real, self.resnorm.real, orthog.real))
        #else:
        #    print "%3d %20.14f %20.14f %20.14f (deflated)" % (self.ciEig, self.cvEig,
        #                                                      self.resnorm, orthog)

        self.iteration += 1

    def updateVecs(self):
        self.vlist[self.currentSize] = self.res.copy()
        self.cv = self.vlist[self.currentSize]

    def checkConvergence(self):
        if self.resnorm < self.tol:
            if VERBOSE:
                print("Eigenvalue %3d converged! (res = %.15g)" % (self.ciEig, self.resnorm))
            self.ciEig += 1
        if self.ciEig == self.nEigen:
            self.converged = True
        if self.resnorm < self.tol and not self.converged:
            if VERBOSE:
                print("")
                print("")
                print("%-3s %-20s %-20s %-8s" % ("#", "  Eigenvalue", "  Res. Norm.", "  Ortho. (should be ~0)"))

    def gramSchmidtCurrentVec(self,northo):
        for k in range(northo):
            fac = np.vdot( self.vlist[k,:], self.cv )
            self.cv -= fac * self.vlist[k,:] #/ np.vdot(self.vlist[i],self.vlist[i])
        cvnorm = np.linalg.norm(self.cv)
        if cvnorm < 1e-4:
            self.cv = np.random.rand(self.size)
            for k in range(northo):
                fac = np.vdot( self.vlist[k,:], self.cv )
                self.cv -= fac * self.vlist[k,:] #/ np.vdot(self.vlist[i],self.vlist[i])
            ########################################################################
            #  Sometimes some of the created vectors are linearly dependent.  i
            #  To get around this I'm not sure what to do other than to throw that
            #  solution out and start at that eigenvalue
            ########################################################################
            print(" ERROR!!!! ... restarting at eigenvalue #%s %s" % (northo, cvnorm))
            self.ciEig = northo
        self.cv /= np.linalg.norm(self.cv)


    def checkDeflate(self):
        if self.currentSize == self.maxM-1:
            self.deflated = 1
            for i in range(self.nEigen):
                self.sol[:self.currentSize] = self.evecs[:self.currentSize,i]
                # Finds the "best" eigenvector for this eigenvalue
                self.constructSolV()
                # Puts this guess in self.Avlist rather than self.vlist for now...
                # since this would mess up self.constructSolV()'s solution
                self.Avlist[i,:] = self.cv.copy()
            for i in range(self.nEigen):
                # This is actually the "best" eigenvector v, not A*v (see above)
                self.cv = self.Avlist[i,:].copy()
                self.gramSchmidtCurrentVec(i)
                self.vlist[i,:] = self.cv.copy()

            orthog = 0.0
            for j in range(self.nEigen):
                for i in range(j):
                    orthog += np.vdot(self.vlist[j,:],self.vlist[i,:])**2.0

            for i in range(self.nEigen):
                # This is actually the "best" eigenvector v, not A*v (see above)
                self.cv = self.vlist[i].copy()
                # Use current vector cv to create cAv
                self.hMult()
                self.Avlist[i] = self.cAv.copy()

    def constructDeflatedSub(self):
        if self.deflated == 1:
            self.currentSize = self.nEigen
            self.constructSubspace()


class DavidsonZL(object):
    def __init__(self):
        self.maxcycle = 200
        self.crit_e = 1.e-7
        self.crit_vec = 1.e-5
        self.crit_demo = 1.e-10
        self.crit_indp = 1.e-10
        # Basic setting
        self.iprt = 1
        self.ndim = 0
        self.neig = 5
        self.matvec = None
        self.v0 = None
        self.diag = None
        self.matrix = None

    def matvecs(self,vlst):
        n = len(vlst)
        wlst = [0]*n
        for i in range(n):
            wlst[i] = self.matvec(vlst[i])
        return wlst

    def genMatrix(self):
        v = np.identity(self.ndim)
        vlst = list(v)
        wlst = self.matvecs(vlst)
        Hmat = np.array(vlst).dot(np.array(wlst).T)
        self.matrix = Hmat
        return Hmat

    def solve_full(self):
        Hmat = self.matrix
        eig,vl,vr = eigGeneral(Hmat)
        return eig,vr

    def genV0(self):
        index = np.argsort(self.diag)[:self.neig]
        self.v0 = [0]*self.neig
        for i in range(self.neig):
            v = np.zeros(self.ndim)
            v[index[i]] = 1.0
            self.v0[i] = v.copy()
        return self.v0

    def solve_iter(self):
        if VERBOSE:
            print('Davidson solver for AX = wX')
            print(' ndim = ', self.ndim)
            print(' neig = ', self.neig)
            print(' maxcycle = ', self.maxcycle)
        #
        # Generate v0
        #
        vlst = self.genV0()
        wlst = self.matvecs(vlst)
        #
        # Begin to solve
        #
        ifconv = False
        neig = self.neig
        iconv = [False]*neig
        eigs = np.zeros(neig)
        ndim = neig
        rlst = []
        for niter in range(self.maxcycle):
            if self.iprt > 0:
                if VERBOSE:
                    print('\n --- niter=',niter,'ndim0=',self.ndim,
                          'ndim=',ndim,'---')

            # List[n,N] -> Max[N,n]
            vbas = np.array(vlst).transpose(1,0)
            wbas = np.array(wlst).transpose(1,0)
            iden = vbas.T.dot(vbas)
            diff = np.linalg.norm(iden-np.identity(ndim))
            if diff > 1.e-10:
                if VERBOSE: print('diff=',diff)
                if VERBOSE: print(iden)
                exit(1)
            tmpH = vbas.T.dot(wbas)
            eig,vl,vr = eigGeneral(tmpH)
            teig = eig[:neig]

            # Eigenvalue convergence
            nconv1 = 0
            for i in range(neig):
                tmp = abs(teig[i]-eigs[i])
                if VERBOSE: print(' i,eold,enew,ediff=',i,eigs[i],teig[i],tmp)
                if tmp <= self.crit_e: nconv1+=1
            if VERBOSE: print(' No. of converged eigval:',nconv1)
            if nconv1 == neig:
                if VERBOSE: print(' Cong: all eignvalues converged ! ')
            eigs = teig.copy()

            # Full Residuals: Res[i]=Res'[i]-w[i]*X[i]
            vr = vr[:,:neig].copy()
            jvec = vbas.dot(vr)
            rbas = wbas.dot(vr) - jvec.dot(np.diag(eigs))
            nconv2 = 0
            for i in range(neig):
                tmp = np.linalg.norm(rbas[:,i])
                if tmp <= self.crit_vec:
                    nconv2 += 1
                    iconv[i] = True
                else:
                    iconv[i] = False
                if VERBOSE: print(' i, norm=', i, tmp, iconv[i])
            if VERBOSE: print(' No. of converged eigvec:', nconv2)
            if nconv2 == neig:
                if VERBOSE: print(' Cong: all eigenvectors converged ! ')

            ifconv = (nconv1 == neig) or (nconv2 == neig)
            if ifconv:
                if VERBOSE: print(' Cong: ALL are converged !\n')
                break

            # Rotated basis to minimal subspace that
            # can give the exact [neig] eigenvalues
            nkeep = ndim #neig
            qbas = realRepresentation(vl,vr,nkeep)
            vbas = vbas.dot(qbas)
            wbas = wbas.dot(qbas)
            vlst = list(vbas.transpose(1,0))
            wlst = list(wbas.transpose(1,0))

            # New directions from residuals
            rlst = []
            for i in range(neig):
                if iconv[i]:
                    continue
                for j in range(self.ndim):
                    tmp = self.diag[j] - eigs[i]
                    if abs(tmp) < self.crit_demo:
                        rbas[j,i] = rbas[j,i]/self.crit_demo
                    else:
                        rbas[j,i] = rbas[j,i]/tmp
                rlst.append(rbas[:,i].real)
                rlst.append(rbas[:,i].imag)

            # Re-orthogonalization and get Nindp
            nindp,vlst2 = mgs_ortho(vlst,rlst,self.crit_indp)

            if nindp != 0:
                wlst2 = self.matvecs(vlst2)
                vlst = vlst + vlst2
                wlst = wlst + wlst2
                ndim = len(vlst)
            else:
                if VERBOSE: print('Convergence Failure: Nindp=0 !')
                exit(1)

        if not ifconv:
            if VERBOSE: print('Convergence Failure: Out of Nmaxcycle !')

        return eigs, jvec


def svd_cut(mat,thresh):
    if len(mat.shape) != 2:
        print("NOT A MATRIX in SVD_CUT !", mat.shape)
        exit(1)
    d1, d2 = mat.shape
    u, sig, v = scipy.linalg.svd(mat, full_matrices=False)
    r = len(sig)
    for i in range(r):
        if sig[i] < thresh*1.01:
            r = i
            break
    # return empty list
    if r == 0: return [[],[],[]]
    bdim = r
    rkep = r
    u2 = np.zeros((d1,bdim))
    s2 = np.zeros((bdim))
    v2 = np.zeros((bdim,d2))
    u2[:,:rkep] = u[:,:rkep]
    s2[:rkep] = sig[:rkep]
    v2[:rkep,:] = v[:rkep,:]
    return u2, s2, v2

def eigGeneral(Hmat):
    n = Hmat.shape[0]
    eig,vl,vr = scipy.linalg.eig(Hmat,left=True)
    order = np.argsort(eig.real)
    for i in range(n-1):
        rdiff = eig[order[i]].real-eig[order[i+1]].real
        idiff = eig[order[i]].imag-eig[order[i+1]].imag
        # swap to a-w,a+w
        if abs(rdiff) < 1.e-14 and idiff>0.0:
            j = i + 1
            torder = order[i]
            order[i] = order[j]
            order[j] = torder
    eig = eig[order]
    vl = vl[:,order]
    vr = vr[:,order]
    # Normalize
    for i in range(n):
        ova = vl[:,i].T.conj().dot(vr[:,i])
        vl[:,i] = vl[:,i]/ova.conj()
        #test
        #ova = vl[:,i].T.conj().dot(vr[:,i])
    # A=RwL^+T
    tmpH = reduce(np.dot,(vr,np.diag(eig),vl.T.conj()))
    diff = np.linalg.norm(tmpH-Hmat)
    if diff > 1.e-8:
        if VERBOSE: print('error: A=R*w*L^+ !',diff)
        #exit(1)
    return eig, vl, vr

def realRepresentation(vl,vr,nred):
    vbas = np.hstack((vl[:,:nred].real,vl[:,:nred].imag,
                      vr[:,:nred].real,vl[:,:nred].imag))
    u,w,v=svd_cut(vbas,thresh=1.e-12)
    vbas =u.copy()
    return vbas

def mgs_ortho(vlst,rlst,crit_indp,iop=0):
    debug = False
    ndim = len(vlst)
    nres = len(rlst)
    maxtimes = 2
    # [N,n]
    vbas = np.array(vlst).transpose(1,0)
    rbas = np.array(rlst).transpose(1,0)
    # res=I-VVt*rlst
    for k in range(maxtimes):
        #rbas = rbas - reduce(np.dot,(vbas,vbas.T,rbas))
        tmp = np.dot(vbas.T,rbas)
        rbas -= np.dot(vbas,tmp)
    nindp = 0
    vlst2 = []
    if iop == 1:
        u,w,v = svd_cut(rbas,thresh=1.e-12)
        nindp = len(w)
        if nindp == 0: return nindp,vlst2
        vbas = np.hstack((vbas,u))
        vlst2 = list(u.transpose(1,0))
    else:
        # orthogonalization
        # - MORE STABLE since SVD sometimes does not converge !!!
        for k in range(maxtimes):
            for i in range(nres):
                rvec = rbas[:,i].copy()
                rii = np.linalg.norm(rbas[:,i])
                if debug: print(' ktime,i,rii=', k, i, rii)
                # TOO SMALL
                if rii <= crit_indp*10.0**(-k):
                    if debug:
                        print(' unable to normalize:', i, ' norm=', rii,
                              ' thresh=', crit_indp)
                    continue
                # NORMALIZE
                rvec = rvec / rii
                rii = np.linalg.norm(rvec)
                rvec = rvec / rii
                nindp = nindp + 1
                vlst2.append(rvec)
                # Substract all things
                # [N,n]
                vbas = np.hstack((vbas,rvec.reshape(-1,1)))
                for j in range(i,nres):
                    #rbas[:,j]=rbas[:,j]-reduce(np.dot,(vbas,vbas.T,rbas[:,j]))
                    tmp = np.dot(vbas.T,rbas[:,j])
                    rbas[:,j] -= np.dot(vbas,tmp)
        # iden
    iden = vbas.T.dot(vbas)
    diff = np.linalg.norm(iden-np.identity(ndim+nindp))
    if diff > 1.e-10:
        if VERBOSE: print(' error in mgs_ortho: diff=', diff)
        if VERBOSE: print(iden)
        raise RuntimeError
    else:
        if VERBOSE: print(' final nindp from mgs_ortho =', nindp,
                          ' diffIden=', diff)
    return nindp, vlst2

if __name__ == '__main__':
    N = 200
    neig = 4
    A = np.zeros((N,N))
    k = N/2
    for ii in range(N):
        i = ii+1
        for jj in range(N):
            j = jj+1
            if j <= k:
                A[ii,jj] = i*(i==j)-(i-j-k**2)
            else:
                A[ii,jj] = i*(i==j)+(i-j-k**2)
    def matvec(x):
        return np.dot(A,x)

    e,c = eigs(matvec,N,neig,Adiag=np.diag(A))
    print("# davidson evals =", e)

    e,c = diagonalize_asymm(A)
    print("# numpy evals =", e.real[:neig])


# pyscf/pbc/mpicc/__init__.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pyscf.pbc.cc import ccsd
from pyscf.pbc import scf

def KRCCSD(mf, frozen=None, mo_coeff=None, mo_occ=None):
    from pyscf.pbc.mpicc import kccsd_rhf
    mf = scf.addons.convert_to_rhf(mf)
    return kccsd_rhf.RCCSD(mf, frozen)

KCCSD = KRCCSD


# pyscf/pbc/mpicc/kccsd_rhf.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors: James D. McClain
#          Timothy Berkelbach <tim.berkelbach@gmail.com>
#

import copy
from functools import reduce

import numpy
import os
import numpy as np
import h5py

from pyscf.pbc.mpicc import mpi_kpoint_helper
from pyscf import lib
import pyscf.ao2mo
from pyscf.lib import logger
import pyscf.cc
import pyscf.cc.ccsd
from pyscf.pbc.mp.kmp2 import get_frozen_mask, get_nocc, get_nmo, padded_mo_coeff, padding_k_idx
from pyscf.pbc.mpicc import kintermediates_rhf as imdk
from pyscf.pbc.lib.linalg_helper import eigs
from pyscf.lib.linalg_helper import eig
from pyscf.pbc.mpitools.mpi_helper import generate_max_task_list, safeAllreduceInPlace, safeNormDiff, safeBcastInPlace
from pyscf.lib.parameters import LOOSE_ZERO_TOL, LARGE_DENOM
from pyscf.lib.numpy_helper import cartesian_prod
from pyscf.pbc.mpitools import mpi_load_balancer, mpi
from pyscf.pbc.tools.tril import tril_index, unpack_tril
from pyscf.pbc.lib import kpts_helper
import pyscf.pbc.cc.kccsd_rhf
from pyscf.pbc.cc.eom_kccsd_rhf_ea import mask_frozen as mask_frozen_ea
from pyscf.pbc.cc.eom_kccsd_rhf_ip import mask_frozen as mask_frozen_ip

from mpi4py import MPI

#einsum = np.einsum
einsum = lib.einsum

rank = MPI.COMM_WORLD.Get_rank()
size = MPI.COMM_WORLD.Get_size()
comm = MPI.COMM_WORLD

def read_amplitudes(t1_shape, t2_shape, t1=None, t2=None, filename="t_amplitudes.hdf5"):
    task_list = generate_max_task_list(t2_shape)
    read_success = False
    if os.path.isfile(filename):
        if t1 is None:
            t1 = np.empty(t1_shape)
        if t2 is None:
            t2 = np.empty(t2_shape)
        print("reading t amplitudes from file")
        feri = h5py.File(filename, 'r', driver='mpio', comm=MPI.COMM_WORLD)
        saved_t1 = feri['t1']
        saved_t2 = feri['t2']
        assert(saved_t1.shape == t1_shape)
        assert(saved_t2.shape == t2_shape)

        task_list = generate_max_task_list(t1.shape)
        for block in task_list:
            which_slice = [slice(*x) for x in block]
            t1[tuple(which_slice)] = saved_t1[tuple(which_slice)]
        task_list = generate_max_task_list(t2.shape)
        for block in task_list:
            which_slice = [slice(*x) for x in block]
            t2[tuple(which_slice)] = saved_t2[tuple(which_slice)]
        feri.close()
        read_success = True
    return read_success, t1, t2

def write_amplitudes(t1, t2, filename="t_amplitudes.hdf5"):
    task_list = generate_max_task_list(t2.shape)
    if rank == 0:
        print("writing t amplitudes to file")
        feri = h5py.File(filename, 'w')
        ds_type = t2.dtype
        out_t1  = feri.create_dataset('t1', t1.shape, dtype=ds_type)
        out_t2  = feri.create_dataset('t2', t2.shape, dtype=ds_type)

        task_list = generate_max_task_list(t1.shape)
        for block in task_list:
            which_slice = [slice(*x) for x in block]
            out_t1[tuple(which_slice)] = t1[tuple(which_slice)]
        task_list = generate_max_task_list(t2.shape)
        for block in task_list:
            which_slice = [slice(*x) for x in block]
            out_t2[tuple(which_slice)] = t2[tuple(which_slice)]
        feri.close()
    return

def read_eom_amplitudes(vec_shape, filename="reom_amplitudes.hdf5", vec=None):
    task_list = generate_max_task_list(vec_shape)
    read_success = False

    # TODO: find a way to make the amplitudes are consistent
    # with the signs of the eris/t-amplitudes when restarting
    return False, None

    print("attempting to read in eom amplitudes from file ", filename)
    if os.path.isfile(filename):
        print("reading eom amplitudes from file. shape=", vec_shape)
        feri = h5py.File(filename, 'r', driver='mpio', comm=MPI.COMM_WORLD)
        saved_v = feri['v']
        if vec is None:
            vec = np.empty(vec_shape,dtype=saved_v.dtype)
        assert(saved_v.shape == vec_shape)
        task_list = generate_max_task_list(vec.shape)
        for block in task_list:
            which_slice = [slice(*x) for x in block]
            vec[tuple(which_slice)] = saved_v[tuple(which_slice)]
        feri.close()
        read_success = True
    if vec is not None and vec_shape[-1] == 1:
        vec = vec.reshape(vec_shape[:-1])
    return read_success, vec

def write_eom_amplitudes(vec, filename="reom_amplitudes.hdf5"):
    task_list = generate_max_task_list(vec.shape)
    if rank == 0:
        print("writing eom amplitudes to file")
        feri = h5py.File(filename, 'w')
        ds_type = vec.dtype
        out_v  = feri.create_dataset('v', vec.shape, dtype=ds_type)

        task_list = generate_max_task_list(vec.shape)
        for block in task_list:
            which_slice = [slice(*x) for x in block]
            out_v[tuple(which_slice)] = vec[tuple(which_slice)]
        feri.close()
    return

def restore_from_diis_(mycc, diis_file, inplace=True):
    '''Reuse an existed DIIS object in the CCSD calculation.
    The CCSD amplitudes will be restored from the DIIS object to generate t1
    and t2 amplitudes. The t1/t2 amplitudes of the CCSD object will be
    overwritten by the generated t1 and t2 amplitudes. The amplitudes vector
    and error vector will be reused in the CCSD calculation.
    '''
    adiis = lib.diis.DIIS(mycc, mycc.diis_file, incore=mycc.incore_complete)
    if rank == 0:
        adiis.restore(diis_file, inplace=inplace)

        ccvec = adiis.extrapolate()
        t1, t2 = mycc.vector_to_amplitudes(ccvec)
    info = None
    if rank == 0:
        info = (t1.shape, t2.shape, np.result_type(t1, t2))
    info = MPI.COMM_WORLD.bcast(info)

    if rank != 0:  # Create empty arrays for master to bcast into
        t1_shape, t2_shape, dtype = info
        t1 = np.empty(t1_shape, dtype=dtype)
        t2 = np.empty(t2_shape, dtype=dtype)
    safeBcastInPlace(MPI.COMM_WORLD, t1)
    safeBcastInPlace(MPI.COMM_WORLD, t2)
    mycc.t1, mycc.t2 = t1, t2
    if inplace:
        mycc.diis = adiis
    return mycc

def get_normt_diff(cc, t1, t2, t1new, t2new):
    '''Calculates norm(t1 - t1new) + norm(t2 - t2new).'''
    normt = safeNormDiff(t1new, t1) + safeNormDiff(t2new, t2)  # Blocking; saves memory
    return normt

def update_t1(cc,t1,t2,eris,ints1e):
    nkpts, nocc, nvir = t1.shape
    fock = eris.fock

    # Get location of padded elements in occupied and virtual space
    nonzero_opadding, nonzero_vpadding = padding_k_idx(cc, kind="split")

    fov = fock[:,:nocc,nocc:]

    Foo,Fvv,Fov,Loo,Lvv = ints1e

    kconserv = cc.kconserv
    t1new = numpy.zeros((nkpts,nocc,nvir),dtype=t1.dtype)

    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(numpy.floor(numpy.sqrt(int(numpy.floor(mem/pre))))),1),nkpts)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts_blksize,))
    loader.set_ranges((range(nkpts),))

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0 = loader.get_blocks_from_data(data)

        s0 = slice(min(ranges0),max(ranges0)+1)

        eris_voov_aXi = _cp(eris.voov[s0,:,s0])
        eris_ovov_Xai = _cp(eris.ovov[:,s0,s0])

        for iterka,ka in enumerate(ranges0):
            ki = ka
            iterki = iterka
            # kc == ki; kk == ka
            t1new[ka] = fov[ka].conj().copy()
            t1new[ka] += -2.*einsum('kc,ka,ic->ia',fov[ki],t1[ka],t1[ki])
            t1new[ka] += einsum('ac,ic->ia',Fvv[ka],t1[ki])
            t1new[ka] += -einsum('ki,ka->ia',Foo[ki],t1[ka])

            tau_term = numpy.empty((nkpts,nocc,nocc,nvir,nvir),dtype=t1.dtype)
            for kk in range(nkpts):
                #tau_term[kk] = 2*t2[kk,ki,kk]
                tau_term[kk] = 2*unpack_tril(t2,nkpts,kk,ki,kk,kconserv[kk,kk,ki])
                #tau_term[kk] -= t2[ki,kk,kk].transpose(1,0,2,3)
                tau_term[kk] -= unpack_tril(t2,nkpts,ki,kk,kk,kconserv[ki,kk,kk]).transpose(1,0,2,3)
            tau_term[ka] += einsum('ic,ka->kica',t1[ki],t1[ka])

            t1new[ka] += einsum('kc,kica->ia',
                                Fov[:].reshape(nocc*nkpts,nvir),
                                tau_term[:].reshape(nocc*nkpts,nocc,nvir,nvir))

            t1new[ka] += einsum('akic,kc->ia',
                                eris_voov_aXi[iterka,:,iterki].transpose(1,0,2,3,4).reshape(nvir,nocc*nkpts,nocc,nvir),
                                2*t1[:].reshape(nocc*nkpts,nvir))
            t1new[ka] += einsum('kaic,kc->ia',
                                eris_ovov_Xai[:,iterka,iterki].reshape(nocc*nkpts,nvir,nocc,nvir),
                                -t1[:].reshape(nocc*nkpts,nvir))
        loader.slave_finished()

    comm.Barrier()

    mem = 0.5e9
    pre = 1.*nocc*nvir*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(numpy.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(numpy.floor(mem/(nkpts_blksize*pre))),1),nkpts)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts_blksize,nkpts_blksize2,))
    loader.set_ranges((range(nkpts),range(nkpts),))

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0,ranges1 = loader.get_blocks_from_data(data)

        s0,s1= [slice(min(x),max(x)+1) for x in (ranges0,ranges1)]

        eris_ovvv_kaX = _cp(eris.ovvv[s1,s0,:])
        eris_ooov_kXi = _cp(eris.ooov[s1,:,s0])
        eris_ooov_Xki = _cp(eris.ooov[:,s1,s0])

        for iterka,ka in enumerate(ranges0):
            ki = ka
            iterki = iterka
            for iterkk,kk in enumerate(ranges1):
                kd_list = _cp(kconserv[ka,range(nkpts),kk])
                kc_list = _cp(range(nkpts))
                Svovv = (2*eris_ovvv_kaX[iterkk,iterka,kd_list].transpose(0,2,1,4,3)
                         - eris_ovvv_kaX[iterkk,iterka,kc_list].transpose(0,2,1,3,4))
                #tau_term_1 = t2[ki,kk,:].copy()
                tau_term_1 = unpack_tril(t2,nkpts,ki,kk,range(nkpts),kconserv[ki,range(nkpts),kk]).copy()
                tau_term_1[ki] += einsum('ic,kd->ikcd',t1[ki],t1[kk])
                t1new[ka] += einsum('ak,ik->ia',Svovv.transpose(1,2,0,3,4).reshape(nvir,-1),
                                                tau_term_1.transpose(1,2,0,3,4).reshape(nocc,-1))

                kl_list = _cp(kconserv[ki,kk,range(nkpts)])
                Sooov = (2*eris_ooov_kXi[iterkk,kl_list,iterki]
                         - eris_ooov_Xki[kl_list,iterkk,iterki].transpose(0,2,1,3,4))
                #tau_term_1 = t2[kk,kl_list,ka].copy()
                tau_term_1 = unpack_tril(t2,nkpts,kk,kl_list,ka,kconserv[kk,ka,kl_list]).copy()
                if kk == ka:
                    tau_term_1[kc_list==kl_list] += einsum('ka,xlc->xklac',t1[ka],t1[kc_list==kl_list])
                t1new[ka] += -einsum('ki,ka->ia',Sooov.transpose(0,1,2,4,3).reshape(-1,nocc),
                                     tau_term_1.transpose(0,1,2,4,3).reshape(-1,nvir))

        loader.slave_finished()

    comm.Allreduce(MPI.IN_PLACE, t1new, op=MPI.SUM)
    return t1new

# This is restricted (R)CCSD
# following Hirata, ..., Barlett, J. Chem. Phys. 120, 2581 (2004)

def update_amps(cc, t1, t2, eris, max_memory=2000):
    time0 = logger.process_clock(), logger.perf_counter()
    log = logger.Logger(cc.stdout, cc.verbose)
    nkpts, nocc, nvir = t1.shape
    tril_shape = ((nkpts)*(nkpts+1))//2

    mo_e_o = [e[:nocc] for e in eris.mo_energy]
    mo_e_v = [e[nocc:] + cc.level_shift for e in eris.mo_energy]

    # Get location of padded elements in occupied and virtual space
    nonzero_opadding, nonzero_vpadding = padding_k_idx(cc, kind="split")
    #t2tmp = numpy.zeros((tril_shape,nkpts,nocc,nocc,nvir,nvir),dtype=t2.dtype)
    #for ki in range(nkpts):
    #    for kj in range(nkpts):
    #        for ka in range(nkpts):
    #            if ki <= kj:
    #                t2tmp[tril_index(ki,kj),ka] = t2[ki,kj,ka]
    #t2 = t2tmp

    #mo_e = eris.fock.diagonal()
    #eia = mo_e[:nocc,None] - mo_e[None,nocc:]
    #eijab = lib.direct_sum('ia,jb->ijab',eia,eia)

    ds_type = t1.dtype

    #Woooo = imdk.cc_Woooo(cc,t1,t2,eris)
    #Wvvvv = imdk.cc_Wvvvv(cc,t1,t2,eris)
    #Wvoov = imdk.cc_Wvoov(cc,t1,t2,eris)
    #Wvovo = imdk.cc_Wvovo(cc,t1,t2,eris)

    Foo = imdk.cc_Foo(cc,t1,t2,eris)
    Fvv = imdk.cc_Fvv(cc,t1,t2,eris)
    Fov = imdk.cc_Fov(cc,t1,t2,eris)
    Loo = imdk.Loo(cc,t1,t2,eris)
    Lvv = imdk.Lvv(cc,t1,t2,eris)

    if rank == 0:
        print("done making intermediates...")
    # Move energy terms to the other side
    for k in range(nkpts):
        Foo[k][np.diag_indices(nocc)] -= mo_e_o[k]
        Fvv[k][np.diag_indices(nvir)] -= mo_e_v[k]
        Loo[k][np.diag_indices(nocc)] -= mo_e_o[k]
        Lvv[k][np.diag_indices(nvir)] -= mo_e_v[k]

    kconserv = cc.kconserv

    if rank == 0:
        print("t1 equation")
    # T1 equation
    t1new = update_t1(cc,t1,t2,eris,[Foo,Fvv,Fov,Loo,Lvv])

    if rank == 0:
        print("t2 equation")
    # T2 equation
    #t2new = numpy.array(eris.oovv, copy=True).conj()
    #t2new = numpy.zeros((nkpts,nkpts,nkpts,nocc,nocc,nvir,nvir),dtype=ds_type)
    t2new_tril = numpy.zeros((tril_shape,nkpts,nocc,nocc,nvir,nvir),dtype=ds_type)

    cput2 = logger.process_clock(), logger.perf_counter()
    loader = mpi_load_balancer.load_balancer(BLKSIZE=(1,nkpts,nkpts,))
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    # Figuring out number of kpoints we can have in our oovv terms below
    # TODO : clean this up- just temporary
    mem = 0.5e9
    pre = 1.*nkpts*nkpts*nocc*nocc*nvir*nvir*16
    nkpts_blksize = max(int(numpy.floor(mem/pre)),1)
    BLKSIZE2 = min(nkpts,nkpts_blksize)
    BLKSIZE2_ranges = [(BLKSIZE2*i,min(nkpts,BLKSIZE2*(i+1)))
                       for i in range(int(numpy.ceil(1.*nkpts/BLKSIZE2)))]

    # Making Woooo terms
    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        if min(ranges0) > max(ranges1): #continue if ki > kj
            loader.slave_finished()
            continue

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_oovv = _cp(eris.oovv[s0,s1,s2])

        eris_oooo = _cp(eris.oooo[s1,s0])
        eris_ovoo_ij = _cp(eris.ovoo[s0,s1])
        eris_ovoo_ji = _cp(eris.ovoo[s1,s0])

        for iterki,ki in enumerate(ranges0):
            for iterkj,kj in enumerate(ranges1):
                if ki <= kj:
                    #t2new[ki,kj,ranges2] += _cp(eris_oovv[iterki,iterkj,s2]).conj()
                    t2new_tril[tril_index(ki,kj),ranges2] += _cp(eris_oovv[iterki,iterkj,s2]).conj()

        for kblock in BLKSIZE2_ranges:
            kl_block_size = kblock[1]-kblock[0]
            # Find out how large of a block_size we need
            kklist = []
            for iterkl,kl in enumerate(range(kblock[0],kblock[1])):
                for iterki,ki in enumerate(ranges0):
                    for iterkj,kj in enumerate(ranges1):
                        kk = kconserv[kj,kl,ki]
                        kklist.append(kk)
                        iterkk = numpy.where(numpy.asarray(kklist)==kk)[0]
                        if len(iterkk)==0: #if not found, append
                            kklist.append(kk)

            kk_block_size = len(kklist)
            eris_oovv1 = numpy.empty((kk_block_size,kl_block_size,nkpts,nocc,nocc,nvir,nvir),
                                     dtype=t2.dtype)
            #  Now fill in the matrix elements
            for iterkl,kl in enumerate(range(kblock[0],kblock[1])):
                for iterki,ki in enumerate(ranges0):
                    for iterkj,kj in enumerate(ranges1):
                        kk = kconserv[kj,kl,ki]
                        iterkk = numpy.where(numpy.asarray(kklist)==kk)[0][0]
                        eris_oovv1[iterkk,iterkl,:] = _cp(eris.oovv[kk,kl,:])

            for iterkl,kl in enumerate(range(kblock[0],kblock[1])):
                for iterki,ki in enumerate(ranges0):
                    for iterkj,kj in enumerate(ranges1):
                        if ki <= kj:
                            kk = kconserv[kj,kl,ki]
                            iterkk = numpy.where(kklist==kk)[0][0]

                            #wOOoo = numpy.empty((nkpts,nocc,nocc,nocc,nocc),dtype=t2.dtype)
                            #tau1_ooVV = t2[ki,kj,:].copy()
                            tau1_ooVV = unpack_tril(t2,nkpts,ki,kj,range(nkpts),kconserv[ki,range(nkpts),kj])
                            tau1_ooVV[ki] += einsum('ic,jd->ijcd',t1[ki],t1[kj])

                            # TODO read only packed oovv terms and unpack after reading
                            wOOoo = _cp(eris_oooo[iterkj,iterki,kl].transpose(3,2,1,0)).conj()
                            wOOoo += einsum('klic,jc->klij',eris_ovoo_ij[iterki,iterkj,kk].transpose(2,3,0,1).conj(),t1[kj])
                            wOOoo += einsum('lkjc,ic->klij',eris_ovoo_ji[iterkj,iterki,kl].transpose(2,3,0,1).conj(),t1[ki])
                            wOOoo += einsum('klcd,ijcd->klij',
                                            eris_oovv1[iterkk,iterkl,:].transpose(1,2,0,3,4).reshape(nocc,nocc,nkpts*nvir,nvir),
                                            tau1_ooVV.transpose(1,2,0,3,4).reshape(nocc,nocc,nkpts*nvir,nvir))

                            for iterka,ka in enumerate(ranges2):
                                # Chemist's notation for momentum conserving t2(ki,kj,ka,kb)
                                kb = kconserv[ki,ka,kj]
                                kn = kconserv[kj,kl,ki]
                                #tau1_OOvv = t2[kn,kl,ka].copy()
                                tau1_OOvv = unpack_tril(t2,nkpts,kn,kl,ka,kconserv[kn,ka,kl])
                                if ka == kk and kl == kb:
                                    tau1_OOvv += einsum('ka,lb->klab',t1[ka],t1[kb])
                                tmp = einsum('klij,klab->ijab',wOOoo,tau1_OOvv) #kl combined into one
                                t2new_tril[tril_index(ki,kj),ka] += tmp
                                #t2new[ki,kj,ka] += tmp

        loader.slave_finished()
    comm.Barrier()
    cput2 = log.timer_debug1('transforming Woooo', *cput2)

    cput2 = logger.process_clock(), logger.perf_counter()

    mem = 0.5e9
    pre = 1.*nvir*nvir*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(numpy.floor(mem/pre)),1),nkpts)
    if rank == 0:
        print("vvvv blocksize = ", nkpts_blksize)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts,1,nkpts_blksize,))
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    # Making Wvvvv terms. Notice the change of for loops.
    def func3():
        good2go = True
        while(good2go):
            good2go, data = loader.slave_set()
            if good2go is False:
                break
            ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
            if min(ranges1) > max(ranges2): #continue if ka > kb
                loader.slave_finished()
                continue

            s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

            eris_ovvv_ab = _cp(eris.ovvv[s1,s2])
            eris_vovv_ab = _cp(eris.vovv[s1,s2])
            eris_vvvv_ab = _cp(eris.vvvv[s1,s2])

            for iterka,ka in enumerate(ranges1):
                for iterkb,kb in enumerate(ranges2):
                    if ka <= kb:
                        # Wvvvv term
                        ovVV = eris_ovvv_ab[iterka,iterkb,:].transpose(1,2,0,3,4).reshape(nocc,nvir,-1)
                        voVV = eris_vovv_ab[iterka,iterkb,:].transpose(1,2,0,3,4).reshape(nvir,nocc,-1)
                        wvvVV = einsum('akd,kb->abd',voVV,-t1[kb])
                        wvvVV += einsum('ak,kbd->abd',-t1[ka].T,ovVV)
                        wvvVV += eris_vvvv_ab[iterka,iterkb].transpose(1,2,0,3,4).reshape(nvir,nvir,-1)
                        wvvVV = wvvVV.transpose(2,0,1)

                        kj_list = kconserv[kb,ranges0,ka]
                        tau1_ooVV = numpy.zeros((len(ranges0),nkpts,nocc,nocc,nvir,nvir),dtype=t2.dtype)
                        for iterki,ki in enumerate(ranges0):
                            kj = kj_list[iterki]
                            #tau1_ooVV[iterki]    += t2[ki,kj,:]
                            tau1_ooVV[iterki] += unpack_tril(t2,nkpts,ki,kj,range(nkpts),kconserv[ki,range(nkpts),kj])
                            tau1_ooVV[iterki,ki] += einsum('ic,jd->ijcd',t1[ki],t1[kj])
                        tau1_ooVV = tau1_ooVV.transpose(0,2,3,1,4,5).reshape(len(ranges0),nocc,nocc,-1)
                        tmp = einsum('kijd,dab->kijab',tau1_ooVV,wvvVV)


                        for iterki,ki in enumerate(ranges0):
                            kj = kj_list[iterki]
                            if ki == kj:
                                t2new_tril[tril_index(ki,kj),ka] += tmp[iterki]
                                #t2new[ki,kj,ka] += tmp[iterki]
                                if ka < kb:
                                    t2new_tril[tril_index(kj,ki),kb] += tmp[iterki].transpose(1,0,3,2)
                                    #t2new[kj,ki,kb] += tmp[iterki].transpose(1,0,3,2)
                            elif ki < kj:
                                t2new_tril[tril_index(ki,kj),ka] += tmp[iterki]
                                #t2new[ki,kj,ka] += tmp[iterki]
                            elif ki > kj:
                                if ka < kb:
                                    t2new_tril[tril_index(kj,ki),kb] += tmp[iterki].transpose(1,0,3,2)
                                    #t2new[kj,ki,kb] += tmp[iterki].transpose(1,0,3,2)
            loader.slave_finished()
    func3()
    comm.Barrier()
    cput2 = log.timer_debug1('transforming Wvvvv', *cput2)

    # Making Wvoov and Wovov terms. (part 1/2)
    cput2 = logger.process_clock(), logger.perf_counter()

    mem = 0.5e9
    pre = 1.*nocc*nvir*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(numpy.floor(mem/pre)),1),nkpts)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=(1,nkpts,nkpts_blksize,))
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        if min(ranges0) > max(ranges1): #continue if ki > kj
            loader.slave_finished()
            continue

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        # TODO this can sometimes not be optimal for ooov, calls for all kb,
        # but in most block set-ups you only need 1 index
        eris_ooov_ji = _cp(eris.ooov[s1,s0])

        eris_voovR1_aXi = _cp(eris.voovR1[s0,s2,:])
        eris_ooovR1_aXi = _cp(eris.ooovR1[s0,s2,:])
        eris_vovvR1_aXi = _cp(eris.vovvR1[s0,s2,:])

        eris_ovovRev_Xai = _cp(eris.ovovRev[s0,s2,:])
        eris_ooovRev_Xai = _cp(eris.ooovRev[s0,s2,:])
        eris_ovvvRev_Xai = _cp(eris.ovvvRev[s0,s2,:])

        for iterki,ki in enumerate(ranges0):
            for iterkj,kj in enumerate(ranges1):
                if ki <= kj:
                    for iterka,ka in enumerate(ranges2):
                        kb = kconserv[ki,ka,kj]
                        # t2 with 1-electron terms (1/2)
                        #tmp = einsum('ac,ijcb->ijab',Lvv[ka],t2[ki,kj,ka])
                        tmp = einsum('ac,ijcb->ijab',Lvv[ka],unpack_tril(t2,nkpts,ki,kj,ka,kconserv[ki,ka,kj]))
                        #tmp += einsum('ki,kjab->ijab',-Loo[ki],t2[ki,kj,ka])
                        tmp += einsum('ki,kjab->ijab',-Loo[ki],unpack_tril(t2,nkpts,ki,kj,ka,kconserv[ki,ka,kj]))
                        # t1 with ooov terms       (1/2)
                        tmp2 = (eris_ooov_ji[iterkj,iterki,kb].transpose(3,2,1,0).conj() +
                                einsum('akic,jc->akij',eris_voovR1_aXi[iterki,iterka,kb],t1[kj])) #ooov[kj,ki,kb,ka] ovvo[kb,ka,kj,ki]
                        tmp -= einsum('akij,kb->ijab',tmp2,t1[kb])
                        if ki == kj:
                            t2new_tril[tril_index(ki,kj),ka] += tmp
                            t2new_tril[tril_index(ki,kj),kb] += tmp.transpose(1,0,3,2)
                            #t2new[ki,kj,ka] += tmp
                            #t2new[ki,kj,kb] += tmp.transpose(1,0,3,2)
                        else:
                            t2new_tril[tril_index(ki,kj),ka] += tmp
                            #t2new[ki,kj,ka] += tmp

        for kblock in BLKSIZE2_ranges:
            kk_block_size = kblock[1]-kblock[0]
            kk_slice = slice(kblock[0],kblock[1])
            kk_range = range(kblock[0],kblock[1])

            oOVv   = numpy.empty((nkpts,kk_block_size,nocc,nocc,nvir,nvir),dtype=t2.dtype)
            oOvV   = numpy.empty((nkpts,kk_block_size,nocc,nocc,nvir,nvir),dtype=t2.dtype)

            for iterki,ki in enumerate(ranges0):
                for iterka,ka in enumerate(ranges2):
                    kc_list = kconserv[kk_slice,ki,ka]
                    # Wvoov term (ka,kk,ki,kc)
                    #    a) the Soovv and oovv contribution to Wvoov is done after the Wovov term, where
                    #        Soovv = 2*oovv[l,k,c,d] - oovv[l,k,d,c]
                    _WvOoV  = _cp(eris_voovR1_aXi[iterki,iterka,kk_slice]).transpose(1,3,0,2,4).reshape(nvir,nocc,-1)                               #voov[ka,*,ki,*]
                    _WvOoV -= einsum('lic,la->aic',eris_ooovR1_aXi[iterki,iterka,kk_slice].transpose(1,3,0,2,4).reshape(nocc,nocc,-1),t1[ka])       #ooov[ka,*,ki,*]
                    _WvOoV += einsum('adc,id->aic',eris_vovvR1_aXi[iterki,iterka,kk_slice].transpose(1,3,0,2,4).reshape(nvir,nvir,-1),t1[ki])       #vovv[ka,*,ki,*]
                    # Wovov term (kk,ka,ki,kc)
                    _WOvoV = _cp(eris_ovovRev_Xai[iterki,iterka,kk_slice]).transpose(2,3,0,1,4).reshape(nvir,nocc,-1)                          #ovov[*,ka,ki,*]
                    _WOvoV -= einsum('lic,la->aic',eris_ooovRev_Xai[iterki,iterka,kk_slice].transpose(2,3,0,1,4).reshape(nocc,nocc,-1),t1[ka]) #ooov[*,ka,ki,*]
                    _WOvoV += einsum('adc,id->aic',eris_ovvvRev_Xai[iterki,iterka,kk_slice].transpose(2,3,0,1,4).reshape(nvir,nvir,-1),t1[ki]) #ovvv[*,ka,ki,*]
                    # Filling in the oovv terms
                    for iterkk,kk in enumerate(kk_range):
                        oOVv[:,iterkk] = _cp(eris.oovv[:,kk,kc_list[iterkk]])
                        oOvV[:,iterkk] = _cp(eris.oovv[kk,:,kc_list[iterkk]])
                    oOVv_f = oOVv.transpose(0,2,5,1,3,4).reshape(nocc*nvir*nkpts,nocc*nvir*kk_block_size)
                    oOvV_f = oOvV.transpose(0,3,5,1,2,4).reshape(nocc*nvir*nkpts,nocc*nvir*kk_block_size)

                    #tau2_OovV  = t2[:,ki,ka].copy()
                    tau2_OovV  = unpack_tril(t2,nkpts,range(nkpts),ki,ka,kconserv[range(nkpts),ka,ki])
                    tau2_OovV[ka] += 2*einsum('id,la->liad',t1[ki],t1[ka])
                    tau2_OovV = tau2_OovV.transpose(2,3,0,1,4).reshape(nocc,nvir,-1)

                    _WvOoV -= 0.5*einsum('dc,iad->aic',oOvV_f,tau2_OovV) # kc consolidated into c, ld consolidated into d
                    _WOvoV -= 0.5*einsum('dc,iad->aic',oOVv_f,tau2_OovV)
                    #_WvOoV += 0.5*einsum('dc,iad->aic',2*oOvV_f-oOVv_f,t2[ki,:,ka].transpose(1,3,0,2,4).reshape(nocc,nvir,-1))
                    _WvOoV += 0.5*einsum('dc,iad->aic', 2*oOvV_f-oOVv_f,
                                          unpack_tril(t2,nkpts,ki,range(nkpts),ka,kconserv[ki,ka,range(nkpts)]).transpose(1,3,0,2,4).reshape(nocc,nvir,-1))

                    for iterkj,kj in enumerate(ranges1):
                        if ki <= kj:
                            kb = kconserv[ki,ka,kj]
                            #tmp = einsum('aic,jbc->ijab',(2*_WvOoV-_WOvoV),t2[kj,kk_slice,kb].transpose(1,3,0,2,4).reshape(nocc,nvir,-1))
                            tmp = einsum('aic,jbc->ijab',2*_WvOoV-_WOvoV,
                                         unpack_tril(t2,nkpts,kj,kk_range,kb,kconserv[kj,kb,kk_range]).transpose(1,3,0,2,4).reshape(nocc,nvir,-1))
                            #tmp -= einsum('aic,jbc->ijab',_WvOoV,t2[kk_slice,kj,kb].transpose(2,3,0,1,4).reshape(nocc,nvir,-1))
                            tmp -= einsum('aic,jbc->ijab',_WvOoV,
                                          unpack_tril(t2,nkpts,kk_range,kj,kb,kconserv[kk_range,kb,kj]).transpose(2,3,0,1,4).reshape(nocc,nvir,-1))
                            if ki == kj:
                                t2new_tril[tril_index(ki,kj),ka] += tmp
                                t2new_tril[tril_index(ki,kj),kb] += tmp.transpose(1,0,3,2)
                                #t2new[ki,kj,ka] += tmp
                                #t2new[ki,kj,kb] += tmp.transpose(1,0,3,2)
                            else:
                                t2new_tril[tril_index(ki,kj),ka] += tmp
                                #t2new[ki,kj,ka] += tmp
                    #kj_ranges = ranges1[ranges1 >= ki]
                    #nkj = kj_ranges.shape[0]
                    #kb_ranges = kconserv[ki,ka,kj_ranges]
                    #t2new[ki,kj_ranges,ka] += einsum('aic,xjbc->xijab',
                    #                                 (2*_WvOoV-_WOvoV),
                    #                                 t2[kj_ranges,kk_slice,kb_ranges].transpose(0,2,4,1,3,5).reshape(nkj,nocc,nvir,-1))
                    #t2new[ki,kj_ranges,ka] -= einsum('aic,xjbc->xijab',_WvOoV,t2[kk_slice,kj_ranges,kb_ranges].transpose(1,3,4,0,2,5).reshape(nkj,nocc,nvir,-1))
        loader.slave_finished()
    comm.Barrier()
    cput2 = log.timer_debug1('transforming Wvoov (ai)', *cput2)

    # Making Wvoov and Wovov terms (part 2/2)

    cput2 = logger.process_clock(), logger.perf_counter()
    loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts,1,nkpts_blksize,))
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        if min(ranges0) >= max(ranges1): #continue if ki >= kj
            loader.slave_finished()
            continue

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

        # TODO this is not optimal for ooov, calls for all ka, but in most block set-ups you only need 1 index
        eris_ooov_ij = _cp(eris.ooov[:,s1])

        eris_voovR1_bXj = _cp(eris.voovR1[s1,s2,:])
        eris_ooovR1_bXj = _cp(eris.ooovR1[s1,s2,:])
        eris_vovvR1_bXj = _cp(eris.vovvR1[s1,s2,:])

        eris_ovovRev_Xbj = _cp(eris.ovovRev[s1,s2,:])
        eris_ooovRev_Xbj = _cp(eris.ooovRev[s1,s2,:])
        eris_ovvvRev_Xbj = _cp(eris.ovvvRev[s1,s2,:])

        for iterki,ki in enumerate(ranges0):
            for iterkj,kj in enumerate(ranges1):
                if ki < kj:
                    for iterkb,kb in enumerate(ranges2):
                        ka = kconserv[ki,kb,kj]
                        # t2 with 1-electron terms (2/2)
                        #tmp = einsum('bc,jica->ijab',Lvv[kb],t2[kj,ki,kb])
                        tmp = einsum('bc,jica->ijab',Lvv[kb],unpack_tril(t2,nkpts,kj,ki,kb,kconserv[kj,kb,ki]))
                        #tmp += einsum('kj,kiba->ijab',-Loo[kj],t2[kj,ki,kb])
                        tmp += einsum('kj,kiba->ijab',-Loo[kj],unpack_tril(t2,nkpts,kj,ki,kb,kconserv[kj,kb,ki]))
                        # t1 with ooov terms (2/2)
                        tmp2 = (eris_ooov_ij[iterki,iterkj,ka].transpose(3,2,1,0).conj() +
                                einsum('bkjc,ic->bkji',eris_voovR1_bXj[iterkj,iterkb,ka],t1[ki])) #ooov[ki,kj,ka,kb] ovvo[ka,kb,ki,kj]
                        tmp -= einsum('bkji,ka->ijab',tmp2,t1[ka])
                        t2new_tril[tril_index(ki,kj),ka] += tmp
                        #t2new[ki,kj,ka] += tmp

        for kblock in BLKSIZE2_ranges:
            kk_block_size = kblock[1]-kblock[0]
            kk_slice = slice(kblock[0],kblock[1])
            kk_range = range(kblock[0],kblock[1])

            oOVv   = numpy.empty((nkpts,kk_block_size,nocc,nocc,nvir,nvir),dtype=t2.dtype)
            oOvV   = numpy.empty((nkpts,kk_block_size,nocc,nocc,nvir,nvir),dtype=t2.dtype)

            for iterkj,kj in enumerate(ranges1):
                for iterkb,kb in enumerate(ranges2):
                    kc_list = kconserv[kk_slice,kj,kb]
                    # Wvoov term (kb,kk,kj,kc)
                    _WvOoV  = _cp(eris_voovR1_bXj[iterkj,iterkb,kk_slice]).transpose(1,3,0,2,4).reshape(nvir,nocc,-1)                          #voov[kb,*,kj,*]
                    _WvOoV -= einsum('ljc,lb->bjc',eris_ooovR1_bXj[iterkj,iterkb,kk_slice].transpose(1,3,0,2,4).reshape(nocc,nocc,-1),t1[kb])  #ooov[kb,*,kj,*]
                    _WvOoV += einsum('bdc,jd->bjc',eris_vovvR1_bXj[iterkj,iterkb,kk_slice].transpose(1,3,0,2,4).reshape(nvir,nvir,-1),t1[kj])  #vovv[kb,*,kj,*]
                    # Wovov term (kk,kb,kj,kc)
                    _WOvoV = _cp(eris_ovovRev_Xbj[iterkj,iterkb,kk_slice]).transpose(2,3,0,1,4).reshape(nvir,nocc,-1)                          #ovov[*,kb,kj,*]
                    _WOvoV -= einsum('ljc,lb->bjc',eris_ooovRev_Xbj[iterkj,iterkb,kk_slice].transpose(2,3,0,1,4).reshape(nocc,nocc,-1),t1[kb]) #ooov[*,kb,kj,*]
                    _WOvoV += einsum('bdc,jd->bjc',eris_ovvvRev_Xbj[iterkj,iterkb,kk_slice].transpose(2,3,0,1,4).reshape(nvir,nvir,-1),t1[kj]) #ovvv[*,kb,kj,*]
                    #
                    # Filling in the oovv terms
                    #
                    for iterkk,kk in enumerate(kk_range):
                        oOVv[:,iterkk] = _cp(eris.oovv[:,kk,kc_list[iterkk]])
                        oOvV[:,iterkk] = _cp(eris.oovv[kk,:,kc_list[iterkk]])
                    oOVv_f = oOVv.transpose(0,2,5,1,3,4).reshape(nocc*nvir*nkpts,nocc*nvir*kk_block_size)
                    oOvV_f = oOvV.transpose(0,3,5,1,2,4).reshape(nocc*nvir*nkpts,nocc*nvir*kk_block_size)

                    #tau2_OovV  = t2[:,kj,kb].copy()
                    tau2_OovV  = unpack_tril(t2,nkpts,range(nkpts),kj,kb,kconserv[range(nkpts),kb,kj])
                    tau2_OovV[kb] += 2*einsum('jd,lb->ljbd',t1[kj],t1[kb])
                    tau2_OovV = tau2_OovV.transpose(2,3,0,1,4).reshape(nocc,nvir,-1)

                    _WvOoV -= 0.5*einsum('dc,jbd->bjc',oOvV_f,tau2_OovV) # kc consolidated into c, ld consolidated into d
                    _WOvoV -= 0.5*einsum('dc,jbd->bjc',oOVv_f,tau2_OovV)
                    #_WvOoV += 0.5*einsum('dc,jbd->bjc',2*oOvV_f-oOVv_f,t2[kj,:,kb].transpose(1,3,0,2,4).reshape(nocc,nvir,-1))
                    _WvOoV += 0.5*einsum('dc,jbd->bjc',2*oOvV_f-oOVv_f,
                                         unpack_tril(t2,nkpts,kj,range(nkpts),kb,kconserv[kj,kb,range(nkpts)]).transpose(1,3,0,2,4).reshape(nocc,nvir,-1))

                    for iterki,ki in enumerate(ranges0):
                        if ki < kj:
                            ka = kconserv[ki,kb,kj]
                            #tmp = einsum('bjc,iac->ijab',(2*_WvOoV-_WOvoV),t2[ki,kk_slice,ka].transpose(1,3,0,2,4).reshape(nocc,nvir,-1))
                            tmp = einsum('bjc,iac->ijab',2*_WvOoV-_WOvoV,
                                         unpack_tril(t2,nkpts,ki,kk_range,ka,kconserv[ki,ka,kk_range]).transpose(1,3,0,2,4).reshape(nocc,nvir,-1))
                            #tmp -= einsum('bjc,iac->ijab',_WvOoV,t2[kk_slice,ki,ka].transpose(2,3,0,1,4).reshape(nocc,nvir,-1))
                            tmp -= einsum('bjc,iac->ijab',_WvOoV,
                                          unpack_tril(t2,nkpts,kk_range,ki,ka,kconserv[kk_range,ka,ki]).transpose(2,3,0,1,4).reshape(nocc,nvir,-1))
                            t2new_tril[tril_index(ki,kj),ka] += tmp
                            #t2new[ki,kj,ka] += tmp
        loader.slave_finished()
    comm.Barrier()
    cput2 = log.timer_debug1('transforming Wvoov (bj)', *cput2)

    cput2 = logger.process_clock(), logger.perf_counter()
    loader = mpi_load_balancer.load_balancer(BLKSIZE=(1,nkpts,nkpts_blksize,))
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    # Making last of the Wovov terms (part 1/2)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        if min(ranges0) > max(ranges1): #continue if ki > kj
            loader.slave_finished()
            continue

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

        eris_ovovRev_Xbi = _cp(eris.ovovRev[s0,s2,:])
        eris_ooovRev_Xbi = _cp(eris.ooovRev[s0,s2,:])
        eris_ovvvRev_Xbi = _cp(eris.ovvvRev[s0,s2,:])

        eris_vovvL1_jib = _cp(eris.vovvL1[s0,s2,:])

        for iterki,ki in enumerate(ranges0):
            for iterkj,kj in enumerate(ranges1):
                if ki <= kj:
                    for iterkb,kb in enumerate(ranges2):
                        ka = kconserv[ki,kb,kj]
                        # t1 with ovvv terms  (part 1/2)
                        tmp2 = (eris_vovvL1_jib[iterki,iterkb,iterkj].transpose(3,2,1,0).conj() -
                                einsum('kbic,ka->abic',eris_ovovRev_Xbi[iterki,iterkb,ka],t1[ka])) #ovvv[ki,kj,ka,kb]  ovov[ka,kb,ki,kj]
                        tmp  = einsum('abic,jc->ijab',tmp2,t1[kj])
                        if ki == kj:
                            t2new_tril[tril_index(ki,kj),ka] += tmp
                            t2new_tril[tril_index(ki,kj),kb] += tmp.transpose(1,0,3,2)
                            #t2new[ki,kj,ka] += tmp
                            #t2new[ki,kj,kb] += tmp.transpose(1,0,3,2)
                        else:
                            t2new_tril[tril_index(ki,kj),ka] += tmp
                            #t2new[ki,kj,ka] += tmp

        for kblock in BLKSIZE2_ranges:
            kk_block_size = kblock[1]-kblock[0]
            kk_slice = slice(kblock[0],kblock[1])
            kk_range = range(kblock[0],kblock[1])

            oOVv   = numpy.empty((nkpts,kk_block_size,nocc,nocc,nvir,nvir),dtype=t2.dtype)

            for iterki,ki in enumerate(ranges0):
                for iterkb,kb in enumerate(ranges2):
                    kc_list = kconserv[kk_slice,ki,kb]
                    # Wovov term (kk,kb,ki,kc)
                    _WOvoV = _cp(eris_ovovRev_Xbi[iterki,iterkb,kk_slice]).transpose(2,3,0,1,4).reshape(nvir,nocc,-1)                          #ovov[*,kb,ki,*]
                    _WOvoV -= einsum('lic,lb->bic',eris_ooovRev_Xbi[iterki,iterkb,kk_slice].transpose(2,3,0,1,4).reshape(nocc,nocc,-1),t1[kb]) #ooov[*,kb,ki,*]
                    _WOvoV += einsum('bdc,id->bic',eris_ovvvRev_Xbi[iterki,iterkb,kk_slice].transpose(2,3,0,1,4).reshape(nvir,nvir,-1),t1[ki]) #ovvv[*,kb,ki,*]
                    #
                    # Filling in the oovv terms
                    #
                    for iterkk,kk in enumerate(kk_range):
                        oOVv[:,iterkk] = _cp(eris.oovv[:,kk,kc_list[iterkk]])
                    oOVv_f = oOVv.transpose(0,2,5,1,3,4).reshape(nocc*nvir*nkpts,nocc*nvir*kk_block_size)

                    #tau2_OovV  = t2[:,ki,kb].copy()
                    tau2_OovV  = unpack_tril(t2,nkpts,range(nkpts),ki,kb,kconserv[range(nkpts),kb,ki])
                    tau2_OovV[kb] += 2*einsum('id,lb->libd',t1[ki],t1[kb])
                    _WOvoV -= 0.5*einsum('dc,ibd->bic',oOVv_f,tau2_OovV.transpose(2,3,0,1,4).reshape(nocc,nvir,-1))

                    for iterkj,kj in enumerate(ranges1):
                        if ki <= kj:
                            ka = kconserv[ki,kb,kj]
                            #tmp = einsum('bic,jac->ijab',_WOvoV,t2[kk_slice,kj,ka].transpose(2,3,0,1,4).reshape(nocc,nvir,-1))
                            tmp = einsum('bic,jac->ijab',_WOvoV,
                                         unpack_tril(t2,nkpts,kk_range,kj,ka,kconserv[kk_range,ka,kj]).transpose(2,3,0,1,4).reshape(nocc,nvir,-1))
                            if ki == kj:
                                t2new_tril[tril_index(ki,kj),ka] -= tmp
                                t2new_tril[tril_index(ki,kj),kb] -= tmp.transpose(1,0,3,2)
                                #t2new[ki,kj,ka] -= tmp
                                #t2new[ki,kj,kb] -= tmp.transpose(1,0,3,2)
                            else:
                                t2new_tril[tril_index(ki,kj),ka] -= tmp
                                #t2new[ki,kj,ka] -= tmp
        loader.slave_finished()
    comm.Barrier()
    cput2 = log.timer_debug1('transforming Wovov (bi)', *cput2)

    cput2 = logger.process_clock(), logger.perf_counter()
    loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts,1,nkpts_blksize,))
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    # Making last of the Wovov terms (part 2/2)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        if min(ranges0) >= max(ranges1): #continue if ki >= kj
            loader.slave_finished()
            continue

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

        eris_ovovRev_Xaj = _cp(eris.ovovRev[s1,s2,:])
        eris_ooovRev_Xaj = _cp(eris.ooovRev[s1,s2,:])
        eris_ovvvRev_Xaj = _cp(eris.ovvvRev[s1,s2,:])

        eris_vovvL1_ija = _cp(eris.vovvL1[s1,s2,:])

        for iterki,ki in enumerate(ranges0):
            for iterkj,kj in enumerate(ranges1):
                if ki < kj:
                    for iterka,ka in enumerate(ranges2):
                        kb = kconserv[ki,ka,kj]
                        # t1 with ovvv terms  (part 2/2)
                        tmp2 = (eris_vovvL1_ija[iterkj,iterka,iterki].transpose(3,2,1,0).conj() -
                                einsum('kajc,kb->bajc',eris_ovovRev_Xaj[iterkj,iterka,kb],t1[kb])) #ovvv[kj,ki,kb,ka]  ovov[kb,ka,kj,ki]
                        tmp  = einsum('bajc,ic->ijab',tmp2,t1[ki])
                        t2new_tril[tril_index(ki,kj),ka] += tmp
                        #t2new[ki,kj,ka] += tmp

        for kblock in BLKSIZE2_ranges:
            kk_block_size = kblock[1]-kblock[0]
            kk_slice = slice(kblock[0],kblock[1])
            kk_range = range(kblock[0],kblock[1])

            oOVv   = numpy.empty((nkpts,kk_block_size,nocc,nocc,nvir,nvir),dtype=t2.dtype)

            for iterkj,kj in enumerate(ranges1):
                for iterka,ka in enumerate(ranges2):
                    kc_list = kconserv[kk_slice,kj,ka]
                    # Wovov term (kk,ka,kj,kc)
                    _WOvoV = _cp(eris_ovovRev_Xaj[iterkj,iterka,kk_slice]).transpose(2,3,0,1,4).reshape(nvir,nocc,-1)                          #ovov[*,ka,kj,*]
                    _WOvoV -= einsum('ljc,la->ajc',eris_ooovRev_Xaj[iterkj,iterka,kk_slice].transpose(2,3,0,1,4).reshape(nocc,nocc,-1),t1[ka]) #ooov[*,ka,kj,*]
                    _WOvoV += einsum('adc,jd->ajc',eris_ovvvRev_Xaj[iterkj,iterka,kk_slice].transpose(2,3,0,1,4).reshape(nvir,nvir,-1),t1[kj]) #ovvv[*,ka,kj,*]
                    #
                    # Filling in the oovv terms
                    #
                    for iterkk,kk in enumerate(kk_range):
                        oOVv[:,iterkk] = _cp(eris.oovv[:,kk,kc_list[iterkk]])
                    oOVv_f = oOVv.transpose(0,2,5,1,3,4).reshape(nocc*nvir*nkpts,nocc*nvir*kk_block_size)

                    #tau2_OovV  = t2[:,kj,ka].copy()
                    tau2_OovV  = unpack_tril(t2,nkpts,range(nkpts),kj,ka,kconserv[range(nkpts),ka,kj])
                    tau2_OovV[ka] += 2*einsum('jd,la->ljad',t1[kj],t1[ka])
                    _WOvoV -= 0.5*einsum('dc,jad->ajc',oOVv_f,tau2_OovV.transpose(2,3,0,1,4).reshape(nocc,nvir,-1))

                    for iterki,ki in enumerate(ranges0):
                        if ki < kj:
                            kb = kconserv[ki,ka,kj]
                            #tmp = einsum('ajc,ibc->ijab',_WOvoV,t2[kk_slice,ki,kb].transpose(2,3,0,1,4).reshape(nocc,nvir,-1))
                            tmp = einsum('ajc,ibc->ijab', _WOvoV,
                                         unpack_tril(t2,nkpts,kk_range,ki,kb,kconserv[kk_range,kb,ki]).transpose(2,3,0,1,4).reshape(nocc,nvir,-1))
                            t2new_tril[tril_index(ki,kj),ka] -= tmp
                            #t2new[ki,kj,ka] -= tmp
        loader.slave_finished()
    comm.Barrier()
    cput2 = log.timer_debug1('transforming Wovov (aj)', *cput2)

    comm.Barrier()
    safeAllreduceInPlace(comm, t2new_tril)

    for ki in range(nkpts):
        ka = ki
        # Remove zero/padded elements from denominator
        eia = LARGE_DENOM * np.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
        n0_ovp_ia = np.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
        eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]
        t1new[ki] /= eia

    for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
        if ki > kj:
            continue

        kb = kconserv[ki, ka, kj]
        # For LARGE_DENOM, see t1new update above
        eia = LARGE_DENOM * np.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
        n0_ovp_ia = np.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
        eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]

        ejb = LARGE_DENOM * np.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
        n0_ovp_jb = np.ix_(nonzero_opadding[kj], nonzero_vpadding[kb])
        ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb]
        eijab = eia[:, None, :, None] + ejb[:, None, :]

        t2new_tril[tril_index(ki,kj),ka] /= eijab

    time0 = log.timer_debug1('update t1 t2', *time0)

    comm.Barrier()
    return t1new, t2new_tril


def energy(cc, t1, t2, eris):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv
    fock = eris.fock
    e = 0.0 + 1j*0.0
    for ki in range(nkpts):
        e += 2*einsum('ia,ia', fock[ki,:nocc,nocc:], t1[ki])
    t1t1 = numpy.zeros(shape=t2.shape,dtype=t2.dtype)
    for ki in range(nkpts):
        ka = ki
        for kj in range(nkpts):
            #kb = kj
            t1t1[ki,kj,ka] = einsum('ia,jb->ijab',t1[ki],t1[kj])
    tau = t2 + t1t1
    for ki in range(nkpts):
        for kj in range(nkpts):
            for ka in range(nkpts):
                kb = kconserv[ki,ka,kj]
                e += einsum('ijab,ijab', 2*tau[ki,kj,ka], eris.oovv[ki,kj,ka])
                e += einsum('ijab,ijba',  -tau[ki,kj,ka], eris.oovv[ki,kj,kb])
    e /= nkpts
    return e.real

def energy_tril(cc, t1, t2, eris):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv
    fock = eris.fock
    e = numpy.array(0.0,dtype=numpy.complex128)
    for ki in range(nkpts):
        e += 2*einsum('ia,ia', fock[ki,:nocc,nocc:], t1[ki])
    t1t1 = numpy.zeros(shape=t2.shape,dtype=t2.dtype)
    for ki in range(nkpts):
        ka = ki
        for kj in range(nkpts):
            if ki <= kj:
                t1t1[tril_index(ki,kj),ka] = einsum('ia,jb->ijab',t1[ki],t1[kj])
    tau = t2 + t1t1
    for ki in range(nkpts):
        for kj in range(nkpts):
            for ka in range(nkpts):
                if ki <= kj:
                    kb = kconserv[ki,ka,kj]
                    e += einsum('ijab,ijab', tau[tril_index(ki,kj),ka],
                                (2.*eris.oovv[ki,kj,ka]-eris.oovv[ki,kj,kb].transpose(0,1,3,2)))
                if kj < ki:
                    kb = kconserv[ki,ka,kj]
                    e += einsum('ijab,ijab', tau[tril_index(kj,ki),kb].transpose(1,0,3,2),
                                (2.*eris.oovv[ki,kj,ka]-eris.oovv[ki,kj,kb].transpose(0,1,3,2)))
    comm.Barrier()
    e /= nkpts
    return e.real

def _update_procs_mf(mf):
    '''Update mean-field objects to be the same on all processors'''
    mf1 = copy.copy(mf)

    mo_coeff  = comm.bcast(mf.mo_coeff, root=0)
    mo_energy = comm.bcast(mf.mo_energy, root=0)
    mo_occ    = comm.bcast(mf.mo_occ, root=0)
    kpts      = comm.bcast(mf.kpts, root=0)

    mf1.mo_coeff = mo_coeff
    mf1.mo_energy = mo_energy
    mf1.mo_occ = mo_occ
    mf1.kpts  = kpts
    comm.Barrier()
    return mf1

class RCCSD(pyscf.pbc.cc.kccsd_rhf.RCCSD):

    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
        mf = _update_procs_mf(mf)
        pyscf.pbc.cc.kccsd_rhf.RCCSD.__init__(self, mf, frozen, mo_coeff, mo_occ)
        self.kconserv = kpts_helper.get_kconserv(mf.cell, mf.kpts)
        self.__imds__ = None

    def _init_amps_tril(self, eris):
        time0 = logger.process_clock(), logger.perf_counter()
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        t1 = numpy.zeros((nkpts,nocc,nvir), dtype=numpy.complex128)
        tril_shape = ((nkpts)*(nkpts+1))//2
        t2_tril = numpy.zeros((tril_shape,nkpts,nocc,nocc,nvir,nvir),dtype=numpy.complex128)
        local_mp2 = numpy.array(0.0,dtype=numpy.complex128)
        self.emp2 = 0
        eia = numpy.zeros((nocc,nvir))
        eijab = numpy.zeros((nocc,nocc,nvir,nvir))

        mo_e_o = [e[:nocc] for e in eris.mo_energy]
        mo_e_v = [e[nocc:] for e in eris.mo_energy]

        # Get location of padded elements in occupied and virtual space
        nonzero_opadding, nonzero_vpadding = padding_k_idx(self, kind="split")

        kconserv = self.kconserv
        loader = mpi_load_balancer.load_balancer(BLKSIZE=(1,1,nkpts,))
        loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

        good2go = True
        while(good2go):
            good2go, data = loader.slave_set()
            if good2go is False:
                break
            ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
            for ki in ranges0:
                for kj in ranges1:
                    if ki <= kj:
                        for ka in ranges2:
                            kb = kconserv[ki,ka,kj]
                            # For discussion of LARGE_DENOM, see t1new update above
                            eia = LARGE_DENOM * np.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
                            n0_ovp_ia = np.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
                            eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]

                            ejb = LARGE_DENOM * np.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
                            n0_ovp_jb = np.ix_(nonzero_opadding[kj], nonzero_vpadding[kb])
                            ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb]
                            eijab = eia[:, None, :, None] + ejb[:, None, :]

                            oovv_ijab = numpy.array(eris.oovv[ki,kj,ka])
                            oovv_ijba = numpy.array(eris.oovv[ki,kj,kb]).transpose(0,1,3,2)
                            woovv = 2.*oovv_ijab - oovv_ijba

                            t2_tril[tril_index(ki,kj),ka] = numpy.conj(oovv_ijab / eijab)
                            local_mp2 += numpy.dot(t2_tril[tril_index(ki,kj),ka].flatten(),woovv.flatten())
                    if kj < ki:
                        for ka in ranges2:
                            kb = kconserv[ki,ka,kj]
                            # For discussion of LARGE_DENOM, see t1new update above
                            eia = LARGE_DENOM * np.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
                            n0_ovp_ia = np.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
                            eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]

                            ejb = LARGE_DENOM * np.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
                            n0_ovp_jb = np.ix_(nonzero_opadding[kj], nonzero_vpadding[kb])
                            ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb]
                            eijab = eia[:, None, :, None] + ejb[:, None, :]

                            idx = abs(eijab) < LOOSE_ZERO_TOL
                            eijab[idx] = LARGE_DENOM

                            oovv_ijab = numpy.array(eris.oovv[ki,kj,ka])
                            oovv_ijba = numpy.array(eris.oovv[ki,kj,kb]).transpose(0,1,3,2)
                            woovv = 2.*oovv_ijab - oovv_ijba

                            tmp = numpy.conj(oovv_ijab / eijab)
                            local_mp2 += numpy.dot(tmp.flatten(),woovv.flatten())
            loader.slave_finished()

        comm.Allreduce(MPI.IN_PLACE, local_mp2, op=MPI.SUM)
        safeAllreduceInPlace(comm, t2_tril)
        self.emp2 = local_mp2.real
        self.emp2 /= nkpts

        if rank == 0:
            logger.info(self, 'Init t2, MP2 energy (with fock eigenvalue shift) = %.15g', self.emp2)
            logger.timer(self, 'init mp2', *time0)
        return self.emp2, t1, t2_tril

    def _init_amps(self, eris):
        time0 = logger.process_clock(), logger.perf_counter()
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        t1 = numpy.zeros((nkpts,nocc,nvir), dtype=numpy.complex128)
        t2 = numpy.zeros((nkpts,nkpts,nkpts,nocc,nocc,nvir,nvir), dtype=numpy.complex128)
        woovv = numpy.empty((nkpts,nkpts,nkpts,nocc,nocc,nvir,nvir), dtype=numpy.complex128)
        self.emp2 = 0
        eris_oovv = eris.oovv.copy()
        eia = numpy.zeros((nocc,nvir))
        eijab = numpy.zeros((nocc,nocc,nvir,nvir))

        mo_e_o = [e[:nocc] for e in eris.mo_energy]
        mo_e_v = [e[nocc:] for e in eris.mo_energy]

        # Get location of padded elements in occupied and virtual space
        nonzero_opadding, nonzero_vpadding = padding_k_idx(self, kind="split")

        kconserv = self.kconserv
        for ki, kj, ka in kpts_helper.loop_kkk(nkpts):
            kb = kconserv[ki,ka,kj]
            # For discussion of LARGE_DENOM, see t1new update above
            eia = LARGE_DENOM * np.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
            n0_ovp_ia = np.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
            eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]

            ejb = LARGE_DENOM * np.ones((nocc, nvir), dtype=eris.mo_energy[0].dtype)
            n0_ovp_jb = np.ix_(nonzero_opadding[kj], nonzero_vpadding[kb])
            ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb]
            eijab = eia[:, None, :, None] + ejb[:, None, :]

            woovv[ki,kj,ka] = (2*eris_oovv[ki,kj,ka] - eris_oovv[ki,kj,kb].transpose(0,1,3,2))
            t2[ki,kj,ka] = eris_oovv[ki,kj,ka] / eijab

        t2 = numpy.conj(t2)
        emp2 = numpy.einsum('pqrijab,pqrijab',t2,woovv).real
        emp2 /= nkpts
        logger.info(self, 'Init t2, MP2 energy = %.15g', self.emp2)
        logger.timer(self, 'init mp2', *time0)
        return self.emp2, t1, t2

    def init_amps(self, eris):
        return self._init_amps_tril(eris)

    get_normt_diff = get_normt_diff

    def energy(self, t1, t2, eris):
        return energy_tril(self, t1, t2, eris)

    def ao2mo(self, mo_coeff=None):
        return _ERIS(self, mo_coeff)

    def update_amps(self, t1, t2, eris, max_memory=2000):
        return update_amps(self, t1, t2, eris, max_memory)

    def ipccsd_diag(self):
        t1,t2 = self.t1, self.t2
        nkpts, nocc, nvir = t1.shape
        kshift = self.kshift
        kconserv = self.kconserv

        if not self.made_ip_imds:
            if not getattr(self, 'imds', None):
                self.imds = _IMDS(self)
            self.imds.make_ip(self)
            self.made_ip_imds = True

        imds = self.imds

        Hr1 = -numpy.diag(imds.Loo[kshift])

        Hr2 = numpy.zeros((nkpts,nkpts,nocc,nocc,nvir),dtype=t1.dtype)
        mem = 0.5e9
        pre = 1.*nocc*nvir*nvir*nvir*nkpts*16
        nkpts_blksize  = min(max(int(numpy.floor(mem/pre)),1),nkpts)
        nkpts_blksize2 = min(max(int(numpy.floor(mem/(nkpts_blksize*pre))),1),nkpts)
        loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts_blksize2,nkpts_blksize,))
        loader.set_ranges((range(nkpts),range(nkpts),))

        good2go = True
        while(good2go):
            good2go, data = loader.slave_set()
            if good2go is False:
                break
            ranges0, ranges1 = loader.get_blocks_from_data(data)

            s0,s1 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1)]

            for iterki,ki in enumerate(ranges0):
                for iterkj,kj in enumerate(ranges1):
                    kb = kconserv[ki,kshift,kj]

                    Woooo_iji = _cp(imds.Woooo[ki,kj,ki])
                    Wvoov_bjj = _cp(imds.Wvoov[kb,kj,kj])
                    Wovov_jbj = _cp(imds.Wovov[kj,kb,kj])
                    Wovov_ibi = _cp(imds.Wovov[ki,kb,ki])

                    for b in range(nvir):
                        Hr2[ki,kj,:,:,b] = imds.Lvv[kb,b,b]
                    for i in range(nocc):
                        Hr2[ki,kj,i,:,:] -= imds.Loo[ki,i,i]
                    for j in range(nocc):
                        Hr2[ki,kj,:,j,:] -= imds.Loo[kj,j,j]
                    for i in range(nocc):
                        for j in range(nocc):
                            Hr2[ki,kj,i,j,:] += Woooo_iji[i,j,i,j]
                    for j in range(nocc):
                        for b in range(nvir):
                            Hr2[ki,kj,:,j,b] += 2.*Wvoov_bjj[b,j,j,b]
                            Hr2[ki,kj,j,j,b] += -Wvoov_bjj[b,j,j,b]
                            Hr2[ki,kj,:,j,b] += -Wovov_jbj[j,b,j,b]

                            Hr2[ki,kj,j,:,b] += -Wovov_ibi[j,b,j,b]

                            for i in range(nocc):
                                kd = kconserv[kj,kshift,ki]
                                Hr2[ki,kj,i,j,b] += -numpy.dot(unpack_tril(t2,nkpts,ki,kj,kshift,kconserv[ki,kshift,kj])[i,j,:,b],
                                                               -2.*imds.Woovv[kj,ki,kd,j,i,b,:] + imds.Woovv[ki,kj,kd,i,j,b,:])

            loader.slave_finished()
        comm.Allreduce(MPI.IN_PLACE, Hr2, op=MPI.SUM)

        vector = self.amplitudes_to_vector_ip(Hr1,Hr2)
        return vector

    def ipccsd(self, nroots=2*4, kptlist=None):
        time0 = logger.process_clock(), logger.perf_counter()
        log = logger.Logger(self.stdout, self.verbose)
        nkpts = self.nkpts
        if kptlist is None:
            kptlist = range(nkpts)
        size = self.vector_size_ip()
        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            nfrozen = np.sum(self.mask_frozen_ip(np.zeros(size, dtype=int), kshift, const=1))
            nroots = min(nroots, size - nfrozen)
        evals = np.zeros((len(kptlist),nroots), np.float)
        evecs = np.zeros((len(kptlist),nroots,size), np.complex128)

        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            diag = self.ipccsd_diag()
            diag = self.mask_frozen_ip(diag, kshift, const=LARGE_DENOM)
            # Initial guess from file
            amplitude_filename = "__ripccsd" + str(kshift) + "__.hdf5"
            rsuccess, x0 = read_eom_amplitudes((nroots,size),filename=amplitude_filename)
            if x0 is not None:
                x0 = x0.T
            #if not rsuccess:
            #    x0 = np.zeros_like(diag)
            #    x0[np.argmin(diag)] = 1.0

            conv, evals_k, evecs_k = eigs(self.ipccsd_matvec, size, nroots, x0=x0, Adiag=diag, verbose=self.verbose)

            evals_k = evals_k.real
            evals[k] = evals_k
            evecs[k] = evecs_k.T

            write_eom_amplitudes(evecs[k],filename=amplitude_filename)
        time0 = log.timer_debug1('converge ip-ccsd', *time0)
        comm.Barrier()
        return evals.real, evecs

    restore_from_diis_ = restore_from_diis_

    def run_diis(self, t1, t2, istep, normt, de, adiis):
        if rank == 0:
            if (adiis and
                istep >= self.diis_start_cycle and
                abs(de) < self.diis_start_energy_diff):
                vec = self.amplitudes_to_vector(t1, t2)
                t1, t2 = self.vector_to_amplitudes(adiis.update(vec))
                logger.debug1(self, 'DIIS for step %d', istep)
        safeBcastInPlace(MPI.COMM_WORLD, t1)
        safeBcastInPlace(MPI.COMM_WORLD, t2)
        return t1, t2

    def ipccsd_matvec(self, vector):
        kshift = self.kshift
        # Ref: Z. Tu, F. Wang, and X. Li, J. Chem. Phys. 136, 174102 (2012) Eqs.(8)-(9)
        vector = self.mask_frozen_ip(vector, kshift, const=0.0)
        r1,r2 = self.vector_to_amplitudes_ip(vector)
        r1 = comm.bcast(r1, root=0)
        r2 = comm.bcast(r2, root=0)

        t1,t2 = self.t1, self.t2
        nkpts,nocc,nvir = self.t1.shape
        nkpts = self.nkpts
        kconserv = self.kconserv

        if not self.made_ip_imds:
            if not getattr(self, 'imds', None):
                self.imds = _IMDS(self)
            self.imds.make_ip(self)
            self.made_ip_imds = True

        imds = self.imds

        Hr1 = numpy.zeros(r1.shape,dtype=t1.dtype)
        loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts,))
        loader.set_ranges((range(nkpts),))

        good2go = True
        while(good2go):
            good2go, data = loader.slave_set()
            if good2go is False:
                break
            ranges0 = loader.get_blocks_from_data(data)

            s0 = slice(min(ranges0),max(ranges0)+1)

            Wooov_Xls = _cp(imds.Wooov[:,s0,kshift])
            Wooov_lXs = _cp(imds.Wooov[s0,:,kshift])

            for iterkl,kl in enumerate(ranges0):
                Hr1 += einsum('ld,ild->i',imds.Fov[kl],2.*r2[kshift,kl]-r2[kl,kshift].transpose(1,0,2))
                Hr1 += einsum('xklid,xkld->i',-2.*Wooov_Xls[:,iterkl]+Wooov_lXs[iterkl,:].transpose(0,2,1,3,4),r2[:,kl])
            loader.slave_finished()
        comm.Allreduce(MPI.IN_PLACE, Hr1, op=MPI.SUM)
        Hr1 -= einsum('ki,k->i',imds.Loo[kshift],r1)

        Hr2 = numpy.zeros(r2.shape,dtype=t1.dtype)
        loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts,1,))
        loader.set_ranges((range(nkpts),range(nkpts),))

        good2go = True
        while(good2go):
            good2go, data = loader.slave_set()
            if good2go is False:
                break
            ranges0, ranges1 = loader.get_blocks_from_data(data)

            s0,s1 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1)]
            Wovoo_sXi  = _cp(imds.Wovoo[kshift,:,s0])
            WooooS_Xij = _cp(imds.WooooS[:,s0,s1])

            tmp = numpy.zeros(nvir,dtype=t2.dtype)
            for kl in range(nkpts):
                kk_list = range(nkpts)
                kd_list = kconserv[kl,kshift,kk_list]
                tmp += einsum('lc,l->c', (2.*imds.Woovv[kl,kk_list,kd_list].transpose(1,0,2,3,4) -
                                          imds.Woovv[kk_list,kl,kd_list].transpose(2,0,1,3,4)).reshape(-1,nvir),
                                         r2[kk_list,kl].transpose(2,0,1,3).reshape(-1))
            for iterki, ki in enumerate(ranges0):
                for iterkj, kj in enumerate(ranges1):
                    Hr2[ki,kj] += -einsum('c,ijcb->ijb',tmp,unpack_tril(t2,nkpts,ki,kj,kshift,kconserv[ki,kshift,kj]))

            for iterki, ki in enumerate(ranges0):
                for iterkj, kj in enumerate(ranges1):
                    kb = kconserv[ki,kshift,kj]
                    Hr2[ki,kj] += einsum('bd,ijd->ijb',imds.Lvv[kb],r2[ki,kj])
                    Hr2[ki,kj] -= einsum('li,ljb->ijb',imds.Loo[ki],r2[ki,kj])
                    Hr2[ki,kj] -= einsum('lj,ilb->ijb',imds.Loo[kj],r2[ki,kj])
                    Hr2[ki,kj] -= einsum('kbij,k->ijb',Wovoo_sXi[kb,iterki],r1)

                    kl_list = range(nkpts)
                    kk_list = kconserv[ki,kl_list,kj]
                    Hr2[ki,kj] += einsum('klij,klb->ijb',
                                         WooooS_Xij[kl_list,iterki,iterkj].transpose(1,0,2,3,4).reshape(nocc,nocc*nkpts,nocc,nocc),
                                         r2[kk_list,kl_list].transpose(1,0,2,3).reshape(nocc,nocc*nkpts,nvir))

            Wovov_Xbi = _cp(imds.Wovov[:,s1,s0])

            for iterki,ki in enumerate(ranges0):
                for iterkb,kb in enumerate(ranges1):
                    kj = kconserv[kshift,ki,kb]
                    Hr2[ki,kj] += -einsum('lbid,ljd->ijb',Wovov_Xbi[:,iterkb,iterki].reshape(nocc*nkpts,nvir,nocc,nvir),
                                                          r2[:,kj].reshape(nocc*nkpts,nocc,nvir))
            Wvoov_bXj = _cp(imds.Wvoov[s1,:,s0])
            Wovov_Xbj = _cp(imds.Wovov[:,s1,s0])

            for iterkj,kj in enumerate(ranges0):
                for iterkb,kb in enumerate(ranges1):
                    ki = kconserv[kshift,kj,kb]
                    Hr2[ki,kj] += einsum('bljd,ild->ijb',
                                         Wvoov_bXj[iterkb,:,iterkj].transpose(1,0,2,3,4).reshape(nvir,nocc*nkpts,nocc,nvir),
                                         (2.*r2[ki,:].transpose(1,0,2,3)-
                                          r2[:,ki].transpose(2,0,1,3)).reshape(nocc,nocc*nkpts,nvir))
                    Hr2[ki,kj] += -einsum('lbjd,ild->ijb',
                                          Wovov_Xbj[:,iterkb,iterkj].reshape(nocc*nkpts,nvir,nocc,nvir),
                                          r2[ki,:].transpose(1,0,2,3).reshape(nocc,nocc*nkpts,nvir))  #typo in nooijen's paper
            loader.slave_finished()
        comm.Allreduce(MPI.IN_PLACE, Hr2, op=MPI.SUM)

        vector = self.amplitudes_to_vector_ip(Hr1,Hr2)
        vector = self.mask_frozen_ip(vector, kshift, const=0.0)
        return vector

    def lipccsd(self, nroots=2*4, kptlist=None):
        time0 = logger.process_clock(), logger.perf_counter()
        log = logger.Logger(self.stdout, self.verbose)
        nkpts = self.nkpts
        if kptlist is None:
            kptlist = range(nkpts)
        size = self.vector_size_ip()
        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            nfrozen = np.sum(self.mask_frozen_ip(np.zeros(size, dtype=int), kshift, const=1))
            nroots = min(nroots, size - nfrozen)
        evals = np.zeros((len(kptlist),nroots), np.float)
        evecs = np.zeros((len(kptlist),nroots,size), np.complex128)

        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            diag = self.ipccsd_diag()
            # Initial guess from file
            amplitude_filename = "__lipccsd" + str(kshift) + "__.hdf5"
            rsuccess, x0 = read_eom_amplitudes((nroots,size),filename=amplitude_filename)
            if x0 is not None:
                x0 = x0.T
            #if not rsuccess:
            #    x0 = np.zeros_like(diag)
            #    x0[np.argmin(diag)] = 1.0

            conv, evals_k, evecs_k = eigs(self.lipccsd_matvec, size, nroots, x0=x0, Adiag=diag, verbose=self.verbose)

            evals_k = evals_k.real
            evals[k] = evals_k
            evecs[k] = evecs_k.T

            write_eom_amplitudes(evecs[k],filename=amplitude_filename)
        time0 = log.timer_debug1('converge ip-ccsd', *time0)
        comm.Barrier()
        return evals.real, evecs

    def lipccsd_matvec(self, vector):
        # Ref: Z. Tu, F. Wang, and X. Li, J. Chem. Phys. 136, 174102 (2012) Eqs.(8)-(9)
        kshift = self.kshift
        vector = self.mask_frozen_ip(vector, kshift, const=0.0)
        r1,r2 = self.vector_to_amplitudes_ip(vector)
        r1 = comm.bcast(r1, root=0)
        r2 = comm.bcast(r2, root=0)

        t1,t2 = self.t1, self.t2
        nkpts,nocc,nvir = self.t1.shape
        nkpts = self.nkpts
        kconserv = self.kconserv

        if not self.made_ip_imds:
            if not getattr(self, 'imds', None):
                self.imds = _IMDS(self)
            self.imds.make_ip(self)
            self.made_ip_imds = True

        imds = self.imds

        Hr1 = numpy.zeros(r1.shape,dtype=t1.dtype)
        Hr2 = numpy.zeros(r2.shape,dtype=t1.dtype)

        def mem_usage_ovoo(nocc, nvir, nkpts):
            return nocc**3 * nvir**1 * 16
        array_size = [nkpts,nkpts]
        #chunk_size = get_max_blocksize_from_mem(0.3e9, mem_usage_ovoo(nocc,nvir,nkpts),
        #                                        array_size, priority_list=[1,1])
        #task_list = generate_task_list(chunk_size,array_size)
        task_list = generate_max_task_list(array_size,blk_mem_size=mem_usage_ovoo(nocc,nvir,nkpts),priority_list=[1,1])

        for kbrange, kirange in mpi.work_stealing_partition(task_list):
            Wovoo_sbi = _cp(imds.Wovoo[kshift,slice(*kbrange),slice(*kirange)])

            for iterkb, kb in enumerate(range(*kbrange)):
                for iterki, ki in enumerate(range(*kirange)):
                    kj = kconserv[kshift,ki,kb]
                    Hr1 -= einsum('kbij,ijb->k',Wovoo_sbi[iterkb,iterki],r2[ki,kj])

        comm.Allreduce(MPI.IN_PLACE, Hr1, op=MPI.SUM)
        Hr1 -= einsum('ki,i->k',imds.Loo[kshift],r1)
        #Hr1 = ( - einsum('ki,i->k',Loo,r1)
        #        - einsum('kbij,ijb->k',Wovoo,r2)
        #        )

        # Using same task_list as before
        for klrange, kkrange in mpi.work_stealing_partition(task_list):
            Wooov_kls = _cp(imds.Wooov[slice(*kkrange),slice(*klrange),kshift])
            Wooov_lks = _cp(imds.Wooov[slice(*klrange),slice(*kkrange),kshift])

            for iterkk, kk in enumerate(range(*kkrange)):
                for iterkl, kl in enumerate(range(*klrange)):
                    kd = kconserv[kk,kshift,kl]
                    Hr2[kk,kl] -= 2.*einsum('klid,i->kld',Wooov_kls[iterkk,iterkl],r1)
                    Hr2[kk,kl] += einsum('lkid,i->kld',Wooov_lks[iterkl,iterkk],r1)
                    Hr2[kk,kl] -= einsum('ki,ild->kld',imds.Loo[kk],r2[kk,kl])
                    Hr2[kk,kl] -= einsum('lj,kjd->kld',imds.Loo[kl],r2[kk,kl])
                    Hr2[kk,kl] += einsum('bd,klb->kld',imds.Lvv[kd],r2[kk,kl])
                    Hr2[kk,kshift] -= (kk==kd)*einsum('kd,l->kld',imds.Fov[kk],r1)
                    Hr2[kshift,kl] += (kl==kd)*2.*einsum('ld,k->kld',imds.Fov[kl],r1)

        def mem_usage_ovvok(nocc, nvir, nkpts):
            return nocc**2 * nvir**2 * nkpts *  16
        array_size = [nkpts,nkpts]
        #chunk_size = get_max_blocksize_from_mem(0.3e9, 2.*mem_usage_ovvok(nocc,nvir,nkpts),
        #                                        array_size, priority_list=[1,1])
        #task_list = generate_task_list(chunk_size,array_size)
        task_list = generate_max_task_list(array_size,blk_mem_size=mem_usage_ovvok(nocc,nvir,nkpts),priority_list=[1,1])

        for kbrange, klrange in mpi.work_stealing_partition(task_list):

            Wvoov_blX = _cp(imds.Wvoov[slice(*kbrange),slice(*klrange),:])
            Wovov_lbX = _cp(imds.Wovov[slice(*klrange),slice(*kbrange),:])

            for iterkb, kb in enumerate(range(*kbrange)):
                for iterkl, kl in enumerate(range(*klrange)):
                    for iterkj, kj in enumerate(range(nkpts)):
                        kd = kconserv[kb,kj,kl]
                        kk = kconserv[kshift,kl,kd]
                        tmp = einsum('bljd,kjb->kld',Wvoov_blX[iterkb,iterkl,kj],r2[kk,kj])
                        Hr2[kk,kl] += 2.*tmp
                        Hr2[kk,kl] -= einsum('lbjd,kjb->kld', Wovov_lbX[iterkl,iterkb,kj],
                                             r2[kk,kj])  # typo in nooijen's paper

                        # Notice we switch around the variable kk and kl
                        kd = kconserv[kl,kj,kb]
                        kk = kconserv[kshift,kl,kd]
                        Hr2[kl,kk] -= einsum('kbjd,jlb->kld',Wovov_lbX[iterkl,iterkb,kj],r2[kj,kk])
                        Hr2[kl,kk] -= einsum('bkjd,ljb->kld',Wvoov_blX[iterkb,iterkl,kj],r2[kk,kj])

        def mem_usage_ovvok(nocc, nvir, nkpts):
            return nocc**2 * nvir**2 * nkpts *  16
        array_size = [nkpts,nkpts]
        #chunk_size = get_max_blocksize_from_mem(0.3e9, 3.*mem_usage_ovvok(nocc,nvir,nkpts),
        #                                        array_size, priority_list=[1,1])
        #task_list = generate_task_list(chunk_size,array_size)
        task_list = generate_max_task_list(array_size,blk_mem_size=mem_usage_ovvok(nocc,nvir,nkpts),priority_list=[1,1])

        # TODO tmp2 only needs to create tmp2[kshift], but should wait for the fix in the mpi.stealing
        # as defined for the analogous quantity in the ipccsd_matvec
        tmp2 = numpy.zeros((nkpts,nvir),dtype=t1.dtype)
        for kirange, kjrange in mpi.work_stealing_partition(task_list):
            for iterki, ki in enumerate(range(*kirange)):
                for iterkj, kj in enumerate(range(*kjrange)):
                    for iterkc, kc in enumerate(range(nkpts)):
                        t2_tmp = unpack_tril(t2,nkpts,ki,kj,kc,kconserv[ki,kc,kj])
                        tmp2[kc] += einsum('ijcb,ijb->c',t2_tmp,r2[ki,kj])
        comm.Allreduce(MPI.IN_PLACE, tmp2, op=MPI.SUM)

        for kkrange, klrange in mpi.work_stealing_partition(task_list):

            Woooo_klX = _cp(imds.Woooo[slice(*kkrange),slice(*klrange),:])
            Woovv_klX = _cp(imds.Woovv[slice(*kkrange),slice(*klrange),:])

            for iterkk, kk in enumerate(range(*kkrange)):
                for iterkl, kl in enumerate(range(*klrange)):
                    for iterki, ki in enumerate(range(nkpts)):
                        kj = kconserv[kk,ki,kl]
                        Hr2[kk,kl] += einsum('klij,ijd->kld',Woooo_klX[iterkk,iterkl,ki],r2[ki,kj])
                    kd = kconserv[kk,kshift,kl]
                    tmp3 = einsum('kldc,c->kld',Woovv_klX[iterkk,iterkl,kd],tmp2[kshift])
                    Hr2[kk,kl] +=    tmp3
                    Hr2[kl,kk] -= 2.*tmp3.transpose(1,0,2) # Notice change of kl,kk in Hr2

        #tmp = einsum('ijcb,ijb->c',t2,r2)
        #Hr2 = ( - einsum('kd,l->kld',Fov,r1)
        #        + 2.*einsum('ld,k->kld',Fov,r1)
        #        - 2.*einsum('klid,i->kld',Wooov,r1)
        #        + einsum('lkid,i->kld',Wooov,r1)
        #        - einsum('ki,ild->kld',Loo,r2)
        #        - einsum('lj,kjd->kld',Loo,r2)
        #        + einsum('bd,klb->kld',Lvv,r2)
        #        + 2.*einsum('lbdj,kjb->kld',Wovvo,r2)
        #        - einsum('kbdj,ljb->kld',Wovvo,r2)
        #        - einsum('lbjd,kjb->kld',Wovov,r2) #typo in nooijen's paper
        #        + einsum('klij,ijd->kld',Woooo,r2)
        #        - einsum('kbid,ilb->kld',Wovov,r2)
        #        + einsum('kldc,c->kld',Woovv,tmp)
        #        - 2.*einsum('lkdc,c->kld',Woovv,tmp)
        #        )
        comm.Allreduce(MPI.IN_PLACE, Hr2, op=MPI.SUM)

        vector = self.amplitudes_to_vector_ip(Hr1,Hr2)
        vector = self.mask_frozen_ip(vector, kshift, const=0.0)
        return vector

    def vector_size_ip(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts

        size = nocc + nkpts ** 2 * nocc ** 2 * nvir
        return size

    mask_frozen_ip = mask_frozen_ip

    def vector_to_amplitudes_ip(self,vector):
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts

        r1 = vector[:nocc].copy()
        r2 = vector[nocc:].copy().reshape(nkpts,nkpts,nocc,nocc,nvir)
        #r2 = np.zeros((nkpts,nkpts,nocc,nocc,nvir), vector.dtype)
        #index = nocc
        #for ki in range(nkpts):
        #    for kj in range(nkpts):
        #        for i in range(nocc):
        #            for j in range(nocc):
        #                for a in range(nvir):
        #                    r2[ki,kj,i,j,a] =  vector[index]
        #                    index += 1
        return [r1,r2]

    def amplitudes_to_vector_ip(self,r1,r2):
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        size = nocc + nkpts*nkpts*nocc*nocc*nvir

        vector = np.zeros((size), r1.dtype)
        vector[:nocc] = r1.copy()
        vector[nocc:] = r2.copy().reshape(nkpts*nkpts*nocc*nocc*nvir)
        #index = nocc
        #for ki in range(nkpts):
        #    for kj in range(nkpts):
        #        for i in range(nocc):
        #            for j in range(nocc):
        #                for a in range(nvir):
        #                    vector[index] = r2[ki,kj,i,j,a]
        #                    index += 1
        return vector

    def ipccsd_star_contract(self, ipccsd_evals, ipccsd_evecs, lipccsd_evecs):
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        kconserv = self.kconserv
        eris = self.eris
        kshift = self.kshift
        t2 = self.t2
        foo = eris.fock[:,:nocc,:nocc]
        fvv = eris.fock[:,nocc:,nocc:]

        e = []
        assert len(ipccsd_evecs.shape) == 2  # Done at a single k-point, kshift
        for _eval, _evec, _levec in zip(ipccsd_evals, ipccsd_evecs, lipccsd_evecs):
            l1,l2 = self.vector_to_amplitudes_ip(_levec)
            r1,r2 = self.vector_to_amplitudes_ip(_evec)

            # Transposing the l2 operator

            l2_t = numpy.zeros_like(l2)
            for ki in range(nkpts):
                for kj in range(nkpts):
                    ka = kconserv[ki,kshift,kj]
                    l2_t[ki,kj] = l2[kj,ki].transpose(1,0,2)

            # Normalization constant for left-hand eigenvector
            ldotr = numpy.dot(l1.ravel(),r1.ravel()) + numpy.dot(l2.ravel(),r2.ravel())
            l2 = (l2 + 2.*l2_t)/3.
            l2_t = None

            l1 /= ldotr
            l2 /= ldotr
            l1 = l1.reshape(-1)
            r1 = r1.reshape(-1)

            lijkab_tmp = numpy.zeros((nocc,nocc,nocc,nvir,nvir),dtype=t2.dtype)
            rijkab_tmp = numpy.zeros((nocc,nocc,nocc,nvir,nvir),dtype=t2.dtype)
            def mem_usage_ovvvkk(nocc, nvir, nkpts):
                return nocc * nvir**3 * nkpts * 16
            array_size = [nkpts,nkpts]
            # TODO: figure out a good chunk size
            #task_list = generate_max_task_list(array_size,priority_list=[1,1])
            task_list = generate_max_task_list(array_size,blk_mem_size=1e9*2.*mem_usage_ovvvkk(nocc,nvir,nkpts),priority_list=[1,1])

            star_energy = numpy.array(0.0 + 1j*0.0)
            for kirange, kjrange in mpi.work_stealing_partition(task_list):
                oovv_ijX = _cp(eris.oovv[slice(*kirange),slice(*kjrange),range(nkpts)])
                oovv_jiX = _cp(eris.oovv[slice(*kjrange),slice(*kirange),range(nkpts)])

                kklist = kpts_helper.get_kconserv3(self._scf.cell, self.kpts, [range(nkpts),range(nkpts),kshift,range(*kirange),range(*kjrange)])

                for iterki, ki in enumerate(range(*kirange)):
                    for iterkj, kj in enumerate(range(*kjrange)):
                        for iterka, ka in enumerate(range(nkpts)):

                            # Starting the left amplitude equations

                            for iterkb, kb in enumerate(range(nkpts)):
                                if ka > kb:
                                    continue
                                fac = (2.0 - float(ka==kb))

                                if ka == kb:
                                    if ki > kj:
                                        continue
                                    else:
                                        fac *= (2.0 - float(ki==kj))

                                lijkab_tmp *= 0.0
                                rijkab_tmp *= 0.0
                                kk = kklist[ka,kb,iterki,iterkj]

                                # (i,j,k) -> (i,j,k)
                                if kk == kshift and kb == kconserv[ki,ka,kj]:
                                    tmp = 0.5*einsum('ijab,k->ijkab',oovv_ijX[iterki,iterkj,iterka],l1)
                                    lijkab_tmp += 8.* tmp

                                if kk == kshift and kb == kconserv[ki,ka,kj]:
                                    # (j,i,k) -> (i,j,k)
                                    tmp = 0.5*einsum('jiab,k->ijkab',oovv_jiX[iterkj,iterki,ka],l1)
                                    lijkab_tmp -= 4.* tmp

                                # (k,j,i) -> (i,j,k)
                                if ki == kshift and kb == kconserv[kj,ka,kk]:
                                    tmp = 0.5*einsum('kjab,i->ijkab',eris.oovv[kk,kj,ka],l1)
                                    lijkab_tmp -= 4.* tmp

                                # (i,k,j) -> (i,j,k)
                                if kj == kshift and kb == kconserv[ki,ka,kk]:
                                    tmp = 0.5*einsum('ikab,j->ijkab',eris.oovv[ki,kk,ka],l1)
                                    lijkab_tmp -= 4.* tmp

                                # (j,k,i) -> (i,j,k)
                                if kj == kshift and kb == kconserv[ki,ka,kk]:
                                    tmp = 0.5*einsum('kiab,j->ijkab',eris.oovv[kk,ki,ka],l1)
                                    lijkab_tmp += 2.* tmp

                                # (k,i,j) -> (i,j,k)
                                if ki == kshift and kb == kconserv[kj,ka,kk]:
                                    tmp = 0.5*einsum('jkab,i->ijkab',eris.oovv[kj,kk,ka],l1)
                                    lijkab_tmp += 2.* tmp

                                # Beginning of ovvv terms

                                ke = kconserv[ka,ki,kb]
                                tmp = einsum('ieab,jke->ijkab',eris.ovvv[ki,ke,ka],2.*l2[kj,kk]-l2[kk,kj].transpose(1,0,2))
                                lijkab_tmp += 2.*tmp

                                ke = kconserv[ka,kj,kb]
                                tmp = einsum('jeba,ike->ijkab',eris.ovvv[kj,ke,kb],2.*l2[ki,kk]-l2[kk,ki].transpose(1,0,2))
                                lijkab_tmp += 2.*tmp

                                # (j,i,k) -> (i,j,k)
                                ke = kconserv[ka,kj,kb]
                                tmp = einsum('jeab,ike->ijkab',eris.ovvv[kj,ke,ka],2.*l2[ki,kk]-l2[kk,ki].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                                ke = kconserv[ka,ki,kb]
                                tmp = einsum('ieba,jke->ijkab',eris.ovvv[ki,ke,kb],2.*l2[kj,kk]-l2[kk,kj].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                                # (k,j,i) -> (i,j,k)
                                ke = kconserv[ki,kshift,kj]
                                tmp = einsum('keab,jie->ijkab',eris.ovvv[kk,ke,ka],2.*l2[kj,ki]-l2[ki,kj].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                                ke = kconserv[ki,kshift,kj]
                                tmp = einsum('keba,ije->ijkab',eris.ovvv[kk,ke,kb],2.*l2[ki,kj]-l2[kj,ki].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                                # ooov part 1

                                km = kconserv[kshift,ki,ka]
                                tmp = -einsum('kjmb,ima->ijkab',eris.ooov[kk,kj,km],2.*l2[ki,km]-l2[km,ki].transpose(1,0,2))
                                lijkab_tmp += 2.*tmp

                                km = kconserv[kshift,kj,kb]
                                tmp = -einsum('kima,jmb->ijkab',eris.ooov[kk,ki,km],2.*l2[kj,km]-l2[km,kj].transpose(1,0,2))
                                lijkab_tmp += 2.*tmp

                                # (j,i,k) -> (i,j,k)
                                km = kconserv[kshift,kj,ka]
                                tmp = -einsum('kimb,jma->ijkab',eris.ooov[kk,ki,km],2.*l2[kj,km]-l2[km,kj].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                                km = kconserv[kshift,ki,kb]
                                tmp = -einsum('kjma,imb->ijkab',eris.ooov[kk,kj,km],2.*l2[ki,km]-l2[km,ki].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                                # (k,j,i) -> (i,j,k)
                                km = kconserv[kshift,kj,kb]
                                tmp = -einsum('ikma,jmb->ijkab',eris.ooov[ki,kk,km],2.*l2[kj,km]-l2[km,kj].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                                km = kconserv[kshift,ki,ka]
                                tmp = -einsum('jkmb,ima->ijkab',eris.ooov[kj,kk,km],2.*l2[ki,km]-l2[km,ki].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                                # ooov part 2

                                km = kconserv[ki,kb,kj]
                                tmp = -einsum('ijmb,mka->ijkab',eris.ooov[ki,kj,km],2.*l2[km,kk]-l2[kk,km].transpose(1,0,2))
                                lijkab_tmp += 2.*tmp

                                km = kconserv[kj,ka,ki]
                                tmp = -einsum('jima,mkb->ijkab',eris.ooov[kj,ki,km],2.*l2[km,kk]-l2[kk,km].transpose(1,0,2))
                                lijkab_tmp += 2.*tmp

                                # (j,i,k) -> (i,j,k)
                                km = kconserv[ki,kb,kj]
                                tmp = -einsum('jimb,mka->ijkab',eris.ooov[kj,ki,km],2.*l2[km,kk]-l2[kk,km].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                                km = kconserv[kj,ka,ki]
                                tmp = -einsum('ijma,mkb->ijkab',eris.ooov[ki,kj,km],2.*l2[km,kk]-l2[kk,km].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                                # (k,j,i) -> (i,j,k)
                                km = kconserv[kshift,ki,kb]
                                tmp = -einsum('jkma,mib->ijkab',eris.ooov[kj,kk,km],2.*l2[km,ki]-l2[ki,km].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                                km = kconserv[kshift,kj,ka]
                                tmp = -einsum('ikmb,mja->ijkab',eris.ooov[ki,kk,km],2.*l2[km,kj]-l2[kj,km].transpose(1,0,2))
                                lijkab_tmp -= 1.*tmp

                            # Starting the right amplitude equations #

                                kk = kklist[ka,kb,iterki,iterkj]
                                km = kshift
                                ke = kconserv[ki,ka,kj]
                                tmp = einsum('mbke,m->bke',eris.ovov[km,kb,kk],r1)
                                tril = unpack_tril(t2,nkpts,ki,kj,ka,kconserv[ki,ka,kj])
                                tmp2 = -einsum('bke,ijae->ijkab',tmp,tril)
                                rijkab_tmp += tmp2

                                ke = kconserv[ki,kb,kj]
                                tmp = einsum('make,m->ake',eris.ovov[km,ka,kk],r1)
                                tril = unpack_tril(t2,nkpts,kj,ki,kb,kconserv[kj,kb,ki])
                                tmp2 = -einsum('ake,jibe->ijkab',tmp,tril)
                                rijkab_tmp += tmp2

                                kk = kklist[ka,kb,iterki,iterkj]
                                km = kshift
                                ke = kconserv[km,kj,kb]
                                tmp = einsum('mbej,m->bej',eris.ovvo[km,kb,ke],r1)
                                tril = unpack_tril(t2,nkpts,ki,kk,ka,kconserv[ki,ka,kk])
                                tmp2 = -einsum('bej,ikae->ijkab',tmp,tril)
                                rijkab_tmp += tmp2

                                ke = kconserv[km,ki,ka]
                                tmp = einsum('maei,m->aei',eris.ovvo[km,ka,ke],r1)
                                tril = unpack_tril(t2,nkpts,kj,kk,kb,kconserv[kj,kb,kk])
                                tmp2 = -einsum('aei,jkbe->ijkab',tmp,tril)
                                rijkab_tmp += tmp2

                                # ki is free, kj is free, ka is free, kk is free
                                # km, kb -> fixed
                                # (i,j,k) -> (i,j,k)
                                kk = kklist[ka,kb,iterki,iterkj]
                                km = kconserv[ka,ki,kb]
                                tmp = einsum('mnjk,n->mjk',eris.oooo[km,kshift,kj],r1)
                                tril = unpack_tril(t2,nkpts,ki,km,ka,kconserv[ki,ka,km])
                                tmp2 = einsum('mjk,imab->ijkab',tmp,tril)
                                rijkab_tmp += tmp2

                                km = kconserv[kb,kj,ka]
                                tmp = einsum('mnik,n->mik',eris.oooo[km,kshift,ki],r1)
                                tril = unpack_tril(t2,nkpts,kj,km,kb,kconserv[kj,kb,km])
                                tmp2 = einsum('mik,jmba->ijkab',tmp,tril)
                                rijkab_tmp += tmp2

                                ke = kconserv[ka,ki,kb]
                                #if kk == kconserv[kshift,kj,ke]:
                                # Here we need vvov terms, but instead use ovvv.transpose(2,3,0,1).conj()
                                # terms, where the transpose was done in the einsum.
                                #
                                # See corresponding rccsd_eom equation
                                tmp2 = einsum('ieab,kje->ijkab',eris.ovvv[ki,ke,ka].conj(),r2[kk,kj])
                                rijkab_tmp += tmp2

                                ke = kconserv[kb,kj,ka]
                                tmp2 = einsum('jeba,kie->ijkab',eris.ovvv[kj,ke,kb].conj(),r2[kk,ki])
                                rijkab_tmp += tmp2

                                km = kconserv[kshift,ki,ka]
                                tmp2 = -einsum('kjmb,mia->ijkab',eris.ooov[kk,kj,km].conj(),r2[km,ki])
                                rijkab_tmp += tmp2

                                km = kconserv[kshift,kj,kb]
                                tmp2 = -einsum('kima,mjb->ijkab',eris.ooov[kk,ki,km].conj(),r2[km,kj])
                                rijkab_tmp += tmp2

                                km = kconserv[ki,kb,kj]
                                tmp2 = -einsum('mbij,kma->ijkab',eris.ovoo[km,kb,ki],r2[kk,km])
                                rijkab_tmp += tmp2

                                km = kconserv[kj,ka,ki]
                                tmp2 = -einsum('maji,kmb->ijkab',eris.ovoo[km,ka,kj],r2[kk,km])
                                rijkab_tmp += tmp2

                                eia = numpy.diagonal(foo[ki]).reshape(-1,1) - numpy.diagonal(fvv[ka])
                                eia += _eval

                                ejb = numpy.diagonal(foo[kj]).reshape(-1,1) - numpy.diagonal(fvv[kb])
                                eijab = pyscf.lib.direct_sum('ia,jb->ijab',eia,ejb)

                                eijkab = pyscf.lib.direct_sum('ijab,k->ijkab',eijab,numpy.diagonal(foo[kk]))
                                eijkab = 1./eijkab

                                star_energy += fac*0.5*einsum('ijkab,ijkab,ijkab',lijkab_tmp,rijkab_tmp,eijkab)

            comm.Allreduce(MPI.IN_PLACE, star_energy, op=MPI.SUM)
            if rank == 0:
                logger.info(self, 'EOM-IPCCSD* delta energy = %.15g (Imag = %.15g)',
                            star_energy.real,star_energy.imag)
            e.append(star_energy+_eval)

        return np.array(e)

    def eaccsd_diag(self):
        t1,t2 = self.t1, self.t2
        nkpts, nocc, nvir = t1.shape
        kshift = self.kshift
        kconserv = self.kconserv

        if not self.made_ea_imds:
            if not getattr(self, 'imds', None):
                self.imds = _IMDS(self)
            self.imds.make_ea(self)
            self.made_ea_imds = True

        imds = self.imds

        Hr1 = numpy.diag(imds.Lvv[kshift])

        Hr2 = numpy.zeros((nkpts,nkpts,nocc,nvir,nvir),dtype=t1.dtype)
        mem = 0.5e9
        pre = 1.*nvir*nvir*nvir*nvir*nkpts*16
        nkpts_blksize  = min(max(int(numpy.floor(mem/pre)),1),nkpts)
        nkpts_blksize2 = min(max(int(numpy.floor(mem/(nkpts_blksize*pre))),1),nkpts)
        loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts_blksize2,nkpts_blksize,))
        loader.set_ranges((range(nkpts),range(nkpts),))

        good2go = True
        while(good2go):
            good2go, data = loader.slave_set()
            if good2go is False:
                break
            ranges0, ranges1 = loader.get_blocks_from_data(data)

            s0,s1 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1)]

            for iterkj,kj in enumerate(ranges0):
                for iterka,ka in enumerate(ranges1):
                    kb = kconserv[kshift,ka,kj]

                    Wvvvv_aba    = _cp(imds.Wvvvv[ka,kb,ka])
                    WvoovR1_jbj  = _cp(imds.WvoovR1[kj,kb,kj])
                    WovovRev_jbj = _cp(imds.WovovRev[kj,kb,kj])
                    WovovRev_jaj = _cp(imds.WovovRev[kj,ka,kj])

                    for j in range(nocc):
                        Hr2[kj,ka,j,:,:] -= imds.Loo[kj,j,j]
                    for a in range(nvir):
                        Hr2[kj,ka,:,a,:] += imds.Lvv[ka,a,a]
                    for b in range(nvir):
                        Hr2[kj,ka,:,:,b] += imds.Lvv[kb,b,b]

                    for a in range(nvir):
                        for b in range(nvir):
                            Hr2[kj,ka,j,a,b] += Wvvvv_aba[a,b,a,b]

                    for j in range(nocc):
                        for b in range(nvir):
                            Hr2[kj,ka,j,:,b] += 2.*WvoovR1_jbj[b,j,j,b]
                            Hr2[kj,ka,j,:,b] += -WovovRev_jbj.transpose(1,0,3,2)[b,j,b,j]
                            Hr2[kj,ka,j,b,b] += -WvoovR1_jbj[b,j,j,b]

                            Hr2[kj,ka,j,b,:] += -WovovRev_jaj.transpose(1,0,3,2)[b,j,b,j]
                            for a in range(nvir):
                                Hr2[kj,ka,j,a,b] += numpy.dot(unpack_tril(t2,nkpts,kshift,kj,ka,kconserv[kshift,ka,kj])[:,j,a,b],
                                                              -2.*imds.Woovv[kshift,kj,ka,:,j,a,b] + imds.Woovv[kshift,kj,kb,:,j,b,a])

            loader.slave_finished()
        comm.Allreduce(MPI.IN_PLACE, Hr2, op=MPI.SUM)

        vector = self.amplitudes_to_vector_ea(Hr1,Hr2)
        return vector

    def eaccsd(self, nroots=2*4, kptlist=None):
        time0 = logger.process_clock(), logger.perf_counter()
        log = logger.Logger(self.stdout, self.verbose)
        nkpts = self.nkpts
        if kptlist is None:
            kptlist = range(nkpts)
        size = self.vector_size_ea()
        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            nfrozen = np.sum(self.mask_frozen_ea(np.zeros(size, dtype=int), kshift, const=1))
            nroots = min(nroots, size - nfrozen)
        evals = np.zeros((len(kptlist),nroots), np.float)
        evecs = np.zeros((len(kptlist),nroots,size), np.complex128)

        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            diag = self.eaccsd_diag()
            diag = self.mask_frozen_ea(diag, kshift, const=LARGE_DENOM)
            # Initial guess from file
            amplitude_filename = "__reaccsd" + str(kshift) + "__.hdf5"
            rsuccess, x0 = read_eom_amplitudes((nroots,size),filename=amplitude_filename)
            if x0 is not None:
                x0 = x0.T
            #if not rsuccess:
            #    x0 = np.zeros_like(diag)
            #    x0[np.argmin(diag)] = 1.0

            conv, evals_k, evecs_k = eigs(self.eaccsd_matvec, size, nroots, x0=x0, Adiag=diag, verbose=self.verbose)

            evals_k = evals_k.real
            evals[k] = evals_k
            evecs[k] = evecs_k.T

            write_eom_amplitudes(evecs[k],filename=amplitude_filename)
        time0 = log.timer_debug1('converge ea-ccsd', *time0)
        comm.Barrier()
        return evals.real, evecs

    def eaccsd_matvec(self, vector):
        # Ref: Nooijen and Bartlett, J. Chem. Phys. 102, 3629 (1994) Eqs.(30)-(31)
        kshift = self.kshift
        vector = self.mask_frozen_ea(vector, kshift, const=0.0)
        r1,r2 = self.vector_to_amplitudes_ea(vector)
        r1 = comm.bcast(r1, root=0)
        r2 = comm.bcast(r2, root=0)

        t1,t2 = self.t1, self.t2
        nkpts,nocc,nvir = self.t1.shape
        nkpts = self.nkpts
        kconserv = self.kconserv

        if not self.made_ea_imds:
            if not getattr(self, 'imds', None):
                self.imds = _IMDS(self)
            self.imds.make_ea(self)
            self.made_ea_imds = True

        imds = self.imds

        Hr1 = numpy.zeros(r1.shape,dtype=t1.dtype)
        mem = 0.5e9
        pre = 1.*nocc*nvir*nvir*nvir*nkpts*16
        nkpts_blksize = min(max(int(numpy.floor(mem/pre)),1),nkpts)
        loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts_blksize,))
        loader.set_ranges((range(nkpts),))

        good2go = True
        while(good2go):
            good2go, data = loader.slave_set()
            if good2go is False:
                break
            ranges0 = loader.get_blocks_from_data(data)

            s0 = slice(min(ranges0),max(ranges0)+1)

            Wvovv_slX = _cp(imds.Wvovv[kshift,s0,:])

            for iterkl,kl in enumerate(ranges0):
                Hr1 += 2.*einsum('ld,lad->a',imds.Fov[kl],r2[kl,kshift])
                Hr1 +=   -einsum('ld,lda->a',imds.Fov[kl],r2[kl,kl])
                kd_list = numpy.array(kconserv[kshift,range(nkpts),kl])
                Hr1 += einsum('alxcd,lxcd->a',
                              (2.*Wvovv_slX[iterkl,:].transpose(1,2,0,3,4) -
                               Wvovv_slX[iterkl,kd_list].transpose(1,2,0,4,3)),
                              r2[kl,:].transpose(1,0,2,3))
            loader.slave_finished()
        comm.Allreduce(MPI.IN_PLACE, Hr1, op=MPI.SUM)
        Hr1 += einsum('ac,c->a',imds.Lvv[kshift],r1)

        Hr2 = numpy.zeros(r2.shape,dtype=t1.dtype)
        mem = 0.5e9
        pre = 1.*nvir*nvir*nvir*nvir*nkpts*16
        nkpts_blksize  = min(max(int(numpy.floor(mem/pre)),1),nkpts)
        nkpts_blksize2 = min(max(int(numpy.floor(mem/(nkpts_blksize*pre))),1),nkpts)
        loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts_blksize2,nkpts_blksize,))
        loader.set_ranges((range(nkpts),range(nkpts),))

        good2go = True
        while(good2go):
            good2go, data = loader.slave_set()
            if good2go is False:
                break
            ranges0, ranges1 = loader.get_blocks_from_data(data)

            s0,s1 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1)]

            for iterkj,kj in enumerate(ranges0):
                for iterka,ka in enumerate(ranges1):
                    kb = kconserv[kshift,ka,kj]
                    Hr2[kj,ka] -= einsum('lj,lab->jab',imds.Loo[kj],r2[kj,ka])
                    Hr2[kj,ka] += einsum('ac,jcb->jab',imds.Lvv[ka],r2[kj,ka])
                    Hr2[kj,ka] += einsum('bd,jad->jab',imds.Lvv[kb],r2[kj,ka])

            WvvvoR1_abX = _cp(imds.WvvvoR1[kshift,s0,s1])
            for iterka,ka in enumerate(ranges0):
                for iterkb,kb in enumerate(ranges1):
                    kj = kconserv[ka,kshift,kb]
                    Hr2[kj,ka] += einsum('abcj,c->jab',WvvvoR1_abX[iterka,iterkb],r1)

            WovovRev_Xaj = _cp(imds.WovovRev[s0,s1,:])
            for iterkj,kj in enumerate(ranges0):
                for iterka,ka in enumerate(ranges1):
                    kb = kconserv[kshift,ka,kj]
                    kl_range = range(nkpts)
                    kd_range = kconserv[ka,kj,kl_range]
                    Hr2[kj,ka] += -einsum('axj,xb->jab',
                                          WovovRev_Xaj[iterkj,iterka,:].transpose(2,0,1,4,3).reshape(nvir,nocc*nvir*nkpts,nocc),
                                          r2[kl_range,kd_range].reshape(nkpts*nocc*nvir,nvir))

            tmp = numpy.zeros(nocc,dtype=t2.dtype)
            for kl in range(nkpts):
                kd_range = _cp(range(nkpts))
                kc_range = _cp(kconserv[kshift,kd_range,kl])
                tmp += einsum('kl,l->k',(2.*imds.Woovv[kshift,kl,kc_range].transpose(1,2,0,3,4)-
                                            imds.Woovv[kshift,kl,kd_range].transpose(1,2,0,4,3)).reshape(nocc,-1),
                                        r2[kl,kc_range].transpose(1,0,2,3).reshape(-1))

            for iterkj,kj in enumerate(ranges0):
                for iterka,ka in enumerate(ranges1):
                    kb = kconserv[kshift,ka,kj]
                    Hr2[kj,ka] += -einsum('k,kjab->jab', tmp, unpack_tril(t2,nkpts,kshift,kj,ka,kb))

            WovovRev_Xbj = _cp(imds.WovovRev[s0,s1,:])
            WvoovR1_bXj = _cp(imds.WvoovR1[s0,s1,:])
            for iterkj,kj in enumerate(ranges0):
                for iterkb,kb in enumerate(ranges1):
                    ka = kconserv[kshift,kb,kj]
                    Hr2[kj,ka] += -einsum('bldj,lad->jab',
                                          WovovRev_Xbj[iterkj,iterkb,:].transpose(2,0,1,4,3).reshape(nvir,nkpts*nocc,nvir,nocc),
                                          r2[:,ka].reshape(nkpts*nocc,nvir,nvir))
                    kl_range = _cp(range(nkpts))
                    kd_range = _cp(kconserv[kb,kj,kl_range])
                    Hr2[kj,ka] += einsum('bljd,lad->jab',
                                         WvoovR1_bXj[iterkj,iterkb,:].transpose(1,0,2,3,4).reshape(nvir,nocc*nkpts,nocc,nvir),
                                         (2.*r2[:,ka]-r2[kl_range,kd_range].transpose(0,1,3,2)).reshape(nocc*nkpts,nvir,nvir))

            Wvvvv_abX = _cp(imds.Wvvvv[s0,s1])
            for iterka,ka in enumerate(ranges0):
                for iterkb,kb in enumerate(ranges1):
                    kj = kconserv[ka,kshift,kb]
                    Hr2[kj,ka] += einsum('abx,jx->jab',
                                         Wvvvv_abX[iterka,iterkb,:].transpose(1,2,0,3,4).reshape(nvir,nvir,nvir*nkpts*nvir),
                                         r2[kj,:].transpose(1,0,2,3).reshape(nocc,nvir*nkpts*nvir))
            loader.slave_finished()

        comm.Allreduce(MPI.IN_PLACE, Hr2, op=MPI.SUM)

        vector = self.amplitudes_to_vector_ea(Hr1,Hr2)
        vector = self.mask_frozen_ea(vector, kshift, const=0.0)
        return vector

    def leaccsd(self, nroots=2*4, kptlist=None):
        time0 = logger.process_clock(), logger.perf_counter()
        log = logger.Logger(self.stdout, self.verbose)
        nkpts = self.nkpts
        if kptlist is None:
            kptlist = range(nkpts)
        size = self.vector_size_ea()
        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            nfrozen = np.sum(self.mask_frozen_ea(np.zeros(size, dtype=int), kshift, const=1))
            nroots = min(nroots, size - nfrozen)
        evals = np.zeros((len(kptlist),nroots), np.float)
        evecs = np.zeros((len(kptlist),nroots,size), np.complex128)

        for k,kshift in enumerate(kptlist):
            self.kshift = kshift
            diag = self.eaccsd_diag()
            # Initial guess from file
            amplitude_filename = "__leaccsd" + str(kshift) + "__.hdf5"
            rsuccess, x0 = read_eom_amplitudes((nroots,size),filename=amplitude_filename)
            if x0 is not None:
                x0 = x0.T
            #if not rsuccess:
            #    x0 = np.zeros_like(diag)
            #    x0[np.argmin(diag)] = 1.0

            conv, evals_k, evecs_k = eigs(self.leaccsd_matvec, size, nroots, x0=x0, Adiag=diag, verbose=self.verbose)

            evals_k = evals_k.real
            evals[k] = evals_k
            evecs[k] = evecs_k.T

            write_eom_amplitudes(evecs[k],amplitude_filename)
        time0 = log.timer_debug1('converge lea-ccsd', *time0)
        comm.Barrier()
        return evals.real, evecs

    def leaccsd_matvec(self, vector):
        # See relevant equations in cc/rccsd_slow.py
        # Does not follow the equations in Nooijen's paper
        # for eaccsd, uses a different left basis.
        r1,r2 = self.vector_to_amplitudes_ea(vector)
        r1 = comm.bcast(r1, root=0)
        r2 = comm.bcast(r2, root=0)

        t1,t2 = self.t1, self.t2
        nkpts,nocc,nvir = self.t1.shape
        nkpts = self.nkpts
        kshift = self.kshift
        kconserv = self.kconserv

        if not self.made_ea_imds:
            if not getattr(self, 'imds', None):
                self.imds = _IMDS(self)
            self.imds.make_ea(self)
            self.made_ea_imds = True

        imds = self.imds

        # Beginning HR1 multiplication
        Hr1 = numpy.zeros(r1.shape,dtype=t2.dtype)
        def mem_usage_vvvo(nocc, nvir, nkpts):
            return nocc**1 * nvir**3 * 16.
        array_size = [nkpts,nkpts]
        task_list = generate_max_task_list(array_size,blk_mem_size=mem_usage_vvvo(nocc,nvir,nkpts),priority_list=[1,1])
        for karange, kbrange in mpi.work_stealing_partition(task_list):
            WvvvoR1_sab = _cp(imds.WvvvoR1[kshift,slice(*karange),slice(*kbrange)])
            for iterka, ka in enumerate(range(*karange)):
                for iterkb, kb in enumerate(range(*kbrange)):
                    kj = kconserv[ka,kshift,kb]
                    Hr1 += einsum('abcj,jab->c',WvvvoR1_sab[iterka,iterkb],r2[kj,ka])
        comm.Allreduce(MPI.IN_PLACE, Hr1, op=MPI.SUM)
        Hr1 += einsum('ac,a->c',imds.Lvv[kshift],r1)

        # Beginning HR2 multiplication
        Hr2 = numpy.zeros(r2.shape,dtype=t2.dtype)
        # use same task list as before
        for klrange, kcrange in mpi.work_stealing_partition(task_list):
            Wvovv_slc = _cp(imds.Wvovv[kshift,slice(*klrange),slice(*kcrange)])
            for iterkl, kl in enumerate(range(*klrange)):
                for iterkc, kc in enumerate(range(*kcrange)):
                    kd = kconserv[kl,kc,kshift]
                    Hr2[kl,kc] +=  einsum('lad,ac->lcd',r2[kl,kc],imds.Lvv[kc])
                    Hr2[kl,kc] +=  einsum('lcb,bd->lcd',r2[kl,kc],imds.Lvv[kd])
                    Hr2[kl,kc] += -einsum('jcd,lj->lcd',r2[kl,kc],imds.Loo[kl])
                    tmp = einsum('a,alcd->lcd',r1,Wvovv_slc[iterkl,iterkc])
                    Hr2[kl,kc] += 2.*tmp
                    Hr2[kl,kd] +=   -tmp.transpose(0,2,1)
                    tmp = einsum('c,ld->lcd',r1,imds.Fov[kl])
                    Hr2[kl,kshift] += 2.*tmp*(kl==kd)
                    Hr2[kl,kd]     +=   -tmp.transpose(0,2,1)*(kl==kd)

        def mem_usage_voovk(nocc, nvir, nkpts):
            return nocc**2 * nvir**2 * nkpts * 16.
        array_size = [nkpts,nkpts]
        task_list = generate_max_task_list(array_size, mem_usage_voovk(nocc,nvir,nkpts), priority_list=[1,1])
        for kjrange, kbrange in mpi.work_stealing_partition(task_list):
            WvoovR1_jbX  = _cp(imds.WvoovR1[slice(*kjrange),slice(*kbrange),:])
            WovovRev_jbX = _cp(imds.WovovRev[slice(*kjrange),slice(*kbrange),:])

            for iterkj, kj in enumerate(range(*kjrange)):
                for iterkb, kb in enumerate(range(*kbrange)):
                    for iterkl, kl in enumerate(range(nkpts)):
                        kc = kconserv[kj,kb,kshift]
                        Hr2[kl,kc] += 2.*einsum('jcb,bljd->lcd',r2[kj,kc],WvoovR1_jbX[iterkj,iterkb,kl])
                        Hr2[kl,kc] -=    einsum('jcb,lbjd->lcd',r2[kj,kc],WovovRev_jbX[iterkj,iterkb,kl])
                        kc = kconserv[kl,kj,kb]
                        kd = kconserv[kshift,kc,kl]
                        Hr2[kl,kc] -=    einsum('jdb,bljc->lcd',r2[kj,kd],WvoovR1_jbX[iterkj,iterkb,kl])
                        Hr2[kl,kc] -=    einsum('jbd,lbjc->lcd',r2[kj,kb],WovovRev_jbX[iterkj,iterkb,kl])

        def mem_usage_vvvvk(nocc, nvir, nkpts):
            return nocc**0 * nvir**4 * nkpts * 16.
        array_size = [nkpts,nkpts]
        task_list = generate_max_task_list(array_size, mem_usage_vvvvk(nocc,nvir,nkpts),priority_list=[1,1])
        for karange, kbrange in mpi.work_stealing_partition(task_list):
            Wvvvv_abX = _cp(imds.Wvvvv[slice(*karange),slice(*kbrange),:])

            for iterka, ka in enumerate(range(*karange)):
                for iterkb, kb in enumerate(range(*kbrange)):
                    for iterkc, kc in enumerate(range(nkpts)):
                        kl = kconserv[ka,kshift,kb]
                        Hr2[kl,kc] += einsum('lab,abcd->lcd',r2[kl,ka],Wvvvv_abX[iterka,iterkb,kc])

        def mem_usage_oovvk(nocc, nvir, nkpts):
            return nocc**2 * nvir**2 * nkpts * 16.
        array_size = [nkpts,nkpts]
        task_list = generate_max_task_list(array_size, mem_usage_oovvk(nocc,nvir,nkpts),priority_list=[1,1])
        # TODO mpi.work_stealing_partition returns [[0,1]] for one dimension
        # and this doesn't work with the kjrange
        # Making term for Woovv contraction
        tmp_occ = numpy.zeros(nocc,dtype=t1.dtype)
        for kirange, kcrange in mpi.work_stealing_partition(task_list):
            for iterki, ki in enumerate(range(*kirange)):
                for iterkc, kc in enumerate(range(*kcrange)):
                    kb = kconserv[ki,kc,kshift]
                    tmp2 = unpack_tril(t2,nkpts,ki,kshift,kc,kconserv[ki,kc,kshift])
                    tmp_occ += einsum('ijcb,ibc->j',tmp2,r2[ki,kb])
        comm.Allreduce(MPI.IN_PLACE, tmp_occ, op=MPI.SUM)

        for kkrange, kfrange in mpi.work_stealing_partition(task_list):
            Woovv_kf = _cp(imds.Woovv[slice(*kkrange),kshift,slice(*kfrange)])
            for iterkk, kk in enumerate(range(*kkrange)):
                for iterkf, kf in enumerate(range(*kfrange)):
                    ke = kconserv[kshift,kf,kk]
                    tmp = einsum('kjfe,j->kef',Woovv_kf[iterkk,iterkf],tmp_occ)
                    Hr2[kk,ke] += -2.*tmp
                    Hr2[kk,kf] += tmp.transpose(0,2,1)

        comm.Allreduce(MPI.IN_PLACE, Hr2, op=MPI.SUM)

        vector = self.amplitudes_to_vector_ea(Hr1,Hr2)
        return vector

    def vector_size_ea(self):
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts

        size = nvir + nkpts ** 2 * nvir ** 2 * nocc
        return size

    mask_frozen_ea = mask_frozen_ea

    def vector_to_amplitudes_ea(self,vector):
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts

        r1 = vector[:nvir].copy()
        r2 = vector[nvir:].copy().reshape(nkpts,nkpts,nocc,nvir,nvir)
        #r2 = np.zeros((nkpts,nkpts,nocc,nvir,nvir), vector.dtype)
        #index = nvir
        #for kj in range(nkpts):
        #    for ka in range(nkpts):
        #        for j in range(nocc):
        #            for a in range(nvir):
        #                for b in range(nvir):
        #                    r2[kj,ka,j,a,b] = vector[index]
        #                    index += 1
        return [r1,r2]

    def amplitudes_to_vector_ea(self,r1,r2):
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        size = nvir + nkpts*nkpts*nocc*nvir*nvir

        vector = np.zeros((size), r1.dtype)
        vector[:nvir] = r1.copy()
        vector[nvir:] = r2.copy().reshape(nkpts*nkpts*nocc*nvir*nvir)
        #index = nvir
        #for kj in range(nkpts):
        #    for ka in range(nkpts):
        #        for j in range(nocc):
        #            for a in range(nvir):
        #                for b in range(nvir):
        #                    vector[index] = r2[kj,ka,j,a,b]
        #                    index += 1
        return vector

    def eaccsd_star_contract(self, eaccsd_evals, eaccsd_evecs, leaccsd_evecs):
        nocc = self.nocc
        nvir = self.nmo - nocc
        nkpts = self.nkpts
        kconserv = self.kconserv
        eris = self.eris
        kshift = self.kshift
        t2 = self.t2
        foo = eris.fock[:,:nocc,:nocc]
        fvv = eris.fock[:,nocc:,nocc:]

        e = []
        assert len(eaccsd_evecs.shape) == 2  # Done at a single k-point, kshift
        for _eval, _evec, _levec in zip(eaccsd_evals, eaccsd_evecs, leaccsd_evecs):
            print(_eval, _evec.shape, _levec.shape)
            l1,l2 = self.vector_to_amplitudes_ea(_levec)
            r1,r2 = self.vector_to_amplitudes_ea(_evec)

            # Transposing the l2 operator
            # (equations were worked out for the transpose of these operators!)

            l2_t = numpy.zeros_like(l2)
            r2_t = numpy.zeros_like(r2)
            for kj in range(nkpts):
                for ka in range(nkpts):
                    kb = kconserv[kshift,ka,kj]
                    l2_t[kj,kb] = l2[kj,ka].transpose(0,2,1)
                    r2_t[kj,kb] = r2[kj,ka].transpose(0,2,1)

            # Normalization for the left eigenvector
            ldotr = numpy.dot(l1.ravel(),r1.ravel()) + numpy.dot(l2.ravel(),r2.ravel())
            l2 = (l2 + 2.*l2_t)/3.
            r2 = r2_t.copy()
            r2_t = None
            l2_t = None

            l1 /= ldotr
            l2 /= ldotr
            l1 = l1.reshape(-1)
            r1 = r1.reshape(-1)

            lijabc_tmp  = numpy.zeros((nocc,nocc,nvir,nvir,nvir),dtype=t2.dtype)
            rijabc_tmp  = numpy.zeros((nocc,nocc,nvir,nvir,nvir),dtype=t2.dtype)
            def mem_usage_oovvk(nocc, nvir, nkpts):
                return nocc**2 * nvir**2 * nkpts * 16
            array_size = [nkpts,nkpts]
            # FIXME: find a good blocksize for this. right now just makes the smallest blocksize possible
            task_list = generate_max_task_list(array_size,blk_mem_size=1e12,priority_list=[1,1])
            #nproc = comm.Get_size()
            #chunk_size = get_max_blocksize_from_mem(0.3e9, 2.*mem_usage_oovvk(nocc,nvir,nkpts),
            #                                        array_size, priority_list=[1,1])
            #max_chunk_size = tuple(nkpts//numpy.array([nproc,nproc]))
            #if max_chunk_size[0] < chunk_size[0]:
            #    chunk_size = max_chunk_size
            #task_list = generate_task_list(chunk_size,array_size)

            star_energy = numpy.array(0.0 + 1j*0.0)
            for kirange, kjrange in mpi.work_stealing_partition(task_list):
                #eris_oovv_jiX = _cp(eris.oovv[slice(*kjrange),slice(*kirange),range(nkpts)])
                eris_oovv_ijX = _cp(eris.oovv[slice(*kirange),slice(*kjrange),range(nkpts)])

                eris_ooov_jiX = _cp(eris.ooov[slice(*kjrange),slice(*kirange),range(nkpts)])
                eris_ooov_ijX = _cp(eris.ooov[slice(*kirange),slice(*kjrange),range(nkpts)])

                kclist = kpts_helper.get_kconserv3(self._scf.cell, self.kpts,
                                                   [range(*kirange),range(*kjrange),kshift,range(nkpts),range(nkpts)])

                for iterki, ki in enumerate(range(*kirange)):
                    for iterkj, kj in enumerate(range(*kjrange)):

                        # Starting the left amplitude equations

                        for iterka, ka in enumerate(range(nkpts)):
                            for iterkb, kb in enumerate(range(nkpts)):
                                if ka > kb:
                                    continue
                                kc = kclist[iterki,iterkj,ka,kb]
                                fac = (2.0 - float(ka==kb))

                                if ka == kb:
                                    if ki > kj:
                                        continue
                                    else:
                                        fac *= (2.0 - float(ki==kj))

                                lijabc_tmp *= 0.0
                                rijabc_tmp *= 0.0

                                #(a,b,c) -> (a,b,c)
                                if kc == kshift and kb == kconserv[ki,ka,kj]:
                                    tmp = -0.5*einsum('ijab,c->ijabc',eris_oovv_ijX[iterki,iterkj,ka],l1)
                                    lijabc_tmp += 8.*tmp

                                #(b,a,c) -> (a,b,c)
                                if kc == kshift and ka == kconserv[ki,kb,kj]:
                                    tmp = -0.5*einsum('ijba,c->ijabc',eris_oovv_ijX[iterki,iterkj,kb],l1)
                                    lijabc_tmp += -4.* tmp

                                #(c,b,a) -> (a,b,c)
                                if ka == kshift and kb == kconserv[ki,kc,kj]:
                                    tmp = -0.5*einsum('ijcb,a->ijabc',eris_oovv_ijX[iterki,iterkj,kc],l1)
                                    lijabc_tmp += -4.* tmp

                                #(a,c,b) -> (a,b,c)
                                if kb == kshift and kc == kconserv[ki,ka,kj]:
                                    tmp = -0.5*einsum('ijac,b->ijabc',eris_oovv_ijX[iterki,iterkj,ka],l1)
                                    lijabc_tmp += -4.* tmp

                                #(b,c,a) -> (a,b,c)
                                if kb == kshift and ka == kconserv[ki,kc,kj]:
                                    tmp = -0.5*einsum('ijca,b->ijabc',eris_oovv_ijX[iterki,iterkj,kc],l1)
                                    lijabc_tmp += 2.* tmp

                                #(c,a,b) -> (a,b,c)
                                if ka == kshift and kc == kconserv[ki,kb,kj]:
                                    tmp = -0.5*einsum('ijbc,a->ijabc',eris_oovv_ijX[iterki,iterkj,kb],l1)
                                    lijabc_tmp += 2.* tmp

                                # end of oovv terms

                                #(a,b,c) -> (a,b,c)
                                km = kconserv[ki,ka,kj]
                                tmp = einsum('jima,mbc->ijabc',eris_ooov_jiX[iterkj,iterki,km],2.*l2[km,kb] - l2[km,kc].transpose(0,2,1))
                                lijabc_tmp += 2.* tmp

                                km = kconserv[kj,kb,ki]
                                tmp = einsum('ijmb,mac->ijabc',eris_ooov_ijX[iterki,iterkj,km],2.*l2[km,ka] - l2[km,kc].transpose(0,2,1))
                                lijabc_tmp += 2.* tmp

                                #(b,a,c) -> (a,b,c)
                                km = kconserv[ki,kb,kj]
                                tmp = einsum('jimb,mac->ijabc',eris_ooov_jiX[iterkj,iterki,km],2.*l2[km,ka] - l2[km,kc].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp

                                km = kconserv[kj,ka,ki]
                                tmp = einsum('ijma,mbc->ijabc',eris_ooov_ijX[iterki,iterkj,km],2.*l2[km,kb] - l2[km,kc].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp

                                #(c,b,a) -> (a,b,c)
                                km = kconserv[ki,kc,kj]
                                tmp = einsum('jimc,mba->ijabc',eris_ooov_jiX[iterkj,iterki,km],2.*l2[km,kb] - l2[km,ka].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp

                                km = kconserv[kj,kc,ki]
                                tmp = einsum('ijmc,mab->ijabc',eris_ooov_ijX[iterki,iterkj,km],2.*l2[km,ka] - l2[km,kb].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp

                                # ovvv term 1

                                # ke = ks - kc + kj
                                # kb = ki + ke - ka
                                #    = ki + ks - kc + kj - ka
                                #    = ki + kj + ks - ka - kc

                                #(c,b,a) -> (a,b,c)
                                ke = kconserv[kshift,ka,ki]
                                tmp = -einsum('jebc,iae->ijabc',eris.ovvv[kj,ke,kb],2.*l2[ki,ka]-l2[ki,ke].transpose(0,2,1))
                                lijabc_tmp += 2.* tmp

                                ke = kconserv[kshift,kb,kj]
                                tmp = -einsum('ieac,jbe->ijabc',eris.ovvv[ki,ke,ka],2.*l2[kj,kb]-l2[kj,ke].transpose(0,2,1))
                                lijabc_tmp += 2.* tmp

                                #(b,c,a) -> (a,b,c)
                                ke = kconserv[kshift,kb,kj]
                                tmp = -einsum('ieca,jbe->ijabc',eris.ovvv[ki,ke,kc],2.*l2[kj,kb]-l2[kj,ke].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp

                                ke = kconserv[kshift,kb,ki]
                                tmp = -einsum('jeac,ibe->ijabc',eris.ovvv[kj,ke,ka],2.*l2[ki,kb]-l2[ki,ke].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp

                                #(c,a,b) -> (a,b,c)
                                ke = kconserv[kshift,ka,kj]
                                tmp = -einsum('iebc,jae->ijabc',eris.ovvv[ki,ke,kb],2.*l2[kj,ka]-l2[kj,ke].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp

                                ke = kconserv[kshift,ka,ki]
                                tmp = -einsum('jecb,iae->ijabc',eris.ovvv[kj,ke,kc],2.*l2[ki,ka]-l2[ki,ke].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp

                                # ovvv term 2

                                # ke = kshift + ki - ka
                                # kb = kj + ke - kc
                                #    = kj + ki + kshift - ka - kc

                                #(c,b,a) -> (a,b,c)
                                ke = kconserv[kshift,kc,ki]
                                tmp = -einsum('jeba,iec->ijabc',eris.ovvv[kj,ke,kb],2.*l2[ki,ke]-l2[ki,kc].transpose(0,2,1))
                                lijabc_tmp += 2.* tmp

                                ke = kconserv[kshift,kc,kj]
                                tmp = -einsum('ieab,jec->ijabc',eris.ovvv[ki,ke,ka],2.*l2[kj,ke]-l2[kj,kc].transpose(0,2,1))
                                lijabc_tmp += 2.* tmp

                                #(b,c,a) -> (a,b,c)
                                ke = kconserv[kshift,kc,ki]
                                tmp = -einsum('jeab,iec->ijabc',eris.ovvv[kj,ke,ka],2.*l2[ki,ke]-l2[ki,kc].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp

                                ke = kconserv[kshift,ka,kj]
                                tmp = -einsum('iecb,jea->ijabc',eris.ovvv[ki,ke,kc],2.*l2[kj,ke]-l2[kj,ka].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp

                                #(c,a,b) -> (a,b,c)
                                ke = kconserv[kshift,kb,ki]
                                tmp = -einsum('jeca,ieb->ijabc',eris.ovvv[kj,ke,kc],2.*l2[ki,ke]-l2[ki,kb].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp

                                ke = kconserv[kshift,kc,kj]
                                tmp = -einsum('ieba,jec->ijabc',eris.ovvv[ki,ke,kb],2.*l2[kj,ke]-l2[kj,kc].transpose(0,2,1))
                                lijabc_tmp += -1.* tmp


                                # Starting the right amplitude equations

                                # kf = ks
                                # ke = ki + kj - ka
                                # kc = kf + ke - kb
                                #    = ks + ki + kj - ka - kb
                                ke = kconserv[ki,ka,kj]
                                vvvv = _cp(eris.vvvv[kb,kc,ke])
                                tmp2 = einsum('bcef,f->bce',vvvv,r1)
                                tril = unpack_tril(t2,nkpts,ki,kj,ka,kconserv[ki,ka,kj])
                                tmp = - einsum('bce,ijae->ijabc',tmp2,tril)
                                rijabc_tmp += tmp

                                ke = kconserv[kj,kb,ki]
                                vvvv = _cp(eris.vvvv[ka,kc,ke])
                                tmp2 = einsum('acef,f->ace',vvvv,r1)
                                tril = unpack_tril(t2,nkpts,kj,ki,kb,kconserv[kj,kb,ki])
                                tmp = - einsum('ace,jibe->ijabc',tmp2,tril)
                                rijabc_tmp += tmp
                                ke = kshift
                                km = kconserv[ke,kc,kj]
                                ovov = _cp(eris.ovov[km,kc,kj])
                                tmp2 = einsum('mcje,e->mcj',ovov,r1)
                                tril = unpack_tril(t2,nkpts,ki,km,ka,kconserv[ki,ka,km])
                                tmp = einsum('mcj,imab->ijabc',tmp2,tril)
                                rijabc_tmp += tmp

                                ke = kshift
                                km = kconserv[ke,kc,ki]
                                ovov = _cp(eris.ovov[km,kc,ki])
                                tmp2 = einsum('mcie,e->mci',ovov,r1)
                                tril = unpack_tril(t2,nkpts,kj,km,kb,kconserv[kj,kb,km])
                                tmp = einsum('mci,jmba->ijabc',tmp2,tril)
                                rijabc_tmp += tmp

                                ke = kshift
                                km = kconserv[kc,ki,ka]
                                ovvo = _cp(eris.ovvo[km,kb,ke])
                                tmp2 = einsum('mbej,e->mbj',ovvo,r1)
                                tril = unpack_tril(t2,nkpts,ki,km,ka,kconserv[ki,ka,km])
                                tmp = einsum('mbj,imac->ijabc',tmp2,tril)
                                rijabc_tmp += tmp

                                ke = kshift
                                km = kconserv[kc,kj,kb]
                                ovvo = _cp(eris.ovvo[km,ka,ke])
                                tmp2 = einsum('maei,e->mai',ovvo,r1)
                                tril = unpack_tril(t2,nkpts,kj,km,kb,kconserv[kj,kb,km])
                                tmp = einsum('mai,jmbc->ijabc',tmp2,tril)
                                rijabc_tmp += tmp

                                ks = kshift
                                km = kconserv[ki,ka,kj]
                                tmp = einsum('maji,mbc->ijabc',eris.ovoo[km,ka,kj],r2[km,kb])
                                rijabc_tmp += tmp

                                ks = kshift
                                km = kconserv[kj,kb,ki]
                                tmp = einsum('mbij,mac->ijabc',eris.ovoo[km,kb,ki],r2[km,ka])
                                rijabc_tmp += tmp

                                ks = kshift
                                ke = kconserv[ks,ka,ki]
                                tmp = -einsum('ejcb,iae->ijabc',eris.vovv[ke,kj,kc].conj(),r2[ki,ka])
                                rijabc_tmp += tmp

                                ks = kshift
                                ke = kconserv[ks,kb,kj]
                                tmp = -einsum('eica,jbe->ijabc',eris.vovv[ke,ki,kc].conj(),r2[kj,kb])
                                rijabc_tmp += tmp

                                ks = kshift
                                ke = kconserv[ks,kc,kj]
                                tmp = -einsum('eiba,jec->ijabc',eris.vovv[ke,ki,kb].conj(),r2[kj,ke])
                                rijabc_tmp += tmp

                                ks = kshift
                                ke = kconserv[ks,kc,ki]
                                tmp = -einsum('ejab,iec->ijabc',eris.vovv[ke,kj,ka].conj(),r2[ki,ke])
                                rijabc_tmp += tmp

                                eia = numpy.diagonal(foo[ki]).reshape(-1,1) - numpy.diagonal(fvv[ka])
                                eia += _eval
                                ejb = numpy.diagonal(foo[kj]).reshape(-1,1) - numpy.diagonal(fvv[kb])
                                eijab = pyscf.lib.direct_sum('ia,jb->ijab',eia,ejb)

                                kc = kclist[iterki,iterkj,ka,kb]
                                eijabc = pyscf.lib.direct_sum('ijab,c->ijabc',eijab,-numpy.diagonal(fvv[kc]))
                                eijabc = 1./eijabc

                                star_energy += fac*0.5*einsum('ijabc,ijabc,ijabc',lijabc_tmp,rijabc_tmp,eijabc)

            comm.Allreduce(MPI.IN_PLACE, star_energy, op=MPI.SUM)
            if rank == 0:
                logger.info(self, 'EOM-EACCSD* delta energy = %.15g (Imag = %.15g)',
                            star_energy.real,star_energy.imag)
            e.append(star_energy+_eval)

        return np.array(e)

    def amplitudes_to_vector(self, t1, t2):
        return np.hstack((t1.ravel(), t2.ravel()))

    def vector_to_amplitudes(self, vec, nmo=None, nocc=None):
        if nocc is None: nocc = self.nocc
        if nmo is None: nmo = self.nmo
        nvir = nmo - nocc
        nkpts = self.nkpts
        nov = nkpts*nocc*nvir
        t1 = vec[:nov].reshape(nkpts,nocc,nvir)
        t2 = vec[nov:].reshape(nkpts*(nkpts+1)//2,nkpts,nocc,nocc,nvir,nvir)
        return t1, t2

    @property
    def imds(self):
        if self.__imds__ is None:
            self.__imds__ = _IMDS(self)
        return self.__imds__

class _ERIS:
    def __init__(self, cc, mo_coeff=None, method='incore'):
        from pyscf.pbc import tools
        from pyscf.pbc.cc.ccsd import _adjust_occ
        cput0 = (logger.process_clock(), logger.perf_counter())
        cell = cc._scf.cell
        kpts = cc.kpts
        nkpts = cc.nkpts
        nocc = cc.nocc
        nmo = cc.nmo
        nvir = nmo - nocc

        if mo_coeff is None:
            # If mo_coeff is not canonical orbital
            # TODO does this work for k-points? changed to conjugate.
            raise NotImplementedError
            mo_coeff = cc.mo_coeff
        dtype = mo_coeff[0].dtype

        mo_coeff = self.mo_coeff = padded_mo_coeff(cc, mo_coeff)

        # Re-make our fock MO matrix elements from density and fock AO
        dm = cc._scf.make_rdm1(cc.mo_coeff, cc.mo_occ)
        with lib.temporary_env(cc._scf, exxdiv=None):
            # _scf.exxdiv affects eris.fock. HF exchange correction should be
            # excluded from the Fock matrix.
            fockao = cc._scf.get_hcore() + cc._scf.get_veff(cell, dm)
        self.fock = np.asarray([reduce(np.dot, (mo.T.conj(), fockao[k], mo))
                                for k, mo in enumerate(mo_coeff)])
        self.fock = comm.bcast(self.fock, root=0)  # Ensure all processes have same fock

        self.mo_energy = [self.fock[k].diagonal().real for k in range(nkpts)]
        # Add HFX correction in the self.mo_energy to improve convergence in
        # CCSD iteration. It is useful for the 2D systems since their occupied and
        # the virtual orbital energies may overlap which may lead to numerical
        # issue in the CCSD iterations.
        # FIXME: Whether to add this correction for other exxdiv treatments?
        # Without the correction, MP2 energy may be largely off the correct value.
        madelung = tools.madelung(cell, kpts)
        self.mo_energy = [_adjust_occ(mo_e, nocc, -madelung)
                          for k, mo_e in enumerate(self.mo_energy)]

        # Get location of padded elements in occupied and virtual space.
        nocc_per_kpt = get_nocc(cc, per_kpoint=True)
        nonzero_padding = padding_k_idx(cc, kind="joint")

        # Check direct and indirect gaps for possible issues with CCSD convergence.
        mo_e = [self.mo_energy[kp][nonzero_padding[kp]] for kp in range(nkpts)]
        mo_e = np.sort([y for x in mo_e for y in x])  # Sort de-nested array
        gap = mo_e[np.sum(nocc_per_kpt)] - mo_e[np.sum(nocc_per_kpt)-1]
        if gap < 1e-5:
            logger.warn(cc, 'HOMO-LUMO gap %s too small for KCCSD. '
                            'May cause issues in convergence.', gap)

        fao2mo = cc._scf.with_df.ao2mo

        kconserv = cc.kconserv
        khelper = cc.khelper

        log = logger.Logger(cc.stdout, cc.verbose)
        #mem_now = pyscf.lib.current_memory()[0]
        #if (method == 'incore' and (mem_incore+mem_now < cc.max_memory) or cc.mol.incore_anyway):
        if False:
            log.info('using incore ERI storage')
            eri = numpy.zeros((nkpts,nkpts,nkpts,nmo,nmo,nmo,nmo), dtype=numpy.complex128)
            unique_klist = list(khelper.symm_map.keys())
            nUnique_klist = len(unique_klist)

            # Looping over unique list of k-vectors
            loader = mpi_load_balancer.load_balancer(BLKSIZE=(nkpts,))
            loader.set_ranges((range(nUnique_klist),))

            good2go = True
            while(good2go):
                good2go, data = loader.slave_set()
                if good2go is False:
                    break
                index = 0
                block = data[index]
                ranges = loader.outblocks[index][block]
                for indices in ranges:
                    kp, kq, kr = unique_klist[indices]
                    ks = kconserv[kp,kq,kr]
                    eri_kpt = fao2mo((mo_coeff[kp],mo_coeff[kq],mo_coeff[kr],mo_coeff[ks]),
                                     (cc.kpts[kp],cc.kpts[kq],cc.kpts[kr],cc.kpts[ks]))
                    eri_kpt = eri_kpt.reshape(nmo,nmo,nmo,nmo)
                    eri[kp,kq,kr] = eri_kpt.copy()
                    loader.slave_finished()

            comm.Barrier()
            comm.Allreduce(MPI.IN_PLACE, eri, op=MPI.SUM)
            comm.Barrier()

            for kp in range(nkpts):
                for kq in range(nkpts):
                    for kr in range(nkpts):
                        ikp, ikq, ikr = khelper.get_irrVec(kp,kq,kr)
                        irr_eri = eri[ikp,ikq,ikr]
                        eri[kp,kq,kr] = khelper.transform_irr2full(irr_eri,kp,kq,kr)
            comm.Barrier()

            # Chemist -> physics notation
            eri = eri.transpose(0,2,1,3,5,4,6)

            self.dtype = eri.dtype
            self.oooo = eri[:,:,:,:nocc,:nocc,:nocc,:nocc].copy() / nkpts
            self.ooov = eri[:,:,:,:nocc,:nocc,:nocc,nocc:].copy() / nkpts
            self.ovoo = eri[:,:,:,:nocc,nocc:,:nocc,:nocc].copy() / nkpts
            self.oovv = eri[:,:,:,:nocc,:nocc,nocc:,nocc:].copy() / nkpts
            self.ovov = eri[:,:,:,:nocc,nocc:,:nocc,nocc:].copy() / nkpts
            self.ovvv = eri[:,:,:,:nocc,nocc:,nocc:,nocc:].copy() / nkpts
            self.vvvv = eri[:,:,:,nocc:,nocc:,nocc:,nocc:].copy() / nkpts
            #ovvv = eri1[:nocc,nocc:,nocc:,nocc:].copy()
            #self.ovvv = numpy.empty((nocc,nvir,nvir*(nvir+1)//2))
            #for i in range(nocc):
            #    for j in range(nvir):
            #        self.ovvv[i,j] = lib.pack_tril(ovvv[i,j])
            #self.vvvv = pyscf.ao2mo.restore(4, eri1[nocc:,nocc:,nocc:,nocc:], nvir)
        else:
            _tmpfile1_name = None
            if rank == 0:
                _tmpfile1_name = "eris1.hdf5"
            _tmpfile1_name = comm.bcast(_tmpfile1_name, root=0)

            read_feri=False
            if rank == 0:
                if os.path.isfile(_tmpfile1_name):
                    read_feri=True
            read_feri = comm.bcast(read_feri,root=0)

            if read_feri is True:
                self.feri1 = h5py.File(_tmpfile1_name, 'r', driver='mpio', comm=MPI.COMM_WORLD)
                self.oooo  = self.feri1['oooo']
                self.ooov  = self.feri1['ooov']
                self.ovoo  = self.feri1['ovoo']
                self.oovv  = self.feri1['oovv']
                self.ovov  = self.feri1['ovov']
                self.ovvo  = self.feri1['ovvo']
                self.voov  = self.feri1['voov']
                self.ovvv  = self.feri1['ovvv']
                self.vovv  = self.feri1['vovv']
                self.vvvv  = self.feri1['vvvv']

                self.ovovL1  = self.feri1['ovovL1']
                self.ooovL1  = self.feri1['ooovL1']
                #self.ovvvL1  = self.feri1['ovvvL1']
                self.voovR1  = self.feri1['voovR1']
                self.ooovR1  = self.feri1['ooovR1']
                self.vovvR1  = self.feri1['vovvR1']
                self.vovvL1  = self.feri1['vovvL1']
                self.ovovRev  = self.feri1['ovovRev']
                self.ooovRev  = self.feri1['ooovRev']
                self.ovvvRev  = self.feri1['ovvvRev']

                log.warn('using oovv integrals in memory')
                new_oovv = numpy.empty( (nkpts,nkpts,nkpts,nocc,nocc,nvir,nvir), dtype=dtype)
                for kp in range(nkpts):
                    for kq in range(nkpts):
                        for kr in range(nkpts):
                            new_oovv[kp,kq,kr] = self.oovv[kp,kq,kr].copy()
                self.oovv = new_oovv

                #print("lower triangular oovv")
                #self.triu_oovv = numpy.empty( ((nkpts*(nkpts+1))/2,nkpts,nocc,nocc,nvir,nvir), dtype=mo_coeff.dtype)
                #triu_indices = [list(x) for x in numpy.triu_indices(nkpts)]
                #self.triu_oovv = self.oovv[triu_indices]
                return
            comm.Barrier()

            self.feri1 = h5py.File(_tmpfile1_name, 'w', driver='mpio', comm=MPI.COMM_WORLD)

            ds_type = dtype

            self.oooo = self.feri1.create_dataset('oooo', (nkpts,nkpts,nkpts,nocc,nocc,nocc,nocc), dtype=ds_type)
            self.ooov = self.feri1.create_dataset('ooov', (nkpts,nkpts,nkpts,nocc,nocc,nocc,nvir), dtype=ds_type)
            self.ovoo = self.feri1.create_dataset('ovoo', (nkpts,nkpts,nkpts,nocc,nvir,nocc,nocc), dtype=ds_type)
            self.oovv = self.feri1.create_dataset('oovv', (nkpts,nkpts,nkpts,nocc,nocc,nvir,nvir), dtype=ds_type)

            self.ovov = self.feri1.create_dataset('ovov', (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), dtype=ds_type)
            self.ovvo = self.feri1.create_dataset('ovvo', (nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc), dtype=ds_type)
            self.voov = self.feri1.create_dataset('voov', (nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir), dtype=ds_type)
            self.ovvv = self.feri1.create_dataset('ovvv', (nkpts,nkpts,nkpts,nocc,nvir,nvir,nvir), dtype=ds_type)
            self.vovv = self.feri1.create_dataset('vovv', (nkpts,nkpts,nkpts,nvir,nocc,nvir,nvir), dtype=ds_type)
            self.vvvv = self.feri1.create_dataset('vvvv', (nkpts,nkpts,nkpts,nvir,nvir,nvir,nvir), dtype=ds_type)

            self.ovovL1 = self.feri1.create_dataset('ovovL1', (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), dtype=ds_type)
            self.ooovL1 = self.feri1.create_dataset('ooovL1', (nkpts,nkpts,nkpts,nocc,nocc,nocc,nvir), dtype=ds_type)
            #self.ovvvL1 = self.feri1.create_dataset('ovvvL1', (nkpts,nkpts,nkpts,nocc,nvir,nvir,nvir), dtype=ds_type)
            self.ovovRev = self.feri1.create_dataset('ovovRev', (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), dtype=ds_type)
            self.ooovRev = self.feri1.create_dataset('ooovRev', (nkpts,nkpts,nkpts,nocc,nocc,nocc,nvir), dtype=ds_type)
            self.ovvvRev = self.feri1.create_dataset('ovvvRev', (nkpts,nkpts,nkpts,nocc,nvir,nvir,nvir), dtype=ds_type)

            self.voovR1 = self.feri1.create_dataset('voovR1', (nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir), dtype=ds_type)
            self.ooovR1 = self.feri1.create_dataset('ooovR1', (nkpts,nkpts,nkpts,nocc,nocc,nocc,nvir), dtype=ds_type)
            self.vovvR1 = self.feri1.create_dataset('vovvR1', (nkpts,nkpts,nkpts,nvir,nocc,nvir,nvir), dtype=ds_type)
            self.vovvL1 = self.feri1.create_dataset('vovvL1', (nkpts,nkpts,nkpts,nvir,nocc,nvir,nvir), dtype=ds_type)

            mem = 0.5e9
            pre = 1.*nocc*nocc*nmo*nmo*nkpts*16
            nkpts_blksize = min(max(int(numpy.floor(mem/pre)),1),nkpts)

            BLKSIZE = (1,nkpts_blksize,nkpts,)
            if rank == 0:
                log.info("ERI oopq blksize = (%3d %3d %3d)" % BLKSIZE)
            loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
            loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

            def get_norbitals_from_string(string):
                '''Converts a string representing a block of integrals to number of orbitals in that block

                A list of conversions is as follows:
                    'p'    - number of molecule orbitals
                    'o'    - number of occupied orbitals
                    'v'    - number of virtual orbitals
                    'k'    - number of k-points
                '''
                if len(string) == 0:
                    return 0

                chars, counts = np.unique(list(string), return_counts=True)
                norbitals = 1.
                for ic, c in enumerate(chars):
                    if c == 'p':
                        norbitals *= nmo
                    elif c == 'o':
                        norbitals *= nocc
                    elif c == 'v':
                        norbitals *= nvir
                    elif c == 'k':
                        norbitals *= nkpts
                    else:
                        raise ValueError('String %s not allowed.')

                return norbitals

            #array_size = [nkpts, nkpts, nkpts]
            #task_list = generate_max_task_list(array_size,
            #                blk_mem_size=1e9*2*16*get_norbitals_from_string('kkoopp'),
            #                priority_list=[3,2,1])
            #                        for kirange, kjrange in mpi.work_stealing_partition(task_list):

            tmp_block_shape = BLKSIZE + (nocc,nocc,nmo,nmo)
            tmp_block = numpy.empty(shape=tmp_block_shape,dtype=ds_type)
            cput1 = logger.process_clock(), logger.perf_counter()
            good2go = True
            while(good2go):
                good2go, data = loader.slave_set()
                if good2go is False:
                    break
                ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
                #rslice = [slice(0,max(x)-min(x)) for x in (ranges0,ranges1,ranges2)]
                for kp in ranges0:
                    for kq in ranges2:
                        for kr in ranges1:
                            ks = kconserv[kp,kq,kr]
                            orbo_p = mo_coeff[kp][:,:nocc]
                            orbo_r = mo_coeff[kr][:,:nocc]
                            eri_kpt = fao2mo((orbo_p,mo_coeff[kq][:,:],orbo_r,mo_coeff[ks][:,:]),
                                             (cc.kpts[kp],cc.kpts[kq],cc.kpts[kr],cc.kpts[ks]))
                            eri_kpt = eri_kpt.reshape(nocc,nmo,nocc,nmo)
                            eri_kpt = eri_kpt.transpose(0,2,1,3) / nkpts
                            tmp_block[kp-ranges0[0],kr-ranges1[0],kq-ranges2[0]] = eri_kpt

                self.oooo    [min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,:nocc,:nocc]
                self.ooov    [min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,:nocc,nocc:]
                self.ooovL1  [min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1,min(ranges0):max(ranges0)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,:nocc,nocc:].transpose(1,2,0,3,4,5,6)
                self.ooovR1  [min(ranges2):max(ranges2)+1,min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,:nocc,nocc:].transpose(2,0,1,3,4,5,6)
                self.ooovRev [min(ranges2):max(ranges2)+1,min(ranges1):max(ranges1)+1,min(ranges0):max(ranges0)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,:nocc,nocc:].transpose(2,1,0,3,4,5,6)
                self.oovv    [min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,nocc:,nocc:]

                loader.slave_finished()

            comm.Barrier()
            cput1 = log.timer_debug1('transforming oopq', *cput1)

            mem = 0.5e9
            pre = 1.*nocc*nvir*nmo*nmo*nkpts*16
            nkpts_blksize = min(max(int(numpy.floor(mem/pre)),1),nkpts)

            BLKSIZE = (1,nkpts_blksize,nkpts,)
            if rank == 0:
                log.info("ERI ovpq blksize = (%3d %3d %3d)" % BLKSIZE)
            loader1 = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
            loader1.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

            tmp_block_shape = BLKSIZE + (nocc,nvir,nmo,nmo)
            tmp_block  = numpy.empty(shape=tmp_block_shape,dtype=ds_type)

            cput1 = logger.process_clock(), logger.perf_counter()
            good2go = True
            while(good2go):
                good2go, data = loader1.slave_set()
                if good2go is False:
                    break
                ranges0, ranges1, ranges2 = loader1.get_blocks_from_data(data)
                #rslice = [slice(0,len(x)) for x in (ranges0,ranges1,ranges2)]
                for kp in ranges0:
                    for kq in ranges2:
                        for kr in ranges1:
                            ks = kconserv[kp,kq,kr]
                            orbo_p = mo_coeff[kp][:,:nocc]
                            orbv_r = mo_coeff[kr][:,nocc:]
                            eri_kpt = fao2mo((orbo_p,mo_coeff[kq][:,:],orbv_r,mo_coeff[ks][:,:]),
                                             (cc.kpts[kp],cc.kpts[kq],cc.kpts[kr],cc.kpts[ks]))
                            eri_kpt = eri_kpt.reshape(nocc,nmo,nvir,nmo)
                            eri_kpt = eri_kpt.transpose(0,2,1,3) / nkpts
                            tmp_block[kp-ranges0[0],kr-ranges1[0],kq-ranges2[0]] = eri_kpt
                            self.voov[kr,kp,ks] = eri_kpt.transpose(1,0,3,2)[:,:,:nocc,nocc:]
                            self.voovR1[ks,kr,kp] = eri_kpt.transpose(1,0,3,2)[:,:,:nocc,nocc:]
                            self.vovv[kr,kp,ks] = eri_kpt.transpose(1,0,3,2)[:,:,nocc:,nocc:]
                            self.vovvR1[ks,kr,kp] = eri_kpt.transpose(1,0,3,2)[:,:,nocc:,nocc:]
                            self.vovvL1[kp,ks,kr] = eri_kpt.transpose(1,0,3,2)[:,:,nocc:,nocc:]

                self.ovoo[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,:nocc,:nocc]
                self.ovov[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,:nocc,nocc:]
                self.ovovL1[min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1,min(ranges0):max(ranges0)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,:nocc,nocc:].transpose(1,2,0,3,4,5,6)
                self.ovovRev[min(ranges2):max(ranges2)+1,min(ranges1):max(ranges1)+1,min(ranges0):max(ranges0)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,:nocc,nocc:].transpose(2,1,0,3,4,5,6)
                self.ovvo[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,nocc:,:nocc]
                self.ovvv[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,nocc:,nocc:]
                #self.ovvvL1[min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1,min(ranges0):max(ranges0)+1] = \
                #        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,nocc:,nocc:].transpose(1,2,0,3,4,5,6)
                self.ovvvRev[min(ranges2):max(ranges2)+1,min(ranges1):max(ranges1)+1,min(ranges0):max(ranges0)+1] = \
                        tmp_block[:len(ranges0),:len(ranges1),:len(ranges2),:,:,nocc:,nocc:].transpose(2,1,0,3,4,5,6)
                loader1.slave_finished()

            comm.Barrier()
            cput1 = log.timer_debug1('transforming ovpq', *cput1)

            # Here we can exploit the full 4-permutational symmetry
            # for 'vvvv' unlike in the cases above

            mem = 0.5e9
            pre = 1.*nvir*nvir*nvir*nvir*16
            unique_klist = list(khelper.symm_map.keys())
            nUnique_klist = len(unique_klist)
            nkpts_blksize = min(max(int(numpy.floor(mem/pre)),1),nUnique_klist)

            BLKSIZE = (nkpts_blksize,)
            if rank == 0:
                log.info("ERI vvvv blksize = %3d" % nkpts_blksize)
            loader2 = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
            loader2.set_ranges((range(nUnique_klist),))

            tmp_block_shape = BLKSIZE + (nvir,nvir,nvir,nvir)
            tmp_block = numpy.empty(shape=tmp_block_shape,dtype=ds_type)

            good2go = True
            while(good2go):
                good2go, data = loader2.slave_set()
                if good2go is False:
                    break
                ranges = loader2.get_blocks_from_data(data)
                chkpts = [int(numpy.ceil(nUnique_klist/10))*i for i in range(10)]
                for indices in ranges:
                    if indices in chkpts:
                        log.info("vvvv transform is %4.2f percent complete" % (1.*indices/nUnique_klist*100))
                    kp, kq, kr = unique_klist[indices]
                    ks = kconserv[kp,kq,kr]
                    orbv_p = mo_coeff[kp][:,nocc:]
                    orbv_q = mo_coeff[kq][:,nocc:]
                    orbv_r = mo_coeff[kr][:,nocc:]
                    orbv_s = mo_coeff[ks][:,nocc:]
                    eri_kpt = fao2mo((orbv_p,orbv_q,orbv_r,orbv_s),
                                     (cc.kpts[kp],cc.kpts[kq],cc.kpts[kr],cc.kpts[ks]))
                    eri_kpt = eri_kpt.reshape(nvir,nvir,nvir,nvir)
                    eri_kpt = eri_kpt.transpose(0,2,1,3) / nkpts

                    self.vvvv[kp,kr,kq] = eri_kpt.copy()
                    # Store symmetric permutations
                    self.vvvv[kr,kp,ks] = eri_kpt.transpose(1,0,3,2).copy()
                    self.vvvv[kq,ks,kp] = eri_kpt.transpose(2,3,0,1).conj().copy()
                    self.vvvv[ks,kq,kr] = eri_kpt.transpose(3,2,1,0).conj().copy()
                loader2.slave_finished()

            comm.Barrier()
            cput1 = log.timer_debug1('transforming vvvv', *cput1)

            self.feri1.close()
            self.feri1 = h5py.File(_tmpfile1_name, 'r', driver='mpio', comm=MPI.COMM_WORLD)
            self.oooo  = self.feri1['oooo']
            self.ooov  = self.feri1['ooov']
            self.ovoo  = self.feri1['ovoo']
            self.oovv  = self.feri1['oovv']
            self.ovov  = self.feri1['ovov']
            self.ovvo  = self.feri1['ovvo']
            self.voov  = self.feri1['voov']
            self.ovvv  = self.feri1['ovvv']
            self.vovv  = self.feri1['vovv']
            self.vvvv  = self.feri1['vvvv']

            self.ovovL1  = self.feri1['ovovL1']
            self.ooovL1  = self.feri1['ooovL1']
            #self.ovvvL1  = self.feri1['ovvvL1']
            self.voovR1  = self.feri1['voovR1']
            self.ooovR1  = self.feri1['ooovR1']
            self.vovvR1  = self.feri1['vovvR1']
            self.vovvL1  = self.feri1['vovvL1']
            self.ovovRev  = self.feri1['ovovRev']
            self.ooovRev  = self.feri1['ooovRev']
            self.ovvvRev  = self.feri1['ovvvRev']

            log.warn('using oovv integrals in memory')
            new_oovv = numpy.empty( (nkpts,nkpts,nkpts,nocc,nocc,nvir,nvir), dtype=dtype)
            for kp in range(nkpts):
                for kq in range(nkpts):
                    for kr in range(nkpts):
                        new_oovv[kp,kq,kr] = self.oovv[kp,kq,kr].copy()
            self.oovv = new_oovv

        log.timer('CCSD integral transformation', *cput0)

    def __del__(self):
        if getattr(self, 'feri1', None):
            #for key in self.feri1.keys(): del(self.feri1[key])
            self.feri1.close()

CCSD = RCCSD

class _IMDS:
    def __init__(self, cc):
        return

    def make_ip(self,cc):
        #cc = self.cc
        t1,t2,eris = cc.t1, cc.t2, cc.eris
        nkpts,nocc,nvir = t1.shape

        if not getattr(self, 'fint1', None):
            self.fint1 = None

        tmpfile1_name = "eom_intermediates_IP.hdf5"
        self.fint1 = h5py.File(tmpfile1_name, 'w', driver='mpio', comm=MPI.COMM_WORLD)

        ds_type = t2.dtype

        self.Wooov  = self.fint1.create_dataset('Wooov',  (nkpts,nkpts,nkpts,nocc,nocc,nocc,nvir), dtype=ds_type)
        self.Woooo  = self.fint1.create_dataset('Woooo',  (nkpts,nkpts,nkpts,nocc,nocc,nocc,nocc), dtype=ds_type)
        self.WooooS = self.fint1.create_dataset('WooooS',  (nkpts,nkpts,nkpts,nocc,nocc,nocc,nocc), dtype=ds_type)
        self.W1voov = self.fint1.create_dataset('W1voov', (nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir), dtype=ds_type)
        self.W2voov = self.fint1.create_dataset('W2voov', (nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir), dtype=ds_type)
        self.Wvoov  = self.fint1.create_dataset('Wvoov',  (nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir), dtype=ds_type)
        #self.W1ovvo = self.fint1.create_dataset('W1ovvo', (nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc), dtype=ds_type)
        #self.W2ovvo = self.fint1.create_dataset('W2ovvo', (nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc), dtype=ds_type)
        #self.Wovvo  = self.fint1.create_dataset('Wovvo',  (nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc), dtype=ds_type)
        self.W1ovov = self.fint1.create_dataset('W1ovov',  (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), dtype=ds_type)
        self.W2ovov = self.fint1.create_dataset('W2ovov',  (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), dtype=ds_type)
        self.Wovov  = self.fint1.create_dataset('Wovov',   (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), dtype=ds_type)
        self.Wovoo  = self.fint1.create_dataset('Wovoo',  (nkpts,nkpts,nkpts,nocc,nvir,nocc,nocc), dtype=ds_type)

        self.Lvv = imdk.Lvv(cc,t1,t2,eris)
        self.Loo = imdk.Loo(cc,t1,t2,eris)
        self.Fov = imdk.cc_Fov(cc,t1,t2,eris)

        #
        # Order matters here for array creation
        self.Wooov = imdk.Wooov(cc,t1,t2,eris,self.fint1)

        self.W1voov = imdk.W1voov(cc,t1,t2,eris,self.fint1)
        self.W2voov = imdk.W2voov(cc,t1,t2,eris,self.fint1)
        self.Wvoov = imdk.Wvoov(cc,t1,t2,eris,self.fint1)

        #self.W1ovvo = imdk.W1ovvo(cc,t1,t2,eris,self.fint1)
        #self.W2ovvo = imdk.W2ovvo(cc,t1,t2,eris,self.fint1)
        #self.Wovvo = imdk.Wovvo(cc,t1,t2,eris,self.fint1)

        self.Woooo = imdk.Woooo(cc,t1,t2,eris,self.fint1)
        self.WooooS = imdk.WooooS(cc,t1,t2,eris,self.fint1)

        self.W1ovov = imdk.W1ovov(cc,t1,t2,eris,self.fint1)
        self.W2ovov = imdk.W2ovov(cc,t1,t2,eris,self.fint1)
        self.Wovov  = imdk.Wovov(cc,t1,t2,eris,self.fint1)

        self.Woovv = eris.oovv

        self.Wovoo = imdk.Wovoo(cc,t1,t2,eris,self.fint1)

        self.fint1.close()
        self.fint1 = h5py.File(tmpfile1_name, 'r', driver='mpio', comm=MPI.COMM_WORLD)

        self.Wooov  = self.fint1['Wooov' ]
        self.Woooo  = self.fint1['Woooo' ]
        self.WooooS = self.fint1['WooooS' ]
        #self.W1ovvo = self.fint1['W1ovvo']
        #self.W2ovvo = self.fint1['W2ovvo']
        #self.Wovvo  = self.fint1['Wovvo' ]
        self.W1voov = self.fint1['W1voov']
        self.W2voov = self.fint1['W2voov']
        self.Wvoov  = self.fint1['Wvoov' ]
        self.W1ovov = self.fint1['W1ovov']
        self.W2ovov = self.fint1['W2ovov']
        self.Wovov  = self.fint1['Wovov' ]
        self.Wovoo  = self.fint1['Wovoo' ]

    def close_ip(self,cc):
        self.fint1.close()

    def __del__(self):
        if getattr(self, 'fint1', None):
            #for key in self.feri1.keys(): del(self.feri1[key])
            self.fint1.close()
        if getattr(self, 'fint2', None):
            #for key in self.feri1.keys(): del(self.feri1[key])
            self.fint2.close()

    def make_ea(self,cc):
        t1,t2,eris = cc.t1, cc.t2, cc.eris
        nkpts,nocc,nvir = t1.shape

        if not getattr(self, 'fint2', None):
            self.fint2 = None

        tmpfile1_name = "eom_intermediates_EA.hdf5"
        self.fint2 = h5py.File(tmpfile1_name, 'w', driver='mpio', comm=MPI.COMM_WORLD)

        ds_type = t2.dtype

        self.Wooov  = self.fint2.create_dataset('Wooov',  (nkpts,nkpts,nkpts,nocc,nocc,nocc,nvir), dtype=ds_type)
        self.Wvovv  = self.fint2.create_dataset('Wvovv',  (nkpts,nkpts,nkpts,nvir,nocc,nvir,nvir), dtype=ds_type)

        #self.W1ovvo = self.fint2.create_dataset('W1ovvo', (nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc), dtype=ds_type)
        #self.W2ovvo = self.fint2.create_dataset('W2ovvo', (nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc), dtype=ds_type)
        #self.Wovvo  = self.fint2.create_dataset('Wovvo',  (nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc), dtype=ds_type)

        self.W1voov = self.fint2.create_dataset('W1voov', (nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir), dtype=ds_type)
        self.W2voov = self.fint2.create_dataset('W2voov', (nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir), dtype=ds_type)
#
        #self.Wvoov  = self.fint2.create_dataset('Wvoov',  (nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir), dtype=ds_type)
        self.WvoovR1 = self.fint2.create_dataset('WvoovR1',  (nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir), dtype=ds_type)
        self.Wvvvv   = self.fint2.create_dataset('Wvvvv',  (nkpts,nkpts,nkpts,nvir,nvir,nvir,nvir), dtype=ds_type)
        self.W1ovov  = self.fint2.create_dataset('W1ovov',  (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), dtype=ds_type)
        self.W2ovov  = self.fint2.create_dataset('W2ovov',  (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), dtype=ds_type)
#
        #self.Wovov  = self.fint2.create_dataset('Wovov',   (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), dtype=ds_type)
        self.WovovRev  = self.fint2.create_dataset('WovovRev',   (nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir), dtype=ds_type)
#
        #self.Wvvvo  = self.fint2.create_dataset('Wvvvo',  (nkpts,nkpts,nkpts,nvir,nvir,nvir,nocc), dtype=ds_type)
        self.WvvvoR1  = self.fint2.create_dataset('WvvvoR1',  (nkpts,nkpts,nkpts,nvir,nvir,nvir,nocc), dtype=ds_type)

        self.Lvv = imdk.Lvv(cc,t1,t2,eris)
        self.Loo = imdk.Loo(cc,t1,t2,eris)
        self.Fov = imdk.cc_Fov(cc,t1,t2,eris)

        #
        # Order matters here for array creation
        self.Wooov = imdk.Wooov(cc,t1,t2,eris,self.fint2)

        self.Wvovv = imdk.Wvovv(cc,t1,t2,eris,self.fint2)

        self.W1voov = imdk.W1voov(cc,t1,t2,eris,self.fint2)
        self.W2voov = imdk.W2voov(cc,t1,t2,eris,self.fint2)
#
        #self.Wvoov = imdk.Wvoov(cc,t1,t2,eris,self.fint2)
        self.WvoovR1 = imdk.WvoovR1(cc,t1,t2,eris,self.fint2)

        #self.W1ovvo = imdk.W1ovvo(cc,t1,t2,eris,self.fint2)
        #self.W2ovvo = imdk.W2ovvo(cc,t1,t2,eris,self.fint2)
        #self.Wovvo = imdk.Wovvo(cc,t1,t2,eris,self.fint2)

        self.Wvvvv = imdk.Wvvvv(cc,t1,t2,eris,self.fint2)

        self.Woovv = eris.oovv

        self.W1ovov = imdk.W1ovov(cc,t1,t2,eris,self.fint2)
        self.W2ovov = imdk.W2ovov(cc,t1,t2,eris,self.fint2)
#
        #self.Wovov  = imdk.Wovov(cc,t1,t2,eris,self.fint2)
        self.WovovRev  = imdk.WovovRev(cc,t1,t2,eris,self.fint2)

#
        #self.Wvvvo = imdk.Wvvvo(cc,t1,t2,eris,self.fint2)
        self.WvvvoR1 = imdk.WvvvoR1(cc,t1,t2,eris,self.fint2)

        self.fint2.close()
        self.fint2 = h5py.File(tmpfile1_name, 'r', driver='mpio', comm=MPI.COMM_WORLD)

        self.Wooov  = self.fint2['Wooov' ]
        self.Wvovv  = self.fint2['Wvovv' ]
        #self.W1ovvo = self.fint2['W1ovvo']
        #self.W2ovvo = self.fint2['W2ovvo']
        #self.Wovvo  = self.fint2['Wovvo' ]
        self.W1voov = self.fint2['W1voov']
        self.W2voov = self.fint2['W2voov']
#
        #self.Wvoov  = self.fint2['Wvoov' ]
        self.WvoovR1  = self.fint2['WvoovR1' ]
        self.Wvvvv  = self.fint2['Wvvvv' ]
        self.W1ovov = self.fint2['W1ovov']
        self.W2ovov = self.fint2['W2ovov']
#
        #self.Wovov  = self.fint2['Wovov' ]
        self.WovovRev  = self.fint2['WovovRev' ]
#
        #self.Wvvvo  = self.fint2['Wvvvo' ]
        self.WvvvoR1  = self.fint2['WvvvoR1' ]

def _cp(a):
    return np.array(a, copy=False, order='C')



# pyscf/pbc/mpicc/kintermediates_rhf.py

# Copyright 2014-2018,2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy as np

import pyscf.pbc.tools as tools
from mpi4py import MPI
from pyscf.lib import logger
from pyscf.pbc.mpitools import mpi_load_balancer
from pyscf import lib
from pyscf.pbc import lib as pbclib
from pyscf.pbc.tools.tril import tril_index, unpack_tril

comm = MPI.COMM_WORLD

#einsum = np.einsum
einsum = lib.einsum
dot = np.dot

# This is restricted (R)CCSD
# Ref: Hirata et al., J. Chem. Phys. 120, 2581 (2004)

### Eqs. (37)-(39) "kappa"

def cc_tau1(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    BLKSIZE = (1,1,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    tmp_oovv_shape = BLKSIZE + (nocc,nocc,nvir,nvir)
    tmp_oovv  = np.empty(tmp_oovv_shape,dtype=t1.dtype)
    tmp2_oovv = np.zeros_like(tmp_oovv)

    tau1_ooVv = feri2['tau1_ooVv']
    tau1_oOvv = feri2['tau1_oOvv']
    tau1_oovv_rev = feri2['tau1_oovv_rev']
    tau2_Oovv = feri2['tau2_Oovv']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        for iterki, ki in enumerate(ranges0):
            for iterkj, kj in enumerate(ranges1):
                for iterka, ka in enumerate(ranges2):
                    kb = kconserv[ki,ka,kj]
                    tmp_oovv[iterki,iterkj,iterka] = t2[ki,kj,ka].copy()
                    tmp2_oovv[iterki,iterkj,iterka] *= 0.0
                    if ki == ka and kj == kb:
                        tmp2_oovv[iterki,iterkj,iterka] = einsum('ia,jb->ijab',t1[ki],t1[kj])

                    tau1_oovv_rev[kj,ka,kb] = (tmp_oovv[iterki,iterkj,iterka] + tmp2_oovv[iterki,iterkj,iterka])

        tau1_ooVv[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,nvir*min(ranges2):nvir*(max(ranges2)+1)] = \
                (tmp_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] +
                 tmp2_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] ).transpose(0,1,2,5,3,4,6).reshape(len(ranges0),len(ranges1),len(ranges2)*nvir,nocc,nocc,nvir)
        tau1_oOvv[min(ranges0):max(ranges0)+1,min(ranges2):max(ranges2)+1,nocc*min(ranges1):nocc*(max(ranges1)+1)] = \
                (tmp_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] +
                 tmp2_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] ).transpose(0,2,1,4,3,5,6).reshape(len(ranges0),len(ranges2),len(ranges1)*nocc,nocc,nvir,nvir)
        tau2_Oovv[min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1,nocc*min(ranges0):nocc*(max(ranges0)+1)] = \
                (tmp_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] +
                 2*tmp2_oovv[:len(ranges0),:len(ranges1),:len(ranges2)] ).transpose(1,2,0,3,4,5,6).reshape(len(ranges1),len(ranges2),len(ranges0)*nocc,nocc,nvir,nvir)
        loader.slave_finished()
    comm.Barrier()
    return

def cc_Foo(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv
    Fki = np.empty((nkpts,nocc,nocc),dtype=t2.dtype)

    for ki in range(nkpts):
        kk = ki
        Fki[ki] = eris.fock[ki,:nocc,:nocc].copy()
        for kl in range(nkpts):
            for kc in range(nkpts):
                #Fki[ki] += einsum('lkcd,licd->ki',eris.SoOvv[kk,kc],tau1_oOvv[ki,kc])
                kd = kconserv[kk,kc,kl]
                Soovv = 2*eris.oovv[kk,kl,kc] - eris.oovv[kk,kl,kd].transpose(0,1,3,2)
                #Fki[ki] += einsum('klcd,ilcd->ki',Soovv,t2[ki,kl,kc])
                Fki[ki] += einsum('klcd,ilcd->ki',Soovv,unpack_tril(t2,nkpts,ki,kl,kc,kconserv[ki,kc,kl]))
            #if ki == kc:
            kd = kconserv[kk,ki,kl]
            Soovv = 2*eris.oovv[kk,kl,ki] - eris.oovv[kk,kl,kd].transpose(0,1,3,2)
            Fki[ki] += einsum('klcd,ic,ld->ki',Soovv,t1[ki],t1[kl])
    return Fki

def cc_Fvv(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv
    Fac = np.empty((nkpts,nvir,nvir),dtype=t2.dtype)

    for ka in range(nkpts):
        kc = ka
        Fac[ka] = eris.fock[ka,nocc:,nocc:].copy()
        #for kk in range(nkpts):
        #    Fac[ka] += -einsum('lkcd,lkad->ac',eris.SoOvv[kk,kc],tau1_oOvv[kk,ka])
        for kl in range(nkpts):
            for kk in range(nkpts):
                kd = kconserv[kk,kc,kl]
                Soovv = 2*eris.oovv[kk,kl,kc] - eris.oovv[kk,kl,kd].transpose(0,1,3,2)
                #Fac[ka] += -einsum('klcd,klad->ac',Soovv,t2[kk,kl,ka])
                Fac[ka] += -einsum('klcd,klad->ac',Soovv,unpack_tril(t2,nkpts,kk,kl,ka,kconserv[kk,ka,kl]))
            #if kk == ka
            kd = kconserv[ka,kc,kl]
            Soovv = 2*eris.oovv[ka,kl,kc] - eris.oovv[ka,kl,kd].transpose(0,1,3,2)
            Fac[ka] += -einsum('klcd,ka,ld->ac',Soovv,t1[ka],t1[kl])
    return Fac

def cc_Fov(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    Fkc = np.empty((nkpts,nocc,nvir),dtype=t2.dtype)
    Fkc[:] = eris.fock[:,:nocc,nocc:].copy()
    for kk in range(nkpts):
        for kl in range(nkpts):
            Soovv = 2.*eris.oovv[kk,kl,kk] - eris.oovv[kk,kl,kl].transpose(0,1,3,2)
            Fkc[kk] += einsum('klcd,ld->kc',Soovv,t1[kl])
    return Fkc

### Eqs. (40)-(41) "lambda"

def Loo(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    fov = eris.fock[:,:nocc,nocc:]
    Lki = cc_Foo(cc,t1,t2,eris,feri2)
    for ki in range(nkpts):
        Lki[ki] += einsum('kc,ic->ki',fov[ki],t1[ki])
        SoOov = (2*eris.ooov[ki,:,ki] -
                 eris.ooov[:,ki,ki].transpose(0,2,1,3,4)).transpose(0,2,1,3,4).reshape(nkpts*nocc,nocc,nocc,nvir)
        Lki[ki] += einsum('lkic,lc->ki',SoOov,t1.reshape(nkpts*nocc,nvir))
    return Lki

def Lvv(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    fov = eris.fock[:,:nocc,nocc:]
    Lac = cc_Fvv(cc,t1,t2,eris,feri2)
    for ka in range(nkpts):
        Lac[ka] += -einsum('kc,ka->ac',fov[ka],t1[ka])
        for kk in range(nkpts):
            Svovv = 2*eris.ovvv[kk,ka,kk].transpose(1,0,3,2) - eris.ovvv[kk,ka,ka].transpose(1,0,2,3)
            Lac[ka] += einsum('akcd,kd->ac',Svovv,t1[kk])
    return Lac

### Eqs. (42)-(45) "chi"

def cc_Woooo(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    #Wklij = np.array(eris.oooo, copy=True)
    #for pqr in range(nUnique_klist):
    #    kk, kl, ki = unique_klist[pqr]
    BLKSIZE = (1,1,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    oooo_tmp_shape = BLKSIZE + (nocc,nocc,nocc,nocc)
    oooo_tmp = np.empty(shape=oooo_tmp_shape,dtype=t1.dtype)

    tau1_ooVv = feri2['tau1_ooVv']
    #Woooo     = feri2['Woooo']
    Woooo_rev = feri2['Woooo_rev']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        for iterkk, kk in enumerate(ranges0):
            for iterkl, kl in enumerate(ranges1):
                for iterki, ki in enumerate(ranges2):
                    kj = kconserv[kk,ki,kl]
                    oooo_tmp[iterkk,iterkl,iterki] = np.array(eris.oooo[kk,kl,ki],copy=True)
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('klic,jc->klij',eris.ooov[kk,kl,ki],t1[kj])
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('klcj,ic->klij',eris.ooov[kl,kk,kj].transpose(1,0,3,2),t1[ki])

                    # Note the indices and way the tau1 is stored : instead of a loop over kpt='kc' and
                    # loop over mo='c', the (kc,k,l,c,d) index is changed instead to (nkpts*nvir,k,l,d)
                    # so that we only have to loop over the first index, saving read operations.
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('ckld,cijd->klij',eris.ooVv[kk,kl],tau1_ooVv[ki,kj])

                    #for kc in range(nkpts):
                    #    oooo_tmp[iterkk,iterkl,iterki] += einsum('klcd,ijcd->klij',eris.oovv[kk,kl,kc],t2[ki,kj,kc])
                    #oooo_tmp[iterkk,iterkl,iterki] += einsum('klcd,ic,jd->klij',eris.oovv[kk,kl,ki],t1[ki],t1[kj])
                    #Woooo[kk,kl,ki] = oooo_tmp[iterkk,iterkl,iterki]
                    #Woooo[kl,kk,kj] = oooo_tmp[iterkk,iterkl,iterki].transpose(1,0,3,2)
                    Woooo_rev[kl,ki,kj] = oooo_tmp[iterkk,iterkl,iterki]

        #Woooo[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
        #                oooo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]

        # for if you want to take into account symmetry of Woooo integral
        #feri2.Woooo[min(ranges1):max(ranges1)+1,min(ranges0):max(ranges0)+1,min(ranges2):max(ranges2)+1] = \
        #                oooo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)].transpose(0,1,2,4,3,6,5)
        loader.slave_finished()
    comm.Barrier()
    return

def cc_Wvvvv(cc,t1,t2,eris,feri2=None):
    ## Slow:
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv
    vvvv_tmp = np.empty((nvir,nvir,nvir,nvir),dtype=t1.dtype)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=(1,nkpts,))
    loader.set_ranges((range(nkpts),range(nkpts),))

    Wvvvv = feri2['Wvvvv']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1 = loader.get_blocks_from_data(data)
        for ka in ranges0:
            for kc in ranges1:
                for kb in range(ka+1):
                    kd = kconserv[ka,kc,kb]
                    vvvv_tmp = np.array(eris.vvvv[ka,kb,kc],copy=True)
                    vvvv_tmp += einsum('akcd,kb->abcd',eris.ovvv[kb,ka,kd].transpose(1,0,3,2),-t1[kb])
                    vvvv_tmp += einsum('kbcd,ka->abcd',eris.ovvv[ka,kb,kc],-t1[ka])
                    Wvvvv[ka,kb,kc] = vvvv_tmp
                    Wvvvv[kb,ka,kd] = vvvv_tmp.transpose(1,0,3,2)
        loader.slave_finished()

    ## Fast
    #nocc,nvir = t1.shape
    #Wabcd = np.empty((nvir,)*4)
    #for a in range(nvir):
    #    Wabcd[a,:] = einsum('kcd,kb->bcd',eris.vovv[a],-t1)
    ##Wabcd += einsum('kbcd,ka->abcd',eris.ovvv,-t1)
    #Wabcd += lib.dot(-t1.T,eris.ovvv.reshape(nocc,-1)).reshape((nvir,)*4)
    #Wabcd += np.asarray(eris.vvvv)

    comm.Barrier()
    return

def cc_Wvoov(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv
    #Wakic = np.empty((nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir),dtype=t1.dtype)

    BLKSIZE = (1,1,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    voov_tmp_shape = BLKSIZE + (nvir,nocc,nocc,nvir)
    voov_tmp = np.empty(voov_tmp_shape,dtype=t1.dtype)

    tau2_Oovv = feri2['tau2_Oovv']
    #Wvoov     = feri2['Wvoov']
    WvOov     = feri2['WvOov']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        #ix = sum([[min(x),max(x)+1] for x in (ranges0,ranges1,ranges2)], [])
        #eris_ooov = eris.ooov[ix[0]:ix[1], ix[2]:ix[3], ix[4]:ix[5]]
        for iterka, ka in enumerate(ranges0):
            for iterkk, kk in enumerate(ranges1):
                for iterki, ki in enumerate(ranges2):
                    kc = kconserv[ka,ki,kk]
                    voov_tmp[iterka,iterkk,iterki] = np.array(eris.ovvo[kk,ka,kc]).transpose(1,0,3,2)
                    voov_tmp[iterka,iterkk,iterki] -= einsum('lkic,la->akic',eris.ooov[ka,kk,ki],t1[ka])
                    voov_tmp[iterka,iterkk,iterki] += einsum('akdc,id->akic',eris.ovvv[kk,ka,kc].transpose(1,0,3,2),t1[ki])
                    # Beginning of change
                    #for kl in range(nkpts):
                    #    # kl - kd + kk = kc
                    #    # => kd = kl - kc + kk
                    #    kd = kconserv[kl,kc,kk]
                    #    Soovv = 2*np.array(eris.oovv[kl,kk,kd]) - np.array(eris.oovv[kl,kk,kc]).transpose(0,1,3,2)
                    #    voov_tmp[iterka,iterkk,iterki] += 0.5*einsum('lkdc,ilad->akic',Soovv,t2[ki,kl,ka])
                    #    voov_tmp[iterka,iterkk,iterki] -= 0.5*einsum('lkdc,ilda->akic',eris.oovv[kl,kk,kd],t2[ki,kl,kd])
                    #voov_tmp[iterka,iterkk,iterki] -= einsum('lkdc,id,la->akic',eris.oovv[ka,kk,ki],t1[ki],t1[ka])
                    #Wvoov[ka,kk,ki] = voov_tmp[iterka,iterkk,iterki]

                    # Making various intermediates...
                    #t2_oOvv = t2[ki,:,ka].transpose(0,2,1,3,4).reshape(nkpts*nocc,nocc,nvir,nvir)
                    t2_oOvv = unpack_tril(t2,nkpts,ki,range(nkpts),ka,kconserv[ki,ka,range(nkpts)]).transpose(0,2,1,3,4).reshape(nkpts*nocc,nocc,nvir,nvir)
                    #eris_oOvv = eris.oovv[kk,:,kc].transpose(0,2,1,3,4).reshape(nkpts*nocc,nocc,nvir,nvir)

                    voov_tmp[iterka,iterkk,iterki] += 0.5*einsum('lkcd,liad->akic',eris.SoOvv[kk,kc],t2_oOvv)
                    voov_tmp[iterka,iterkk,iterki] -= 0.5*einsum('lkcd,liad->akic',eris.oOvv[kk,kc],tau2_Oovv[ki,ka])

                    # End of change
        #Wvoov[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
        #        voov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        WvOov[min(ranges0):max(ranges0)+1,min(ranges2):max(ranges2)+1,nocc*min(ranges1):nocc*(max(ranges1)+1)] = \
                voov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)].transpose(0,2,1,4,3,5,6).reshape(
                    len(ranges1),len(ranges2),len(ranges0)*nocc,nvir,nocc,nvir)
        loader.slave_finished()
    comm.Barrier()
    return

def cc_Wvovo(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    BLKSIZE = (1,1,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    vovo_tmp_shape = BLKSIZE + (nvir,nocc,nvir,nocc)
    vovo_tmp = np.empty(shape=vovo_tmp_shape,dtype=t1.dtype)

    Wvovo = feri2['Wvovo']
    WvOVo = feri2['WvOVo']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        for iterka, ka in enumerate(ranges0):
            for iterkk, kk in enumerate(ranges1):
                for iterkc, kc in enumerate(ranges2):
                    ki = kconserv[ka,kc,kk]
                    vovo_tmp[iterka,iterkk,iterkc] = np.array(eris.ovov[kk,ka,ki]).transpose(1,0,3,2)
                    vovo_tmp[iterka,iterkk,iterkc] -= einsum('lkci,la->akci',eris.ooov[kk,ka,ki].transpose(1,0,3,2),t1[ka])
                    vovo_tmp[iterka,iterkk,iterkc] += einsum('akcd,id->akci',eris.ovvv[kk,ka,ki].transpose(1,0,3,2),t1[ki])
                    # Beginning of change
                    #for kl in range(nkpts):
                    #    kd = kconserv[kl,kc,kk]
                    #    vovo_tmp[iterka,iterkk,iterkc] -= 0.5*einsum('lkcd,ilda->akci',eris.oovv[kl,kk,kc],t2[ki,kl,kd])
                    #vovo_tmp[iterka,iterkk,iterkc] -= einsum('lkcd,id,la->akci',eris.oovv[ka,kk,kc],t1[ki],t1[ka])
                    #Wvovo[ka,kk,kc] = vovo_tmp[iterka,iterkk,iterkc]

                    oovvf = eris.oovv[:,kk,kc].reshape(nkpts*nocc,nocc,nvir,nvir)
                    #t2f = t2[:,ki,ka].copy() #This is a tau like term
                    t2f = unpack_tril(t2,nkpts,range(nkpts),ki,ka,kconserv[range(nkpts),ka,ki]).copy() #This is a tau like term
                    #for kl in range(nkpts):
                    #    kd = kconserv[kl,kc,kk]
                    #    if ki == kd and kl == ka:
                    #        t2f[kl] += 2*einsum('id,la->liad',t1[ki],t1[ka])
                    kd = kconserv[ka,kc,kk]
                    t2f[ka] += 2*einsum('id,la->liad',t1[kd],t1[ka])
                    t2f = t2f.reshape(nkpts*nocc,nocc,nvir,nvir)
                    vovo_tmp[iterka,iterkk,iterkc] -= 0.5*einsum('lkcd,liad->akci',oovvf,t2f)

        Wvovo[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
                vovo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        WvOVo[min(ranges0):max(ranges0)+1,nocc*min(ranges1):nocc*(max(ranges1)+1),nvir*min(ranges2):nvir*(max(ranges2)+1)] = \
                vovo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)].transpose(0,1,4,2,5,3,6).reshape(
                    len(ranges0),len(ranges1)*nocc,len(ranges2)*nvir,nvir,nocc)
        loader.slave_finished()
    comm.Barrier()
    return

def cc_Wovov(cc,t1,t2,eris,feri2=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    BLKSIZE = (1,1,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))

    ovov_tmp_shape = BLKSIZE + (nocc,nvir,nocc,nvir)
    ovov_tmp = np.empty(shape=ovov_tmp_shape,dtype=t1.dtype)

    #Wovov = feri2['Wovov']
    WOvov = feri2['WOvov']

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)
        for iterkk, kk in enumerate(ranges0):
            for iterka, ka in enumerate(ranges1):
                for iterki, ki in enumerate(ranges2):
                    kc = kconserv[kk,ki,ka]
                    ovov_tmp[iterkk,iterka,iterki] = np.array(eris.ovov[kk,ka,ki],copy=True)
                    ovov_tmp[iterkk,iterka,iterki] -= einsum('lkci,la->kaic',eris.ooov[kk,ka,ki].transpose(1,0,3,2),t1[ka])
                    ovov_tmp[iterkk,iterka,iterki] += einsum('akcd,id->kaic',eris.ovvv[kk,ka,ki].transpose(1,0,3,2),t1[ki])
                    # Beginning of change
                    #for kl in range(nkpts):
                    #    kd = kconserv[kl,kc,kk]
                    #    ovov_tmp[iterka,iterkk,iterkc] -= 0.5*einsum('lkcd,ilda->akci',eris.oovv[kl,kk,kc],t2[ki,kl,kd])
                    #ovov_tmp[iterka,iterkk,iterkc] -= einsum('lkcd,id,la->akci',eris.oovv[ka,kk,kc],t1[ki],t1[ka])
                    #Wvovo[ka,kk,kc] = ovov_tmp[iterka,iterkk,iterkc]

                    oovvf = eris.oovv[:,kk,kc].reshape(nkpts*nocc,nocc,nvir,nvir)
                    #t2f = t2[:,ki,ka].copy() #This is a tau like term
                    t2f = unpack_tril(t2,nkpts,range(nkpts),ki,ka,kconserv[range(nkpts),ka,ki]).copy() #This is a tau like term
                    #for kl in range(nkpts):
                    #    kd = kconserv[kl,kc,kk]
                    #    if ki == kd and kl == ka:
                    #        t2f[kl] += 2*einsum('id,la->liad',t1[ki],t1[ka])
                    kd = kconserv[ka,kc,kk]
                    t2f[ka] += 2*einsum('id,la->liad',t1[kd],t1[ka])
                    t2f = t2f.reshape(nkpts*nocc,nocc,nvir,nvir)
                    ovov_tmp[iterkk,iterka,iterki] -= 0.5*einsum('lkcd,liad->kaic',oovvf,t2f)

        #Wovov[min(ranges0):max(ranges0)+1,min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1] = \
        #        ovov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        WOvov[min(ranges1):max(ranges1)+1,min(ranges2):max(ranges2)+1,nocc*min(ranges0):nocc*(max(ranges0)+1)] = \
                ovov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)].transpose(1,2,0,3,4,5,6).reshape(
                    len(ranges1),len(ranges2),len(ranges0)*nocc,nvir,nocc,nvir)
        loader.slave_finished()
    comm.Barrier()
    return

# EOM Intermediates w/ k-points

def Wooov(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape

    if fint is None:
        Wklid = np.zeros((nkpts,nkpts,nkpts,nocc,nocc,nocc,nvir),dtype=t2.dtype)
    else:
        Wklid = fint['Wooov']

    # TODO can do much better than this... call recursive function
    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    ooov_tmp_size = BLKSIZE + (nocc,nocc,nocc,nvir)
    ooov_tmp = np.empty(ooov_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_ooov_kli = _cp(eris.ooov[s0,s1,s2])
        eris_oovv_kli = _cp(eris.oovv[s0,s1,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterkl,kl in enumerate(ranges1):
                for iterki,ki in enumerate(ranges2):
                    ooov_tmp[iterkk,iterkl,iterki] = eris_ooov_kli[iterkk,iterkl,iterki].copy()
                    ooov_tmp[iterkk,iterkl,iterki] += einsum('ic,klcd->klid', t1[ki],
                                                             eris_oovv_kli[iterkk,iterkl,iterki])
        Wklid[s0,s1,s2] = ooov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wklid, op=MPI.SUM)

    return Wklid


def Wvovv(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape

    if fint is None:
        Walcd = np.zeros((nkpts,nkpts,nkpts,nvir,nocc,nvir,nvir),dtype=t2.dtype)
    else:
        Walcd = fint['Wvovv']

    # TODO can do much better than this... call recursive function
    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nvir*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    vovv_tmp_size = BLKSIZE + (nvir,nocc,nvir,nvir)
    vovv_tmp = np.empty(vovv_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_vovv_alc = _cp(eris.vovv[s0,s1,s2])
        eris_oovv_alc = _cp(eris.oovv[s0,s1,s2])

        for iterka,ka in enumerate(ranges0):
            for iterkl,kl in enumerate(ranges1):
                for iterkc,kc in enumerate(ranges2):
                    # vovv[ka,kl,kc,kd] <= ovvv[kl,ka,kd,kc].transpose(1,0,3,2)
                    vovv_tmp[iterka,iterkl,iterkc] = eris_vovv_alc[iterka,iterkl,iterkc] #np.array(eris.ovvv[kl,ka,kd]).transpose(1,0,3,2)
                    vovv_tmp[iterka,iterkl,iterkc] += -einsum('ka,klcd->alcd',t1[ka],eris_oovv_alc[iterka,iterkl,iterkc])
        Walcd[s0,s1,s2] = vovv_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Walcd, op=MPI.SUM)

    return Walcd

def W1ovvo(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wkaci  = np.zeros((nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc),dtype=t1.dtype)
    else:
        Wkaci  = fint['W1ovvo']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    ovvo_tmp_size = BLKSIZE + (nocc,nvir,nvir,nocc)
    ovvo_tmp = np.empty(ovvo_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

        eris_ovvo_kac = _cp(eris.ovvo[s0,s1,s2])
        eris_oovv_kXc = _cp(eris.oovv[s0,:,s2])
        eris_oovv_Xkc = _cp(eris.oovv[:,s0,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterka,ka in enumerate(ranges1):
                for iterkc,kc in enumerate(ranges2):
                    ki = kconserv[kk,kc,ka]
                    ovvo_tmp[iterkk,iterka,iterkc] = _cp(eris_ovvo_kac[iterkk,iterka,iterkc])
                    #St2 = 2.*t2[ki,:,ka]
                    St2 = 2.*unpack_tril(t2,nkpts,ki,range(nkpts),ka,kconserv[ki,ka,range(nkpts)])
                    #St2 -= t2[:,ki,ka].transpose(0,2,1,3,4)
                    St2 -= unpack_tril(t2,nkpts,range(nkpts),ki,ka,kconserv[range(nkpts),ka,ki]).transpose(0,2,1,3,4)
                    ovvo_tmp[iterkk,iterka,iterkc] += einsum('klcd,ilad->kaci',
                                                             eris_oovv_kXc[iterkk,:,iterkc].transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir),
                                                             St2.transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
                    ovvo_tmp[iterkk,iterka,iterkc] += -einsum('lkcd,ilad->kaci',
                                                              eris_oovv_Xkc[:,iterkk,iterkc].reshape(nocc*nkpts,nocc,nvir,nvir),
                                                              unpack_tril(t2,nkpts,ki,range(nkpts),ka,kconserv[ki,ka,range(nkpts)]).transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
#                                                                t2[ki,:,ka].transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
        Wkaci[s0,s1,s2] = ovvo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]

        loader.slave_finished()

    comm.Barrier()

    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkaci, op=MPI.SUM)

    return Wkaci

def W2ovvo(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wkaci  = np.zeros((nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc),dtype=t1.dtype)
        WWooov = Wooov(cc,t1,t2,eris)
    else:
        Wkaci  = fint['W2ovvo']
        WWooov = fint['Wooov']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nvir*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    ovvo_tmp_size = BLKSIZE + (nocc,nvir,nvir,nocc)
    ovvo_tmp = np.empty(ovvo_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

        Wooov_akX     = _cp(WWooov[s1,s0])
        eris_ovvv_kac = _cp(eris.ovvv[s0,s1,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterka,ka in enumerate(ranges1):
                for iterkc,kc in enumerate(ranges2):
                    ki = kconserv[kk,kc,ka]
                    ovvo_tmp[iterkk,iterka,iterkc] = einsum('la,lkic->kaci',-t1[ka],Wooov_akX[iterka,iterkk,ki])
                    ovvo_tmp[iterkk,iterka,iterkc] += einsum('akdc,id->kaci',
                                                             eris_ovvv_kac[iterkk,iterka,iterkc].transpose(1,0,3,2),t1[ki])

        Wkaci[s0,s1,s2] = ovvo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]

        loader.slave_finished()

    comm.Barrier()

    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkaci, op=MPI.SUM)

    return Wkaci


def Wovvo(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape

    if fint is None:
        Wkaci = np.zeros((nkpts,nkpts,nkpts,nocc,nvir,nvir,nocc),dtype=t2.dtype)
        W1kaci = W1ovvo(cc,t1,t2,eris,fint)
        W2kaci = W2ovvo(cc,t1,t2,eris,fint)
    else:
        Wkaci = fint['Wovvo']
        W1kaci = fint['W1ovvo']
        W2kaci = fint['W2ovvo']

    # TODO can do much better than this... call recursive function
    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        Wkaci[s0,s1,s2] = _cp(W1kaci[s0,s1,s2]) + _cp(W2kaci[s0,s1,s2])

        loader.slave_finished()

    comm.Barrier()

    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkaci, op=MPI.SUM)

    return Wkaci

def W1ovov(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wkbid = np.zeros((nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir),dtype=t2.dtype)
    else:
        Wkbid = fint['W1ovov']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    ovov_tmp_size = BLKSIZE + (nocc,nvir,nocc,nvir)
    ovov_tmp = np.empty(ovov_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_ovov = _cp(eris.ovov[s0,s1,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterkb,kb in enumerate(ranges1):
                for iterki,ki in enumerate(ranges2):
                    kd = kconserv[kk,ki,kb]
                    ovov_tmp[iterkk,iterkb,iterki] = eris_ovov[iterkk,iterkb,iterki].copy()
                    ovov_tmp[iterkk,iterkb,iterki] += -einsum('lkdc,libc->kbid',
                                                              eris.oovv[:,kk,kd].reshape(nkpts*nocc,nocc,nvir,nvir),
                                                              #t2[:,ki,kb].reshape(nkpts*nocc,nocc,nvir,nvir))
                                                              unpack_tril(t2,nkpts,range(nkpts),ki,kb,kconserv[range(nkpts),kb,ki]).reshape(nkpts*nocc,nocc,nvir,nvir))
        Wkbid[s0,s1,s2] = ovov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkbid, op=MPI.SUM)

    return Wkbid

def W2ovov(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape

    if fint is None:
        Wkbid = np.zeros((nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir),dtype=t2.dtype)
        WWooov = Wooov(cc,t1,t2,eris)
    else:
        Wkbid = fint['W2ovov']
        WWooov = fint['Wooov']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nvir*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    ovov_tmp_size = BLKSIZE + (nocc,nvir,nocc,nvir)
    ovov_tmp = np.empty(ovov_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_ovvv  = _cp(eris.ovvv[s0,s1,s2])
        WWooov_kbi = _cp(WWooov[s0,s1,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterkb,kb in enumerate(ranges1):
                for iterki,ki in enumerate(ranges2):
                    ovov_tmp[iterkk,iterkb,iterki] = einsum('klid,lb->kbid',WWooov_kbi[iterkk,iterkb,iterki],-t1[kb])
                    ovov_tmp[iterkk,iterkb,iterki] += einsum('kbcd,ic->kbid',eris_ovvv[iterkk,iterkb,iterki],t1[ki])
        Wkbid[s0,s1,s2] = ovov_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkbid, op=MPI.SUM)

    return Wkbid

def Wovov(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape

    if fint is None:
        Wkbid = np.zeros((nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir),dtype=t2.dtype)
        WW1ovov = W1ovov(cc,t1,t2,eris)
        WW2ovov = W2ovov(cc,t1,t2,eris)
    else:
        Wkbid = fint['Wovov']
        WW1ovov = fint['W1ovov']
        WW2ovov = fint['W2ovov']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

        Wkbid[s0,s1,s2] = _cp(WW1ovov[s0,s1,s2]) + _cp(WW2ovov[s0,s1,s2])

        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkbid, op=MPI.SUM)

    return Wkbid


def WovovRev(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape

    if fint is None:
        Wkbid = np.zeros((nkpts,nkpts,nkpts,nocc,nvir,nocc,nvir),dtype=t2.dtype)
        WW1ovov = W1ovov(cc,t1,t2,eris)
        WW2ovov = W2ovov(cc,t1,t2,eris)
    else:
        Wkbid = fint['WovovRev']
        WW1ovov = fint['W1ovov']
        WW2ovov = fint['W2ovov']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

        Wkbid[s2,s1,s0] = (_cp(WW1ovov[s0,s1,s2]) + _cp(WW2ovov[s0,s1,s2])).transpose(2,1,0,3,4,5,6)

        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkbid, op=MPI.SUM)

    return Wkbid



# This is the same Woooo intermediate used in cc

def Woooo(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wklij = np.zeros((nkpts,nkpts,nkpts,nocc,nocc,nocc,nocc),dtype=t2.dtype)
    else:
        Wklij = fint['Woooo']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    nkpts_blksize2 = 1
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    oooo_tmp_size = BLKSIZE + (nocc,nocc,nocc,nocc)
    oooo_tmp = np.empty(oooo_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_oovv_klX = _cp(eris.oovv[s0,s1,s2])
        eris_oooo_kli = _cp(eris.oooo[s0,s1,s2])
        eris_ooov_klX = _cp(eris.ooov[s0,s1,s2])
        eris_ooov_lkX = _cp(eris.ooov[s1,s0,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterkl,kl in enumerate(ranges1):
                for iterki,ki in enumerate(ranges2):
                    kj = kconserv[kk,ki,kl]
                    #tau1 = t2[ki,kj,:].copy()
                    tau1 = unpack_tril(t2,nkpts,ki,kj,range(nkpts),kconserv[ki,range(nkpts),kj]).copy()
                    tau1[ki] += einsum('ic,jd->ijcd',t1[ki],t1[kj])
                    oooo_tmp[iterkk,iterkl,iterki] = eris_oooo_kli[iterkk,iterkl,iterki].copy()
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('kld,ijd->klij',eris_oovv_klX[iterkk,iterkl,:].transpose(1,2,0,3,4).reshape(nocc,nocc,-1),
                                                             tau1.transpose(1,2,0,3,4).reshape(nocc,nocc,-1))
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('klid,jd->klij',eris_ooov_klX[iterkk,iterkl,ki],t1[kj])
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('lkjc,ic->klij',eris_ooov_lkX[iterkl,iterkk,kj],t1[ki])
        Wklij[s0,s1,s2] = oooo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wklij, op=MPI.SUM)

    return Wklij

# This has different storage compared to Woooo, more amenable to I/O
# Instead of calling Woooo[kk,kl,ki] to get Woooo[kk,kl,ki,kj] you call
# WooooS[kl,ki,kj]
def WooooS(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wklij = np.zeros((nkpts,nkpts,nkpts,nocc,nocc,nocc,nocc),dtype=t2.dtype)
    else:
        Wklij = fint['WooooS']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    nkpts_blksize2 = 1
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    oooo_tmp_size = BLKSIZE + (nocc,nocc,nocc,nocc)
    oooo_tmp = np.empty(oooo_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_oovv_klX = _cp(eris.oovv[s0,s1,s2])
        eris_oooo_kli = _cp(eris.oooo[s0,s1,s2])
        eris_ooov_klX = _cp(eris.ooov[s0,s1,s2])
        eris_ooov_lkX = _cp(eris.ooov[s1,s0,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterkl,kl in enumerate(ranges1):
                for iterki,ki in enumerate(ranges2):
                    kj = kconserv[kk,ki,kl]
                    #tau1 = t2[ki,kj,:].copy()
                    tau1 = unpack_tril(t2,nkpts,ki,kj,range(nkpts),kconserv[ki,range(nkpts),kj]).copy()
                    tau1[ki] += einsum('ic,jd->ijcd',t1[ki],t1[kj])
                    oooo_tmp[iterkk,iterkl,iterki] = eris_oooo_kli[iterkk,iterkl,iterki].copy()
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('kld,ijd->klij',
                                                             eris_oovv_klX[iterkk,iterkl,:].transpose(1,2,0,3,4).reshape(nocc,nocc,-1),
                                                             tau1.transpose(1,2,0,3,4).reshape(nocc,nocc,-1))
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('klid,jd->klij',eris_ooov_klX[iterkk,iterkl,ki],t1[kj])
                    oooo_tmp[iterkk,iterkl,iterki] += einsum('lkjc,ic->klij',eris_ooov_lkX[iterkl,iterkk,kj],t1[ki])
                    Wklij[kl,ki,kj] = oooo_tmp[iterkk,iterkl,iterki]
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wklij, op=MPI.SUM)

    return Wklij

def Wvvvv(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wabcd = np.zeros((nkpts,nkpts,nkpts,nvir,nvir,nvir,nvir),dtype=t2.dtype)
    else:
        Wabcd = fint['Wvvvv']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nvir*nvir*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    nkpts_blksize2 = 1
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    vvvv_tmp_size = BLKSIZE + (nvir,nvir,nvir,nvir)
    vvvv_tmp = np.empty(vvvv_tmp_size,dtype=t2.dtype)

    print("vvvv blksize")
    print(BLKSIZE)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_vovv = _cp(eris.vovv[s0,s1,s2])
        eris_ovvv = _cp(eris.ovvv[s0,s1,s2])
        eris_oovv_abc = _cp(eris.oovv[s0,s1,s2])

        vvvv_tmp = _cp(eris.vvvv[s0,s1,s2])
        for iterka,ka in enumerate(ranges0):
            for iterkb,kb in enumerate(ranges1):
                for iterkc,kc in enumerate(ranges2):
                    kd = kconserv[ka,kc,kb]
                    vvvv_tmp[iterka,iterkb,iterkc] += einsum('klcd,ka,lb->abcd',
                                                             eris_oovv_abc[iterka,iterkb,iterkc],t1[ka],t1[kb])

                    OOvv   = np.empty( (nkpts,nocc,nocc,nvir,nvir), dtype=t2.dtype)
                    t2_tmp = np.empty( (nkpts,nocc,nocc,nvir,nvir), dtype=t2.dtype)
                    #for kk in range(nkpts):
                    #    # kk + kl - kc - kd = 0
                    #    # => kl = kc - kk + kd
                    #    kl = kconserv[kc,kk,kd]
                    #    vvvv_tmp[iterka,iterkb,iterkc] += einsum('klcd,klab->abcd',eris.oovv[kk,kl,kc],t2[kk,kl,ka])
                    for kk in range(nkpts):
                        # kk + kl - kc - kd = 0
                        kl = kconserv[kc,kk,kd]
                        OOvv[kk]   = eris.oovv[kk,kl,kc]
                        #t2_tmp[kk] = t2[kk,kl,ka]
                        t2_tmp[kk] = unpack_tril(t2,nkpts,kk,kl,ka,kconserv[kk,ka,kl])
                    OOvv   = OOvv.reshape(-1,nvir,nvir)
                    t2_tmp = t2_tmp.reshape(-1,nvir,nvir)
                    vvvv_tmp[iterka,iterkb,iterkc] += einsum('xcd,xab->abcd',OOvv,t2_tmp)

                    vvvv_tmp[iterka,iterkb,iterkc] += einsum('alcd,lb->abcd',eris_vovv[iterka,iterkb,iterkc],-t1[kb])
                    #vvvv_tmp[iterka,iterkb,iterkc] += einsum('bkdc,ka->abcd',eris.vovv[kb,ka,kd],-t1[ka])
                    vvvv_tmp[iterka,iterkb,iterkc] += einsum('kbcd,ka->abcd',eris_ovvv[iterka,iterkb,iterkc],-t1[ka])
        Wabcd[s0,s1,s2] = vvvv_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wabcd, op=MPI.SUM)

    return Wabcd

def Wvvvo(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wabcj = np.zeros((nkpts,nkpts,nkpts,nvir,nvir,nvir,nocc),dtype=t2.dtype)
        WWvvvv = Wvvvv(cc,t1,t2,eris)
        WW1ovov = W1ovov(cc,t1,t2,eris)
        FFov = cc_Fov(cc,t1,t2,eris)
    else:
        Wabcj = fint['Wvvvo']
        WWvvvv = fint['Wvvvv']
        WW1ovov = fint['W1ovov']
        WW1voov = fint['W1voov']
        FFov = cc_Fov(cc,t1,t2,eris)

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nvir*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    nkpts_blksize2 = 1
    BLKSIZE = (1,nkpts_blksize2,nkpts_blksize,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    vvvo_tmp_size = BLKSIZE + (nvir,nvir,nvir,nocc)
    vvvo_tmp = np.empty(vvvo_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_vovv_aXc = _cp(eris.vovv[s0,:,s2])
        eris_ovvv_Xac = _cp(eris.ovvv[:,s0,s2])
        eris_ovvv_Xbc = _cp(eris.ovvv[:,s1,s2])

        Wvvvv_abc = _cp(WWvvvv[s0,s1,s2])
        W1voov_abc = _cp(WW1voov[s1,s0,:])
        W1ovov_baX = _cp(WW1ovov[s1,s0,:])

        for iterka,ka in enumerate(ranges0):
            for iterkb,kb in enumerate(ranges1):
                for iterkc,kc in enumerate(ranges2):
                    kj = kconserv[ka,kc,kb]
                    vvvo_tmp[iterka,iterkb,iterkc] = np.array(eris.vovv[kc,kj,ka]).transpose(2,3,0,1).conj()
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('abcd,jd->abcj',Wvvvv_abc[iterka,iterkb,iterkc],t1[kj])
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('lajc,lb->abcj',W1ovov_baX[iterkb,iterka,kj],-t1[kb])
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('bkjc,ka->abcj',W1voov_abc[iterkb,iterka,kj],-t1[ka])

                    St2 = np.empty((nkpts,nocc,nocc,nvir,nvir),dtype=t2.dtype)
                    for kl in range(nkpts):
                        # ka + kl - kc - kd = 0
                        # => kd = ka - kc + kl
                        kd = kconserv[ka,kc,kl]
                        #St2[kl] = 2.*t2[kl,kj,kd]
                        St2[kl] = 2.*unpack_tril(t2,nkpts,kl,kj,kd,kconserv[kl,kd,kj])
                        #St2[kl] -= t2[kl,kj,kb].transpose(0,1,3,2)
                        St2[kl] -= unpack_tril(t2,nkpts,kl,kj,kb,kconserv[kl,kb,kj]).transpose(0,1,3,2)
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('alcd,ljdb->abcj',
                                                             eris_vovv_aXc[iterka,:,iterkc].transpose(1,0,2,3,4).reshape(nvir,nkpts*nocc,nvir,nvir),
                                                             St2.reshape(nkpts*nocc,nocc,nvir,nvir))
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('lacd,jlbd->abcj',
                                                             eris_ovvv_Xac[:,iterka,iterkc].reshape(nkpts*nocc,nvir,nvir,nvir),
                                                             #-t2[kj,:,kb].transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
                                                             -unpack_tril(t2,nkpts,kj,range(nkpts),kb,kconserv[kj,kb,range(nkpts)]).transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('lbcd,ljad->abcj',
                                                             eris_ovvv_Xbc[:,iterkb,iterkc].reshape(nkpts*nocc,nvir,nvir,nvir),
                                                             #-t2[:,kj,ka].reshape(nkpts*nocc,nocc,nvir,nvir))
                                                             -unpack_tril(t2,nkpts,range(nkpts),kj,ka,kconserv[range(nkpts),ka,kj]).reshape(nkpts*nocc,nocc,nvir,nvir))
                    for kl in range(nkpts):
                        kk = kconserv[kb,kl,ka]
                        #vvvo_tmp[iterka,iterkb,iterkc] += einsum('jclk,lkba->abcj',eris.ovoo[kj,kc,kl].conj(),t2[kl,kk,kb])
                        vvvo_tmp[iterka,iterkb,iterkc] += einsum('jclk,lkba->abcj',
                                                                 eris.ovoo[kj,kc,kl].conj(),
                                                                 unpack_tril(t2,nkpts,kl,kk,kb,kconserv[kl,kb,kk]))

                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('lkjc,lb,ka->abcj',eris.ooov[kb,ka,kj],t1[kb],t1[ka])
                    #vvvo_tmp[iterka,iterkb,iterkc] += einsum('lc,ljab->abcj',-FFov[kc],t2[kc,kj,ka])
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('lc,ljab->abcj', -FFov[kc],
                                                             unpack_tril(t2,nkpts,kc,kj,ka,kconserv[kc,ka,kj]))

        Wabcj[s0,s1,s2] = vvvo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wabcj, op=MPI.SUM)

    return Wabcj

def WvvvoR1(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wabcj = np.zeros((nkpts,nkpts,nkpts,nvir,nvir,nvir,nocc),dtype=t2.dtype)
        WWvvvv = Wvvvv(cc,t1,t2,eris)
        WW1ovov = W1ovov(cc,t1,t2,eris)
        FFov = cc_Fov(cc,t1,t2,eris)
    else:
        Wabcj = fint['WvvvoR1']
        WWvvvv = fint['Wvvvv']
        WW1ovov = fint['W1ovov']
        WW1voov = fint['W1voov']
        FFov = cc_Fov(cc,t1,t2,eris)

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nvir*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (1,nkpts_blksize2,nkpts_blksize,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    vvvo_tmp_size = BLKSIZE + (nvir,nvir,nvir,nocc)
    vvvo_tmp = np.empty(vvvo_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        eris_vovv_aXc = _cp(eris.vovv[s0,:,s2])
        eris_ovvv_Xac = _cp(eris.ovvv[:,s0,s2])
        eris_ovvv_Xbc = _cp(eris.ovvv[:,s1,s2])
        eris_vovvR1_cXa = _cp(eris.vovvR1[s0,s2,:])

        Wvvvv_abc = _cp(WWvvvv[s0,s1,s2])
        W1voov_baX = _cp(WW1voov[s1,s0,:])
        W1ovov_baX = _cp(WW1ovov[s1,s0,:])

        for iterka,ka in enumerate(ranges0):
            for iterkb,kb in enumerate(ranges1):
                for iterkc,kc in enumerate(ranges2):
                    kj = kconserv[ka,kc,kb]
                    #vvvo_tmp[iterka,iterkb,iterkc] = np.array(eris.vovv[kc,kj,ka]).transpose(2,3,0,1).conj()
                    vvvo_tmp[iterka,iterkb,iterkc] = np.array(eris_vovvR1_cXa[iterka,iterkc,kj]).transpose(2,3,0,1).conj()
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('abcd,jd->abcj',Wvvvv_abc[iterka,iterkb,iterkc],t1[kj])
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('lajc,lb->abcj',W1ovov_baX[iterkb,iterka,kj],-t1[kb])
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('bkjc,ka->abcj',W1voov_baX[iterkb,iterka,kj],-t1[ka])

                    St2 = np.empty((nkpts,nocc,nocc,nvir,nvir),dtype=t2.dtype)
                    for kl in range(nkpts):
                        # ka + kl - kc - kd = 0
                        # => kd = ka - kc + kl
                        kd = kconserv[ka,kc,kl]
                        #St2[kl] = 2.*t2[kl,kj,kd]
                        St2[kl] = 2.*unpack_tril(t2,nkpts,kl,kj,kd,kconserv[kl,kd,kj])
                        #St2[kl] -= t2[kl,kj,kb].transpose(0,1,3,2)
                        St2[kl] -= unpack_tril(t2,nkpts,kl,kj,kb,kconserv[kl,kb,kj]).transpose(0,1,3,2)
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('alcd,ljdb->abcj',
                                                             eris_vovv_aXc[iterka,:,iterkc].transpose(1,0,2,3,4).reshape(nvir,nkpts*nocc,nvir,nvir),
                                                             St2.reshape(nkpts*nocc,nocc,nvir,nvir))
                    #vvvo_tmp[iterka,iterkb,iterkc] += einsum('alcd,ljdb->abcj',
                    #                                    eris_vovvR1_aXc[iterkc,iterka,:].transpose(1,0,2,3,4).reshape(nvir,nkpts*nocc,nvir,nvir),
                    #                                    St2.reshape(nkpts*nocc,nocc,nvir,nvir))
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('lacd,jlbd->abcj',
                                                             eris_ovvv_Xac[:,iterka,iterkc].reshape(nkpts*nocc,nvir,nvir,nvir),
                                                             -unpack_tril(t2,nkpts,kj,range(nkpts),kb,kconserv[kj,kb,range(nkpts)]).transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
                    #vvvo_tmp[iterka,iterkb,iterkc] += einsum('lacd,jlbd->abcj',
                    #                                    eris_ovvvRev_Xac[iterkc,iterka,:].reshape(nkpts*nocc,nvir,nvir,nvir),
                    #                                    -t2[kj,:,kb].transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('lbcd,ljad->abcj',
                                                             eris_ovvv_Xbc[:,iterkb,iterkc].reshape(nkpts*nocc,nvir,nvir,nvir),
                                                             -unpack_tril(t2,nkpts,range(nkpts),kj,ka,kconserv[range(nkpts),ka,kj]).reshape(nkpts*nocc,nocc,nvir,nvir))
                    #vvvo_tmp[iterka,iterkb,iterkc] += einsum('lbcd,ljad->abcj',
                    #                                    eris_ovvvRev_Xbc[iterkc,iterkb,:].reshape(nkpts*nocc,nvir,nvir,nvir),
                    #                                    -t2[:,kj,ka].reshape(nkpts*nocc,nocc,nvir,nvir))
                    #for kl in range(nkpts):
                    #    kk = kconserv[kb,kl,ka]
                    #    vvvo_tmp[iterka,iterkb,iterkc] += einsum('jclk,lkba->abcj',eris.ovoo[kj,kc,kl].conj(),t2[kl,kk,kb])
                    eris_ovoo_jcX = _cp(eris.ovoo[kj,kc,:])
                    t2_tmp = np.empty( (nkpts,nocc,nocc,nvir,nvir), dtype=t2.dtype)
                    for kl in range(nkpts):
                        kk = kconserv[kb,kl,ka]
                        t2_tmp[kl] = unpack_tril(t2,nkpts,kl,kk,kb,kconserv[kl,kb,kk])
                    #vvvo_tmp[iterka,iterkb,iterkc] += einsum('xjclk,xlkba->abcj',eris.ovoo[kj,kc,kl].conj(),t2[kl,kk,kb])
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('jcx,xba->abcj',
                                                             eris_ovoo_jcX.transpose(1,2,0,3,4).reshape(nocc,nvir,-1).conj(),
                                                             t2_tmp.reshape(-1,nvir,nvir))

                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('lkjc,lb,ka->abcj', eris.ooov[kb,ka,kj],t1[kb],t1[ka])
                    #vvvo_tmp[iterka,iterkb,iterkc] += einsum('lc,ljab->abcj',-FFov[kc],t2[kc,kj,ka])
                    vvvo_tmp[iterka,iterkb,iterkc] += einsum('lc,ljab->abcj', -FFov[kc],
                                                             unpack_tril(t2,nkpts,kc,kj,ka,kconserv[kc,ka,kj]))

        Wabcj[s2,s0,s1] = vvvo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)].transpose(2,0,1,3,4,5,6)
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wabcj, op=MPI.SUM)

    return Wabcj

def Wovoo(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    FFov = cc_Fov(cc,t1,t2,eris)
    if fint is None:
        Wkbij = np.zeros((nkpts,nkpts,nkpts,nocc,nvir,nocc,nocc),dtype=t2.dtype)
        WW1ovov = W1ovov(cc,t1,t2,eris)
        WWoooo = Woooo(cc,t1,t2,eris)
        #WW1ovvo = W1ovvo(cc,t1,t2,eris)
        WW1voov = W1voov(cc,t1,t2,eris)
    else:
        Wkbij = fint['Wovoo']
        WW1ovov = fint['W1ovov']
        WWoooo  = fint['Woooo']
        #WW1ovvo = fint['W1ovvo']
        WW1voov = fint['W1voov']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nvir*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    nkpts_blksize2 = 1
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    print(BLKSIZE)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    ovoo_tmp_size = BLKSIZE + (nocc,nvir,nocc,nocc)
    ovoo_tmp = np.empty(ovoo_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        WW1ovov_kbi = _cp(WW1ovov[s0,s1,s2])
        WWoooo_kbi  = _cp(WWoooo[s0,s1,s2])
        #WW1ovvo_kbi = _cp(WW1ovvo[s0,s1,s2])
        WW1voov_bkX = _cp(WW1voov[s1,s0,:])

        eris_vovv_bkX = _cp(eris.vovv[s1,s0,:])
        eris_ooov_XkX = _cp(eris.ooov[:,s0,:])
        eris_ooov_kXi = _cp(eris.ooov[s0,:,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterkb,kb in enumerate(ranges1):
                for iterki,ki in enumerate(ranges2):
                    kj = kconserv[kk,ki,kb]

                    ovoo_tmp[iterkk,iterkb,iterki] = np.array(eris.ovoo[kk,kb,ki],copy=True)
                    ovoo_tmp[iterkk,iterkb,iterki] += einsum('kbid,jd->kbij',WW1ovov_kbi[iterkk,iterkb,iterki], t1[kj])
                    ovoo_tmp[iterkk,iterkb,iterki] += einsum('klij,lb->kbij', WWoooo_kbi[iterkk,iterkb,iterki],-t1[kb])
                    #ovoo_tmp[iterkk,iterkb,iterki] += einsum('kbcj,ic->kbij',WW1ovvo_kbi[iterkk,iterkb,iterki],t1[ki])
                    ovoo_tmp[iterkk,iterkb,iterki] += einsum('bkjc,ic->kbij',WW1voov_bkX[iterkb,iterkk,kj],t1[ki])

                    ovoo_tmp[iterkk,iterkb,iterki] += einsum('lkid,jlbd->kbij',
                                                             -eris_ooov_XkX[:,iterkk,ki].reshape(nkpts*nocc,nocc,nocc,nvir),
                                                             unpack_tril(t2,nkpts,kj,range(nkpts),kb,kconserv[kj,kb,range(nkpts)]).transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
                    ovoo_tmp[iterkk,iterkb,iterki] += einsum('lkjd,libd->kbij',
                                                             -eris_ooov_XkX[:,iterkk,kj].reshape(nkpts*nocc,nocc,nocc,nvir),
                                                             unpack_tril(t2,nkpts,range(nkpts),ki,kb,kconserv[range(nkpts),kb,ki]).reshape(nkpts*nocc,nocc,nvir,nvir))

                    #St2 = 2.*t2[kj,:,kb]
                    St2 = 2.*unpack_tril(t2,nkpts,kj,range(nkpts),kb,kconserv[kj,kb,range(nkpts)])
                    #St2 -= t2[:,kj,kb].transpose(0,2,1,3,4)
                    St2 -= unpack_tril(t2,nkpts,range(nkpts),kj,kb,kconserv[range(nkpts),kb,kj]).transpose(0,2,1,3,4)
                    St2 = St2.transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir)
                    ovoo_tmp[iterkk,iterkb,iterki] += einsum('klid,jlbd->kbij',
                                                             eris_ooov_kXi[iterkk,:,iterki].transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nocc,nvir),
                                                             St2)

                    #tau1 = t2[kj,ki,:].copy()
                    tau1 = unpack_tril(t2,nkpts,kj,ki,range(nkpts),kconserv[kj,range(nkpts),ki]).copy()
                    tau1[kj] += einsum('jd,ic->jidc',t1[kj],t1[ki])
                    ovoo_tmp[iterkk,iterkb,iterki] += einsum('bkdc,jidc->kbij',
                                                             eris_vovv_bkX[iterkb,iterkk,:].transpose(1,2,0,3,4).reshape(nvir,nocc,nvir*nkpts,nvir),
                                                             tau1.transpose(1,2,0,3,4).reshape(nocc,nocc,nkpts*nvir,nvir))
                    #ovoo_tmp[iterkk,iterkb,iterki] += einsum('kc,ijcb->kbij', FFov[kk],t2[ki,kj,kk])
                    ovoo_tmp[iterkk,iterkb,iterki] += einsum('kc,ijcb->kbij', FFov[kk],
                                                             unpack_tril(t2,nkpts,ki,kj,kk,kconserv[ki,kk,kj]))

        Wkbij[s0,s1,s2] = ovoo_tmp[:len(ranges0),:len(ranges1),:len(ranges2)]
        loader.slave_finished()

    comm.Barrier()
    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkbij, op=MPI.SUM)

    return Wkbij

def W1voov(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wkaci  = np.zeros((nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir),dtype=t1.dtype)
    else:
        Wkaci  = fint['W1voov']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    ovvo_tmp_size = BLKSIZE + (nocc,nvir,nvir,nocc)
    ovvo_tmp = np.empty(ovvo_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

        eris_ovvo_kac = _cp(eris.ovvo[s0,s1,s2])
        eris_oovv_kXc = _cp(eris.oovv[s0,:,s2])
        eris_oovv_Xkc = _cp(eris.oovv[:,s0,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterka,ka in enumerate(ranges1):
                for iterkc,kc in enumerate(ranges2):
                    ki = kconserv[kk,kc,ka]
                    ovvo_tmp[iterkk,iterka,iterkc] = _cp(eris_ovvo_kac[iterkk,iterka,iterkc])
                    #St2 = 2.*t2[ki,:,ka]
                    St2 = 2.*unpack_tril(t2,nkpts,ki,range(nkpts),ka,kconserv[ki,ka,range(nkpts)])
                    #St2 -= t2[:,ki,ka].transpose(0,2,1,3,4)
                    St2 -= unpack_tril(t2,nkpts,range(nkpts),ki,ka,kconserv[range(nkpts),ka,ki]).transpose(0,2,1,3,4)
                    ovvo_tmp[iterkk,iterka,iterkc] += einsum('klcd,ilad->kaci',
                                                             eris_oovv_kXc[iterkk,:,iterkc].transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir),
                                                             St2.transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))
                    ovvo_tmp[iterkk,iterka,iterkc] += -einsum('lkcd,ilad->kaci',
                                                              eris_oovv_Xkc[:,iterkk,iterkc].reshape(nocc*nkpts,nocc,nvir,nvir),
                                                              unpack_tril(t2,nkpts,ki,range(nkpts),ka,kconserv[ki,ka,range(nkpts)]).transpose(1,0,2,3,4).reshape(nocc,nkpts*nocc,nvir,nvir))

                    Wkaci[ka,kk,ki] = ovvo_tmp[iterkk,iterka,iterkc].transpose(1,0,3,2)

        loader.slave_finished()

    comm.Barrier()

    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkaci, op=MPI.SUM)

    return Wkaci

def W2voov(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape
    kconserv = cc.kconserv

    if fint is None:
        Wkaci  = np.zeros((nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir),dtype=t1.dtype)
        WWooov = Wooov(cc,t1,t2,eris)
    else:
        Wkaci  = fint['W2voov']
        WWooov = fint['Wooov']

    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nvir*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here
    ovvo_tmp_size = BLKSIZE + (nocc,nvir,nvir,nocc)
    ovvo_tmp = np.empty(ovvo_tmp_size,dtype=t2.dtype)

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]

        Wooov_akX     = _cp(WWooov[s1,s0])
        eris_ovvv_kac = _cp(eris.ovvv[s0,s1,s2])

        for iterkk,kk in enumerate(ranges0):
            for iterka,ka in enumerate(ranges1):
                for iterkc,kc in enumerate(ranges2):
                    ki = kconserv[kk,kc,ka]
                    ovvo_tmp[iterkk,iterka,iterkc] = einsum('la,lkic->kaci',-t1[ka],Wooov_akX[iterka,iterkk,ki])
                    ovvo_tmp[iterkk,iterka,iterkc] += \
                            einsum('akdc,id->kaci', eris_ovvv_kac[iterkk,iterka,iterkc].transpose(1,0,3,2),t1[ki])

                    Wkaci[ka,kk,ki] = ovvo_tmp[iterkk,iterka,iterkc].transpose(1,0,3,2)

        loader.slave_finished()

    comm.Barrier()

    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkaci, op=MPI.SUM)

    return Wkaci


def Wvoov(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape

    if fint is None:
        Wkaci  = np.zeros((nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir),dtype=t1.dtype)
        W1kaci = W1voov(cc,t1,t2,eris,fint)
        W2kaci = W2voov(cc,t1,t2,eris,fint)
    else:
        Wkaci = fint['Wvoov']
        W1kaci = fint['W1voov']
        W2kaci = fint['W2voov']

    # TODO can do much better than this... call recursive function
    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        Wkaci[s0,s1,s2] = _cp(W1kaci[s0,s1,s2]) + _cp(W2kaci[s0,s1,s2])

        loader.slave_finished()

    comm.Barrier()

    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkaci, op=MPI.SUM)

    return Wkaci

def WvoovR1(cc,t1,t2,eris,fint=None):
    nkpts, nocc, nvir = t1.shape

    if fint is None:
        Wkaci  = np.zeros((nkpts,nkpts,nkpts,nvir,nocc,nocc,nvir),dtype=t1.dtype)
        W1kaci = W1voov(cc,t1,t2,eris,fint)
        W2kaci = W2voov(cc,t1,t2,eris,fint)
    else:
        Wkaci = fint['WvoovR1']
        W1kaci = fint['W1voov']
        W2kaci = fint['W2voov']

    # TODO can do much better than this... call recursive function
    # Adaptive blocking begins here
    mem = 0.5e9
    pre = 1.*nocc*nocc*nvir*nvir*nkpts*16
    nkpts_blksize = min(max(int(np.floor(mem/pre)),1),nkpts)
    nkpts_blksize2 = min(max(int(np.floor(mem/(pre*nkpts_blksize))),1),nkpts)
    BLKSIZE = (nkpts_blksize2,nkpts_blksize,nkpts,)
    loader = mpi_load_balancer.load_balancer(BLKSIZE=BLKSIZE)
    loader.set_ranges((range(nkpts),range(nkpts),range(nkpts),))
    # Adaptive blocking ends here

    good2go = True
    while(good2go):
        good2go, data = loader.slave_set()
        if good2go is False:
            break
        ranges0, ranges1, ranges2 = loader.get_blocks_from_data(data)

        s0,s1,s2 = [slice(min(x),max(x)+1) for x in (ranges0,ranges1,ranges2)]
        Wkaci[s2,s0,s1] = (_cp(W1kaci[s0,s1,s2]) + _cp(W2kaci[s0,s1,s2])).transpose(2,0,1,3,4,5,6)

        loader.slave_finished()

    comm.Barrier()

    if fint is None:
        comm.Allreduce(MPI.IN_PLACE, Wkaci, op=MPI.SUM)

    return Wkaci

def _cp(a):
    return np.array(a, copy=False, order='C')


# pyscf/pbc/mpicc/mpi_kpoint_helper.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors: James D. McClain
#          Timothy Berkelbach <tim.berkelbach@gmail.com>
#

import numpy
import pyscf.pbc.ao2mo
import pyscf.lib
from pyscf.pbc.lib import kpts_helper

DEBUG = 0

class unique_pqr_list:
    #####################################################################################
    # The following only computes the integrals not related by permutational symmetries.
    # Wasn't sure how to do this 'cleanly', but it's fairly straightforward
    #####################################################################################
    def __init__(self,cell,kpts):
        kconserv = kpts_helper.get_kconserv(cell,kpts)
        nkpts = len(kpts)
        temp = range(0,nkpts)
        klist = pyscf.lib.cartesian_prod((temp,temp,temp))
        completed = numpy.zeros((nkpts,nkpts,nkpts),dtype=int)

        self.operations = numpy.zeros((nkpts,nkpts,nkpts),dtype=int)
        self.equivalentList = numpy.zeros((nkpts,nkpts,nkpts,3),dtype=int)
        self.nUnique = 0
        self.uniqueList = numpy.array([],dtype=int)

        ivec = 0
        not_done = True
        while( not_done ):
            current_kvec = klist[ivec]
            # check to see if it's been done...
            kp = current_kvec[0]
            kq = current_kvec[1]
            kr = current_kvec[2]
            #print "computing ",kp,kq,kr
            if completed[kp,kq,kr] == 0:
                self.nUnique += 1
                self.uniqueList = numpy.append(self.uniqueList,current_kvec)
                ks = kconserv[kp,kq,kr]

                # Now find all equivalent kvectors by permuting it all possible ways...
                # and then storing how its related by symmetry
                completed[kp,kq,kr] = 1
                self.operations[kp,kq,kr] = 0
                self.equivalentList[kp,kq,kr] = current_kvec.copy()

                completed[kr,ks,kp] = 1
                self.operations[kr,ks,kp] = 1 #.transpose(2,3,0,1)
                self.equivalentList[kr,ks,kp] = current_kvec.copy()

                completed[kq,kp,ks] = 1
                self.operations[kq,kp,ks] = 2 #numpy.conj(.transpose(1,0,3,2))
                self.equivalentList[kq,kp,ks] = current_kvec.copy()

                completed[ks,kr,kq] = 1
                self.operations[ks,kr,kq] = 3 #numpy.conj(.transpose(3,2,1,0))
                self.equivalentList[ks,kr,kq] = current_kvec.copy()

            ivec += 1
            if ivec == len(klist):
                not_done = False

        self.uniqueList = self.uniqueList.reshape(self.nUnique,-1)
        if DEBUG == 1:
            print("::: kpoint helper :::")
            print("kvector list (in)")
            print("   shape = ", klist.shape)
            print("kvector list (out)")
            print("   shape  = ", self.uniqueList.shape)
            print("   unique list =")
            print(self.uniqueList)
            print("transformation =")
            for i in range(klist.shape[0]):
                pqr = klist[i]
                irr_pqr = self.equivalentList[pqr[0],pqr[1],pqr[2]]
                print("%3d %3d %3d   ->  %3d %3d %3d" % (pqr[0],pqr[1],pqr[2],
                                                         irr_pqr[0],irr_pqr[1],irr_pqr[2]))

    def get_uniqueList(self):
        return self.uniqueList

    def get_irrVec(self,kp,kq,kr):
        return self.equivalentList[kp,kq,kr]

    def get_transformation(self,kp,kq,kr):
        return self.operations[kp,kq,kr]

    ######################################################
    # for the invec created out of our unique list from
    # the irreducible brillouin zone, we transform it to
    # arbitrary kp,kq,kr
    ######################################################
    def transform_irr2full(self,invec,kp,kq,kr):
        operation = self.get_transformation(kp,kq,kr)
        if operation == 0:
            return invec
        if operation == 1:
            return invec.transpose(2,3,0,1)
        if operation == 2:
            return numpy.conj(invec.transpose(1,0,3,2))
        if operation == 3:
            return numpy.conj(invec.transpose(3,2,1,0))


# pyscf/pbc/mpitools/__init__.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# pyscf/pbc/mpitools/mpi.py

#!/usr/bin/env python
# Copyright 2014-2018,2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import sys

import time
import threading
import traceback
import numpy
from mpi4py import MPI
from . import mpi_pool
from .mpi_pool import MPIPool

_registry = {}

if 'pool' not in _registry:
    import atexit
    pool = MPIPool(debug=False)
    _registry['pool'] = pool
    atexit.register(pool.close)

comm = pool.comm
rank = pool.rank

def static_partition(tasks):
    size = len(tasks)
    segsize = (size+pool.size-1) // pool.size
    start = pool.rank * segsize
    stop = min(size, start+segsize)
    return tasks[start:stop]

def work_balanced_partition(tasks, costs=None):
    if costs is None:
        costs = numpy.ones(tasks)
    if rank == 0:
        segsize = float(sum(costs)) / pool.size
        loads = []
        cum_costs = numpy.cumsum(costs)
        start_id = 0
        for k in range(pool.size):
            stop_id = numpy.argmin(abs(cum_costs - (k+1)*segsize)) + 1
            stop_id = max(stop_id, start_id+1)
            loads.append([start_id,stop_id])
            start_id = stop_id
        comm.bcast(loads)
    else:
        loads = comm.bcast()
    if rank < len(loads):
        start, stop = loads[rank]
        return tasks[start:stop]
    else:
        return tasks[:0]

INQUIRY = 50050
TASK = 50051
def work_share_partition(tasks, interval=.02, loadmin=1):
    loadmin = max(loadmin, len(tasks)//50//pool.size)
    rest_tasks = [x for x in tasks[loadmin*pool.size:]]
    tasks = tasks[loadmin*rank:loadmin*rank+loadmin]
    def distribute_task():
        while True:
            load = len(tasks)
            if rank == 0:
                for i in range(pool.size):
                    if i != 0:
                        load = comm.recv(source=i, tag=INQUIRY)
                    if rest_tasks:
                        if load <= loadmin:
                            task = rest_tasks.pop(0)
                            comm.send(task, i, tag=TASK)
                    else:
                        comm.send('OUT_OF_TASK', i, tag=TASK)
            else:
                comm.send(load, 0, tag=INQUIRY)
            if comm.Iprobe(source=0, tag=TASK):
                tasks.append(comm.recv(source=0, tag=TASK))
                if isinstance(tasks[-1], str) and tasks[-1] == 'OUT_OF_TASK':
                    return
            time.sleep(interval)

    tasks_handler = threading.Thread(target=distribute_task)
    tasks_handler.start()

    while True:
        if tasks:
            task = tasks.pop(0)
            if isinstance(task, str) and task == 'OUT_OF_TASK':
                tasks_handler.join()
                return
            yield task

def work_stealing_partition(tasks, interval=.0001):
    tasks = static_partition(tasks)
    out_of_task = [False]
    def task_daemon():
        while True:
            time.sleep(interval)
            while comm.Iprobe(source=MPI.ANY_SOURCE, tag=INQUIRY):
                src, req = comm.recv(source=MPI.ANY_SOURCE, tag=INQUIRY)
                if isinstance(req, str) and req == 'STOP_DAEMON':
                    return
                elif tasks:
                    comm.send(tasks.pop(), src, tag=TASK)
                elif src == 0 and isinstance(req, str) and req == 'ALL_DONE':
                    comm.send(out_of_task[0], src, tag=TASK)
                elif out_of_task[0]:
                    comm.send('OUT_OF_TASK', src, tag=TASK)
                else:
                    comm.send('BYPASS', src, tag=TASK)
    def prepare_to_stop():
        out_of_task[0] = True
        if rank == 0:
            while True:
                done = []
                for i in range(1, pool.size):
                    comm.send((0,'ALL_DONE'), i, tag=INQUIRY)
                    done.append(comm.recv(source=i, tag=TASK))
                if all(done):
                    break
                time.sleep(interval)
            for i in range(pool.size):
                comm.send((0,'STOP_DAEMON'), i, tag=INQUIRY)
        tasks_handler.join()

    if pool.size > 1:
        tasks_handler = threading.Thread(target=task_daemon)
        tasks_handler.start()

    while tasks:
        task = tasks.pop(0)
        yield task

    if pool.size > 1:
        def next_proc(proc):
            proc = (proc+1) % pool.size
            if proc == rank:
                proc = (proc+1) % pool.size
            return proc
        proc_last = (rank + 1) % pool.size
        proc = next_proc(proc_last)

        while True:
            comm.send((rank,None), proc, tag=INQUIRY)
            task = comm.recv(source=proc, tag=TASK)
            if isinstance(task, str) and task == 'OUT_OF_TASK':
                prepare_to_stop()
                return
            elif isinstance(task, str) and task == 'BYPASS':
                if proc == proc_last:
                    prepare_to_stop()
                    return
                else:
                    proc = next_proc(proc)
            else:
                if proc != proc_last:
                    proc_last, proc = proc, next_proc(proc)
                yield task

def bcast(buf, root=0):
    buf = numpy.asarray(buf, order='C')
    shape, dtype = comm.bcast((buf.shape, buf.dtype.char))
    if rank != root:
        buf = numpy.empty(shape, dtype=dtype)
    comm.Bcast(buf, root)
    return buf

## Useful when sending large batches of arrays
#def safe_bcast(buf, root=0):


def reduce(sendbuf, op=MPI.SUM, root=0):
    sendbuf = numpy.asarray(sendbuf, order='C')
    shape, mpi_dtype = comm.bcast((sendbuf.shape, sendbuf.dtype.char))
    _assert(sendbuf.shape == shape and sendbuf.dtype.char == mpi_dtype)

    recvbuf = numpy.zeros_like(sendbuf)
    comm.Reduce(sendbuf, recvbuf, op, root)
    if rank == root:
        return recvbuf
    else:
        return sendbuf

def allreduce(sendbuf, op=MPI.SUM):
    sendbuf = numpy.asarray(sendbuf, order='C')
    shape, mpi_dtype = comm.bcast((sendbuf.shape, sendbuf.dtype.char))
    _assert(sendbuf.shape == shape and sendbuf.dtype.char == mpi_dtype)

    recvbuf = numpy.zeros_like(sendbuf)
    comm.Allreduce(sendbuf, recvbuf, op)
    return recvbuf

def gather(sendbuf, root=0):
    #if pool.debug:
    #    if rank == 0:
    #        res = [sendbuf]
    #        for k in range(1, pool.size):
    #            dat = comm.recv(source=k)
    #            res.append(dat)
    #        return numpy.vstack([x for x in res if len(x) > 0])
    #    else:
    #        comm.send(sendbuf, dest=0)
    #        return sendbuf

    sendbuf = numpy.asarray(sendbuf, order='C')
    mpi_dtype = sendbuf.dtype.char
    if rank == root:
        size_dtype = comm.gather((sendbuf.size, mpi_dtype), root=root)
        _assert(all(x[1] == mpi_dtype for x in size_dtype if x[0] > 0))
        counts = numpy.array([x[0] for x in size_dtype])
        displs = numpy.append(0, numpy.cumsum(counts[:-1]))
        recvbuf = numpy.empty(sum(counts), dtype=sendbuf.dtype)
        comm.Gatherv([sendbuf.ravel(), mpi_dtype],
                     [recvbuf.ravel(), counts, displs, mpi_dtype], root)
        return recvbuf.reshape((-1,) + sendbuf[0].shape)
    else:
        comm.gather((sendbuf.size, mpi_dtype), root=root)
        comm.Gatherv([sendbuf.ravel(), mpi_dtype], None, root)
        return sendbuf

def allgather(sendbuf):
    sendbuf = numpy.asarray(sendbuf, order='C')
    shape, mpi_dtype = comm.bcast((sendbuf.shape, sendbuf.dtype.char))
    _assert(sendbuf.dtype.char == mpi_dtype or sendbuf.size == 0)
    counts = numpy.array(comm.allgather(sendbuf.size))
    displs = numpy.append(0, numpy.cumsum(counts[:-1]))
    recvbuf = numpy.empty(sum(counts), dtype=sendbuf.dtype)
    comm.Allgatherv([sendbuf.ravel(), mpi_dtype],
                    [recvbuf.ravel(), counts, displs, mpi_dtype])
    return recvbuf.reshape((-1,) + shape[1:])

def alltoall(sendbuf, split_recvbuf=False):
    if isinstance(sendbuf, numpy.ndarray):
        mpi_dtype = comm.bcast(sendbuf.dtype.char)
        sendbuf = numpy.asarray(sendbuf, mpi_dtype, 'C')
        nrow = sendbuf.shape[0]
        ncol = sendbuf.size // nrow
        segsize = (nrow+pool.size-1) // pool.size * ncol
        sdispls = numpy.arange(0, pool.size*segsize, segsize)
        sdispls[sdispls>sendbuf.size] = sendbuf.size
        scounts = numpy.append(sdispls[1:]-sdispls[:-1], sendbuf.size-sdispls[-1])
    else:
        assert(len(sendbuf) == pool.size)
        mpi_dtype = comm.bcast(sendbuf[0].dtype.char)
        sendbuf = [numpy.asarray(x, mpi_dtype).ravel() for x in sendbuf]
        scounts = numpy.asarray([x.size for x in sendbuf])
        sdispls = numpy.append(0, numpy.cumsum(scounts[:-1]))
        sendbuf = numpy.hstack(sendbuf)

    rcounts = numpy.asarray(comm.alltoall(scounts))
    rdispls = numpy.append(0, numpy.cumsum(rcounts[:-1]))

    recvbuf = numpy.empty(sum(rcounts), dtype=mpi_dtype)
    comm.Alltoallv([sendbuf.ravel(), scounts, sdispls, mpi_dtype],
                   [recvbuf.ravel(), rcounts, rdispls, mpi_dtype])
    if split_recvbuf:
        return [recvbuf[p0:p0+c] for p0,c in zip(rdispls,rcounts)]
    else:
        return recvbuf

def sendrecv(sendbuf, source=0, dest=0):
    if source == dest:
        return sendbuf

    if rank == source:
        sendbuf = numpy.asarray(sendbuf, order='C')
        comm.send((sendbuf.shape, sendbuf.dtype), dest=dest)
        comm.Send(sendbuf, dest=dest)
        return sendbuf
    elif rank == dest:
        shape, dtype = comm.recv(source=source)
        recvbuf = numpy.empty(shape, dtype=dtype)
        comm.Recv(recvbuf, source=source)
        return recvbuf

def _assert(condition):
    if not condition:
        sys.stderr.write(''.join(traceback.format_stack()[:-1]))
        comm.Abort()

def register_for(obj):
    global _registry
    key = id(obj)
    # Keep track of the object in a global registry.  On slave nodes, the
    # object can be accessed from global registry.
    _registry[key] = obj
    keys = comm.gather(key)
    if rank == 0:
        obj._reg_keys = keys
    return obj

def del_registry(reg_keys):
    if reg_keys:
        def f(reg_keys):
            from mpi4pyscf.tools import mpi
            mpi._registry.pop(reg_keys[mpi.rank])
        pool.apply(f, reg_keys, reg_keys)
    return []


# pyscf/pbc/mpitools/mpi_blksize.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy
from mpi4py import MPI

comm = MPI.COMM_WORLD
nproc = comm.Get_size()

##########################################################################
# TODO: currently if the mem_per_blocksize is too
#       small then it will make everything just be one gigantic block...
#       which may be bad computationally in that if each processor is
#       responsible for work on a block, only one processor will be used
#       for work.
##########################################################################

def get_max_blocksize_from_mem(array_size, mem_per_block, mem, priority_list=None):
    '''
    Args:
        priority_list : list of importance of indices for the blocksizes,
            i.e. an index with priority 1 will be made to be the
            maximum it can be (according to array_size and mem),
            afterwhich an index with priority 2 will be made the maximum size
            it can be.  If two priorities are the same, the block size
            for those indices will be forced to be equal.
    '''
    #assert((priority_list is not None and hasattr(priority_list, '__iter__')) and
    #        "nchunks (int) or priority_list (iterable) must be specified.")
    #print("memory max = %.8e" % mem)
    nindices = len(array_size)
    if priority_list is None:
        _priority_list = [1]*nindices
    else:
        # still fails for a numpy.array(5), with shape == 0 and no len()
        _priority_list = priority_list
    cmem = mem/mem_per_block # current memory to distribute over blocks
    _priority_list = numpy.array(_priority_list)
    _array_size = numpy.array(array_size)
    idx = numpy.argsort(_priority_list)
    idxinv = numpy.argsort(idx) # maps sorted indices back to original
    _priority_list = _priority_list[idx]
    _array_size = _array_size[idx]
    iprior = 0
    chunksize = []
    loop = True
    while( loop ):
        ib = _priority_list[iprior]
        len_b = 1
        for jprior in range(iprior+1,nindices):
            jb = _priority_list[jprior]
            if jb == ib:
                len_b += 1
            else:
                break
        jprior = iprior+len_b
        index_chunks = int(min(min(_array_size[iprior:jprior]),cmem**(1./len_b)))
        iprior = jprior
        index_chunks = max(index_chunks,1)
        for index in range(len_b):
            chunksize.append(index_chunks)
        cmem /= (index_chunks ** len_b)
        if iprior == nindices:
            loop = False
    chunksize = numpy.array(chunksize)[idxinv]
    #print("chunks = ", chunksize)
    #print("mem_per_chunk = %.8e" % (numpy.prod(numpy.asarray(chunksize))*mem_per_block))
    return tuple(chunksize)



# pyscf/pbc/mpitools/mpi_helper.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from .mpi_blksize import get_max_blocksize_from_mem
from pyscf.lib.numpy_helper import cartesian_prod
import numpy
from mpi4py import MPI

MEM_SIZE = 0.5e9

def generate_task_list(chunk_size, array_size):
    segs = [range(int(numpy.ceil(array_size[i]*1./chunk_size[i]))) for i in range(len(array_size))]
    task_id = numpy.array(cartesian_prod(segs))
    task_ranges_lower = task_id * numpy.array(chunk_size)
    task_ranges_upper = numpy.minimum((task_id+1)*numpy.array(chunk_size),array_size)
    return list(numpy.dstack((task_ranges_lower,task_ranges_upper)))

def generate_max_task_list(array_size,blk_mem_size=16.,memory=MEM_SIZE,priority_list=None):
    shape = list(array_size)
    length = len(shape)
    if priority_list is None:
        priority_list = numpy.arange(length)[::-1]
    chunk_size = get_max_blocksize_from_mem(shape,blk_mem_size,memory,priority_list)
    return generate_task_list(chunk_size,shape)

def safeNormDiff(in_array1, in_array2):
    shape = in_array1.shape
    assert shape == in_array2.shape
    length = len(shape)
    # just use 16 for blocksize, size of complex(double)
    chunk_size = get_max_blocksize_from_mem(list(shape),16.,MEM_SIZE,priority_list=numpy.arange(length)[::-1])
    task_list = generate_task_list(chunk_size,shape)
    norm = 0.0
    for block in task_list:
        which_slice = [slice(*x) for x in block]
        norm += numpy.linalg.norm(in_array1[tuple(which_slice)] - in_array2[tuple(which_slice)])
    return norm

def safeAllreduceInPlace(comm, in_array):
    shape = in_array.shape
    length = len(shape)
    # just use 16 for blocksize, size of complex(double)
    chunk_size = get_max_blocksize_from_mem(list(shape),16.,MEM_SIZE,priority_list=numpy.arange(length)[::-1])
    task_list = generate_task_list(chunk_size,shape)
    for block in task_list:
        which_slice = [slice(*x) for x in block]
        tmp = in_array[tuple(which_slice)].copy()
        comm.Allreduce(MPI.IN_PLACE, tmp, op=MPI.SUM)
        in_array[tuple(which_slice)] = tmp

def safeBcastInPlace(comm, in_array, root=0):
    shape = in_array.shape
    length = len(shape)
    # just use 16 for blocksize, size of complex(double)
    chunk_size = get_max_blocksize_from_mem(list(shape),16.,MEM_SIZE,priority_list=numpy.arange(length)[::-1])
    task_list = generate_task_list(chunk_size,shape)
    for block in task_list:
        which_slice = [slice(*x) for x in block]
        tmp = in_array[tuple(which_slice)].copy()
        tmp = comm.bcast(tmp,root=0)
        in_array[tuple(which_slice)] = tmp


# pyscf/pbc/mpitools/mpi_load_balancer.py

# Copyright 2014-2018,2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from mpi4py import MPI
import pyscf.lib
import numpy as np

import sys
from time import sleep

def enum(**enums):
    return type('Enum', (), enums)

tags = enum(WORK=1, WORK_DONE=2, KILL=3)

class load_balancer:
    def __init__(self,BLKSIZE):
        self.rank = MPI.COMM_WORLD.Get_rank()
        self.size = MPI.COMM_WORLD.Get_size()
        self.COMM = MPI.COMM_WORLD
        self.BLKSIZE = BLKSIZE
        self.curr_block = None

    def set_ranges(self,inindices,BLKSIZE=None):
        rank = self.rank
#        if rank == 0:
#            print("Starting new mpi_load_balance...")
#        print("proc ", rank, " setting ranges ", inindices)
        if BLKSIZE is None:
            BLKSIZE = self.BLKSIZE
        if len(BLKSIZE) != len(inindices):
            print("BLKSIZE AND ININDICES MUST HAVE SAME SHAPE!!!!")
            print("rank, BLKSIZE, inindices")
            print(self.rank, BLKSIZE, inindices)
            sys.exit()
        self.nindices = len(inindices)
        ##print("nindices = ", self.nindices)
        ##print("inindices = ", inindices)
        ##print("BLKSIZE = ", BLKSIZE)
        outblocks = []
        for i in range(self.nindices):

            max_range = max(inindices[i])
            min_range = min(inindices[i])
            ranges_size = len(inindices[i])
            nblocks = int(np.ceil(ranges_size/(1.*BLKSIZE[i])))
            ##print("nblocks for range (%3d-%3d) = %3d" % (min_range, max_range, nblocks))
            segment_blocks = []
            for j in range(nblocks):
                block_j_min = min_range + BLKSIZE[i]*j
                block_j_max = min(max_range+1,min_range+BLKSIZE[i]*(j+1))
                segment_blocks.append(range(block_j_min,block_j_max))
#            if rank == 0:
#                print("segment blocks [segment %3d]: " % i)
#                print(segment_blocks)
            outblocks.append(segment_blocks)
        ###print("index 0 block 1")
        ###print(outblocks[0][1])
        self.outblocks = np.asarray(outblocks)
        #if rank == 0:
        #    ##print("final block structure")
        #    ##print(outblocks)
        if rank == 0:
            self.master()

    def master(self):
        status = MPI.Status()
        work = []
        for index in range(0,self.nindices):
            work.append(range(len(self.outblocks[index])))
        block_indices = pyscf.lib.cartesian_prod(work)
        nwork = len(block_indices)
        iwork = 0
        iwork_recieved = 0
        working_procs=[]
        for i in range(1,self.size):
            data = 0
            tag = tags.KILL
            if iwork < nwork:
                data = block_indices[iwork]
                tag = tags.WORK
                working_procs.append(i)
                #print("MASTER : sending out msg to processor ", i, data)
            self.COMM.isend(obj=data, dest=i, tag=tag)
            iwork += 1

        for i in range(iwork,nwork+1):
            self.COMM.Probe(MPI.ANY_SOURCE, tag=tags.WORK_DONE, status=status)
            recieved_from = status.Get_source()
            data = self.COMM.recv(source=recieved_from, tag=tags.WORK_DONE)
            iwork_recieved += 1
            #print("MASTER : just recieved work_done from processor ", recieved_from, " (",iwork_recieved,"/",nwork,")")
            if i == nwork:
                #print("MASTER : returning...")
                break

            data = block_indices[i]
            tag = tags.WORK
            #print("MASTER : sending out new work to processor ", recieved_from, data)
            self.COMM.isend(obj=data, dest=recieved_from, tag=tag)

        for i in range(iwork_recieved,nwork):
            #print("waiting on work...")
            self.COMM.Probe(MPI.ANY_SOURCE, tag=tags.WORK_DONE, status=status)
            recieved_from = status.Get_source()
            data = self.COMM.recv(source=recieved_from, tag=tags.WORK_DONE)

        # You only have to send a kill to the processors that actually did work, otherwise they were sent a kill
        # at the very first loop
        for i in working_procs:
            data = 0
            tag = tags.KILL
            #print("MASTER (ALL_WORK_DONE): sending kill out to rank ", i, data)
            self.COMM.isend(obj=data, dest=i, tag=tag)

        return

    def slave_set(self):
        if self.rank > 0:
            #print("SLAVE : ", self.rank, " starting...")
            status = MPI.Status()
            #print("SLAVE : ", self.rank, " probing for message...")
            self.COMM.Probe(0, MPI.ANY_TAG, status=status)
            #print("SLAVE : ", self.rank, " recieved a message... ", status.Get_tag())
            self.working = True
            if status.Get_tag() == tags.WORK:
                workingBlock = self.COMM.recv(source=0, tag=tags.WORK)
                #print("SLAVE : ", self.rank, " just recieved ", workingBlock)
                self.curr_block = workingBlock
                self.working = True
                return True, workingBlock
            else:
                self.working = False
                workingBlock = self.COMM.recv(source=0, tag=tags.KILL)
                #print("SLAVE : ", self.rank, " dying...")
                return False, 0
        else:
            return False, 0

    def get_blocks_from_data(self,data):
        nindices = len(data)
        if nindices == 1:
            index = 0
            block = data[index]
            ranges = self.outblocks[index][block]
            return ranges
        outranges = []
        for i in range(nindices):
            index = i
            block = data[i]
            ranges = self.outblocks[index][block]
            outranges.append(ranges)
        return outranges

    def slave_finished(self):
        if self.rank > 0:
            if self.working is True:
                ###print(self.rank, " just finished", self.curr_block)
                self.COMM.isend(obj=self.curr_block, dest=0, tag=tags.WORK_DONE)
#                print("SLAVE : ", self.rank, " just finished work on ", self.curr_block, ". sending tag...")
        else:
            return



# pyscf/pbc/mpitools/mpi_pool.py

# -*- coding: utf-8 -*-
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Modified from  https://github.com/adrn/mpipool
#

from __future__ import (division, print_function, absolute_import,
                        unicode_literals)

__all__ = ["MPIPool", "MPIPoolException"]

import os
import sys
import types
import marshal
#import traceback
from mpi4py import MPI

class MPIPool(object):
    """
    A pool that distributes tasks over a set of MPI processes using
    mpi4py. MPI is an API for distributed memory parallelism, used
    by large cluster computers. This class provides a similar interface
    to Python's multiprocessing Pool, but currently only supports the
    :func:`map` method.

    Contributed initially by `Joe Zuntz <https://github.com/joezuntz>`_.

    Parameters
    ----------
    comm : (optional)
        The ``mpi4py`` communicator.

    debug : bool (optional)
        If ``True``, print out a lot of status updates at each step.
    """
    def __init__(self, comm=None, debug=False):
        self.comm = MPI.COMM_WORLD if comm is None else comm
        self.rank = self.comm.Get_rank()
        self.size = self.comm.Get_size()
        self.debug = debug
        self.function = _error_function

        if self.debug:
            import platform
            node = platform.node()
            if self.rank == 0:
                print('Master: host {0} PID {1}'.format(node, os.getpid()))
            else:
                print('Worker: host {0} PID {1}'.format(node, os.getpid()))

    def is_master(self):
        """
        Is the current process the master?

        """
        return self.rank == 0

    def wait(self):
        """
        If this isn't the master process, wait for instructions.

        """
        if self.is_master():
            return

        MPI.Status()

        while True:
            # Event loop.
            # Sit here and await instructions.
            if self.debug:
                print("Worker {0} waiting for task.".format(self.rank))

            # Blocking receive to wait for instructions.
            task = self.comm.bcast()
            if self.debug:
                print("Worker {0} got task {1}.".format(self.rank, task))

            # Check if message is special sentinel signaling end.
            # If so, stop.
            if isinstance(task, _close_pool_message):
                if self.debug:
                    print("Worker {0} close.".format(self.rank))
# Handle global import lock for multithreading, see
#   http://stackoverflow.com/questions/12389526/import-inside-of-a-python-thread
#   https://docs.python.org/3.4/library/imp.html#imp.lock_held
# Global import lock affects the ctypes module.  It leads to deadlock when
# ctypes function is called in new threads created by threading module.
                if sys.version_info < (3,4):
                    import imp
                    if not imp.lock_held():
                        imp.acquire_lock()
                sys.exit(0)

            # Check if message is special type containing new function
            # to be applied
            elif isinstance(task, _function_wrapper):
                code = marshal.loads(task.func_code)
                if self.debug:
                    print('function {0}'.format(code))
                    print("Worker {0} replaced its task function: {1}."
                          .format(self.rank, self.function))
                self.function = types.FunctionType(code, globals())

            else:  # message are function args
                self.function(task)

    def apply(self, function, master_args, worker_args):
        if not self.is_master():
            self.wait()
            exit(0)

        if function is not self.function:
            if self.debug:
                print("Master replacing pool function with {0}."
                      .format(function))

            # Tell all the workers what function to use.
            self.function = function
            F = _function_wrapper(function)
            self.comm.bcast(F)

        # Send all the tasks off and wait for them to be received.
        self.comm.bcast(worker_args)

        result = function(master_args)
        return result

    def close(self):
        """
        Just send a message off to all the pool members which contains
        the special :class:`_close_pool_message` sentinel.

        """
        if self.is_master():
            self.comm.bcast(_close_pool_message())
            if self.debug:
                print('master close')

    def __enter__(self):
        return self

    def __exit__(self):
        self.close()


class _close_pool_message(object):
    def __repr__(self):
        return "<Close pool message>"


class _function_wrapper(object):
    def __init__(self, function):
        #print(function.__closure__)
        self.func_code = marshal.dumps(function.func_code)

def _error_function(task):
    raise RuntimeError("Pool was sent tasks before being told what "
                       "function to apply.")

class MPIPoolException(Exception):
    def __init__(self, tb):
        self.traceback = tb


# pyscf/pbc/scf/__init__.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''Hartree-Fock for periodic systems
'''

from pyscf.pbc.scf import hf
rhf = hf
from pyscf.pbc.scf import uhf
from pyscf.pbc.scf import rohf
from pyscf.pbc.scf import ghf
from pyscf.pbc.scf import khf
krhf = khf
from pyscf.pbc.scf import kuhf
from pyscf.pbc.scf import krohf
from pyscf.pbc.scf import kghf
from pyscf.pbc.scf import newton_ah
from pyscf.pbc.scf import addons

UHF = uhf.UHF
ROHF = rohf.ROHF
GHF = ghf.GHF

def RHF(cell, *args, **kwargs):
    if cell.spin == 0:
        return rhf.RHF(cell, *args, **kwargs)
    else:
        return rohf.ROHF(cell, *args, **kwargs)
RHF.__doc__ = rhf.RHF.__doc__

KRHF = krhf.KRHF  # KRHF supports cell.spin != 0 if number of k-points is even
KUHF = kuhf.KUHF
KROHF = krohf.KROHF
KGHF = kghf.KGHF

newton = newton_ah.newton

def HF(cell, *args, **kwargs):
    if cell.spin == 0:
        return rhf.RHF(cell, *args, **kwargs)
    else:
        return uhf.UHF(cell, *args, **kwargs)

def KHF(cell, *args, **kwargs):
    if cell.spin == 0:
        return krhf.KRHF(cell, *args, **kwargs)
    else:
        return kuhf.KUHF(cell, *args, **kwargs)


def KS(cell, *args, **kwargs):
    from pyscf.pbc import dft
    return dft.KS(cell, *args, **kwargs)

def KKS(cell, *args, **kwargs):
    from pyscf.pbc import dft
    return dft.KKS(cell, *args, **kwargs)

def RKS(cell, *args, **kwargs):
    from pyscf.pbc import dft
    return dft.RKS(cell, *args, **kwargs)

def ROKS(cell, *args, **kwargs):
    from pyscf.pbc import dft
    return dft.ROKS(cell, *args, **kwargs)

def UKS(cell, *args, **kwargs):
    from pyscf.pbc import dft
    return dft.UKS(cell, *args, **kwargs)

def KRKS(cell, *args, **kwargs):
    from pyscf.pbc import dft
    return dft.KRKS(cell, *args, **kwargs)

def KROKS(cell, *args, **kwargs):
    from pyscf.pbc import dft
    return dft.KROKS(cell, *args, **kwargs)

def KUKS(cell, *args, **kwargs):
    from pyscf.pbc import dft
    return dft.KUKS(cell, *args, **kwargs)



# pyscf/pbc/scf/_response_functions.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Generate KSCF response functions
'''

import numpy
from pyscf import lib
from pyscf.pbc.scf import khf, kuhf, krohf, kghf

def _gen_rhf_response(mf, mo_coeff=None, mo_occ=None,
                      singlet=None, hermi=0, max_memory=None):
    from pyscf.pbc.dft import numint, multigrid
    assert(isinstance(mf, khf.KRHF))

    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if mo_occ is None: mo_occ = mf.mo_occ
    cell = mf.cell
    kpts = mf.kpts
    if getattr(mf, 'xc', None) and getattr(mf, '_numint', None):
        ni = mf._numint
        ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)

        omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, spin=cell.spin)
        hybrid = abs(hyb) > 1e-10
        if abs(omega) > 1e-10:  # For range separated Coulomb
            raise NotImplementedError

        if not hybrid and isinstance(mf.with_df, multigrid.MultiGridFFTDF):
            dm0 = mf.make_rdm1(mo_coeff, mo_occ)
            return multigrid._gen_rhf_response(mf, dm0, singlet, hermi)

        if singlet is None:  # for newton solver
            rho0, vxc, fxc = ni.cache_xc_kernel(cell, mf.grids, mf.xc, mo_coeff,
                                                mo_occ, 0, kpts)
        else:
            if isinstance(mo_occ, numpy.ndarray):
                mo_occ = mo_occ*.5
            else:
                mo_occ = [x*.5 for x in mo_occ]
            rho0, vxc, fxc = ni.cache_xc_kernel(cell, mf.grids, mf.xc,
                                                [mo_coeff]*2, [mo_occ]*2,
                                                spin=1, kpts=kpts)
        dm0 = None #mf.make_rdm1(mo_coeff, mo_occ)

        if max_memory is None:
            mem_now = lib.current_memory()[0]
            max_memory = max(2000, mf.max_memory*.8-mem_now)

        if singlet is None:  # Without specify singlet, general case
            def vind(dm1):
                # The singlet hessian
                if hermi == 2:
                    v1 = numpy.zeros_like(dm1)
                else:
                    v1 = ni.nr_rks_fxc(cell, mf.grids, mf.xc, dm0, dm1, 0, hermi,
                                       rho0, vxc, fxc, kpts, max_memory=max_memory)
                if hybrid:
                    if hermi != 2:
                        vj, vk = mf.get_jk(cell, dm1, hermi=hermi, kpts=kpts)
                        v1 += vj - .5 * hyb * vk
                    else:
                        v1 -= .5 * hyb * mf.get_k(cell, dm1, hermi=hermi, kpts=kpts)
                elif hermi != 2:
                    v1 += mf.get_j(cell, dm1, hermi=hermi, kpts=kpts)
                return v1

        elif singlet:
            def vind(dm1):
                if hermi == 2:
                    v1 = numpy.zeros_like(dm1)
                else:
                    # nr_rks_fxc_st requires alpha of dm1
                    v1 = numint.nr_rks_fxc_st(ni, cell, mf.grids, mf.xc, dm0, dm1, 0,
                                              True, rho0, vxc, fxc, kpts,
                                              max_memory=max_memory)
                    v1 *= .5
                if hybrid:
                    if hermi != 2:
                        vj, vk = mf.get_jk(cell, dm1, hermi=hermi, kpts=kpts)
                        v1 += vj - .5 * hyb * vk
                    else:
                        v1 -= .5 * hyb * mf.get_k(cell, dm1, hermi=hermi, kpts=kpts)
                elif hermi != 2:
                    v1 += mf.get_j(cell, dm1, hermi=hermi, kpts=kpts)
                return v1
        else:  # triplet
            def vind(dm1):
                if hermi == 2:
                    v1 = numpy.zeros_like(dm1)
                else:
                    # nr_rks_fxc_st requires alpha of dm1
                    v1 = numint.nr_rks_fxc_st(ni, cell, mf.grids, mf.xc, dm0, dm1, 0,
                                              False, rho0, vxc, fxc, kpts,
                                              max_memory=max_memory)
                    v1 *= .5
                if hybrid:
                    v1 += -.5 * hyb * mf.get_k(cell, dm1, hermi=hermi, kpts=kpts)
                return v1

    else:  # HF
        if (singlet is None or singlet) and hermi != 2:
            def vind(dm1):
                vj, vk = mf.get_jk(cell, dm1, hermi=hermi, kpts=kpts)
                return vj - .5 * vk
        else:
            def vind(dm1):
                return -.5 * mf.get_k(cell, dm1, hermi=hermi, kpts=kpts)

    return vind

def _gen_uhf_response(mf, mo_coeff=None, mo_occ=None,
                      with_j=True, hermi=0, max_memory=None):
    from pyscf.pbc.dft import multigrid
    assert(isinstance(mf, (kuhf.KUHF, krohf.KROHF)))

    if mo_coeff is None: mo_coeff = mf.mo_coeff
    if mo_occ is None: mo_occ = mf.mo_occ
    cell = mf.cell
    kpts = mf.kpts
    if getattr(mf, 'xc', None) and getattr(mf, '_numint', None):
        ni = mf._numint
        ni.libxc.test_deriv_order(mf.xc, 2, raise_error=True)

        omega, alpha, hyb = ni.rsh_and_hybrid_coeff(mf.xc, spin=cell.spin)
        hybrid = abs(hyb) > 1e-10
        if abs(omega) > 1e-10:  # For range separated Coulomb
            raise NotImplementedError

        if not hybrid and isinstance(mf.with_df, multigrid.MultiGridFFTDF):
            dm0 = mf.make_rdm1(mo_coeff, mo_occ)
            return multigrid._gen_uhf_response(mf, dm0, with_j, hermi)

        rho0, vxc, fxc = ni.cache_xc_kernel(cell, mf.grids, mf.xc,
                                            mo_coeff, mo_occ, 1, kpts)
        #dm0 =(numpy.dot(mo_coeff[0]*mo_occ[0], mo_coeff[0].conj().T),
        #      numpy.dot(mo_coeff[1]*mo_occ[1], mo_coeff[1].conj().T))
        dm0 = None

        if max_memory is None:
            mem_now = lib.current_memory()[0]
            max_memory = max(2000, mf.max_memory*.8-mem_now)

        def vind(dm1):
            if hermi == 2:
                v1 = numpy.zeros_like(dm1)
            else:
                v1 = ni.nr_uks_fxc(cell, mf.grids, mf.xc, dm0, dm1, 0, hermi,
                                   rho0, vxc, fxc, kpts, max_memory=max_memory)
            if not hybrid:
                if with_j:
                    vj = mf.get_j(cell, dm1, hermi=hermi, kpts=kpts)
                    v1 += vj[0] + vj[1]
            else:
                if with_j:
                    vj, vk = mf.get_jk(cell, dm1, hermi=hermi, kpts=kpts)
                    v1 += vj[0] + vj[1] - vk * hyb
                else:
                    v1 -= hyb * mf.get_k(cell, dm1, hermi=hermi, kpts=kpts)
            return v1

    elif with_j:
        def vind(dm1):
            vj, vk = mf.get_jk(cell, dm1, hermi=hermi, kpts=kpts)
            v1 = vj[0] + vj[1] - vk
            return v1

    else:
        def vind(dm1):
            return -mf.get_k(cell, dm1, hermi=hermi, kpts=kpts)

    return vind

def _gen_ghf_response(mf, mo_coeff=None, mo_occ=None,
                      with_j=True, hermi=0, max_memory=None):
    '''Generate a function to compute the product of KGHF response function and
    KGHF density matrices.
    '''
    raise NotImplementedError


khf.KRHF.gen_response = _gen_rhf_response
kuhf.KUHF.gen_response = _gen_uhf_response
kghf.KGHF.gen_response = _gen_ghf_response
krohf.KROHF.gen_response = _gen_uhf_response



# pyscf/pbc/scf/addons.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#         Timothy Berkelbach <tim.berkelbach@gmail.com>
#

import copy
from functools import reduce
import numpy
import scipy.linalg
import scipy.special
import scipy.optimize
from pyscf import lib
from pyscf.pbc import gto as pbcgto
from pyscf.lib import logger
from pyscf.scf import addons as mol_addons
from pyscf import __config__

SMEARING_METHOD = getattr(__config__, 'pbc_scf_addons_smearing_method', 'fermi')


def project_mo_nr2nr(cell1, mo1, cell2, kpts=None):
    r''' Project orbital coefficients

    .. math::

        |\psi1> = |AO1> C1

        |\psi2> = P |\psi1> = |AO2>S^{-1}<AO2| AO1> C1 = |AO2> C2

        C2 = S^{-1}<AO2|AO1> C1
    '''
    s22 = cell2.pbc_intor('int1e_ovlp', hermi=1, kpts=kpts)
    s21 = pbcgto.intor_cross('int1e_ovlp', cell2, cell1, kpts=kpts)
    if kpts is None or numpy.shape(kpts) == (3,):  # A single k-point
        return scipy.linalg.solve(s22, s21.dot(mo1), sym_pos=True)
    else:
        assert(len(kpts) == len(mo1))
        return [scipy.linalg.solve(s22[k], s21[k].dot(mo1[k]), sym_pos=True)
                for k, kpt in enumerate(kpts)]


def smearing_(mf, sigma=None, method=SMEARING_METHOD, mu0=None):
    '''Fermi-Dirac or Gaussian smearing'''
    from pyscf.scf import uhf
    from pyscf.scf import ghf
    from pyscf.pbc.scf import khf
    mf_class = mf.__class__
    is_uhf = isinstance(mf, uhf.UHF)
    is_ghf = isinstance(mf, ghf.GHF)
    is_rhf = (not is_uhf) and (not is_ghf)
    is_khf = isinstance(mf, khf.KSCF)

    def fermi_smearing_occ(m, mo_energy_kpts, sigma):
        occ = numpy.zeros_like(mo_energy_kpts)
        de = (mo_energy_kpts - m) / sigma
        occ[de<40] = 1./(numpy.exp(de[de<40])+1.)
        return occ
    def gaussian_smearing_occ(m, mo_energy_kpts, sigma):
        return 0.5 * scipy.special.erfc((mo_energy_kpts - m) / sigma)

    def partition_occ(mo_occ, mo_energy_kpts):
        mo_occ_kpts = []
        p1 = 0
        for e in mo_energy_kpts:
            p0, p1 = p1, p1 + e.size
            occ = mo_occ[p0:p1]
            mo_occ_kpts.append(occ)
        return mo_occ_kpts

    def get_occ(mo_energy_kpts=None, mo_coeff_kpts=None):
        '''Label the occupancies for each orbital for sampled k-points.

        This is a k-point version of scf.hf.SCF.get_occ
        '''
        mo_occ_kpts = mf_class.get_occ(mf, mo_energy_kpts, mo_coeff_kpts)
        if (mf.sigma == 0) or (not mf.sigma) or (not mf.smearing_method):
            return mo_occ_kpts

        if is_khf:
            nkpts = len(mf.kpts)
        else:
            nkpts = 1
        if isinstance(mf.mol, pbcgto.Cell):
            nelectron = mf.mol.tot_electrons(nkpts)
        else:
            nelectron = mf.mol.tot_electrons()
        if is_uhf:
            nocc = nelectron
            mo_es = numpy.append(numpy.hstack(mo_energy_kpts[0]),
                                 numpy.hstack(mo_energy_kpts[1]))
        elif is_ghf:
            nocc = nelectron
            mo_es = numpy.hstack(mo_energy_kpts)
        else:
            nocc = nelectron // 2
            mo_es = numpy.hstack(mo_energy_kpts)

        if mf.smearing_method.lower() == 'fermi':  # Fermi-Dirac smearing
            f_occ = fermi_smearing_occ
        else:  # Gaussian smearing
            f_occ = gaussian_smearing_occ

        mo_energy = numpy.sort(mo_es.ravel())

        # If mu0 is given, fix mu instead of electron number. XXX -Chong Sun
        sigma = mf.sigma
        fermi = mo_energy[nocc-1]
        if mu0 is None:
            def nelec_cost_fn(m):
                mo_occ_kpts = f_occ(m, mo_es, sigma)
                if is_rhf:
                    mo_occ_kpts *= 2
                return (mo_occ_kpts.sum() - nelectron)**2
            res = scipy.optimize.minimize(nelec_cost_fn, fermi, method='Powell')
            mu = res.x
            mo_occs = f = f_occ(mu, mo_es, sigma)
        else:
            mu = mu0
            mo_occs = f = f_occ(mu, mo_es, sigma)

        # See https://www.vasp.at/vasp-workshop/slides/k-points.pdf
        if mf.smearing_method.lower() == 'fermi':
            f = f[(f>0) & (f<1)]
            mf.entropy = -(f*numpy.log(f) + (1-f)*numpy.log(1-f)).sum() / nkpts
        else:
            mf.entropy = (numpy.exp(-((mo_es-mu)/mf.sigma)**2).sum()
                          / (2*numpy.sqrt(numpy.pi)) / nkpts)
        if is_rhf:
            mo_occs *= 2
            mf.entropy *= 2

        # DO NOT use numpy.array for mo_occ_kpts and mo_energy_kpts, they may
        # have different dimensions for different k-points
        if is_uhf:
            if is_khf:
                nao_tot = mo_occs.size//2
                mo_occ_kpts =(partition_occ(mo_occs[:nao_tot], mo_energy_kpts[0]),
                              partition_occ(mo_occs[nao_tot:], mo_energy_kpts[1]))
            else:
                mo_occ_kpts = partition_occ(mo_occs, mo_energy_kpts)
        else: # rhf and ghf
            if is_khf:
                mo_occ_kpts = partition_occ(mo_occs, mo_energy_kpts)
            else:
                mo_occ_kpts = mo_occs

        logger.debug(mf, '    Fermi level %g  Sum mo_occ_kpts = %s  should equal nelec = %s',
                     fermi, mo_occs.sum(), nelectron)
        logger.info(mf, '    sigma = %g  Optimized mu = %.12g  entropy = %.12g',
                    mf.sigma, mu, mf.entropy)

        return mo_occ_kpts

    def get_grad_tril(mo_coeff_kpts, mo_occ_kpts, fock):
        if is_khf:
            grad_kpts = []
            for k, mo in enumerate(mo_coeff_kpts):
                f_mo = reduce(numpy.dot, (mo.T.conj(), fock[k], mo))
                nmo = f_mo.shape[0]
                grad_kpts.append(f_mo[numpy.tril_indices(nmo, -1)])
            return numpy.hstack(grad_kpts)
        else:
            f_mo = reduce(numpy.dot, (mo_coeff_kpts.T.conj(), fock, mo_coeff_kpts))
            nmo = f_mo.shape[0]
            return f_mo[numpy.tril_indices(nmo, -1)]

    def get_grad(mo_coeff_kpts, mo_occ_kpts, fock=None):
        if (mf.sigma == 0) or (not mf.sigma) or (not mf.smearing_method):
            return mf_class.get_grad(mf, mo_coeff_kpts, mo_occ_kpts, fock)
        if fock is None:
            dm1 = mf.make_rdm1(mo_coeff_kpts, mo_occ_kpts)
            fock = mf.get_hcore() + mf.get_veff(mf.mol, dm1)
        if is_uhf:
            ga = get_grad_tril(mo_coeff_kpts[0], mo_occ_kpts[0], fock[0])
            gb = get_grad_tril(mo_coeff_kpts[1], mo_occ_kpts[1], fock[1])
            return numpy.hstack((ga,gb))
        else: # rhf and ghf
            return get_grad_tril(mo_coeff_kpts, mo_occ_kpts, fock)

    def energy_tot(dm=None, h1e=None, vhf=None):
        e_tot = mf.energy_elec(dm, h1e, vhf)[0] + mf.energy_nuc()
        if (mf.sigma and mf.smearing_method and
            mf.entropy is not None):
            mf.e_free = e_tot - mf.sigma * mf.entropy
            mf.e_zero = e_tot - mf.sigma * mf.entropy * .5
            logger.info(mf, '    Total E(T) = %.15g  Free energy = %.15g  E0 = %.15g',
                        e_tot, mf.e_free, mf.e_zero)
        return e_tot

    mf.sigma = sigma
    mf.smearing_method = method
    mf.entropy = None
    mf.e_free = None
    mf.e_zero = None
    mf._keys = mf._keys.union(['sigma', 'smearing_method',
                               'entropy', 'e_free', 'e_zero'])

    mf.get_occ = get_occ
    mf.energy_tot = energy_tot
    mf.get_grad = get_grad
    return mf

def canonical_occ_(mf, nelec=None):
    '''Label the occupancies for each orbital for sampled k-points.
    This is for KUHF objects.
    Each k-point has a fixed number of up and down electrons in this,
    which results in a finite size error for metallic systems
    but can accelerate convergence.
    '''
    from pyscf.pbc.scf import kuhf
    assert(isinstance(mf, kuhf.KUHF))

    def get_occ(mo_energy_kpts=None, mo_coeff=None):
        if mo_energy_kpts is None: mo_energy_kpts = mf.mo_energy

        if nelec is None:
            cell_nelec = mf.cell.nelec
        else:
            cell_nelec = nelec

        homo=[-1e8,-1e8]
        lumo=[1e8,1e8]
        mo_occ_kpts = [[], []]
        for s in [0,1]:
            for k, mo_energy in enumerate(mo_energy_kpts[s]):
                e_idx = numpy.argsort(mo_energy)
                e_sort = mo_energy[e_idx]
                n = cell_nelec[s]
                mo_occ = numpy.zeros_like(mo_energy)
                mo_occ[e_idx[:n]] = 1
                homo[s] = max(homo[s], e_sort[n-1])
                lumo[s] = min(lumo[s], e_sort[n])
                mo_occ_kpts[s].append(mo_occ)

        for nm,s in zip(['alpha','beta'],[0,1]):
            logger.info(mf, nm+' HOMO = %.12g  LUMO = %.12g', homo[s], lumo[s])
            if homo[s] > lumo[s]:
                logger.warn(mf, "WARNING! HOMO is greater than LUMO! "
                            "This may lead to incorrect canonical occupation.")

        return mo_occ_kpts

    mf.get_occ = get_occ
    return mf
canonical_occ = canonical_occ_


def convert_to_uhf(mf, out=None):
    '''Convert the given mean-field object to the corresponding unrestricted
    HF/KS object
    '''
    from pyscf.pbc import scf
    from pyscf.pbc import dft

    if out is None:
        if isinstance(mf, (scf.uhf.UHF, scf.kuhf.KUHF)):
            return copy.copy(mf)
        else:
            unknown_cls = [scf.kghf.KGHF]
            for i, cls in enumerate(mf.__class__.__mro__):
                if cls in unknown_cls:
                    raise NotImplementedError(
                        "No conversion from %s to uhf object" % cls)

            known_cls = {dft.krks.KRKS  : dft.kuks.KUKS,
                         dft.kroks.KROKS: dft.kuks.KUKS,
                         scf.khf.KRHF   : scf.kuhf.KUHF,
                         scf.krohf.KROHF: scf.kuhf.KUHF,
                         dft.rks.RKS    : dft.uks.UKS  ,
                         dft.roks.ROKS  : dft.uks.UKS  ,
                         scf.hf.RHF     : scf.uhf.UHF  ,
                         scf.rohf.ROHF  : scf.uhf.UHF  ,}
            # .with_df should never be removed or changed during the conversion.
            # It is needed to compute JK matrix in all pbc SCF objects
            out = mol_addons._object_without_soscf(mf, known_cls, remove_df=False)
    else:
        assert(isinstance(out, (scf.uhf.UHF, scf.kuhf.KUHF)))
        if isinstance(mf, scf.khf.KSCF):
            assert(isinstance(out, scf.khf.KSCF))
        else:
            assert(not isinstance(out, scf.khf.KSCF))

    out = mol_addons.convert_to_uhf(mf, out, False)
    # Manually update .with_df because this attribute may not be passed to the
    # output object correctly in molecular convert function
    out.with_df = mf.with_df
    return out

def convert_to_rhf(mf, out=None):
    '''Convert the given mean-field object to the corresponding restricted
    HF/KS object
    '''
    from pyscf.pbc import scf
    from pyscf.pbc import dft

    if getattr(mf, 'nelec', None) is None:
        nelec = mf.cell.nelec
    else:
        nelec = mf.nelec

    if out is not None:
        assert(isinstance(out, (scf.hf.RHF, scf.khf.KRHF)))
        if isinstance(mf, scf.khf.KSCF):
            assert(isinstance(out, scf.khf.KSCF))
        else:
            assert(not isinstance(out, scf.khf.KSCF))

    elif nelec[0] != nelec[1] and isinstance(mf, scf.rohf.ROHF):
        if getattr(mf, '_scf', None):
            return mol_addons._update_mf_without_soscf(mf, copy.copy(mf._scf), False)
        else:
            return copy.copy(mf)

    else:
        if isinstance(mf, (scf.hf.RHF, scf.khf.KRHF)):
            return copy.copy(mf)
        else:
            unknown_cls = [scf.kghf.KGHF]
            for i, cls in enumerate(mf.__class__.__mro__):
                if cls in unknown_cls:
                    raise NotImplementedError(
                        "No conversion from %s to rhf object" % cls)

            if nelec[0] == nelec[1]:
                known_cls = {dft.kuks.KUKS : dft.krks.KRKS,
                             scf.kuhf.KUHF : scf.khf.KRHF ,
                             dft.uks.UKS   : dft.rks.RKS  ,
                             scf.uhf.UHF   : scf.hf.RHF   ,
                             dft.kroks.KROKS : dft.krks.KRKS,
                             scf.krohf.KROHF : scf.khf.KRHF ,
                             dft.roks.ROKS   : dft.rks.RKS  ,
                             scf.rohf.ROHF   : scf.hf.RHF   }
            else:
                known_cls = {dft.kuks.KUKS : dft.krks.KROKS,
                             scf.kuhf.KUHF : scf.khf.KROHF ,
                             dft.uks.UKS   : dft.rks.ROKS  ,
                             scf.uhf.UHF   : scf.hf.ROHF   }
            # .with_df should never be removed or changed during the conversion.
            # It is needed to compute JK matrix in all pbc SCF objects
            out = mol_addons._object_without_soscf(mf, known_cls, remove_df=False)

    out = mol_addons.convert_to_rhf(mf, out, False)
    # Manually update .with_df because this attribute may not be passed to the
    # output object correctly in molecular convert function
    out.with_df = mf.with_df
    return out

def convert_to_ghf(mf, out=None):
    '''Convert the given mean-field object to the generalized HF/KS object

    Args:
        mf : SCF object

    Returns:
        An generalized SCF object
    '''
    from pyscf.pbc import scf

    if out is not None:
        assert(isinstance(out, (scf.ghf.GHF, scf.kghf.KGHF)))
        if isinstance(mf, scf.khf.KSCF):
            assert(isinstance(out, scf.khf.KSCF))
        else:
            assert(not isinstance(out, scf.khf.KSCF))

    if isinstance(mf, scf.ghf.GHF):
        if out is None:
            return copy.copy(mf)
        else:
            out.__dict__.update(mf.__dict__)
            return out

    elif isinstance(mf, scf.khf.KSCF):

        def update_mo_(mf, mf1):
            _keys = mf._keys.union(mf1._keys)
            mf1.__dict__.update(mf.__dict__)
            mf1._keys = _keys
            if mf.mo_energy is not None:
                mf1.mo_energy = []
                mf1.mo_occ = []
                mf1.mo_coeff = []
                nkpts = len(mf.kpts)
                is_rhf = isinstance(mf, scf.hf.RHF)
                for k in range(nkpts):
                    if is_rhf:
                        mo_a = mo_b = mf.mo_coeff[k]
                        ea = getattr(mf.mo_energy[k], 'mo_ea', mf.mo_energy[k])
                        eb = getattr(mf.mo_energy[k], 'mo_eb', mf.mo_energy[k])
                        occa = mf.mo_occ[k] > 0
                        occb = mf.mo_occ[k] == 2
                        orbspin = mol_addons.get_ghf_orbspin(ea, mf.mo_occ[k], True)
                    else:
                        mo_a = mf.mo_coeff[0][k]
                        mo_b = mf.mo_coeff[1][k]
                        ea = mf.mo_energy[0][k]
                        eb = mf.mo_energy[1][k]
                        occa = mf.mo_occ[0][k]
                        occb = mf.mo_occ[1][k]
                        orbspin = mol_addons.get_ghf_orbspin((ea, eb), (occa, occb), False)

                    nao, nmo = mo_a.shape

                    mo_energy = numpy.empty(nmo*2)
                    mo_energy[orbspin==0] = ea
                    mo_energy[orbspin==1] = eb
                    mo_occ = numpy.empty(nmo*2)
                    mo_occ[orbspin==0] = occa
                    mo_occ[orbspin==1] = occb

                    mo_coeff = numpy.zeros((nao*2,nmo*2), dtype=mo_a.dtype)
                    mo_coeff[:nao,orbspin==0] = mo_a
                    mo_coeff[nao:,orbspin==1] = mo_b
                    mo_coeff = lib.tag_array(mo_coeff, orbspin=orbspin)

                    mf1.mo_energy.append(mo_energy)
                    mf1.mo_occ.append(mo_occ)
                    mf1.mo_coeff.append(mo_coeff)

            return mf1

        if out is None:
            out = scf.kghf.KGHF(mf.cell)
        return update_mo_(mf, out)

    else:
        if out is None:
            out = scf.ghf.GHF(mf.cell)
        out = mol_addons.convert_to_ghf(mf, out, remove_df=False)
        # Manually update .with_df because this attribute may not be passed to the
        # output object correctly in molecular convert function
        out.with_df = mf.with_df
        return out

def convert_to_khf(mf, out=None):
    '''Convert gamma point SCF object to k-point SCF object
    '''
    from pyscf.pbc import scf, dft
    if not isinstance(mf, scf.khf.KSCF):
        known_cls = {dft.uks.UKS   : dft.kuks.KUKS,
                     scf.uhf.UHF   : scf.kuhf.KUHF,
                     dft.rks.RKS   : dft.krks.KRKS,
                     scf.hf.RHF    : scf.khf.KRHF,
                     dft.roks.ROKS : dft.kroks.KROKS,
                     scf.rohf.ROHF : scf.krohf.KROHF,
                     #TODO: dft.gks.GKS   : dft.kgks.KGKS,
                     scf.ghf.GHF   : scf.kghf.KGHF}
        # Keep the attribute with_df
        with_df = mf.with_df
        mf = mol_addons._object_without_soscf(mf, known_cls, remove_df=False)
        mf.with_df = with_df

    if out is None:
        return mf
    else:
        out.__dict__.update(mf.__dict__)
        return out

del(SMEARING_METHOD)


if __name__ == '__main__':
    import pyscf.pbc.scf as pscf
    cell = pbcgto.Cell()
    cell.atom = '''
    He 0 0 1
    He 1 0 1
    '''
    cell.basis = 'ccpvdz'
    cell.a = numpy.eye(3) * 4
    cell.mesh = [17] * 3
    cell.verbose = 4
    cell.build()
    nks = [2,1,1]
    mf = pscf.KUHF(cell, cell.make_kpts(nks))
    #mf = smearing_(mf, .1) # -5.86052594663696
    mf = smearing_(mf, .1, mu0=0.280911009667) # -5.86052594663696
    #mf = smearing_(mf, .1, method='gauss')
    mf.kernel()


# pyscf/pbc/scf/chkfile.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

from pyscf.lib.chkfile import load_chkfile_key, load
from pyscf.lib.chkfile import dump_chkfile_key, dump, save
from pyscf.pbc.lib.chkfile import load_cell, save_cell
from pyscf.scf.chkfile import dump_scf

def load_scf(chkfile):
    return load_cell(chkfile), load(chkfile, 'scf')



# pyscf/pbc/scf/cphf.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Xing Zhang <zhangxing.nju@gmail.com>
#

'''
Restricted coupled pertubed Hartree-Fock solver
Modified from pyscf.scf.cphf
'''


import numpy as np
from pyscf import lib
from pyscf.lib import logger


def solve(fvind, mo_energy, mo_occ, h1, s1=None,
          max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
    '''
    Args:
        fvind : function
            Given density matrix, compute (ij|kl)D_{lk}*2 - (ij|kl)D_{jk}

    Kwargs:
        hermi : boolean
            Whether the matrix defined by fvind is Hermitian or not.
    '''
    if s1 is None:
        return solve_nos1(fvind, mo_energy, mo_occ, h1,
                          max_cycle, tol, hermi, verbose)
    else:
        return solve_withs1(fvind, mo_energy, mo_occ, h1, s1,
                            max_cycle, tol, hermi, verbose)
kernel = solve

# h1 shape is (:,nvir,nocc)
def solve_nos1(fvind, mo_energy, mo_occ, h1,
               max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
    '''For field independent basis. First order overlap matrix is zero'''
    log = logger.new_logger(verbose=verbose)
    t0 = (logger.process_clock(), logger.perf_counter())

    nkpt = len(h1)
    moloc = np.zeros([nkpt+1], dtype=int)
    for k in range(nkpt):
        moloc[k+1] = moloc[k] + h1[k].size

    occidx = []
    viridx = []
    for k in range(nkpt):
        occidx.append(mo_occ[k] > 0)
        viridx.append(mo_occ[k] == 0)

    e_a = [mo_energy[k][viridx[k]] for k in range(nkpt)]
    e_i = [mo_energy[k][occidx[k]] for k in range(nkpt)]
    e_ai = [1 / lib.direct_sum('a-i->ai', e_a[k], e_i[k]) for k in range(nkpt)]
    mo1base = []
    for k in range(nkpt):
        mo1base.append((h1[k] * -e_ai[k]).ravel())
    mo1base = np.hstack(mo1base)

    def vind_vo(mo1):
        mo1 = mo1.flatten()
        tmp = []
        for k in range(nkpt):
            tmp.append(mo1[moloc[k]:moloc[k+1]].reshape(h1[k].shape))
        v = fvind(tmp)
        for k in range(nkpt):
            v[k] *= e_ai[k]
            v[k] = v[k].ravel()
        return np.hstack(v)

    _mo1 = lib.krylov(vind_vo, mo1base,
                      tol=tol, max_cycle=max_cycle, hermi=hermi, verbose=log).flatten()
    log.timer('krylov solver in CPHF', *t0)
    mo1 = []
    for k in range(nkpt):
        mo1.append(_mo1[moloc[k]:moloc[k+1]].reshape(h1[k].shape))
    return mo1, None

# h1 shape is (:,nocc+nvir,nocc)
def solve_withs1(fvind, mo_energy, mo_occ, h1, s1,
                 max_cycle=20, tol=1e-9, hermi=False, verbose=logger.WARN):
    '''For field dependent basis. First order overlap matrix is non-zero.
    The first order orbitals are set to
    C^1_{ij} = -1/2 S1
    e1 = h1 - s1*e0 + (e0_j-e0_i)*c1 + vhf[c1]

    Kwargs:
        hermi : boolean
            Whether the matrix defined by fvind is Hermitian or not.

    Returns:
        First order orbital coefficients (in MO basis) and first order orbital
        energy matrix
    '''
    log = logger.new_logger(verbose=verbose)
    t0 = (logger.process_clock(), logger.perf_counter())

    nkpt = len(h1)
    ncomp = h1[0].shape[0]
    occidx = []
    viridx = []
    for k in range(nkpt):
        occidx.append(mo_occ[k] > 0)
        viridx.append(mo_occ[k] == 0)

    e_a = [mo_energy[k][viridx[k]] for k in range(nkpt)]
    e_i = [mo_energy[k][occidx[k]] for k in range(nkpt)]
    e_ai = [1 / lib.direct_sum('a-i->ai', e_a[k], e_i[k]) for k in range(nkpt)]
    nocc = np.zeros([nkpt], dtype=int)
    nvir = np.zeros([nkpt], dtype=int)
    nmo  = np.zeros([nkpt], dtype=int)
    moloc = np.zeros([nkpt+1], dtype=int)
    for k in range(nkpt):
        nvir_k, nocc_k = e_ai[k].shape
        nmo_k = nvir_k + nocc_k
        nvir[k] = nvir_k
        nocc[k] = nocc_k
        nmo[k]  = nmo_k
        moloc[k+1] = moloc[k] + nmo_k * nocc_k * ncomp

    mo1base = []
    _mo1base = []
    mo_e1 = []
    for k in range(nkpt):
        mo1base.append(h1[k] - s1[k] * e_i[k])
        mo_e1.append(mo1base[k][:,occidx[k],:].copy())
        mo1base[k][:,viridx[k]] *= -e_ai[k]
        mo1base[k][:,occidx[k]] = -s1[k][:,occidx[k]] * .5
        _mo1base.append(mo1base[k].ravel())
    _mo1base = np.hstack(_mo1base)

    def vind_vo(mo1):
        mo1 = mo1.ravel()
        tmp = []
        for k in range(nkpt):
            tmp.append(mo1[moloc[k]:moloc[k+1]].reshape(-1,nmo[k],nocc[k]))
        v = fvind(tmp)
        for k in range(nkpt):
            v[k][:,viridx[k],:] *= e_ai[k]
            v[k][:,occidx[k],:] = 0
            v[k] = v[k].ravel()
        return np.hstack(v)
    _mo1 = lib.krylov(vind_vo, _mo1base,
                      tol=tol, max_cycle=max_cycle, hermi=hermi, verbose=log)
    mo1 = []
    for k in range(nkpt):
        mo1.append(_mo1[moloc[k]:moloc[k+1]].reshape(-1,nmo[k],nocc[k]))
    log.timer('krylov solver in CPHF', *t0)

    v1mo = fvind(mo1)
    for k in range(nkpt):
        mo1[k][:,viridx[k]] = mo1base[k][:,viridx[k]] - \
                             v1mo[k][:,viridx[k]]*e_ai[k]

    # mo_e1 has the same symmetry as the first order Fock matrix (hermitian or
    # anti-hermitian). mo_e1 = v1mo + u1*lib.direct_sum('i-j->ij',e_i,e_i)
    for k in range(nkpt):
        mo_e1[k] += mo1[k][:,occidx[k]] * lib.direct_sum('i-j->ij', e_i[k], e_i[k])
        mo_e1[k] += v1mo[k][:,occidx[k]]

    return mo1, mo_e1


# pyscf/pbc/scf/ghf.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Generalized Hartree-Fock for periodic systems at a single k-point
'''

import numpy as np
import scipy.linalg
import pyscf.scf.ghf as mol_ghf
from pyscf.pbc.scf import hf as pbchf
from pyscf.pbc.scf import addons
from pyscf.pbc.scf import chkfile  # noqa


def get_jk(mf, cell=None, dm=None, hermi=0, kpt=None, kpts_band=None,
           with_j=True, with_k=True, **kwargs):
    if cell is None: cell = mf.cell
    if dm is None: dm = mf.make_rdm1()
    if kpt is None: kpt = mf.kpt

    dm = np.asarray(dm)
    nso = dm.shape[-1]
    nao = nso // 2
    dms = dm.reshape(-1,nso,nso)
    n_dm = dms.shape[0]

    dmaa = dms[:,:nao,:nao]
    dmab = dms[:,nao:,:nao]
    dmbb = dms[:,nao:,nao:]
    dms = np.vstack((dmaa, dmbb, dmab))

    j1, k1 = mf.with_df.get_jk(dms, hermi, kpt, kpts_band, with_j, with_k,
                               exxdiv=mf.exxdiv)
    j1 = j1.reshape(3,n_dm,nao,nao)
    k1 = k1.reshape(3,n_dm,nao,nao)

    vj = vk = None
    if with_j:
        vj = np.zeros((n_dm,nso,nso), j1.dtype)
        vj[:,:nao,:nao] = vj[:,nao:,nao:] = j1[0] + j1[1]
        vj = vj.reshape(dm.shape)

    if with_k:
        vk = np.zeros((n_dm,nso,nso), k1.dtype)
        vk[:,:nao,:nao] = k1[0]
        vk[:,nao:,nao:] = k1[1]
        vk[:,:nao,nao:] = k1[2]
        vk[:,nao:,:nao] = k1[2].transpose(0,2,1).conj()
        vk = vk.reshape(dm.shape)

    return vj, vk


class GHF(pbchf.SCF, mol_ghf.GHF):
    '''GHF class for PBCs.
    '''

    def get_hcore(self, cell=None, kpt=None):
        hcore = pbchf.SCF.get_hcore(self, cell, kpt)
        return scipy.linalg.block_diag(hcore, hcore)

    def get_ovlp(self, cell=None, kpt=None):
        s = pbchf.SCF.get_ovlp(self, cell, kpt)
        return scipy.linalg.block_diag(s, s)

    get_jk = get_jk
    get_occ = mol_ghf.get_occ
    get_grad = mol_ghf.GHF.get_grad

    def get_j(self, cell=None, dm=None, hermi=0, kpt=None, kpts_band=None,
              **kwargs):
        return self.get_jk(cell, dm, hermi, kpt, kpts_band, True, False)[0]

    def get_k(self, cell=None, dm=None, hermi=0, kpt=None, kpts_band=None,
              **kwargs):
        return self.get_jk(cell, dm, hermi, kpt, kpts_band, False, True)[1]

    def get_veff(self, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
                 kpt=None, kpts_band=None):
        vj, vk = self.get_jk(cell, dm, hermi, kpt, kpts_band, True, True)
        vhf = vj - vk
        return vhf

    def get_bands(self, kpts_band, cell=None, dm=None, kpt=None):
        '''Get energy bands at the given (arbitrary) 'band' k-points.

        Returns:
            mo_energy : (nmo,) ndarray or a list of (nmo,) ndarray
                Bands energies E_n(k)
            mo_coeff : (nao, nmo) ndarray or a list of (nao,nmo) ndarray
                Band orbitals psi_n(k)
        '''
        raise NotImplementedError

    def get_init_guess(self, cell=None, key='minao'):
        if cell is None: cell = self.cell
        dm = mol_ghf.GHF.get_init_guess(self, cell, key)
        dm = pbchf.normalize_dm_(self, dm)
        return dm

    def convert_from_(self, mf):
        '''Convert given mean-field object to RHF/ROHF'''
        addons.convert_to_ghf(mf, self)
        return self

    stability = None
    nuc_grad_method = None


if __name__ == '__main__':
    from pyscf.pbc import gto
    from pyscf.pbc import scf

    cell = gto.Cell()
    cell.atom = '''
    H 0 0 0
    H 1 0 0
    H 0 1 0
    H 0 1 1
    '''
    cell.a = np.eye(3)*2
    cell.basis = [[0, [1.2, 1]]]
    cell.verbose = 4
    cell.build()

    kpts = cell.make_kpts([2,2,2])
    mf = scf.RHF(cell, kpt=kpts[7]).run()
    mf = GHF(cell, kpt=kpts[7])
    mf.kernel()


# pyscf/pbc/scf/hf.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors: Garnet Chan <gkc1000@gmail.com>
#          Timothy Berkelbach <tim.berkelbach@gmail.com>
#          Qiming Sun <osirpt.sun@gmail.com>
#

'''
Hartree-Fock for periodic systems at a single k-point

See Also:
    pyscf.pbc.scf.khf.py : Hartree-Fock for periodic systems with k-point sampling
'''

import sys

import numpy as np
import h5py
from pyscf.scf import hf as mol_hf
from pyscf import lib
from pyscf.lib import logger
from pyscf.data import nist
from pyscf.pbc import gto
from pyscf.pbc import tools
from pyscf.pbc.gto import ecp
from pyscf.pbc.gto.pseudo import get_pp
from pyscf.pbc.scf import chkfile  # noqa
from pyscf.pbc.scf import addons
from pyscf.pbc import df
from pyscf.pbc.scf.rsjk import RangeSeparationJKBuilder
from pyscf.pbc.lib.kpts_helper import gamma_point
from pyscf import __config__


def get_ovlp(cell, kpt=np.zeros(3)):
    '''Get the overlap AO matrix.
    '''
    # Avoid pbcopt's prescreening in the lattice sum, for better accuracy
    s = cell.pbc_intor('int1e_ovlp', hermi=0, kpts=kpt,
                       pbcopt=lib.c_null_ptr())
    s = lib.asarray(s)
    hermi_error = abs(s - np.rollaxis(s.conj(), -1, -2)).max()
    if hermi_error > cell.precision and hermi_error > 1e-12:
        logger.warn(cell, '%.4g error found in overlap integrals. '
                    'cell.precision  or  cell.rcut  can be adjusted to '
                    'improve accuracy.')

    cond = np.max(lib.cond(s))
    if cond * cell.precision > 1e2:
        prec = 1e2 / cond
        rmin = max([cell.bas_rcut(ib, prec) for ib in range(cell.nbas)])
        if cell.rcut < rmin:
            logger.warn(cell, 'Singularity detected in overlap matrix.  '
                        'Integral accuracy may be not enough.\n      '
                        'You can adjust  cell.precision  or  cell.rcut  to '
                        'improve accuracy.  Recommended values are\n      '
                        'cell.precision = %.2g  or smaller.\n      '
                        'cell.rcut = %.4g  or larger.', prec, rmin)
    return s


def get_hcore(cell, kpt=np.zeros(3)):
    '''Get the core Hamiltonian AO matrix.
    '''
    hcore = get_t(cell, kpt)
    if cell.pseudo:
        hcore += get_pp(cell, kpt)
    else:
        hcore += get_nuc(cell, kpt)
    if len(cell._ecpbas) > 0:
        hcore += ecp.ecp_int(cell, kpt)
    return hcore


def get_t(cell, kpt=np.zeros(3)):
    '''Get the kinetic energy AO matrix.
    '''
    return cell.pbc_intor('int1e_kin', hermi=1, kpts=kpt)


def get_nuc(cell, kpt=np.zeros(3)):
    '''Get the bare periodic nuc-el AO matrix, with G=0 removed.

    See Martin (12.16)-(12.21).
    '''
    return df.FFTDF(cell).get_nuc(kpt)


def get_j(cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3), kpts_band=None):
    '''Get the Coulomb (J) AO matrix for the given density matrix.

    Args:
        dm : ndarray or list of ndarrays
            A density matrix or a list of density matrices

    Kwargs:
        hermi : int
            Whether J, K matrix is hermitian
            | 0 : no hermitian or symmetric
            | 1 : hermitian
            | 2 : anti-hermitian
        vhfopt :
            A class which holds precomputed quantities to optimize the
            computation of J, K matrices
        kpt : (3,) ndarray
            The "inner" dummy k-point at which the DM was evaluated (or
            sampled).
        kpts_band : (3,) ndarray or (*,3) ndarray
            An arbitrary "band" k-point at which J is evaluated.

    Returns:
        The function returns one J matrix, corresponding to the input
        density matrix (both order and shape).
    '''
    return df.FFTDF(cell).get_jk(dm, hermi, kpt, kpts_band, with_k=False)[0]


def get_jk(mf, cell, dm, hermi=1, vhfopt=None, kpt=np.zeros(3),
           kpts_band=None, with_j=True, with_k=True, omega=None, **kwargs):
    '''Get the Coulomb (J) and exchange (K) AO matrices for the given density matrix.

    Args:
        dm : ndarray or list of ndarrays
            A density matrix or a list of density matrices

    Kwargs:
        hermi : int
            Whether J, K matrix is hermitian
            | 0 : no hermitian or symmetric
            | 1 : hermitian
            | 2 : anti-hermitian
        vhfopt :
            A class which holds precomputed quantities to optimize the
            computation of J, K matrices
        kpt : (3,) ndarray
            The "inner" dummy k-point at which the DM was evaluated (or
            sampled).
        kpts_band : (3,) ndarray or (*,3) ndarray
            An arbitrary "band" k-point at which J and K are evaluated.

    Returns:
        The function returns one J and one K matrix, corresponding to the input
        density matrix (both order and shape).
    '''
    return df.FFTDF(cell).get_jk(dm, hermi, kpt, kpts_band, with_j, with_k,
                                 omega, exxdiv=mf.exxdiv)


def get_bands(mf, kpts_band, cell=None, dm=None, kpt=None):
    '''Get energy bands at the given (arbitrary) 'band' k-points.

    Returns:
        mo_energy : (nmo,) ndarray or a list of (nmo,) ndarray
            Bands energies E_n(k)
        mo_coeff : (nao, nmo) ndarray or a list of (nao,nmo) ndarray
            Band orbitals psi_n(k)
    '''
    if cell is None: cell = mf.cell
    if dm is None: dm = mf.make_rdm1()
    if kpt is None: kpt = mf.kpt

    kpts_band = np.asarray(kpts_band)
    single_kpt_band = (getattr(kpts_band, 'ndim', None) == 1)
    kpts_band = kpts_band.reshape(-1,3)

    fock = mf.get_hcore(cell, kpts_band)
    fock = fock + mf.get_veff(cell, dm, kpt=kpt, kpts_band=kpts_band)
    s1e = mf.get_ovlp(cell, kpts_band)
    nkpts = len(kpts_band)
    mo_energy = []
    mo_coeff = []
    for k in range(nkpts):
        e, c = mf.eig(fock[k], s1e[k])
        mo_energy.append(e)
        mo_coeff.append(c)

    if single_kpt_band:
        mo_energy = mo_energy[0]
        mo_coeff = mo_coeff[0]
    return mo_energy, mo_coeff


def init_guess_by_chkfile(cell, chkfile_name, project=None, kpt=None):
    '''Read the HF results from checkpoint file, then project it to the
    basis defined by ``cell``

    Returns:
        Density matrix, (nao,nao) ndarray
    '''
    from pyscf.pbc.scf import uhf
    dm = uhf.init_guess_by_chkfile(cell, chkfile_name, project, kpt)
    return dm[0] + dm[1]

get_fock = mol_hf.get_fock
get_occ = mol_hf.get_occ
get_grad = mol_hf.get_grad
make_rdm1 = mol_hf.make_rdm1
energy_elec = mol_hf.energy_elec


def dip_moment(cell, dm, unit='Debye', verbose=logger.NOTE,
               grids=None, rho=None, kpt=np.zeros(3), origin=None):
    ''' Dipole moment in the unit cell (is it well defined)?

    Args:
         cell : an instance of :class:`Cell`

         dm (ndarray) : density matrix

    Return:
        A list: the dipole moment on x, y and z components
    '''
    from pyscf.pbc import tools
    from pyscf.pbc.dft import gen_grid
    from pyscf.pbc.dft import numint
    if cell.dimension != 3:
        # raise NotImplementedError
        logger.warn(cell, 'Dipole moment for low-dimension system is not supported.')
        return np.zeros(3)

    log = logger.new_logger(cell, verbose)
    a = cell.lattice_vectors()
    b = np.linalg.inv(a).T

    if grids is None:
        grids = gen_grid.UniformGrids(cell)
        #? FIXME: Less requirements on the density accuracy.
        #ke_cutoff = gto.estimate_ke_cutoff(cell, 1e-5)
        #grids.mesh = tools.cutoff_to_mesh(a, ke_cutoff)
    if rho is None:
        rho = numint.NumInt().get_rho(cell, dm, grids, kpt, cell.max_memory)

    if origin is None:
        origin = _search_dipole_gauge_origin(cell, grids, rho, log)

    # Move the unit cell to the position around the origin.
    def shift_grids(r):
        r_frac = lib.dot(r - origin, b.T)
        # Grids on the boundary (r_frac == +/-0.5) of the new cell may lead to
        # unbalanced contributions to the dipole moment. Exclude them from the
        # dipole and quadrupole
        r_frac[r_frac== 0.5] = 0
        r_frac[r_frac==-0.5] = 0
        r_frac[r_frac > 0.5] -= 1
        r_frac[r_frac <-0.5] += 1
        r = lib.dot(r_frac, a)
        return r
    r = shift_grids(grids.coords)
    e_dip = np.einsum('g,g,gx->x', rho, grids.weights, r)

    charges = cell.atom_charges()
    r = shift_grids(cell.atom_coords())
    nuc_dip = np.einsum('g,gx->x', charges, r)
    dip = nuc_dip - e_dip

    if unit.upper() == 'DEBYE':
        dip *= nist.AU2DEBYE
        log.note('Dipole moment(X, Y, Z, Debye): %8.5f, %8.5f, %8.5f', *dip)
    else:
        log.note('Dipole moment(X, Y, Z, A.U.): %8.5f, %8.5f, %8.5f', *dip)
    return dip

def _search_dipole_gauge_origin(cell, grids, rho, log):
    '''Optimize the position of the unit cell in crystal. With a proper unit
    cell, the nuclear charge center and the electron density center are at the
    same point. This function returns the origin of such unit cell.
    '''
    from pyscf.pbc.dft import gen_grid
    a = cell.lattice_vectors()
    b = np.linalg.inv(a).T
    charges = cell.atom_charges()
    coords  = cell.atom_coords()

    origin = np.einsum('i,ix->x', charges, coords) / charges.sum()
    log.debug1('Initial guess of origin center %s', origin)
    nelec = np.dot(rho, grids.weights)

    # The dipole moment in the crystal is not uniquely defined. Depending on
    # the position and the shape of the unit cell, the value of dipole moment
    # can be very different. The optimization below searches the boundary of
    # cell inside which the nuclear charge center and electron density charge
    # center locate at the same point. The gauge origin will be chosen at the
    # center of unit cell.
    if (cell.dimension == 3 and
        # For orthogonal lattice only
        abs(np.diag(a.diagonal())).max() and
        isinstance(grids, gen_grid.UniformGrids)):
        gridxyz = grids.coords.T.reshape(3, *grids.mesh)
        gridbase = (gridxyz[0,:,0,0], gridxyz[1,0,:,0], gridxyz[2,0,0,:])
        wxyz = grids.weights.reshape(grids.mesh)
        rhoxyz = rho.reshape(grids.mesh)

        def search_orig(ix, origin):
            nx = grids.mesh[ix]
            Lx = a[ix,ix]
            g = gridbase[ix] - origin
            g_frac = (g * b[ix,ix]).round(6)
            g[g_frac == .5]  = 0
            g[g_frac ==-.5]  = 0
            g[g_frac >  .5] -= Lx
            g[g_frac < -.5] += Lx
            meanx = np.einsum('xyz,xyz->'+('xyz'[ix]), rhoxyz, wxyz) / nelec
            ex = meanx * g

            r_nuc = coords[:,ix] - origin
            r_frac = (r_nuc * b[ix,ix]).round(6)
            r_nuc[r_frac == .5]  = 0
            r_nuc[r_frac ==-.5]  = 0
            r_nuc[r_frac >  .5] -= Lx
            r_nuc[r_frac < -.5] += Lx
            nuc_dip = np.dot(charges, r_nuc) / charges.sum()

            # ex.sum() ~ electron dipole wrt the given origin
            dipx = nuc_dip - ex.sum()

            g = gridbase[ix] - origin
            sorted_meanx = np.hstack((meanx[g >= Lx/2],
                                      meanx[(g < Lx/2) & (g >= -Lx/2)],
                                      meanx[g < -Lx/2]))
            if abs(dipx) < 1e-3:
                offx = 0
            elif dipx > 0:
                # To cancel the positive dipole, move electrons to the right side
                rcum_dip = np.append(0, np.cumsum(sorted_meanx * Lx))
                idx = np.where(rcum_dip > dipx)[0][0]
                dx = (rcum_dip[idx] - dipx) / (rcum_dip[idx] - rcum_dip[idx-1])
                offx = (idx - dx) * Lx/nx
            else:
                # To cancel the negative dipole, move electrons to the left side
                lcum_dip = np.append(0, np.cumsum(sorted_meanx[::-1] * Lx))
                idx = np.where(lcum_dip > -dipx)[0][0]
                dx = (lcum_dip[idx] - -dipx) / (lcum_dip[idx] - lcum_dip[idx-1])
                offx = -(idx - dx) * Lx/nx
            return origin + offx

        wbar = grids.weights[0]**(1./3)
        for i in range(4):
            orig_last = origin
            origin[0] = search_orig(0, origin[0])
            origin[1] = search_orig(1, origin[1])
            origin[2] = search_orig(2, origin[2])
            if abs(origin - orig_last).max() < wbar:
                break
            log.debug1('iter %d: origin %s', i, origin)
    else:
        # If the grids are non-cubic grids, regenerating the grids is expensive if
        # the position or the shape of the unit cell is changed. The position of
        # the unit cell is not optimized. The gauge origin is set to the nuclear
        # charge center of the original unit cell.
        pass

    return origin

def get_rho(mf, dm=None, grids=None, kpt=None):
    '''Compute density in real space
    '''
    from pyscf.pbc.dft import gen_grid
    from pyscf.pbc.dft import numint
    if dm is None:
        dm = mf.make_rdm1()
    if getattr(dm, 'ndim', None) != 2:  # UHF
        dm = dm[0] + dm[1]
    if grids is None:
        grids = gen_grid.UniformGrids(mf.cell)
    if kpt is None:
        kpt = mf.kpt
    ni = numint.NumInt()
    return ni.get_rho(mf.cell, dm, grids, kpt, mf.max_memory)

def _dip_correction(mf):
    '''Makov-Payne corrections for charged systems.'''
    from pyscf.pbc import tools
    from pyscf.pbc.dft import gen_grid
    log = logger.new_logger(mf)
    cell = mf.cell
    a = cell.lattice_vectors()
    b = np.linalg.inv(a).T

    grids = gen_grid.UniformGrids(cell)
    ke_cutoff = gto.estimate_ke_cutoff(cell, 1e-5)
    grids.mesh = tools.cutoff_to_mesh(a, ke_cutoff)

    dm = mf.make_rdm1()
    rho = mf.get_rho(dm, grids, mf.kpt)
    origin = _search_dipole_gauge_origin(cell, grids, rho, log)

    def shift_grids(r):
        r_frac = lib.dot(r - origin, b.T)
        # Grids on the boundary (r_frac == +/-0.5) of the new cell may lead to
        # unbalanced contributions to the dipole moment. Exclude them from the
        # dipole and quadrupole
        r_frac[r_frac== 0.5] = 0
        r_frac[r_frac==-0.5] = 0
        r_frac[r_frac > 0.5] -= 1
        r_frac[r_frac <-0.5] += 1
        r = lib.dot(r_frac, a)
        return r

    #           SC              BCC             FCC
    madelung = (-2.83729747948, -3.63923344951, -4.58486207411)
    vol = cell.vol
    L = vol ** (1./3)
    chg = cell.charge

    # epsilon is the dielectric constant of the system. For systems
    # surrounded by vacuum, epsilon = 1.
    epsilon = 1

    # Energy correction of point charges of a simple cubic lattice.
    de_mono = - chg**2 * np.array(madelung) / (2 * L * epsilon)

    # dipole energy correction
    r_e = shift_grids(grids.coords)
    r_nuc = shift_grids(cell.atom_coords())
    charges = cell.atom_charges()
    e_dip = np.einsum('g,g,gx->x', rho, grids.weights, r_e)
    nuc_dip = np.einsum('g,gx->x', charges, r_nuc)
    dip = nuc_dip - e_dip
    de_dip = -2.*np.pi/(3*cell.vol) * np.linalg.norm(dip)**2

    # quadrupole energy correction
    if abs(a - np.eye(3)*L).max() > 1e-5:
        logger.warn(mf, 'System is not cubic cell. Quadrupole energy '
                    'correction is inaccurate since it is developed based on '
                    'cubic cell.')
    e_quad = np.einsum('g,g,gx,gx->', rho, grids.weights, r_e, r_e)
    nuc_quad = np.einsum('g,gx,gx->', charges, r_nuc, r_nuc)
    quad = nuc_quad - e_quad
    de_quad = 2.*np.pi/(3*cell.vol) * quad

    de = de_mono + de_dip + de_quad
    return de_mono, de_dip, de_quad, de

def makov_payne_correction(mf):
    '''Makov-Payne correction (Phys. Rev. B, 51, 4014)
    '''
    cell = mf.cell
    logger.note(mf, 'Makov-Payne correction for charged 3D PBC systems')
    # PRB 51 (1995), 4014
    # PRB 77 (2008), 115139
    if cell.dimension != 3:
        logger.warn(mf, 'Correction for low-dimension PBC systems'
                    'is not available.')
        return 0

    de_mono, de_dip, de_quad, de = _dip_correction(mf)

    if mf.verbose >= logger.NOTE:
        write = mf.stdout.write
        write('Corrections (AU)\n')
        write('       Monopole      Dipole          Quadrupole    total\n')
        write('SC   %12.8f   %12.8f   %12.8f   %12.8f\n' %
              (de_mono[0], de_dip   , de_quad   , de[0]))
        write('BCC  %12.8f   %12.8f   %12.8f   %12.8f\n' %
              (de_mono[1], de_dip   , de_quad   , de[1]))
        write('FCC  %12.8f   %12.8f   %12.8f   %12.8f\n' %
              (de_mono[2], de_dip   , de_quad   , de[2]))
    return de


class SCF(mol_hf.SCF):
    '''SCF base class adapted for PBCs.

    Attributes:
        kpt : (3,) ndarray
            The AO k-point in Cartesian coordinates, in units of 1/Bohr.

        exxdiv : str
            Exchange divergence treatment, can be one of

            | None : ignore G=0 contribution in exchange
            | 'ewald' : Ewald probe charge correction [JCP 122, 234102 (2005); DOI:10.1063/1.1926272]

        with_df : density fitting object
            Default is the FFT based DF model. For all-electron calculation,
            MDF model is favored for better accuracy.  See also :mod:`pyscf.pbc.df`.

        direct_scf : bool
            When this flag is set to true, the J/K matrices will be computed
            directly through the underlying with_df methods.  Otherwise,
            depending the available memory, the 4-index integrals may be cached
            and J/K matrices are computed based on the 4-index integrals.
    '''

    direct_scf = getattr(__config__, 'pbc_scf_SCF_direct_scf', False)

    def __init__(self, cell, kpt=np.zeros(3),
                 exxdiv=getattr(__config__, 'pbc_scf_SCF_exxdiv', 'ewald')):
        if not cell._built:
            sys.stderr.write('Warning: cell.build() is not called in input\n')
            cell.build()
        self.cell = cell
        mol_hf.SCF.__init__(self, cell)

        self.with_df = df.FFTDF(cell)
        # Range separation JK builder
        self.rsjk = None

        self.exxdiv = exxdiv
        self.kpt = kpt
        self.conv_tol = cell.precision * 10

        self._keys = self._keys.union(['cell', 'exxdiv', 'with_df', 'rsjk'])

    @property
    def kpt(self):
        if 'kpt' in self.__dict__:
            # To handle the attribute kpt loaded from chkfile
            self.kpt = self.__dict__.pop('kpt')
        return self.with_df.kpts.reshape(3)
    @kpt.setter
    def kpt(self, x):
        self.with_df.kpts = np.reshape(x, (-1, 3))
        if self.rsjk:
            self.rsjk.kpts = self.with_df.kpts

    def build(self, cell=None):
        if 'kpt' in self.__dict__:
            # To handle the attribute kpt loaded from chkfile
            self.kpt = self.__dict__.pop('kpt')

        if self.rsjk:
            if not np.all(self.rsjk.kpts == self.kpt):
                self.rsjk = self.rsjk.__class__(cell, self.kpt.reshape(1,3))
            self.rsjk.build(direct_scf_tol=self.direct_scf_tol)

        if self.verbose >= logger.WARN:
            self.check_sanity()
        return self

    def reset(self, cell=None):
        '''Reset cell and relevant attributes associated to the old cell object'''
        if cell is not None:
            self.cell = cell
            self.mol = cell # used by hf kernel
        self.with_df.reset(cell)
        return self

    def dump_flags(self, verbose=None):
        mol_hf.SCF.dump_flags(self, verbose)
        logger.info(self, '******** PBC SCF flags ********')
        logger.info(self, 'kpt = %s', self.kpt)
        logger.info(self, 'Exchange divergence treatment (exxdiv) = %s', self.exxdiv)
        cell = self.cell
        if ((cell.dimension >= 2 and cell.low_dim_ft_type != 'inf_vacuum') and
            isinstance(self.exxdiv, str) and self.exxdiv.lower() == 'ewald'):
            madelung = tools.pbc.madelung(cell, [self.kpt])
            logger.info(self, '    madelung (= occupied orbital energy shift) = %s', madelung)
            logger.info(self, '    Total energy shift due to Ewald probe charge'
                        ' = -1/2 * Nelec*madelung = %.12g',
                        madelung*cell.nelectron * -.5)
        logger.info(self, 'DF object = %s', self.with_df)
        if not getattr(self.with_df, 'build', None):
            # .dump_flags() is called in pbc.df.build function
            self.with_df.dump_flags(verbose)
        return self

    def check_sanity(self):
        mol_hf.SCF.check_sanity(self)
        self.with_df.check_sanity()
        if (isinstance(self.exxdiv, str) and self.exxdiv.lower() != 'ewald' and
            isinstance(self.with_df, df.df.DF)):
            logger.warn(self, 'exxdiv %s is not supported in DF or MDF',
                        self.exxdiv)
        return self

    def get_hcore(self, cell=None, kpt=None):
        if cell is None: cell = self.cell
        if kpt is None: kpt = self.kpt
        if cell.pseudo:
            nuc = self.with_df.get_pp(kpt)
        else:
            nuc = self.with_df.get_nuc(kpt)
        if len(cell._ecpbas) > 0:
            nuc += ecp.ecp_int(cell, kpt)
        return nuc + cell.pbc_intor('int1e_kin', 1, 1, kpt)

    def get_ovlp(self, cell=None, kpt=None):
        if cell is None: cell = self.cell
        if kpt is None: kpt = self.kpt
        return get_ovlp(cell, kpt)

    def get_jk(self, cell=None, dm=None, hermi=1, kpt=None, kpts_band=None,
               with_j=True, with_k=True, omega=None, **kwargs):
        r'''Get Coulomb (J) and exchange (K) following :func:`scf.hf.RHF.get_jk_`.
        for particular k-point (kpt).

        When kpts_band is given, the J, K matrices on kpts_band are evaluated.

            J_{pq} = \sum_{rs} (pq|rs) dm[s,r]
            K_{pq} = \sum_{rs} (pr|sq) dm[r,s]

        where r,s are orbitals on kpt. p and q are orbitals on kpts_band
        if kpts_band is given otherwise p and q are orbitals on kpt.
        '''
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        if kpt is None: kpt = self.kpt

        cpu0 = (logger.process_clock(), logger.perf_counter())
        dm = np.asarray(dm)
        nao = dm.shape[-1]

        if (not omega and kpts_band is None and
            # TODO: generate AO integrals with rsjk algorithm
            not self.rsjk and
            (self.exxdiv == 'ewald' or not self.exxdiv) and
            (self._eri is not None or cell.incore_anyway or
             (not self.direct_scf and self._is_mem_enough()))):
            if self._eri is None:
                logger.debug(self, 'Building PBC AO integrals incore')
                self._eri = self.with_df.get_ao_eri(kpt, compact=True)
            vj, vk = mol_hf.dot_eri_dm(self._eri, dm, hermi, with_j, with_k)

            if with_k and self.exxdiv == 'ewald':
                from pyscf.pbc.df.df_jk import _ewald_exxdiv_for_G0
                # G=0 is not inculded in the ._eri integrals
                _ewald_exxdiv_for_G0(self.cell, kpt, dm.reshape(-1,nao,nao),
                                     vk.reshape(-1,nao,nao))
        elif self.rsjk:
            vj, vk = self.rsjk.get_jk(dm.reshape(-1,nao,nao), hermi, kpt, kpts_band,
                                      with_j, with_k, omega, exxdiv=self.exxdiv)
        else:
            vj, vk = self.with_df.get_jk(dm.reshape(-1,nao,nao), hermi, kpt, kpts_band,
                                         with_j, with_k, omega, exxdiv=self.exxdiv)

        if with_j:
            vj = _format_jks(vj, dm, kpts_band)
        if with_k:
            vk = _format_jks(vk, dm, kpts_band)
        logger.timer(self, 'vj and vk', *cpu0)
        return vj, vk

    def get_j(self, cell=None, dm=None, hermi=1, kpt=None, kpts_band=None,
              omega=None):
        r'''Compute J matrix for the given density matrix and k-point (kpt).
        When kpts_band is given, the J matrices on kpts_band are evaluated.

            J_{pq} = \sum_{rs} (pq|rs) dm[s,r]

        where r,s are orbitals on kpt. p and q are orbitals on kpts_band
        if kpts_band is given otherwise p and q are orbitals on kpt.
        '''
        return self.get_jk(cell, dm, hermi, kpt, kpts_band, with_k=False,
                           omega=omega)[0]

    def get_k(self, cell=None, dm=None, hermi=1, kpt=None, kpts_band=None,
              omega=None):
        '''Compute K matrix for the given density matrix.
        '''
        return self.get_jk(cell, dm, hermi, kpt, kpts_band, with_j=False,
                           omega=omega)[1]

    def get_veff(self, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
                 kpt=None, kpts_band=None):
        '''Hartree-Fock potential matrix for the given density matrix.
        See :func:`scf.hf.get_veff` and :func:`scf.hf.RHF.get_veff`
        '''
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        if kpt is None: kpt = self.kpt
        if self.rsjk and self.direct_scf:
            # Enable direct-SCF for real space JK builder
            ddm = dm - dm_last
            vj, vk = self.get_jk(cell, ddm, hermi, kpt, kpts_band)
            vhf = vj - vk * .5
            vhf += vhf_last
        else:
            vj, vk = self.get_jk(cell, dm, hermi, kpt, kpts_band)
            vhf = vj - vk * .5
        return vhf

    def get_jk_incore(self, cell=None, dm=None, hermi=1, kpt=None, omega=None,
                      **kwargs):
        '''Get Coulomb (J) and exchange (K) following :func:`scf.hf.RHF.get_jk_`.

        *Incore* version of Coulomb and exchange build only.
        Currently RHF always uses PBC AO integrals (unlike RKS), since
        exchange is currently computed by building PBC AO integrals.
        '''
        if omega:
            raise NotImplementedError
        if cell is None: cell = self.cell
        if kpt is None: kpt = self.kpt
        if self._eri is None:
            self._eri = self.with_df.get_ao_eri(kpt, compact=True)
        return self.get_jk(cell, dm, hermi, kpt)

    def energy_nuc(self):
        return self.cell.energy_nuc()

    get_bands = get_bands

    get_rho = get_rho

    @lib.with_doc(dip_moment.__doc__)
    def dip_moment(self, cell=None, dm=None, unit='Debye', verbose=logger.NOTE,
                   **kwargs):
        rho = kwargs.pop('rho', None)
        if rho is None:
            rho = self.get_rho(dm)
        if cell is None:
            cell = self.cell
        return dip_moment(cell, dm, unit, verbose, rho=rho, kpt=self.kpt, **kwargs)

    def _finalize(self):
        '''Hook for dumping results and clearing up the object.'''
        mol_hf.SCF._finalize(self)

        if self.cell.charge != 0:
            makov_payne_correction(self)
        return self

    def get_init_guess(self, cell=None, key='minao'):
        if cell is None: cell = self.cell
        dm = mol_hf.SCF.get_init_guess(self, cell, key)
        dm = normalize_dm_(self, dm)
        return dm

    def init_guess_by_1e(self, cell=None):
        if cell is None: cell = self.cell
        if cell.dimension < 3:
            logger.warn(self, 'Hcore initial guess is not recommended in '
                        'the SCF of low-dimensional systems.')
        return mol_hf.SCF.init_guess_by_1e(self, cell)

    def init_guess_by_chkfile(self, chk=None, project=None, kpt=None):
        if chk is None: chk = self.chkfile
        if kpt is None: kpt = self.kpt
        return init_guess_by_chkfile(self.cell, chk, project, kpt)
    def from_chk(self, chk=None, project=None, kpt=None):
        return self.init_guess_by_chkfile(chk, project, kpt)

    def dump_chk(self, envs):
        if self.chkfile:
            mol_hf.SCF.dump_chk(self, envs)
            with h5py.File(self.chkfile, 'a') as fh5:
                fh5['scf/kpt'] = self.kpt
        return self

    def _is_mem_enough(self):
        nao = self.cell.nao_nr()
        if abs(self.kpt).sum() < 1e-9:
            mem_need = nao**4*8/4/1e6
        else:
            mem_need = nao**4*16/1e6
        return mem_need + lib.current_memory()[0] < self.max_memory*.95

    def density_fit(self, auxbasis=None, with_df=None):
        from pyscf.pbc.df import df_jk
        return df_jk.density_fit(self, auxbasis, with_df=with_df)

    def rs_density_fit(self, auxbasis=None, with_df=None):
        from pyscf.pbc.df import rsdf_jk
        return rsdf_jk.density_fit(self, auxbasis, with_df=with_df)

    def mix_density_fit(self, auxbasis=None, with_df=None):
        from pyscf.pbc.df import mdf_jk
        return mdf_jk.density_fit(self, auxbasis, with_df=with_df)

    def sfx2c1e(self):
        from pyscf.pbc.x2c import sfx2c1e
        return sfx2c1e.sfx2c1e(self)
    x2c = x2c1e = sfx2c1e

    def to_rhf(self, mf):
        '''Convert the input mean-field object to a RHF/ROHF/RKS/ROKS object'''
        return addons.convert_to_rhf(mf)

    def to_uhf(self, mf):
        '''Convert the input mean-field object to a UHF/UKS object'''
        return addons.convert_to_uhf(mf)

    def to_ghf(self, mf):
        '''Convert the input mean-field object to a GHF/GKS object'''
        return addons.convert_to_ghf(mf)

    def nuc_grad_method(self, *args, **kwargs):
        raise NotImplementedError

    def jk_method(self, J='FFTDF', K=None):
        '''
        Set up the schemes to evaluate Coulomb and exchange matrix

        FFTDF: planewave density fitting using Fast Fourier Transform
        AFTDF: planewave density fitting using analytic Fourier Transform
        GDF: Gaussian density fitting
        MDF: Gaussian and planewave mix density fitting
        RS: range-separation JK builder
        RSDF: range-separation density fitting
        '''
        if K is None:
            K = J

        if J != K:
            raise NotImplementedError('J != K')

        if 'DF' in J or 'DF' in K:
            if 'DF' in J and 'DF' in K:
                assert J == K
            else:
                df_method = J if 'DF' in J else K
                self.with_df = getattr(df, df_method)(self.cell, self.kpt)

        if 'RS' in J or 'RS' in K:
            if not gamma_point(self.kpt):
                raise NotImplementedError('Single k-point must be gamma point')
            self.rsjk = RangeSeparationJKBuilder(self.cell, self.kpt)
            self.rsjk.verbose = self.verbose

        # For nuclear attraction
        if J == 'RS' and K == 'RS' and not isinstance(self.with_df, df.GDF):
            self.with_df = df.GDF(self.cell, self.kpt)

        nuc = self.with_df.__class__.__name__
        logger.debug1(self, 'Apply %s for J, %s for K, %s for nuc', J, K, nuc)
        return self


class RHF(SCF, mol_hf.RHF):

    stability = mol_hf.RHF.stability

    def convert_from_(self, mf):
        '''Convert given mean-field object to RHF'''
        addons.convert_to_rhf(mf, self)
        return self


def _format_jks(vj, dm, kpts_band):
    if kpts_band is None:
        vj = vj.reshape(dm.shape)
    elif kpts_band.ndim == 1:  # a single k-point on bands
        vj = vj.reshape(dm.shape)
    elif getattr(dm, "ndim", 0) == 2:
        vj = vj[0]
    return vj

def normalize_dm_(mf, dm):
    '''
    Scale density matrix to make it produce the correct number of electrons.
    '''
    cell = mf.cell
    if isinstance(dm, np.ndarray) and dm.ndim == 2:
        ne = np.einsum('ij,ji->', dm, mf.get_ovlp(cell)).real
    else:
        ne = np.einsum('xij,ji->', dm, mf.get_ovlp(cell)).real
    if abs(ne - cell.nelectron).sum() > 1e-7:
        logger.debug(mf, 'Big error detected in the electron number '
                     'of initial guess density matrix (Ne/cell = %g)!\n'
                     '  This can cause huge error in Fock matrix and '
                     'lead to instability in SCF for low-dimensional '
                     'systems.\n  DM is normalized wrt the number '
                     'of electrons %s', ne, cell.nelectron)
        dm *= cell.nelectron / ne
    return dm


# pyscf/pbc/scf/kghf.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Generalized Hartree-Fock for periodic systems with k-point sampling
'''

from functools import reduce
import numpy as np
import scipy.linalg
import pyscf.scf.ghf as mol_ghf  # noqa
from pyscf import lib
from pyscf.lib import logger
from pyscf.pbc.scf import khf
from pyscf.pbc.scf import ghf as pbcghf
from pyscf.pbc.scf import addons
from pyscf.pbc.df.df_jk import _format_jks
from pyscf import __config__


def get_jk(mf, cell=None, dm_kpts=None, hermi=0, kpts=None, kpts_band=None,
           with_j=True, with_k=True, **kwargs):
    if cell is None: cell = mf.cell
    if dm_kpts is None: dm_kpts = mf.make_rdm1()
    if kpts is None: kpts = mf.kpts
    if kpts_band is None: kpts_band = kpts
    nkpts = len(kpts)
    nband = len(kpts_band)

    dm_kpts = np.asarray(dm_kpts)
    nso = dm_kpts.shape[-1]
    nao = nso // 2
    dms = dm_kpts.reshape(-1,nkpts,nso,nso)
    n_dm = dms.shape[0]

    dmaa = dms[:,:,:nao,:nao]
    dmab = dms[:,:,nao:,:nao]
    dmbb = dms[:,:,nao:,nao:]
    dms = np.vstack((dmaa, dmbb, dmab))

    j1, k1 = mf.with_df.get_jk(dms, hermi, kpts, kpts_band, with_j, with_k,
                               exxdiv=mf.exxdiv)
    j1 = j1.reshape(3,n_dm,nband,nao,nao)
    k1 = k1.reshape(3,n_dm,nband,nao,nao)

    vj = vk = None
    if with_j:
        vj = np.zeros((n_dm,nband,nso,nso), j1.dtype)
        vj[:,:,:nao,:nao] = vj[:,:,nao:,nao:] = j1[0] + j1[1]
        vj = _format_jks(vj, dm_kpts, kpts_band, kpts)

    if with_k:
        vk = np.zeros((n_dm,nband,nso,nso), k1.dtype)
        vk[:,:,:nao,:nao] = k1[0]
        vk[:,:,nao:,nao:] = k1[1]
        vk[:,:,:nao,nao:] = k1[2]
        vk[:,:,nao:,:nao] = k1[2].transpose(0,1,3,2).conj()
        vk = _format_jks(vk, dm_kpts, kpts_band, kpts)

    return vj, vk

def get_occ(mf, mo_energy_kpts=None, mo_coeff_kpts=None):
    '''Label the occupancies for each orbital for sampled k-points.

    This is a k-point version of scf.hf.SCF.get_occ
    '''
    if mo_energy_kpts is None: mo_energy_kpts = mf.mo_energy

    nkpts = len(mo_energy_kpts)
    nocc = mf.cell.nelectron * nkpts

    mo_energy = np.sort(np.hstack(mo_energy_kpts))
    fermi = mo_energy[nocc-1]
    mo_occ_kpts = []
    for mo_e in mo_energy_kpts:
        mo_occ_kpts.append((mo_e <= fermi).astype(np.double))

    if nocc < mo_energy.size:
        logger.info(mf, 'HOMO = %.12g  LUMO = %.12g',
                    mo_energy[nocc-1], mo_energy[nocc])
        if mo_energy[nocc-1]+1e-3 > mo_energy[nocc]:
            logger.warn(mf, 'HOMO %.12g == LUMO %.12g',
                        mo_energy[nocc-1], mo_energy[nocc])
    else:
        logger.info(mf, 'HOMO = %.12g', mo_energy[nocc-1])

    if mf.verbose >= logger.DEBUG:
        np.set_printoptions(threshold=len(mo_energy))
        logger.debug(mf, '     k-point                  mo_energy')
        for k,kpt in enumerate(mf.cell.get_scaled_kpts(mf.kpts)):
            logger.debug(mf, '  %2d (%6.3f %6.3f %6.3f)   %s %s',
                         k, kpt[0], kpt[1], kpt[2],
                         mo_energy_kpts[k][mo_occ_kpts[k]> 0],
                         mo_energy_kpts[k][mo_occ_kpts[k]==0])
        np.set_printoptions(threshold=1000)

    return mo_occ_kpts


class KGHF(pbcghf.GHF, khf.KSCF):
    '''GHF class for PBCs.
    '''
    def __init__(self, cell, kpts=np.zeros((1,3)),
                 exxdiv=getattr(__config__, 'pbc_scf_SCF_exxdiv', 'ewald')):
        khf.KSCF.__init__(self, cell, kpts, exxdiv)

    def get_hcore(self, cell=None, kpts=None):
        hcore = khf.KSCF.get_hcore(self, cell, kpts)
        return lib.asarray([scipy.linalg.block_diag(h, h) for h in hcore])

    def get_ovlp(self, cell=None, kpts=None):
        s = khf.KSCF.get_ovlp(self, cell, kpts)
        return lib.asarray([scipy.linalg.block_diag(x, x) for x in s])

    get_jk = get_jk
    get_occ = get_occ
    energy_elec = khf.KSCF.energy_elec

    def get_j(self, cell=None, dm_kpts=None, hermi=0, kpts=None, kpts_band=None):
        return self.get_jk(cell, dm_kpts, hermi, kpts, kpts_band, True, False)[0]

    def get_k(self, cell=None, dm_kpts=None, hermi=0, kpts=None, kpts_band=None):
        return self.get_jk(cell, dm_kpts, hermi, kpts, kpts_band, False, True)[1]

    def get_veff(self, cell=None, dm_kpts=None, dm_last=0, vhf_last=0, hermi=1,
                 kpts=None, kpts_band=None):
        vj, vk = self.get_jk(cell, dm_kpts, hermi, kpts, kpts_band, True, True)
        vhf = vj - vk
        return vhf

    def get_grad(self, mo_coeff_kpts, mo_occ_kpts, fock=None):
        '''
        returns 1D array of gradients, like non K-pt version
        note that occ and virt indices of different k pts now occur
        in sequential patches of the 1D array
        '''
        if fock is None:
            dm1 = self.make_rdm1(mo_coeff_kpts, mo_occ_kpts)
            fock = self.get_hcore(self.cell, self.kpts) + self.get_veff(self.cell, dm1)

        def grad(mo, mo_occ, fock):
            occidx = mo_occ > 0
            viridx = ~occidx
            g = reduce(np.dot, (mo[:,viridx].conj().T, fock, mo[:,occidx]))
            return g.ravel()

        grad_kpts = [grad(mo, mo_occ_kpts[k], fock[k])
                     for k, mo in enumerate(mo_coeff_kpts)]
        return np.hstack(grad_kpts)

    def get_bands(self, kpts_band, cell=None, dm_kpts=None, kpts=None):
        '''Get energy bands at the given (arbitrary) 'band' k-points.

        Returns:
            mo_energy : (nmo,) ndarray or a list of (nmo,) ndarray
                Bands energies E_n(k)
            mo_coeff : (nao, nmo) ndarray or a list of (nao,nmo) ndarray
                Band orbitals psi_n(k)
        '''
        raise NotImplementedError

    get_init_guess = khf.KSCF.get_init_guess

    def _finalize(self):
        if self.converged:
            logger.note(self, 'converged SCF energy = %.15g', self.e_tot)
        else:
            logger.note(self, 'SCF not converged.')
            logger.note(self, 'SCF energy = %.15g after %d cycles',
                        self.e_tot, self.max_cycle)
        return self

    def convert_from_(self, mf):
        '''Convert given mean-field object to RHF/ROHF'''
        addons.convert_to_ghf(mf, self)
        return self

    density_fit = khf.KSCF.density_fit
    rs_density_fit = khf.KSCF.rs_density_fit
    newton = khf.KSCF.newton

    x2c = None
    stability = None
    nuc_grad_method = None


if __name__ == '__main__':
    from pyscf.pbc import gto

    cell = gto.Cell()
    cell.atom = '''
    H 0 0 0
    H 1 0 0
    H 0 1 0
    H 0 1 1
    '''
    cell.a = np.eye(3)*2
    cell.basis = [[0, [1.2, 1]]]
    cell.verbose = 4
    cell.build()

    kpts = cell.make_kpts([2,1,1])
    mf = KGHF(cell, kpts=kpts)
    mf.kernel()


# pyscf/pbc/scf/khf.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors: Garnet Chan <gkc1000@gmail.com>
#          Timothy Berkelbach <tim.berkelbach@gmail.com>
#          Qiming Sun <osirpt.sun@gmail.com>
#

'''
Hartree-Fock for periodic systems with k-point sampling

See Also:
    hf.py : Hartree-Fock for periodic systems at a single k-point
'''

import sys

from functools import reduce
import numpy as np
import scipy.linalg
import h5py
from pyscf.pbc.scf import hf as pbchf
from pyscf import lib
from pyscf.scf import hf as mol_hf
from pyscf.lib import logger
from pyscf.pbc.gto import ecp
from pyscf.pbc.scf import addons
from pyscf.pbc.scf import chkfile  # noqa
from pyscf.pbc import tools
from pyscf.pbc import df
from pyscf.pbc.scf.rsjk import RangeSeparationJKBuilder
from pyscf import __config__

WITH_META_LOWDIN = getattr(__config__, 'pbc_scf_analyze_with_meta_lowdin', True)
PRE_ORTH_METHOD = getattr(__config__, 'pbc_scf_analyze_pre_orth_method', 'ANO')
CHECK_COULOMB_IMAG = getattr(__config__, 'pbc_scf_check_coulomb_imag', True)


def get_ovlp(mf, cell=None, kpts=None):
    '''Get the overlap AO matrices at sampled k-points.

    Args:
        kpts : (nkpts, 3) ndarray

    Returns:
        ovlp_kpts : (nkpts, nao, nao) ndarray
    '''
    if cell is None: cell = mf.cell
    if kpts is None: kpts = mf.kpts
    # Avoid pbcopt's prescreening in the lattice sum, for better accuracy
    s = cell.pbc_intor('int1e_ovlp', hermi=0, kpts=kpts,
                       pbcopt=lib.c_null_ptr())
    s = lib.asarray(s)
    hermi_error = abs(s - s.conj().transpose(0,2,1)).max()
    if hermi_error > cell.precision and hermi_error > 1e-12:
        logger.warn(mf, '%.4g error found in overlap integrals. '
                    'cell.precision  or  cell.rcut  can be adjusted to '
                    'improve accuracy.')

    cond = np.max(lib.cond(s))
    if cond * cell.precision > 1e2:
        prec = 1e2 / cond
        rmin = max([cell.bas_rcut(ib, prec) for ib in range(cell.nbas)])
        if cell.rcut < rmin:
            logger.warn(mf, 'Singularity detected in overlap matrix.  '
                        'Integral accuracy may be not enough.\n      '
                        'You can adjust  cell.precision  or  cell.rcut  to '
                        'improve accuracy.  Recommended values are\n      '
                        'cell.precision = %.2g  or smaller.\n      '
                        'cell.rcut = %.4g  or larger.', prec, rmin)
    return s


def get_hcore(mf, cell=None, kpts=None):
    '''Get the core Hamiltonian AO matrices at sampled k-points.

    Args:
        kpts : (nkpts, 3) ndarray

    Returns:
        hcore : (nkpts, nao, nao) ndarray
    '''
    if cell is None: cell = mf.cell
    if kpts is None: kpts = mf.kpts
    if cell.pseudo:
        nuc = lib.asarray(mf.with_df.get_pp(kpts))
    else:
        nuc = lib.asarray(mf.with_df.get_nuc(kpts))
    if len(cell._ecpbas) > 0:
        nuc += lib.asarray(ecp.ecp_int(cell, kpts))
    t = lib.asarray(cell.pbc_intor('int1e_kin', 1, 1, kpts))
    return nuc + t


def get_j(mf, cell, dm_kpts, kpts, kpts_band=None):
    '''Get the Coulomb (J) AO matrix at sampled k-points.

    Args:
        dm_kpts : (nkpts, nao, nao) ndarray or a list of (nkpts,nao,nao) ndarray
            Density matrix at each k-point.  If a list of k-point DMs, eg,
            UHF alpha and beta DM, the alpha and beta DMs are contracted
            separately.  It needs to be Hermitian.

    Kwargs:
        kpts_band : (k,3) ndarray
            A list of arbitrary "band" k-points at which to evalute the matrix.

    Returns:
        vj : (nkpts, nao, nao) ndarray
        or list of vj if the input dm_kpts is a list of DMs
    '''
    return df.FFTDF(cell).get_jk(dm_kpts, kpts, kpts_band, with_k=False)[0]


def get_jk(mf, cell, dm_kpts, kpts, kpts_band=None, with_j=True, with_k=True,
           omega=None, **kwargs):
    '''Get the Coulomb (J) and exchange (K) AO matrices at sampled k-points.

    Args:
        dm_kpts : (nkpts, nao, nao) ndarray
            Density matrix at each k-point. It needs to be Hermitian.

    Kwargs:
        kpts_band : (3,) ndarray
            A list of arbitrary "band" k-point at which to evalute the matrix.

    Returns:
        vj : (nkpts, nao, nao) ndarray
        vk : (nkpts, nao, nao) ndarray
        or list of vj and vk if the input dm_kpts is a list of DMs
    '''
    return df.FFTDF(cell).get_jk(dm_kpts, kpts, kpts_band, with_j, with_k,
                                 omega, exxdiv=mf.exxdiv)

def get_fock(mf, h1e=None, s1e=None, vhf=None, dm=None, cycle=-1, diis=None,
             diis_start_cycle=None, level_shift_factor=None, damp_factor=None):
    h1e_kpts, s_kpts, vhf_kpts, dm_kpts = h1e, s1e, vhf, dm
    if h1e_kpts is None: h1e_kpts = mf.get_hcore()
    if vhf_kpts is None: vhf_kpts = mf.get_veff(mf.cell, dm_kpts)
    f_kpts = h1e_kpts + vhf_kpts
    if cycle < 0 and diis is None:  # Not inside the SCF iteration
        return f_kpts

    if diis_start_cycle is None:
        diis_start_cycle = mf.diis_start_cycle
    if level_shift_factor is None:
        level_shift_factor = mf.level_shift
    if damp_factor is None:
        damp_factor = mf.damp
    if s_kpts is None: s_kpts = mf.get_ovlp()
    if dm_kpts is None: dm_kpts = mf.make_rdm1()

    if 0 <= cycle < diis_start_cycle-1 and abs(damp_factor) > 1e-4:
        f_kpts = [mol_hf.damping(s1e, dm_kpts[k] * 0.5, f_kpts[k], damp_factor)
                  for k, s1e in enumerate(s_kpts)]
    if diis and cycle >= diis_start_cycle:
        f_kpts = diis.update(s_kpts, dm_kpts, f_kpts, mf, h1e_kpts, vhf_kpts)
    if abs(level_shift_factor) > 1e-4:
        f_kpts = [mol_hf.level_shift(s, dm_kpts[k], f_kpts[k], level_shift_factor)
                  for k, s in enumerate(s_kpts)]
    return lib.asarray(f_kpts)

def get_fermi(mf, mo_energy_kpts=None, mo_occ_kpts=None):
    '''Fermi level
    '''
    if mo_energy_kpts is None: mo_energy_kpts = mf.mo_energy
    if mo_occ_kpts is None: mo_occ_kpts = mf.mo_occ

    # mo_energy_kpts and mo_occ_kpts are k-point RHF quantities
    assert(mo_energy_kpts[0].ndim == 1)
    assert(mo_occ_kpts[0].ndim == 1)

    # occ array in mo_occ_kpts may have different size. See issue #250
    nocc = sum(mo_occ.sum() for mo_occ in mo_occ_kpts) / 2
    # nocc may not be perfect integer when smearing is enabled
    nocc = int(nocc.round(3))
    fermi = np.sort(np.hstack(mo_energy_kpts))[nocc-1]

    for k, mo_e in enumerate(mo_energy_kpts):
        mo_occ = mo_occ_kpts[k]
        if mo_occ[mo_e > fermi].sum() > 1.:
            logger.warn(mf, 'Occupied band above Fermi level: \n'
                        'k=%d, mo_e=%s, mo_occ=%s', k, mo_e, mo_occ)
    return fermi

def get_occ(mf, mo_energy_kpts=None, mo_coeff_kpts=None):
    '''Label the occupancies for each orbital for sampled k-points.

    This is a k-point version of scf.hf.SCF.get_occ
    '''
    if mo_energy_kpts is None: mo_energy_kpts = mf.mo_energy

    nkpts = len(mo_energy_kpts)
    nocc = mf.cell.tot_electrons(nkpts) // 2

    mo_energy = np.sort(np.hstack(mo_energy_kpts))
    fermi = mo_energy[nocc-1]
    mo_occ_kpts = []
    for mo_e in mo_energy_kpts:
        mo_occ_kpts.append((mo_e <= fermi).astype(np.double) * 2)

    if nocc < mo_energy.size:
        logger.info(mf, 'HOMO = %.12g  LUMO = %.12g',
                    mo_energy[nocc-1], mo_energy[nocc])
        if mo_energy[nocc-1]+1e-3 > mo_energy[nocc]:
            logger.warn(mf, 'HOMO %.12g == LUMO %.12g',
                        mo_energy[nocc-1], mo_energy[nocc])
    else:
        logger.info(mf, 'HOMO = %.12g', mo_energy[nocc-1])

    if mf.verbose >= logger.DEBUG:
        np.set_printoptions(threshold=len(mo_energy))
        logger.debug(mf, '     k-point                  mo_energy')
        for k,kpt in enumerate(mf.cell.get_scaled_kpts(mf.kpts)):
            logger.debug(mf, '  %2d (%6.3f %6.3f %6.3f)   %s %s',
                         k, kpt[0], kpt[1], kpt[2],
                         mo_energy_kpts[k][mo_occ_kpts[k]> 0],
                         mo_energy_kpts[k][mo_occ_kpts[k]==0])
        np.set_printoptions(threshold=1000)

    return mo_occ_kpts


def get_grad(mo_coeff_kpts, mo_occ_kpts, fock):
    '''
    returns 1D array of gradients, like non K-pt version
    note that occ and virt indices of different k pts now occur
    in sequential patches of the 1D array
    '''
    nkpts = len(mo_occ_kpts)
    grad_kpts = [mol_hf.get_grad(mo_coeff_kpts[k], mo_occ_kpts[k], fock[k])
                 for k in range(nkpts)]
    return np.hstack(grad_kpts)


def make_rdm1(mo_coeff_kpts, mo_occ_kpts, **kwargs):
    '''One particle density matrices for all k-points.

    Returns:
        dm_kpts : (nkpts, nao, nao) ndarray
    '''
    nkpts = len(mo_occ_kpts)
    dm_kpts = [mol_hf.make_rdm1(mo_coeff_kpts[k], mo_occ_kpts[k])
               for k in range(nkpts)]
    return lib.asarray(dm_kpts)


def energy_elec(mf, dm_kpts=None, h1e_kpts=None, vhf_kpts=None):
    '''Following pyscf.scf.hf.energy_elec()
    '''
    if dm_kpts is None: dm_kpts = mf.make_rdm1()
    if h1e_kpts is None: h1e_kpts = mf.get_hcore()
    if vhf_kpts is None: vhf_kpts = mf.get_veff(mf.cell, dm_kpts)

    nkpts = len(dm_kpts)
    e1 = 1./nkpts * np.einsum('kij,kji', dm_kpts, h1e_kpts)
    e_coul = 1./nkpts * np.einsum('kij,kji', dm_kpts, vhf_kpts) * 0.5
    mf.scf_summary['e1'] = e1.real
    mf.scf_summary['e2'] = e_coul.real
    logger.debug(mf, 'E1 = %s  E_coul = %s', e1, e_coul)
    if CHECK_COULOMB_IMAG and abs(e_coul.imag > mf.cell.precision*10):
        logger.warn(mf, "Coulomb energy has imaginary part %s. "
                    "Coulomb integrals (e-e, e-N) may not converge !",
                    e_coul.imag)
    return (e1+e_coul).real, e_coul.real


def analyze(mf, verbose=logger.DEBUG, with_meta_lowdin=WITH_META_LOWDIN,
            **kwargs):
    '''Analyze the given SCF object:  print orbital energies, occupancies;
    print orbital coefficients; Mulliken population analysis; Dipole moment
    '''
    mf.dump_scf_summary(verbose)

    mo_occ = mf.mo_occ
    mo_coeff = mf.mo_coeff
    ovlp_ao = mf.get_ovlp()
    dm = mf.make_rdm1(mo_coeff, mo_occ)
    if with_meta_lowdin:
        return mf.mulliken_meta(mf.cell, dm, s=ovlp_ao, verbose=verbose)
    else:
        raise NotImplementedError
        #return mf.mulliken_pop(mf.cell, dm, s=ovlp_ao, verbose=verbose)


def mulliken_meta(cell, dm_ao_kpts, verbose=logger.DEBUG,
                  pre_orth_method=PRE_ORTH_METHOD, s=None):
    '''A modified Mulliken population analysis, based on meta-Lowdin AOs.

    Note this function only computes the Mulliken population for the gamma
    point density matrix.
    '''
    from pyscf.lo import orth
    if s is None:
        s = get_ovlp(cell)
    log = logger.new_logger(cell, verbose)
    log.note('Analyze output for *gamma point*')
    log.info('    To include the contributions from k-points, transform to a '
             'supercell then run the population analysis on the supercell\n'
             '        from pyscf.pbc.tools import k2gamma\n'
             '        k2gamma.k2gamma(mf).mulliken_meta()')
    log.note("KRHF mulliken_meta")
    dm_ao_gamma = dm_ao_kpts[0,:,:].real
    s_gamma = s[0,:,:].real
    orth_coeff = orth.orth_ao(cell, 'meta_lowdin', pre_orth_method, s=s_gamma)
    c_inv = np.dot(orth_coeff.T, s_gamma)
    dm = reduce(np.dot, (c_inv, dm_ao_gamma, c_inv.T.conj()))

    log.note(' ** Mulliken pop on meta-lowdin orthogonal AOs **')
    return mol_hf.mulliken_pop(cell, dm, np.eye(orth_coeff.shape[0]), log)


def canonicalize(mf, mo_coeff_kpts, mo_occ_kpts, fock=None):
    if fock is None:
        dm = mf.make_rdm1(mo_coeff_kpts, mo_occ_kpts)
        fock = mf.get_fock(dm=dm)
    mo_coeff = []
    mo_energy = []
    for k, mo in enumerate(mo_coeff_kpts):
        mo1 = np.empty_like(mo)
        mo_e = np.empty_like(mo_occ_kpts[k])
        occidx = mo_occ_kpts[k] == 2
        viridx = ~occidx
        for idx in (occidx, viridx):
            if np.count_nonzero(idx) > 0:
                orb = mo[:,idx]
                f1 = reduce(np.dot, (orb.T.conj(), fock[k], orb))
                e, c = scipy.linalg.eigh(f1)
                mo1[:,idx] = np.dot(orb, c)
                mo_e[idx] = e
        mo_coeff.append(mo1)
        mo_energy.append(mo_e)
    return mo_energy, mo_coeff


def init_guess_by_chkfile(cell, chkfile_name, project=None, kpts=None):
    '''Read the KHF results from checkpoint file, then project it to the
    basis defined by ``cell``

    Returns:
        Density matrix, 3D ndarray
    '''
    from pyscf.pbc.scf import kuhf
    dm = kuhf.init_guess_by_chkfile(cell, chkfile_name, project, kpts)
    return dm[0] + dm[1]


def dip_moment(cell, dm_kpts, unit='Debye', verbose=logger.NOTE,
               grids=None, rho=None, kpts=np.zeros((1,3))):
    ''' Dipole moment in the unit cell (is it well defined)?

    Args:
         cell : an instance of :class:`Cell`

         dm_kpts (a list of ndarrays) : density matrices of k-points

    Return:
        A list: the dipole moment on x, y and z components
    '''
    from pyscf.pbc.dft import gen_grid
    from pyscf.pbc.dft import numint
    if grids is None:
        grids = gen_grid.UniformGrids(cell)
    if rho is None:
        rho = numint.KNumInt().get_rho(cell, dm_kpts, grids, kpts, cell.max_memory)
    return pbchf.dip_moment(cell, dm_kpts, unit, verbose, grids, rho, kpts)

def get_rho(mf, dm=None, grids=None, kpts=None):
    '''Compute density in real space
    '''
    from pyscf.pbc.dft import gen_grid
    from pyscf.pbc.dft import numint
    if dm is None:
        dm = mf.make_rdm1()
    if getattr(dm[0], 'ndim', None) != 2:  # KUHF
        dm = dm[0] + dm[1]
    if grids is None:
        grids = gen_grid.UniformGrids(mf.cell)
    if kpts is None:
        kpts = mf.kpts
    ni = numint.KNumInt()
    return ni.get_rho(mf.cell, dm, grids, kpts, mf.max_memory)

def as_scanner(mf):
    import copy
    if isinstance(mf, lib.SinglePointScanner):
        return mf

    logger.info(mf, 'Create scanner for %s', mf.__class__)

    class SCF_Scanner(mf.__class__, lib.SinglePointScanner):
        def __init__(self, mf_obj):
            self.__dict__.update(mf_obj.__dict__)

        def __call__(self, cell_or_geom, **kwargs):
            from pyscf.pbc import gto
            if isinstance(cell_or_geom, gto.Cell):
                cell = cell_or_geom
            else:
                cell = self.cell.set_geom_(cell_or_geom, inplace=False)

            # Cleanup intermediates associated to the pervious mol object
            self.reset(cell)

            if 'dm0' in kwargs:
                dm0 = kwargs.pop('dm0')
            elif self.mo_coeff is None:
                dm0 = None
            elif self.chkfile and h5py.is_hdf5(self.chkfile):
                dm0 = self.from_chk(self.chkfile)
            else:
                dm0 = self.make_rdm1()
                # dm0 form last calculation cannot be used in the current
                # calculation if a completely different system is given.
                # Obviously, the systems are very different if the number of
                # basis functions are different.
                # TODO: A robust check should include more comparison on
                # various attributes between current `mol` and the `mol` in
                # last calculation.
                if dm0.shape[-1] != cell.nao_nr():
                    #TODO:
                    #from pyscf.scf import addons
                    #if numpy.any(last_mol.atom_charges() != mol.atom_charges()):
                    #    dm0 = None
                    #elif non-relativistic:
                    #    addons.project_dm_nr2nr(last_mol, dm0, last_mol)
                    #else:
                    #    addons.project_dm_r2r(last_mol, dm0, last_mol)
                    dm0 = None
            self.mo_coeff = None  # To avoid last mo_coeff being used by SOSCF
            e_tot = self.kernel(dm0=dm0, **kwargs)
            return e_tot

    return SCF_Scanner(mf)


class KSCF(pbchf.SCF):
    '''SCF base class with k-point sampling.

    Compared to molecular SCF, some members such as mo_coeff, mo_occ
    now have an additional first dimension for the k-points,
    e.g. mo_coeff is (nkpts, nao, nao) ndarray

    Attributes:
        kpts : (nks,3) ndarray
            The sampling k-points in Cartesian coordinates, in units of 1/Bohr.
    '''
    conv_tol_grad = getattr(__config__, 'pbc_scf_KSCF_conv_tol_grad', None)
    direct_scf = getattr(__config__, 'pbc_scf_SCF_direct_scf', True)

    def __init__(self, cell, kpts=np.zeros((1,3)),
                 exxdiv=getattr(__config__, 'pbc_scf_SCF_exxdiv', 'ewald')):
        if not cell._built:
            sys.stderr.write('Warning: cell.build() is not called in input\n')
            cell.build()
        self.cell = cell
        mol_hf.SCF.__init__(self, cell)

        self.with_df = df.FFTDF(cell)
        # Range separation JK builder
        self.rsjk = None

        self.exxdiv = exxdiv
        self.kpts = kpts
        self.conv_tol = cell.precision * 10

        self.exx_built = False
        self._keys = self._keys.union(['cell', 'exx_built', 'exxdiv', 'with_df', 'rsjk'])

    @property
    def kpts(self):
        if 'kpts' in self.__dict__:
            # To handle the attribute kpt loaded from chkfile
            self.kpt = self.__dict__.pop('kpts')
        return self.with_df.kpts
    @kpts.setter
    def kpts(self, x):
        self.with_df.kpts = np.reshape(x, (-1,3))
        if self.rsjk:
            self.rsjk.kpts = self.with_df.kpts

    @property
    def mo_energy_kpts(self):
        return self.mo_energy

    @property
    def mo_coeff_kpts(self):
        return self.mo_coeff

    @property
    def mo_occ_kpts(self):
        return self.mo_occ

    def dump_flags(self, verbose=None):
        mol_hf.SCF.dump_flags(self, verbose)
        logger.info(self, '\n')
        logger.info(self, '******** PBC SCF flags ********')
        logger.info(self, 'N kpts = %d', len(self.kpts))
        logger.debug(self, 'kpts = %s', self.kpts)
        logger.info(self, 'Exchange divergence treatment (exxdiv) = %s', self.exxdiv)
        # "vcut_ws" precomputing is triggered by pbc.tools.pbc.get_coulG
        #if self.exxdiv == 'vcut_ws':
        #    if self.exx_built is False:
        #        self.precompute_exx()
        #    logger.info(self, 'WS alpha = %s', self.exx_alpha)
        cell = self.cell
        if ((cell.dimension >= 2 and cell.low_dim_ft_type != 'inf_vacuum') and
            isinstance(self.exxdiv, str) and self.exxdiv.lower() == 'ewald'):
            madelung = tools.pbc.madelung(cell, [self.kpts])
            logger.info(self, '    madelung (= occupied orbital energy shift) = %s', madelung)
            nkpts = len(self.kpts)
            # FIXME: consider the fractional num_electron or not? This maybe
            # relates to the charged system.
            nelectron = float(self.cell.tot_electrons(nkpts)) / nkpts
            logger.info(self, '    Total energy shift due to Ewald probe charge'
                        ' = -1/2 * Nelec*madelung = %.12g',
                        madelung*nelectron * -.5)
        logger.info(self, 'DF object = %s', self.with_df)
        if not getattr(self.with_df, 'build', None):
            # .dump_flags() is called in pbc.df.build function
            self.with_df.dump_flags(verbose)
        return self

    def check_sanity(self):
        mol_hf.SCF.check_sanity(self)
        self.with_df.check_sanity()
        if (isinstance(self.exxdiv, str) and self.exxdiv.lower() != 'ewald' and
            isinstance(self.with_df, df.df.DF)):
            logger.warn(self, 'exxdiv %s is not supported in DF or MDF',
                        self.exxdiv)
        return self

    def build(self, cell=None):
        if cell is None:
            cell = self.cell
        #if self.exxdiv == 'vcut_ws':
        #    self.precompute_exx()

        if 'kpts' in self.__dict__:
            # To handle the attribute kpts loaded from chkfile
            self.kpts = self.__dict__.pop('kpts')

        if self.rsjk:
            if not np.all(self.rsjk.kpts == self.kpts):
                self.rsjk = self.rsjk.__class__(cell, self.kpts)
            self.rsjk.build(direct_scf_tol=self.direct_scf_tol)

        if self.verbose >= logger.WARN:
            self.check_sanity()
        return self

    def get_init_guess(self, cell=None, key='minao'):
        if cell is None:
            cell = self.cell
        dm_kpts = None
        key = key.lower()
        if key == '1e' or key == 'hcore':
            dm_kpts = self.init_guess_by_1e(cell)
        elif getattr(cell, 'natm', 0) == 0:
            logger.info(self, 'No atom found in cell. Use 1e initial guess')
            dm_kpts = self.init_guess_by_1e(cell)
        elif key == 'atom':
            dm = self.init_guess_by_atom(cell)
        elif key[:3] == 'chk':
            try:
                dm_kpts = self.from_chk()
            except (IOError, KeyError):
                logger.warn(self, 'Fail to read %s. Use MINAO initial guess',
                            self.chkfile)
                dm = self.init_guess_by_minao(cell)
        else:
            dm = self.init_guess_by_minao(cell)

        if dm_kpts is None:
            dm_kpts = lib.asarray([dm]*len(self.kpts))

        ne = np.einsum('kij,kji->', dm_kpts, self.get_ovlp(cell)).real
        # FIXME: consider the fractional num_electron or not? This maybe
        # relate to the charged system.
        nkpts = len(self.kpts)
        nelectron = float(self.cell.tot_electrons(nkpts))
        if abs(ne - nelectron) > 1e-7*nkpts:
            logger.debug(self, 'Big error detected in the electron number '
                         'of initial guess density matrix (Ne/cell = %g)!\n'
                         '  This can cause huge error in Fock matrix and '
                         'lead to instability in SCF for low-dimensional '
                         'systems.\n  DM is normalized wrt the number '
                         'of electrons %s', ne/nkpts, nelectron/nkpts)
            dm_kpts *= (nelectron / ne).reshape(-1,1,1)
        return dm_kpts

    def init_guess_by_1e(self, cell=None):
        if cell is None: cell = self.cell
        if cell.dimension < 3:
            logger.warn(self, 'Hcore initial guess is not recommended in '
                        'the SCF of low-dimensional systems.')
        return mol_hf.SCF.init_guess_by_1e(self, cell)

    get_hcore = get_hcore
    get_ovlp = get_ovlp
    get_fock = get_fock
    get_occ = get_occ
    energy_elec = energy_elec
    get_fermi = get_fermi

    def get_j(self, cell=None, dm_kpts=None, hermi=1, kpts=None,
              kpts_band=None, omega=None):
        return self.get_jk(cell, dm_kpts, hermi, kpts, kpts_band,
                           with_k=False, omega=omega)[0]

    def get_k(self, cell=None, dm_kpts=None, hermi=1, kpts=None,
              kpts_band=None, omega=None):
        return self.get_jk(cell, dm_kpts, hermi, kpts, kpts_band,
                           with_j=False, omega=omega)[1]

    def get_jk(self, cell=None, dm_kpts=None, hermi=1, kpts=None, kpts_band=None,
               with_j=True, with_k=True, omega=None, **kwargs):
        if cell is None: cell = self.cell
        if kpts is None: kpts = self.kpts
        if dm_kpts is None: dm_kpts = self.make_rdm1()
        cpu0 = (logger.process_clock(), logger.perf_counter())
        if self.rsjk:
            vj, vk = self.rsjk.get_jk(dm_kpts, hermi, kpts, kpts_band,
                                      with_j, with_k, omega, self.exxdiv)
        else:
            vj, vk = self.with_df.get_jk(dm_kpts, hermi, kpts, kpts_band,
                                         with_j, with_k, omega, self.exxdiv)
        logger.timer(self, 'vj and vk', *cpu0)
        return vj, vk

    def get_veff(self, cell=None, dm_kpts=None, dm_last=0, vhf_last=0, hermi=1,
                 kpts=None, kpts_band=None):
        '''Hartree-Fock potential matrix for the given density matrix.
        See :func:`scf.hf.get_veff` and :func:`scf.hf.RHF.get_veff`
        '''
        if dm_kpts is None:
            dm_kpts = self.make_rdm1()
        if self.rsjk and self.direct_scf:
            # Enable direct-SCF for real space JK builder
            ddm = dm_kpts - dm_last
            vj, vk = self.get_jk(cell, ddm, hermi, kpts, kpts_band)
            return vhf_last + vj - vk * .5
        else:
            vj, vk = self.get_jk(cell, dm_kpts, hermi, kpts, kpts_band)
            return vj - vk * .5

    def analyze(self, verbose=None, with_meta_lowdin=WITH_META_LOWDIN,
                **kwargs):
        if verbose is None: verbose = self.verbose
        return analyze(self, verbose, with_meta_lowdin, **kwargs)

    def get_grad(self, mo_coeff_kpts, mo_occ_kpts, fock=None):
        '''
        returns 1D array of gradients, like non K-pt version
        note that occ and virt indices of different k pts now occur
        in sequential patches of the 1D array
        '''
        if fock is None:
            dm1 = self.make_rdm1(mo_coeff_kpts, mo_occ_kpts)
            fock = self.get_hcore(self.cell, self.kpts) + self.get_veff(self.cell, dm1)
        return get_grad(mo_coeff_kpts, mo_occ_kpts, fock)

    def eig(self, h_kpts, s_kpts):
        nkpts = len(h_kpts)
        eig_kpts = []
        mo_coeff_kpts = []

        for k in range(nkpts):
            e, c = self._eigh(h_kpts[k], s_kpts[k])
            eig_kpts.append(e)
            mo_coeff_kpts.append(c)
        return eig_kpts, mo_coeff_kpts

    def make_rdm1(self, mo_coeff_kpts=None, mo_occ_kpts=None, **kwargs):
        if mo_coeff_kpts is None:
            # Note: this is actually "self.mo_coeff_kpts"
            # which is stored in self.mo_coeff of the scf.hf.RHF superclass
            mo_coeff_kpts = self.mo_coeff
        if mo_occ_kpts is None:
            # Note: this is actually "self.mo_occ_kpts"
            # which is stored in self.mo_occ of the scf.hf.RHF superclass
            mo_occ_kpts = self.mo_occ

        return make_rdm1(mo_coeff_kpts, mo_occ_kpts, **kwargs)

    def get_bands(self, kpts_band, cell=None, dm_kpts=None, kpts=None):
        '''Get energy bands at the given (arbitrary) 'band' k-points.

        Returns:
            mo_energy : (nmo,) ndarray or a list of (nmo,) ndarray
                Bands energies E_n(k)
            mo_coeff : (nao, nmo) ndarray or a list of (nao,nmo) ndarray
                Band orbitals psi_n(k)
        '''
        if cell is None: cell = self.cell
        if dm_kpts is None: dm_kpts = self.make_rdm1()
        if kpts is None: kpts = self.kpts

        kpts_band = np.asarray(kpts_band)
        single_kpt_band = (kpts_band.ndim == 1)
        kpts_band = kpts_band.reshape(-1,3)

        fock = self.get_hcore(cell, kpts_band)
        fock = fock + self.get_veff(cell, dm_kpts, kpts=kpts, kpts_band=kpts_band)
        s1e = self.get_ovlp(cell, kpts_band)
        mo_energy, mo_coeff = self.eig(fock, s1e)
        if single_kpt_band:
            mo_energy = mo_energy[0]
            mo_coeff = mo_coeff[0]
        return mo_energy, mo_coeff

    def init_guess_by_chkfile(self, chk=None, project=None, kpts=None):
        if chk is None: chk = self.chkfile
        if kpts is None: kpts = self.kpts
        return init_guess_by_chkfile(self.cell, chk, project, kpts)
    def from_chk(self, chk=None, project=None, kpts=None):
        return self.init_guess_by_chkfile(chk, project, kpts)

    def dump_chk(self, envs):
        if self.chkfile:
            mol_hf.SCF.dump_chk(self, envs)
            with h5py.File(self.chkfile, 'a') as fh5:
                fh5['scf/kpts'] = self.kpts
        return self

    def mulliken_meta(self, cell=None, dm=None, verbose=logger.DEBUG,
                      pre_orth_method=PRE_ORTH_METHOD, s=None):
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        if s is None: s = self.get_ovlp(cell)
        return mulliken_meta(cell, dm, s=s, verbose=verbose,
                             pre_orth_method=pre_orth_method)

    def mulliken_pop(self):
        raise NotImplementedError

    get_rho = get_rho

    @lib.with_doc(dip_moment.__doc__)
    def dip_moment(self, cell=None, dm=None, unit='Debye', verbose=logger.NOTE,
                   **kwargs):
        rho = kwargs.pop('rho', None)
        if rho is None:
            rho = self.get_rho(dm)
        if cell is None:
            cell = self.cell
        return dip_moment(cell, dm, unit, verbose, rho=rho, kpts=self.kpts, **kwargs)

    canonicalize = canonicalize

    def density_fit(self, auxbasis=None, with_df=None):
        from pyscf.pbc.df import df_jk
        return df_jk.density_fit(self, auxbasis, with_df=with_df)

    def rs_density_fit(self, auxbasis=None, with_df=None):
        from pyscf.pbc.df import rsdf_jk
        return rsdf_jk.density_fit(self, auxbasis, with_df=with_df)

    def mix_density_fit(self, auxbasis=None, with_df=None):
        from pyscf.pbc.df import mdf_jk
        return mdf_jk.density_fit(self, auxbasis, with_df=with_df)

    def jk_method(self, J='FFTDF', K=None):
        '''
        Set up the schemes to evaluate Coulomb and exchange matrix

        FFTDF: planewave density fitting using Fast Fourier Transform
        AFTDF: planewave density fitting using analytic Fourier Transform
        GDF: Gaussian density fitting
        MDF: Gaussian and planewave mix density fitting
        RS: range-separation JK builder
        RSDF: range-separation density fitting
        '''
        if K is None:
            K = J

        if J != K:
            raise NotImplementedError('J != K')

        if 'DF' in J or 'DF' in K:
            if 'DF' in J and 'DF' in K:
                assert J == K
            else:
                df_method = J if 'DF' in J else K
                self.with_df = getattr(df, df_method)(self.cell, self.kpts)

        if 'RS' in J or 'RS' in K:
            self.rsjk = RangeSeparationJKBuilder(self.cell, self.kpts)
            self.rsjk.verbose = self.verbose

        # For nuclear attraction
        if J == 'RS' and K == 'RS' and not isinstance(self.with_df, df.GDF):
            self.with_df = df.GDF(self.cell, self.kpts)

        nuc = self.with_df.__class__.__name__
        logger.debug1(self, 'Apply %s for J, %s for K, %s for nuc', J, K, nuc)
        return self

    def stability(self,
                  internal=getattr(__config__, 'pbc_scf_KSCF_stability_internal', True),
                  external=getattr(__config__, 'pbc_scf_KSCF_stability_external', False),
                  verbose=None):
        from pyscf.pbc.scf.stability import rhf_stability
        return rhf_stability(self, internal, external, verbose)

    def newton(self):
        from pyscf.pbc.scf import newton_ah
        return newton_ah.newton(self)

    def sfx2c1e(self):
        from pyscf.pbc.x2c import sfx2c1e
        return sfx2c1e.sfx2c1e(self)
    x2c = x2c1e = sfx2c1e

    def to_rhf(self, mf):
        '''Convert the input mean-field object to a KRHF/KROHF/KRKS/KROKS object'''
        return addons.convert_to_rhf(mf)

    def to_uhf(self, mf):
        '''Convert the input mean-field object to a KUHF/KUKS object'''
        return addons.convert_to_uhf(mf)

    def to_ghf(self, mf):
        '''Convert the input mean-field object to a KGHF/KGKS object'''
        return addons.convert_to_ghf(mf)

    as_scanner = as_scanner


class KRHF(KSCF, pbchf.RHF):
    def check_sanity(self):
        cell = self.cell
        if cell.spin != 0 and len(self.kpts) % 2 != 0:
            logger.warn(self, 'Problematic nelec %s and number of k-points %d '
                        'found in KRHF method.', cell.nelec, len(self.kpts))
        return KSCF.check_sanity(self)

    def convert_from_(self, mf):
        '''Convert given mean-field object to KRHF'''
        addons.convert_to_rhf(mf, self)
        return self

    def nuc_grad_method(self):
        from pyscf.pbc.grad import krhf
        return krhf.Gradients(self)

del(WITH_META_LOWDIN, PRE_ORTH_METHOD)


if __name__ == '__main__':
    from pyscf.pbc import gto
    cell = gto.Cell()
    cell.atom = '''
    He 0 0 1
    He 1 0 1
    '''
    cell.basis = '321g'
    cell.a = np.eye(3) * 3
    cell.mesh = [11] * 3
    cell.verbose = 5
    cell.build()
    mf = KRHF(cell, [2,1,1])
    mf.kernel()
    mf.analyze()


# pyscf/pbc/scf/krohf.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Restricted open-shell Hartree-Fock for periodic systems with k-point sampling
'''

from functools import reduce
import numpy as np
import scipy.linalg
from pyscf.scf import hf as mol_hf
from pyscf.pbc.scf import khf
from pyscf.pbc.scf import kuhf
from pyscf.pbc.scf import rohf as pbcrohf
from pyscf import lib
from pyscf.lib import logger
from pyscf.pbc.scf import addons
from pyscf import __config__

WITH_META_LOWDIN = getattr(__config__, 'pbc_scf_analyze_with_meta_lowdin', True)
PRE_ORTH_METHOD = getattr(__config__, 'pbc_scf_analyze_pre_orth_method', 'ANO')


def make_rdm1(mo_coeff_kpts, mo_occ_kpts, **kwargs):
    '''Alpha and beta spin one particle density matrices for all k-points.

    Returns:
        dm_kpts : (2, nkpts, nao, nao) ndarray
    '''
    dma = []
    dmb = []
    for k, occ in enumerate(mo_occ_kpts):
        mo_a = mo_coeff_kpts[k][:,occ> 0]
        mo_b = mo_coeff_kpts[k][:,occ==2]
        dma.append(np.dot(mo_a, mo_a.conj().T))
        dmb.append(np.dot(mo_b, mo_b.conj().T))
    return lib.asarray((dma,dmb))

def get_fock(mf, h1e=None, s1e=None, vhf=None, dm=None, cycle=-1, diis=None,
             diis_start_cycle=None, level_shift_factor=None, damp_factor=None):
    h1e_kpts, s_kpts, vhf_kpts, dm_kpts = h1e, s1e, vhf, dm
    if h1e_kpts is None: h1e_kpts = mf.get_hcore()
    if vhf_kpts is None: vhf_kpts = mf.get_veff(mf.cell, dm_kpts)
    focka = h1e_kpts + vhf_kpts[0]
    fockb = h1e_kpts + vhf_kpts[1]
    f_kpts = get_roothaan_fock((focka,fockb), dm, s1e)
    if cycle < 0 and diis is None:  # Not inside the SCF iteration
        return f_kpts

    if diis_start_cycle is None:
        diis_start_cycle = mf.diis_start_cycle
    if level_shift_factor is None:
        level_shift_factor = mf.level_shift
    if damp_factor is None:
        damp_factor = mf.damp
    if s_kpts is None: s_kpts = mf.get_ovlp()
    if dm_kpts is None: dm_kpts = mf.make_rdm1()

    dm_sf = dm_kpts[0] + dm_kpts[1]
    if 0 <= cycle < diis_start_cycle-1 and abs(damp_factor) > 1e-4:
        raise NotImplementedError('ROHF Fock-damping')
    if diis and cycle >= diis_start_cycle:
        f_kpts = diis.update(s_kpts, dm_sf, f_kpts, mf, h1e_kpts, vhf_kpts)
    if abs(level_shift_factor) > 1e-4:
        f_kpts = [mol_hf.level_shift(s, dm_sf[k]*.5, f_kpts[k], level_shift_factor)
                  for k, s in enumerate(s_kpts)]
    f_kpts = lib.tag_array(lib.asarray(f_kpts), focka=focka, fockb=fockb)
    return f_kpts

def get_roothaan_fock(focka_fockb, dma_dmb, s):
    '''Roothaan's effective fock.

    ======== ======== ====== =========
    space     closed   open   virtual
    ======== ======== ====== =========
    closed      Fc      Fb     Fc
    open        Fb      Fc     Fa
    virtual     Fc      Fa     Fc
    ======== ======== ====== =========

    where Fc = (Fa + Fb) / 2

    Returns:
        Roothaan effective Fock matrix
    '''
    nkpts = len(s)
    nao = s[0].shape[0]
    focka, fockb = focka_fockb
    dma, dmb = dma_dmb
    fock_kpts = []
    for k in range(nkpts):
        fc = (focka[k] + fockb[k]) * .5
        pc = np.dot(dmb[k], s[k])
        po = np.dot(dma[k]-dmb[k], s[k])
        pv = np.eye(nao) - np.dot(dma[k], s[k])
        fock  = reduce(np.dot, (pc.conj().T, fc, pc)) * .5
        fock += reduce(np.dot, (po.conj().T, fc, po)) * .5
        fock += reduce(np.dot, (pv.conj().T, fc, pv)) * .5
        fock += reduce(np.dot, (po.conj().T, fockb[k], pc))
        fock += reduce(np.dot, (po.conj().T, focka[k], pv))
        fock += reduce(np.dot, (pv.conj().T, fc, pc))
        fock_kpts.append(fock + fock.conj().T)
    fock_kpts = lib.tag_array(np.asarray(fock_kpts), focka=focka, fockb=fockb)
    return fock_kpts

def get_occ(mf, mo_energy_kpts=None, mo_coeff_kpts=None):
    '''Label the occupancies for each orbital for sampled k-points.

    This is a k-point version of scf.hf.SCF.get_occ
    '''

    if mo_energy_kpts is None: mo_energy_kpts = mf.mo_energy
    if getattr(mo_energy_kpts[0], 'mo_ea', None) is not None:
        mo_ea_kpts = [x.mo_ea for x in mo_energy_kpts]
        mo_eb_kpts = [x.mo_eb for x in mo_energy_kpts]
    else:
        mo_ea_kpts = mo_eb_kpts = mo_energy_kpts

    nocc_a, nocc_b = mf.nelec
    mo_energy_kpts1 = np.hstack(mo_energy_kpts)
    mo_energy = np.sort(mo_energy_kpts1)
    if nocc_b > 0:
        core_level = mo_energy[nocc_b-1]
    else:
        core_level = -1e9
    if nocc_a == nocc_b:
        fermi = core_level
    else:
        mo_ea_kpts1 = np.hstack(mo_ea_kpts)
        mo_ea = np.sort(mo_ea_kpts1[mo_energy_kpts1 > core_level])
        fermi = mo_ea[nocc_a - nocc_b - 1]

    mo_occ_kpts = []
    for k, mo_e in enumerate(mo_energy_kpts):
        occ = np.zeros_like(mo_e)
        occ[mo_e <= core_level] = 2
        if nocc_a != nocc_b:
            occ[(mo_e > core_level) & (mo_ea_kpts[k] <= fermi)] = 1
        mo_occ_kpts.append(occ)

    if nocc_a < len(mo_energy):
        logger.info(mf, 'HOMO = %.12g  LUMO = %.12g',
                    mo_energy[nocc_a-1], mo_energy[nocc_a])
    else:
        logger.info(mf, 'HOMO = %.12g', mo_energy[nocc_a-1])

    np.set_printoptions(threshold=len(mo_energy))
    if mf.verbose >= logger.DEBUG:
        logger.debug(mf, '                  Roothaan           | alpha              | beta')
        for k,kpt in enumerate(mf.cell.get_scaled_kpts(mf.kpts)):
            core_idx = mo_occ_kpts[k] == 2
            open_idx = mo_occ_kpts[k] == 1
            vir_idx = mo_occ_kpts[k] == 0
            logger.debug(mf, '  kpt %2d (%6.3f %6.3f %6.3f)',
                         k, kpt[0], kpt[1], kpt[2])
            if np.count_nonzero(core_idx) > 0:
                logger.debug(mf, '  Highest 2-occ = %18.15g | %18.15g | %18.15g',
                             max(mo_energy_kpts[k][core_idx]),
                             max(mo_ea_kpts[k][core_idx]), max(mo_eb_kpts[k][core_idx]))
            if np.count_nonzero(vir_idx) > 0:
                logger.debug(mf, '  Lowest 0-occ =  %18.15g | %18.15g | %18.15g',
                             min(mo_energy_kpts[k][vir_idx]),
                             min(mo_ea_kpts[k][vir_idx]), min(mo_eb_kpts[k][vir_idx]))
            for i in np.where(open_idx)[0]:
                logger.debug(mf, '  1-occ =         %18.15g | %18.15g | %18.15g',
                             mo_energy_kpts[k][i], mo_ea_kpts[k][i], mo_eb_kpts[k][i])

        logger.debug(mf, '     k-point                  Roothaan mo_energy')
        for k,kpt in enumerate(mf.cell.get_scaled_kpts(mf.kpts)):
            logger.debug(mf, '  %2d (%6.3f %6.3f %6.3f)   %s %s',
                         k, kpt[0], kpt[1], kpt[2],
                         mo_energy_kpts[k][mo_occ_kpts[k]> 0],
                         mo_energy_kpts[k][mo_occ_kpts[k]==0])

    if mf.verbose >= logger.DEBUG1:
        logger.debug1(mf, '     k-point                  alpha mo_energy')
        for k,kpt in enumerate(mf.cell.get_scaled_kpts(mf.kpts)):
            logger.debug1(mf, '  %2d (%6.3f %6.3f %6.3f)   %s %s',
                          k, kpt[0], kpt[1], kpt[2],
                          mo_ea_kpts[k][mo_occ_kpts[k]> 0],
                          mo_ea_kpts[k][mo_occ_kpts[k]==0])
        logger.debug1(mf, '     k-point                  beta  mo_energy')
        for k,kpt in enumerate(mf.cell.get_scaled_kpts(mf.kpts)):
            logger.debug1(mf, '  %2d (%6.3f %6.3f %6.3f)   %s %s',
                          k, kpt[0], kpt[1], kpt[2],
                          mo_eb_kpts[k][mo_occ_kpts[k]==2],
                          mo_eb_kpts[k][mo_occ_kpts[k]!=2])
    np.set_printoptions(threshold=1000)

    return mo_occ_kpts


energy_elec = kuhf.energy_elec
dip_moment = kuhf.dip_moment
get_rho = kuhf.get_rho


@lib.with_doc(khf.mulliken_meta.__doc__)
def mulliken_meta(cell, dm_ao_kpts, verbose=logger.DEBUG,
                  pre_orth_method=PRE_ORTH_METHOD, s=None):
    '''Mulliken population analysis, based on meta-Lowdin AOs.

    Note this function only computes the Mulliken population for the gamma
    point density matrix.
    '''
    dm = dm_ao_kpts[0] + dm_ao_kpts[1]
    return khf.mulliken_meta(cell, dm, verbose, pre_orth_method, s)


def canonicalize(mf, mo_coeff_kpts, mo_occ_kpts, fock=None):
    '''Canonicalization diagonalizes the ROHF Fock matrix within occupied,
    virtual subspaces separatedly (without change occupancy).
    '''
    if fock is None:
        dm = mf.make_rdm1(mo_coeff_kpts, mo_occ_kpts)
        fock = mf.get_fock(dm=dm)

    mo_coeff = []
    mo_energy = []
    for k, mo in enumerate(mo_coeff_kpts):
        mo1 = np.empty_like(mo)
        mo_e = np.empty_like(mo_occ_kpts[k])
        coreidx = mo_occ_kpts[k] == 2
        openidx = mo_occ_kpts[k] == 1
        viridx = mo_occ_kpts[k] == 0
        for idx in (coreidx, openidx, viridx):
            if np.count_nonzero(idx) > 0:
                orb = mo[:,idx]
                f1 = reduce(np.dot, (orb.T.conj(), fock[k], orb))
                e, c = scipy.linalg.eigh(f1)
                mo1[:,idx] = np.dot(orb, c)
                mo_e[idx] = e
        if getattr(fock, 'focka', None) is not None:
            fa, fb = fock.focka[k], fock.fockb[k]
            mo_ea = np.einsum('pi,pi->i', mo1.conj(), fa.dot(mo1)).real
            mo_eb = np.einsum('pi,pi->i', mo1.conj(), fb.dot(mo1)).real
            mo_e = lib.tag_array(mo_e, mo_ea=mo_ea, mo_eb=mo_eb)
        mo_coeff.append(mo1)
        mo_energy.append(mo_e)
    return mo_energy, mo_coeff

init_guess_by_chkfile = kuhf.init_guess_by_chkfile


class KROHF(khf.KRHF, pbcrohf.ROHF):
    '''UHF class with k-point sampling.
    '''
    conv_tol = getattr(__config__, 'pbc_scf_KSCF_conv_tol', 1e-7)
    conv_tol_grad = getattr(__config__, 'pbc_scf_KSCF_conv_tol_grad', None)
    direct_scf = getattr(__config__, 'pbc_scf_SCF_direct_scf', True)

    def __init__(self, cell, kpts=np.zeros((1,3)),
                 exxdiv=getattr(__config__, 'pbc_scf_SCF_exxdiv', 'ewald')):
        khf.KSCF.__init__(self, cell, kpts, exxdiv)
        self.nelec = None

    @property
    def nelec(self):
        if self._nelec is not None:
            return self._nelec
        else:
            cell = self.cell
            nkpts = len(self.kpts)
            ne = cell.tot_electrons(nkpts)
            nalpha = (ne + cell.spin) // 2
            nbeta = nalpha - cell.spin
            if nalpha + nbeta != ne:
                raise RuntimeError('Electron number %d and spin %d are not consistent\n'
                                   'Note cell.spin = 2S = Nalpha - Nbeta, not 2S+1' %
                                   (ne, cell.spin))
            return nalpha, nbeta
    @nelec.setter
    def nelec(self, x):
        self._nelec = x

    def dump_flags(self, verbose=None):
        khf.KSCF.dump_flags(self, verbose)
        logger.info(self, 'number of electrons per unit cell  '
                    'alpha = %d beta = %d', *self.nelec)
        return self

#?    def get_init_guess(self, cell=None, key='minao'):
#?        dm_kpts = khf.KSCF.get_init_guess(self, cell, key)
#?        if dm_kpts.ndim != 4:  # The KRHF initial guess
#?            # dm_kpts shape should be (spin, nkpts, nao, nao)
#?            dm_kpts = lib.asarray([dm_kpts*.5,]*2)
#?        return dm_kpts
#?
    def get_init_guess(self, cell=None, key='minao'):
        if cell is None:
            cell = self.cell
        dm_kpts = None
        key = key.lower()
        if key == '1e' or key == 'hcore':
            dm_kpts = self.init_guess_by_1e(cell)
        elif getattr(cell, 'natm', 0) == 0:
            logger.info(self, 'No atom found in cell. Use 1e initial guess')
            dm_kpts = self.init_guess_by_1e(cell)
        elif key == 'atom':
            dm = self.init_guess_by_atom(cell)
        elif key[:3] == 'chk':
            try:
                dm_kpts = self.from_chk()
            except (IOError, KeyError):
                logger.warn(self, 'Fail to read %s. Use MINAO initial guess',
                            self.chkfile)
                dm = self.init_guess_by_minao(cell)
        else:
            dm = self.init_guess_by_minao(cell)

        if dm_kpts is None:
            nkpts = len(self.kpts)
            # dm[spin,nao,nao] at gamma point -> dm_kpts[spin,nkpts,nao,nao]
            dm_kpts = np.repeat(dm[:,None,:,:], nkpts, axis=1)

        ne = np.einsum('xkij,kji->', dm_kpts, self.get_ovlp(cell)).real
        # FIXME: consider the fractional num_electron or not? This maybe
        # relates to the charged system.
        nkpts = len(self.kpts)
        nelec = float(sum(self.nelec))
        if np.any(abs(ne - nelec) > 1e-7*nkpts):
            logger.debug(self, 'Big error detected in the electron number '
                         'of initial guess density matrix (Ne/cell = %g)!\n'
                         '  This can cause huge error in Fock matrix and '
                         'lead to instability in SCF for low-dimensional '
                         'systems.\n  DM is normalized wrt the number '
                         'of electrons %g', ne/nkpts, nelec/nkpts)
            dm_kpts *= nelec / ne
        return dm_kpts

    init_guess_by_minao  = pbcrohf.ROHF.init_guess_by_minao
    init_guess_by_atom   = pbcrohf.ROHF.init_guess_by_atom
    init_guess_by_huckel = pbcrohf.ROHF.init_guess_by_huckel

    get_rho = get_rho

    get_fock = get_fock
    get_occ = get_occ
    energy_elec = energy_elec

    def get_veff(self, cell=None, dm_kpts=None, dm_last=0, vhf_last=0, hermi=1,
                 kpts=None, kpts_band=None):
        if dm_kpts is None:
            dm_kpts = self.make_rdm1()
        if getattr(dm_kpts, 'mo_coeff', None) is not None:
            mo_coeff = dm_kpts.mo_coeff
            mo_occ_a = [(x > 0).astype(np.double) for x in dm_kpts.mo_occ]
            mo_occ_b = [(x ==2).astype(np.double) for x in dm_kpts.mo_occ]
            dm_kpts = lib.tag_array(dm_kpts, mo_coeff=(mo_coeff,mo_coeff),
                                    mo_occ=(mo_occ_a,mo_occ_b))
        if self.rsjk and self.direct_scf:
            ddm = dm_kpts - dm_last
            vj, vk = self.get_jk(cell, ddm, hermi, kpts, kpts_band)
            vhf = vj[0] + vj[1] - vk
            vhf += vhf_last
        else:
            vj, vk = self.get_jk(cell, dm_kpts, hermi, kpts, kpts_band)
            vhf = vj[0] + vj[1] - vk
        return vhf

    def get_grad(self, mo_coeff_kpts, mo_occ_kpts, fock=None):
        if fock is None:
            dm1 = self.make_rdm1(mo_coeff_kpts, mo_occ_kpts)
            fock = self.get_hcore(self.cell, self.kpts) + self.get_veff(self.cell, dm1)

        if getattr(fock, 'focka', None) is not None:
            focka = fock.focka
            fockb = fock.fockb
        elif getattr(fock, 'ndim', None) == 4:
            focka, fockb = fock
        else:
            focka = fockb = fock

        def grad(k):
            mo_occ = mo_occ_kpts[k]
            mo_coeff = mo_coeff_kpts[k]
            return pbcrohf.get_grad(mo_coeff, mo_occ, (focka[k], fockb[k]))

        nkpts = len(self.kpts)
        grad_kpts = np.hstack([grad(k) for k in range(nkpts)])
        return grad_kpts

    def eig(self, fock, s):
        e, c = khf.KSCF.eig(self, fock, s)
        if getattr(fock, 'focka', None) is not None:
            for k, mo in enumerate(c):
                fa, fb = fock.focka[k], fock.fockb[k]
                mo_ea = np.einsum('pi,pi->i', mo.conj(), fa.dot(mo)).real
                mo_eb = np.einsum('pi,pi->i', mo.conj(), fb.dot(mo)).real
                e[k] = lib.tag_array(e[k], mo_ea=mo_ea, mo_eb=mo_eb)
        return e, c

    def make_rdm1(self, mo_coeff_kpts=None, mo_occ_kpts=None, **kwargs):
        if mo_coeff_kpts is None: mo_coeff_kpts = self.mo_coeff
        if mo_occ_kpts is None: mo_occ_kpts = self.mo_occ
        return make_rdm1(mo_coeff_kpts, mo_occ_kpts, **kwargs)

    def init_guess_by_chkfile(self, chk=None, project=True, kpts=None):
        if chk is None: chk = self.chkfile
        if kpts is None: kpts = self.kpts
        return init_guess_by_chkfile(self.cell, chk, project, kpts)


    def analyze(self, verbose=None, with_meta_lowdin=WITH_META_LOWDIN,
                **kwargs):
        if verbose is None: verbose = self.verbose
        return khf.analyze(self, verbose, with_meta_lowdin, **kwargs)

    def mulliken_meta(self, cell=None, dm=None, verbose=logger.DEBUG,
                      pre_orth_method=PRE_ORTH_METHOD, s=None):
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        if s is None: s = self.get_ovlp(cell)
        return mulliken_meta(cell, dm, s=s, verbose=verbose,
                             pre_orth_method=pre_orth_method)

    @lib.with_doc(dip_moment.__doc__)
    def dip_moment(self, cell=None, dm=None, unit='Debye', verbose=logger.NOTE,
                   **kwargs):
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        rho = kwargs.pop('rho', None)
        if rho is None:
            rho = self.get_rho(dm)
        return dip_moment(cell, dm, unit, verbose, rho=rho, kpts=self.kpts, **kwargs)

    spin_square = pbcrohf.ROHF.spin_square

    canonicalize = canonicalize

    def stability(self,
                  internal=getattr(__config__, 'pbc_scf_KSCF_stability_internal', True),
                  external=getattr(__config__, 'pbc_scf_KSCF_stability_external', False),
                  verbose=None):
        raise NotImplementedError

    def convert_from_(self, mf):
        '''Convert given mean-field object to KUHF'''
        addons.convert_to_rhf(mf, self)
        return self

del(WITH_META_LOWDIN, PRE_ORTH_METHOD)


if __name__ == '__main__':
    from pyscf.pbc import gto
    cell = gto.Cell()
    cell.atom = '''
    He 0 0 1
    He 1 0 1
    '''
    cell.basis = '321g'
    cell.a = np.eye(3) * 3
    cell.mesh = [11] * 3
    cell.verbose = 5
    cell.spin = 2
    cell.build()
    mf = KROHF(cell, [2,1,1])
    mf.kernel()
    mf.analyze()



# pyscf/pbc/scf/kuhf.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Hartree-Fock for periodic systems with k-point sampling

See Also:
    hf.py : Hartree-Fock for periodic systems at a single k-point
'''

from functools import reduce
import numpy as np
import scipy.linalg
from pyscf.scf import hf as mol_hf
from pyscf.scf import uhf as mol_uhf
from pyscf.pbc.scf import khf
from pyscf.pbc.scf import uhf as pbcuhf
from pyscf import lib
from pyscf.lib import logger
from pyscf.pbc.scf import addons
from pyscf.pbc.scf import chkfile  # noqa
from pyscf import __config__

WITH_META_LOWDIN = getattr(__config__, 'pbc_scf_analyze_with_meta_lowdin', True)
PRE_ORTH_METHOD = getattr(__config__, 'pbc_scf_analyze_pre_orth_method', 'ANO')
CHECK_COULOMB_IMAG = getattr(__config__, 'pbc_scf_check_coulomb_imag', True)


canonical_occ = canonical_occ_ = addons.canonical_occ_


def make_rdm1(mo_coeff_kpts, mo_occ_kpts, **kwargs):
    '''Alpha and beta spin one particle density matrices for all k-points.

    Returns:
        dm_kpts : (2, nkpts, nao, nao) ndarray
    '''
    nkpts = len(mo_occ_kpts[0])
    nao, nmo = mo_coeff_kpts[0][0].shape
    def make_dm(mos, occs):
        return [np.dot(mos[k]*occs[k], mos[k].T.conj()) for k in range(nkpts)]
    dm_kpts =(make_dm(mo_coeff_kpts[0], mo_occ_kpts[0]) +
              make_dm(mo_coeff_kpts[1], mo_occ_kpts[1]))
    return lib.asarray(dm_kpts).reshape(2,nkpts,nao,nao)

def get_fock(mf, h1e=None, s1e=None, vhf=None, dm=None, cycle=-1, diis=None,
             diis_start_cycle=None, level_shift_factor=None, damp_factor=None):
    h1e_kpts, s_kpts, vhf_kpts, dm_kpts = h1e, s1e, vhf, dm
    if h1e_kpts is None: h1e_kpts = mf.get_hcore()
    if vhf_kpts is None: vhf_kpts = mf.get_veff(mf.cell, dm_kpts)
    f_kpts = h1e_kpts + vhf_kpts
    if cycle < 0 and diis is None:  # Not inside the SCF iteration
        return f_kpts

    if diis_start_cycle is None:
        diis_start_cycle = mf.diis_start_cycle
    if level_shift_factor is None:
        level_shift_factor = mf.level_shift
    if damp_factor is None:
        damp_factor = mf.damp
    if s_kpts is None: s_kpts = mf.get_ovlp()
    if dm_kpts is None: dm_kpts = mf.make_rdm1()

    if isinstance(level_shift_factor, (tuple, list, np.ndarray)):
        shifta, shiftb = level_shift_factor
    else:
        shifta = shiftb = level_shift_factor
    if isinstance(damp_factor, (tuple, list, np.ndarray)):
        dampa, dampb = damp_factor
    else:
        dampa = dampb = damp_factor

    if 0 <= cycle < diis_start_cycle-1 and abs(dampa)+abs(dampb) > 1e-4:
        f_a = []
        f_b = []
        for k, s1e in enumerate(s_kpts):
            f_a.append(mol_hf.damping(s1e, dm_kpts[0][k], f_kpts[0][k], dampa))
            f_b.append(mol_hf.damping(s1e, dm_kpts[1][k], f_kpts[1][k], dampb))
        f_kpts = [f_a, f_b]
    if diis and cycle >= diis_start_cycle:
        f_kpts = diis.update(s_kpts, dm_kpts, f_kpts, mf, h1e_kpts, vhf_kpts)
    if abs(level_shift_factor) > 1e-4:
        f_kpts =([mol_hf.level_shift(s, dm_kpts[0,k], f_kpts[0,k], shifta)
                  for k, s in enumerate(s_kpts)],
                 [mol_hf.level_shift(s, dm_kpts[1,k], f_kpts[1,k], shiftb)
                  for k, s in enumerate(s_kpts)])
    return lib.asarray(f_kpts)

def get_fermi(mf, mo_energy_kpts=None, mo_occ_kpts=None):
    '''A pair of Fermi level for spin-up and spin-down orbitals
    '''
    if mo_energy_kpts is None: mo_energy_kpts = mf.mo_energy
    if mo_occ_kpts is None: mo_occ_kpts = mf.mo_occ

    # mo_energy_kpts and mo_occ_kpts are k-point UHF quantities
    assert(mo_energy_kpts[0][0].ndim == 1)
    assert(mo_occ_kpts[0][0].ndim == 1)

    nocca = sum(mo_occ.sum() for mo_occ in mo_occ_kpts[0])
    noccb = sum(mo_occ.sum() for mo_occ in mo_occ_kpts[1])
    # nocc may not be perfect integer when smearing is enabled
    nocca = int(nocca.round(3))
    noccb = int(noccb.round(3))

    fermi_a = np.sort(np.hstack(mo_energy_kpts[0]))[nocca-1]
    fermi_b = np.sort(np.hstack(mo_energy_kpts[1]))[noccb-1]

    for k, mo_e in enumerate(mo_energy_kpts[0]):
        mo_occ = mo_occ_kpts[0][k]
        if mo_occ[mo_e > fermi_a].sum() > 0.5:
            logger.warn(mf, 'Alpha occupied band above Fermi level: \n'
                        'k=%d, mo_e=%s, mo_occ=%s', k, mo_e, mo_occ)
    for k, mo_e in enumerate(mo_energy_kpts[1]):
        mo_occ = mo_occ_kpts[1][k]
        if mo_occ[mo_e > fermi_b].sum() > 0.5:
            logger.warn(mf, 'Beta occupied band above Fermi level: \n'
                        'k=%d, mo_e=%s, mo_occ=%s', k, mo_e, mo_occ)
    return (fermi_a, fermi_b)

def get_occ(mf, mo_energy_kpts=None, mo_coeff_kpts=None):
    '''Label the occupancies for each orbital for sampled k-points.

    This is a k-point version of scf.hf.SCF.get_occ
    '''

    if mo_energy_kpts is None: mo_energy_kpts = mf.mo_energy

    nocc_a, nocc_b = mf.nelec
    mo_energy = np.sort(np.hstack(mo_energy_kpts[0]))
    fermi_a = mo_energy[nocc_a-1]
    mo_occ_kpts = [[], []]
    for mo_e in mo_energy_kpts[0]:
        mo_occ_kpts[0].append((mo_e <= fermi_a).astype(np.double))
    if nocc_a < len(mo_energy):
        logger.info(mf, 'alpha HOMO = %.12g  LUMO = %.12g', fermi_a, mo_energy[nocc_a])
    else:
        logger.info(mf, 'alpha HOMO = %.12g  (no LUMO because of small basis) ', fermi_a)

    if nocc_b > 0:
        mo_energy = np.sort(np.hstack(mo_energy_kpts[1]))
        fermi_b = mo_energy[nocc_b-1]
        for mo_e in mo_energy_kpts[1]:
            mo_occ_kpts[1].append((mo_e <= fermi_b).astype(np.double))
        if nocc_b < len(mo_energy):
            logger.info(mf, 'beta HOMO = %.12g  LUMO = %.12g', fermi_b, mo_energy[nocc_b])
        else:
            logger.info(mf, 'beta HOMO = %.12g  (no LUMO because of small basis) ', fermi_b)
    else:
        for mo_e in mo_energy_kpts[1]:
            mo_occ_kpts[1].append(np.zeros_like(mo_e))

    if mf.verbose >= logger.DEBUG:
        np.set_printoptions(threshold=len(mo_energy))
        logger.debug(mf, '     k-point                  alpha mo_energy')
        for k,kpt in enumerate(mf.cell.get_scaled_kpts(mf.kpts)):
            if (np.count_nonzero(mo_occ_kpts[0][k]) > 0 and
                np.count_nonzero(mo_occ_kpts[0][k] == 0) > 0):
                logger.debug(mf, '  %2d (%6.3f %6.3f %6.3f)   %s %s',
                             k, kpt[0], kpt[1], kpt[2],
                             mo_energy_kpts[0][k][mo_occ_kpts[0][k]> 0],
                             mo_energy_kpts[0][k][mo_occ_kpts[0][k]==0])
            else:
                logger.debug(mf, '  %2d (%6.3f %6.3f %6.3f)   %s',
                             k, kpt[0], kpt[1], kpt[2], mo_energy_kpts[0][k])
        logger.debug(mf, '     k-point                  beta  mo_energy')
        for k,kpt in enumerate(mf.cell.get_scaled_kpts(mf.kpts)):
            if (np.count_nonzero(mo_occ_kpts[1][k]) > 0 and
                np.count_nonzero(mo_occ_kpts[1][k] == 0) > 0):
                logger.debug(mf, '  %2d (%6.3f %6.3f %6.3f)   %s %s',
                             k, kpt[0], kpt[1], kpt[2],
                             mo_energy_kpts[1][k][mo_occ_kpts[1][k]> 0],
                             mo_energy_kpts[1][k][mo_occ_kpts[1][k]==0])
            else:
                logger.debug(mf, '  %2d (%6.3f %6.3f %6.3f)   %s',
                             k, kpt[0], kpt[1], kpt[2], mo_energy_kpts[1][k])
        np.set_printoptions(threshold=1000)

    return mo_occ_kpts


def energy_elec(mf, dm_kpts=None, h1e_kpts=None, vhf_kpts=None):
    '''Following pyscf.scf.hf.energy_elec()
    '''
    if dm_kpts is None: dm_kpts = mf.make_rdm1()
    if h1e_kpts is None: h1e_kpts = mf.get_hcore()
    if vhf_kpts is None: vhf_kpts = mf.get_veff(mf.cell, dm_kpts)

    nkpts = len(h1e_kpts)
    e1 = 1./nkpts * np.einsum('kij,kji', dm_kpts[0], h1e_kpts)
    e1+= 1./nkpts * np.einsum('kij,kji', dm_kpts[1], h1e_kpts)
    e_coul = 1./nkpts * np.einsum('kij,kji', dm_kpts[0], vhf_kpts[0]) * 0.5
    e_coul+= 1./nkpts * np.einsum('kij,kji', dm_kpts[1], vhf_kpts[1]) * 0.5
    mf.scf_summary['e1'] = e1.real
    mf.scf_summary['e2'] = e_coul.real
    logger.debug(mf, 'E1 = %s  E_coul = %s', e1, e_coul)
    if CHECK_COULOMB_IMAG and abs(e_coul.imag > mf.cell.precision*10):
        logger.warn(mf, "Coulomb energy has imaginary part %s. "
                    "Coulomb integrals (e-e, e-N) may not converge !",
                    e_coul.imag)
    return (e1+e_coul).real, e_coul.real


def mulliken_meta(cell, dm_ao_kpts, verbose=logger.DEBUG,
                  pre_orth_method=PRE_ORTH_METHOD, s=None):
    '''A modified Mulliken population analysis, based on meta-Lowdin AOs.

    Note this function only computes the Mulliken population for the gamma
    point density matrix.
    '''
    from pyscf.lo import orth
    if s is None:
        s = khf.get_ovlp(cell)
    log = logger.new_logger(cell, verbose)
    log.note('Analyze output for *gamma point*.')
    log.info('    To include the contributions from k-points, transform to a '
             'supercell then run the population analysis on the supercell\n'
             '        from pyscf.pbc.tools import k2gamma\n'
             '        k2gamma.k2gamma(mf).mulliken_meta()')
    log.note("KUHF mulliken_meta")
    dm_ao_gamma = dm_ao_kpts[:,0,:,:].real
    s_gamma = s[0,:,:].real
    orth_coeff = orth.orth_ao(cell, 'meta_lowdin', pre_orth_method, s=s_gamma)
    c_inv = np.dot(orth_coeff.T, s_gamma)
    dm_a = reduce(np.dot, (c_inv, dm_ao_gamma[0], c_inv.T.conj()))
    dm_b = reduce(np.dot, (c_inv, dm_ao_gamma[1], c_inv.T.conj()))

    log.note(' ** Mulliken pop alpha/beta on meta-lowdin orthogonal AOs **')
    return mol_uhf.mulliken_pop(cell, (dm_a,dm_b), np.eye(orth_coeff.shape[0]), log)


def canonicalize(mf, mo_coeff_kpts, mo_occ_kpts, fock=None):
    '''Canonicalization diagonalizes the UHF Fock matrix within occupied,
    virtual subspaces separatedly (without change occupancy).
    '''
    if fock is None:
        dm = mf.make_rdm1(mo_coeff_kpts, mo_occ_kpts)
        fock = mf.get_fock(dm=dm)

    def eig_(fock, mo_coeff, idx, es, cs):
        if np.count_nonzero(idx) > 0:
            orb = mo_coeff[:,idx]
            f1 = reduce(np.dot, (orb.T.conj(), fock, orb))
            e, c = scipy.linalg.eigh(f1)
            es[idx] = e
            cs[:,idx] = np.dot(orb, c)

    mo_coeff = [[], []]
    mo_energy = [[], []]
    for k, mo in enumerate(mo_coeff_kpts[0]):
        mo1 = np.empty_like(mo)
        mo_e = np.empty_like(mo_occ_kpts[0][k])
        occidxa = mo_occ_kpts[0][k] == 1
        viridxa = ~occidxa
        eig_(fock[0][k], mo, occidxa, mo_e, mo1)
        eig_(fock[0][k], mo, viridxa, mo_e, mo1)
        mo_coeff[0].append(mo1)
        mo_energy[0].append(mo_e)
    for k, mo in enumerate(mo_coeff_kpts[1]):
        mo1 = np.empty_like(mo)
        mo_e = np.empty_like(mo_occ_kpts[1][k])
        occidxb = mo_occ_kpts[1][k] == 1
        viridxb = ~occidxb
        eig_(fock[1][k], mo, occidxb, mo_e, mo1)
        eig_(fock[1][k], mo, viridxb, mo_e, mo1)
        mo_coeff[1].append(mo1)
        mo_energy[1].append(mo_e)
    return mo_energy, mo_coeff

def init_guess_by_chkfile(cell, chkfile_name, project=None, kpts=None):
    '''Read the KHF results from checkpoint file, then project it to the
    basis defined by ``cell``

    Returns:
        Density matrix, 3D ndarray
    '''
    from pyscf import gto
    chk_cell, scf_rec = chkfile.load_scf(chkfile_name)
    if project is None:
        project = not gto.same_basis_set(chk_cell, cell)

    if kpts is None:
        kpts = scf_rec['kpts']

    if 'kpt' in scf_rec:
        chk_kpts = scf_rec['kpt'].reshape(-1,3)
    elif 'kpts' in scf_rec:
        chk_kpts = scf_rec['kpts']
    else:
        chk_kpts = np.zeros((1,3))

    mo = scf_rec['mo_coeff']
    mo_occ = scf_rec['mo_occ']
    if 'kpts' not in scf_rec:  # gamma point or single k-point
        if mo.ndim == 2:
            mo = np.expand_dims(mo, axis=0)
            mo_occ = np.expand_dims(mo_occ, axis=0)
        else:  # UHF
            mo = [np.expand_dims(mo[0], axis=0),
                  np.expand_dims(mo[1], axis=0)]
            mo_occ = [np.expand_dims(mo_occ[0], axis=0),
                      np.expand_dims(mo_occ[1], axis=0)]

    if project:
        s = cell.pbc_intor('int1e_ovlp', kpts=kpts)
    def fproj(mo, kpts):
        if project:
            mo = addons.project_mo_nr2nr(chk_cell, mo, cell, kpts)
            for k, c in enumerate(mo):
                norm = np.einsum('pi,pi->i', c.conj(), s[k].dot(c))
                mo[k] /= np.sqrt(norm)
        return mo

    if kpts.shape == chk_kpts.shape and np.allclose(kpts, chk_kpts):
        def makedm(mos, occs):
            moa, mob = mos
            mos =([fproj(mo, None) for mo in moa],
                  [fproj(mo, None) for mo in mob])
            return make_rdm1(mos, occs)
    else:
        def makedm(mos, occs):
            where = [np.argmin(lib.norm(chk_kpts-kpt, axis=1)) for kpt in kpts]
            moa, mob = mos
            occa, occb = occs
            dkpts = [chk_kpts[w]-kpts[i] for i,w in enumerate(where)]
            mos = (fproj([moa[w] for w in where], dkpts),
                   fproj([mob[w] for w in where], dkpts))
            occs = ([occa[i] for i in where], [occb[i] for i in where])
            return make_rdm1(mos, occs)

    if getattr(mo[0], 'ndim', None) == 2:  # KRHF
        mo_occa = [(occ>1e-8).astype(np.double) for occ in mo_occ]
        mo_occb = [occ-mo_occa[k] for k,occ in enumerate(mo_occ)]
        dm = makedm((mo, mo), (mo_occa, mo_occb))
    else:  # KUHF
        dm = makedm(mo, mo_occ)

    # Real DM for gamma point
    if np.allclose(kpts, 0):
        dm = dm.real
    return dm


def dip_moment(cell, dm_kpts, unit='Debye', verbose=logger.NOTE,
               grids=None, rho=None, kpts=np.zeros((1,3))):
    ''' Dipole moment in the unit cell.

    Args:
         cell : an instance of :class:`Cell`

         dm_kpts (two lists of ndarrays) : KUHF density matrices of k-points

    Return:
        A list: the dipole moment on x, y and z components
    '''
    dm_kpts = dm_kpts[0] + dm_kpts[1]
    return khf.dip_moment(cell, dm_kpts, unit, verbose, grids, rho, kpts)

get_rho = khf.get_rho


class KUHF(khf.KSCF, pbcuhf.UHF):
    '''UHF class with k-point sampling.
    '''
    conv_tol = getattr(__config__, 'pbc_scf_KSCF_conv_tol', 1e-7)
    conv_tol_grad = getattr(__config__, 'pbc_scf_KSCF_conv_tol_grad', None)
    direct_scf = getattr(__config__, 'pbc_scf_SCF_direct_scf', True)

    def __init__(self, cell, kpts=np.zeros((1,3)),
                 exxdiv=getattr(__config__, 'pbc_scf_SCF_exxdiv', 'ewald')):
        khf.KSCF.__init__(self, cell, kpts, exxdiv)
        self.nelec = None
        self.init_guess_breaksym = None
        self._keys = self._keys.union(["init_guess_breaksym"])

    @property
    def nelec(self):
        if self._nelec is not None:
            return self._nelec
        else:
            cell = self.cell
            nkpts = len(self.kpts)
            ne = cell.tot_electrons(nkpts)
            nalpha = (ne + cell.spin) // 2
            nbeta = nalpha - cell.spin
            if nalpha + nbeta != ne:
                raise RuntimeError('Electron number %d and spin %d are not consistent\n'
                                   'Note cell.spin = 2S = Nalpha - Nbeta, not 2S+1' %
                                   (ne, cell.spin))
            return nalpha, nbeta
    @nelec.setter
    def nelec(self, x):
        self._nelec = x

    def dump_flags(self, verbose=None):
        khf.KSCF.dump_flags(self, verbose)
        logger.info(self, 'number of electrons per unit cell  '
                    'alpha = %d beta = %d', *self.nelec)
        return self

    def get_init_guess(self, cell=None, key='minao'):
        if cell is None:
            cell = self.cell
        dm_kpts = None
        key = key.lower()
        if key == '1e' or key == 'hcore':
            dm_kpts = self.init_guess_by_1e(cell)
        elif getattr(cell, 'natm', 0) == 0:
            logger.info(self, 'No atom found in cell. Use 1e initial guess')
            dm_kpts = self.init_guess_by_1e(cell)
        elif key == 'atom':
            dm = self.init_guess_by_atom(cell)
        elif key[:3] == 'chk':
            try:
                dm_kpts = self.from_chk()
            except (IOError, KeyError):
                logger.warn(self, 'Fail to read %s. Use MINAO initial guess',
                            self.chkfile)
                dm = self.init_guess_by_minao(cell)
        else:
            dm = self.init_guess_by_minao(cell)

        if dm_kpts is None:
            nkpts = len(self.kpts)
            # dm[spin,nao,nao] at gamma point -> dm_kpts[spin,nkpts,nao,nao]
            dm_kpts = np.repeat(dm[:,None,:,:], nkpts, axis=1)
            dm_kpts[0,:] *= 1.01
            dm_kpts[1,:] *= 0.99  # To slightly break spin symmetry
            assert dm_kpts.shape[0]==2

        ne = np.einsum('xkij,kji->x', dm_kpts, self.get_ovlp(cell)).real
        # FIXME: consider the fractional num_electron or not? This maybe
        # relates to the charged system.
        nkpts = len(self.kpts)
        nelec = np.asarray(self.nelec)
        if np.any(abs(ne - nelec) > 1e-7*nkpts):
            logger.debug(self, 'Big error detected in the electron number '
                         'of initial guess density matrix (Ne/cell = %g)!\n'
                         '  This can cause huge error in Fock matrix and '
                         'lead to instability in SCF for low-dimensional '
                         'systems.\n  DM is normalized wrt the number '
                         'of electrons %s', ne.mean()/nkpts, nelec/nkpts)
            dm_kpts *= (nelec / ne).reshape(2,-1,1,1)
        return dm_kpts

    get_fock = get_fock
    get_fermi = get_fermi
    get_occ = get_occ
    energy_elec = energy_elec

    get_rho = get_rho

    def get_veff(self, cell=None, dm_kpts=None, dm_last=0, vhf_last=0, hermi=1,
                 kpts=None, kpts_band=None):
        if dm_kpts is None:
            dm_kpts = self.make_rdm1()
        if self.rsjk and self.direct_scf:
            ddm = dm_kpts - dm_last
            vj, vk = self.get_jk(cell, ddm, hermi, kpts, kpts_band)
            vhf = vj[0] + vj[1] - vk
            vhf += vhf_last
        else:
            vj, vk = self.get_jk(cell, dm_kpts, hermi, kpts, kpts_band)
            vhf = vj[0] + vj[1] - vk
        return vhf


    def analyze(self, verbose=None, with_meta_lowdin=WITH_META_LOWDIN,
                **kwargs):
        if verbose is None: verbose = self.verbose
        return khf.analyze(self, verbose, with_meta_lowdin, **kwargs)


    def get_grad(self, mo_coeff_kpts, mo_occ_kpts, fock=None):
        if fock is None:
            dm1 = self.make_rdm1(mo_coeff_kpts, mo_occ_kpts)
            fock = self.get_hcore(self.cell, self.kpts) + self.get_veff(self.cell, dm1)

        def grad(mo, mo_occ, fock):
            occidx = mo_occ > 0
            viridx = ~occidx
            g = reduce(np.dot, (mo[:,viridx].T.conj(), fock, mo[:,occidx]))
            return g.ravel()

        nkpts = len(self.kpts)
        grad_kpts = [grad(mo_coeff_kpts[0][k], mo_occ_kpts[0][k], fock[0][k])
                     for k in range(nkpts)]
        grad_kpts+= [grad(mo_coeff_kpts[1][k], mo_occ_kpts[1][k], fock[1][k])
                     for k in range(nkpts)]
        return np.hstack(grad_kpts)

    def eig(self, h_kpts, s_kpts):
        e_a, c_a = khf.KSCF.eig(self, h_kpts[0], s_kpts)
        e_b, c_b = khf.KSCF.eig(self, h_kpts[1], s_kpts)
        return (e_a,e_b), (c_a,c_b)

    def make_rdm1(self, mo_coeff_kpts=None, mo_occ_kpts=None, **kwargs):
        if mo_coeff_kpts is None: mo_coeff_kpts = self.mo_coeff
        if mo_occ_kpts is None: mo_occ_kpts = self.mo_occ
        return make_rdm1(mo_coeff_kpts, mo_occ_kpts, **kwargs)

    def get_bands(self, kpts_band, cell=None, dm_kpts=None, kpts=None):
        '''Get energy bands at the given (arbitrary) 'band' k-points.

        Returns:
            mo_energy : (nmo,) ndarray or a list of (nmo,) ndarray
                Bands energies E_n(k)
            mo_coeff : (nao, nmo) ndarray or a list of (nao,nmo) ndarray
                Band orbitals psi_n(k)
        '''
        if cell is None: cell = self.cell
        if dm_kpts is None: dm_kpts = self.make_rdm1()
        if kpts is None: kpts = self.kpts

        kpts_band = np.asarray(kpts_band)
        single_kpt_band = (kpts_band.ndim == 1)
        kpts_band = kpts_band.reshape(-1,3)

        fock = self.get_hcore(cell, kpts_band)
        fock = fock + self.get_veff(cell, dm_kpts, kpts=kpts, kpts_band=kpts_band)
        s1e = self.get_ovlp(cell, kpts_band)
        (e_a,e_b), (c_a,c_b) = self.eig(fock, s1e)
        if single_kpt_band:
            e_a = e_a[0]
            e_b = e_b[0]
            c_a = c_a[0]
            c_b = c_b[0]
        return (e_a,e_b), (c_a,c_b)

    def init_guess_by_chkfile(self, chk=None, project=True, kpts=None):
        if chk is None: chk = self.chkfile
        if kpts is None: kpts = self.kpts
        return init_guess_by_chkfile(self.cell, chk, project, kpts)

    init_guess_by_1e     = pbcuhf.UHF.init_guess_by_1e
    init_guess_by_minao  = pbcuhf.UHF.init_guess_by_minao
    init_guess_by_atom   = pbcuhf.UHF.init_guess_by_atom
    init_guess_by_huckel = pbcuhf.UHF.init_guess_by_huckel

    @lib.with_doc(mulliken_meta.__doc__)
    def mulliken_meta(self, cell=None, dm=None, verbose=logger.DEBUG,
                      pre_orth_method=PRE_ORTH_METHOD, s=None):
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        if s is None: s = self.get_ovlp(cell)
        return mulliken_meta(cell, dm, s=s, verbose=verbose,
                             pre_orth_method=pre_orth_method)

    def mulliken_pop(self):
        raise NotImplementedError

    @lib.with_doc(dip_moment.__doc__)
    def dip_moment(self, cell=None, dm=None, unit='Debye', verbose=logger.NOTE,
                   **kwargs):
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        rho = kwargs.pop('rho', None)
        if rho is None:
            rho = self.get_rho(dm)
        return dip_moment(cell, dm, unit, verbose, rho=rho, kpts=self.kpts, **kwargs)

    @lib.with_doc(mol_uhf.spin_square.__doc__)
    def spin_square(self, mo_coeff=None, s=None):
        '''Treating the k-point sampling wfn as a giant Slater determinant,
        the spin_square value is the <S^2> of the giant determinant.
        '''
        nkpts = len(self.kpts)
        if mo_coeff is None:
            mo_a = [self.mo_coeff[0][k][:,self.mo_occ[0][k]>0] for k in range(nkpts)]
            mo_b = [self.mo_coeff[1][k][:,self.mo_occ[1][k]>0] for k in range(nkpts)]
        else:
            mo_a, mo_b = mo_coeff
        if s is None:
            s = self.get_ovlp()

        nelec_a = sum([mo_a[k].shape[1] for k in range(nkpts)])
        nelec_b = sum([mo_b[k].shape[1] for k in range(nkpts)])
        ssxy = (nelec_a + nelec_b) * .5
        for k in range(nkpts):
            sij = reduce(np.dot, (mo_a[k].T.conj(), s[k], mo_b[k]))
            ssxy -= np.einsum('ij,ij->', sij.conj(), sij).real
        ssz = (nelec_b-nelec_a)**2 * .25
        ss = ssxy + ssz
        s = np.sqrt(ss+.25) - .5
        return ss, s*2+1

    canonicalize = canonicalize

    def stability(self,
                  internal=getattr(__config__, 'pbc_scf_KSCF_stability_internal', True),
                  external=getattr(__config__, 'pbc_scf_KSCF_stability_external', False),
                  verbose=None):
        from pyscf.pbc.scf.stability import uhf_stability
        return uhf_stability(self, internal, external, verbose)

    def convert_from_(self, mf):
        '''Convert given mean-field object to KUHF'''
        addons.convert_to_uhf(mf, self)
        return self

    def nuc_grad_method(self):
        from pyscf.pbc.grad import kuhf
        return kuhf.Gradients(self)

del(WITH_META_LOWDIN, PRE_ORTH_METHOD)


if __name__ == '__main__':
    from pyscf.pbc import gto
    cell = gto.Cell()
    cell.atom = '''
    He 0 0 1
    He 1 0 1
    '''
    cell.basis = '321g'
    cell.a = np.eye(3) * 3
    cell.mesh = [11] * 3
    cell.verbose = 5
    cell.build()
    mf = KUHF(cell, [2,1,1])
    mf.kernel()
    mf.analyze()


# pyscf/pbc/scf/newton_ah.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Co-iterative augmented hessian second order SCF solver (CIAH-SOSCF)
'''

from functools import reduce
import numpy
from pyscf import lib
from pyscf.pbc.scf import _response_functions  # noqa

def gen_g_hop_rhf(mf, mo_coeff, mo_occ, fock_ao=None, h1e=None):
    cell = mf.cell
    nkpts = len(mo_occ)
    occidx = [numpy.where(mo_occ[k]==2)[0] for k in range(nkpts)]
    viridx = [numpy.where(mo_occ[k]==0)[0] for k in range(nkpts)]
    orbo = [mo_coeff[k][:,occidx[k]] for k in range(nkpts)]
    orbv = [mo_coeff[k][:,viridx[k]] for k in range(nkpts)]

    if fock_ao is None:
        if h1e is None: h1e = mf.get_hcore()
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        fock_ao = h1e + mf.get_veff(cell, dm0)
    fock = [reduce(numpy.dot, (mo_coeff[k].conj().T, fock_ao[k], mo_coeff[k]))
            for k in range(nkpts)]

    g = [fock[k][viridx[k][:,None],occidx[k]] * 2 for k in range(nkpts)]

    foo = [fock[k][occidx[k][:,None],occidx[k]] for k in range(nkpts)]
    fvv = [fock[k][viridx[k][:,None],viridx[k]] for k in range(nkpts)]

    h_diag = [(fvv[k].diagonal().real[:,None]-foo[k].diagonal().real) * 2
              for k in range(nkpts)]

    vind = mf.gen_response(mo_coeff, mo_occ, singlet=None, hermi=1)

    def h_op(x1):
        x1 = _unpack(x1, mo_occ)
        dm1 = []
        for k in range(nkpts):
            # *2 for double occupancy
            d1 = reduce(numpy.dot, (orbv[k], x1[k]*2, orbo[k].conj().T))
            dm1.append(d1+d1.conj().T)

        v1 = vind(lib.asarray(dm1))
        x2 = [0] * nkpts
        for k in range(nkpts):
            x2[k] = numpy.einsum('ps,sq->pq', fvv[k], x1[k]) * 2
            x2[k]-= numpy.einsum('ps,rp->rs', foo[k], x1[k]) * 2
            x2[k] += reduce(numpy.dot, (orbv[k].conj().T, v1[k], orbo[k])) * 2
        return numpy.hstack([x.ravel() for x in x2])

    return (numpy.hstack([x.ravel() for x in g]), h_op,
            numpy.hstack([x.ravel() for x in h_diag]))

def gen_g_hop_uhf(mf, mo_coeff, mo_occ, fock_ao=None, h1e=None):
    cell = mf.cell
    nkpts = len(mo_occ[0])
    occidxa = [numpy.where(mo_occ[0][k]>0)[0] for k in range(nkpts)]
    occidxb = [numpy.where(mo_occ[1][k]>0)[0] for k in range(nkpts)]
    viridxa = [numpy.where(mo_occ[0][k]==0)[0] for k in range(nkpts)]
    viridxb = [numpy.where(mo_occ[1][k]==0)[0] for k in range(nkpts)]
    moa, mob = mo_coeff
    orboa = [moa[k][:,occidxa[k]] for k in range(nkpts)]
    orbva = [moa[k][:,viridxa[k]] for k in range(nkpts)]
    orbob = [mob[k][:,occidxb[k]] for k in range(nkpts)]
    orbvb = [mob[k][:,viridxb[k]] for k in range(nkpts)]
    tot_vopair_a = sum(len(occidxa[k])*len(viridxa[k]) for k in range(nkpts))

    if fock_ao is None:
        if h1e is None: h1e = mf.get_hcore()
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        fock_ao = h1e + mf.get_veff(cell, dm0)
    focka = [reduce(numpy.dot, (moa[k].conj().T, fock_ao[0][k], moa[k]))
             for k in range(nkpts)]
    fockb = [reduce(numpy.dot, (mob[k].conj().T, fock_ao[1][k], mob[k]))
             for k in range(nkpts)]
    fooa = [focka[k][occidxa[k][:,None],occidxa[k]] for k in range(nkpts)]
    fvva = [focka[k][viridxa[k][:,None],viridxa[k]] for k in range(nkpts)]
    foob = [fockb[k][occidxb[k][:,None],occidxb[k]] for k in range(nkpts)]
    fvvb = [fockb[k][viridxb[k][:,None],viridxb[k]] for k in range(nkpts)]

    g = ([focka[k][viridxa[k][:,None],occidxa[k]] for k in range(nkpts)] +
         [fockb[k][viridxb[k][:,None],occidxb[k]] for k in range(nkpts)])

    h_diag = ([fvva[k].diagonal().real[:,None]-fooa[k].diagonal().real for k in range(nkpts)] +
              [fvvb[k].diagonal().real[:,None]-foob[k].diagonal().real for k in range(nkpts)])

    vind = mf.gen_response(mo_coeff, mo_occ, hermi=1)
    nao = orboa[0].shape[0]

    def h_op(x1):
        x1a = _unpack(x1[:tot_vopair_a], mo_occ[0])
        x1b = _unpack(x1[tot_vopair_a:], mo_occ[1])
        dm1 = numpy.empty((2,nkpts,nao,nao), dtype=x1.dtype)
        for k in range(nkpts):
            d1 = reduce(numpy.dot, (orbva[k], x1a[k], orboa[k].conj().T))
            dm1[0,k] = d1+d1.conj().T
        for k in range(nkpts):
            d1 = reduce(numpy.dot, (orbvb[k], x1b[k], orbob[k].conj().T))
            dm1[1,k] = d1+d1.conj().T
        v1 = vind(dm1)

        x2a = [0] * nkpts
        x2b = [0] * nkpts
        for k in range(nkpts):
            x2a[k] = numpy.einsum('ps,sq->pq', fvva[k], x1a[k])
            x2a[k]-= numpy.einsum('ps,rp->rs', fooa[k], x1a[k])
            x2b[k] = numpy.einsum('ps,sq->pq', fvvb[k], x1b[k])
            x2b[k]-= numpy.einsum('ps,rp->rs', foob[k], x1b[k])

            x2a[k] += reduce(numpy.dot, (orbva[k].conj().T, v1[0][k], orboa[k]))
            x2b[k] += reduce(numpy.dot, (orbvb[k].conj().T, v1[1][k], orbob[k]))
        return numpy.hstack([x.ravel() for x in (x2a+x2b)])

    return (numpy.hstack([x.ravel() for x in g]), h_op,
            numpy.hstack([x.ravel() for x in h_diag]))

def gen_g_hop_rohf(mf, mo_coeff, mo_occ, fock_ao=None, h1e=None):
    if getattr(fock_ao, 'focka', None) is None:
        dm0 = mf.make_rdm1(mo_coeff, mo_occ)
        fock_ao = mf.get_fock(h1e, dm=dm0)
    fock_ao = fock_ao.focka, fock_ao.fockb
    mo_occa = occidxa = [occ > 0 for occ in mo_occ]
    mo_occb = occidxb = [occ ==2 for occ in mo_occ]
    ug, uh_op, uh_diag = gen_g_hop_uhf(mf, (mo_coeff,)*2, (mo_occa,mo_occb),
                                       fock_ao, None)

    nkpts = len(mo_occ)
    idx_var_a = []
    idx_var_b = []
    p0 = 0
    for k in range(nkpts):
        viridxa = ~occidxa[k]
        viridxb = ~occidxb[k]
        uniq_var_a = viridxa[:,None] & occidxa[k]
        uniq_var_b = viridxb[:,None] & occidxb[k]
        uniq_ab = uniq_var_a | uniq_var_b
        nmo = len(mo_occ[k])

        n_uniq_ab = numpy.count_nonzero(uniq_ab)
        idx_array = numpy.zeros((nmo,nmo), dtype=int)
        idx_array[uniq_ab] = numpy.arange(n_uniq_ab)
        idx_var_a.append(p0 + idx_array[uniq_var_a])
        idx_var_b.append(p0 + idx_array[uniq_var_b])
        p0 += n_uniq_ab

    idx_var_a = numpy.hstack(idx_var_a)
    idx_var_b = numpy.hstack(idx_var_b)
    nvars = p0

    def sum_ab(x):
        x1 = numpy.zeros(nvars, dtype=x.dtype)
        x1[idx_var_a]  = x[:len(idx_var_a)]
        x1[idx_var_b] += x[len(idx_var_a):]
        return x1

    g = sum_ab(ug)
    h_diag = sum_ab(uh_diag)
    def h_op(x):
        # unpack ROHF rotation parameters
        x1 = numpy.hstack((x[idx_var_a], x[idx_var_b]))
        return sum_ab(uh_op(x1))

    return g, h_op, h_diag


# Be careful with the parameter ordering conventions are different for the
# _unpack function here and the one in tdscf.krhf
def _unpack(vo, mo_occ):
    z = []
    p1 = 0
    for k, occ in enumerate(mo_occ):
        no = numpy.count_nonzero(occ > 0)
        nv = occ.size - no
        p0, p1 = p1, p1 + nv * no
        z.append(vo[p0:p1].reshape(nv,no))
    return z


def newton(mf):
    from pyscf.soscf import newton_ah
    from pyscf.pbc import scf as pscf
    if not isinstance(mf, pscf.khf.KSCF):
        # Note for single k-point other than gamma point (mf.kpt != 0) mf object,
        # orbital hessian is approximated by gamma point hessian.
        return newton_ah.newton(mf)

    if isinstance(mf, newton_ah._CIAH_SOSCF):
        return mf

    if mf.__doc__ is None:
        mf_doc = ''
    else:
        mf_doc = mf.__doc__

    class SecondOrderKRHF(mf.__class__, newton_ah._CIAH_SOSCF):
        __doc__ = mf_doc + newton_ah._CIAH_SOSCF.__doc__
        __init__ = newton_ah._CIAH_SOSCF.__init__
        dump_flags = newton_ah._CIAH_SOSCF.dump_flags
        build = newton_ah._CIAH_SOSCF.build
        kernel = newton_ah._CIAH_SOSCF.kernel

        gen_g_hop = gen_g_hop_rhf

        def update_rotate_matrix(self, dx, mo_occ, u0=1, mo_coeff=None):
            p0 = 0
            u = []
            for k, occ in enumerate(mo_occ):
                occidx = occ > 0
                viridx = ~occidx
                nocc = occidx.sum()
                nvir = viridx.sum()
                nmo = nocc + nvir
                dr = numpy.zeros((nmo,nmo), dtype=dx.dtype)
                dr[viridx[:,None] & occidx] = dx[p0:p0+nocc*nvir]
                dr = dr - dr.conj().T
                p0 += nocc * nvir

                u1 = newton_ah.expmat(dr)
                if isinstance(u0, int) and u0 == 1:
                    u.append(u1)
                else:
                    u.append(numpy.dot(u0[k], u1))
            return lib.asarray(u)

        def rotate_mo(self, mo_coeff, u, log=None):
            return lib.asarray([numpy.dot(mo, u[k]) for k,mo in enumerate(mo_coeff)])

    if isinstance(mf, pscf.kuhf.KUHF):
        class SecondOrderKUHF(mf.__class__, newton_ah._CIAH_SOSCF):
            __doc__ = mf_doc + newton_ah._CIAH_SOSCF.__doc__
            __init__ = newton_ah._CIAH_SOSCF.__init__
            dump_flags = newton_ah._CIAH_SOSCF.dump_flags
            build = newton_ah._CIAH_SOSCF.build
            kernel = newton_ah._CIAH_SOSCF.kernel

            gen_g_hop = gen_g_hop_uhf

            def update_rotate_matrix(self, dx, mo_occ, u0=1, mo_coeff=None):
                nkpts = len(mo_occ[0])
                p0 = 0
                u = []
                for occ in mo_occ:
                    ua = []
                    for k in range(nkpts):
                        occidx = occ[k] > 0
                        viridx = ~occidx
                        nocc = occidx.sum()
                        nvir = viridx.sum()
                        nmo = nocc + nvir
                        dr = numpy.zeros((nmo,nmo), dtype=dx.dtype)
                        dr[viridx[:,None] & occidx] = dx[p0:p0+nocc*nvir]
                        dr = dr - dr.conj().T
                        p0 += nocc * nvir
                        u1 = newton_ah.expmat(dr)
                        if isinstance(u0, int) and u0 == 1:
                            ua.append(u1)
                        else:
                            ua.append(numpy.dot(u0[k], u1))
                    u.append(ua)
                return lib.asarray(u)

            def rotate_mo(self, mo_coeff, u, log=None):
                mo = ([numpy.dot(mo, u[0][k]) for k, mo in enumerate(mo_coeff[0])],
                      [numpy.dot(mo, u[1][k]) for k, mo in enumerate(mo_coeff[1])])
                return lib.asarray(mo)

        return SecondOrderKUHF(mf)

    elif isinstance(mf, pscf.krohf.KROHF):
        class SecondOrderKROHF(SecondOrderKRHF):
            gen_g_hop = gen_g_hop_rohf

        return SecondOrderKROHF(mf)

    elif isinstance(mf, pscf.kghf.KGHF):
        raise NotImplementedError

    else:
        return SecondOrderKRHF(mf)

if __name__ == '__main__':
    import pyscf.pbc.gto as pbcgto
    import pyscf.pbc.scf as pscf
    cell = pbcgto.Cell()
    cell.atom = '''
    He 0 0 1
    He 1 0 1
    '''
    cell.basis = 'ccpvdz'
    cell.a = numpy.eye(3) * 4
    cell.mesh = [11] * 3
    cell.verbose = 4
    cell.build()
    nks = [2,1,1]
    mf = pscf.KRHF(cell, cell.make_kpts(nks))
    mf.max_cycle = 2
    mf.kernel()
    mf.max_cycle = 5
    pscf.newton(mf).kernel()

    mf = pscf.KUHF(cell, cell.make_kpts(nks))
    mf.max_cycle = 2
    mf.kernel()
    mf.max_cycle = 5
    pscf.newton(mf).kernel()




# pyscf/pbc/scf/rohf.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Restricted open-shell Hartree-Fock for periodic systems at a single k-point

See Also:
    pyscf/pbc/scf/khf.py : Hartree-Fock for periodic systems with k-point sampling
'''

import numpy as np
from pyscf import lib
from pyscf.lib import logger
from pyscf.scf import rohf as mol_rohf
from pyscf.pbc.scf import hf as pbchf
from pyscf.pbc.scf import uhf as pbcuhf
from pyscf import __config__


get_fock = mol_rohf.get_fock
get_occ = mol_rohf.get_occ
get_grad = mol_rohf.get_grad
make_rdm1 = mol_rohf.make_rdm1
energy_elec = mol_rohf.energy_elec
dip_moment = pbcuhf.dip_moment
get_rho = pbcuhf.get_rho

class ROHF(pbchf.RHF, mol_rohf.ROHF):
    '''ROHF class for PBCs.
    '''

    direct_scf = getattr(__config__, 'pbc_scf_SCF_direct_scf', False)

    def __init__(self, cell, kpt=np.zeros(3),
                 exxdiv=getattr(__config__, 'pbc_scf_SCF_exxdiv', 'ewald')):
        pbchf.SCF.__init__(self, cell, kpt, exxdiv)
        self.nelec = None

    @property
    def nelec(self):
        if self._nelec is not None:
            return self._nelec
        else:
            cell = self.cell
            ne = cell.nelectron
            nalpha = (ne + cell.spin) // 2
            nbeta = nalpha - cell.spin
            if nalpha + nbeta != ne:
                raise RuntimeError('Electron number %d and spin %d are not consistent\n'
                                   'Note cell.spin = 2S = Nalpha - Nbeta, not 2S+1' %
                                   (ne, self.spin))
            return nalpha, nbeta
    @nelec.setter
    def nelec(self, x):
        self._nelec = x

    def dump_flags(self, verbose=None):
        pbchf.SCF.dump_flags(self, verbose)
        logger.info(self, 'number of electrons per unit cell  '
                    'alpha = %d beta = %d', *self.nelec)
        return self

    get_rho = get_rho

    eig = mol_rohf.ROHF.eig

    get_fock = mol_rohf.ROHF.get_fock
    get_occ = mol_rohf.ROHF.get_occ
    get_grad = mol_rohf.ROHF.get_grad
    make_rdm1 = mol_rohf.ROHF.make_rdm1
    energy_elec = mol_rohf.ROHF.energy_elec

    def get_veff(self, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
                 kpt=None, kpts_band=None):
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        if kpt is None: kpt = self.kpt
        if isinstance(dm, np.ndarray) and dm.ndim == 2:
            dm = np.asarray((dm*.5,dm*.5))
        if getattr(dm, 'mo_coeff', None) is not None:
            mo_coeff = dm.mo_coeff
            mo_occ_a = (dm.mo_occ > 0).astype(np.double)
            mo_occ_b = (dm.mo_occ ==2).astype(np.double)
            dm = lib.tag_array(dm, mo_coeff=(mo_coeff,mo_coeff),
                               mo_occ=(mo_occ_a,mo_occ_b))
        if self.rsjk and self.direct_scf:
            # Enable direct-SCF for real space JK builder
            ddm = dm - dm_last
            vj, vk = self.get_jk(cell, ddm, hermi, kpt, kpts_band)
            vhf = vj[0] + vj[1] - vk
            vhf += vhf_last
        else:
            vj, vk = self.get_jk(cell, dm, hermi, kpt, kpts_band)
            vhf = vj[0] + vj[1] - vk
        return vhf

    def get_bands(self, kpts_band, cell=None, dm=None, kpt=None):
        '''Get energy bands at the given (arbitrary) 'band' k-points.

        Returns:
            mo_energy : (nmo,) ndarray or a list of (nmo,) ndarray
                Bands energies E_n(k)
            mo_coeff : (nao, nmo) ndarray or a list of (nao,nmo) ndarray
                Band orbitals psi_n(k)
        '''
        raise NotImplementedError

    @lib.with_doc(dip_moment.__doc__)
    def dip_moment(self, cell=None, dm=None, unit='Debye', verbose=logger.NOTE,
                   **kwargs):
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        rho = kwargs.pop('rho', None)
        if rho is None:
            rho = self.get_rho(dm)
        return dip_moment(cell, dm, unit, verbose, rho=rho, kpt=self.kpt, **kwargs)

    def get_init_guess(self, cell=None, key='minao'):
        if cell is None: cell = self.cell
        dm = mol_rohf.ROHF.get_init_guess(self, cell, key)
        dm = pbchf.normalize_dm_(self, dm)
        return dm

    def init_guess_by_1e(self, cell=None):
        if cell is None: cell = self.cell
        if cell.dimension < 3:
            logger.warn(self, 'Hcore initial guess is not recommended in '
                        'the SCF of low-dimensional systems.')
        return mol_rohf.ROHF.init_guess_by_1e(self, cell)

    def init_guess_by_chkfile(self, chk=None, project=True, kpt=None):
        if chk is None: chk = self.chkfile
        if kpt is None: kpt = self.kpt
        return pbcuhf.init_guess_by_chkfile(self.cell, chk, project, kpt)

    init_guess_by_minao  = mol_rohf.ROHF.init_guess_by_minao
    init_guess_by_atom   = mol_rohf.ROHF.init_guess_by_atom
    init_guess_by_huckel = mol_rohf.ROHF.init_guess_by_huckel

    analyze = mol_rohf.ROHF.analyze
    canonicalize = mol_rohf.ROHF.canonicalize
    spin_square = mol_rohf.ROHF.spin_square
    stability = mol_rohf.ROHF.stability

    def convert_from_(self, mf):
        '''Convert given mean-field object to RHF/ROHF'''
        from pyscf.pbc.scf import addons
        addons.convert_to_rhf(mf, self)
        return self



# pyscf/pbc/scf/rsjk.py

#!/usr/bin/env python
# Copyright 2020-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Range separation JK builder

Ref:
    Q. Sun, arXiv:2012.07929
'''

import copy
import ctypes
import numpy as np
import scipy.linalg
from pyscf import gto
from pyscf import lib
from pyscf.lib import logger
from pyscf.scf import _vhf
from pyscf.pbc.df import aft, aft_jk
from pyscf.pbc import tools as pbctools
from pyscf.pbc.tools import k2gamma
from pyscf.pbc.df.df_jk import zdotNN, zdotCN, zdotNC, _ewald_exxdiv_for_G0
from pyscf.pbc.df.df_jk import _format_dms, _format_kpts_band, _format_jks
from pyscf.pbc.lib.kpts_helper import is_zero, gamma_point
from pyscf import __config__

# Threshold of steep bases and local bases
RCUT_THRESHOLD = getattr(__config__, 'pbc_scf_rsjk_rcut_threshold', 3.2)


libpbc = lib.load_library('libpbc')

class RangeSeparationJKBuilder(object):
    def __init__(self, cell, kpts=np.zeros((1,3))):
        self.cell = cell
        self.stdout = cell.stdout
        self.verbose = cell.verbose
        self.max_memory = cell.max_memory
        self.mesh = None
        self.kpts = np.reshape(kpts, (-1, 3))
        self.purify = True

        self.omega = None
        self.cell_rs = None
        # Born-von Karman supercell
        self.bvkcell = None
        self.bvkmesh_Ls = None
        self.bvk_kmesh = None
        self.phase = None
        self.supmol = None
        self.supmol_Ls = None
        # bvk_bas_mask[bvk_cell_id, bas_id]: if basis in bvkcell presented in supmol
        self.bvk_bas_mask = None
        # For shells in the supmol, cell0_shl_id is the shell ID in cell0
        self.cell0_shl_id = None
        # bvk_cell_id is the Id of image in BvK supercell
        self.bvk_cell_id = None
        # For shells in bvkcell, use overlap mask to remove d-d block
        self.ovlp_mask = None
        self.lr_aft = None
        self.ke_cutoff = None
        self.vhfopt = None

        # to mimic molecular DF object
        self.blockdim = getattr(__config__, 'pbc_df_df_DF_blockdim', 240)

        self._keys = set(self.__dict__.keys())

    def dump_flags(self, verbose=None):
        logger.info(self, '\n')
        logger.info(self, '******** %s ********', self.__class__)
        logger.info(self, 'mesh = %s (%d PWs)', self.mesh, np.prod(self.mesh))
        logger.info(self, 'omega = %s', self.omega)
        #logger.info(self, 'len(kpts) = %d', len(self.kpts))
        #logger.debug1(self, '    kpts = %s', self.kpts)
        if self.lr_aft is not None:
            self.lr_aft.dump_flags(verbose)
        return self

    def reset(self, cell=None):
        if cell is not None:
            self.cell = cell
        self.supmol = None
        return self

    def build(self, omega=None, direct_scf_tol=None):
        cpu0 = logger.process_clock(), logger.perf_counter()
        cell = self.cell
        kpts = self.kpts

        k_scaled = cell.get_scaled_kpts(kpts).sum(axis=0)
        k_mod_to_half = k_scaled * 2 - (k_scaled * 2).round(0)
        if abs(k_mod_to_half).sum() > 1e-5:
            raise NotImplementedError('k-points must be symmetryic')

        if omega is not None:
            self.omega = omega

        if self.omega is None:
            # Search a proper range-separation parameter omega that can balance the
            # computational cost between the real space integrals and moment space
            # integrals
            self.omega, self.mesh, self.ke_cutoff = _guess_omega(cell, kpts, self.mesh)
        else:
            self.ke_cutoff = aft.estimate_ke_cutoff_for_omega(cell, self.omega)
            self.mesh = pbctools.cutoff_to_mesh(cell.lattice_vectors(), self.ke_cutoff)

        logger.info(self, 'omega = %.15g  ke_cutoff = %s  mesh = %s',
                    self.omega, self.ke_cutoff, self.mesh)

        if direct_scf_tol is None:
            direct_scf_tol = cell.precision**1.5
            logger.debug(self, 'Set direct_scf_tol %g', direct_scf_tol)

        self.cell_rs = cell_rs = _re_contract_cell(cell, self.ke_cutoff)
        self.bvk_kmesh = kmesh = k2gamma.kpts_to_kmesh(cell_rs, kpts)
        bvkcell, phase = k2gamma.get_phase(cell_rs, kpts, kmesh)
        self.bvkmesh_Ls = Ks = k2gamma.translation_vectors_for_kmesh(cell_rs, kmesh)
        self.bvkcell = bvkcell
        self.phase = phase

        # Given ke_cutoff, eta corresponds to the most steep Gaussian basis
        # of which the Coulomb integrals can be accurately computed in moment
        # space.
        eta = aft.estimate_eta_for_ke_cutoff(cell, self.ke_cutoff,
                                             precision=cell.precision)
        # * Assuming the most steep function in smooth basis has exponent eta,
        # with attenuation parameter omega, rcut_sr is the distance of which
        # the value of attenuated Coulomb integrals of four shells |eta> is
        # smaller than the required precision.
        # * The attenuated coulomb integrals between four s-type Gaussians
        # (2*a/pi)^{3/4}exp(-a*r^2) is
        #   (erfc(omega*a^0.5/(omega^2+a)^0.5*R) - erfc(a^0.5*R)) / R
        # if two Gaussians on one center and the other two on another center
        # and the distance between the two centers are R.
        # * The attenuated coulomb integrals between two spherical charge
        # distributions is
        #   ~(pi/eta)^3/2 (erfc(tau*(eta/2)^0.5*R) - erfc((eta/2)^0.5*R)) / R
        #       tau = omega/sqrt(omega^2 + eta/2)
        # if the spherical charge distribution is the product of above s-type
        # Gaussian with exponent eta and a very smooth function.
        # When R is large, the attenuated Coulomb integral is
        #   ~= (pi/eta)^3/2 erfc(tau*(eta/2)^0.5*R) / R
        #   ~= pi/(tau*eta^2*R^2) exp(-tau^2*eta*R^2/2)
        tau = self.omega / (self.omega**2 + eta/2)**.5
        rcut_sr = 10  # initial guess
        rcut_sr = (-np.log(direct_scf_tol * tau * (eta * rcut_sr)**2/np.pi) / (tau**2*eta/2))**.5
        logger.debug(self, 'eta = %g  rcut_sr = %g', eta, rcut_sr)

        # Ls is the translation vectors to mimic periodicity of a cell
        Ls = bvkcell.get_lattice_Ls(rcut=cell.rcut+rcut_sr)
        self.supmol_Ls = Ls = Ls[np.linalg.norm(Ls, axis=1).argsort()]

        supmol = _make_extended_mole(cell_rs, Ls, Ks, self.omega, direct_scf_tol)
        self.supmol = supmol

        nkpts = len(self.bvkmesh_Ls)
        nbas = cell_rs.nbas
        n_steep, n_local, n_diffused = cell_rs._nbas_each_set
        n_compact = n_steep + n_local
        bas_mask = supmol._bas_mask

        self.bvk_bas_mask = bvk_bas_mask = bas_mask.any(axis=2)
        # Some basis in bvk-cell are not presented in the supmol. They can be
        # skipped when computing SR integrals
        self.bvkcell._bas = bvkcell._bas[bvk_bas_mask.ravel()]

        # Record the mapping between the dense bvkcell basis and the
        # original sparse bvkcell basis
        bvk_cell_idx = np.repeat(np.arange(nkpts)[:,None], nbas, axis=1)
        self.bvk_cell_id = bvk_cell_idx[bvk_bas_mask].astype(np.int32)
        cell0_shl_idx = np.repeat(np.arange(nbas)[None,:], nkpts, axis=0)
        self.cell0_shl_id = cell0_shl_idx[bvk_bas_mask].astype(np.int32)

        logger.timer_debug1(self, 'initializing supmol', *cpu0)
        logger.info(self, 'sup-mol nbas = %d cGTO = %d pGTO = %d',
                    supmol.nbas, supmol.nao, supmol.npgto_nr())

        supmol.omega = -self.omega  # Set short range coulomb
        with supmol.with_integral_screen(direct_scf_tol**2):
            vhfopt = _vhf.VHFOpt(supmol, 'int2e_sph',
                                 qcondname=libpbc.PBCVHFsetnr_direct_scf)
        vhfopt.direct_scf_tol = direct_scf_tol
        self.vhfopt = vhfopt
        logger.timer(self, 'initializing vhfopt', *cpu0)

        q_cond = vhfopt.get_q_cond((supmol.nbas, supmol.nbas))
        idx = supmol._images_loc
        bvk_q_cond = lib.condense('NP_absmax', q_cond, idx, idx)
        ovlp_mask = bvk_q_cond > direct_scf_tol
        # Remove diffused-diffused block
        if n_diffused > 0:
            diffused_mask = np.zeros_like(bvk_bas_mask)
            diffused_mask[:,n_compact:] = True
            diffused_mask = diffused_mask[bvk_bas_mask]
            ovlp_mask[diffused_mask[:,None] & diffused_mask] = False
        self.ovlp_mask = ovlp_mask.astype(np.int8)

        # mute rcut_threshold, divide basis into two sets only
        cell_lr_aft = _re_contract_cell(cell, self.ke_cutoff, -1, verbose=0)
        self.lr_aft = lr_aft = _LongRangeAFT(cell_lr_aft, kpts,
                                             self.omega, self.bvk_kmesh)
        lr_aft.ke_cutoff = self.ke_cutoff
        lr_aft.mesh = self.mesh
        lr_aft.eta = eta
        return self

    def get_jk(self, dm_kpts, hermi=1, kpts=None, kpts_band=None,
               with_j=True, with_k=True, omega=None, exxdiv=None):
        if omega is not None:  # J/K for RSH functionals
            # TODO: call AFTDF.get_jk function
            raise NotImplementedError

        # Does not support to specify arbitrary kpts
        if kpts is not None and abs(kpts-self.kpts).max() > 1e-7:
            raise RuntimeError('kpts error')
        kpts = self.kpts

        if kpts_band is not None:
            raise NotImplementedError

        cpu0 = logger.process_clock(), logger.perf_counter()
        if self.supmol is None:
            self.build()

        nkpts = kpts.shape[0]
        vhfopt = self.vhfopt
        supmol = self.supmol
        bvkcell = self.bvkcell
        phase = self.phase
        cell = self.cell_rs
        nao = cell.nao
        orig_nao = self.cell.nao

        # * dense_bvk_ao_loc are the AOs which appear in supmol (some basis
        # are removed)
        # * sparse_ao_loc has dimension (Nk,nbas), corresponding to the
        # bvkcell with all basis
        dense_bvk_ao_loc = bvkcell.ao_loc
        sparse_ao_loc = nao * np.arange(nkpts)[:,None] + cell.ao_loc[:-1]
        sparse_ao_loc = np.append(sparse_ao_loc.ravel(), nao * nkpts)
        nbands = nkpts

        if dm_kpts.ndim != 4:
            dm = dm_kpts.reshape(-1, nkpts, orig_nao, orig_nao)
        else:
            dm = dm_kpts
        n_dm = dm.shape[0]

        rs_c_coeff = cell._contr_coeff
        sc_dm = lib.einsum('nkij,pi,qj->nkpq', dm, rs_c_coeff, rs_c_coeff)
        # Utilized symmetry sc_dm[R,S] = sc_dm[S-R] = sc_dm[(S-R)%N]
        #:sc_dm = lib.einsum('Rk,nkuv,Sk->nRuSv', phase, sc_dm, phase.conj())
        sc_dm = lib.einsum('k,Sk,nkuv->nSuv', phase[0], phase.conj(), sc_dm)
        dm_translation = k2gamma.double_translation_indices(self.bvk_kmesh).astype(np.int32)
        dm_imag_max = abs(sc_dm.imag).max()
        is_complex_dm = dm_imag_max > 1e-6
        if is_complex_dm:
            if dm_imag_max < 1e-2:
                logger.warn(self, 'DM in (BvK) cell has small imaginary part.  '
                            'It may be a signal of symmetry broken in k-point symmetry')
            sc_dm = np.vstack([sc_dm.real, sc_dm.imag])
        else:
            sc_dm = sc_dm.real
        sc_dm = np.asarray(sc_dm.reshape(-1, nkpts, nao, nao), order='C')
        n_sc_dm = sc_dm.shape[0]

        dm_cond = [lib.condense('NP_absmax', d, sparse_ao_loc, sparse_ao_loc[:cell.nbas+1])
                   for d in sc_dm]
        dm_cond = np.asarray(np.max(dm_cond, axis=0), order='C')
        libpbc.CVHFset_dm_cond(vhfopt._this,
                               dm_cond.ctypes.data_as(ctypes.c_void_p), dm_cond.size)
        dm_cond = None

        bvk_nbas = bvkcell.nbas
        shls_slice = (0, cell.nbas, 0, bvk_nbas, 0, bvk_nbas, 0, bvk_nbas)

        if hermi:
            fdot_suffix = 's2kl'
        else:
            fdot_suffix = 's1'
        if with_j and with_k:
            fdot = 'PBCVHF_contract_jk_' + fdot_suffix
            vs = np.zeros((2, n_sc_dm, nao, nkpts, nao))
        elif with_j:
            fdot = 'PBCVHF_contract_j_' + fdot_suffix
            vs = np.zeros((1, n_sc_dm, nao, nkpts, nao))
        else:  # with_k
            fdot = 'PBCVHF_contract_k_' + fdot_suffix
            vs = np.zeros((1, n_sc_dm, nao, nkpts, nao))

        drv = libpbc.PBCVHF_direct_drv
        drv(getattr(libpbc, fdot), vs.ctypes.data_as(ctypes.c_void_p),
            sc_dm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(n_dm),
            ctypes.c_int(nkpts), ctypes.c_int(nbands), ctypes.c_int(cell.nbas),
            self.ovlp_mask.ctypes.data_as(ctypes.c_void_p),
            self.bvk_cell_id.ctypes.data_as(ctypes.c_void_p),
            self.cell0_shl_id.ctypes.data_as(ctypes.c_void_p),
            supmol._images_loc.ctypes.data_as(ctypes.c_void_p),
            (ctypes.c_int*8)(*shls_slice),
            dense_bvk_ao_loc.ctypes.data_as(ctypes.c_void_p),
            dm_translation.ctypes.data_as(ctypes.c_void_p),
            vhfopt._cintopt, vhfopt._this,
            supmol._atm.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(supmol.natm),
            supmol._bas.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(supmol.nbas),
            supmol._env.ctypes.data_as(ctypes.c_void_p))

        if is_complex_dm:
            vs = vs[:,:n_dm] + vs[:,n_dm:] * 1j

        if with_j and with_k:
            vj, vk = vs
        elif with_j:
            vj, vk = vs[0], None
        else:
            vj, vk = None, vs[0]
        cpu1 = logger.timer(self, 'short range part vj and vk', *cpu0)

        lr_c_coeff = self.lr_aft.cell._contr_coeff
        lr_dm = lib.einsum('nkij,pi,qj->nkpq', dm, lr_c_coeff, lr_c_coeff)
        # For rho product other than diffused-diffused block, construct LR
        # parts in terms of full ERIs and SR ERIs
        vj1, vk1 = self.lr_aft.get_jk(lr_dm, hermi, kpts, kpts_band,
                                      with_j, with_k, exxdiv=exxdiv)
        cpu1 = logger.timer(self, 'AFT-vj and AFT-vk', *cpu1)

        # expRk is almost the same to phase, except a normalization factor
        expRk = np.exp(1j*np.dot(self.bvkmesh_Ls, kpts.T))

        if with_j:
            vj = lib.einsum('npRq,pi,qj,Rk->nkij', vj, rs_c_coeff, rs_c_coeff, expRk)
            vj += lib.einsum('nkpq,pi,qj->nkij', vj1, lr_c_coeff, lr_c_coeff)
            if self.purify and kpts_band is None:
                vj = _purify(vj, phase)
            if gamma_point(kpts) and dm_kpts.dtype == np.double:
                vj = vj.real
            if hermi:
                vj = (vj + vj.conj().transpose(0,1,3,2)) * .5
            vj = vj.reshape(dm_kpts.shape)

        if with_k:
            vk = lib.einsum('npRq,pi,qj,Rk->nkij', vk, rs_c_coeff, rs_c_coeff, expRk)
            vk += lib.einsum('nkpq,pi,qj->nkij', vk1, lr_c_coeff, lr_c_coeff)
            if self.purify and kpts_band is None:
                vk = _purify(vk, phase)
            if gamma_point(kpts) and dm_kpts.dtype == np.double:
                vk = vk.real
            if hermi:
                vk = (vk + vk.conj().transpose(0,1,3,2)) * .5
            vk = vk.reshape(dm_kpts.shape)

        return vj, vk

def _purify(mat_kpts, phase):
    #:mat_bvk = np.einsum('Rk,nkij,Sk->nRSij', phase, mat_kpts, phase.conj())
    #:return np.einsum('Rk,nRSij,Sk->nkij', phase.conj(), mat_bvk.real, phase)
    nkpts = phase.shape[1]
    mat_bvk = lib.einsum('k,Sk,nkuv->nSuv', phase[0], phase.conj(), mat_kpts)
    return lib.einsum('S,Sk,nSuv->nkuv', nkpts*phase[:,0].conj(), phase, mat_bvk.real)

def _make_extended_mole(cell, Ls, bvkmesh_Ls, omega, precision=None, verbose=None):
    from pyscf.pbc.df.ft_ao import _estimate_overlap
    if precision is None:
        precision = cell.precision * 1e-4
    LKs = Ls[:,None,:] + bvkmesh_Ls
    nimgs, nk = LKs.shape[:2]
    nbas = cell.nbas

    supmol = cell.to_mol()
    supmol = pbctools.pbc._build_supcell_(supmol, cell, LKs.reshape(nimgs*nk, 3))

    n_steep, n_local, n_diffused = cell._nbas_each_set
    n_compact = n_steep + n_local

    exps = np.array([cell.bas_exp(ib).min() for ib in range(nbas)])
    aij = exps[:,None] * exps / (exps[:,None] + exps)
    s0 = _estimate_overlap(cell, LKs.reshape(nimgs*nk, 3))
    # (bas_i, bas_j, image_id, bvk_cell_for_j)
    s0max = (2 * (aij[:,:,None]/np.pi)**.5 * s0).max(axis=0)

    if n_diffused > 0 and n_compact > 0:
        # For remotely separated smooth functions |i> and |l>, the eri-type
        # (ij|kl) can have small contributions if |j> and |k> are steep functions
        # and close in real space. Since we are focusing on the SR attenuated
        # Coulomb integrals, we don't have to consider the steep functions j and k
        # remotely separated since SR integrals decays exponently wrt their
        # separation.
        atom_coords = cell.atom_coords()
        bas_coords = atom_coords[cell._bas[:,gto.ATOM_OF]]
        exps_d = exps[n_compact:]  # diffused
        exps_c_min = exps[:n_compact].min()
        aij_min = aij[:n_compact,n_compact:].min(axis=0)
        # The product of two diffused functions
        rij_dd = bas_coords[n_compact:,None,:] - bas_coords[n_compact:]
        # The product of two diffused functions from different cells
        dijL = np.linalg.norm(rij_dd[:,:,None,:] - LKs.reshape(nimgs*nk,3), axis=-1)
        fac = 16/(2*np.pi)**.5/(exps_d + exps_c_min) * aij_min**1.5
        # Estimation of (DC|DC) from four s-type shells
        eri_dcdc = fac[:,None,None] * np.exp(-.5*aij_min[:,None,None] * dijL**2)
        s0max[n_compact:] = eri_dcdc.max(axis=1)

    if n_diffused == 0:
        # When diffused functions exist, ranges of compact functions are
        # determined by the diffused functions. Without diffused functions,
        # assuming two charge distributions are well separated, their
        # interactions are erfc(omega*r12)/r12
        exps_c = exps[:n_compact]
        r2_LKs = np.einsum('lkx,lkx->lk', LKs, LKs).ravel()
        # Exclude cell 0
        assert r2_LKs[0] == 0, 'Ls not sorted'
        r2_LKs = r2_LKs[1:]
        fac = -omega**2*exps_c/(2*omega**2+exps_c)
        upper_bounds = np.exp(fac[:,None] * r2_LKs) / (2*omega*r2_LKs)
        s0max[:n_compact,1:] = np.max([s0max[:n_compact,1:], upper_bounds], axis=0)

    # (bas_id, image_id, bvk_cell_id) -> (bvk_cell_id, bas_id, image_id)
    s0max = s0max.reshape(nbas, nimgs, nk).transpose(2,0,1)
    bas_mask = s0max > precision
    _bas_reordered = supmol._bas.reshape(nimgs,nk,nbas,gto.BAS_SLOTS).transpose(1,2,0,3)
    supmol._bas = np.asarray(_bas_reordered[bas_mask], dtype=np.int32, order='C')

    images_count = np.count_nonzero(bas_mask, axis=2)
    # Some bases are completely local inside the bvk-cell. Exclude them from
    # lattice sum.
    images_loc = np.append(0, np.cumsum(images_count[images_count != 0]))
    supmol._images_loc = images_loc.astype(np.int32)
    supmol._bas_mask = bas_mask

    log = logger.new_logger(cell, verbose)
    log.debug('Steep basis in sup-mol %d', np.count_nonzero(bas_mask[:,:n_steep,:]))
    log.debug('Local basis in sup-mol %d', np.count_nonzero(bas_mask[:,n_steep:n_compact,:]))
    log.debug('Diffused basis in sup-mol %d', np.count_nonzero(bas_mask[:,n_compact:,:]))
    return supmol

def _re_contract_cell(cell, ke_cut_threshold, rcut_threshold=RCUT_THRESHOLD, verbose=None):
    from pyscf.gto import NPRIM_OF, NCTR_OF, PTR_EXP, PTR_COEFF, ATOM_OF
    from pyscf.pbc.dft.multigrid import _primitive_gto_cutoff
    log = logger.new_logger(cell, verbose)

    # Split shells based on rcut
    rcuts, kecuts = _primitive_gto_cutoff(cell, cell.precision)

    def transform_(cint_coeff, orig_bas, bas_to_append, env, pexp, pcoeff):
        np1, nc1 = cint_coeff.shape
        l = orig_bas[gto.ANG_OF]
        if cell.cart:
            degen = (l + 1) * (l + 2) // 2
        else:
            degen = 2 * l + 1
        if np1 >= nc1:
            bas = orig_bas.copy()
            bas[NPRIM_OF] = np1
            bas[PTR_EXP] = pexp
            bas[PTR_COEFF] = pcoeff
            bas_to_append.append(bas)
            coeff = np.eye(nc1 * degen)
        else:
            bas = np.repeat(orig_bas.copy()[None,:], np1, axis=0)
            bas[:,NPRIM_OF] = 1
            bas[:,NCTR_OF] = 1
            bas[:,PTR_EXP] = pexp + np.arange(np1)
            bas[:,PTR_COEFF] = pcoeff + np.arange(np1)
            bas_to_append.extend(bas)

            exps = _env[pexp:pexp+np1]
            cs = gto.gto_norm(l, exps)
            _env[pcoeff:pcoeff+np1] = cs
            unit_coeff = np.eye(degen)
            coeff = np.einsum('p,mn,pc->pmcn', 1/cs, unit_coeff, cint_coeff)
            coeff = coeff.reshape(np1*degen, nc1*degen)
        return coeff

    ao_loc = cell.ao_loc_nr()
    _env = cell._env.copy()
    steep_bas = []
    local_bas = []
    smooth_bas = []
    steep_p2c = []
    local_p2c = []
    smooth_p2c = []
    ke_cutoff = 0
    for ib, orig_bas in enumerate(cell._bas):
        nprim = orig_bas[NPRIM_OF]
        nctr = orig_bas[NCTR_OF]
        ke = kecuts[ib]

        smooth_mask = ke < ke_cut_threshold
        steep_mask = (~smooth_mask) & (rcuts[ib] < rcut_threshold)
        local_mask = (~steep_mask) & (~smooth_mask)
        if log.verbose >= logger.DEBUG3:
            log.debug3('bas %d rcuts %s', ib, rcuts)
            log.debug3('bas %d kecuts %s', ib, kecuts)
            log.debug3('steep %s, local %s, smooth', np.where(steep_mask),
                       np.where(local_mask), np.where(smooth_mask))

        pexp = orig_bas[PTR_EXP]
        pcoeff = orig_bas[PTR_COEFF]
        es = cell.bas_exp(ib)
        cs = cell._libcint_ctr_coeff(ib)

        c_steep = cs[steep_mask]
        c_local = cs[local_mask]
        c_smooth = cs[smooth_mask]
        _env[pcoeff:pcoeff+nprim*nctr] = np.hstack([
            c_steep.T.ravel(),
            c_local.T.ravel(),
            c_smooth.T.ravel(),
        ])
        _env[pexp:pexp+nprim] = np.hstack([
            es[steep_mask],
            es[local_mask],
            es[smooth_mask],
        ])

        if c_steep.size > 0:
            steep_p2c.append(transform_(c_steep, orig_bas, steep_bas, _env,
                                        pexp, pcoeff))
        else:
            steep_p2c.append(np.zeros((0, ao_loc[ib+1]-ao_loc[ib])))

        if c_local.size > 0:
            local_p2c.append(transform_(c_local, orig_bas, local_bas, _env,
                                        pexp + c_steep.shape[0],
                                        pcoeff + c_steep.size))
        else:
            local_p2c.append(np.zeros((0, ao_loc[ib+1]-ao_loc[ib])))

        if c_smooth.size > 0:
            smooth_p2c.append(transform_(c_smooth, orig_bas, smooth_bas, _env,
                                         pexp + c_steep.shape[0] + c_local.shape[0],
                                         pcoeff + c_steep.size + c_local.size))
            ke_cutoff = max(ke_cutoff, ke[smooth_mask].max())
        else:
            smooth_p2c.append(np.zeros((0, ao_loc[ib+1]-ao_loc[ib])))

    cell_rs = copy.copy(cell)
    cell_rs._bas = np.asarray(steep_bas + local_bas + smooth_bas,
                              dtype=np.int32, order='C').reshape(-1, gto.BAS_SLOTS)
    cell_rs._env = _env
    cell_rs._contr_coeff = np.vstack([scipy.linalg.block_diag(*steep_p2c),
                                      scipy.linalg.block_diag(*local_p2c),
                                      scipy.linalg.block_diag(*smooth_p2c)])
    cell_rs._nbas_each_set = (len(steep_bas), len(local_bas), len(smooth_bas))
    log.debug('No. steep_bas %d', len(steep_bas))
    log.debug('No. local_bas %d', len(local_bas))
    log.debug('No. smooth_bas %d', len(smooth_bas))
    return cell_rs

def _guess_omega(cell, kpts, mesh=None):
    nao = cell.npgto_nr()
    nkpts = len(kpts)
    nkk = nkpts**(1./3) * 2 - 1
    if mesh is None:
        mesh = [max(5, int(cell.rcut * nao ** (1./3) / nkk + 1))] * 3
        mesh = np.min([cell.mesh, mesh], axis=0)
    ke_cutoff = min(pbctools.mesh_to_cutoff(cell.lattice_vectors(), mesh[:cell.dimension]))
    omega = aft.estimate_omega_for_ke_cutoff(cell, ke_cutoff)
    return omega, mesh, ke_cutoff

class _LongRangeAFT(aft.AFTDF):
    '''
    Regular Coulomb metric for (DD|DD), (DD|CD), (DD|CC)
    Short-range Coulomb metric for (CD|CD), (CD|CC), (CC|CC)
    '''
    def __init__(self, cell, kpts=np.zeros((1,3)), omega=None, bvk_kmesh=None):
        self.omega = omega
        self.bvk_kmesh = bvk_kmesh
        aft.AFTDF.__init__(self, cell, kpts)

    def weighted_coulG_LR(self, kpt=np.zeros(3), exx=False, mesh=None):
        cell = self.cell
        if mesh is None:
            mesh = self.mesh
        if cell.omega != 0:
            raise RuntimeError('RangeSeparationJKBuilder cannot be used to evaluate '
                               'the long-range HF exchange in RSH functional')
        Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
        coulG = pbctools.get_coulG(cell, kpt, False, self, mesh, Gv,
                                   omega=self.omega)
        coulG *= kws
        return coulG

    def weighted_coulG_SR(self, kpt=np.zeros(3), exx=False, mesh=None):
        cell = self.cell
        if mesh is None:
            mesh = self.mesh
        if cell.omega != 0:
            raise RuntimeError('RangeSeparationJKBuilder cannot be used to evaluate '
                               'the long-range HF exchange in RSH functional')
        Gv, Gvbase, kws = cell.get_Gv_weights(mesh)
        coulG = pbctools.get_coulG(cell, kpt, False, self, mesh, Gv,
                                   omega=-self.omega)
        coulG *= kws
        return coulG

    def ft_loop(self, mesh=None, q=np.zeros(3), kpts=None, shls_slice=None,
                max_memory=4000, aosym='s1', intor='GTO_ft_ovlp', comp=1):
        for dat, p0, p1 in aft.AFTDF.ft_loop(self, mesh, q, kpts, shls_slice,
                                             max_memory, aosym, intor, comp,
                                             bvk_kmesh=self.bvk_kmesh):
            yield dat, p0, p1

    def get_jk(self, dm, hermi=1, kpts=None, kpts_band=None,
               with_j=True, with_k=True, omega=None, exxdiv=None):
        if omega is not None:  # J/K for RSH functionals
            raise NotImplementedError

        if kpts is None:
            if np.all(self.kpts == 0):
                # Gamma-point calculation by default
                kpts = np.zeros(3)
            else:
                kpts = self.kpts
        kpts = np.asarray(kpts)

        is_single_kpt = kpts.ndim == 1
        if is_single_kpt:
            kpts = kpts.reshape(1,3)

        vj = vk = None
        if with_k:
            vk = self.get_k_kpts(dm, hermi, kpts, kpts_band, exxdiv)
            if is_single_kpt:
                vk = vk[...,0,:,:]
        if with_j:
            vj = self.get_j_kpts(dm, hermi, kpts, kpts_band)
            if is_single_kpt:
                vj = vj[...,0,:,:]
        return vj, vk

    def get_j_kpts(self, dm_kpts, hermi=1, kpts=np.zeros((1,3)), kpts_band=None):
        '''
        C ~ compact basis, D ~ diffused basis

        Compute J matrix with coulG_LR:
        (CC|CC) (CC|CD) (CC|DC) (CD|CC) (CD|CD) (CD|DC) (DC|CC) (DC|CD) (DC|DC)

        Compute J matrix with full coulG:
        (CC|DD) (CD|DD) (DC|DD) (DD|CC) (DD|CD) (DD|DC) (DD|DD)
        '''
        if kpts_band is not None:
            return self.get_j_for_bands(dm_kpts, hermi, kpts, kpts_band)

        if len(kpts) == 1 and not is_zero(kpts):
            raise NotImplementedError('Single k-point get-j')

        cell = self.cell
        dm_kpts = lib.asarray(dm_kpts, order='C')
        dms = _format_dms(dm_kpts, kpts)
        n_dm, nkpts, nao = dms.shape[:3]

        n_diffused = cell._nbas_each_set[2]
        nao_compact = cell.ao_loc[cell.nbas-n_diffused]

        vj_kpts = np.zeros((n_dm,nkpts,nao,nao), dtype=np.complex128)
        kpt_allow = np.zeros(3)
        mesh = self.mesh
        coulG = self.weighted_coulG(kpt_allow, False, mesh)
        coulG_LR = self.weighted_coulG_LR(kpt_allow, False, mesh)
        coulG_SR = coulG - coulG_LR
        max_memory = (self.max_memory - lib.current_memory()[0]) * .8
        weight = 1./len(kpts)
        for aoaoks, p0, p1 in self.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
            if nao_compact < nao:
                aoaoks = [aoao.reshape(-1,nao,nao) for aoao in aoaoks]
                aft_jk._update_vj_(vj_kpts, aoaoks, dms, coulG[p0:p1], weight)
                for aoao in aoaoks:
                    aoao[:,nao_compact:,nao_compact:] = 0
                aft_jk._update_vj_(vj_kpts, aoaoks, dms, coulG_SR[p0:p1], -weight)
            else:
                aft_jk._update_vj_(vj_kpts, aoaoks, dms, coulG_LR[p0:p1], weight)
            aoao = aoaoks = p0 = p1 = None

        # G=0 contribution, associated to 2e integrals in real-space
        if cell.dimension >= 2:
            ovlp = np.asarray(cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpts))
            if nao_compact < nao:
                ovlp[:,nao_compact:,nao_compact:] = 0
            kws = cell.get_Gv_weights(mesh)[2]
            G0_weight = kws[0] if isinstance(kws, np.ndarray) else kws
            vj_G0 = lib.einsum('kpq,nkqp,lrs->nlrs', ovlp, dm_kpts, ovlp)
            vj_kpts -= np.pi/self.omega**2 * weight * G0_weight * vj_G0

        if gamma_point(kpts):
            vj_kpts = vj_kpts.real.copy()
        return _format_jks(vj_kpts, dm_kpts, kpts_band, kpts)

    def get_j_for_bands(self, dm_kpts, hermi=1, kpts=np.zeros((1,3)), kpts_band=None):
        log = logger.Logger(self.stdout, self.verbose)
        t1 = logger.process_clock(), logger.perf_counter()

        cell = self.cell
        dm_kpts = lib.asarray(dm_kpts, order='C')
        dms = _format_dms(dm_kpts, kpts)
        n_dm, nkpts, nao = dms.shape[:3]

        n_diffused = cell._nbas_each_set[2]
        nao_compact = cell.ao_loc[cell.nbas-n_diffused]

        kpt_allow = np.zeros(3)
        mesh = self.mesh
        coulG = self.weighted_coulG(kpt_allow, False, mesh)
        coulG_LR = self.weighted_coulG_LR(kpt_allow, False, mesh)
        coulG_SR = coulG - coulG_LR
        ngrids = len(coulG)
        vG = np.zeros((n_dm,ngrids), dtype=np.complex128)
        vG_SR = np.zeros((n_dm,ngrids), dtype=np.complex128)
        max_memory = (self.max_memory - lib.current_memory()[0]) * .8

        for aoaoks, p0, p1 in self.ft_loop(mesh, kpt_allow, kpts, max_memory=max_memory):
            #:rho = np.einsum('lkL,lk->L', pqk.conj(), dm)
            for k, aoao in enumerate(aoaoks):
                aoao = aoao.reshape(-1,nao,nao)
                if nao_compact < nao:
                    for i in range(n_dm):
                        rho = np.einsum('ij,Lji->L', dms[i,k], aoao.conj())
                        vG[i,p0:p1] += rho * coulG[p0:p1]
                    aoao[:,nao_compact:,nao_compact:] = 0
                    for i in range(n_dm):
                        rho = np.einsum('ij,Lji->L', dms[i,k], aoao.conj())
                        vG_SR[i,p0:p1] += rho * coulG_SR[p0:p1]
                else:
                    for i in range(n_dm):
                        rho = np.einsum('ij,Lji->L', dms[i,k], aoao.conj())
                        vG[i,p0:p1] += rho * coulG_LR[p0:p1]
            aoao = aoaoks = p0 = p1 = None
        weight = 1./len(kpts)
        vG *= weight
        vG_SR *= weight
        t1 = log.timer_debug1('get_j pass 1 to compute J(G)', *t1)

        kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
        nband = len(kpts_band)
        vj_kpts = np.zeros((n_dm,nband,nao,nao), dtype=np.complex128)
        for aoaoks, p0, p1 in self.ft_loop(mesh, kpt_allow, kpts_band,
                                            max_memory=max_memory):
            for k, aoao in enumerate(aoaoks):
                aoao = aoao.reshape(-1,nao,nao)
                if nao_compact < nao:
                    for i in range(n_dm):
                        vj_kpts[i,k] += np.einsum('L,Lij->ij', vG[i,p0:p1], aoao)
                    aoao[:,nao_compact:,nao_compact:] = 0
                    for i in range(n_dm):
                        vj_kpts[i,k] -= np.einsum('L,Lij->ij', vG_SR[i,p0:p1], aoao)
                else:
                    for i in range(n_dm):
                        vj_kpts[i,k] += np.einsum('L,Lij->ij', vG[i,p0:p1], aoao)
            aoao = aoaoks = p0 = p1 = None

        # G=0 contribution, associated to 2e integrals in real-space
        if cell.dimension >= 2:
            ovlp = np.asarray(cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpts))
            ovlp[:,nao_compact:,nao_compact:] = 0
            ovlp_b = np.asarray(cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpts_band))
            ovlp_b[:,nao_compact:,nao_compact:] = 0
            kws = cell.get_Gv_weights(mesh)[2]
            G0_weight = kws[0] if isinstance(kws, np.ndarray) else kws
            vj_G0 = lib.einsum('kpq,nkqp,lrs->nlrs', ovlp, dm_kpts, ovlp_b)
            vj_kpts -= np.pi/self.omega**2 * weight * G0_weight * vj_G0

        if gamma_point(kpts_band):
            vj_kpts = vj_kpts.real.copy()
        t1 = log.timer_debug1('get_j pass 2', *t1)
        return _format_jks(vj_kpts, dm_kpts, input_band, kpts)

    def get_k_kpts(self, dm_kpts, hermi=1, kpts=np.zeros((1,3)), kpts_band=None,
                   exxdiv=None):
        '''
        C ~ compact basis, D ~ diffused basis

        Compute K matrix with coulG_LR:
        (CC|CC) (CC|CD) (CC|DC) (CD|CC) (CD|CD) (CD|DC) (DC|CC) (DC|CD) (DC|DC)

        Compute K matrix with full coulG:
        (CC|DD) (CD|DD) (DC|DD) (DD|CC) (DD|CD) (DD|DC) (DD|DD)
        '''
        cell = self.cell
        log = logger.Logger(self.stdout, self.verbose)
        t1 = logger.process_clock(), logger.perf_counter()

        mesh = self.mesh
        dm_kpts = lib.asarray(dm_kpts, order='C')
        dms = _format_dms(dm_kpts, kpts)
        nset, nkpts, nao = dms.shape[:3]

        swap_2e = (kpts_band is None)
        kpts_band, input_band = _format_kpts_band(kpts_band, kpts), kpts_band
        nband = len(kpts_band)
        kk_table = kpts_band.reshape(-1,1,3) - kpts.reshape(1,-1,3)
        kk_todo = np.ones(kk_table.shape[:2], dtype=bool)
        vkR = np.zeros((nset,nband,nao,nao))
        vkI = np.zeros((nset,nband,nao,nao))
        dmsR = np.asarray(dms.real, order='C')
        dmsI = np.asarray(dms.imag, order='C')
        weight = 1. / nkpts

        n_diffused = cell._nbas_each_set[2]
        nao_compact = cell.ao_loc[cell.nbas-n_diffused]

        mem_now = lib.current_memory()[0]
        max_memory = max(2000, (self.max_memory - mem_now)) * .8
        log.debug1('max_memory = %d MB (%d in use)', max_memory, mem_now)
        # K_pq = ( p{k1} i{k2} | i{k2} q{k1} )
        def make_kpt(kpt):  # kpt = kptj - kpti
            # search for all possible ki and kj that has ki-kj+kpt=0
            kk_match = np.einsum('ijx->ij', abs(kk_table + kpt)) < 1e-9
            kpti_idx, kptj_idx = np.where(kk_todo & kk_match)
            nkptj = len(kptj_idx)
            log.debug1('kpt = %s', kpt)
            log.debug2('kpti_idx = %s', kpti_idx)
            log.debug2('kptj_idx = %s', kptj_idx)
            kk_todo[kpti_idx,kptj_idx] = False
            if swap_2e and not is_zero(kpt):
                kk_todo[kptj_idx,kpti_idx] = False

            max_memory1 = max_memory * (nkptj+1)/(nkptj+5)
            #blksize = max(int(max_memory1*4e6/(nkptj+5)/16/nao**2), 16)

            #bufR = np.empty((blksize*nao**2))
            #bufI = np.empty((blksize*nao**2))
            # Use DF object to mimic KRHF/KUHF object in function get_coulG
            vkcoulG = self.weighted_coulG(kpt, exxdiv, mesh)
            coulG_SR = self.weighted_coulG_SR(kpt, False, mesh)
            coulG_LR = vkcoulG - coulG_SR
            kptjs = kpts[kptj_idx]
            perm_sym = swap_2e and not is_zero(kpt)
            for aoaoks, p0, p1 in self.ft_loop(mesh, kpt, kptjs, max_memory=max_memory1):
                if nao_compact < nao:
                    aoaoks = [aoao.reshape(-1,nao,nao) for aoao in aoaoks]
                    aft_jk._update_vk_((vkR, vkI), aoaoks, (dmsR, dmsI), vkcoulG[p0:p1],
                                       weight, kpti_idx, kptj_idx, perm_sym)
                    for aoao in aoaoks:
                        aoao[:,nao_compact:,nao_compact:] = 0
                    aft_jk._update_vk_((vkR, vkI), aoaoks, (dmsR, dmsI), coulG_SR[p0:p1],
                                       -weight, kpti_idx, kptj_idx, perm_sym)
                else:
                    aft_jk._update_vk_((vkR, vkI), aoaoks, (dmsR, dmsI), coulG_LR[p0:p1],
                                       weight, kpti_idx, kptj_idx, perm_sym)

        for ki, kpti in enumerate(kpts_band):
            for kj, kptj in enumerate(kpts):
                if kk_todo[ki,kj]:
                    make_kpt(kptj-kpti)
            t1 = log.timer_debug1('get_k_kpts: make_kpt (%d,*)'%ki, *t1)

        if (gamma_point(kpts) and gamma_point(kpts_band) and
            not np.iscomplexobj(dm_kpts)):
            vk_kpts = vkR
        else:
            vk_kpts = vkR + vkI * 1j

        # G=0 associated to 2e integrals in real-space
        if cell.dimension >= 2:
            ovlp = np.asarray(cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpts))
            ovlp[:,nao_compact:,nao_compact:] = 0
            kws = cell.get_Gv_weights(mesh)[2]
            G0_weight = kws[0] if isinstance(kws, np.ndarray) else kws
            vk_G0 = lib.einsum('kpq,nkqr,krs->nkps', ovlp, dm_kpts, ovlp)
            vk_kpts -= np.pi/self.omega**2 * weight * G0_weight * vk_G0

        # Add ewald_exxdiv contribution because G=0 was not included in the
        # non-uniform grids
        if (exxdiv == 'ewald' and
            (cell.dimension < 2 or  # 0D and 1D are computed with inf_vacuum
             (cell.dimension == 2 and cell.low_dim_ft_type == 'inf_vacuum'))):
            _ewald_exxdiv_for_G0(cell, kpts_band, dms, vk_kpts, kpts_band)

        return _format_jks(vk_kpts, dm_kpts, input_band, kpts)

if __name__ == '__main__':
    from pyscf.pbc.gto import Cell
    cells = []

    cell = Cell()
    cell.a = np.eye(3)*1.8
    cell.atom = '''#He     0.      0.      0.
                   He     0.4917  0.4917  0.4917'''
    cell.basis = {'He': [[0, [2.5, 1]]]}
    cell.build()
    cells.append(cell)

    if 1:
        cell = Cell()
        cell.a = np.eye(3)*2.4
        cell.atom = '''He     0.      0.      0.
                       He     0.4917  0.4917  0.4917'''
        cell.basis = {'He': [[0, [4.1, 1, -.2],
                                 [0.5, .2, .5],
                                 [0.15, .5, .5]],
                             #[1, [1.5, 1]],
                             [1, [0.3, 1]],]}
        cell.build()
        cell.verbose = 6
        cells.append(cell)

    if 1:
        cell = Cell().build(a = '''
3.370137329, 0.000000000, 0.000000000
0.000000000, 3.370137329, 0.000000000
0.000000000, 0.000000000, 3.370137329
                            ''',
                            unit = 'Ang',
                            atom = '''
C 0.000000000000   0.000000000000   0.000000000000
C 1.685068664391   1.685068664391   1.685068664391''',
                            basis='''
C S
4.3362376436      0.1490797872
1.2881838513      -0.0292640031
0.4037767149      -0.688204051
0.1187877657      -0.3964426906
C P
4.3362376436      -0.0878123619
1.2881838513      -0.27755603
0.4037767149      -0.4712295093
0.1187877657      -0.4058039291
''')
        cell.verbose = 6
        cells.append(cell)

    for cell in cells:
        kpts = cell.make_kpts([3,1,1])
        mf = cell.KRHF(kpts=kpts)#.run()
        #dm = mf.make_rdm1()
        np.random.seed(1)
        dm = (np.random.rand(len(kpts), cell.nao, cell.nao) +
              np.random.rand(len(kpts), cell.nao, cell.nao) * 1j)
        dm = dm + dm.transpose(0,2,1).conj()
        kmesh = k2gamma.kpts_to_kmesh(cell, kpts)
        phase = k2gamma.get_phase(cell, kpts, kmesh)[1]
        dm = lib.einsum('Rk,kuv,Sk->RSuv', phase.conj().T, dm, phase.T)
        dm = lib.einsum('Rk,RSuv,Sk->kuv', phase, dm.real, phase.conj())

        jref, kref = mf.get_jk(cell, dm, kpts=kpts)
        ej = np.einsum('kij,kji->', jref, dm)
        ek = np.einsum('kij,kji->', kref, dm) * .5

        jk_builder = RangeSeparationJKBuilder(cell, kpts)
        jk_builder.build(omega=0.5)
        #jk_builder.mesh = [6,6,6]
        #print(jk_builder.omega, jk_builder.mesh)
        vj, vk = jk_builder.get_jk(dm, kpts=kpts, exxdiv=mf.exxdiv)
        print(abs(vj - jref).max())
        print(abs(vk - kref).max())
        print('ej_ref', ej, 'ek_ref', ek)
        print('ej', np.einsum('kij,kji->', vj, dm).real,
              'ek', np.einsum('kij,kji->', vk, dm).real * .5)


# pyscf/pbc/scf/scfint.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors: Garnet Chan <gkc1000@gmail.com>
#          Timothy Berkelbach <tim.berkelbach@gmail.com>
#          Qiming Sun <osirpt.sun@gmail.com>
#

'''
SCF (Hartree-Fock and DFT) tools for periodic systems at a *single* k-point,
    using analytical GTO integrals instead of PWs.

See Also:
    kscf.py : SCF tools for periodic systems with k-point *sampling*.

'''

import numpy as np
import pyscf.pbc.scf
from pyscf.pbc import gto as pgto

print('This module is deporacted and will be removed in future release.  '
      'Please use cell.pbc_intor and pbc.hf.get_hcore.')

def get_hcore(cell, kpt=None):
    '''Get the core Hamiltonian AO matrix, following :func:`dft.rks.get_veff_`.'''
    if kpt is None:
        kpt = np.zeros(3)

    # TODO: these are still on grid
    if cell.pseudo is None:
        hcore = pyscf.pbc.scf.hf.get_nuc(cell, kpt)
    else:
        hcore = pyscf.pbc.scf.hf.get_pp(cell, kpt)
    hcore += get_t(cell, kpt)

    return hcore

def get_int1e_cross(intor, cell1, cell2, kpt=None, comp=1):
    r'''1-electron integrals from two molecules like

    .. math::

        \langle \mu | intor | \nu \rangle, \mu \in cell1, \nu \in cell2
    '''
    return pgto.intor_cross(intor, cell1, cell2, comp, 0, kpt)

def get_int1e(intor, cell, kpt=None):
    '''Get the one-electron integral defined by `intor` using lattice sums.'''
    return cell.pbc_intor(intor, kpts=kpt)

def get_ovlp(cell, kpt=None):
    '''Get the overlap AO matrix.'''
    return cell.pbc_intor('cint1e_ovlp_sph', hermi=1, kpts=kpt)

def get_t(cell, kpt=None):
    '''Get the kinetic energy AO matrix.'''
    return cell.pbc_intor('cint1e_kin_sph', hermi=1, kpts=kpt)



# pyscf/pbc/scf/stability.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''Wave Function Stability Analysis

Ref.
JCP 66, 3045 (1977); DOI:10.1063/1.434318
JCP 104, 9047 (1996); DOI:10.1063/1.471637

See also tddft/rhf.py and scf/newton_ah.py
'''

from functools import reduce
import numpy
import scipy.linalg
from pyscf import lib
from pyscf.lib import logger
from pyscf.pbc.scf import newton_ah
from pyscf.pbc.scf import _response_functions  # noqa

def rhf_stability(mf, internal=True, external=False, verbose=None):
    mo_i = mo_e = None
    if internal:
        mo_i = rhf_internal(mf, verbose=verbose)
    if external:
        mo_e = rhf_external(mf, verbose=verbose)
    return mo_i, mo_e

def uhf_stability(mf, internal=True, external=False, verbose=None):
    mo_i = mo_e = None
    if internal:
        mo_i = uhf_internal(mf, verbose=verbose)
    if external:
        mo_e = uhf_external(mf, verbose=verbose)
    return mo_i, mo_e

def rhf_internal(mf, verbose=None):
    log = logger.new_logger(mf, verbose)
    g, hop, hdiag = newton_ah.gen_g_hop_rhf(mf, mf.mo_coeff, mf.mo_occ)
    def precond(dx, e, x0):
        hdiagd = hdiag*2 - e
        hdiagd[abs(hdiagd)<1e-8] = 1e-8
        return dx/hdiagd
    # The results of hop(x) corresponds to a displacement that reduces
    # gradients g.  It is the vir-occ block of the matrix vector product
    # (Hessian*x). The occ-vir block equals to x2.T.conj(). The overall
    # Hessian for internal reotation is x2 + x2.T.conj(). This is
    # the reason we apply (.real * 2) below
    def hessian_x(x):
        return hop(x).real * 2

    x0 = numpy.zeros_like(g)
    x0[g!=0] = 1. / hdiag[g!=0]
    e, v = lib.davidson(hessian_x, x0, precond, tol=1e-4, verbose=log)
    if e < -1e-5:
        log.log('KRHF/KRKS wavefunction has an internal instability')
        mo = _rotate_mo(mf.mo_coeff, mf.mo_occ, v)
    else:
        log.log('KRHF/KRKS wavefunction is stable in the internal stability analysis')
        mo = mf.mo_coeff
    return mo

def _rotate_mo(mo_coeff, mo_occ, dx):
    mo = []
    p1 = 0
    dtype = numpy.result_type(dx, *mo_coeff)
    for k, occ in enumerate(mo_occ):
        nmo = occ.size
        no = numpy.count_nonzero(occ > 0)
        nv = nmo - no
        p0, p1 = p1, p1 + nv * no
        dr = numpy.zeros((nmo,nmo), dtype=dtype)
        dr[no:,:no] = dx[p0:p1].reshape(nv,no)
        dr[:no,no:] =-dx[p0:p1].reshape(nv,no).conj().T
        mo.append(numpy.dot(mo_coeff[k], scipy.linalg.expm(dr)))
    return mo

#FIXME: numerically unstable with small mesh?
#TODO: Add a warning message for small mesh.
def _gen_hop_rhf_external(mf, verbose=None):
    from pyscf.pbc.scf.newton_ah import _unpack
    cell = mf.cell

    mo_coeff = mf.mo_coeff
    mo_occ = mf.mo_occ
    nkpts = len(mo_occ)
    occidx = [numpy.where(mo_occ[k]==2)[0] for k in range(nkpts)]
    viridx = [numpy.where(mo_occ[k]==0)[0] for k in range(nkpts)]
    orbo = [mo_coeff[k][:,occidx[k]] for k in range(nkpts)]
    orbv = [mo_coeff[k][:,viridx[k]] for k in range(nkpts)]

    h1e = mf.get_hcore()
    dm0 = mf.make_rdm1(mo_coeff, mo_occ)
    fock_ao = h1e + mf.get_veff(cell, dm0)
    fock = [reduce(numpy.dot, (mo_coeff[k].T.conj(), fock_ao[k], mo_coeff[k]))
            for k in range(nkpts)]
    foo = [fock[k][occidx[k][:,None],occidx[k]] for k in range(nkpts)]
    fvv = [fock[k][viridx[k][:,None],viridx[k]] for k in range(nkpts)]

    hdiag = [(fvv[k].diagonal().real[:,None]-foo[k].diagonal().real) * 2
             for k in range(nkpts)]
    hdiag = numpy.hstack([x.ravel() for x in hdiag])

    vresp1 = mf.gen_response(singlet=False, hermi=1)
    def hop_rhf2uhf(x1):
        x1 = _unpack(x1, mo_occ)
        dmvo = []
        for k in range(nkpts):
            # *2 for double occupancy
            dm1 = reduce(numpy.dot, (orbv[k], x1[k]*2, orbo[k].T.conj()))
            dmvo.append(dm1 + dm1.T.conj())
        dmvo = lib.asarray(dmvo)

        v1ao = vresp1(dmvo)
        x2 = [0] * nkpts
        for k in range(nkpts):
            x2[k] = numpy.einsum('ps,sq->pq', fvv[k], x1[k])
            x2[k]-= numpy.einsum('ps,rp->rs', foo[k], x1[k])
            x2[k]+= reduce(numpy.dot, (orbv[k].T.conj(), v1ao[k], orbo[k]))

        # The displacement x2 corresponds to the response of rotation for bra.
        # Hessian*x also provides the rotation for ket which equals to
        # x2.T.conj(). The overall displacement is x2 + x2.T.conj(). This is
        # the reason of x2.real below
        return numpy.hstack([x.real.ravel() for x in x2])

    return hop_rhf2uhf, hdiag


def rhf_external(mf, verbose=None):
    log = logger.new_logger(mf, verbose)
    hop2, hdiag2 = _gen_hop_rhf_external(mf)

    def precond(dx, e, x0):
        hdiagd = hdiag2 - e
        hdiagd[abs(hdiagd)<1e-8] = 1e-8
        return dx/hdiagd
    x0 = numpy.zeros_like(hdiag2)
    x0[hdiag2>1e-5] = 1. / hdiag2[hdiag2>1e-5]
    e3, v3 = lib.davidson(hop2, x0, precond, tol=1e-4, verbose=log)
    if e3 < -1e-5:
        log.log('KRHF/KRKS wavefunction has an KRHF/KRKS -> KUHF/KUKS instability.')
        mo = (_rotate_mo(mf.mo_coeff, mf.mo_occ, v3), mf.mo_coeff)
    else:
        log.log('KRHF/KRKS wavefunction is stable in the KRHF/KRKS -> KUHF/KUKS stability analysis')
        mo = (mf.mo_coeff, mf.mo_coeff)
    return mo

def uhf_internal(mf, verbose=None):
    log = logger.new_logger(mf, verbose)
    g, hop, hdiag = newton_ah.gen_g_hop_uhf(mf, mf.mo_coeff, mf.mo_occ)
    def precond(dx, e, x0):
        hdiagd = hdiag*2 - e
        hdiagd[abs(hdiagd)<1e-8] = 1e-8
        return dx/hdiagd
    def hessian_x(x): # See comments in function rhf_internal
        return hop(x).real * 2

    x0 = numpy.zeros_like(g)
    x0[g!=0] = 1. / hdiag[g!=0]
    e, v = lib.davidson(hessian_x, x0, precond, tol=1e-4, verbose=log)
    if e < -1e-5:
        log.log('KUHF/KUKS wavefunction has an internal instability.')
        tot_x_a = sum((occ>0).sum()*(occ==0).sum() for occ in mf.mo_occ[0])
        mo = (_rotate_mo(mf.mo_coeff[0], mf.mo_occ[0], v[:tot_x_a]),
              _rotate_mo(mf.mo_coeff[1], mf.mo_occ[1], v[tot_x_a:]))
    else:
        log.log('KUHF/KUKS wavefunction is stable in the internal stability analysis')
        mo = mf.mo_coeff
    return mo

def _gen_hop_uhf_external(mf, verbose=None):
    cell = mf.cell

    mo_coeff = mf.mo_coeff
    mo_occ = mf.mo_occ
    nkpts = len(mo_occ[0])
    occidxa = [numpy.where(mo_occ[0][k]>0)[0] for k in range(nkpts)]
    occidxb = [numpy.where(mo_occ[1][k]>0)[0] for k in range(nkpts)]
    viridxa = [numpy.where(mo_occ[0][k]==0)[0] for k in range(nkpts)]
    viridxb = [numpy.where(mo_occ[1][k]==0)[0] for k in range(nkpts)]
    nocca = [len(occidxa[k]) for k in range(nkpts)]
    noccb = [len(occidxb[k]) for k in range(nkpts)]
    nvira = [len(viridxa[k]) for k in range(nkpts)]
    nvirb = [len(viridxb[k]) for k in range(nkpts)]
    moa, mob = mo_coeff
    orboa = [moa[k][:,occidxa[k]] for k in range(nkpts)]
    orbva = [moa[k][:,viridxa[k]] for k in range(nkpts)]
    orbob = [mob[k][:,occidxb[k]] for k in range(nkpts)]
    orbvb = [mob[k][:,viridxb[k]] for k in range(nkpts)]

    h1e = mf.get_hcore()
    dm0 = mf.make_rdm1(mo_coeff, mo_occ)
    fock_ao = h1e + mf.get_veff(cell, dm0)
    focka = [reduce(numpy.dot, (moa[k].T.conj(), fock_ao[0][k], moa[k]))
             for k in range(nkpts)]
    fockb = [reduce(numpy.dot, (mob[k].T.conj(), fock_ao[1][k], mob[k]))
             for k in range(nkpts)]
    fooa = [focka[k][occidxa[k][:,None],occidxa[k]] for k in range(nkpts)]
    fvva = [focka[k][viridxa[k][:,None],viridxa[k]] for k in range(nkpts)]
    foob = [fockb[k][occidxb[k][:,None],occidxb[k]] for k in range(nkpts)]
    fvvb = [fockb[k][viridxb[k][:,None],viridxb[k]] for k in range(nkpts)]

    hdiagab = [fvva[k].diagonal().real[:,None] - foob[k].diagonal().real for k in range(nkpts)]
    hdiagba = [fvvb[k].diagonal().real[:,None] - fooa[k].diagonal().real for k in range(nkpts)]
    hdiag2 = numpy.hstack([x.ravel() for x in (hdiagab + hdiagba)])

    vresp1 = mf.gen_response(with_j=False, hermi=0)
    def hop_uhf2ghf(x1):
        x1ab = []
        x1ba = []
        ip = 0
        for k in range(nkpts):
            nv = nvira[k]
            no = noccb[k]
            x1ab.append(x1[ip:ip+nv*no].reshape(nv,no))
            ip += nv * no
        for k in range(nkpts):
            nv = nvirb[k]
            no = nocca[k]
            x1ba.append(x1[ip:ip+nv*no].reshape(nv,no))
            ip += nv * no

        dm1ab = []
        dm1ba = []
        for k in range(nkpts):
            d1ab = reduce(numpy.dot, (orbva[k], x1ab[k], orbob[k].T.conj()))
            d1ba = reduce(numpy.dot, (orbvb[k], x1ba[k], orboa[k].T.conj()))
            dm1ab.append(d1ab+d1ba.T.conj())
            dm1ba.append(d1ba+d1ab.T.conj())

        v1ao = vresp1(lib.asarray([dm1ab,dm1ba]))
        x2ab = [0] * nkpts
        x2ba = [0] * nkpts
        for k in range(nkpts):
            x2ab[k] = numpy.einsum('pr,rq->pq', fvva[k], x1ab[k])
            x2ab[k]-= numpy.einsum('sq,ps->pq', foob[k], x1ab[k])
            x2ba[k] = numpy.einsum('pr,rq->pq', fvvb[k], x1ba[k])
            x2ba[k]-= numpy.einsum('qs,ps->pq', fooa[k], x1ba[k])
            x2ab[k] += reduce(numpy.dot, (orbva[k].T.conj(), v1ao[0][k], orbob[k]))
            x2ba[k] += reduce(numpy.dot, (orbvb[k].T.conj(), v1ao[1][k], orboa[k]))
        return numpy.hstack([x.real.ravel() for x in (x2ab+x2ba)])

    return hop_uhf2ghf, hdiag2


def uhf_external(mf, verbose=None):
    log = logger.new_logger(mf, verbose)
    hop2, hdiag2 = _gen_hop_uhf_external(mf)

    def precond(dx, e, x0):
        hdiagd = hdiag2 - e
        hdiagd[abs(hdiagd)<1e-8] = 1e-8
        return dx/hdiagd
    x0 = numpy.zeros_like(hdiag2)
    x0[hdiag2>1e-5] = 1. / hdiag2[hdiag2>1e-5]
    e3, v = lib.davidson(hop2, x0, precond, tol=1e-4, verbose=log)
    log.debug('uhf_external: lowest eigs of H = %s', e3)
    mo = None
    if e3 < -1e-5:
        log.log('KUHF/KUKS wavefunction has an KUHF/KUKS -> KGHF/KGKS instability.')
    else:
        log.log('KUHF/KUKS wavefunction is stable in the KUHF/KUKS -> KGHF/KGKS stability analysis')
    return mo


if __name__ == '__main__':
    from pyscf.pbc import gto
    from pyscf.pbc import scf, dft
    from pyscf.pbc import df
    cell = gto.Cell()
    cell.unit = 'B'
    cell.atom = '''
    C  0.          0.          0.
    C  1.68506879  1.68506879  1.68506879
    '''
    cell.a = '''
    0.          3.37013758  3.37013758
    3.37013758  0.          3.37013758
    3.37013758  3.37013758  0.
    '''

    cell.basis = 'gth-szv'
    cell.pseudo = 'gth-pade'
    cell.mesh = [25]*3
    cell.build()
    kpts = cell.make_kpts([2,1,1])
    mf = scf.KRHF(cell, kpts[1:]).set(exxdiv=None).run()
    #mf.with_df = df.DF(cell, kpts)
    #mf.with_df.auxbasis = 'weigend'
    #mf.with_df._cderi = 'eri3d-df.h5'
    #mf.with_df.build(with_j3c=False)
    rhf_stability(mf, True, True, verbose=5)

    mf = scf.KUHF(cell, kpts).set(exxdiv=None).run()
    uhf_stability(mf, True, True, verbose=5)

    mf = dft.KRKS(cell, kpts).set(xc='bp86').run()
    rhf_stability(mf, True, True, verbose=5)

    mf = dft.KUKS(cell, kpts).set(xc='bp86').run()
    uhf_stability(mf, True, True, verbose=5)


# pyscf/pbc/scf/uhf.py

#!/usr/bin/env python
# Copyright 2014-2019 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

'''
Unrestricted Hartree-Fock for periodic systems at a single k-point

See Also:
    pyscf/pbc/scf/khf.py : Hartree-Fock for periodic systems with k-point sampling
'''

import numpy as np
from pyscf import lib
from pyscf.lib import logger
from pyscf.scf import uhf as mol_uhf
from pyscf.pbc.scf import hf as pbchf
from pyscf.pbc.scf import addons
from pyscf.pbc.scf import chkfile
from pyscf import __config__


def init_guess_by_chkfile(cell, chkfile_name, project=None, kpt=None):
    '''Read the HF results from checkpoint file and make the density matrix
    for UHF initial guess.

    Returns:
        Density matrix, (nao,nao) ndarray
    '''
    from pyscf import gto
    chk_cell, scf_rec = chkfile.load_scf(chkfile_name)
    if project is None:
        project = not gto.same_basis_set(chk_cell, cell)

    mo = scf_rec['mo_coeff']
    mo_occ = scf_rec['mo_occ']
    if kpt is None:
        kpt = np.zeros(3)
    if 'kpt' in scf_rec:
        chk_kpt = scf_rec['kpt']
    elif 'kpts' in scf_rec:
        kpts = scf_rec['kpts'] # the closest kpt from KRHF results
        where = np.argmin(lib.norm(kpts-kpt, axis=1))
        chk_kpt = kpts[where]
        if getattr(mo[0], 'ndim', None) == 2:  # KRHF
            mo = mo[where]
            mo_occ = mo_occ[where]
        else:  # KUHF
            mo = [mo[0][where], mo[1][where]]
            mo_occ = [mo_occ[0][where], mo_occ[1][where]]
    else:  # from molecular code
        chk_kpt = np.zeros(3)

    if project:
        s = cell.pbc_intor('int1e_ovlp', kpt=kpt)
    def fproj(mo):
        if project:
            mo = addons.project_mo_nr2nr(chk_cell, mo, cell, chk_kpt-kpt)
            norm = np.einsum('pi,pi->i', mo.conj(), s.dot(mo))
            mo /= np.sqrt(norm)
        return mo

    if getattr(mo, 'ndim', None) == 2:
        mo = fproj(mo)
        mo_occa = (mo_occ>1e-8).astype(np.double)
        mo_occb = mo_occ - mo_occa
        dm = mol_uhf.make_rdm1([mo,mo], [mo_occa,mo_occb])
    else:  # UHF
        dm = mol_uhf.make_rdm1([fproj(mo[0]),fproj(mo[1])], mo_occ)

    # Real DM for gamma point
    if kpt is None or np.allclose(kpt, 0):
        dm = dm.real
    return dm


def dip_moment(cell, dm, unit='Debye', verbose=logger.NOTE,
               grids=None, rho=None, kpt=np.zeros(3)):
    ''' Dipole moment in the unit cell.

    Args:
         cell : an instance of :class:`Cell`

         dm_kpts (a list of ndarrays) : density matrices of k-points

    Return:
        A list: the dipole moment on x, y and z components
    '''
    dm = dm[0] + dm[1]
    return pbchf.dip_moment(cell, dm, unit, verbose, grids, rho, kpt)

get_rho = pbchf.get_rho


class UHF(pbchf.SCF, mol_uhf.UHF):
    '''UHF class for PBCs.
    '''

    direct_scf = getattr(__config__, 'pbc_scf_SCF_direct_scf', False)

    def __init__(self, cell, kpt=np.zeros(3),
                 exxdiv=getattr(__config__, 'pbc_scf_SCF_exxdiv', 'ewald')):
        pbchf.SCF.__init__(self, cell, kpt, exxdiv)
        self.nelec = None
        self.init_guess_breaksym = None
        self._keys = self._keys.union(["init_guess_breaksym"])

    @property
    def nelec(self):
        if self._nelec is not None:
            return self._nelec
        else:
            cell = self.cell
            ne = cell.nelectron
            nalpha = (ne + cell.spin) // 2
            nbeta = nalpha - cell.spin
            if nalpha + nbeta != ne:
                raise RuntimeError('Electron number %d and spin %d are not consistent\n'
                                   'Note cell.spin = 2S = Nalpha - Nbeta, not 2S+1' %
                                   (ne, cell.spin))
            return nalpha, nbeta
    @nelec.setter
    def nelec(self, x):
        self._nelec = x

    def dump_flags(self, verbose=None):
        pbchf.SCF.dump_flags(self, verbose)
        logger.info(self, 'number of electrons per unit cell  '
                    'alpha = %d beta = %d', *self.nelec)
        return self

    get_rho = get_rho

    eig = mol_uhf.UHF.eig

    get_fock = mol_uhf.UHF.get_fock
    get_grad = mol_uhf.UHF.get_grad
    get_occ = mol_uhf.UHF.get_occ
    make_rdm1 = mol_uhf.UHF.make_rdm1
    energy_elec = mol_uhf.UHF.energy_elec

    def get_veff(self, cell=None, dm=None, dm_last=0, vhf_last=0, hermi=1,
                 kpt=None, kpts_band=None):
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        if kpt is None: kpt = self.kpt
        if isinstance(dm, np.ndarray) and dm.ndim == 2:
            dm = np.asarray((dm*.5,dm*.5))
        if self.rsjk and self.direct_scf:
            # Enable direct-SCF for real space JK builder
            ddm = dm - dm_last
            vj, vk = self.get_jk(cell, ddm, hermi, kpt, kpts_band)
            vhf = vj[0] + vj[1] - vk
            vhf += vhf_last
        else:
            vj, vk = self.get_jk(cell, dm, hermi, kpt, kpts_band)
            vhf = vj[0] + vj[1] - vk
        return vhf

    def get_bands(self, kpts_band, cell=None, dm=None, kpt=None):
        '''Get energy bands at the given (arbitrary) 'band' k-points.

        Returns:
            mo_energy : (nmo,) ndarray or a list of (nmo,) ndarray
                Bands energies E_n(k)
            mo_coeff : (nao, nmo) ndarray or a list of (nao,nmo) ndarray
                Band orbitals psi_n(k)
        '''
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        if kpt is None: kpt = self.kpt

        kpts_band = np.asarray(kpts_band)
        single_kpt_band = (getattr(kpts_band, 'ndim', None) == 1)
        kpts_band = kpts_band.reshape(-1,3)

        fock = self.get_hcore(cell, kpts_band)
        focka, fockb = fock + self.get_veff(cell, dm, kpt=kpt, kpts_band=kpts_band)
        s1e = self.get_ovlp(cell, kpts_band)
        nkpts = len(kpts_band)
        e_a = []
        e_b = []
        c_a = []
        c_b = []
        for k in range(nkpts):
            e, c = self.eig((focka[k], fockb[k]), s1e[k])
            e_a.append(e[0])
            e_b.append(e[1])
            c_a.append(c[0])
            c_b.append(c[1])
        mo_energy = (e_a, e_b)
        mo_coeff = (c_a, c_b)

        if single_kpt_band:
            mo_energy = (mo_energy[0][0], mo_energy[1][0])
            mo_coeff = (mo_coeff[0][0], mo_coeff[1][0])
        return mo_energy, mo_coeff

    @lib.with_doc(dip_moment.__doc__)
    def dip_moment(self, cell=None, dm=None, unit='Debye', verbose=logger.NOTE,
                   **kwargs):
        if cell is None: cell = self.cell
        if dm is None: dm = self.make_rdm1()
        rho = kwargs.pop('rho', None)
        if rho is None:
            rho = self.get_rho(dm)
        return dip_moment(cell, dm, unit, verbose, rho=rho, kpt=self.kpt, **kwargs)

    def get_init_guess(self, cell=None, key='minao'):
        if cell is None: cell = self.cell
        dm = mol_uhf.UHF.get_init_guess(self, cell, key)
        dm = pbchf.normalize_dm_(self, dm)
        return dm

    def init_guess_by_1e(self, cell=None):
        if cell is None: cell = self.cell
        if cell.dimension < 3:
            logger.warn(self, 'Hcore initial guess is not recommended in '
                        'the SCF of low-dimensional systems.')
        return mol_uhf.UHF.init_guess_by_1e(self, cell)

    def init_guess_by_chkfile(self, chk=None, project=True, kpt=None):
        if chk is None: chk = self.chkfile
        if kpt is None: kpt = self.kpt
        return init_guess_by_chkfile(self.cell, chk, project, kpt)

    init_guess_by_minao  = mol_uhf.UHF.init_guess_by_minao
    init_guess_by_atom   = mol_uhf.UHF.init_guess_by_atom
    init_guess_by_huckel = mol_uhf.UHF.init_guess_by_huckel

    analyze = mol_uhf.UHF.analyze
    mulliken_pop = mol_uhf.UHF.mulliken_pop
    mulliken_meta = mol_uhf.UHF.mulliken_meta
    spin_square = mol_uhf.UHF.spin_square
    canonicalize = mol_uhf.UHF.canonicalize
    stability = mol_uhf.UHF.stability

    def convert_from_(self, mf):
        '''Convert given mean-field object to RHF/ROHF'''
        addons.convert_to_uhf(mf, self)
        return self


# pyscf/pbc/tddft/__init__.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

from pyscf.pbc.tdscf import rks
from pyscf.pbc.tdscf import rhf


# pyscf/pbc/tdscf/__init__.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

from pyscf.pbc.tdscf import rhf
from pyscf.pbc.tdscf import uhf
from pyscf.pbc.tdscf import krhf
from pyscf.pbc.tdscf import kuhf

try:
    from pyscf.pbc.tdscf import rks
    from pyscf.pbc.tdscf import uks
    from pyscf.pbc.tdscf import krks
    from pyscf.pbc.tdscf import kuks
except (ImportError, IOError):
    pass

def TDHF(mf):
    import numpy
    if numpy.abs(getattr(mf, 'kpt', 0)).max() > 1e-9:
        raise NotImplementedError
    from pyscf.pbc import scf
    if getattr(mf, 'xc', None):
        raise RuntimeError('TDHF does not support DFT object %s' % mf)
    if isinstance(mf, scf.uhf.UHF):
        #mf = scf.addons.convert_to_uhf(mf) # To remove newton decoration
        return uhf.TDHF(mf)
    else:
        #mf = scf.addons.convert_to_rhf(mf)
        return rhf.TDHF(mf)

def TDA(mf):
    import numpy
    if numpy.abs(getattr(mf, 'kpt', 0)).max() > 1e-9:
        raise NotImplementedError
    from pyscf.pbc import scf
    if isinstance(mf, scf.uhf.UHF):
        #mf = scf.addons.convert_to_uhf(mf)
        if getattr(mf, 'xc', None):
            return uks.TDA(mf)
        else:
            return uhf.TDA(mf)
    else:
        #mf = scf.addons.convert_to_rhf(mf)
        if getattr(mf, 'xc', None):
            return rks.TDA(mf)
        else:
            return rhf.TDA(mf)

def TDDFT(mf):
    import numpy
    if numpy.abs(getattr(mf, 'kpt', 0)).max() > 1e-9:
        raise NotImplementedError
    from pyscf.pbc import scf
    if isinstance(mf, scf.uhf.UHF):
        #mf = scf.addons.convert_to_uhf(mf)
        if getattr(mf, 'xc', None):
            return uks.tddft(mf)
        else:
            return uhf.TDHF(mf)
    else:
        #mf = scf.addons.convert_to_rhf(mf)
        if getattr(mf, 'xc', None):
            return rks.tddft(mf)
        else:
            return rhf.TDHF(mf)

def KTDHF(mf):
    from pyscf.pbc import scf
    if getattr(mf, 'xc', None):
        raise RuntimeError('TDHF does not support DFT object %s' % mf)
    if isinstance(mf, scf.uhf.UHF):
        return kuhf.TDHF(mf)
    else:
        return krhf.TDHF(mf)

def KTDA(mf):
    from pyscf.pbc import scf
    if isinstance(mf, scf.uhf.UHF):
        if getattr(mf, 'xc', None):
            return kuks.TDA(mf)
        else:
            return kuhf.TDA(mf)
    else:
        if getattr(mf, 'xc', None):
            return krks.TDA(mf)
        else:
            return krhf.TDA(mf)

def KTDDFT(mf):
    from pyscf.pbc import scf
    if isinstance(mf, scf.uhf.UHF):
        if getattr(mf, 'xc', None):
            return kuks.tddft(mf)
        else:
            return kuhf.TDHF(mf)
    else:
        if getattr(mf, 'xc', None):
            return krks.tddft(mf)
        else:
            return krhf.TDHF(mf)

KTD = KTDDFT



# pyscf/pbc/tdscf/kproxy.py

#  Author: Artem Pulkin
# flake8: noqa
"""
This and other `proxy` modules implement the time-dependent mean-field procedure using the existing pyscf
implementations as a black box. The main purpose of these modules is to overcome the existing limitations in pyscf
(i.e. real-only orbitals, davidson diagonalizer, incomplete Bloch space, etc). The primary performance drawback is that,
unlike the original pyscf routines with an implicit construction of the eigenvalue problem, these modules construct TD
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 TD roots. Several variants of proxy-TD are available:

 * `pyscf.tdscf.proxy`: the molecular implementation;
 * `pyscf.pbc.tdscf.proxy`: PBC (periodic boundary condition) Gamma-point-only implementation;
 * `pyscf.pbc.tdscf.kproxy_supercell`: PBC implementation constructing supercells. Works with an arbitrary number of
   k-points but has an overhead due to ignoring the momentum conservation law. In addition, works only with
   time reversal invariant (TRI) models: i.e. the k-point grid has to be aligned and contain at least one TRI momentum.
 * (this module) `pyscf.pbc.tdscf.kproxy`: same as the above but respect the momentum conservation and, thus, diagonlizes smaller
   matrices (the performance gain is the total number of k-points in the model).
"""

# Convention for these modules:
# * PhysERI is the proxying class constructing time-dependent matrices
# * vector_to_amplitudes reshapes and normalizes the solution
# * TDProxy provides a container

from functools import reduce
from pyscf.pbc.tdscf import kproxy_supercell, krhf_slow

import numpy


def kov2ov(nocc, nmo, k):
    """
    Converts k point pairs into ov mask.
    Args:
        nocc (Iterable): occupation numbers per k-point;
        nmo (Iterable): numbers of orbitals per k-point;
        k (ndarray): k-point pairs;

    Returns:
        An ov-mask. Basis order: [k_o, o, k_v, v].
    """
    nocc = numpy.asanyarray(nocc)
    nmo = numpy.asanyarray(nmo)

    nvirt = nmo - nocc

    mask = numpy.zeros((sum(nocc), sum(nvirt)), dtype=bool)

    o_e = numpy.cumsum(nocc)
    o_s = o_e - o_e[0]

    v_e = numpy.cumsum(nvirt)
    v_s = v_e - v_e[0]

    for k1, k2 in enumerate(k):
        mask[o_s[k1]:o_e[k1], v_s[k2]:v_e[k2]] = True

    return mask.reshape(-1)


class PhysERI(kproxy_supercell.PhysERI):
    def __init__(self, model, proxy, x, mf_constructor, frozen=None, **kwargs):
        """
        A proxy class for calculating the TD matrix blocks (k-point version).

        Args:
            model: the base model with a time reversal-invariant k-point grid;
            proxy: a pyscf proxy with TD response function, one of 'hf', 'dft';
            x (Iterable): the original k-grid dimensions (numbers of k-points per each axis);
            mf_constructor (Callable): a function constructing the mean-field object;
            frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
            **kwargs: arguments to `k2s` function constructing supercells;
        """
        super(PhysERI, self).__init__(model, proxy, x, mf_constructor, frozen=frozen, **kwargs)

    def get_ov_space_mask(self):
        """
        Prepares the space mask in the ov form.
        Returns:
            The mask in the ov form.
        """
        return kproxy_supercell.orb2ov(numpy.concatenate(self.space), self.nocc_full, self.nmo_full)

    def kov2ov(self, k):
        """
        Converts k-ov mask into ov mask.
        Args:
            k (ndarray): k-point pairs;

        Returns:
            An ov-mask. Basis order: [k_o, o, k_v, v].
        """
        mask = self.get_ov_space_mask()
        return numpy.logical_and(mask, kov2ov(self.nocc_full, self.nmo_full, k))

    def proxy_response_ov_batch(self, k_row, k_col):
        """
        A raw response submatrix corresponding to specific k-points.
        Args:
            k_row (ndarray): sets of k-point pairs (row index);
            k_col (ndarray): sets of k-point pairs (column index);

        Returns:
            A raw response matrix.
        """
        masks_row = tuple(self.kov2ov(i) for i in k_row)
        masks_col = tuple(self.kov2ov(i) for i in k_col)

        full_mask_row = reduce(numpy.logical_or, masks_row)
        full_mask_col = reduce(numpy.logical_or, masks_col)

        big = kproxy_supercell.supercell_response_ov(
            self.proxy_vind,
            (full_mask_row, full_mask_col),
            self.nocc_full,
            self.nmo_full,
            self.proxy_is_double(),
            self.model_super.supercell_inv_rotation,
            self.model,
        )

        result = []

        for m_row, m_col in zip(masks_row, masks_col):
            m_row_red = m_row[full_mask_row]
            m_col_red = m_col[full_mask_col]
            result.append(tuple(i[m_row_red][:, m_col_red] for i in big))

        return tuple(result)

    # This is needed for krhf_slow.get_block_k_ix
    get_k_ix = krhf_slow.PhysERI.get_k_ix.im_func

    def tdhf_primary_form(self, k):
        """
        A primary form of TD matrices (full).
        Args:
            k (tuple, int): momentum transfer: either a pair of k-point indexes specifying the momentum transfer
            vector or a single integer with the second index assuming the first index being zero;

        Returns:
            Output type: "full", and the corresponding matrix.
        """
        r1, r2, c1, c2 = krhf_slow.get_block_k_ix(self, k)
        (a, _), (_, b), (_, b_star), (a_star, _) = self.proxy_response_ov_batch((r1, r1, r2, r2), (c1, c2, c1, c2))
        return "full", numpy.block([[a, b], [-b_star.conj(), -a_star.conj()]])


vector_to_amplitudes = krhf_slow.vector_to_amplitudes


class TDProxy(kproxy_supercell.TDProxy):
    v2a = staticmethod(vector_to_amplitudes)
    proxy_eri = PhysERI

    def __init__(self, mf, proxy, x, mf_constructor, frozen=None, **kwargs):
        """
        Performs TD calculation. Roots and eigenvectors are stored in `self.e`, `self.xy`.
        Args:
            mf: the base model with a time-reversal invariant k-point grid;
            proxy: a pyscf proxy with TD response function, one of 'hf', 'dft';
            x (Iterable): the original k-grid dimensions (numbers of k-points per each axis);
            mf_constructor (Callable): a function constructing the mean-field object;
            frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
            **kwargs: arguments to `k2s` function constructing supercells;
        """
        super(TDProxy, self).__init__(mf, proxy, x, mf_constructor, frozen=frozen, **kwargs)
        self.e = {}
        self.xy = {}

    def kernel(self, k=None):
        """
        Calculates eigenstates and eigenvalues of the TDHF problem.
        Args:
            k (tuple, int): momentum transfer: either an index specifying the momentum transfer or a list of such
            indexes;

        Returns:
            Positive eigenvalues and eigenvectors.
        """
        if k is None:
            k = numpy.arange(len(self._scf.kpts))

        if isinstance(k, int):
            k = [k]

        for kk in k:
            self.e[kk], self.xy[kk] = self.__kernel__(k=kk)
        return self.e, self.xy


# pyscf/pbc/tdscf/kproxy_supercell.py

#  Author: Artem Pulkin
# flake8: noqa
"""
This and other `proxy` modules implement the time-dependent mean-field procedure using the existing pyscf
implementations as a black box. The main purpose of these modules is to overcome the existing limitations in pyscf
(i.e. real-only orbitals, davidson diagonalizer, incomplete Bloch space, etc). The primary performance drawback is that,
unlike the original pyscf routines with an implicit construction of the eigenvalue problem, these modules construct TD
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 TD roots. Several variants of proxy-TD are available:

 * `pyscf.tdscf.proxy`: the molecular implementation;
 * `pyscf.pbc.tdscf.proxy`: PBC (periodic boundary condition) Gamma-point-only implementation;
 * (this module) `pyscf.pbc.tdscf.kproxy_supercell`: PBC implementation constructing supercells. Works with an arbitrary number of
   k-points but has an overhead due to ignoring the momentum conservation law. In addition, works only with
   time reversal invariant (TRI) models: i.e. the k-point grid has to be aligned and contain at least one TRI momentum.
 * `pyscf.pbc.tdscf.kproxy`: same as the above but respect the momentum conservation and, thus, diagonlizes smaller
   matrices (the performance gain is the total number of k-points in the model).
"""

# Convention for these modules:
# * PhysERI is the proxying class constructing time-dependent matrices
# * vector_to_amplitudes reshapes and normalizes the solution
# * TDProxy provides a container

from functools import reduce
from pyscf.tdscf.common_slow import TDProxyMatrixBlocks, PeriodicMFMixin
from pyscf.tdscf import proxy as mol_proxy
from pyscf.pbc.tdscf import krhf_slow_supercell, KTDDFT, KTDHF
from pyscf.lib import einsum, cartesian_prod, norm, logger
from pyscf.pbc.tools.pbc import super_cell

import numpy
from scipy import sparse

from warnings import warn


def minus_k(model, threshold=None, degeneracy_threshold=None):
    """
    Retrieves an array of indexes of negative k.
    Args:
        model: a 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-6
    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}, "
                               "threshold: {:.3e}). Use the 'threshold' keyword to loosen the threshold or revise"
                               "your model".format(
                repr(k), id_k, delta[i], threshold,
            ))
        delta = abs(model.mo_energy[id_k] - model.mo_energy[i]).max()
        if delta > degeneracy_threshold:
            raise RuntimeError("Non-symmetric band structure (time-reversal) at k={} (ID: {:d}) and k={} (ID: {:d}), "
                               "max difference: {:.3e}, threshold: {:.3e}. This prevents composing real-valued "
                               "orbitals. Use the 'degeneracy_threshold' keyword to loosen the threshold or revise"
                               "your model".format(
                                    repr(k), id_k, repr(kpts[i]), i, delta, degeneracy_threshold,
                                ))
        result.append(i)
    return result


def assert_scf_converged(model, threshold=1e-7):
    """
    Tests if scf is converged.
    Args:
        model: a mean-field 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())
    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:
            warn("{:d} vectors at k={:d} are not converged, max difference: {:.3e}; orthogonality error: {:.3e} "
                 "warning threshold: {:.3e}".format(
                    nabove, k, max(delta), delta_o, threshold,
                 ))


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-point model into a supercell with real orbitals.
    Args:
        model: a mean-field pbc model;
        grid_spec (Iterable): integer dimensions of the k-grid in the mean-field model;
        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.
    """
    # This hack works as follows. Provided TRS Hamiltonian
    #    H(k) = H(-k)*,
    # with same real eigenvalues and eigenfunctions related as
    #    psi(k) = c psi(-k)*,
    # c - arbitrary phase, it is easy to construct real (non-Bloch) eigenvectors of the whole Hamiltonian
    #    real1(|k|) = c* psi(k) + psi(-k) = psi(-k)* + psi(-k)
    # and
    #    real2(|k|) = 1.j * (c* psi(k) - psi(-k)) = 1.j* (psi(-k)* - psi(-k)).
    # The coefficient c is determined as
    #    psi(k) * psi(-k) = c psi(-k)* * psi(-k) = c
    if imaginary_threshold is None:
        imaginary_threshold = 1e-7

    mk = minus_k(model, threshold=threshold, degeneracy_threshold=degeneracy_threshold)

    # Fix phases
    ovlp = model.get_ovlp()
    phases = {}
    for k1, k2 in enumerate(mk):
        if k1 <= k2:
            c1 = model.mo_coeff[k1]
            c2 = model.mo_coeff[k2]
            o = ovlp[k1]
            r = reduce(numpy.dot, (c2.T, o, c1))
            delta = abs(abs(r) - numpy.eye(r.shape[0])).max()
            if delta > imaginary_threshold:
                raise RuntimeError("K-points connected by time reversal {:d} and {:d} are not complex conjugate: "
                                   "the difference {:.3e} is larger than the threshold {:.3e}".format(
                                        k1, k2, delta, imaginary_threshold,
                                    ))
            p = numpy.angle(numpy.diag(r))
            if k1 == k2:
                phases[k1] = numpy.exp(- .5j * p)[numpy.newaxis, :]
            else:
                phases[k1] = numpy.exp(- 1.j * p)[numpy.newaxis, :]

    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)
    k_spaces = tuple(numpy.arange(i, j) for i, j in zip(nvecs[:-1], nvecs[1:]))

    for k in range(nk):

        i = k_spaces[k]
        j = k_spaces[mk[k]]

        if k == mk[k]:
            rotation_matrix[i, i] = phases[k]
            inv_rotation_matrix[i, i] = phases[k].conj()

        elif k < mk[k]:
            rotation_matrix[i, i] = .5**.5 * phases[k]
            rotation_matrix[j, i] = .5**.5
            rotation_matrix[i, j] = -1.j * .5**.5 * phases[k]
            rotation_matrix[j, j] = 1.j * .5**.5

            inv_rotation_matrix[i, i] = .5**.5 * phases[k].conj()
            inv_rotation_matrix[j, i] = 1.j * .5**.5 * phases[k].conj()
            inv_rotation_matrix[i, j] = .5**.5
            inv_rotation_matrix[j, j] = -1.j * .5**.5

        else:
            pass

    rotation_matrix = rotation_matrix.tocsc()
    inv_rotation_matrix = inv_rotation_matrix.tocsc()

    moc = sparse_transform(moc, 1, rotation_matrix)
    max_imag = abs(moc.imag).max()
    if max_imag > imaginary_threshold:
        raise RuntimeError("Failed to compose real-valued orbitals: imaginary part is {:.3e}".format(max_imag))
    moc = moc.real

    mok = numpy.concatenate(tuple([i] * len(j) for i, j in enumerate(model.mo_energy)))
    moi = numpy.concatenate(tuple(numpy.arange(len(j)) for j in model.mo_energy))

    order = numpy.argsort(moe)

    moe = moe[order]
    moc = moc[:, order]
    moo = moo[order]
    mok = mok[order]
    moi = moi[order]
    rotation_matrix = rotation_matrix[:, order]
    inv_rotation_matrix = inv_rotation_matrix[order, :]

    result.mo_occ = moo,
    result.mo_energy = moe,
    result.mo_coeff = moc,
    result.supercell_rotation = rotation_matrix
    result.supercell_inv_rotation = inv_rotation_matrix
    result.supercell_orig_k = mok
    result.supercell_orig_i = moi

    assert_scf_converged(result, model.conv_tol ** .5)

    p1 = abs(result.supercell_rotation.dot(result.supercell_inv_rotation) - numpy.eye(rotation_matrix.shape[0])).max()
    p2 = abs(result.supercell_inv_rotation.dot(result.supercell_rotation) - numpy.eye(rotation_matrix.shape[0])).max()
    if p1 > 1e-14 or p2 > 1e-14:
        raise RuntimeError("Rotation matrix error: {:.3e}, {:.3e}".format(p1, p2))

    return result


def ko_mask(nocc, nmo):
    """
    Prepares a mask of an occupied space.
    Args:
        nocc (Iterable): occupation numbers per k-point;
        nmo (Iterable): numbers of orbitals per k-point;

    Returns:
        The mask where `True` denotes occupied orbitals. Basis order: [k, orb=o+v]
    """
    result = numpy.zeros(sum(nmo), dtype=bool)
    offset = 0
    for no, nm in zip(nocc, nmo):
        result[offset:offset+no] = True
        offset += nm
    return result


def split_transform(transform, nocc, nmo, tolerance=1e-14):
    """
    Splits the transform into `oo` and `vv` blocks.
    Args:
        transform (numpy.ndarray): the original transform. The basis order for the transform is [real orb=o+v; k, orb=o+v];
        nocc (Iterable): occupation numbers per k-point;
        nmo (Iterable): the number of orbitals per k-point;
        tolerance (float): tolerance to check zeros at the `ov` block;

    Returns:
        `oo` and `vv` blocks of the transform.
    """
    o_mask = ko_mask(nocc, nmo)
    v_mask = ~o_mask
    ov = transform[:sum(nocc), v_mask]
    vo = transform[sum(nocc):, o_mask]
    if abs(ov).max() > tolerance or abs(vo).max() > tolerance:
        raise ValueError("Occupied and virtual spaces are coupled by the transformation")
    return transform[:sum(nocc), o_mask], transform[sum(nocc):, v_mask]


def supercell_space_required(transform_oo, transform_vv, final_space):
    """
    For a given orbital transformation and a given `ov` mask in the transformed space, calculates a minimal `ov` mask
    in the original space required to achieve this transform.
    Args:
        transform_oo (ndarray): the transformation in the occupied space;
        transform_vv (ndarray): the transformation in the virtual space;
        final_space (ndarray): the final `ov` space. Basis order: [k_o, o, k_v, v];

    Returns:
        The initial active space. Basis order: [k_o, o, k_v, v].
    """
    final_space = numpy.asanyarray(final_space)
    final_space = final_space.reshape(final_space.shape[:-1] + (transform_oo.shape[1], transform_vv.shape[1]))
    result = einsum(
        "ao,bv,...ov->...ab",
        (transform_oo.toarray() != 0).astype(int),
        (transform_vv.toarray() != 0).astype(int),
        final_space.astype(int),
    ) != 0
    return result.reshape(result.shape[:-2] + (-1,))


def get_sparse_ov_transform(oo, vv):
    """
    Retrieves a sparse `ovov` transform out of sparse `oo` and `vv` transforms.
    Args:
        oo (ndarray): the transformation in the occupied space;
        vv (ndarray): the transformation in the virtual space;

    Returns:
        The resulting matrix representing the sparse transform in the `ov` space.
    """
    i, a = oo.shape
    j, b = vv.shape

    # If the input is dense the result is simply
    # return (oo[:, numpy.newaxis, :, numpy.newaxis] * vv[numpy.newaxis, :, numpy.newaxis, :]).reshape(i*j, a*b)

    result_data = numpy.zeros(oo.nnz * vv.nnz, dtype=numpy.common_type(oo, vv))
    result_indices = numpy.zeros(len(result_data), dtype=int)
    result_indptr = numpy.zeros(a * b + 1, dtype=int)

    ptr_counter = 0
    for i_a in range(a):
        oo_col = oo.getcol(i_a)
        assert tuple(oo_col.indptr.tolist()) == (0, len(oo_col.data))
        i_i, oo_col_v = oo_col.indices, oo_col.data
        for i_b in range(b):
            vv_col = vv.getcol(i_b)
            assert tuple(vv_col.indptr.tolist()) == (0, len(vv_col.data))
            i_j, vv_col_v = vv_col.indices, vv_col.data

            data_length = len(i_i) * len(i_j)
            result_indices[ptr_counter:ptr_counter + data_length] = ((i_i * j)[:, numpy.newaxis] + i_j[numpy.newaxis, :]).reshape(-1)
            result_data[ptr_counter:ptr_counter + data_length] = (oo_col_v[:, numpy.newaxis] * vv_col_v[numpy.newaxis, :]).reshape(-1)
            result_indptr[i_a * b + i_b] = ptr_counter

            ptr_counter += data_length

    result_indptr[-1] = ptr_counter
    return sparse.csc_matrix((result_data, result_indices, result_indptr))


def ov2orb(space, nocc, nmo):
    """
    Converts ov-pairs active space specification into orbital space spec.
    Args:
        space (ndarray): the ov space. Basis order: [k_o, o, k_v, v];
        nocc (Iterable): the numbers of occupied orbitals per k-point;
        nmo (Iterable): the total numbers of orbitals per k-point;

    Returns:
        The orbital space specification. Basis order: [k, orb=o+v].
    """
    nocc = numpy.asanyarray(nocc)
    nmo = numpy.asanyarray(nmo)
    nvirt = nmo - nocc

    space = numpy.asanyarray(space)
    space = space.reshape(space.shape[:-1] + (sum(nocc), sum(nvirt)))  # [k_o, o; k_v, v]

    s_o = numpy.any(space, axis=-1)  # [k_o, o]
    s_v = numpy.any(space, axis=-2)  # [k_v, v]

    o_offset = numpy.cumsum(numpy.concatenate(([0], nocc)))
    v_offset = numpy.cumsum(numpy.concatenate(([0], nvirt)))
    result = []
    for o_fr, o_to, v_fr, v_to in zip(o_offset[:-1], o_offset[1:], v_offset[:-1], v_offset[1:]):
        result.append(s_o[..., o_fr:o_to])
        result.append(s_v[..., v_fr:v_to])
    return numpy.concatenate(result, axis=-1)  # [k, orb=o+v]


def supercell_response_ov(vind, space_ov, nocc, nmo, double, rot_bloch, log_dest):
    """
    Retrieves a raw response matrix.
    Args:
        vind (Callable): a pyscf matvec routine;
        space_ov (ndarray): the active `ov` space mask: either the same mask for both rows and columns (1D array) or
        separate `ov` masks for rows and columns (2D array). Basis order: [k_o, o, k_v, v];
        nocc (ndarray): the numbers of occupied orbitals (frozen and active) per k-point;
        nmo (ndarray): the total number of orbitals per k-point;
        double (bool): set to True if `vind` returns the double-sized (i.e. full) matrix;
        rot_bloch (ndarray): a matrix specifying the rotation from real orbitals returned from pyscf to Bloch
        functions;
        log_dest (object): pyscf logging;

    Returns:
        The TD matrix.
    """
    if not double:
        raise NotImplementedError("Not implemented for MK-type matrixes")

    nocc_full = sum(nocc)
    nmo_full = sum(nmo)
    nvirt_full = nmo_full - nocc_full
    size_full = nocc_full * nvirt_full

    space_ov = numpy.array(space_ov)

    if space_ov.shape == (size_full,):
        space_ov = numpy.repeat(space_ov[numpy.newaxis, :], 2, axis=0)

    elif space_ov.shape != (2, size_full):
        raise ValueError(
            "The 'space_ov' argument should be a 1D array with dimension {size_full:d} or a 2D array with"
            " dimensions 2x{size_full:d}, found: {actual}".format(
                size_full=size_full,
                actual=space_ov.shape,
            ))

    oo, vv = split_transform(rot_bloch, nocc, nmo)
    space_real_ov = supercell_space_required(oo, vv, space_ov)

    logger.debug1(log_dest, "Performing a supercell proxy response calculation ...")
    logger.debug1(log_dest, "  Total ov space size: {:d} requested elements: {} real elements to calculate: {}".format(
        size_full,
        "x".join(map(str, space_ov.sum(axis=-1))),
        "x".join(map(str, space_real_ov.sum(axis=-1))),
    ))

    logger.debug1(log_dest, "  collecting the A, B matrices ...")
    response_real_a, response_real_b = mol_proxy.molecular_response_ov(
        vind,
        space_real_ov,
        nocc_full,
        nmo_full,
        double,
        log_dest,
    )
    logger.debug1(log_dest, "  done, shapes: {} and {}".format(
        response_real_a.shape,
        response_real_b.shape,
    ))

    logger.debug1(log_dest, "Transforming into Bloch basis ...")

    ovov_nc = get_sparse_ov_transform(oo, vv.conj())
    ovov_cn = get_sparse_ov_transform(oo.conj(), vv)

    ovov_row = ovov_nc[:, space_ov[0]][space_real_ov[0]]
    ovov_col_a = ovov_cn[:, space_ov[1]][space_real_ov[1]]
    ovov_col_b = ovov_nc[:, space_ov[1]][space_real_ov[1]]

    # Rotate
    logger.debug1(log_dest, "  rotating A ...")
    response_bloch_a = sparse_transform(response_real_a, 0, ovov_row, 1, ovov_col_a)
    logger.debug1(log_dest, "  rotating B ...")
    response_bloch_b = sparse_transform(response_real_b, 0, ovov_row, 1, ovov_col_b)
    logger.debug1(log_dest, "  shapes: {} and {}".format(response_bloch_a.shape, response_bloch_b.shape))

    return response_bloch_a, response_bloch_b


def orb2ov(space, nocc, nmo):
    """
    Converts orbital active space specification into ov-pairs space spec.
    Args:
        space (ndarray): the obital space. Basis order: [k, orb=o+v];
        nocc (Iterable): the numbers of occupied orbitals per k-point;
        nmo (Iterable): the total numbers of orbitals per k-point;

    Returns:
        The ov space specification. Basis order: [k_o, o, k_v, v].
    """
    space = numpy.asanyarray(space)
    m = ko_mask(nocc, nmo)  # [k, orb=o+v]
    o = space[...,  m]  # [k, o]
    v = space[..., ~m]  # [k, v]
    return (o[..., numpy.newaxis] * v[..., numpy.newaxis, :]).reshape(space.shape[:-1] + (-1,))  # [k_o, o, k_v, v]


def supercell_response(vind, space, nocc, nmo, double, rot_bloch, log_dest):
    """
    Retrieves a raw response matrix.
    Args:
        vind (Callable): a pyscf matvec routine;
        space (ndarray): the active space: either for both rows and columns (1D array) or for rows and columns
        separately (2D array). Basis order: [k, orb=o+v];
        nocc (ndarray): the numbers of occupied orbitals (frozen and active) per k-point;
        nmo (ndarray): the total number of orbitals per k-point;
        double (bool): set to True if `vind` returns the double-sized (i.e. full) matrix;
        rot_bloch (ndarray): a matrix specifying the rotation from real orbitals returned from pyscf to Bloch
        functions;
        log_dest (object): pyscf logging;

    Returns:
        The TD matrix.
    """
    if not double:
        raise NotImplementedError("Not implemented for MK-type matrixes")

    # Full space dims
    nmo_full = sum(nmo)
    space = numpy.array(space)

    if space.shape == (nmo_full,):
        space = numpy.repeat(space[numpy.newaxis, :], 2, axis=0)
    elif space.shape != (2, nmo_full):
        raise ValueError("The 'space' argument should a 1D array with dimension {:d} or a 2D array with dimensions {},"
                         " found: {}".format(nmo_full, (2, nmo_full), space.shape))

    return supercell_response_ov(vind, orb2ov(space, nocc, nmo), nocc, nmo, double, rot_bloch, log_dest)


class PhysERI(PeriodicMFMixin, TDProxyMatrixBlocks):
    proxy_choices = {
        "hf": KTDHF,
        "dft": KTDDFT,
    }

    def __init__(self, model, proxy, x, mf_constructor, frozen=None, **kwargs):
        """
        A proxy class for calculating TD matrix blocks (supercell version).

        Args:
            model: the base model with a time reversal-invariant k-point grid;
            proxy: a pyscf proxy with TD response function, one of 'hf', 'dft';
            x (Iterable): the original k-grid dimensions (numbers of k-points per each axis);
            mf_constructor (Callable): a function constructing the mean-field object;
            frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
            **kwargs: arguments to `k2s` function constructing supercells;
        """
        model_super = k2s(model, x, mf_constructor, **kwargs)
        TDProxyMatrixBlocks.__init__(self, self.proxy_choices[proxy](model_super))
        PeriodicMFMixin.__init__(self, model, frozen=frozen)
        self.model_super = model_super

    def proxy_is_double(self):
        """
        Determines if double-sized matrices are proxied.
        Returns:
            True if double-sized matrices are proxied.
        """
        nocc_full = sum(self.nocc_full)
        nmo_full = sum(self.nmo_full)
        size_full = nocc_full * (nmo_full - nocc_full)
        size_hdiag = len(self.proxy_diag)

        if size_full == size_hdiag:
            return False

        elif 2 * size_full == size_hdiag:
            return True

        else:
            raise RuntimeError("Do not recognize the size of TD diagonal: {:d}. The size of ov-space is {:d}".format(
                size_hdiag, size_full
            ))

    def proxy_response(self):
        """
        A raw response matrix.
        Returns:
            A raw response matrix.
        """
        return supercell_response(
            self.proxy_vind,
            numpy.concatenate(self.space),
            self.nocc_full,
            self.nmo_full,
            self.proxy_is_double(),
            self.model_super.supercell_inv_rotation,
            self.model,
        )

    def tdhf_primary_form(self, *args, **kwargs):
        """
        A primary form of TD matrixes.

        Returns:
            Output type: "full", "ab", or "mk" and the corresponding matrix(es).
        """
        a, b = self.proxy_response()

        # Transform into supercell convention: [k_o, o, k_v, v] -> [k_o, k_v, o, v]
        nocc_k = self.nocc[0]
        nk = len(self.nocc)
        nvirt_k = self.nmo[0] - self.nocc[0]
        size = nk * nk * nocc_k * nvirt_k

        a = a.reshape((nk, nocc_k, nk, nvirt_k) * 2).transpose(0, 2, 1, 3, 4, 6, 5, 7)
        b = b.reshape((nk, nocc_k, nk, nvirt_k) * 2).transpose(0, 2, 1, 3, 4, 6, 5, 7)
        return "ab", a.reshape(size, size), b.reshape(size, size)


vector_to_amplitudes = krhf_slow_supercell.vector_to_amplitudes


class TDProxy(mol_proxy.TDProxy):
    v2a = staticmethod(vector_to_amplitudes)
    proxy_eri = PhysERI

    def __init__(self, mf, proxy, x, mf_constructor, frozen=None, **kwargs):
        """
        Performs TD calculation. Roots and eigenvectors are stored in `self.e`, `self.xy`.
        Args:
            mf: the base model with a time reversal-invariant k-point grid;
            proxy: a pyscf proxy with TD response function, one of 'hf', 'dft';
            x (Iterable): the original k-grid dimensions (numbers of k-points per each axis);
            mf_constructor (Callable): a function constructing the mean-field object;
            frozen (int, Iterable): the number of frozen valence orbitals or the list of frozen orbitals;
            **kwargs: arguments to `k2s` function constructing supercells;
        """
        super(TDProxy, self).__init__(mf, proxy, frozen=frozen)
        self.fast = False
        self.x = x
        self.mf_constructor = mf_constructor
        self.__k2s_kwargs__ = kwargs

    def ao2mo(self):
        """
        Prepares ERI.

        Returns:
            A suitable ERI.
        """
        return self.proxy_eri(
            self._scf,
            self.__proxy__,
            x=self.x,
            mf_constructor=self.mf_constructor,
            frozen=self.frozen,
            **self.__k2s_kwargs__
        )



# pyscf/pbc/tdscf/krhf.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#
# Ref:
# Chem Phys Lett, 256, 454
# J. Mol. Struct. THEOCHEM, 914, 3
#

from functools import reduce
import numpy
from pyscf import lib
from pyscf.lib import linalg_helper
from pyscf.lib import logger
from pyscf.tdscf import rhf
from pyscf.pbc.scf import _response_functions  # noqa
from pyscf.pbc.lib.kpts_helper import gamma_point
from pyscf import __config__

REAL_EIG_THRESHOLD = getattr(__config__, 'pbc_tdscf_rhf_TDDFT_pick_eig_threshold', 1e-3)
POSTIVE_EIG_THRESHOLD = getattr(__config__, 'pbc_tdscf_rhf_TDDFT_positive_eig_threshold', 1e-3)

class TDA(rhf.TDA):
    conv_tol = getattr(__config__, 'pbc_tdscf_rhf_TDA_conv_tol', 1e-6)

    def __init__(self, mf):
        from pyscf.pbc import scf
        assert(isinstance(mf, scf.khf.KSCF))
        self.cell = mf.cell
        rhf.TDA.__init__(self, mf)
        from pyscf.pbc.df.df_ao2mo import warn_pbc2d_eri
        warn_pbc2d_eri(mf)

    def gen_vind(self, mf):
        singlet = self.singlet

        mo_coeff = mf.mo_coeff
        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        nkpts = len(mo_occ)
        nao, nmo = mo_coeff[0].shape
        occidx = [numpy.where(mo_occ[k]==2)[0] for k in range(nkpts)]
        viridx = [numpy.where(mo_occ[k]==0)[0] for k in range(nkpts)]
        orbo = [mo_coeff[k][:,occidx[k]] for k in range(nkpts)]
        orbv = [mo_coeff[k][:,viridx[k]] for k in range(nkpts)]
        e_ia = _get_e_ia(mo_energy, mo_occ)
        hdiag = numpy.hstack([x.ravel() for x in e_ia])

        mem_now = lib.current_memory()[0]
        max_memory = max(2000, self.max_memory*.8-mem_now)
        vresp = mf.gen_response(singlet=singlet, hermi=0, max_memory=max_memory)

        def vind(zs):
            nz = len(zs)
            z1s = [_unpack(z, mo_occ) for z in zs]
            dmov = numpy.empty((nz,nkpts,nao,nao), dtype=numpy.complex128)
            for i in range(nz):
                for k in range(nkpts):
                    # *2 for double occupancy
                    dm1 = z1s[i][k] * 2
                    dmov[i,k] = reduce(numpy.dot, (orbo[k], dm1, orbv[k].conj().T))

            with lib.temporary_env(mf, exxdiv=None):
                v1ao = vresp(dmov)
            v1s = []
            for i in range(nz):
                dm1 = z1s[i]
                for k in range(nkpts):
                    v1vo = reduce(numpy.dot, (orbo[k].conj().T, v1ao[i,k], orbv[k]))
                    v1vo += e_ia[k] * dm1[k]
                    v1s.append(v1vo.ravel())
            return lib.asarray(v1s).reshape(nz,-1)
        return vind, hdiag

    def init_guess(self, mf, nstates=None):
        if nstates is None: nstates = self.nstates

        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        e_ia = numpy.hstack([x.ravel() for x in _get_e_ia(mo_energy, mo_occ)])

        e_ia_max = e_ia.max()
        nov = e_ia.size
        nstates = min(nstates, nov)
        e_threshold = min(e_ia_max, e_ia[numpy.argsort(e_ia)[nstates-1]])
        # Handle degeneracy, include all degenerated states in initial guess
        e_threshold += 1e-6

        idx = numpy.where(e_ia <= e_threshold)[0]
        x0 = numpy.zeros((idx.size, nov))
        for i, j in enumerate(idx):
            x0[i, j] = 1  # Koopmans' excitations
        return x0

    def kernel(self, x0=None):
        '''TDA diagonalization solver
        '''
        self.check_sanity()
        self.dump_flags()

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)

        def pickeig(w, v, nroots, envs):
            idx = numpy.where(w > POSTIVE_EIG_THRESHOLD**2)[0]
            return w[idx], v[:,idx], idx

        log = logger.Logger(self.stdout, self.verbose)
        precision = self.cell.precision * 1e-2

        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)
        self.converged, self.e, x1 = \
                lib.davidson1(vind, x0, precond,
                              tol=self.conv_tol,
                              nroots=self.nstates, lindep=self.lindep,
                              max_space=self.max_space, pick=pickeig,
                              fill_heff=purify_krlyov_heff(precision, 0, log),
                              verbose=self.verbose)

        mo_occ = self._scf.mo_occ
# 1/sqrt(2) because self.x is for alpha excitation amplitude and 2(X^+*X) = 1
        self.xy = [(_unpack(xi*numpy.sqrt(.5), mo_occ), 0) for xi in x1]
        return self.e, self.xy
CIS = KTDA = TDA


class TDHF(TDA):
    def gen_vind(self, mf):
        '''
        [ A   B ][X]
        [-B* -A*][Y]
        '''
        singlet = self.singlet

        mo_coeff = mf.mo_coeff
        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        nkpts = len(mo_occ)
        nao, nmo = mo_coeff[0].shape
        occidx = [numpy.where(mo_occ[k]==2)[0] for k in range(nkpts)]
        viridx = [numpy.where(mo_occ[k]==0)[0] for k in range(nkpts)]
        orbo = [mo_coeff[k][:,occidx[k]] for k in range(nkpts)]
        orbv = [mo_coeff[k][:,viridx[k]] for k in range(nkpts)]
        e_ia = _get_e_ia(mo_energy, mo_occ)
        hdiag = numpy.hstack([x.ravel() for x in e_ia])
        tot_x = hdiag.size
        hdiag = numpy.hstack((hdiag, -hdiag))

        mem_now = lib.current_memory()[0]
        max_memory = max(2000, self.max_memory*.8-mem_now)
        vresp = mf.gen_response(singlet=singlet, hermi=0, max_memory=max_memory)

        def vind(xys):
            nz = len(xys)
            z1xs = [_unpack(xy[:tot_x], mo_occ) for xy in xys]
            z1ys = [_unpack(xy[tot_x:], mo_occ) for xy in xys]
            dmov = numpy.empty((nz,nkpts,nao,nao), dtype=numpy.complex128)
            for i in range(nz):
                for k in range(nkpts):
                    # *2 for double occupancy
                    dmx = z1xs[i][k] * 2
                    dmy = z1ys[i][k] * 2
                    dmov[i,k] = reduce(numpy.dot, (orbo[k], dmx, orbv[k].T.conj()))
                    dmov[i,k]+= reduce(numpy.dot, (orbv[k], dmy.T, orbo[k].T.conj()))

            with lib.temporary_env(mf, exxdiv=None):
                v1ao = vresp(dmov)
            v1s = []
            for i in range(nz):
                dmx = z1xs[i]
                dmy = z1ys[i]
                v1xs = []
                v1ys = []
                for k in range(nkpts):
                    v1x = reduce(numpy.dot, (orbo[k].T.conj(), v1ao[i,k], orbv[k]))
                    v1y = reduce(numpy.dot, (orbv[k].T.conj(), v1ao[i,k], orbo[k])).T
                    v1x+= e_ia[k] * dmx[k]
                    v1y+= e_ia[k] * dmy[k]
                    v1xs.append(v1x.ravel())
                    v1ys.append(-v1y.ravel())
                v1s += v1xs + v1ys
            return lib.asarray(v1s).reshape(nz,-1)
        return vind, hdiag

    def init_guess(self, mf, nstates=None):
        x0 = TDA.init_guess(self, mf, nstates)
        y0 = numpy.zeros_like(x0)
        return numpy.hstack((x0,y0))

    def kernel(self, x0=None):
        '''TDHF diagonalization with non-Hermitian eigenvalue solver
        '''
        logger.warn(self, 'PBC-TDDFT is an experimental feature. '
                    'It is numerically sensitive to the accuracy of integrals '
                    '(relating to cell.precision).')

        self.check_sanity()
        self.dump_flags()

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)
        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        real_system = (gamma_point(self._scf.kpts) and
                       self._scf.mo_coeff[0].dtype == numpy.double)

        # We only need positive eigenvalues
        def pickeig(w, v, nroots, envs):
            realidx = numpy.where((abs(w.imag) < REAL_EIG_THRESHOLD) &
                                  (w.real > POSTIVE_EIG_THRESHOLD))[0]
            return lib.linalg_helper._eigs_cmplx2real(w, v, realidx, real_system)

        log = logger.Logger(self.stdout, self.verbose)
        precision = self.cell.precision * 1e-2
        hermi = 0

        self.converged, w, x1 = \
                lib.davidson_nosym1(vind, x0, precond,
                                    tol=self.conv_tol,
                                    nroots=self.nstates, lindep=self.lindep,
                                    max_space=self.max_space, pick=pickeig,
                                    fill_heff=purify_krlyov_heff(precision, hermi, log),
                                    verbose=self.verbose)
        mo_occ = self._scf.mo_occ
        self.e = w
        def norm_xy(z):
            x, y = z.reshape(2,-1)
            norm = 2*(lib.norm(x)**2 - lib.norm(y)**2)
            norm = 1/numpy.sqrt(norm)
            x *= norm
            y *= norm
            return _unpack(x, mo_occ), _unpack(y, mo_occ)
        self.xy = [norm_xy(z) for z in x1]

        return self.e, self.xy
RPA = KTDHF = TDHF


def _get_e_ia(mo_energy, mo_occ):
    e_ia = []
    for k, occ in enumerate(mo_occ):
        occidx = occ >  0
        viridx = occ == 0
        e_ia.append(mo_energy[k][viridx] - mo_energy[k][occidx,None])
    return e_ia

def _unpack(vo, mo_occ):
    z = []
    p1 = 0
    for k, occ in enumerate(mo_occ):
        no = numpy.count_nonzero(occ > 0)
        nv = occ.size - no
        p0, p1 = p1, p1 + no * nv
        z.append(vo[p0:p1].reshape(no,nv))
    return z

def purify_krlyov_heff(precision, hermi, log):
    def fill_heff(heff, xs, ax, xt, axt, dot):
        if hermi == 1:
            heff = linalg_helper._fill_heff_hermitian(heff, xs, ax, xt, axt, dot)
        else:
            heff = linalg_helper._fill_heff(heff, xs, ax, xt, axt, dot)
        space = len(axt)
        # TODO: PBC integrals has larger errors than molecule systems.
        # purify the effective Hamiltonian with symmetry and other
        # possible conditions.
        if abs(heff[:space,:space].imag).max() < precision:
            log.debug('Remove imaginary part of the Krylov space effective Hamiltonian')
            heff[:space,:space].imag = 0
        return heff
    return fill_heff


from pyscf.pbc import scf
scf.khf.KRHF.TDA  = lib.class_as_method(KTDA)
scf.khf.KRHF.TDHF = lib.class_as_method(KTDHF)
scf.krohf.KROHF.TDA  = None
scf.krohf.KROHF.TDHF = None


if __name__ == '__main__':
    from pyscf.pbc import gto
    from pyscf.pbc import scf
    from pyscf.pbc import df
    cell = gto.Cell()
    cell.unit = 'B'
    cell.atom = '''
    C  0.          0.          0.
    C  1.68506879  1.68506879  1.68506879
    '''
    cell.a = '''
    0.          3.37013758  3.37013758
    3.37013758  0.          3.37013758
    3.37013758  3.37013758  0.
    '''

    cell.basis = 'gth-szv'
    cell.pseudo = 'gth-pade'
    cell.mesh = [25]*3
    cell.build()
    mf = scf.KRHF(cell, cell.make_kpts([2,1,1])).set(exxdiv=None)
    #mf.with_df = df.MDF(cell, cell.make_kpts([2,1,1]))
    #mf.with_df.auxbasis = 'weigend'
    #mf.with_df._cderi = 'eri3d-mdf.h5'
    #mf.with_df.build(with_j3c=False)
    mf.run()
    #mesh=9  -8.65192427146353
    #mesh=12 -8.65192352289817
    #mesh=15 -8.6519235231529
    #MDF mesh=5 -8.6519301815144

    td = TDA(mf)
    td.verbose = 5
    print(td.kernel()[0] * 27.2114)
    #mesh=9  [ 6.0073749   6.09315355  6.3479901 ]
    #mesh=12 [ 6.00253282  6.09317929  6.34799109]
    #mesh=15 [ 6.00253396  6.09317949  6.34799109]
    #MDF mesh=5 [ 6.09317489  6.09318265  6.34798637]

    #from pyscf.pbc import tools
    #scell = tools.super_cell(cell, [2,1,1])
    #mf = scf.RHF(scell).run()
    #td = rhf.TDA(mf)
    #td.verbose = 5
    #print(td.kernel()[0] * 27.2114)

    td = TDHF(mf)
    td.verbose = 5
    print(td.kernel()[0] * 27.2114)
    #mesh=9  [ 6.03860914  6.21664545  8.20305225]
    #mesh=12 [ 6.03868259  6.03860343  6.2167623 ]
    #mesh=15 [ 6.03861321  6.03861324  6.21675868]
    #MDF mesh=5 [ 6.03861693  6.03861775  6.21675694]



# pyscf/pbc/tdscf/krhf_slow.py

#  Author: Artem Pulkin
# flake8: noqa
"""
This and other `_slow` modules implement the time-dependent Hartree-Fock 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 via an AO-MO transformation, i.e. with a O(N^5) complexity scaling. 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. Several variants of TDHF are available:

 * `pyscf.tdscf.rhf_slow`: the molecular implementation;
 * `pyscf.pbc.tdscf.rhf_slow`: PBC (periodic boundary condition) implementation for RHF objects of `pyscf.pbc.scf`
   modules;
 * `pyscf.pbc.tdscf.krhf_slow_supercell`: PBC implementation for KRHF 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.krhf_slow_gamma`: A Gamma-point calculation resembling the original `pyscf.pbc.tdscf.krhf`
   module. Despite its name, it accepts KRHF objects with an arbitrary number of k-points but finds only few TDHF roots
   corresponding to collective oscillations without momentum transfer;
 * (this module) `pyscf.pbc.tdscf.krhf_slow`: PBC implementation for KRHF objects of `pyscf.pbc.scf` modules. Works with
   an arbitrary number of k-points and employs k-point conservation (diagonalizes matrix blocks separately).
"""

from pyscf.pbc.tdscf import krhf_slow_supercell as td
from pyscf.tdscf import rhf_slow
from pyscf.tdscf.common_slow import mknj2i

import numpy


# Convention for these modules:
# * PhysERI, PhysERI4, PhysERI8 are 2-electron integral routines computed directly (for debug purposes), with a 4-fold
#   symmetry and with an 8-fold symmetry
# * vector_to_amplitudes reshapes and normalizes the solution
# * TDRHF provides a container


class PhysERI(td.PhysERI):
    primary_driver = "full"

    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing a full transformation of integrals to Bloch functions. No symmetries are
        employed in this class. The ERIs are returned in blocks of k-points.

        Args:
            model (KRHF): the base model;
            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;
        """
        super(PhysERI, self).__init__(model, frozen=frozen)

    def get_k_ix(self, item, like):
        """
        Retrieves block indexes: row and column.
        Args:
            item (str): a string of 'mknj' letters;
            like (tuple): a 2-tuple with sample pair of k-points;

        Returns:
            Row and column indexes of a sub-block with conserving momentum.
        """
        item_i = numpy.argsort(mknj2i(item))
        item_code = ''.join("++--"[i] for i in item_i)
        if item_code[0] == item_code[1]:
            kc = self.kconserv  # ++-- --++
        elif item_code[0] == item_code[2]:
            kc = self.kconserv.swapaxes(1, 2)  # +-+- -+-+
        elif item_code[1] == item_code[2]:
            kc = self.kconserv.transpose(2, 0, 1)  # +--+ -++-
        else:
            raise RuntimeError("Unknown case: {}".format(item_code))

        y = kc[like]
        x = kc[0, y[0]]

        return x, y

    def tdhf_diag(self, block):
        """
        Retrieves the merged diagonal block only with specific pairs of k-indexes (k, block[k]).
        Args:
            block (Iterable): a k-point pair `k2 = pair[k1]` for each k1;

        Returns:
            The diagonal block.
        """
        return super(PhysERI, self).tdhf_diag(pairs=enumerate(block))

    def eri_mknj(self, item, pair_row, pair_column):
        """
        Retrieves the merged ERI block using 'mknj' notation with pairs of k-indexes (k1, k1, k2, k2).
        Args:
            item (str): a 4-character string of 'mknj' letters;
            pair_row (Iterable): a k-point pair `k2 = pair_row[k1]` for each k1 (row indexes in the final matrix);
            pair_column (Iterable): a k-point pair `k4 = pair_row[k3]` for each k3 (column indexes in the final matrix);

        Returns:
            The corresponding block of ERI (phys notation).
        """
        return super(PhysERI, self).eri_mknj(
            item,
            pairs_row=enumerate(pair_row),
            pairs_column=enumerate(pair_column),
        )

    def tdhf_primary_form(self, k):
        """
        A primary form of TDHF matrixes (full).
        Args:
            k (tuple, int): momentum transfer: either a pair of k-point indexes specifying the momentum transfer
            vector or a single integer with the second index assuming the first index being zero;

        Returns:
            Output type: "full", and the corresponding matrix.
        """
        r1, r2, c1, c2 = get_block_k_ix(self, k)
        d1 = self.tdhf_diag(r1)
        d2 = self.tdhf_diag(r2)
        a = d1 + 2 * self["knmj", r1, c1] - self["knjm", r1, c1]
        b = 2 * self["kjmn", r1, c2] - self["kjnm", r1, c2]
        a_ = d2 + 2 * self["mjkn", r2, c2] - self["mjnk", r2, c2]
        b_ = 2 * self["mnkj", r2, c1] - self["mnjk", r2, c1]
        return "full", numpy.block([[a, b], [-b_, -a_]])


class PhysERI4(PhysERI):

    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing partial transformations of integrals to Bloch functions. A 4-fold
        symmetry of complex-valued functions is employed in this class. The ERIs are returned in blocks of k-points.

        Args:
            model (KRHF): the base model;
            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;
        """
        td.PhysERI4.__init__.im_func(self, model, frozen=frozen)

    symmetries = [
        ((0, 1, 2, 3), False),
        ((1, 0, 3, 2), False),
        ((2, 3, 0, 1), True),
        ((3, 2, 1, 0), True),
    ]

    def __calc_block__(self, item, k):
        if self.kconserv[k[:3]] == k[3]:
            return td.PhysERI4.__calc_block__.im_func(self, item, k)
        else:
            raise ValueError("K is not conserved: {}, expected {}".format(
                repr(k),
                k[:3] + (self.kconserv[k[:3]],),
            ))


class PhysERI8(PhysERI4):
    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing partial transformations of integrals to Bloch functions. An 8-fold
        symmetry of real-valued functions is employed in this class. The ERIs are returned in blocks of k-points.

        Args:
            model (KRHF): the base model;
            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;
        """
        super(PhysERI8, self).__init__(model, frozen=frozen)

    symmetries = [
        ((0, 1, 2, 3), False),
        ((1, 0, 3, 2), False),
        ((2, 3, 0, 1), False),
        ((3, 2, 1, 0), False),

        ((2, 1, 0, 3), False),
        ((3, 0, 1, 2), False),
        ((0, 3, 2, 1), False),
        ((1, 2, 3, 0), False),
    ]


def get_block_k_ix(eri, k):
    """
    Retrieves k indexes of the block with a specific momentum transfer.
    Args:
        eri (TDDFTMatrixBlocks): ERI of the problem;
        k (tuple, int): momentum transfer: either a pair of k-point indexes specifying the momentum transfer
        vector or a single integer with the second index assuming the first index being zero;

    Returns:
        4 arrays: r1, r2, c1, c2 specifying k-indexes of the ERI matrix block.

        +-----------------+-------------+-------------+-----+-----------------+-------------+-------------+-----+-----------------+
        |                 | k34=0,c1[0] | k34=1,c1[1] | ... | k34=nk-1,c1[-1] | k34=0,c2[0] | k34=1,c2[1] | ... | k34=nk-1,c2[-1] |
        +-----------------+-------------+-------------+-----+-----------------+-------------+-------------+-----+-----------------+
        |   k12=0,r1[0]   |                                                   |                                                   |
        +-----------------+                                                   |                                                   |
        |   k12=1,r1[1]   |                                                   |                                                   |
        +-----------------+                  Block r1, c1                     |                  Block r1, c2                     |
        |       ...       |                                                   |                                                   |
        +-----------------+                                                   |                                                   |
        | k12=nk-1,r1[-1] |                                                   |                                                   |
        +-----------------+---------------------------------------------------+---------------------------------------------------+
        |   k12=0,r2[0]   |                                                   |                                                   |
        +-----------------+                                                   |                                                   |
        |   k12=1,r2[1]   |                                                   |                                                   |
        +-----------------+                  Block r2, c1                     |                  Block r2, c2                     |
        |       ...       |                                                   |                                                   |
        +-----------------+                                                   |                                                   |
        | k12=nk-1,r2[-1] |                                                   |                                                   |
        +-----------------+---------------------------------------------------+---------------------------------------------------+
    """
    # All checks here are for debugging purposes
    if isinstance(k, int):
        k = (0, k)
    r1, c1 = eri.get_k_ix("knmj", k)
    assert r1[k[0]] == k[1]
    # knmj and kjmn share row indexes
    _, c2 = eri.get_k_ix("kjmn", (0, r1[0]))
    assert abs(r1 - _).max() == 0
    # knmj and mnkj share column indexes
    _, r2 = eri.get_k_ix("mnkj", (0, c1[0]))
    assert abs(c1 - _).max() == 0
    _r, _c = eri.get_k_ix("mjkn", (0, r2[0]))
    assert abs(r2 - _r).max() == 0
    assert abs(c2 - _c).max() == 0
    _c, _r = eri.get_k_ix("mjkn", (0, c2[0]))
    assert abs(r2 - _r).max() == 0
    assert abs(c2 - _c).max() == 0

    assert abs(r1 - c1).max() == 0
    assert abs(r2 - c2).max() == 0
    assert abs(r1[r2] - numpy.arange(len(r1))).max() == 0
    # The output is, basically, r1, argsort(r1), r1, argsort(r1)
    return r1, r2, c1, c2


def vector_to_amplitudes(vectors, nocc, nmo):
    """
    Transforms (reshapes) and normalizes vectors into amplitudes.
    Args:
        vectors (numpy.ndarray): raw eigenvectors to transform;
        nocc (tuple): numbers of occupied orbitals;
        nmo (int): the total number of orbitals per k-point;

    Returns:
        Amplitudes with the following shape: (# of roots, 2 (x or y), # of kpts, # of occupied orbitals,
        # of virtual orbitals).
    """
    if not all(i == nocc[0] for i in nocc):
        raise NotImplementedError("Varying occupation numbers are not implemented yet")
    nk = len(nocc)
    nocc = nocc[0]
    if not all(i == nmo[0] for i in nmo):
        raise NotImplementedError("Varying AO spaces are not implemented yet")
    nmo = nmo[0]
    vectors = numpy.asanyarray(vectors)
    vectors = vectors.reshape(2, nk, nocc, nmo-nocc, vectors.shape[1])
    norm = (abs(vectors) ** 2).sum(axis=(1, 2, 3))
    norm = 2 * (norm[0] - norm[1])
    vectors /= norm ** .5
    return vectors.transpose(4, 0, 1, 2, 3)


class TDRHF(rhf_slow.TDRHF):
    eri4 = PhysERI4
    eri8 = PhysERI8
    v2a = staticmethod(vector_to_amplitudes)

    def __init__(self, mf, frozen=None):
        """
        Performs TDHF calculation. Roots and eigenvectors are stored in `self.e`, `self.xy`.
        Args:
            mf (RHF): the base restricted Hartree-Fock model;
            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;
        """
        super(TDRHF, self).__init__(mf, frozen=frozen)
        self.e = {}
        self.xy = {}

    def kernel(self, k=None):
        """
        Calculates eigenstates and eigenvalues of the TDHF problem.
        Args:
            k (tuple, int): momentum transfer: either an index specifying the momentum transfer or a list of such
            indexes;

        Returns:
            Positive eigenvalues and eigenvectors.
        """
        if k is None:
            k = numpy.arange(len(self._scf.kpts))

        if isinstance(k, int):
            k = [k]

        for kk in k:
            self.e[kk], self.xy[kk] = self.__kernel__(k=kk)
        return self.e, self.xy


# pyscf/pbc/tdscf/krhf_slow_gamma.py

#  Author: Artem Pulkin
# flake8: noqa
"""
This and other `_slow` modules implement the time-dependent Hartree-Fock 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 via an AO-MO transformation, i.e. with a O(N^5) complexity scaling. 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. Several variants of TDHF are available:

 * `pyscf.tdscf.rhf_slow`: the molecular implementation;
 * `pyscf.pbc.tdscf.rhf_slow`: PBC (periodic boundary condition) implementation for RHF objects of `pyscf.pbc.scf`
   modules;
 * `pyscf.pbc.tdscf.krhf_slow_supercell`: PBC implementation for KRHF 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;
 * (this module) `pyscf.pbc.tdscf.krhf_slow_gamma`: A Gamma-point calculation resembling the original `pyscf.pbc.tdscf.krhf`
   module. Despite its name, it accepts KRHF objects with an arbitrary number of k-points but finds only few TDHF roots
   corresponding to collective oscillations without momentum transfer;
 * `pyscf.pbc.tdscf.krhf_slow`: PBC implementation for KRHF objects of `pyscf.pbc.scf` modules. Works with
   an arbitrary number of k-points and employs k-point conservation (diagonalizes matrix blocks separately).
"""

from pyscf.pbc.tdscf import krhf_slow_supercell as td
from pyscf.tdscf import rhf_slow

import numpy


# Convention for these modules:
# * PhysERI, PhysERI4, PhysERI8 are 2-electron integral routines computed directly (for debug purposes), with a 4-fold
#   symmetry and with an 8-fold symmetry
# * vector_to_amplitudes reshapes and normalizes the solution
# * TDRHF provides a container


class PhysERI(td.PhysERI):

    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing a full transformation of integrals to Bloch functions. No symmetries are
        employed in this class. Only a subset of transformed ERI is returned, corresponding to oscillations without a
        momentum transfer. The k-points of the returned ERIs come in two pairs :math:`(k_1 k_1 | k_2 k_2)`. As a result,
        one of the diagonal blocks of the full TDHF matrix can be reconstructed. For other blocks, look into `PhysERI`
        class of this momdule.

        Args:
            model (KRHF): the base model;
            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;
        """
        super(PhysERI, self).__init__(model, frozen=frozen)

    def tdhf_diag(self):
        """
        Retrieves the merged diagonal block with equal pairs of k-indexes (k, k) only.

        Returns:
            The diagonal block.
        """
        return super(PhysERI, self).tdhf_diag(pairs=((i, i) for i in range(len(self.model.kpts))))

    def __calc_block__(self, item, k):
        if k in self.__full_eri_k__:
            return super(PhysERI, self).__calc_block__(item, k)
        else:
            raise RuntimeError("The block k = {:d} + {:d} - {:d} - {:d} does not conserve momentum".format(*k))

    def eri_mknj(self, item):
        """
        Retrieves the merged ERI block using 'mknj' notation with pairs of k-indexes (k1, k1, k2, k2).
        Args:
            item (str): a 4-character string of 'mknj' letters;

        Returns:
            The corresponding block of ERI (phys notation).
        """
        return super(PhysERI, self).eri_mknj(
            item,
            pairs_row=((i, i) for i in range(len(self.model.kpts))),
            pairs_column=((i, i) for i in range(len(self.model.kpts))),
        )


class PhysERI4(PhysERI):

    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing partial transformations of integrals to Bloch functions. A 4-fold
        symmetry of complex-valued functions is employed in this class. Only a subset of transformed ERI is returned,
        corresponding to oscillations without a momentum transfer. The k-points of the returned ERIs come in two pairs
        :math:`(k_1 k_1 | k_2 k_2)`. As a result, one of the diagonal blocks of the full TDHF matrix can be
        reconstructed. For other blocks, look into `PhysERI` class of this momdule.

        Args:
            model (KRHF): the base model;
            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;
        """
        super(PhysERI4, self).__init__(model, frozen=frozen)

    symmetries = [
        ((0, 1, 2, 3), False),
        ((1, 0, 3, 2), False),
        ((2, 3, 0, 1), True),
        ((3, 2, 1, 0), True),
    ]

    def __calc_block__(self, item, k):
        if k in self.__full_eri_k__:
            return td.PhysERI4.__calc_block__.im_func(self, item, k)
        else:
            raise RuntimeError("The block k = {:d} + {:d} - {:d} - {:d} does not conserve momentum".format(*k))


class PhysERI8(PhysERI4):
    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing partial transformations of integrals to Bloch functions. An 8-fold
        symmetry of real-valued functions is employed in this class. Only a subset of transformed ERI is returned,
        corresponding to oscillations without a momentum transfer. The k-points of the returned ERIs come in two pairs
        :math:`(k_1 k_1 | k_2 k_2)`. As a result, one of the diagonal blocks of the full TDHF matrix can be
        reconstructed. For other blocks, look into `PhysERI` class of this momdule.

        Args:
            model (KRHF): the base model;
            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;
        """
        super(PhysERI8, self).__init__(model, frozen=frozen)

    symmetries = [
        ((0, 1, 2, 3), False),
        ((1, 0, 3, 2), False),
        ((2, 3, 0, 1), False),
        ((3, 2, 1, 0), False),

        ((2, 1, 0, 3), False),
        ((3, 0, 1, 2), False),
        ((0, 3, 2, 1), False),
        ((1, 2, 3, 0), False),
    ]


def vector_to_amplitudes(vectors, nocc, nmo):
    """
    Transforms (reshapes) and normalizes vectors into amplitudes.
    Args:
        vectors (numpy.ndarray): raw eigenvectors to transform;
        nocc (tuple): numbers of occupied orbitals;
        nmo (int): the total number of orbitals per k-point;

    Returns:
        Amplitudes with the following shape: (# of roots, 2 (x or y), # of kpts, # of kpts, # of occupied orbitals,
        # of virtual orbitals).
    """
    if not all(i == nocc[0] for i in nocc):
        raise NotImplementedError("Varying occupation numbers are not implemented yet")
    nk = len(nocc)
    nocc = nocc[0]
    if not all(i == nmo[0] for i in nmo):
        raise NotImplementedError("Varying AO spaces are not implemented yet")
    nmo = nmo[0]
    vectors = numpy.asanyarray(vectors)
    # Compared to krhf_slow_supercell, only one k-point index is present here. The second index corresponds to
    # momentum transfer and is integrated out
    vectors = vectors.reshape(2, nk, nocc, nmo-nocc, vectors.shape[1])
    norm = (abs(vectors) ** 2).sum(axis=(1, 2, 3))
    norm = 2 * (norm[0] - norm[1])
    vectors /= norm ** .5
    return vectors.transpose(4, 0, 1, 2, 3)


class TDRHF(rhf_slow.TDRHF):
    eri4 = PhysERI4
    eri8 = PhysERI8
    v2a = staticmethod(vector_to_amplitudes)


# pyscf/pbc/tdscf/krhf_slow_supercell.py

#  Author: Artem Pulkin
"""
This and other `_slow` modules implement the time-dependent Hartree-Fock 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 via an AO-MO transformation, i.e. with a O(N^5) complexity scaling. 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. Several variants of TDHF are available:

 * `pyscf.tdscf.rhf_slow`: the molecular implementation;
 * `pyscf.pbc.tdscf.rhf_slow`: PBC (periodic boundary condition) implementation for RHF objects of `pyscf.pbc.scf`
   modules;
 * (this module) `pyscf.pbc.tdscf.krhf_slow_supercell`: PBC implementation for KRHF 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.krhf_slow_gamma`: A Gamma-point calculation resembling the original `pyscf.pbc.tdscf.krhf`
   module. Despite its name, it accepts KRHF objects with an arbitrary number of k-points but finds only few TDHF roots
   corresponding to collective oscillations without momentum transfer;
 * `pyscf.pbc.tdscf.krhf_slow`: PBC implementation for KRHF objects of `pyscf.pbc.scf` modules. Works with
   an arbitrary number of k-points and employs k-point conservation (diagonalizes matrix blocks separately).
"""

from pyscf.tdscf.common_slow import TDERIMatrixBlocks, PeriodicMFMixin
from pyscf.tdscf import rhf_slow

from pyscf.pbc.lib.kpts_helper import loop_kkk
from pyscf.lib import logger

import numpy
import scipy
from itertools import product


# Convention for these modules:
# * PhysERI, PhysERI4, PhysERI8 are 2-electron integral routines computed directly (for debug purposes), with a 4-fold
#   symmetry and with an 8-fold symmetry
# * vector_to_amplitudes reshapes and normalizes the solution
# * TDRHF provides a container


class PhysERI(PeriodicMFMixin, TDERIMatrixBlocks):

    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing a full transformation of integrals to Bloch functions. No symmetries are
        employed in this class.

        Args:
            model (KRHF): the base model;
            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;
        """
        TDERIMatrixBlocks.__init__(self)
        PeriodicMFMixin.__init__(self, model, frozen=frozen)
        # Phys representation
        self.__full_eri_k__ = {}
        for k in loop_kkk(len(model.kpts)):
            k = k + (self.kconserv[k],)
            self.__full_eri_k__[k] = self.ao2mo_k(tuple(self.mo_coeff[j] for j in k), k)

    def ao2mo_k(self, coeff, k):
        """
        Phys ERI in MO basis.
        Args:
            coeff (Iterable): MO orbitals;
            k (Iterable): the 4 k-points MOs correspond to;

        Returns:
            ERI in MO basis.
        """
        coeff = (coeff[0], coeff[2], coeff[1], coeff[3])
        k = (k[0], k[2], k[1], k[3])
        result = self.model.with_df.ao2mo(coeff, tuple(self.model.kpts[i] for i in k), compact=False)
        return result.reshape(
            tuple(i.shape[1] for i in coeff)
        ).swapaxes(1, 2)

    def __get_mo_energies__(self, k1, k2):
        """This routine collects occupied and virtual MO energies."""
        return self.mo_energy[k1][:self.nocc[k1]], self.mo_energy[k2][self.nocc[k2]:]

    def tdhf_diag_k(self, k1, k2):
        """
        Retrieves the diagonal block.
        Args:
            k1 (int): first k-index (row);
            k2 (int): second k-index (column);

        Returns:
            The diagonal block.
        """
        # Everything is already implemented in molecular code
        return super(PhysERI, self).tdhf_diag(k1, k2)

    def tdhf_diag(self, pairs=None):
        """
        Retrieves the merged diagonal block with specified or all possible k-index pairs.
        Args:
            pairs (Iterable): pairs of k-points to assmble;

        Returns:
            The diagonal block.
        """
        if pairs is None:
            pairs = product(range(len(self.model.kpts)), range(len(self.model.kpts)))
        result = []
        for k1, k2 in pairs:
            result.append(self.tdhf_diag_k(k1, k2))
        return scipy.linalg.block_diag(*result)

    def __calc_block__(self, item, k):
        if k in self.__full_eri_k__:
            slc = tuple(slice(self.nocc[_k]) if i == 'o' else slice(self.nocc[_k], None) for i, _k in zip(item, k))
            return self.__full_eri_k__[k][slc]
        else:
            return numpy.zeros(tuple(
                self.nocc[_k] if i == 'o' else self.nmo[_k] - self.nocc[_k]
                for i, _k in zip(item, k)
            ))

    def eri_mknj_k(self, item, k):
        """
        Retrieves ERI block using 'mknj' notation.
        Args:
            item (str): a 4-character string of 'mknj' letters;
            k (Iterable): k indexes;

        Returns:
            The corresponding block of ERI (phys notation).
        """
        # Everything is already implemented in molecular code
        return super(PhysERI, self).eri_mknj(item, k)

    def eri_mknj(self, item, pairs_row=None, pairs_column=None):
        """
        Retrieves the merged ERI block using 'mknj' notation with all k-indexes.
        Args:
            item (str): a 4-character string of 'mknj' letters;
            pairs_row (Iterable): iterator for pairs of row k-points (first index in the output matrix);
            pairs_column (Iterable): iterator for pairs of column k-points (second index in the output matrix);

        Returns:
            The corresponding block of ERI (phys notation).
        """
        if pairs_row is None:
            pairs_row = product(range(len(self.model.kpts)), range(len(self.model.kpts)))
        if pairs_column is None:
            pairs_column = product(range(len(self.model.kpts)), range(len(self.model.kpts)))
        # Second index has to support re-iterations
        pairs_column = tuple(pairs_column)
        result = []
        for k1, k2 in pairs_row:
            result.append([])
            for k3, k4 in pairs_column:
                result[-1].append(self.eri_mknj_k(item, (k1, k2, k3, k4)))

        r = numpy.block(result)
        return r / len(self.model.kpts)


class PhysERI4(PhysERI):
    symmetries = [
        ((0, 1, 2, 3), False),
        ((1, 0, 3, 2), False),
        ((2, 3, 0, 1), True),
        ((3, 2, 1, 0), True),
    ]

    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing a partial transformation of integrals to Bloch functions. A 4-fold
        symmetry of complex-valued wavefunctions is employed in this class.

        Args:
            model (KRHF): the base model;
            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;
        """
        TDERIMatrixBlocks.__init__(self)
        PeriodicMFMixin.__init__(self, model, frozen=frozen)

    def __calc_block__(self, item, k):
        if self.kconserv[k[:3]] == k[3]:
            logger.info(self.model, "Computing {} {} ...".format(''.join(item), repr(k)))
            return self.ao2mo_k(tuple(
                self.mo_coeff[_k][:, :self.nocc[_k]] if i == "o" else self.mo_coeff[_k][:, self.nocc[_k]:]
                for i, _k in zip(item, k)
            ), k)
        else:
            return numpy.zeros(tuple(
                self.nocc[_k] if i == 'o' else self.nmo[_k] - self.nocc[_k]
                for i, _k in zip(item, k)
            ))


class PhysERI8(PhysERI4):
    symmetries = [
        ((0, 1, 2, 3), False),
        ((1, 0, 3, 2), False),
        ((2, 3, 0, 1), False),
        ((3, 2, 1, 0), False),

        ((2, 1, 0, 3), False),
        ((3, 0, 1, 2), False),
        ((0, 3, 2, 1), False),
        ((1, 2, 3, 0), False),
    ]

    def __init__(self, model, frozen=None):
        """
        The TDHF ERI implementation performing a partial transformation of integrals to Bloch functions. An 8-fold
        symmetry of real-valued wavefunctions is employed in this class.

        Args:
            model (KRHF): the base model;
            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;
        """
        super(PhysERI8, self).__init__(model, frozen=frozen)


def vector_to_amplitudes(vectors, nocc, nmo):
    """
    Transforms (reshapes) and normalizes vectors into amplitudes.
    Args:
        vectors (numpy.ndarray): raw eigenvectors to transform;
        nocc (tuple): numbers of occupied orbitals;
        nmo (tuple): the total numbers of AOs per k-point;

    Returns:
        Amplitudes with the following shape: (# of roots, 2 (x or y), # of kpts, # of kpts, # of occupied orbitals,
        # of virtual orbitals).
    """
    if not all(i == nocc[0] for i in nocc):
        raise NotImplementedError("Varying occupation numbers are not implemented yet")
    nk = len(nocc)
    nocc = nocc[0]
    if not all(i == nmo[0] for i in nmo):
        raise NotImplementedError("Varying AO spaces are not implemented yet")
    nmo = nmo[0]
    vectors = numpy.asanyarray(vectors)
    vectors = vectors.reshape(2, nk, nk, nocc, nmo-nocc, vectors.shape[1])
    norm = (abs(vectors) ** 2).sum(axis=(1, 2, 3, 4))
    norm = 2 * (norm[0] - norm[1])
    vectors /= norm ** .5
    return vectors.transpose(5, 0, 1, 2, 3, 4)


class TDRHF(rhf_slow.TDRHF):
    eri4 = PhysERI4
    eri8 = PhysERI8
    v2a = staticmethod(vector_to_amplitudes)


# pyscf/pbc/tdscf/krks.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#
# Ref:
# Chem Phys Lett, 256, 454
# J. Mol. Struct. THEOCHEM, 914, 3
#

from pyscf import lib
from pyscf.pbc.tdscf import krhf


KTDA = TDA = krhf.TDA

RPA = KTDDFT = TDDFT = krhf.TDHF

#TODO:
TDDFTNoHybrid = TDDFT


def tddft(mf):
    '''Driver to create TDDFT or TDDFTNoHybrid object'''
    if mf._numint.libxc.is_hybrid_xc(mf.xc):
        return TDDFT(mf)
    else:
        return TDDFTNoHybrid(mf)

from pyscf.pbc import dft
dft.krks.KRKS.TDA   = lib.class_as_method(KTDA)
dft.krks.KRKS.TDHF  = None
dft.krks.KRKS.TDDFT = tddft
dft.kroks.KROKS.TDA   = None
dft.kroks.KROKS.TDHF  = None
dft.kroks.KROKS.TDDFT = None


if __name__ == '__main__':
    from pyscf.pbc import gto
    from pyscf.pbc import dft, df
    cell = gto.Cell()
    cell.unit = 'B'
    cell.atom = '''
    C  0.          0.          0.
    C  1.68506879  1.68506879  1.68506879
    '''
    cell.a = '''
    0.          3.37013758  3.37013758
    3.37013758  0.          3.37013758
    3.37013758  3.37013758  0.
    '''
    cell.basis = 'gth-szv'
    cell.pseudo = 'gth-pade'
    cell.mesh = [25]*3
    cell.build()

    mf = dft.KRKS(cell, cell.make_kpts([2,1,1]))
    #mf.with_df = df.MDF(cell, cell.make_kpts([2,1,1]))
    #mf.with_df.auxbasis = 'weigend'
    #mf.with_df._cderi = 'eri3d-mdf.h5'
    #mf.with_df.build(with_j3c=False)
    mf.xc = 'lda,'
    mf.kernel()
#mesh=12 -10.3077341607895
#mesh=5  -10.3086623157515

    td = TDDFT(mf)
    td.nstates = 5
    td.verbose = 5
    print(td.kernel()[0] * 27.2114)
#mesh=12 [ 6.08108297  6.10231481  6.10231478  6.38355803  6.38355804]
#MDF mesh=5 [ 6.07919157  6.10251718  6.10253961  6.37202499  6.37565246]

    td = TDA(mf)
    td.singlet = False
    td.verbose = 5
    print(td.kernel()[0] * 27.2114)
#mesh=12 [ 4.01539192  5.1750807   5.17508071]
#MDF mesh=5 [ 4.01148649  5.18043397  5.18043459]


# pyscf/pbc/tdscf/kuhf.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

from functools import reduce
import numpy
from pyscf import lib
from pyscf.lib import logger
from pyscf.tdscf import uhf
from pyscf.pbc import scf
from pyscf.pbc.tdscf.krhf import _get_e_ia, purify_krlyov_heff
from pyscf.pbc.lib.kpts_helper import gamma_point
from pyscf.pbc.scf import _response_functions  # noqa
from pyscf import __config__

REAL_EIG_THRESHOLD = getattr(__config__, 'pbc_tdscf_uhf_TDDFT_pick_eig_threshold', 1e-3)
POSTIVE_EIG_THRESHOLD = getattr(__config__, 'pbc_tdscf_uhf_TDDFT_positive_eig_threshold', 1e-3)

class TDA(uhf.TDA):

    conv_tol = getattr(__config__, 'pbc_tdscf_rhf_TDA_conv_tol', 1e-6)

    def __init__(self, mf):
        from pyscf.pbc.df.df_ao2mo import warn_pbc2d_eri
        assert(isinstance(mf, scf.khf.KSCF))
        self.cell = mf.cell
        uhf.TDA.__init__(self, mf)
        warn_pbc2d_eri(mf)

    def gen_vind(self, mf):
        '''Compute Ax'''
        mo_coeff = mf.mo_coeff
        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        nkpts = len(mo_occ[0])
        nao, nmo = mo_coeff[0][0].shape
        occidxa = [numpy.where(mo_occ[0][k]> 0)[0] for k in range(nkpts)]
        occidxb = [numpy.where(mo_occ[1][k]> 0)[0] for k in range(nkpts)]
        viridxa = [numpy.where(mo_occ[0][k]==0)[0] for k in range(nkpts)]
        viridxb = [numpy.where(mo_occ[1][k]==0)[0] for k in range(nkpts)]
        orboa = [mo_coeff[0][k][:,occidxa[k]] for k in range(nkpts)]
        orbob = [mo_coeff[1][k][:,occidxb[k]] for k in range(nkpts)]
        orbva = [mo_coeff[0][k][:,viridxa[k]] for k in range(nkpts)]
        orbvb = [mo_coeff[1][k][:,viridxb[k]] for k in range(nkpts)]

        e_ia_a = _get_e_ia(mo_energy[0], mo_occ[0])
        e_ia_b = _get_e_ia(mo_energy[1], mo_occ[1])
        hdiag = numpy.hstack([x.ravel() for x in (e_ia_a + e_ia_b)])

        mem_now = lib.current_memory()[0]
        max_memory = max(2000, self.max_memory*.8-mem_now)
        vresp = mf.gen_response(hermi=0, max_memory=max_memory)

        def vind(zs):
            nz = len(zs)
            zs = [_unpack(z, mo_occ) for z in zs]
            dmov = numpy.empty((2,nz,nkpts,nao,nao), dtype=numpy.complex128)
            for i in range(nz):
                dm1a, dm1b = zs[i]
                for k in range(nkpts):
                    dmov[0,i,k] = reduce(numpy.dot, (orboa[k], dm1a[k], orbva[k].conj().T))
                    dmov[1,i,k] = reduce(numpy.dot, (orbob[k], dm1b[k], orbvb[k].conj().T))

            with lib.temporary_env(mf, exxdiv=None):
                dmov = dmov.reshape(2*nz,nkpts,nao,nao)
                v1ao = vresp(dmov)
                v1ao = v1ao.reshape(2,nz,nkpts,nao,nao)

            v1s = []
            for i in range(nz):
                dm1a, dm1b = zs[i]
                v1as = []
                v1bs = []
                for k in range(nkpts):
                    v1a = reduce(numpy.dot, (orboa[k].conj().T, v1ao[0,i,k], orbva[k]))
                    v1b = reduce(numpy.dot, (orbob[k].conj().T, v1ao[1,i,k], orbvb[k]))
                    v1a += e_ia_a[k] * dm1a[k]
                    v1b += e_ia_b[k] * dm1b[k]
                    v1as.append(v1a.ravel())
                    v1bs.append(v1b.ravel())
                v1s += v1as + v1bs
            return numpy.hstack(v1s).reshape(nz,-1)

        return vind, hdiag

    def init_guess(self, mf, nstates=None):
        if nstates is None: nstates = self.nstates

        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        e_ia_a = _get_e_ia(mo_energy[0], mo_occ[0])
        e_ia_b = _get_e_ia(mo_energy[1], mo_occ[1])
        e_ia = numpy.hstack([x.ravel() for x in (e_ia_a + e_ia_b)])

        e_ia_max = e_ia.max()
        nov = e_ia.size
        nstates = min(nstates, nov)
        e_threshold = min(e_ia_max, e_ia[numpy.argsort(e_ia)[nstates-1]])
        # Handle degeneracy
        e_threshold += 1e-6
        idx = numpy.where(e_ia <= e_threshold)[0]
        x0 = numpy.zeros((idx.size, nov))
        for i, j in enumerate(idx):
            x0[i, j] = 1  # Koopmans' excitations
        return x0

    def kernel(self, x0=None):
        '''TDA diagonalization solver
        '''
        self.check_sanity()
        self.dump_flags()

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)
        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        def pickeig(w, v, nroots, envs):
            idx = numpy.where(w > POSTIVE_EIG_THRESHOLD)[0]
            return w[idx], v[:,idx], idx

        log = logger.Logger(self.stdout, self.verbose)
        precision = self.cell.precision * 1e-2
        hermi = 1

        self.converged, self.e, x1 = \
                lib.davidson1(vind, x0, precond,
                              tol=self.conv_tol,
                              nroots=self.nstates, lindep=self.lindep,
                              max_space=self.max_space, pick=pickeig,
                              fill_heff=purify_krlyov_heff(precision, hermi, log),
                              verbose=self.verbose)

        mo_occ = self._scf.mo_occ
        self.xy = [(_unpack(xi, mo_occ),  # (X_alpha, X_beta)
                    (0, 0))  # (Y_alpha, Y_beta)
                   for xi in x1]
        #TODO: analyze CIS wfn point group symmetry
        return self.e, self.xy
CIS = KTDA = TDA


class TDHF(TDA):
    def gen_vind(self, mf):
        mo_coeff = mf.mo_coeff
        mo_energy = mf.mo_energy
        mo_occ = mf.mo_occ
        nkpts = len(mo_occ[0])
        nao, nmo = mo_coeff[0][0].shape
        occidxa = [numpy.where(mo_occ[0][k]> 0)[0] for k in range(nkpts)]
        occidxb = [numpy.where(mo_occ[1][k]> 0)[0] for k in range(nkpts)]
        viridxa = [numpy.where(mo_occ[0][k]==0)[0] for k in range(nkpts)]
        viridxb = [numpy.where(mo_occ[1][k]==0)[0] for k in range(nkpts)]
        orboa = [mo_coeff[0][k][:,occidxa[k]] for k in range(nkpts)]
        orbob = [mo_coeff[1][k][:,occidxb[k]] for k in range(nkpts)]
        orbva = [mo_coeff[0][k][:,viridxa[k]] for k in range(nkpts)]
        orbvb = [mo_coeff[1][k][:,viridxb[k]] for k in range(nkpts)]

        e_ia_a = _get_e_ia(mo_energy[0], mo_occ[0])
        e_ia_b = _get_e_ia(mo_energy[1], mo_occ[1])
        hdiag = numpy.hstack([x.ravel() for x in (e_ia_a + e_ia_b)])
        hdiag = numpy.hstack((hdiag, -hdiag))
        tot_x_a = sum(x.size for x in e_ia_a)
        tot_x_b = sum(x.size for x in e_ia_b)
        tot_x = tot_x_a + tot_x_b

        mem_now = lib.current_memory()[0]
        max_memory = max(2000, self.max_memory*.8-mem_now)
        vresp = mf.gen_response(hermi=0, max_memory=max_memory)

        def vind(xys):
            nz = len(xys)
            x1s = [_unpack(x[:tot_x], mo_occ) for x in xys]
            y1s = [_unpack(x[tot_x:], mo_occ) for x in xys]
            dmov = numpy.empty((2,nz,nkpts,nao,nao), dtype=numpy.complex128)
            for i in range(nz):
                xa, xb = x1s[i]
                ya, yb = y1s[i]
                for k in range(nkpts):
                    dmx = reduce(numpy.dot, (orboa[k], xa[k]  , orbva[k].conj().T))
                    dmy = reduce(numpy.dot, (orbva[k], ya[k].T, orboa[k].conj().T))
                    dmov[0,i,k] = dmx + dmy  # AX + BY
                    dmx = reduce(numpy.dot, (orbob[k], xb[k]  , orbvb[k].conj().T))
                    dmy = reduce(numpy.dot, (orbvb[k], yb[k].T, orbob[k].conj().T))
                    dmov[1,i,k] = dmx + dmy  # AX + BY

            with lib.temporary_env(mf, exxdiv=None):
                dmov = dmov.reshape(2*nz,nkpts,nao,nao)
                v1ao = vresp(dmov)
                v1ao = v1ao.reshape(2,nz,nkpts,nao,nao)

            v1s = []
            for i in range(nz):
                xa, xb = x1s[i]
                ya, yb = y1s[i]
                v1xsa = []
                v1xsb = []
                v1ysa = []
                v1ysb = []
                for k in range(nkpts):
                    v1xa = reduce(numpy.dot, (orboa[k].conj().T, v1ao[0,i,k], orbva[k]))
                    v1xb = reduce(numpy.dot, (orbob[k].conj().T, v1ao[1,i,k], orbvb[k]))
                    v1ya = reduce(numpy.dot, (orbva[k].conj().T, v1ao[0,i,k], orboa[k])).T
                    v1yb = reduce(numpy.dot, (orbvb[k].conj().T, v1ao[1,i,k], orbob[k])).T
                    v1xa+= e_ia_a[k] * xa[k]
                    v1xb+= e_ia_b[k] * xb[k]
                    v1ya+= e_ia_a[k] * ya[k]
                    v1yb+= e_ia_b[k] * yb[k]
                    v1xsa.append(v1xa.ravel())
                    v1xsb.append(v1xb.ravel())
                    v1ysa.append(-v1ya.ravel())
                    v1ysb.append(-v1yb.ravel())
                v1s += v1xsa + v1xsb + v1ysa + v1ysb
            return numpy.hstack(v1s).reshape(nz,-1)

        return vind, hdiag

    def init_guess(self, mf, nstates=None, wfnsym=None):
        x0 = TDA.init_guess(self, mf, nstates)
        y0 = numpy.zeros_like(x0)
        return numpy.hstack((x0,y0))

    def kernel(self, x0=None):
        '''TDHF diagonalization with non-Hermitian eigenvalue solver
        '''
        self.check_sanity()
        self.dump_flags()

        vind, hdiag = self.gen_vind(self._scf)
        precond = self.get_precond(hdiag)
        if x0 is None:
            x0 = self.init_guess(self._scf, self.nstates)

        real_system = (gamma_point(self._scf.kpts) and
                       self._scf.mo_coeff[0][0].dtype == numpy.double)

        # We only need positive eigenvalues
        def pickeig(w, v, nroots, envs):
            realidx = numpy.where((abs(w.imag) < REAL_EIG_THRESHOLD) &
                                  (w.real > POSTIVE_EIG_THRESHOLD))[0]
            return lib.linalg_helper._eigs_cmplx2real(w, v, realidx, real_system)

        log = logger.Logger(self.stdout, self.verbose)
        precision = self.cell.precision * 1e-2

        self.converged, w, x1 = \
                lib.davidson_nosym1(vind, x0, precond,
                                    tol=self.conv_tol,
                                    nroots=self.nstates, lindep=self.lindep,
                                    max_space=self.max_space, pick=pickeig,
                                    fill_heff=purify_krlyov_heff(precision, 0, log),
                                    verbose=self.verbose)

        mo_occ = self._scf.mo_occ
        e = []
        xy = []
        for i, z in enumerate(x1):
            xs, ys = z.reshape(2,-1)
            norm = lib.norm(xs)**2 - lib.norm(ys)**2
            if norm > 0:
                norm = 1/numpy.sqrt(norm)
                xs *= norm
                ys *= norm
                e.append(w[i])
                xy.append((_unpack(xs, mo_occ), _unpack(ys, mo_occ)))
        self.e = numpy.array(e)
        self.xy = xy
        return self.e, self.xy
RPA = KTDHF = TDHF

def _unpack(vo, mo_occ):
    za = []
    zb = []
    p1 = 0
    for k, occ in enumerate(mo_occ[0]):
        no = numpy.count_nonzero(occ > 0)
        nv = occ.size - no
        p0, p1 = p1, p1 + no * nv
        za.append(vo[p0:p1].reshape(no,nv))

    for k, occ in enumerate(mo_occ[1]):
        no = numpy.count_nonzero(occ > 0)
        nv = occ.size - no
        p0, p1 = p1, p1 + no * nv
        zb.append(vo[p0:p1].reshape(no,nv))
    return za, zb


scf.kuhf.KUHF.TDA  = lib.class_as_method(KTDA)
scf.kuhf.KUHF.TDHF = lib.class_as_method(KTDHF)


if __name__ == '__main__':
    from pyscf.pbc import gto
    from pyscf.pbc import scf
    from pyscf.pbc import df
    cell = gto.Cell()
    cell.unit = 'B'
    cell.atom = '''
    C  0.          0.          0.
    C  1.68506879  1.68506879  1.68506879
    '''
    cell.a = '''
    0.          3.37013758  3.37013758
    3.37013758  0.          3.37013758
    3.37013758  3.37013758  0.
    '''

    cell.basis = 'gth-szv'
    cell.pseudo = 'gth-pade'
    cell.mesh = [37]*3
    cell.build()
    mf = scf.KUHF(cell, cell.make_kpts([2,1,1])).set(exxdiv=None)
#    mf.with_df = df.DF(cell, cell.make_kpts([2,1,1]))
#    mf.with_df.auxbasis = 'weigend'
#    mf.with_df._cderi = 'eri3d-df.h5'
#    mf.with_df.build(with_j3c=False)
    mf.run()

    td = TDA(mf)
    td.verbose = 5
    td.nstates = 5
    print(td.kernel()[0] * 27.2114)

    td = TDHF(mf)
    td.verbose = 5
    td.nstates = 5
    print(td.kernel()[0] * 27.2114)

    cell.spin = 2
    mf = scf.KUHF(cell, cell.make_kpts([2,1,1])).set(exxdiv=None)
    mf.run()

    td = TDA(mf)
    td.verbose = 5
    td.nstates = 5
    print(td.kernel()[0] * 27.2114)

    td = TDHF(mf)
    td.verbose = 5
    td.nstates = 5
    print(td.kernel()[0] * 27.2114)


# pyscf/pbc/tdscf/kuks.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

from pyscf.pbc.tdscf import kuhf

KTDA = TDA = kuhf.TDA

RPA = KTDDFT = TDDFT = kuhf.TDHF

#TODO:
TDDFTNoHybrid = TDDFT


def tddft(mf):
    '''Driver to create TDDFT or TDDFTNoHybrid object'''
    if mf._numint.libxc.is_hybrid_xc(mf.xc):
        return TDDFT(mf)
    else:
        return TDDFTNoHybrid(mf)

from pyscf import lib
from pyscf.pbc import dft
dft.kuks.KUKS.TDA   = lib.class_as_method(KTDA)
dft.kuks.KUKS.TDHF  = None
dft.kuks.KUKS.TDDFT = tddft


if __name__ == '__main__':
    from pyscf.pbc import gto
    from pyscf.pbc import dft
    from pyscf.pbc import df
    cell = gto.Cell()
    cell.unit = 'B'
    cell.atom = '''
    C  0.          0.          0.
    C  1.68506879  1.68506879  1.68506879
    '''
    cell.a = '''
    0.          3.37013758  3.37013758
    3.37013758  0.          3.37013758
    3.37013758  3.37013758  0.
    '''

    cell.basis = 'gth-szv'
    cell.pseudo = 'gth-pade'
    cell.mesh = [37]*3
    cell.build()
    mf = dft.KUKS(cell, cell.make_kpts([2,1,1])).set(exxdiv=None, xc='b88,p86')
    #mf.with_df = df.MDF(cell, cell.make_kpts([2,1,1]))
    #mf.with_df.auxbasis = 'weigend'
    #mf.with_df._cderi = 'eri3d-mdf.h5'
    #mf.with_df.build(with_j3c=False)
    mf.run()

    td = TDDFT(mf)
    td.verbose = 5
    td.nstates = 5
    print(td.kernel()[0] * 27.2114)

    mf.xc = 'lda,vwn'
    mf.run()
    td = TDA(mf)
    td.verbose = 5
    td.nstates = 5
    print(td.kernel()[0] * 27.2114)

    cell.spin = 2
    mf = dft.KUKS(cell, cell.make_kpts([2,1,1])).set(exxdiv=None, xc='b88,p86')
    mf.run()

    td = TDDFT(mf)
    td.verbose = 5
    td.nstates = 5
    print(td.kernel()[0] * 27.2114)

    mf.xc = 'lda,vwn'
    mf.run()
    td = TDA(mf)
    td.verbose = 5
    td.nstates = 5
    print(td.kernel()[0] * 27.2114)


# pyscf/pbc/tdscf/proxy.py

#  Author: Artem Pulkin
"""
This and other `proxy` modules implement the time-dependent mean-field procedure using the existing pyscf
implementations as a black box. The main purpose of these modules is to overcome the existing limitations in pyscf
(i.e. real-only orbitals, davidson diagonalizer, incomplete Bloch space, etc). The primary performance drawback is that,
unlike the original pyscf routines with an implicit construction of the eigenvalue problem, these modules construct TD
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 TD roots. Several variants of proxy-TD are available:

 * `pyscf.tdscf.proxy`: the molecular implementation;
 * (this module) `pyscf.pbc.tdscf.proxy`: PBC (periodic boundary condition) Gamma-point-only implementation;
 * `pyscf.pbc.tdscf.kproxy_supercell`: PBC implementation constructing supercells. Works with an arbitrary number of
   k-points but has an overhead due to ignoring the momentum conservation law. In addition, works only with
   time reversal invariant (TRI) models: i.e. the k-point grid has to be aligned and contain at least one TRI momentum.
 * `pyscf.pbc.tdscf.kproxy`: same as the above but respect the momentum conservation and, thus, diagonlizes smaller
   matrices (the performance gain is the total number of k-points in the model).
"""

# Convention for these modules:
# * PhysERI is the proxying class constructing time-dependent matrices
# * vector_to_amplitudes reshapes and normalizes the solution
# * TDProxy provides a container

from pyscf.tdscf import proxy as mol_proxy, rhf_slow, common_slow
from pyscf.pbc.tdscf import KTDDFT, KTDHF


class PhysERI(mol_proxy.PhysERI):
    proxy_choices = {
        "hf": KTDHF,
        "dft": KTDDFT,
    }


vector_to_amplitudes = rhf_slow.vector_to_amplitudes


class TDProxy(mol_proxy.TDProxy):
    proxy_eri = PhysERI


# pyscf/pbc/tdscf/rhf.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#
# Ref:
# Chem Phys Lett, 256, 454
# J. Mol. Struct. THEOCHEM, 914, 3
#

from pyscf import lib
from pyscf.tdscf import rhf


class TDA(rhf.TDA):
    def gen_vind(self, mf):
        vind, hdiag = rhf.TDA.gen_vind(self, mf)
        def vindp(x):
            with lib.temporary_env(mf, exxdiv=None):
                return vind(x)
        return vindp, hdiag

    def nuc_grad_method(self):
        raise NotImplementedError

CIS = TDA


class TDHF(rhf.TDHF):
    def gen_vind(self, mf):
        vind, hdiag = rhf.TDHF.gen_vind(self, mf)
        def vindp(x):
            with lib.temporary_env(mf, exxdiv=None):
                return vind(x)
        return vindp, hdiag

    def nuc_grad_method(self):
        raise NotImplementedError

RPA = TDRHF = TDHF


from pyscf.pbc import scf
scf.hf.RHF.TDA = lib.class_as_method(TDA)
scf.hf.RHF.TDHF = lib.class_as_method(TDHF)
scf.rohf.ROHF.TDA = None
scf.rohf.ROHF.TDHF = None



# pyscf/pbc/tdscf/rhf_slow.py

#  Author: Artem Pulkin
"""
This and other `_slow` modules implement the time-dependent Hartree-Fock 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 via an AO-MO transformation, i.e. with a O(N^5) complexity scaling. 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. Several variants of TDHF are available:

 * `pyscf.tdscf.rhf_slow`: the molecular implementation;
 * (this module) `pyscf.pbc.tdscf.rhf_slow`: PBC (periodic boundary condition) implementation for RHF objects of
   `pyscf.pbc.scf` modules;
 * `pyscf.pbc.tdscf.krhf_slow_supercell`: PBC implementation for KRHF 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.krhf_slow_gamma`: A Gamma-point calculation resembling the original `pyscf.pbc.tdscf.krhf`
   module. Despite its name, it accepts KRHF objects with an arbitrary number of k-points but finds only few TDHF roots
   corresponding to collective oscillations without momentum transfer;
 * `pyscf.pbc.tdscf.krhf_slow`: PBC implementation for KRHF 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 2-electron integral routines computed directly (for debug purposes), with a 4-fold
#   symmetry and with an 8-fold symmetry
# * vector_to_amplitudes reshapes and normalizes the solution
# * TDRHF provides a container

# This module is simply an alias of the molecular code
from pyscf.tdscf.rhf_slow import PhysERI, PhysERI4, PhysERI8, vector_to_amplitudes, TDRHF

PhysERI = PhysERI
PhysERI4 = PhysERI4
PhysERI8 = PhysERI8

vector_to_amplitudes = vector_to_amplitudes
TDRHF = TDRHF


# pyscf/pbc/tdscf/rks.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

from pyscf.tdscf import rks
from pyscf.pbc.tdscf.rhf import TDA
from pyscf.pbc.tdscf.rhf import TDHF as TDDFT
from pyscf import lib

RPA = TDRKS = TDDFT

class TDDFTNoHybrid(rks.TDDFTNoHybrid):
    def gen_vind(self, mf):
        vind, hdiag = rks.TDDFTNoHybrid.gen_vind(self, mf)
        def vindp(x):
            with lib.temporary_env(mf, exxdiv=None):
                return vind(x)
        return vindp, hdiag

    def nuc_grad_method(self):
        raise NotImplementedError

def tddft(mf):
    '''Driver to create TDDFT or TDDFTNoHybrid object'''
    if mf._numint.libxc.is_hybrid_xc(mf.xc):
        return TDDFT(mf)
    else:
        return TDDFTNoHybrid(mf)

from pyscf.pbc import dft
dft.rks.RKS.TDA           = lib.class_as_method(TDA)
dft.rks.RKS.TDHF          = None
dft.rks.RKS.TDDFT         = tddft
#dft.rks.RKS.dTDA          = lib.class_as_method(dTDA)
#dft.rks.RKS.dRPA          = lib.class_as_method(dRPA)



# pyscf/pbc/tdscf/uhf.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

from pyscf import lib
from pyscf.tdscf import uhf


class TDA(uhf.TDA):
    def gen_vind(self, mf):
        vind, hdiag = uhf.TDA.gen_vind(self, mf)
        def vindp(x):
            with lib.temporary_env(mf, exxdiv=None):
                return vind(x)
        return vindp, hdiag

    def nuc_grad_method(self):
        raise NotImplementedError

CIS = TDA


class TDHF(uhf.TDHF):
    def gen_vind(self, mf):
        vind, hdiag = uhf.TDHF.gen_vind(self, mf)
        def vindp(x):
            with lib.temporary_env(mf, exxdiv=None):
                return vind(x)
        return vindp, hdiag

    def nuc_grad_method(self):
        raise NotImplementedError

RPA = TDUHF = TDHF


from pyscf.pbc import scf
scf.uhf.UHF.TDA = lib.class_as_method(TDA)
scf.uhf.UHF.TDHF = lib.class_as_method(TDHF)



# pyscf/pbc/tdscf/uks.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Qiming Sun <osirpt.sun@gmail.com>
#

from pyscf import lib
from pyscf.tdscf import uks
from pyscf.pbc.tdscf.uhf import TDA
from pyscf.pbc.tdscf.uhf import TDHF as TDDFT

RPA = TDUKS = TDDFT


class TDDFTNoHybrid(uks.TDDFTNoHybrid):
    def gen_vind(self, mf):
        vind, hdiag = uks.TDDFTNoHybrid.gen_vind(self, mf)
        def vindp(x):
            with lib.temporary_env(mf, exxdiv=None):
                return vind(x)
        return vindp, hdiag

    def nuc_grad_method(self):
        raise NotImplementedError

def tddft(mf):
    '''Driver to create TDDFT or TDDFTNoHybrid object'''
    if mf._numint.libxc.is_hybrid_xc(mf.xc):
        return TDDFT(mf)
    else:
        return TDDFTNoHybrid(mf)

from pyscf.pbc import dft
dft.uks.UKS.TDA           = lib.class_as_method(TDA)
dft.uks.UKS.TDHF          = None
dft.uks.UKS.TDDFT         = tddft
#dft.rks.RKS.dTDA          = lib.class_as_method(dTDA)
#dft.rks.RKS.dRPA          = lib.class_as_method(dRPA)



# pyscf/pbc/tools/__init__.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pyscf.pbc.tools.pbc import *


# pyscf/pbc/tools/k2gamma.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

'''
Convert the k-sampled MO/integrals to corresponding Gamma-point supercell
MO/integrals.

Zhihao Cui <zcui@caltech.edu>
Qiming Sun <osirpt.sun@gmail.com>

See also the original implementation at
https://github.com/zhcui/local-orbital-and-cdft/blob/master/k2gamma.py
'''

from functools import reduce
import numpy as np
import scipy.linalg
from pyscf import lib
from pyscf.pbc import tools


def kpts_to_kmesh(cell, kpts):
    '''Guess kmesh'''
    scaled_k = cell.get_scaled_kpts(kpts).round(8)
    kmesh = (len(np.unique(scaled_k[:,0])),
             len(np.unique(scaled_k[:,1])),
             len(np.unique(scaled_k[:,2])))
    return kmesh

def translation_vectors_for_kmesh(cell, kmesh):
    '''
    Translation vectors to construct super-cell of which the gamma point is
    identical to the k-point mesh of primitive cell
    '''
    latt_vec = cell.lattice_vectors()
    R_rel_a = np.arange(kmesh[0])
    R_rel_b = np.arange(kmesh[1])
    R_rel_c = np.arange(kmesh[2])
    R_vec_rel = lib.cartesian_prod((R_rel_a, R_rel_b, R_rel_c))
    R_vec_abs = np.einsum('nu, uv -> nv', R_vec_rel, latt_vec)
    return R_vec_abs

def get_phase(cell, kpts, kmesh=None):
    '''
    The unitary transformation that transforms the supercell basis k-mesh
    adapted basis.
    '''
    if kmesh is None:
        kmesh = kpts_to_kmesh(cell, kpts)
    R_vec_abs = translation_vectors_for_kmesh(cell, kmesh)

    NR = len(R_vec_abs)
    phase = np.exp(1j*np.einsum('Ru, ku -> Rk', R_vec_abs, kpts))
    phase /= np.sqrt(NR)  # normalization in supercell

    # R_rel_mesh has to be construct exactly same to the Ts in super_cell function
    scell = tools.super_cell(cell, kmesh)
    return scell, phase

def double_translation_indices(kmesh):
    '''Indices to utilize the translation symmetry in the 2D matrix.

    D[M,N] = D[N-M]

    The return index maps the 2D subscripts to 1D subscripts.

    D2 = D1[double_translation_indices()]

    D1 holds all the symmetry unique elements in D2
    '''

    tx = translation_map(kmesh[0])
    ty = translation_map(kmesh[1])
    tz = translation_map(kmesh[2])
    idx = np.ravel_multi_index([tx[:,None,None,:,None,None],
                                ty[None,:,None,None,:,None],
                                tz[None,None,:,None,None,:]], kmesh)
    nk = np.prod(kmesh)
    return idx.reshape(nk, nk)

def translation_map(nk):
    ''' Generate
    [0    1 .. n  ]
    [n    0 .. n-1]
    [n-1  n .. n-2]
    [...  ...  ...]
    [1    2 .. 0  ]
    '''
    idx = np.repeat(np.arange(nk)[None,:], nk-1, axis=0)
    strides = idx.strides
    t_map = np.ndarray((nk, nk), strides=(strides[0]-strides[1], strides[1]),
                       dtype=int, buffer=np.append(idx.ravel(), 0))
    return t_map

def mo_k2gamma(cell, mo_energy, mo_coeff, kpts, kmesh=None):
    scell, phase = get_phase(cell, kpts, kmesh)

    E_g = np.hstack(mo_energy)
    C_k = np.asarray(mo_coeff)
    Nk, Nao, Nmo = C_k.shape
    NR = phase.shape[0]

    # Transform AO indices
    C_gamma = np.einsum('Rk, kum -> Rukm', phase, C_k)
    C_gamma = C_gamma.reshape(Nao*NR, Nk*Nmo)

    E_sort_idx = np.argsort(E_g)
    E_g = E_g[E_sort_idx]
    C_gamma = C_gamma[:,E_sort_idx]
    s = scell.pbc_intor('int1e_ovlp')
    assert(abs(reduce(np.dot, (C_gamma.conj().T, s, C_gamma))
               - np.eye(Nmo*Nk)).max() < 1e-5)

    # For degenerated MOs, the transformed orbitals in super cell may not be
    # real. Construct a sub Fock matrix in super-cell to find a proper
    # transformation that makes the transformed MOs real.
    E_k_degen = abs(E_g[1:] - E_g[:-1]) < 1e-3
    degen_mask = np.append(False, E_k_degen) | np.append(E_k_degen, False)
    if np.any(E_k_degen):
        if abs(C_gamma[:,~degen_mask].imag).max() < 1e-4:
            shift = min(E_g[degen_mask]) - .1
            f = np.dot(C_gamma[:,degen_mask] * (E_g[degen_mask] - shift),
                       C_gamma[:,degen_mask].conj().T)
            assert(abs(f.imag).max() < 1e-4)

            e, na_orb = scipy.linalg.eigh(f.real, s, type=2)
            C_gamma = C_gamma.real
            C_gamma[:,degen_mask] = na_orb[:, e>1e-7]
        else:
            f = np.dot(C_gamma * E_g, C_gamma.conj().T)
            assert(abs(f.imag).max() < 1e-4)
            e, C_gamma = scipy.linalg.eigh(f.real, s, type=2)

    s_k = cell.pbc_intor('int1e_ovlp', kpts=kpts)
    # overlap between k-point unitcell and gamma-point supercell
    s_k_g = np.einsum('kuv,Rk->kuRv', s_k, phase.conj()).reshape(Nk,Nao,NR*Nao)
    # The unitary transformation from k-adapted orbitals to gamma-point orbitals
    mo_phase = lib.einsum('kum,kuv,vi->kmi', C_k.conj(), s_k_g, C_gamma)

    return scell, E_g, C_gamma, mo_phase

def k2gamma(kmf, kmesh=None):
    r'''
    convert the k-sampled mean-field object to the corresponding supercell
    gamma-point mean-field object.

    math:
         C_{\nu ' n'} = C_{\vecR\mu, \veck m} = \frac{1}{\sqrt{N_{\UC}}}
         \e^{\ii \veck\cdot\vecR} C^{\veck}_{\mu  m}
    '''
    from pyscf.pbc import scf
    def transform(mo_energy, mo_coeff, mo_occ):
        scell, E_g, C_gamma = mo_k2gamma(kmf.cell, mo_energy, mo_coeff,
                                         kmf.kpts, kmesh)[:3]
        E_sort_idx = np.argsort(np.hstack(mo_energy))
        mo_occ = np.hstack(mo_occ)[E_sort_idx]
        return scell, E_g, C_gamma, mo_occ

    if isinstance(kmf, scf.khf.KRHF):
        scell, E_g, C_gamma, mo_occ = transform(kmf.mo_energy, kmf.mo_coeff, kmf.mo_occ)
        mf = scf.RHF(scell)
    elif isinstance(kmf, scf.kuhf.KUHF):
        scell, Ea, Ca, occ_a = transform(kmf.mo_energy[0], kmf.mo_coeff[0], kmf.mo_occ[0])
        scell, Eb, Cb, occ_b = transform(kmf.mo_energy[1], kmf.mo_coeff[1], kmf.mo_occ[1])
        mf = scf.UHF(scell)
        E_g = [Ea, Eb]
        C_gamma = [Ca, Cb]
        mo_occ = [occ_a, occ_b]
    else:
        raise NotImplementedError('SCF object %s not supported' % kmf)

    mf.mo_coeff = C_gamma
    mf.mo_energy = E_g
    mf.mo_occ = mo_occ
    return mf


def to_supercell_ao_integrals(cell, kpts, ao_ints):
    '''Transform from the unitcell k-point AO integrals to the supercell
    gamma-point AO integrals.
    '''
    scell, phase = get_phase(cell, kpts)
    NR, Nk = phase.shape
    nao = cell.nao
    scell_ints = np.einsum('Rk,kij,Sk->RiSj', phase, ao_ints, phase.conj())
    return scell_ints.reshape(NR*nao,NR*nao).real


def to_supercell_mo_integrals(kmf, mo_ints):
    '''Transform from the unitcell k-point MO integrals to the supercell
    gamma-point MO integrals.
    '''
    cell = kmf.cell
    kpts = kmf.kpts

    mo_k = np.array(kmf.mo_coeff)
    Nk, nao, nmo = mo_k.shape
    e_k = np.array(kmf.mo_energy)
    scell, E_g, C_gamma, mo_phase = mo_k2gamma(cell, e_k, mo_k, kpts)

    scell_ints = lib.einsum('xui,xuv,xvj->ij', mo_phase.conj(), mo_ints, mo_phase)
    assert(abs(scell_ints.imag).max() < 1e-7)
    return scell_ints.real


if __name__ == '__main__':
    from pyscf.pbc import gto, dft
    cell = gto.Cell()
    cell.atom = '''
    H 0.  0.  0.
    H 0.5 0.3 0.4
    '''

    cell.basis = 'gth-dzvp'
    cell.pseudo = 'gth-pade'
    cell.a = np.eye(3) * 4.
    cell.unit='B'
    cell.build()

    kmesh = [2, 2, 1]
    kpts = cell.make_kpts(kmesh)

    print("Transform k-point integrals to supercell integral")
    scell, phase = get_phase(cell, kpts)
    NR, Nk = phase.shape
    nao = cell.nao
    s_k = cell.pbc_intor('int1e_ovlp', kpts=kpts)
    s = scell.pbc_intor('int1e_ovlp')
    s1 = np.einsum('Rk,kij,Sk->RiSj', phase, s_k, phase.conj())
    print(abs(s-s1.reshape(s.shape)).max())

    s = scell.pbc_intor('int1e_ovlp').reshape(NR,nao,NR,nao)
    s1 = np.einsum('Rk,RiSj,Sk->kij', phase.conj(), s, phase)
    print(abs(s1-s_k).max())

    kmf = dft.KRKS(cell, kpts)
    ekpt = kmf.run()

    mf = k2gamma(kmf, kmesh)
    c_g_ao = mf.mo_coeff

    # The following is to check whether the MO is correctly coverted:

    print("Supercell gamma MO in AO basis from conversion:")
    scell = tools.super_cell(cell, kmesh)
    mf_sc = dft.RKS(scell)

    s = mf_sc.get_ovlp()
    mf_sc.run()
    sc_mo = mf_sc.mo_coeff

    nocc = scell.nelectron // 2
    print("Supercell gamma MO from direct calculation:")
    print(np.linalg.det(c_g_ao[:,:nocc].T.conj().dot(s).dot(sc_mo[:,:nocc])))
    print(np.linalg.svd(c_g_ao[:,:nocc].T.conj().dot(s).dot(sc_mo[:,:nocc]))[1])



# pyscf/pbc/tools/lattice.py

# Copyright 2014-2018,2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

A2B = 1.889725989

def get_ase_atom(formula):
    formula = formula.lower()
    assert formula in ['lih','lif','licl','mgo',
                       'c','si','ge','sic','gaas','gan','cds',
                       'zns','zno','bn','alp']
    if formula == 'lih':
        ase_atom = get_ase_rocksalt('Li','H')
    elif formula == 'lif':
        ase_atom = get_ase_rocksalt('Li','F')
    elif formula == 'licl':
        ase_atom = get_ase_rocksalt('Li','Cl')
    elif formula == 'mgo':
        ase_atom = get_ase_rocksalt('Mg','O')
    elif formula == 'c':
        ase_atom = get_ase_diamond_primitive('C')
    elif formula == 'si':
        ase_atom = get_ase_diamond_primitive('Si')
    elif formula == 'ge':
        ase_atom = get_ase_diamond_primitive('Ge')
    elif formula == 'sic':
        ase_atom = get_ase_zincblende('Si','C')
    elif formula == 'gaas':
        ase_atom = get_ase_zincblende('Ga','As')
    elif formula == 'gan':
        ase_atom = get_ase_zincblende('Ga','N')
    elif formula == 'bn':
        ase_atom = get_ase_zincblende('B','N')
    elif formula == 'alp':
        ase_atom = get_ase_zincblende('Al','P')
    elif formula == 'zno':
        ase_atom = get_ase_wurtzite('Zn','O')
    elif formula == 'cds':
        ase_atom = get_ase_zincblende('Cd','S')
    elif formula == 'zns':
        ase_atom = get_ase_zincblende('Zn','S')

    return ase_atom

def get_ase_wurtzite(A='Zn', B='O'):
    # Lattice constants taken from wikipedia (TODO: is wikipedia a valid
    # citation at this point? en.wikipedia.org/wiki/Lattice_constant)
    assert A in ['Zn']
    assert B in ['O']
    from ase.lattice import bulk
    if A=='Zn' and B=='O':
        ase_atom = bulk('ZnO', 'wurtzite', a=3.25*A2B, c=5.2*A2B)
    else:
        raise NotImplementedError('No formula found for system %s %s. '
                                  'Choose a different system?  Or add it to the list!' % (A, B))
    return ase_atom

def get_bandpath_fcc(ase_atom, npoints=30):
    # Set-up the band-path via special points
    from ase.dft.kpoints import ibz_points, kpoint_convert, get_bandpath
    points = ibz_points['fcc']
    G = points['Gamma']
    X = points['X']
    W = points['W']
    K = points['K']
    L = points['L']
    kpts_reduced, kpath, sp_points = get_bandpath([L, G, X, W, K, G],
                                                  ase_atom.cell, npoints=npoints)
    kpts_cartes = kpoint_convert(ase_atom.cell, skpts_kc=kpts_reduced)

    return kpts_reduced, kpts_cartes, kpath, sp_points

def get_ase_zincblende(A='Ga', B='As'):
    # Lattice constants from Shishkin and Kresse, PRB 75, 235102 (2007)
    assert A in ['Si', 'Ga', 'Cd', 'Zn', 'B', 'Al']
    assert B in ['C', 'As', 'S', 'O', 'N', 'P']
    from ase.lattice import bulk
    if A=='Si' and B=='C':
        ase_atom = bulk('SiC', 'zincblende', a=4.350*A2B)
    elif A=='Ga' and B=='As':
        ase_atom = bulk('GaAs', 'zincblende', a=5.648*A2B)
    elif A=='Ga' and B=='N':
        ase_atom = bulk('GaN', 'zincblende', a=4.520*A2B)
    elif A=='Cd' and B=='S':
        ase_atom = bulk('CdS', 'zincblende', a=5.832*A2B)
    elif A=='Zn' and B=='S':
        ase_atom = bulk('ZnS', 'zincblende', a=5.420*A2B)
    elif A=='Zn' and B=='O':
        ase_atom = bulk('ZnO', 'zincblende', a=4.580*A2B)
    elif A=='B' and B=='N':
        ase_atom = bulk('BN', 'zincblende', a=3.615*A2B)
    elif A=='Al' and B=='P':
        ase_atom = bulk('AlP', 'zincblende', a=5.451*A2B)
    else:
        raise NotImplementedError('No formula found for system %s %s. '
                                  'Choose a different system?  Or add it to the list!' % (A, B))

    return ase_atom

def get_ase_rocksalt(A='Li', B='Cl'):
    assert A in ['Li', 'Mg']
    # Add Na, K
    assert B in ['H', 'F', 'Cl', 'O']
    # Add Br, I
    from ase.lattice import bulk
    if A=='Li':
        if B=='H':
            ase_atom = bulk('LiH', 'rocksalt', a=4.0834*A2B)
        elif B=='F':
            ase_atom = bulk('LiF', 'rocksalt', a=4.0351*A2B)
        elif B=='Cl':
            ase_atom = bulk('LiCl', 'rocksalt', a=5.13*A2B)
    elif A=='Mg' and B=='O':
        ase_atom = bulk('MgO', 'rocksalt', a=4.213*A2B)
    else:
        raise NotImplementedError('No formula found for system %s %s. '
                                  'Choose a different system?  Or add it to the list!' % (A, B))

    return ase_atom

def get_ase_alkali_halide(A='Li', B='Cl'):
    return get_ase_rocksalt(A,B)

def get_ase_diamond_primitive(atom='C'):
    """Get the ASE atoms for primitive (2-atom) diamond unit cell."""
    from ase.build import bulk
    if atom == 'C':
        ase_atom = bulk('C', 'diamond', a=3.5668*A2B)
    elif atom == 'Si':
        ase_atom = bulk('Si', 'diamond', a=5.431*A2B)
    elif atom == 'Ge':
        ase_atom = bulk('Ge', 'diamond', a=5.658*A2B)
    else:
        raise NotImplementedError('No formula found for system %s. '
                                  'Choose a different system?  Or add it to the list!' % atom)
    return ase_atom

def get_ase_diamond_cubic(atom='C'):
    """Get the ASE atoms for cubic (8-atom) diamond unit cell."""
    from ase.lattice.cubic import Diamond
    if atom == 'C':
        ase_atom = Diamond(symbol='C', latticeconstant=3.5668*A2B)
    elif atom == 'Si':
        ase_atom = Diamond(symbol='Si', latticeconstant=5.431*A2B)
    else:
        raise NotImplementedError('No formula found for system %s. '
                                  'Choose a different system?  Or add it to the list!' % atom)
    return ase_atom

def get_ase_graphene(vacuum=5.0):
    """Get the ASE atoms for primitive (2-atom) graphene unit cell."""
    from ase.lattice.hexagonal import Graphene
    ase_atom = Graphene('C', latticeconstant={'a':2.46*A2B,'c':vacuum*A2B})
    return ase_atom

def get_ase_graphene_xxx(vacuum=5.0):
    """Get the ASE atoms for primitive (2-atom) graphene unit cell."""
    from ase.lattice import bulk
    ase_atom = bulk('C', 'hcp', a=2.46*A2B, c=vacuum*A2B)
    ase_atom.positions[1,2] = 0.0
    return ase_atom


# pyscf/pbc/tools/make_test_cell.py

#!/usr/bin/env python
# Copyright 2014-2018,2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Authors: James D. McClain
#          Timothy Berkelbach <tim.berkelbach@gmail.com>
#

#import unittest
import numpy as np

from pyscf import gto, scf

from pyscf.pbc import gto as pbcgto
from pyscf.pbc import scf as pbchf
from pyscf.pbc.ao2mo import eris
import pyscf.pbc.tools
import pyscf.pbc.cc


def make_cell(L, mesh):
    cell = pbcgto.Cell()
    cell.unit = 'B'
    cell.atom.extend([['Be', (L/2.,  L/2., L/2.)]])
    cell.a = L * np.identity(3)

    cell.basis = 'sto-3g'
    cell.pseudo = 'gth-pade'
    cell.mesh = mesh

    cell.output = '/dev/null'
    cell.build()
    return cell


def test_cell_n0(L=5, mesh=[9]*3):
    cell = pbcgto.Cell()
    cell.unit = 'B'
    cell.atom.extend([['Be', (L/2.,  L/2., L/2.)]])
    cell.a = L * np.identity(3)
    cell.a[1,0] = 5.0

    cell.basis = 'gth-szv'
    cell.pseudo = 'gth-pade-q2'
    cell.mesh = mesh

    cell.output = '/dev/null'
    cell.build()
    return cell


def test_cell_n1(L=5, mesh=[9]*3):
    cell = pbcgto.Cell()
    cell.unit = 'B'
    cell.atom.extend([['Be', (L/2.,  L/2., L/2.)]])
    cell.a = L * np.identity(3)

    cell.basis = 'sto-3g'
    cell.pseudo = 'gth-pade-q2'
    cell.mesh = mesh

    cell.output = '/dev/null'
    cell.build()
    return cell


def test_cell_n2(L=5, mesh=[9]*3):
    cell = pbcgto.Cell()
    cell.unit = 'B'
    cell.atom.extend([['O', (L/2., L/2., L/2.)],
                      ['H', (L/2.-0.689440, L/2.+0.578509, L/2.)],
                      ['H', (L/2.+0.689440, L/2.-0.578509, L/2.)], ])
    cell.a = L * np.identity(3)

    cell.basis = 'sto-3g'
    cell.pseudo = 'gth-pade'
    cell.mesh = mesh

    cell.output = '/dev/null'
    cell.build()
    return cell


def test_cell_n3(mesh=[9]*3):
    """
    Take ASE Diamond structure, input into PySCF and run
    """
    cell = pbcgto.Cell()
    cell.unit = 'A'
    cell.atom = 'C 0.,  0.,  0.; C 0.8917,  0.8917,  0.8917'
    cell.a = '''0.      1.7834  1.7834
                1.7834  0.      1.7834
                1.7834  1.7834  0.    '''
    cell.basis = "gth-szv"
    cell.pseudo = "gth-pade"
    cell.mesh = mesh

    cell.output = '/dev/null'
    cell.build()
    return cell

def test_cell_n3_diffuse():
    """
    Take ASE Diamond structure, input into PySCF and run
    """
    cell = pbcgto.Cell()
    cell.unit = 'A'
    cell.atom = 'C 0.,  0.,  0.; C 0.8917,  0.8917,  0.8917'
    cell.a = '''0.      1.7834  1.7834
                1.7834  0.      1.7834
                1.7834  1.7834  0.    '''
    cell.basis = {'C': [[0, (0.1, 1.0)],
                        [1, (0.1, 1.0)]]}
    cell.pseudo = "gth-pade"

    cell.verbose = 7
    cell.output = '/dev/null'
    cell.build()
    return cell


def test_cell_cu_metallic(mesh=[9]*3):
    """
    Copper unit cell w/ special basis giving non-equal number of occupied orbitals per k-point
    """
    cell = pbcgto.Cell()
    cell.pseudo = 'gth-pade'
    cell.atom='''
    Cu 0.        , 0.        , 0.
    Cu 1.6993361, 1.6993361, 1.6993361
    '''
    cell.a = '''
    0.        , 3.39867219, 3.39867219
    3.39867219, 0.        , 3.39867219
    3.39867219, 3.39867219, 0.
    '''
    cell.basis = { 'Cu': [[0, (0.8, 1.0)],
                          [1, (1.0, 1.0)],
                          [2, (1.2, 1.0)]] }
    cell.unit = 'B'
    cell.mesh = mesh
    cell.verbose = 9
    cell.incore_anyway = True
    cell.output = '/dev/null'
    cell.build()
    return cell


# pyscf/pbc/tools/pbc.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import warnings
import copy
import numpy as np
import scipy.linalg
from pyscf import lib
from pyscf.lib import logger
from pyscf.gto import ATM_SLOTS, BAS_SLOTS, ATOM_OF, PTR_COORD
from pyscf.pbc.lib.kpts_helper import get_kconserv, get_kconserv3  # noqa
from pyscf import __config__

FFT_ENGINE = getattr(__config__, 'pbc_tools_pbc_fft_engine', 'BLAS')

def _fftn_blas(f, mesh):
    Gx = np.fft.fftfreq(mesh[0])
    Gy = np.fft.fftfreq(mesh[1])
    Gz = np.fft.fftfreq(mesh[2])
    expRGx = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[0]), Gx))
    expRGy = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[1]), Gy))
    expRGz = np.exp(np.einsum('x,k->xk', -2j*np.pi*np.arange(mesh[2]), Gz))
    out = np.empty(f.shape, dtype=np.complex128)
    buf = np.empty(mesh, dtype=np.complex128)
    for i, fi in enumerate(f):
        buf[:] = fi.reshape(mesh)
        g = lib.dot(buf.reshape(mesh[0],-1).T, expRGx, c=out[i].reshape(-1,mesh[0]))
        g = lib.dot(g.reshape(mesh[1],-1).T, expRGy, c=buf.reshape(-1,mesh[1]))
        g = lib.dot(g.reshape(mesh[2],-1).T, expRGz, c=out[i].reshape(-1,mesh[2]))
    return out.reshape(-1, *mesh)

def _ifftn_blas(g, mesh):
    Gx = np.fft.fftfreq(mesh[0])
    Gy = np.fft.fftfreq(mesh[1])
    Gz = np.fft.fftfreq(mesh[2])
    expRGx = np.exp(np.einsum('x,k->xk', 2j*np.pi*np.arange(mesh[0]), Gx))
    expRGy = np.exp(np.einsum('x,k->xk', 2j*np.pi*np.arange(mesh[1]), Gy))
    expRGz = np.exp(np.einsum('x,k->xk', 2j*np.pi*np.arange(mesh[2]), Gz))
    out = np.empty(g.shape, dtype=np.complex128)
    buf = np.empty(mesh, dtype=np.complex128)
    for i, gi in enumerate(g):
        buf[:] = gi.reshape(mesh)
        f = lib.dot(buf.reshape(mesh[0],-1).T, expRGx, 1./mesh[0], c=out[i].reshape(-1,mesh[0]))
        f = lib.dot(f.reshape(mesh[1],-1).T, expRGy, 1./mesh[1], c=buf.reshape(-1,mesh[1]))
        f = lib.dot(f.reshape(mesh[2],-1).T, expRGz, 1./mesh[2], c=out[i].reshape(-1,mesh[2]))
    return out.reshape(-1, *mesh)

if FFT_ENGINE == 'FFTW':
    # pyfftw is slower than np.fft in most cases
    try:
        import pyfftw
        pyfftw.interfaces.cache.enable()
        nproc = lib.num_threads()
        def _fftn_wrapper(a):
            return pyfftw.interfaces.numpy_fft.fftn(a, axes=(1,2,3), threads=nproc)
        def _ifftn_wrapper(a):
            return pyfftw.interfaces.numpy_fft.ifftn(a, axes=(1,2,3), threads=nproc)
    except ImportError:
        def _fftn_wrapper(a):
            return np.fft.fftn(a, axes=(1,2,3))
        def _ifftn_wrapper(a):
            return np.fft.ifftn(a, axes=(1,2,3))

elif FFT_ENGINE == 'NUMPY':
    def _fftn_wrapper(a):
        return np.fft.fftn(a, axes=(1,2,3))
    def _ifftn_wrapper(a):
        return np.fft.ifftn(a, axes=(1,2,3))

elif FFT_ENGINE == 'NUMPY+BLAS':
    _EXCLUDE = [17, 19, 23, 29, 31, 37, 41, 43, 47, 53, 59, 61, 67, 71, 73, 79,
                83, 89, 97,101,103,107,109,113,127,131,137,139,149,151,157,163,
                167,173,179,181,191,193,197,199,211,223,227,229,233,239,241,251,
                257,263,269,271,277,281,283,293]
    _EXCLUDE = set(_EXCLUDE + [n*2 for n in _EXCLUDE] + [n*3 for n in _EXCLUDE])
    def _fftn_wrapper(a):
        mesh = a.shape[1:]
        if mesh[0] in _EXCLUDE and mesh[1] in _EXCLUDE and mesh[2] in _EXCLUDE:
            return _fftn_blas(a, mesh)
        else:
            return np.fft.fftn(a, axes=(1,2,3))
    def _ifftn_wrapper(a):
        mesh = a.shape[1:]
        if mesh[0] in _EXCLUDE and mesh[1] in _EXCLUDE and mesh[2] in _EXCLUDE:
            return _ifftn_blas(a, mesh)
        else:
            return np.fft.ifftn(a, axes=(1,2,3))

#?elif:  # 'FFTW+BLAS'
else:  # 'BLAS'
    def _fftn_wrapper(a):
        mesh = a.shape[1:]
        return _fftn_blas(a, mesh)
    def _ifftn_wrapper(a):
        mesh = a.shape[1:]
        return _ifftn_blas(a, mesh)


def fft(f, mesh):
    '''Perform the 3D FFT from real (R) to reciprocal (G) space.

    After FFT, (u, v, w) -> (j, k, l).
    (jkl) is in the index order of Gv.

    FFT normalization factor is 1., as in MH and in `numpy.fft`.

    Args:
        f : (nx*ny*nz,) ndarray
            The function to be FFT'd, flattened to a 1D array corresponding
            to the index order of :func:`cartesian_prod`.
        mesh : (3,) ndarray of ints (= nx,ny,nz)
            The number G-vectors along each direction.

    Returns:
        (nx*ny*nz,) ndarray
            The FFT 1D array in same index order as Gv (natural order of
            numpy.fft).

    '''
    if f.size == 0:
        return np.zeros_like(f)

    f3d = f.reshape(-1, *mesh)
    assert(f3d.shape[0] == 1 or f[0].size == f3d[0].size)
    g3d = _fftn_wrapper(f3d)
    ngrids = np.prod(mesh)
    if f.ndim == 1 or (f.ndim == 3 and f.size == ngrids):
        return g3d.ravel()
    else:
        return g3d.reshape(-1, ngrids)

def ifft(g, mesh):
    '''Perform the 3D inverse FFT from reciprocal (G) space to real (R) space.

    Inverse FFT normalization factor is 1./N, same as in `numpy.fft` but
    **different** from MH (they use 1.).

    Args:
        g : (nx*ny*nz,) ndarray
            The function to be inverse FFT'd, flattened to a 1D array
            corresponding to the index order of `span3`.
        mesh : (3,) ndarray of ints (= nx,ny,nz)
            The number G-vectors along each direction.

    Returns:
        (nx*ny*nz,) ndarray
            The inverse FFT 1D array in same index order as Gv (natural order
            of numpy.fft).

    '''
    if g.size == 0:
        return np.zeros_like(g)

    g3d = g.reshape(-1, *mesh)
    assert(g3d.shape[0] == 1 or g[0].size == g3d[0].size)
    f3d = _ifftn_wrapper(g3d)
    ngrids = np.prod(mesh)
    if g.ndim == 1 or (g.ndim == 3 and g.size == ngrids):
        return f3d.ravel()
    else:
        return f3d.reshape(-1, ngrids)


def fftk(f, mesh, expmikr):
    r'''Perform the 3D FFT of a real-space function which is (periodic*e^{ikr}).

    fk(k+G) = \sum_r fk(r) e^{-i(k+G)r} = \sum_r [f(k)e^{-ikr}] e^{-iGr}
    '''
    return fft(f*expmikr, mesh)


def ifftk(g, mesh, expikr):
    r'''Perform the 3D inverse FFT of f(k+G) into a function which is (periodic*e^{ikr}).

    fk(r) = (1/Ng) \sum_G fk(k+G) e^{i(k+G)r} = (1/Ng) \sum_G [fk(k+G)e^{iGr}] e^{ikr}
    '''
    return ifft(g, mesh) * expikr


def get_coulG(cell, k=np.zeros(3), exx=False, mf=None, mesh=None, Gv=None,
              wrap_around=True, omega=None, **kwargs):
    '''Calculate the Coulomb kernel for all G-vectors, handling G=0 and exchange.

    Args:
        k : (3,) ndarray
            k-point
        exx : bool or str
            Whether this is an exchange matrix element.
        mf : instance of :class:`SCF`

    Returns:
        coulG : (ngrids,) ndarray
            The Coulomb kernel.
        mesh : (3,) ndarray of ints (= nx,ny,nz)
            The number G-vectors along each direction.
        omega : float
            Enable Coulomb kernel erf(|omega|*r12)/r12 if omega > 0
            and erfc(|omega|*r12)/r12 if omega < 0.
            Note this parameter is slightly different to setting cell.omega
            for the treatment of exxdiv (at G0).  cell.omega affects Ewald
            probe charge at G0. It is used mostly with RSH functional for
            the long-range part of HF exchange. This parameter is used by
            real-space JK builder which requires Ewald probe charge to be
            computed with regular Coulomb interaction (1/r12) while the rest
            coulG is scaled as long-range Coulomb kernel.
    '''
    exxdiv = exx
    if isinstance(exx, str):
        exxdiv = exx
    elif exx and mf is not None:
        exxdiv = mf.exxdiv

    if mesh is None:
        mesh = cell.mesh
    if 'gs' in kwargs:
        warnings.warn('cell.gs is deprecated.  It is replaced by cell.mesh,'
                      'the number of PWs (=2*gs+1) along each direction.')
        mesh = [2*n+1 for n in kwargs['gs']]
    if Gv is None:
        Gv = cell.get_Gv(mesh)

    if abs(k).sum() > 1e-9:
        kG = k + Gv
    else:
        kG = Gv

    equal2boundary = np.zeros(Gv.shape[0], dtype=bool)
    if wrap_around and abs(k).sum() > 1e-9:
        # Here we 'wrap around' the high frequency k+G vectors into their lower
        # frequency counterparts.  Important if you want the gamma point and k-point
        # answers to agree
        b = cell.reciprocal_vectors()
        box_edge = np.einsum('i,ij->ij', np.asarray(mesh)//2+0.5, b)
        assert(all(np.linalg.solve(box_edge.T, k).round(9).astype(int)==0))
        reduced_coords = np.linalg.solve(box_edge.T, kG.T).T.round(9)
        on_edge = reduced_coords.astype(int)
        if cell.dimension >= 1:
            equal2boundary |= reduced_coords[:,0] == 1
            equal2boundary |= reduced_coords[:,0] ==-1
            kG[on_edge[:,0]== 1] -= 2 * box_edge[0]
            kG[on_edge[:,0]==-1] += 2 * box_edge[0]
        if cell.dimension >= 2:
            equal2boundary |= reduced_coords[:,1] == 1
            equal2boundary |= reduced_coords[:,1] ==-1
            kG[on_edge[:,1]== 1] -= 2 * box_edge[1]
            kG[on_edge[:,1]==-1] += 2 * box_edge[1]
        if cell.dimension == 3:
            equal2boundary |= reduced_coords[:,2] == 1
            equal2boundary |= reduced_coords[:,2] ==-1
            kG[on_edge[:,2]== 1] -= 2 * box_edge[2]
            kG[on_edge[:,2]==-1] += 2 * box_edge[2]

    absG2 = np.einsum('gi,gi->g', kG, kG)

    if getattr(mf, 'kpts', None) is not None:
        kpts = mf.kpts
    else:
        kpts = k.reshape(1,3)
    Nk = len(kpts)

    if exxdiv == 'vcut_sph':  # PRB 77 193110
        Rc = (3*Nk*cell.vol/(4*np.pi))**(1./3)
        with np.errstate(divide='ignore',invalid='ignore'):
            coulG = 4*np.pi/absG2*(1.0 - np.cos(np.sqrt(absG2)*Rc))
        coulG[absG2==0] = 4*np.pi*0.5*Rc**2

        if cell.dimension < 3:
            raise NotImplementedError

    elif exxdiv == 'vcut_ws':  # PRB 87, 165122
        assert(cell.dimension == 3)
        if not getattr(mf, '_ws_exx', None):
            mf._ws_exx = precompute_exx(cell, kpts)
        exx_alpha = mf._ws_exx['alpha']
        exx_kcell = mf._ws_exx['kcell']
        exx_q = mf._ws_exx['q']
        exx_vq = mf._ws_exx['vq']

        with np.errstate(divide='ignore',invalid='ignore'):
            coulG = 4*np.pi/absG2*(1.0 - np.exp(-absG2/(4*exx_alpha**2)))
        coulG[absG2==0] = np.pi / exx_alpha**2
        # Index k+Gv into the precomputed vq and add on
        gxyz = np.dot(kG, exx_kcell.lattice_vectors().T)/(2*np.pi)
        gxyz = gxyz.round(decimals=6).astype(int)
        mesh = np.asarray(exx_kcell.mesh)
        gxyz = (gxyz + mesh)%mesh
        qidx = (gxyz[:,0]*mesh[1] + gxyz[:,1])*mesh[2] + gxyz[:,2]
        #qidx = [np.linalg.norm(exx_q-kGi,axis=1).argmin() for kGi in kG]
        maxqv = abs(exx_q).max(axis=0)
        is_lt_maxqv = (abs(kG) <= maxqv).all(axis=1)
        coulG = coulG.astype(exx_vq.dtype)
        coulG[is_lt_maxqv] += exx_vq[qidx[is_lt_maxqv]]

        if cell.dimension < 3:
            raise NotImplementedError

    else:
        # Ewald probe charge method to get the leading term of the finite size
        # error in exchange integrals

        G0_idx = np.where(absG2==0)[0]
        if cell.dimension != 2 or cell.low_dim_ft_type == 'inf_vacuum':
            with np.errstate(divide='ignore'):
                coulG = 4*np.pi/absG2
                coulG[G0_idx] = 0

        elif cell.dimension == 2:
            # The following 2D analytical fourier transform is taken from:
            # R. Sundararaman and T. Arias PRB 87, 2013
            b = cell.reciprocal_vectors()
            Ld2 = np.pi/np.linalg.norm(b[2])
            Gz = kG[:,2]
            Gp = np.linalg.norm(kG[:,:2], axis=1)
            weights = 1. - np.cos(Gz*Ld2) * np.exp(-Gp*Ld2)
            with np.errstate(divide='ignore', invalid='ignore'):
                coulG = weights*4*np.pi/absG2
            if len(G0_idx) > 0:
                coulG[G0_idx] = -2*np.pi*Ld2**2 #-pi*L_z^2/2

        elif cell.dimension == 1:
            logger.warn(cell, 'No method for PBC dimension 1, dim-type %s.'
                        '  cell.low_dim_ft_type="inf_vacuum"  should be set.',
                        cell.low_dim_ft_type)
            raise NotImplementedError

            # Carlo A. Rozzi, PRB 73, 205119 (2006)
            a = cell.lattice_vectors()
            # Rc is the cylindrical radius
            Rc = np.sqrt(cell.vol / np.linalg.norm(a[0])) / 2
            Gx = abs(kG[:,0])
            Gp = np.linalg.norm(kG[:,1:], axis=1)
            with np.errstate(divide='ignore', invalid='ignore'):
                weights = 1 + Gp*Rc * scipy.special.j1(Gp*Rc) * scipy.special.k0(Gx*Rc)
                weights -= Gx*Rc * scipy.special.j0(Gp*Rc) * scipy.special.k1(Gx*Rc)
                coulG = 4*np.pi/absG2 * weights
                # TODO: numerical integation
                # coulG[Gx==0] = -4*np.pi * (dr * r * scipy.special.j0(Gp*r) * np.log(r)).sum()
            if len(G0_idx) > 0:
                coulG[G0_idx] = -np.pi*Rc**2 * (2*np.log(Rc) - 1)

        # The divergent part of periodic summation of (ii|ii) integrals in
        # Coulomb integrals were cancelled out by electron-nucleus
        # interaction. The periodic part of (ii|ii) in exchange cannot be
        # cancelled out by Coulomb integrals. Its leading term is calculated
        # using Ewald probe charge (the function madelung below)
        if cell.dimension > 0 and exxdiv == 'ewald' and len(G0_idx) > 0:
            coulG[G0_idx] += Nk*cell.vol*madelung(cell, kpts)

    coulG[equal2boundary] = 0

    # Scale the coulG kernel for attenuated Coulomb integrals.
    # * omega is used by RealSpaceJKBuilder which requires ewald probe charge
    # being evaluated with regular Coulomb interaction (1/r12).
    # * cell.omega, which affects the ewald probe charge, is often set by
    # DFT-RSH functionals to build long-range HF-exchange for erf(omega*r12)/r12
    if omega is not None:
        if omega > 0:
            coulG *= np.exp(-.25/omega**2 * absG2)
        else:
            coulG *= (1 - np.exp(-.25/omega**2 * absG2))
    elif cell.omega != 0:
        coulG *= np.exp(-.25/cell.omega**2 * absG2)

    return coulG

def precompute_exx(cell, kpts):
    from pyscf.pbc import gto as pbcgto
    from pyscf.pbc.dft import gen_grid
    log = lib.logger.Logger(cell.stdout, cell.verbose)
    log.debug("# Precomputing Wigner-Seitz EXX kernel")
    Nk = get_monkhorst_pack_size(cell, kpts)
    log.debug("# Nk = %s", Nk)

    kcell = pbcgto.Cell()
    kcell.atom = 'H 0. 0. 0.'
    kcell.spin = 1
    kcell.unit = 'B'
    kcell.verbose = 0
    kcell.a = cell.lattice_vectors() * Nk
    Lc = 1.0/lib.norm(np.linalg.inv(kcell.a), axis=0)
    log.debug("# Lc = %s", Lc)
    Rin = Lc.min() / 2.0
    log.debug("# Rin = %s", Rin)
    # ASE:
    alpha = 5./Rin # sqrt(-ln eps) / Rc, eps ~ 10^{-11}
    log.info("WS alpha = %s", alpha)
    kcell.mesh = np.array([4*int(L*alpha*3.0) for L in Lc])  # ~ [120,120,120]
    # QE:
    #alpha = 3./Rin * np.sqrt(0.5)
    #kcell.mesh = (4*alpha*np.linalg.norm(kcell.a,axis=1)).astype(int)
    log.debug("# kcell.mesh FFT = %s", kcell.mesh)
    rs = gen_grid.gen_uniform_grids(kcell)
    kngs = len(rs)
    log.debug("# kcell kngs = %d", kngs)
    corners_coord = lib.cartesian_prod(([0, 1], [0, 1], [0, 1]))
    corners = np.dot(corners_coord, kcell.a)
    #vR = np.empty(kngs)
    #for i, rv in enumerate(rs):
    #    # Minimum image convention to corners of kcell parallelepiped
    #    r = lib.norm(rv-corners, axis=1).min()
    #    if np.isclose(r, 0.):
    #        vR[i] = 2*alpha / np.sqrt(np.pi)
    #    else:
    #        vR[i] = scipy.special.erf(alpha*r) / r
    r = np.min([lib.norm(rs-c, axis=1) for c in corners], axis=0)
    vR = scipy.special.erf(alpha*r) / (r+1e-200)
    vR[r<1e-9] = 2*alpha / np.sqrt(np.pi)
    vG = (kcell.vol/kngs) * fft(vR, kcell.mesh)

    if abs(vG.imag).max() > 1e-6:
        # vG should be real in regular lattice. If imaginary part is observed,
        # this probably means a ws cell was built from a unconventional
        # lattice. The SR potential erfc(alpha*r) for the charge in the center
        # of ws cell decays to the region out of ws cell. The Ewald-sum based
        # on the minimum image convention cannot be used to build the kernel
        # Eq (12) of PRB 87, 165122
        raise RuntimeError('Unconventional lattice was found')

    ws_exx = {'alpha': alpha,
              'kcell': kcell,
              'q'    : kcell.Gv,
              'vq'   : vG.real.copy()}
    log.debug("# Finished precomputing")
    return ws_exx


def madelung(cell, kpts):
    Nk = get_monkhorst_pack_size(cell, kpts)
    ecell = copy.copy(cell)
    ecell._atm = np.array([[1, cell._env.size, 0, 0, 0, 0]])
    ecell._env = np.append(cell._env, [0., 0., 0.])
    ecell.unit = 'B'
    #ecell.verbose = 0
    ecell.a = np.einsum('xi,x->xi', cell.lattice_vectors(), Nk)
    ecell.mesh = np.asarray(cell.mesh) * Nk

    if cell.omega == 0:
        ew_eta, ew_cut = ecell.get_ewald_params(cell.precision, ecell.mesh)
        lib.logger.debug1(cell, 'Monkhorst pack size %s ew_eta %s ew_cut %s',
                          Nk, ew_eta, ew_cut)
        return -2*ecell.ewald(ew_eta, ew_cut)

    else:
        # cell.ewald function does not use the Coulomb kernel function
        # get_coulG. When computing the nuclear interactions with attenuated
        # Coulomb operator, the Ewald summation technique is not needed
        # because the Coulomb kernel 4pi/G^2*exp(-G^2/4/omega**2) decays
        # quickly.
        Gv, Gvbase, weights = ecell.get_Gv_weights(ecell.mesh)
        coulG = get_coulG(ecell, Gv=Gv)
        ZSI = np.einsum("i,ij->j", ecell.atom_charges(), ecell.get_SI(Gv))
        return -np.einsum('i,i,i->', ZSI.conj(), ZSI, coulG*weights).real


def get_monkhorst_pack_size(cell, kpts):
    skpts = cell.get_scaled_kpts(kpts).round(decimals=6)
    Nk = np.array([len(np.unique(ki)) for ki in skpts.T])
    return Nk


def get_lattice_Ls(cell, nimgs=None, rcut=None, dimension=None, discard=True):
    '''Get the (Cartesian, unitful) lattice translation vectors for nearby images.
    The translation vectors can be used for the lattice summation.'''
    a = cell.lattice_vectors()
    b = cell.reciprocal_vectors(norm_to=1)
    heights_inv = lib.norm(b, axis=1)

    if nimgs is None:
        if rcut is None:
            rcut = cell.rcut
        # For atoms outside the cell, distance between certain basis of nearby
        # images may be smaller than rcut threshold even the corresponding Ls is
        # larger than rcut. The boundary penalty ensures that Ls would be able to
        # cover the basis that sitting out of the cell.
        # See issue https://github.com/pyscf/pyscf/issues/1017
        scaled_atom_coords = cell.atom_coords().dot(b.T)
        boundary_penalty = np.max([abs(scaled_atom_coords).max(axis=0),
                                   abs(1 - scaled_atom_coords).max(axis=0)], axis=0)
        nimgs = np.ceil(rcut * heights_inv + boundary_penalty).astype(int)
    else:
        rcut = max((np.asarray(nimgs))/heights_inv)

    if dimension is None:
        dimension = cell.dimension
    if dimension == 0:
        nimgs = [0, 0, 0]
    elif dimension == 1:
        nimgs = [nimgs[0], 0, 0]
    elif dimension == 2:
        nimgs = [nimgs[0], nimgs[1], 0]

    Ts = lib.cartesian_prod((np.arange(-nimgs[0], nimgs[0]+1),
                             np.arange(-nimgs[1], nimgs[1]+1),
                             np.arange(-nimgs[2], nimgs[2]+1)))
    Ls = np.dot(Ts, a)
    if discard:
        Ls = _discard_edge_images(cell, Ls, rcut)
    return np.asarray(Ls, order='C')

def _discard_edge_images(cell, Ls, rcut):
    '''
    Discard images if no basis in the image would contribute to lattice sum.
    '''
    if rcut <= 0:
        return np.zeros((1, 3))

    a = cell.lattice_vectors()
    scaled_atom_coords = np.linalg.solve(a.T, cell.atom_coords().T).T
    atom_boundary_max = scaled_atom_coords.max(axis=0)
    atom_boundary_min = scaled_atom_coords.min(axis=0)
    # ovlp_penalty ensures the overlap integrals for atoms in the adjcent
    # images are converged.
    ovlp_penalty = atom_boundary_max - atom_boundary_min
    # atom_boundary_min-1 ensures the values of basis at the grids on the edge
    # of the primitive cell converged
    boundary_max = np.ceil(np.max([atom_boundary_max  ,  ovlp_penalty], axis=0)).astype(int)
    boundary_min = np.floor(np.min([atom_boundary_min-1, -ovlp_penalty], axis=0)).astype(int)
    penalty_x = np.arange(boundary_min[0], boundary_max[0]+1)
    penalty_y = np.arange(boundary_min[1], boundary_max[1]+1)
    penalty_z = np.arange(boundary_min[2], boundary_max[2]+1)
    shifts = lib.cartesian_prod([penalty_x, penalty_y, penalty_z]).dot(a)
    Ls_mask = (np.linalg.norm(Ls + shifts[:,None,:], axis=2) < rcut).any(axis=0)
    # cell0 (Ls == 0) should always be included.
    Ls_mask[len(Ls)//2] = True
    return Ls[Ls_mask]


def super_cell(cell, ncopy):
    '''Create an ncopy[0] x ncopy[1] x ncopy[2] supercell of the input cell
    Note this function differs from :fun:`cell_plus_imgs` that cell_plus_imgs
    creates images in both +/- direction.

    Args:
        cell : instance of :class:`Cell`
        ncopy : (3,) array

    Returns:
        supcell : instance of :class:`Cell`
    '''
    a = cell.lattice_vectors()
    #:supcell.atom = []
    #:for Lx in range(ncopy[0]):
    #:    for Ly in range(ncopy[1]):
    #:        for Lz in range(ncopy[2]):
    #:            # Using cell._atom guarantees coord is in Bohr
    #:            for atom, coord in cell._atom:
    #:                L = np.dot([Lx, Ly, Lz], a)
    #:                supcell.atom.append([atom, coord + L])
    Ts = lib.cartesian_prod((np.arange(ncopy[0]),
                             np.arange(ncopy[1]),
                             np.arange(ncopy[2])))
    Ls = np.dot(Ts, a)
    supcell = cell.copy()
    supcell.a = np.einsum('i,ij->ij', ncopy, a)
    supcell.mesh = np.array([ncopy[0]*cell.mesh[0],
                             ncopy[1]*cell.mesh[1],
                             ncopy[2]*cell.mesh[2]])
    return _build_supcell_(supcell, cell, Ls)


def cell_plus_imgs(cell, nimgs):
    '''Create a supercell via nimgs[i] in each +/- direction, as in get_lattice_Ls().
    Note this function differs from :fun:`super_cell` that super_cell only
    stacks the images in + direction.

    Args:
        cell : instance of :class:`Cell`
        nimgs : (3,) array

    Returns:
        supcell : instance of :class:`Cell`
    '''
    a = cell.lattice_vectors()
    Ts = lib.cartesian_prod((np.arange(-nimgs[0], nimgs[0]+1),
                             np.arange(-nimgs[1], nimgs[1]+1),
                             np.arange(-nimgs[2], nimgs[2]+1)))
    Ls = np.dot(Ts, a)
    supcell = cell.copy()
    supcell.a = np.einsum('i,ij->ij', nimgs, a)
    supcell.mesh = np.array([(nimgs[0]*2+1)*cell.mesh[0],
                             (nimgs[1]*2+1)*cell.mesh[1],
                             (nimgs[2]*2+1)*cell.mesh[2]])
    return _build_supcell_(supcell, cell, Ls)

def _build_supcell_(supcell, cell, Ls):
    '''
    Construct supcell ._env directly without calling supcell.build() method.
    This reserves the basis contraction coefficients defined in cell
    '''
    nimgs = len(Ls)
    symbs = [atom[0] for atom in cell._atom] * nimgs
    coords = Ls.reshape(-1,1,3) + cell.atom_coords()
    supcell.atom = supcell._atom = list(zip(symbs, coords.reshape(-1,3).tolist()))
    supcell.unit = 'B'

    # Do not call supcell.build() since it may normalize the basis contraction
    # coefficients
    _env = np.append(cell._env, coords.ravel())
    _atm = np.repeat(cell._atm[None,:,:], nimgs, axis=0)
    _atm = _atm.reshape(-1, ATM_SLOTS)
    # Point to the corrdinates appended to _env
    _atm[:,PTR_COORD] = cell._env.size + np.arange(nimgs * cell.natm) * 3

    _bas = np.repeat(cell._bas[None,:,:], nimgs, axis=0)
    # For atom pointers in each image, shift natm*image_id
    _bas[:,:,ATOM_OF] += np.arange(nimgs)[:,None] * cell.natm

    supcell._atm = np.asarray(_atm, dtype=np.int32)
    supcell._bas = np.asarray(_bas.reshape(-1, BAS_SLOTS), dtype=np.int32)
    supcell._env = _env
    return supcell


def cutoff_to_mesh(a, cutoff):
    r'''
    Convert KE cutoff to FFT-mesh

        uses KE = k^2 / 2, where k_max ~ \pi / grid_spacing

    Args:
        a : (3,3) ndarray
            The real-space unit cell lattice vectors. Each row represents a
            lattice vector.
        cutoff : float
            KE energy cutoff in a.u.

    Returns:
        mesh : (3,) array
    '''
    b = 2 * np.pi * np.linalg.inv(a.T)
    cutoff = cutoff * _cubic2nonorth_factor(a)
    mesh = np.ceil(np.sqrt(2*cutoff)/lib.norm(b, axis=1) * 2).astype(int)
    return mesh

def mesh_to_cutoff(a, mesh):
    '''
    Convert #grid points to KE cutoff
    '''
    b = 2 * np.pi * np.linalg.inv(a.T)
    Gmax = lib.norm(b, axis=1) * np.asarray(mesh) * .5
    ke_cutoff = Gmax**2/2
    # scale down Gmax to get the real energy cutoff for non-orthogonal lattice
    return ke_cutoff / _cubic2nonorth_factor(a)

def _cubic2nonorth_factor(a):
    '''The factors to transform the energy cutoff from cubic lattice to
    non-orthogonal lattice. Energy cutoff is estimated based on cubic lattice.
    It needs to be rescaled for the non-orthogonal lattice to ensure that the
    minimal Gv vector in the reciprocal space is larger than the required
    energy cutoff.
    '''
    # Using ke_cutoff to set up a sphere, the sphere needs to be completely
    # inside the box defined by Gv vectors
    abase = a / np.linalg.norm(a, axis=1)[:,None]
    bbase = np.linalg.inv(abase.T)
    overlap = np.einsum('ix,ix->i', abase, bbase)
    return 1./overlap**2

def cutoff_to_gs(a, cutoff):
    '''Deprecated.  Replaced by function cutoff_to_mesh.'''
    return [n//2 for n in cutoff_to_mesh(a, cutoff)]

def gs_to_cutoff(a, gs):
    '''Deprecated.  Replaced by function mesh_to_cutoff.'''
    return mesh_to_cutoff(a, [2*n+1 for n in gs])


# pyscf/pbc/tools/pyscf_ase.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Garnet Chan <gkc1000@gmail.com>
#

'''
ASE package interface
'''

import numpy as np
from ase.calculators.calculator import Calculator
import ase.dft.kpoints
from ase.lattice import bulk

def ase_atoms_to_pyscf(ase_atoms):
    '''Convert ASE atoms to PySCF atom.

    Note: ASE atoms always use A.
    '''
    return [[atom.symbol, atom.position] for atom in ase_atoms]
atoms_from_ase = ase_atoms_to_pyscf

class PySCF(Calculator):
    implemented_properties = ['energy']

    def __init__(self, restart=None, ignore_bad_restart_file=False,
                 label='PySCF', atoms=None, scratch=None, **kwargs):
        """Construct PySCF-calculator object.

        Parameters
        ==========
        label: str
            Prefix to use for filenames (label.in, label.txt, ...).
            Default is 'PySCF'.

        mfclass: PySCF mean-field class
        molcell: PySCF :Mole: or :Cell:
        """
        Calculator.__init__(self, restart=None, ignore_bad_restart_file=False,
                            label='PySCF', atoms=None, scratch=None, **kwargs)

        # TODO
        # This explicitly refers to "cell". How to refer
        # to both cell and mol together?

        self.mf=None
        self.initialize(**kwargs)

    def initialize(self, molcell, mf_class, mf_dict):
        if not molcell.unit.startswith(('A','a')):
            raise RuntimeError("PySCF unit must be A to work with ASE")

        self.molcell=molcell
        self.mf_class=mf_class
        self.mf_dict=mf_dict

    def set(self, **kwargs):
        changed_parameters = Calculator.set(self, **kwargs)
        if changed_parameters:
            self.reset()

    def calculate(self, atoms=None, properties=['energy'],
                  system_changes=['positions', 'numbers', 'cell',
                                  'pbc', 'charges','magmoms']):

        Calculator.calculate(self, atoms)

        calc_molcell = self.molcell.copy()
        calc_molcell.atom = ase_atoms_to_pyscf(atoms)
        calc_molcell.a = np.asarray(atoms.cell)
        calc_molcell.build(None,None)
        self.mf = self.mf_class(calc_molcell)
        for key in self.mf_dict:
            self.mf.__dict__[key] = self.mf_dict[key]

        self.results['energy']=self.mf.scf()
        self.results['mf']=self.mf

def make_kpts(cell, nks):
    raise DeprecationWarning('Use cell.make_kpts(nks) instead.')



# pyscf/pbc/tools/pywannier90.py

# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
There are two options to call wannier90 in PySCF.  One is the pyWannier90.py
interface as implemented in this file.

(1)
pyWannier90: Wannier90 for PySCF (https://github.com/hungpham2017/pyWannier90)
Hung Q. Pham
email: pqh3.14@gmail.com

(2)
Another wannier90 python interface is available on the repo:
    https://github.com/zhcui/wannier90
Contact its author "Zhihao Cui" <zcui@caltech.edu> for more details of
installation and implementations.
'''

import os, time
import numpy as np
from scipy.io import FortranFile
import pyscf.data.nist as param
from pyscf import lib
from pyscf.pbc import df
from pyscf.pbc.dft import gen_grid, numint

try:
    import libwannier90
except ImportError:
    print('WARNING: Check the installation of libwannier90 and its path in pyscf/pbc/tools/pywannier90.py')
    print('libwannier90 can be found at: https://github.com/hungpham2017/pyWannier90')
    raise


def save_kmf(kmf, chkfile):
    ''' Save a wavefunction'''
    from pyscf.lib.chkfile import save
    kpts = kmf.kpts
    mo_energy_kpts = kmf.mo_energy_kpts
    mo_coeff_kpts = kmf.mo_coeff_kpts

    scf_dic = {'kpts'          : kpts,
               'mo_energy_kpts': mo_energy_kpts,
               'mo_coeff_kpts' : mo_coeff_kpts}
    save(chkfile, 'scf', scf_dic)

def load_kmf(chkfile):
    ''' Load a wavefunction'''
    from pyscf.lib.chkfile import load
    kmf = load(chkfile, 'scf')
    class fake_kmf:
        def __init__(self, kmf):
            self.kpts = kmf['kpts']
            self.mo_energy_kpts = kmf['mo_energy_kpts']
            self.mo_coeff_kpts = kmf['mo_coeff_kpts']
    kmf = fake_kmf(kmf)
    return kmf

def angle(v1, v2):
    '''
    Return the angle (in radiant between v1 and v2)
    '''

    v1 = np.asarray(v1)
    v2 = np.asarray(v2)
    cosa = v1.dot(v2)/ np.linalg.norm(v1) / np.linalg.norm(v2)
    return np.arccos(cosa)

def transform(x_vec, z_vec):
    '''
    Construct a transformation matrix to transform r_vec to the new coordinate system defined by x_vec and z_vec
    '''

    x_vec = x_vec/np.linalg.norm(np.asarray(x_vec))
    z_vec = z_vec/np.linalg.norm(np.asarray(z_vec))
    assert x_vec.dot(z_vec) == 0    # x and z have to be orthogonal to one another
    y_vec = -np.cross(x_vec,z_vec)
    new = np.asarray([x_vec, y_vec, z_vec])
    original = np.asarray([[1,0,0],[0,1,0],[0,0,1]])

    tran_matrix = np.empty([3,3])
    for row in range(3):
        for col in range(3):
            tran_matrix[row,col] = np.cos(angle(original[row],new[col]))

    return tran_matrix.T

def cartesian_prod(arrays, out=None, order='C'):
    '''
    This function is similar to lib.cartesian_prod of PySCF, except the output can be in Fortran or in C order
    '''
    arrays = [np.asarray(x) for x in arrays]
    dtype = np.result_type(*arrays)
    nd = len(arrays)
    dims = [nd] + [len(x) for x in arrays]

    if out is None:
        out = np.empty(dims, dtype)
    else:
        out = np.ndarray(dims, dtype, buffer=out)
    tout = out.reshape(dims)

    shape = [-1] + [1] * nd
    for i, arr in enumerate(arrays):
        tout[i] = arr.reshape(shape[:nd-i])

    return tout.reshape((nd,-1),order=order).T

def periodic_grid(cell, grid=[50,50,50], supercell=[1,1,1], order='C'):
    '''
    Generate a periodic grid for the unit/computational cell in F/C order
    '''
    ngrid = np.asarray(grid)
    qv = cartesian_prod([np.arange(-ngrid[i]*(supercell[i]//2),ngrid[i]*((supercell[i]+1)//2)) for i in range(3)], order=order)
    a_frac = lib.einsum('i,ij->ij', 1./ngrid, cell.lattice_vectors())
    coords = np.dot(qv, a_frac)

    # Compute weight
    ngrids = np.prod(grid)
    ncells = np.prod(supercell)
    weights = np.empty(ngrids*ncells)
    weights[:] = cell.vol / ngrids / ncells
    return coords, weights

def R_r(r_norm, r=1, zona=1):
    r'''
    Radial functions used to compute \Theta_{l,m_r}(\theta,\phi)
    '''

    if r == 1:
        R_r = 2 * zona**(3/2) * np.exp(-zona*r_norm)
    elif r == 2:
        R_r = 1 / 2 / np.sqrt(2) * zona**(3/2) * (2 - zona*r_norm) * np.exp(-zona*r_norm/2)
    else:
        R_r = np.sqrt(4/27) * zona**(3/2) * (1 - 2*zona*r_norm/3 + 2*(zona**2)*(r_norm**2)/27) * np.exp(-zona*r_norm/3)

    return R_r

def theta(func, cost, phi):
    r'''
    Basic angular functions (s,p,d,f) used to compute \Theta_{l,m_r}(\theta,\phi)
    ref: Table 3.1 of the Wannier90 User guide
        Link: https://github.com/wannier-developers/wannier90/raw/v3.1.0/doc/compiled_docs/user_guide.pdf
    '''
    sint = np.sqrt(1 - cost**2)
    if func == 's':
        theta = 1 / np.sqrt(4 * np.pi) * np.ones([cost.shape[0]])
    elif func == 'pz':
        theta = np.sqrt(3 / 4 / np.pi) * cost
    elif func == 'px':
        theta = np.sqrt(3 / 4 / np.pi) * sint * np.cos(phi)
    elif func == 'py':
        theta = np.sqrt(3 / 4 / np.pi) * sint * np.sin(phi)
    elif func == 'dz2':
        theta = np.sqrt(5 / 16 / np.pi) * (3*cost**2 - 1)
    elif func == 'dxz':
        theta = np.sqrt(15 / 4 / np.pi) * sint * cost * np.cos(phi)
    elif func == 'dyz':
        theta = np.sqrt(15 / 4 / np.pi) * sint * cost * np.sin(phi)
    elif func == 'dx2-y2':
        theta = np.sqrt(15 / 16 / np.pi) * (sint**2) * np.cos(2*phi)
    elif func == 'dxy':
        theta = np.sqrt(15 / 16 / np.pi) * (sint**2) * np.sin(2*phi)
    elif func == 'fz3':
        theta = np.sqrt(7) / 4 / np.sqrt(np.pi) * (5*cost**3 - 3*cost)
    elif func == 'fxz2':
        theta = np.sqrt(21) / 4 / np.sqrt(2*np.pi) * (5*cost**2 - 1) * sint * np.cos(phi)
    elif func == 'fyz2':
        theta = np.sqrt(21) / 4 / np.sqrt(2*np.pi) * (5*cost**2 - 1) * sint * np.sin(phi)
    elif func == 'fz(x2-y2)':
        theta = np.sqrt(105) / 4 / np.sqrt(np.pi) * sint**2 * cost * np.cos(2*phi)
    elif func == 'fxyz':
        theta = np.sqrt(105) / 4 / np.sqrt(np.pi) * sint**2 * cost * np.sin(2*phi)
    elif func == 'fx(x2-3y2)':
        theta = np.sqrt(35) / 4 / np.sqrt(2*np.pi) * sint**3 * (np.cos(phi)**2 - 3*np.sin(phi)**2) * np.cos(phi)
    elif func == 'fy(3x2-y2)':
        theta = np.sqrt(35) / 4 / np.sqrt(2*np.pi) * sint**3 * (3*np.cos(phi)**2 - np.sin(phi)**2) * np.sin(phi)

    return theta

def theta_lmr(l, mr, cost, phi):
    r'''
    Compute the value of \Theta_{l,m_r}(\theta,\phi)
    ref: Table 3.1 and 3.2 of the Wannier90 User guide
        Link: https://github.com/wannier-developers/wannier90/raw/v3.1.0/doc/compiled_docs/user_guide.pdf
    '''
    assert l in [0,1,2,3,-1,-2,-3,-4,-5]
    assert mr in [1,2,3,4,5,6,7]

    if l == 0:                           # s
        theta_lmr = theta('s', cost, phi)
    elif (l == 1) and (mr == 1):         # pz
        theta_lmr = theta('pz', cost, phi)
    elif (l == 1) and (mr == 2):         # px
        theta_lmr = theta('px', cost, phi)
    elif (l == 1) and (mr == 3):         # py
        theta_lmr = theta('py', cost, phi)
    elif (l == 2) and (mr == 1):         # dz2
        theta_lmr = theta('dz2', cost, phi)
    elif (l == 2) and (mr == 2):         # dxz
        theta_lmr = theta('dxz', cost, phi)
    elif (l == 2) and (mr == 3):         # dyz
        theta_lmr = theta('dyz', cost, phi)
    elif (l == 2) and (mr == 4):         # dx2-y2
        theta_lmr = theta('dx2-y2', cost, phi)
    elif (l == 2) and (mr == 5):         # dxy
        theta_lmr = theta('dxy', cost, phi)
    elif (l == 3) and (mr == 1):         # fz3
        theta_lmr = theta('fz3', cost, phi)
    elif (l == 3) and (mr == 2):         # fxz2
        theta_lmr = theta('fxz2', cost, phi)
    elif (l == 3) and (mr == 3):         # fyz2
        theta_lmr = theta('fyz2', cost, phi)
    elif (l == 3) and (mr == 4):         # fz(x2-y2)
        theta_lmr = theta('fz(x2-y2)', cost, phi)
    elif (l == 3) and (mr == 5):         # fxyz
        theta_lmr = theta('fxyz', cost, phi)
    elif (l == 3) and (mr == 6):         # fx(x2-3y2)
        theta_lmr = theta('fx(x2-3y2)', cost, phi)
    elif (l == 3) and (mr == 7):         # fy(3x2-y2)
        theta_lmr = theta('fy(3x2-y2)', cost, phi)
    elif (l == -1) and (mr == 1):         # sp-1
        theta_lmr = 1/np.sqrt(2) * (theta('s', cost, phi) + theta('px', cost, phi))
    elif (l == -1) and (mr == 2):         # sp-2
        theta_lmr = 1/np.sqrt(2) * (theta('s', cost, phi) - theta('px', cost, phi))
    elif (l == -2) and (mr == 1):         # sp2-1
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) - 1/np.sqrt(6) *theta('px', cost, phi) + 1/np.sqrt(2) * theta('py', cost, phi)
    elif (l == -2) and (mr == 2):         # sp2-2
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) - 1/np.sqrt(6) *theta('px', cost, phi) - 1/np.sqrt(2) * theta('py', cost, phi)
    elif (l == -2) and (mr == 3):         # sp2-3
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) + 2/np.sqrt(6) *theta('px', cost, phi)
    elif (l == -3) and (mr == 1):         # sp3-1
        theta_lmr = 1/2 * (theta('s', cost, phi) + theta('px', cost, phi) + theta('py', cost, phi) + theta('pz', cost, phi))
    elif (l == -3) and (mr == 2):         # sp3-2
        theta_lmr = 1/2 * (theta('s', cost, phi) + theta('px', cost, phi) - theta('py', cost, phi) - theta('pz', cost, phi))
    elif (l == -3) and (mr == 3):         # sp3-3
        theta_lmr = 1/2 * (theta('s', cost, phi) - theta('px', cost, phi) + theta('py', cost, phi) - theta('pz', cost, phi))
    elif (l == -3) and (mr == 4):         # sp3-4
        theta_lmr = 1/2 * (theta('s', cost, phi) - theta('px', cost, phi) - theta('py', cost, phi) + theta('pz', cost, phi))
    elif (l == -4) and (mr == 1):         # sp3d-1
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) - 1/np.sqrt(6) *theta('px', cost, phi) + 1/np.sqrt(2) * theta('py', cost, phi)
    elif (l == -4) and (mr == 2):         # sp3d-2
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) - 1/np.sqrt(6) *theta('px', cost, phi) - 1/np.sqrt(2) * theta('py', cost, phi)
    elif (l == -4) and (mr == 3):         # sp3d-3
        theta_lmr = 1/np.sqrt(3) * theta('s', cost, phi) + 2/np.sqrt(6) * theta('px', cost, phi)
    elif (l == -4) and (mr == 4):         # sp3d-4
        theta_lmr = 1/np.sqrt(2) * (theta('pz', cost, phi) + theta('dz2', cost, phi))
    elif (l == -4) and (mr == 5):         # sp3d-5
        theta_lmr = 1/np.sqrt(2) * (-theta('pz', cost, phi) + theta('dz2', cost, phi))
    elif (l == -5) and (mr == 1):         # sp3d2-1
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) - 1/np.sqrt(2) *theta('px', cost, phi) - 1/np.sqrt(12) *theta('dz2', cost, phi) \
                    + 1/2 *theta('dx2-y2', cost, phi)
    elif (l == -5) and (mr == 2):         # sp3d2-2
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) + 1/np.sqrt(2) *theta('px', cost, phi) - 1/np.sqrt(12) *theta('dz2', cost, phi) \
                    + 1/2 *theta('dx2-y2', cost, phi)
    elif (l == -5) and (mr == 3):         # sp3d2-3
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) - 1/np.sqrt(2) *theta('py', cost, phi) - 1/np.sqrt(12) *theta('dz2', cost, phi) \
                    - 1/2 *theta('dx2-y2', cost, phi)
    elif (l == -5) and (mr == 4):         # sp3d2-4
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) + 1/np.sqrt(2) *theta('py', cost, phi) - 1/np.sqrt(12) *theta('dz2', cost, phi) \
                    - 1/2 *theta('dx2-y2', cost, phi)
    elif (l == -5) and (mr == 5):         # sp3d2-5
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) - 1/np.sqrt(2) *theta('pz', cost, phi) + 1/np.sqrt(3) *theta('dz2', cost, phi)
    elif (l == -5) and (mr == 6):         # sp3d2-6
        theta_lmr = 1/np.sqrt(6) * theta('s', cost, phi) + 1/np.sqrt(2) *theta('pz', cost, phi) + 1/np.sqrt(3) *theta('dz2', cost, phi)

    return theta_lmr

def g_r(grids_coor, site, l, mr, r, zona, x_axis=[1,0,0], z_axis=[0,0,1], unit='B'):
    r'''
    Evaluate the projection function g(r) or \Theta_{l,m_r}(\theta,\phi) on a grid
    ref: Chapter 3, wannier90 User Guide
    Attributes:
        grids_coor : a grids for the cell of interest
        site       : absolute coordinate (in Borh/Angstrom) of the g(r) in the cell
        l, mr      : l and mr value in the Table 3.1 and 3.2 of the ref
    Return:
        theta_lmr  : an array (ngrid, value) of g(r)

    '''

    unit_conv = 1
    if unit == 'A': unit_conv = param.BOHR

    r_vec = (grids_coor - site)
    r_vec = lib.einsum('iv,uv ->iu', r_vec, transform(x_axis, z_axis))
    r_norm = np.linalg.norm(r_vec,axis=1)
    if (r_norm < 1e-8).any():
        r_vec = (grids_coor - site - 1e-5)
        r_vec = lib.einsum('iv,uv ->iu', r_vec, transform(x_axis, z_axis))
        r_norm = np.linalg.norm(r_vec,axis=1)
    cost = r_vec[:,2]/r_norm

    phi = np.empty_like(r_norm)
    larger_idx = r_vec[:, 0] > 1e-8
    smaller_idx = r_vec[:, 0] < -1e-8
    neither_idx = not larger_idx and not smaller_idx
    phi[larger_idx] = np.arctan(r_vec[larger_idx,1]/r_vec[larger_idx,0])
    phi[smaller_idx] = np.arctan(r_vec[smaller_idx,1]/r_vec[smaller_idx,0])  + np.pi
    phi[neither_idx] = np.sign(r_vec[neither_idx,1]) * 0.5 * np.pi

    return theta_lmr(l, mr, cost, phi) * R_r(r_norm * unit_conv, r = r, zona = zona)

def get_wigner_seitz_supercell(w90, ws_search_size=[2,2,2], ws_distance_tol=1e-6):
    '''
    Return a grid that contains all the lattice within the Wigner-Seitz supercell
    Ref: the hamiltonian_wigner_seitz(count_pts) in wannier90/src/hamittonian.F90
    '''

    real_metric = w90.real_lattice_loc.T.dot(w90.real_lattice_loc)
    dist_dim = np.prod(2 * (np.asarray(ws_search_size) + 1) + 1)
    ndegen = []
    irvec = []
    mp_grid = np.asarray(w90.mp_grid_loc)
    n1_range =  np.arange(-ws_search_size[0] * mp_grid[0], ws_search_size[0]*mp_grid[0] + 1)
    n2_range =  np.arange(-ws_search_size[1] * mp_grid[1], ws_search_size[1]*mp_grid[1] + 1)
    n3_range =  np.arange(-ws_search_size[2] * mp_grid[2], ws_search_size[2]*mp_grid[2] + 1)
    x, y, z = np.meshgrid(n1_range, n2_range, n3_range)
    n_list = np.vstack([z.flatten('F'), x.flatten('F'), y.flatten('F')]).T
    i1 = np.arange(- ws_search_size[0] - 1, ws_search_size[0] + 2)
    i2 = np.arange(- ws_search_size[1] - 1, ws_search_size[1] + 2)
    i3 = np.arange(- ws_search_size[2] - 1, ws_search_size[2] + 2)
    x, y, z = np.meshgrid(i1, i2, i3)
    i_list = np.vstack([z.flatten('F'), x.flatten('F'), y.flatten('F')]).T

    nrpts = 0
    for n in n_list:
        # Calculate |r-R|^2
        ndiff = n - i_list * mp_grid
        dist = (ndiff.dot(real_metric).dot(ndiff.T)).diagonal()

        dist_min = dist.min()
        if abs(dist[(dist_dim + 1)//2 -1] - dist_min) < ws_distance_tol**2:
            temp = 0
            for i in range(0, dist_dim):
                if (abs(dist[i] - dist_min) < ws_distance_tol**2):
                    temp = temp + 1
            ndegen.append(temp)
            irvec.append(n.tolist())
            if (n**2).sum() < 1.e-10: rpt_origin = nrpts
            nrpts = nrpts + 1

    irvec = np.asarray(irvec)
    ndegen = np.asarray(ndegen)

    # Check the "sum rule"
    tot = np.sum(1/np.asarray(ndegen))
    assert tot - np.prod(mp_grid) < 1e-8, "Error in finding Wigner-Seitz points!!!"

    return ndegen, irvec, rpt_origin

def R_wz_sc(w90, R_in, R0, ws_search_size=[2,2,2], ws_distance_tol=1e-6):
    '''
    TODO: document it
    Ref: This is the replication of the R_wz_sc function of ws_distance.F90
    '''
    ndegenx = 8 #max number of unit cells that can touch in a single point (i.e.  vertex of cube)
    R_bz = np.asarray(R_in).reshape(-1, 3)
    nR = R_bz.shape[0]
    R0 = np.asarray(R0)
    ndeg = np.zeros([nR], dtype=np.int32)
    ndeg_ = np.zeros([nR, ndegenx])
    shifts = np.zeros([nR, ndegenx, 3])
    R_out = np.zeros([nR, ndegenx, 3])

    mod2_R_bz = np.sum((R_bz - R0)**2, axis=1)
    R_in_f = R_bz.dot(w90.recip_lattice_loc.T / 2 / np.pi)
    n1_range =  np.arange(-ws_search_size[0] - 1, ws_search_size[0] + 2)
    n2_range =  np.arange(-ws_search_size[1] - 1, ws_search_size[1] + 2)
    n3_range =  np.arange(-ws_search_size[2] - 1, ws_search_size[2] + 2)
    x, y, z = np.meshgrid(n1_range, n2_range, n3_range)
    n_list = np.vstack([z.flatten('F'), x.flatten('F'), y.flatten('F')]).T
    trans_vecs = n_list * w90.mp_grid_loc

    # First loop:
    R_f = np.repeat(R_in_f[:,np.newaxis,:], trans_vecs.shape[0], axis=1) + trans_vecs
    R = R_f.dot(w90.real_lattice_loc)
    mod2_R = np.sum((R - R0)**2, axis=2)
    mod2_R_min = mod2_R.min(axis=1)
    mod2_R_min_idx = np.argmin(mod2_R, axis=1)
    idx = mod2_R_min < mod2_R_bz
    R_bz[idx] = R[idx, mod2_R_min_idx[idx]]
    mod2_R_bz[idx] = mod2_R_min[idx]
    shifts_data = np.repeat(trans_vecs[np.newaxis,:,:], nR, axis=0)[idx, mod2_R_min_idx[idx]]
    shifts[idx] = np.repeat(shifts_data[:,np.newaxis,:], ndegenx, axis=1)

    idx = mod2_R_bz < ws_distance_tol**2
    ndeg[idx] = 1
    R_out[idx, 0] = R0

    # Second loop:
    R_in_f = R_bz.dot(w90.recip_lattice_loc.T / 2 / np.pi)
    R_f = np.repeat(R_in_f[:,np.newaxis,:], trans_vecs.shape[0], axis=1) + trans_vecs
    R = R_f.dot(w90.real_lattice_loc)
    mod2_R = np.sum((R - R0)**2, axis=2)
    mod2_R_bz = np.repeat(mod2_R_bz[:,np.newaxis], trans_vecs.shape[0], axis=1)
    abs_diff = abs(np.sqrt(mod2_R) - np.sqrt(mod2_R_bz))
    idx = abs_diff < ws_distance_tol
    ndeg = idx.sum(axis=1)
    assert (ndeg <= 8).all(), "The degeneracy cannot be larger than 8"
    for i in range(nR):
        R_out[i, :ndeg[i]] = R[i, idx[i]]
        shifts[i, :ndeg[i]] = shifts[i, :ndeg[i]] + trans_vecs[idx[i]]
        ndeg_[i, :ndeg[i]] = 1.0

    return ndeg_, ndeg, R_out, shifts

def ws_translate_dist(w90, irvec, ws_search_size=[2,2,2], ws_distance_tol=1e-6):
    '''
    TODO: document it
    Ref: This is the replication of the ws_translate_dist function of ws_distance.F90
    '''
    nrpts = irvec.shape[0]
    ndegenx = 8 #max number of unit cells that can touch in a single point (i.e.  vertex of cube)
    num_wann = w90.num_wann
    assert ndegenx*num_wann*nrpts > 0, "Unexpected dimensions in ws_translate_dist"

    irvec_ = []
    wann_centres_i = []
    wann_centres_j = []
    for i in range(3):
        x, y, z = np.meshgrid(irvec[:,i], np.zeros(num_wann), np.zeros(num_wann), indexing='ij')
        irvec_.append(x.flatten())
        x, y, z = np.meshgrid(np.zeros(nrpts), np.zeros(num_wann), w90.wann_centres[:,i], indexing='ij')
        wann_centres_i.append(z.flatten())
        x, y, z = np.meshgrid(np.zeros(nrpts), w90.wann_centres[:,i], np.zeros(num_wann), indexing='ij')
        wann_centres_j.append(y.flatten())


    irvec_list = np.vstack(irvec_).T
    irvec_cart_list = irvec_list.dot(w90.real_lattice_loc)
    wann_centres_i_list = np.vstack(wann_centres_i).T
    wann_centres_j_list = np.vstack(wann_centres_j).T
    R_in = irvec_cart_list - wann_centres_i_list + wann_centres_j_list
    wdist_ndeg_, wdist_ndeg, R_out, shifts = w90.R_wz_sc(R_in, [0,0,0], ws_search_size, ws_distance_tol)
    ndegenx = wdist_ndeg_.shape[1]
    irdist_ws = np.repeat(irvec_list[:,np.newaxis,:], ndegenx, axis=1) + shifts
    crdist_ws = irdist_ws.dot(w90.real_lattice_loc)

    # Reformat the matrices for the computational convenience in lib.einsum
    wdist_ndeg = wdist_ndeg.reshape(nrpts, num_wann, num_wann)
    wdist_ndeg_ = wdist_ndeg_.reshape(nrpts, num_wann, num_wann, ndegenx).transpose(3,0,1,2)
    irdist_ws = irdist_ws.reshape(nrpts, num_wann, num_wann, ndegenx, 3).transpose(3,0,1,2,4)
    crdist_ws = crdist_ws.reshape(nrpts, num_wann, num_wann, ndegenx, 3).transpose(3,0,1,2,4)

    return wdist_ndeg, wdist_ndeg_, irdist_ws, crdist_ws


'''Main class of pyWannier90'''
class W90:
    def __init__(self, kmf, cell, mp_grid, num_wann, gamma=False, spinors=False, spin_up=None, other_keywords=None):

        if isinstance(kmf, str):
            self.kmf = load_kmf(kmf)
        else:
            self.kmf = kmf
        self.cell = cell
        self.num_wann = num_wann
        self.keywords = other_keywords

        # Collect the pyscf calculation info
        nao_kpts = []
        for mo_energy in kmf.mo_energy_kpts:
            nao_kpts.append(mo_energy.shape[0])

        self.num_bands_tot = np.min(nao_kpts)
        if self.num_bands_tot < cell.nao_nr():
            print(('The number of bands at different k-point are not the same. '
                   'The first %d bands are used.') % (self.num_bands_tot) )

        self.num_kpts_loc = self.kmf.kpts.shape[0]
        self.mp_grid_loc = mp_grid
        assert self.num_kpts_loc == np.asarray(self.mp_grid_loc).prod()
        self.real_lattice_loc = self.cell.lattice_vectors() * param.BOHR
        self.recip_lattice_loc = self.cell.reciprocal_vectors() / param.BOHR
        self.kpt_latt_loc = self.cell.get_scaled_kpts(self.kmf.kpts)
        self.num_atoms_loc = self.cell.natm
        self.atom_symbols_loc = [atom[0] for atom in self.cell._atom]
        self.atom_atomic_loc = [int(self.cell._atm[atom][0] + self.cell.atom_nelec_core(atom))
                                for atom in range(self.num_atoms_loc)]
        self.atoms_cart_loc = np.asarray([(np.asarray(atom[1])* param.BOHR).tolist()
                                          for atom in self.cell._atom])
        self.gamma_only, self.spinors = (0 , 0)
        if gamma: self.gamma_only = 1
        if spinors: self.spinors = 1

        # Wannier90_setup outputs
        self.num_bands_loc = None
        self.num_wann_loc = None
        self.nntot_loc = None
        self.nn_list = None
        self.proj_site = None
        self.proj_l = None
        self.proj_m = None
        self.proj_radial = None
        self.proj_z = None
        self.proj_x = None
        self.proj_zona = None
        self.exclude_bands = None
        self.proj_s = None
        self.proj_s_qaxis = None

        # Input for Wannier90_run
        self.band_included_list = None
        self.A_matrix_loc = None
        self.M_matrix_loc = None
        self.eigenvalues_loc = None

        # Wannier90_run outputs
        self.U_matrix = None
        self.U_matrix_opt = None
        self.lwindow = None
        self.wann_centres = None
        self.wann_spreads = None
        self.spread = None

        # Others
        self.use_bloch_phases = False
        self.spin_up = spin_up
        self.mo_energy_kpts = []
        self.mo_coeff_kpts = []
        if np.mod(self.cell.nelectron,2) !=0:
            if spin_up:
                for kpt in range(self.num_kpts_loc):
                    self.mo_energy_kpts.append(self.kmf.mo_energy_kpts[0][kpt][:self.num_bands_tot])
                    self.mo_coeff_kpts.append(self.kmf.mo_coeff_kpts[0][kpt][:,:self.num_bands_tot])
            else:
                for kpt in range(self.num_kpts_loc):
                    self.mo_energy_kpts.append(self.kmf.mo_energy_kpts[1][kpt][:self.num_bands_tot])
                    self.mo_coeff_kpts.append(self.kmf.mo_coeff_kpts[1][kpt][:,:self.num_bands_tot])
        else:

            for kpt in range(self.num_kpts_loc):
                self.mo_energy_kpts.append(self.kmf.mo_energy_kpts[kpt][:self.num_bands_tot])
                self.mo_coeff_kpts.append(self.kmf.mo_coeff_kpts[kpt][:,:self.num_bands_tot])

    def kernel(self, external_AME=None):
        '''
        Main kernel for pyWannier90
        '''
        self.make_win()
        self.setup()
        if external_AME is not None:
            self.M_matrix_loc = self.read_M_mat(external_AME + '.mmn')
            self.A_matrix_loc = self.read_A_mat(external_AME + '.amn')
            self.eigenvalues_loc = self.read_epsilon_mat(external_AME + '.eig')
        else:
            self.M_matrix_loc = self.get_M_mat()
            self.A_matrix_loc = self.get_A_mat()
            self.eigenvalues_loc = self.get_epsilon_mat()
        self.run()

    def make_win(self):
        '''
        Make a basic *.win file for wannier90
        '''

        win_file = open('wannier90.win', "w")
        win_file.write('! Basic input generated by the pyWannier90. Date: %s\n' % (time.ctime()))
        win_file.write('\n')
        win_file.write('num_bands       = %d\n' % (self.num_bands_tot))
        win_file.write('num_wann       = %d\n' % (self.num_wann))
        win_file.write('\n')
        win_file.write('Begin Unit_Cell_Cart\n')
        for row in range(3):
            win_file.write('%10.7f  %10.7f  %10.7f\n' %
                           (self.real_lattice_loc[0, row], self.real_lattice_loc[1, row],
                            self.real_lattice_loc[2, row]))
        win_file.write('End Unit_Cell_Cart\n')
        win_file.write('\n')
        win_file.write('Begin atoms_cart\n')
        for atom in range(len(self.atom_symbols_loc)):
            win_file.write('%s  %7.7f  %7.7f  %7.7f\n' %
                           (self.atom_symbols_loc[atom], self.atoms_cart_loc[atom,0],
                            self.atoms_cart_loc[atom,1], self.atoms_cart_loc[atom,2]))
        win_file.write('End atoms_cart\n')
        win_file.write('\n')
        if self.use_bloch_phases: win_file.write('use_bloch_phases = T\n\n')
        if self.keywords is not None:
            win_file.write('!Additional keywords\n')
            win_file.write(self.keywords)
        win_file.write('\n\n\n')
        win_file.write('mp_grid        = %d %d %d\n' %
                       (self.mp_grid_loc[0], self.mp_grid_loc[1], self.mp_grid_loc[2]))
        if self.gamma_only == 1: win_file.write('gamma_only : true\n')
        win_file.write('begin kpoints\n')
        for kpt in range(self.num_kpts_loc):
            win_file.write('%7.7f  %7.7f  %7.7f\n' %
                           (self.kpt_latt_loc[kpt][0], self.kpt_latt_loc[kpt][1], self.kpt_latt_loc[kpt][2]))
        win_file.write('End Kpoints\n')
        win_file.close()

    def get_M_mat(self):
        r'''
        Construct the ovelap matrix: M_{m,n}^{(\mathbf{k,b})}
        Equation (25) in MV, Phys. Rev. B 56, 12847
        '''

        M_matrix_loc = np.empty([self.num_kpts_loc, self.nntot_loc,
                                 self.num_bands_loc, self.num_bands_loc],
                                dtype = np.complex128)

        for k_id in range(self.num_kpts_loc):
            for nn in range(self.nntot_loc):
                k1 = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
                k_id2 = self.nn_list[nn, k_id, 0] - 1
                k2_ = self.kpt_latt_loc[k_id2]
                k2_scaled = k2_ + self.nn_list[nn, k_id, 1:4]
                k2 = self.cell.get_abs_kpts(k2_scaled)
                s_AO = df.ft_ao.ft_aopair(self.cell, -k2+k1, kpti_kptj=[k2,k1], q = np.zeros(3))[0]
                Cm = self.mo_coeff_kpts[k_id][:,self.band_included_list]
                Cn = self.mo_coeff_kpts[k_id2][:,self.band_included_list]
                M_matrix_loc[k_id, nn,:,:] = lib.einsum('nu,vm,uv->nm', Cn.T.conj(), Cm, s_AO).conj()

        return M_matrix_loc

    def read_M_mat(self, filename=None):
        r'''
        Read the ovelap matrix: M_{m,n}^{(\mathbf{k,b})} from seedname.mnn
        '''
        if filename is None: filename = 'wannier90.mmn'
        assert os.path.exists(filename), "Cannot find " + filename

        with open(filename, 'r') as f:
            data = f.readlines()
            num_bands_loc, num_kpts_loc, nntot_loc = np.int64(data[1].split())
            data = data[2:]
            nn_list = []
            nline = num_bands_loc**2 + 1
            M_matrix_loc = np.empty([num_kpts_loc, nntot_loc, num_bands_loc, num_bands_loc], dtype = np.complex128)
            jump = 0
            for kpt in range(num_kpts_loc):
                for nn in range(nntot_loc):
                    temp = data[jump : jump + nline]
                    nn_list.append(np.int64(temp[0].split()))
                    val_in_float = np.float64(" ".join(temp[1:]).split()).reshape(-1,2)
                    val_in_complex = val_in_float[:,0] + 1j * val_in_float[:,1]
                    M_matrix_loc[kpt, nn] = val_in_complex.reshape(num_bands_loc, num_bands_loc)
                    jump += nline

        return M_matrix_loc

    def get_A_mat(self):
        r'''
        Construct the projection matrix: A_{m,n}^{\mathbf{k}}
        Equation (62) in MV, Phys. Rev. B 56, 12847 or equation (22) in SMV, Phys. Rev. B 65, 035109
        '''

        A_matrix_loc = np.empty([self.num_kpts_loc, self.num_wann_loc, self.num_bands_loc], dtype = np.complex128)

        if self.use_bloch_phases:
            Amn = np.zeros([self.num_wann_loc, self.num_bands_loc])
            np.fill_diagonal(Amn, 1)
            A_matrix_loc[:,:,:] = Amn
        else:
            from pyscf.dft import numint as mol_numint
            from pyscf.dft import gen_grid as mol_gen_grid
            grids = mol_gen_grid.Grids(self.cell).build()
            coords = grids.coords
            weights = grids.weights
            for ith_wann in range(self.num_wann_loc):
                frac_site = self.proj_site[ith_wann]
                abs_site = frac_site.dot(self.real_lattice_loc) / param.BOHR
                l = self.proj_l[ith_wann]
                mr = self.proj_m[ith_wann]
                r = self.proj_radial[ith_wann]
                zona = self.proj_zona[ith_wann]
                x_axis = self.proj_x[ith_wann]
                z_axis = self.proj_z[ith_wann]
                gr = g_r(coords, abs_site, l, mr, r, zona, x_axis, z_axis, unit = 'B')
                ao_L0 = mol_numint.eval_ao(self.cell, coords)
                s_aoL0_g = lib.einsum('i,i,iv->v', weights, gr, ao_L0)
                for k_id in range(self.num_kpts_loc):
                    kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
                    mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
                    s_kpt = self.cell.pbc_intor('int1e_ovlp', hermi=1, kpts=kpt, pbcopt=lib.c_null_ptr())
                    A_matrix_loc[k_id,ith_wann,:] = lib.einsum('v,vu,um->m', s_aoL0_g, s_kpt, mo_included,
                                                               optimize=True).conj()

        return A_matrix_loc

    def read_A_mat(self, filename=None):
        r'''
        Read the ovelap matrix: M_{m,n}^{(\mathbf{k,b})} from seedname.mnn
        '''
        if filename is None: filename = 'wannier90.amn'
        assert os.path.exists(filename), "Cannot find " + filename

        with open(filename, 'r') as f:
            data = f.readlines()
            num_bands_loc, num_kpts_loc, num_wann_loc = np.int64(data[1].split())
            data = data[2:]
            A_matrix_loc = np.empty([num_kpts_loc, num_wann_loc, num_bands_loc], dtype = np.complex128)
            val_in_float = np.float64(" ".join(data).split()).reshape(-1,5)
            A_matrix_loc = (val_in_float[:,3] + 1j * val_in_float[:,4]).reshape(num_kpts_loc, num_wann_loc, num_bands_loc)

        return A_matrix_loc

    def get_epsilon_mat(self):
        r'''
        Construct the eigenvalues matrix: \epsilon_{n}^(\mathbf{k})
        '''

        return np.asarray(self.mo_energy_kpts, dtype=np.float64)[:,self.band_included_list] * param.HARTREE2EV

    def read_epsilon_mat(self, filename=None):
        r'''
        Read the eigenvalues matrix: \epsilon_{n}^(\mathbf{k})
        '''
        if filename is None: filename = 'wannier90.eig'
        assert os.path.exists(filename), "Cannot find " + filename
        with open(filename, 'r') as f:
            data = f.read()
            temp = np.float64(data.split()).reshape(-1, 3)
            nbands = int(temp[:,0].max())
            nkpts = int(temp[:,1].max())
            eigenvals = temp[:,2].reshape(nkpts, nbands)

        return eigenvals

    def setup(self):
        '''
        Execute the Wannier90_setup
        '''

        real_lattice_loc = self.real_lattice_loc.T.flatten()
        recip_lattice_loc = self.recip_lattice_loc.T.flatten()
        kpt_latt_loc = self.kpt_latt_loc.flatten()
        atoms_cart_loc = self.atoms_cart_loc.flatten()

        (bands_wann_nntot, nn_list, proj_site, proj_l, proj_m, proj_radial,
         proj_z, proj_x, proj_zona, exclude_bands, proj_s, proj_s_qaxis) = \
                libwannier90.setup(self.mp_grid_loc, self.num_kpts_loc, real_lattice_loc,
                                   recip_lattice_loc, kpt_latt_loc,
                                   self.num_bands_tot, self.num_atoms_loc,
                                   self.atom_atomic_loc, atoms_cart_loc, self.gamma_only, self.spinors)

        # Convert outputs to the correct data type
        self.num_bands_loc, self.num_wann_loc, self.nntot_loc = np.int32(bands_wann_nntot)
        self.nn_list = np.int32(nn_list)
        self.proj_site = proj_site
        self.proj_l = np.int32(proj_l)
        self.proj_m = np.int32(proj_m)
        self.proj_radial = np.int32(proj_radial)
        self.proj_z = proj_z
        self.proj_x = proj_x
        self.proj_zona = proj_zona
        self.exclude_bands = np.int32(exclude_bands)
        self.band_included_list = [i for i in range(self.num_bands_tot) if (i + 1) not in self.exclude_bands]
        self.proj_s = np.int32(proj_s)
        self.proj_s_qaxis = proj_s_qaxis

    def run(self):
        '''
        Execute the Wannier90_run
        '''

        assert self.num_wann_loc is not None
        assert isinstance(self.M_matrix_loc, np.ndarray)
        assert isinstance(self.A_matrix_loc, np.ndarray)
        assert isinstance(self.eigenvalues_loc, np.ndarray)

        real_lattice_loc = self.real_lattice_loc.T.flatten()
        recip_lattice_loc = self.recip_lattice_loc.T.flatten()
        kpt_latt_loc = self.kpt_latt_loc.flatten()
        atoms_cart_loc = self.atoms_cart_loc.flatten()
        M_matrix_loc = self.M_matrix_loc.flatten()
        A_matrix_loc = self.A_matrix_loc.flatten()
        eigenvalues_loc = self.eigenvalues_loc.flatten()

        U_matrix, U_matrix_opt, lwindow, wann_centres, wann_spreads, spread = \
                libwannier90.run(self.mp_grid_loc, self.num_kpts_loc, real_lattice_loc,
                                 recip_lattice_loc, kpt_latt_loc, self.num_bands_loc,
                                 self.num_wann_loc, self.nntot_loc, self.num_atoms_loc,
                                 self.atom_atomic_loc, atoms_cart_loc, self.gamma_only,
                                 M_matrix_loc, A_matrix_loc, eigenvalues_loc)

        # Convert outputs to the correct data typ
        self.U_matrix = U_matrix
        self.U_matrix_opt = U_matrix_opt
        lwindow = np.int32(np.abs(lwindow.real))
        self.lwindow = (lwindow == 1)
        self.wann_centres = wann_centres.real
        self.wann_spreads = wann_spreads.real
        self.spread = spread.real

    get_wigner_seitz_supercell = get_wigner_seitz_supercell
    R_wz_sc = R_wz_sc
    ws_translate_dist = ws_translate_dist

    def get_hamiltonian_kpts(self):
        '''Get the Hamiltonian in k-space, this should be identical to Fock matrix from PySCF'''

        assert self.U_matrix is not None, "You must wannierize first, then you can run this function"
        eigenvals_in_window = []
        for k_id in range(self.num_kpts_loc):
            mo_included = self.mo_energy_kpts[k_id][self.band_included_list]
            orbs_in_win = self.lwindow[k_id]
            mo_in_window = mo_included[orbs_in_win]
            U_matrix_opt = self.U_matrix_opt[k_id][ :, orbs_in_win].T
            eigenvals = lib.einsum('m,mo,mo->o', mo_in_window, U_matrix_opt.conj(), U_matrix_opt)
            eigenvals_in_window.append(eigenvals)

        hamiltonian_kpts = lib.einsum('kso,ko,kto->kst', self.U_matrix.conj(), eigenvals_in_window, self.U_matrix)
        return hamiltonian_kpts

    def get_hamiltonian_Rs(self, Rs, ham_kpts=None):
        '''Get the R-space Hamiltonian H(R0, R) centered at R0 or the first R in Rs list
        '''

        assert self.U_matrix is not None, "You must wannierize first, then you can run this function"
        nkpts = self.kpt_latt_loc.shape[0]
        if ham_kpts is not None:
            hamiltonian_kpts = ham_kpts
        else:
            hamiltonian_kpts = self.get_hamiltonian_kpts()

        # Find the center either R(0,0,0) or the first R in the Rs list
        ngrid = len(Rs)
        center = np.arange(ngrid)[(np.asarray(Rs)**2).sum(axis=1) < 1e-10]
        if center.shape[0] == 1:
            center = center[0]
        else:
            center = 0

        # The phase factor is computed using the exp(1j*R.dot(k)) rather than exp(-1j*R.dot(k)) in wannier90
        phase = 1/np.sqrt(nkpts) * np.exp(1j* 2*np.pi * np.dot(Rs, self.kpt_latt_loc.T))
        hamiltonian_R0 = lib.einsum('k,kst,Rk->Rst', phase[center], hamiltonian_kpts, phase.conj())

        return hamiltonian_R0

    def interpolate_ham_kpts(self, frac_kpts, ham_kpts=None,
                             use_ws_distance=True, ws_search_size=[2,2,2],
                             ws_distance_tol=1e-6):
        ''' Interpolate the band structure using the Slater-Koster scheme
            Return:
                eigenvalues and eigenvectors at the desired kpts
        '''

        assert self.U_matrix is not None, "You must wannierize first, then you can run this function"
        ndegen, Rs, center = self.get_wigner_seitz_supercell(ws_search_size, ws_distance_tol)
        hamiltonian_R0 = self.get_hamiltonian_Rs(Rs, ham_kpts)

        # Interpolate H(kpts) at the desired k-pts
        if use_ws_distance:
            wdist_ndeg, wdist_ndeg_, irdist_ws, crdist_ws = self.ws_translate_dist(Rs)
            temp = lib.einsum('iRstx,kx->iRstk', irdist_ws, frac_kpts)
            phase = lib.einsum('iRstk,iRst->Rstk', np.exp(1j* 2*np.pi * temp), wdist_ndeg_)
            inter_hamiltonian_kpts = \
                    lib.einsum('R,Rst,Rts,Rstk->kst', 1/ndegen, 1/wdist_ndeg, hamiltonian_R0, phase)
        else:
            phase = np.exp(1j* 2*np.pi * np.dot(Rs, frac_kpts.T))
            inter_hamiltonian_kpts = \
                    lib.einsum('R,Rst,Rk->kst', 1/ndegen, hamiltonian_R0, phase)

        return inter_hamiltonian_kpts

    def interpolate_band(self, frac_kpts, ham_kpts=None, use_ws_distance=True,
                         ws_search_size=[2,2,2], ws_distance_tol=1e-6):
        ''' Interpolate the band structure using the Slater-Koster scheme
            Return:
                eigenvalues and eigenvectors at the desired kpts
        '''

        assert self.U_matrix is not None, (
            "You must wannierize first, then you can run this function")
        inter_hamiltonian_kpts = self.interpolate_ham_kpts(
            frac_kpts, ham_kpts, use_ws_distance, ws_search_size, ws_distance_tol)
        # Diagonalize H(kpts) to get eigenvalues and eigenvector
        nkpts = frac_kpts.shape[0]
        eigvals, eigvecs = np.linalg.eigh(inter_hamiltonian_kpts)
        idx_kpts = eigvals.argsort()
        eigvals = np.asarray([eigvals[kpt][idx_kpts[kpt]] for kpt in range(nkpts)])
        eigvecs = np.asarray([eigvecs[kpt][:,idx_kpts[kpt]] for kpt in range(nkpts)])

        return eigvals, eigvecs

    def is_real(self, threshold=1.e-6):
        '''
        Fourier transform the mo coefficients to real space and check if it is real
        '''

        assert self.U_matrix is not None, "You must wannierize first, then you can run this function"
        eigenvecs_in_window = []
        for k_id in range(self.num_kpts_loc):
            mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
            orbs_in_win = self.lwindow[k_id]
            mo_in_window = mo_included[:, orbs_in_win].dot(self.U_matrix_opt[k_id][ :, orbs_in_win].T)
            eigenvecs_in_window.append(mo_in_window)

        # Rotate the mo(kpts) into localized basis
        rotated_mo_coeff_kpts = lib.einsum('kum,ksm->kus', eigenvecs_in_window, self.U_matrix)

        # Fourier transform the localized mo
        nkx, nky, nkz = self.mp_grid_loc
        Ts = lib.cartesian_prod((np.arange(nkx), np.arange(nky), np.arange(nkz)))
        nkpts = self.kpt_latt_loc.shape[0]
        phase = 1/np.sqrt(nkpts) * np.exp(1j* 2*np.pi * np.dot(Ts, self.kpt_latt_loc.T))
        mo_coeff_Rs = lib.einsum('k,kus,Rk->Rus', phase[0], rotated_mo_coeff_kpts, phase.conj())

        return mo_coeff_Rs.imag.max() < threshold

    def export_unk(self, grid=[50,50,50]):
        '''
        Export the periodic part of BF in a real space grid for plotting with wannier90
        '''

        grids_coor, weights = periodic_grid(self.cell, grid, order = 'F')

        for k_id in range(self.num_kpts_loc):
            if self.spin_up:
                spin = '.1'
            else:
                spin = '.2'
            kpt = self.cell.get_abs_kpts(self.kpt_latt_loc[k_id])
            ao = numint.eval_ao(self.cell, grids_coor, kpt = kpt)
            u_ao = lib.einsum('x,xi->xi', np.exp(-1j*np.dot(grids_coor, kpt)), ao)
            unk_file = FortranFile('UNK' + "%05d" % (k_id + 1) + spin, 'w')
            unk_file.write_record(np.asarray([grid[0], grid[1], grid[2], k_id + 1, self.num_bands_loc], dtype = np.int32))
            mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
            u_mo = lib.einsum('xi,in->xn', u_ao, mo_included)
            for band in range(len(self.band_included_list)):
                unk_file.write_record(np.asarray(u_mo[:,band], dtype = np.complex128))
            unk_file.close()

    def export_AME(self, grid=[50,50,50]):
        r'''
        Export A_{m,n}^{\mathbf{k}} and M_{m,n}^{(\mathbf{k,b})} and \epsilon_{n}^(\mathbf{k})
        '''

        if self.A_matrix_loc is None:
            self.make_win()
            self.setup()
            self.M_matrix_loc = self.get_M_mat()
            self.A_matrix_loc = self.get_A_mat()
            self.eigenvalues_loc = self.get_epsilon_mat()
            self.export_unk(self, grid = grid)

        with open('wannier90.mmn', 'w') as f:
            f.write('Generated by the pyWannier90. Date: %s\n' % (time.ctime()))
            f.write('    %d    %d    %d\n' % (self.num_bands_loc, self.num_kpts_loc, self.nntot_loc))

            for k_id in range(self.num_kpts_loc):
                for nn in range(self.nntot_loc):
                    k_id1 = k_id + 1
                    k_id2 = self.nn_list[nn, k_id, 0]
                    nnn, nnm, nnl = self.nn_list[nn, k_id, 1:4]
                    f.write('    %d  %d    %d  %d  %d\n' % (k_id1, k_id2, nnn, nnm, nnl))
                    for m in range(self.num_bands_loc):
                        for n in range(self.num_bands_loc):
                            f.write('    %22.18e  %22.18e\n' % (self.M_matrix_loc[k_id, nn,m,n].real,
                                                                self.M_matrix_loc[k_id, nn,m,n].imag))

        with open('wannier90.amn', 'w') as f:
            f.write('Generated by the pyWannier90. Date: %s\n' % (time.ctime()))
            f.write('    %d    %d    %d\n' % (self.num_bands_loc, self.num_kpts_loc, self.num_wann_loc))

            for k_id in range(self.num_kpts_loc):
                for ith_wann in range(self.num_wann_loc):
                    for band in range(self.num_bands_loc):
                        f.write('    %d    %d    %d    %22.18e    %22.18e\n' %
                                (band+1, ith_wann+1, k_id+1,
                                 self.A_matrix_loc[k_id,ith_wann,band].real,
                                 self.A_matrix_loc[k_id,ith_wann,band].imag))

        with open('wannier90.eig', 'w') as f:
            for k_id in range(self.num_kpts_loc):
                for band in range(self.num_bands_loc):
                    f.write('    %d    %d    %22.18e\n' % (band+1, k_id+1, self.eigenvalues_loc[k_id,band]))

    def get_wannier(self, supercell=[1,1,1], grid=[50,50,50]):
        '''
        Evaluate the MLWF using a periodic grid
        '''

        grids_coor, weights = periodic_grid(self.cell, grid, supercell = [1,1,1], order = 'C')
        kpts = self.cell.get_abs_kpts(self.kpt_latt_loc)

        u_mo  = []
        for k_id in range(self.num_kpts_loc):
            mo_included = self.mo_coeff_kpts[k_id][:,self.band_included_list]
            mo_in_window = self.lwindow[k_id]
            C_opt = mo_included[:,mo_in_window].dot(self.U_matrix_opt[k_id][ :, mo_in_window].T)
            C_tildle = C_opt.dot(self.U_matrix[k_id].T)
            kpt = kpts[k_id]
            ao = numint.eval_ao(self.cell, grids_coor, kpt = kpt)
            u_ao = lib.einsum('x,xi->xi', np.exp(-1j*np.dot(grids_coor, kpt)), ao)
            u_mo.append(lib.einsum('xi,in->xn', u_ao, C_tildle))

        u_mo = np.asarray(u_mo)
        WF0 = libwannier90.get_WF0s(self.kpt_latt_loc.shape[0],self.kpt_latt_loc, supercell, grid, u_mo)

        # Fix the global phase following the pw2wannier90 procedure
        max_index = (WF0*WF0.conj()).real.argmax(axis=0)
        norm_wfs = np.diag(WF0[max_index,:])
        norm_wfs = norm_wfs/np.absolute(norm_wfs)
        WF0 = WF0/norm_wfs/self.num_kpts_loc

        # Check the 'reality' following the pw2wannier90 procedure
        for WF_id in range(self.num_wann_loc):
            ratio_max = np.abs(WF0[np.abs(WF0[:,WF_id].real) >= 0.01,WF_id].imag /
                               WF0[np.abs(WF0[:,WF_id].real) >= 0.01,WF_id].real).max(axis=0)
            print('The maximum imag/real for wannier function ', WF_id,' : ', ratio_max)
        return WF0

    def get_guess_orb(self, frac_site=[0,0,0], l=0, mr=1, r=1,
                      zona=1.0, x_axis=[1,0,0], z_axis=[0,0,1],
                      supercell=[1,1,1], grid=[50,50,50]):
        '''
        Evaluate a guess orbital using a periodic uniform grid
        '''
        grids_coor, weights = periodic_grid(self.cell, grid, supercell = supercell, order = 'C')
        frac_site = np.asarray(frac_site)
        abs_site = frac_site.dot(self.real_lattice_loc) / param.BOHR
        gr = g_r(grids_coor, abs_site, l, mr, r, zona, x_axis, z_axis, unit = 'B')
        return gr

    def plot_wf(self, outfile='MLWF', wf_list=None, supercell=[1,1,1], grid=[50,50,50]):
        '''
        Export Wannier function at cell R
        xsf format: http://web.mit.edu/xcrysden_v1.5.60/www/XCRYSDEN/doc/XSF.html
        Attributes:
            wf_list        : a list of MLWFs to plot
            supercell    : a supercell used for plotting
        '''

        if wf_list is None:
            wf_list = list(range(self.num_wann_loc))

        grid = np.asarray(grid)
        origin = np.asarray([-(grid[i]*(supercell[i]//2) + 1)/grid[i] for i in range(3)]).dot(
            self.cell.lattice_vectors().T)* param.BOHR
        real_lattice_loc = (grid*supercell-1)/grid * self.cell.lattice_vectors() * param.BOHR
        nx, ny, nz = grid*supercell
        WF0 = self.get_wannier(supercell = supercell, grid = grid)


        for wf_id in wf_list:
            assert wf_id in list(range(self.num_wann_loc))
            WF = WF0[:,wf_id].reshape(nx,ny,nz).real

            with open(outfile + '-' + str(wf_id) + '.xsf', 'w') as f:
                f.write('Generated by the pyWannier90. Date: %s\n\n' % (time.ctime()))
                f.write('CRYSTAL\n')
                f.write('PRIMVEC\n')
                for row in range(3):
                    f.write('%10.7f  %10.7f  %10.7f\n' %
                            (self.real_lattice_loc[row,0], self.real_lattice_loc[row,1],
                             self.real_lattice_loc[row,2]))
                f.write('CONVVEC\n')
                for row in range(3):
                    f.write('%10.7f  %10.7f  %10.7f\n' %
                            (self.real_lattice_loc[row,0], self.real_lattice_loc[row,1],
                             self.real_lattice_loc[row,2]))
                f.write('PRIMCOORD\n')
                f.write('%3d %3d\n' % (self.num_atoms_loc, 1))
                for atom in range(len(self.atom_symbols_loc)):
                    f.write('%s  %7.7f  %7.7f  %7.7f\n' %
                            (self.atom_symbols_loc[atom], self.atoms_cart_loc[atom][0],
                             self.atoms_cart_loc[atom][1], self.atoms_cart_loc[atom][2]))
                f.write('\n\n')
                f.write('BEGIN_BLOCK_DATAGRID_3D\n3D_field\nBEGIN_DATAGRID_3D_UNKNOWN\n')
                f.write('   %5d     %5d  %5d\n' % (nx, ny, nz))
                f.write('   %10.7f  %10.7f  %10.7f\n' % (origin[0],origin[1],origin[2]))
                for row in range(3):
                    f.write('   %10.7f  %10.7f  %10.7f\n' %
                            (real_lattice_loc[row,0], real_lattice_loc[row,1], real_lattice_loc[row,2]))

                fmt = ' %13.5e' * nx + '\n'
                for iz in range(nz):
                    for iy in range(ny):
                        f.write(fmt % tuple(WF[:,iy,iz].tolist()))
                f.write('END_DATAGRID_3D\nEND_BLOCK_DATAGRID_3D')

    def plot_guess_orbs(self, outfile='guess_orb', frac_site=[0,0,0], l=0, mr=1, r=1,
                        zona=1.0, x_axis=[1,0,0], z_axis=[0,0,1],
                        supercell=[1,1,1], grid=[50,50,50]):
        '''
        Export Wannier function at cell R
        xsf format: http://web.mit.edu/xcrysden_v1.5.60/www/XCRYSDEN/doc/XSF.html
        Attributes:
            wf_list        : a list of MLWFs to plot
            supercell    : a supercell used for plotting
        '''

        grid = np.asarray(grid)
        origin = np.asarray([-grid[i]*(supercell[i]//2)/grid[i] for i in range(3)]).dot(
            self.cell.lattice_vectors().T) * param.BOHR
        real_lattice_loc = (grid*supercell-1)/grid * self.cell.lattice_vectors() * param.BOHR
        nx, ny, nz = grid*supercell
        guess_orb = self.get_guess_orb(frac_site=frac_site, l=l, mr=mr, r=r,
                                       zona=zona, x_axis=x_axis, z_axis=z_axis,
                                       supercell=supercell, grid=grid)
        guess_orb = guess_orb.reshape(nx,ny,nz).real

        with open(outfile + '.xsf', 'w') as f:
            f.write('Generated by the pyWannier90\n\n')
            f.write('CRYSTAL\n')
            f.write('PRIMVEC\n')
            for row in range(3):
                f.write('%10.7f  %10.7f  %10.7f\n' % (self.real_lattice_loc[row,0], self.real_lattice_loc[row,1],
                                                      self.real_lattice_loc[row,2]))
            f.write('CONVVEC\n')
            for row in range(3):
                f.write('%10.7f  %10.7f  %10.7f\n' % (self.real_lattice_loc[row,0], self.real_lattice_loc[row,1],
                                                      self.real_lattice_loc[row,2]))
            f.write('PRIMCOORD\n')
            f.write('%3d %3d\n' % (self.num_atoms_loc, 1))
            for atom in range(len(self.atom_symbols_loc)):
                f.write('%s  %7.7f  %7.7f  %7.7f\n' % (self.atom_symbols_loc[atom], self.atoms_cart_loc[atom][0],
                                                       self.atoms_cart_loc[atom][1], self.atoms_cart_loc[atom][2]))
            f.write('\n\n')
            f.write('BEGIN_BLOCK_DATAGRID_3D\n3D_field\nBEGIN_DATAGRID_3D_UNKNOWN\n')
            f.write('   %5d     %5d  %5d\n' % (nx, ny, nz))
            f.write('   %10.7f  %10.7f  %10.7f\n' % (origin[0],origin[1],origin[2]))
            for row in range(3):
                f.write('   %10.7f  %10.7f  %10.7f\n' % (real_lattice_loc[row,0], real_lattice_loc[row,1],
                                                         real_lattice_loc[row,2]))

            fmt = ' %13.5e' * nx + '\n'
            for iz in range(nz):
                for iy in range(ny):
                    f.write(fmt % tuple(guess_orb[:,iy,iz].tolist()))
            f.write('END_DATAGRID_3D\nEND_BLOCK_DATAGRID_3D')

if __name__ == '__main__':
    from pyscf.pbc import gto as pgto
    from pyscf.pbc import scf as pscf
    import pywannier90

    # build cell object
    cell = pgto.Cell()
    cell.a = [[0.0, 2.7155, 2.7155], [2.7155, 0.0, 2.7155], [2.7155, 2.7155, 0.0]]
    cell.atom = [['Si',[0.0,0.0,0.0]], ['Si',[1.35775, 1.35775, 1.35775]]]
    cell.basis = 'gth-dzv'
    cell.pseudo = 'gth-pade'
    cell.exp_to_discard = 0.1
    cell.build()

    # build and run scf object
    kmesh = [3, 1, 1]
    kpts = cell.make_kpts(kmesh)
    kmf = pscf.KKS(cell, kpts)
    kmf.xc = 'pbe'
    kmf.run()

    # build and run w90 object
    num_wann = 8
    keywords = '''
    begin projections
    Si:sp3
    end projections
    '''
    w90 = pywannier90.W90(kmf, cell, kmesh, num_wann, other_keywords=keywords)
    w90.kernel()


# pyscf/pbc/tools/tril.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import numpy
import pyscf.lib

def tril_index(ki,kj):
    assert (numpy.array([ki<=kj])).all()
    return (kj*(kj+1))//2 + ki

# TODO: fairly messy and slow
def unpack_tril(in_array,nkpts,kp,kq,kr,ks):
    # We are only dealing with the case that one of kp,kq,kr is a list
    #
    if in_array.shape[0] == nkpts:
        return in_array[kp,kq,kr].copy()
    nints = sum([isinstance(x, (int, numpy.integer)) for x in (kp,kq,kr)])
    assert(nints>=2)

    kp,kq,kr,ks = [[x] if isinstance(x, (int, numpy.integer)) else x for x in (kp,kq,kr,ks)]
    kp,kq,kr,ks = [numpy.array(x) for x in (kp,kq,kr,ks)]
    indices = numpy.array(pyscf.lib.cartesian_prod((kp,kq,kr)))

    not_tril = numpy.array([x[0]>x[1] for x in indices])
    if sum(not_tril) == 0: # check if lower triangular
        indices = [tril_index(indices[:,0],indices[:,1]),indices[:,2]]
        tmp = in_array[indices].copy()
        if nints == 3:
            tmp = tmp.reshape(in_array.shape[2:7])
        return tmp
    # If not lower triangular, permute the k-point indices
    indices[not_tril] = indices[not_tril][:,[1,0,2]]
    indices[:,2][not_tril] = ks[not_tril]

    indices = [tril_index(indices[:,0],indices[:,1]),indices[:,2]]

    tmp = in_array[indices].copy()
    tmp[not_tril] = tmp[not_tril].transpose(0,2,1,4,3)
    if nints == 3:
        tmp = tmp.reshape(in_array.shape[2:7])
    return tmp


# pyscf/pbc/x2c/__init__.py

# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


# pyscf/pbc/x2c/sfx2c1e.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
spin-free X2C correction for extended systems (experimental feature)
'''


from functools import reduce
import copy
import numpy
import scipy.linalg
from pyscf import lib
from pyscf.gto import mole
from pyscf.lib import logger
from pyscf.x2c import x2c
from pyscf.pbc import gto as pbcgto
from pyscf.pbc import tools
from pyscf.pbc.df import aft
from pyscf.pbc.df import aft_jk
from pyscf.pbc.df import ft_ao
from pyscf.pbc.scf import ghf
from pyscf import __config__


def sfx2c1e(mf):
    '''Spin-free X2C.
    For the given SCF object, update the hcore constructor.

    Args:
        mf : an SCF object

    Returns:
        An SCF object

    Examples:

    >>> mol = gto.M(atom='H 0 0 0; F 0 0 1', basis='ccpvdz', verbose=0)
    >>> mf = scf.RHF(mol).sfx2c1e()
    >>> mf.scf()

    >>> mol.symmetry = 1
    >>> mol.build(0, 0)
    >>> mf = scf.UHF(mol).sfx2c1e()
    >>> mf.scf()
    '''
    if isinstance(mf, x2c._X2C_SCF):
        if mf.with_x2c is None:
            return mf.__class__(mf)
        else:
            return mf

    mf_class = mf.__class__
    if mf_class.__doc__ is None:
        doc = ''
    else:
        doc = mf_class.__doc__
    class SFX2C1E_SCF(mf_class, x2c._X2C_SCF):
        __doc__ = doc + '''
        Attributes for spin-free X2C:
            with_x2c : X2C object
        '''
        def __init__(self, mf):
            self.__dict__.update(mf.__dict__)
            self.with_x2c = SpinFreeX2C(mf.mol)
            self._keys = self._keys.union(['with_x2c'])

        def get_hcore(self, cell=None, kpts=None, kpt=None):
            if cell is None: cell = self.cell
            if kpts is None:
                if getattr(self, 'kpts', None) is not None:
                    kpts = self.kpts
                else:
                    if kpt is None:
                        kpts = self.kpt
                    else:
                        kpts = kpt
            if self.with_x2c:
                hcore = self.with_x2c.get_hcore(cell, kpts)
                if isinstance(self, ghf.GHF):
                    if kpts.ndim == 1:
                        hcore = scipy.linalg.block_diag(hcore, hcore)
                    else:
                        hcore = [scipy.linalg.block_diag(h, h) for h in hcore]
                return hcore
            else:
                return mf_class.get_hcore(self, cell, kpts)

    return SFX2C1E_SCF(mf)

sfx2c = sfx2c1e

class X2C(x2c.X2C):

    exp_drop = getattr(__config__, 'pbc_x2c_X2C_exp_drop', 0.2)
    approx = getattr(__config__, 'pbc_x2c_X2C_approx', 'atom1e')
    xuncontract = getattr(__config__, 'pbc_x2c_X2C_xuncontract', True)
    basis = getattr(__config__, 'pbc_x2c_X2C_basis', None)

    def __init__(self, cell, kpts=None):
        self.cell = cell
        x2c.X2C.__init__(self, cell)

class SpinFreeX2C(X2C):
    def get_hcore(self, cell=None, kpts=None):
        if cell is None: cell = self.cell
        if kpts is None:
            kpts_lst = numpy.zeros((1,3))
        else:
            kpts_lst = numpy.reshape(kpts, (-1,3))

        xcell, contr_coeff = self.get_xmol(cell)
        with_df = aft.AFTDF(xcell)
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())
        if 'ATOM' in self.approx.upper():
            atom_slices = xcell.offset_nr_by_atom()
            nao = xcell.nao_nr()
            x = numpy.zeros((nao,nao))
            vloc = numpy.zeros((nao,nao))
            wloc = numpy.zeros((nao,nao))
            for ia in range(xcell.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                t1 = xcell.intor('int1e_kin', shls_slice=shls_slice)
                s1 = xcell.intor('int1e_ovlp', shls_slice=shls_slice)
                with xcell.with_rinv_at_nucleus(ia):
                    z = -xcell.atom_charge(ia)
                    v1 = z * xcell.intor('int1e_rinv', shls_slice=shls_slice)
                    w1 = z * xcell.intor('int1e_prinvp', shls_slice=shls_slice)
                vloc[p0:p1,p0:p1] = v1
                wloc[p0:p1,p0:p1] = w1
                x[p0:p1,p0:p1] = x2c._x2c1e_xmatrix(t1, v1, w1, s1, c)
        else:
            raise NotImplementedError

        t = xcell.pbc_intor('int1e_kin', 1, lib.HERMITIAN, kpts_lst)
        s = xcell.pbc_intor('int1e_ovlp', 1, lib.HERMITIAN, kpts_lst)
        v = with_df.get_nuc(kpts_lst)
        #w = get_pnucp(with_df, kpts_lst)
        if self.basis is not None:
            s22 = s
            s21 = pbcgto.intor_cross('int1e_ovlp', xcell, cell, kpts=kpts_lst)

        h1_kpts = []
        for k in range(len(kpts_lst)):
            # The treatment of pnucp local part has huge effects to hcore
            #h1 = x2c._get_hcore_fw(t[k], vloc, wloc, s[k], x, c) - vloc + v[k]
            #h1 = x2c._get_hcore_fw(t[k], v[k], w[k], s[k], x, c)
            h1 = x2c._get_hcore_fw(t[k], v[k], wloc, s[k], x, c)
            if self.basis is not None:
                c = lib.cho_solve(s22[k], s21[k])
                h1 = reduce(numpy.dot, (c.T, h1, c))
            if self.xuncontract and contr_coeff is not None:
                h1 = reduce(numpy.dot, (contr_coeff.T, h1, contr_coeff))
            h1_kpts.append(h1)

        if kpts is None or numpy.shape(kpts) == (3,):
            h1_kpts = h1_kpts[0]
        return lib.asarray(h1_kpts)

    def get_xmat(self, cell=None, kpts=None):
        if cell is None: cell = self.cell
        xcell, contr_coeff = self.get_xmol(cell)
        c = lib.param.LIGHT_SPEED
        assert('1E' in self.approx.upper())
        if 'ATOM' in self.approx.upper():
            atom_slices = xcell.offset_nr_by_atom()
            nao = xcell.nao_nr()
            x = numpy.zeros((nao,nao))
            for ia in range(xcell.natm):
                ish0, ish1, p0, p1 = atom_slices[ia]
                shls_slice = (ish0, ish1, ish0, ish1)
                t1 = xcell.intor('int1e_kin', shls_slice=shls_slice)
                s1 = xcell.intor('int1e_ovlp', shls_slice=shls_slice)
                with xcell.with_rinv_at_nucleus(ia):
                    z = -xcell.atom_charge(ia)
                    v1 = z * xcell.intor('int1e_rinv', shls_slice=shls_slice)
                    w1 = z * xcell.intor('int1e_prinvp', shls_slice=shls_slice)
                x[p0:p1,p0:p1] = x2c._x2c1e_xmatrix(t1, v1, w1, s1, c)
        else:
            raise NotImplementedError
        return x


# Use Ewald-like technique to compute spVsp.
# spVsp may not be divergent because the numerator spsp and the denominator
# in Coulomb kernel 4pi/G^2 are likely cancelled.  Even a real space lattice
# sum can converge to a finite value, it's difficult to accurately converge
# this value, i.e., large number of images in lattice summation is required.
def get_pnucp(mydf, kpts=None):
    cell = mydf.cell
    if kpts is None:
        kpts_lst = numpy.zeros((1,3))
    else:
        kpts_lst = numpy.reshape(kpts, (-1,3))

    log = logger.Logger(mydf.stdout, mydf.verbose)
    t1 = (logger.process_clock(), logger.perf_counter())

    nkpts = len(kpts_lst)
    nao = cell.nao_nr()
    nao_pair = nao * (nao+1) // 2

    Gv, Gvbase, kws = cell.get_Gv_weights(mydf.mesh)
    charge = -cell.atom_charges()
    kpt_allow = numpy.zeros(3)
    coulG = tools.get_coulG(cell, kpt_allow, mesh=mydf.mesh, Gv=Gv)
    coulG *= kws
    if mydf.eta == 0:
        wj = numpy.zeros((nkpts,nao_pair), dtype=numpy.complex128)
        SI = cell.get_SI(Gv)
        vG = numpy.einsum('i,ix->x', charge, SI) * coulG
        wj = numpy.zeros((nkpts,nao_pair), dtype=numpy.complex128)

    else:
        nuccell = copy.copy(cell)
        half_sph_norm = .5/numpy.sqrt(numpy.pi)
        norm = half_sph_norm/mole.gaussian_int(2, mydf.eta)
        chg_env = [mydf.eta, norm]
        ptr_eta = cell._env.size
        ptr_norm = ptr_eta + 1
        chg_bas = [[ia, 0, 1, 1, 0, ptr_eta, ptr_norm, 0] for ia in range(cell.natm)]
        nuccell._atm = cell._atm
        nuccell._bas = numpy.asarray(chg_bas, dtype=numpy.int32)
        nuccell._env = numpy.hstack((cell._env, chg_env))

        wj = lib.asarray(mydf._int_nuc_vloc(nuccell, kpts_lst, 'int3c2e_pvp1'))
        t1 = log.timer_debug1('pnucp pass1: analytic int', *t1)

        aoaux = ft_ao.ft_ao(nuccell, Gv)
        vG = numpy.einsum('i,xi->x', charge, aoaux) * coulG
        if cell.dimension == 3:
            nucbar = sum([z/nuccell.bas_exp(i)[0] for i,z in enumerate(charge)])
            nucbar *= numpy.pi/cell.vol

            ovlp = cell.pbc_intor('int1e_kin', 1, lib.HERMITIAN, kpts_lst)
            for k in range(nkpts):
                s = lib.pack_tril(ovlp[k])
                # *2 due to the factor 1/2 in T
                wj[k] -= nucbar*2 * s

    max_memory = max(2000, mydf.max_memory-lib.current_memory()[0])
    for aoaoks, p0, p1 in mydf.ft_loop(mydf.mesh, kpt_allow, kpts_lst,
                                       max_memory=max_memory, aosym='s2',
                                       intor='GTO_ft_pdotp'):
        for k, aoao in enumerate(aoaoks):
            if aft_jk.gamma_point(kpts_lst[k]):
                wj[k] += numpy.einsum('k,kx->x', vG[p0:p1].real, aoao.real)
                wj[k] += numpy.einsum('k,kx->x', vG[p0:p1].imag, aoao.imag)
            else:
                wj[k] += numpy.einsum('k,kx->x', vG[p0:p1].conj(), aoao)
    t1 = log.timer_debug1('contracting pnucp', *t1)

    wj_kpts = []
    for k, kpt in enumerate(kpts_lst):
        if aft_jk.gamma_point(kpt):
            wj_kpts.append(lib.unpack_tril(wj[k].real.copy()))
        else:
            wj_kpts.append(lib.unpack_tril(wj[k]))

    if kpts is None or numpy.shape(kpts) == (3,):
        wj_kpts = wj_kpts[0]
    return numpy.asarray(wj_kpts)


if __name__ == '__main__':
    from pyscf.pbc import scf
    cell = pbcgto.Cell()
    cell.build(unit = 'B',
               a = numpy.eye(3)*4,
               mesh = [11]*3,
               atom = 'H 0 0 0; H 0 0 1.8',
               verbose = 4,
               basis='sto3g')
    lib.param.LIGHT_SPEED = 2
    mf = scf.RHF(cell)
    mf.with_df = aft.AFTDF(cell)
    enr = mf.kernel()
    print('E(NR) = %.12g' % enr)

    mf = sfx2c1e(mf)
    esfx2c = mf.kernel()
    print('E(SFX2C1E) = %.12g' % esfx2c)

    mf = scf.KRHF(cell)
    mf.with_df = aft.AFTDF(cell)
    mf.kpts = cell.make_kpts([2,2,1])
    enr = mf.kernel()
    print('E(k-NR) = %.12g' % enr)

    mf = sfx2c1e(mf)
    esfx2c = mf.kernel()
    print('E(k-SFX2C1E) = %.12g' % esfx2c)

#    cell = pbcgto.M(unit = 'B',
#               a = numpy.eye(3)*4,
#               atom = 'H 0 0 0; H 0 0 1.8',
#               mesh = None,
#               dimension = 2,
#               basis='sto3g')
#    with_df = aft.AFTDF(cell)
#    w0 = get_pnucp(with_df, cell.make_kpts([2,2,1]))
#    with_df = aft.AFTDF(cell)
#    with_df.eta = 0
#    w1 = get_pnucp(with_df, cell.make_kpts([2,2,1]))
#    print(abs(w0-w1).max())


# pyscf/__all__.py

from . import lib, gto, scf, ao2mo
from .post_scf import *  # noqa: F403
from . import grad
from . import gw
from . import hessian
from . import lo
from . import qmmm
#from . import semiempirical
from . import sgx
from . import solvent
from . import tools

from . import geomopt

#try:
#    from . import mrpt
#except ImportError:
#    pass
#try:
#    from . import prop
#except ImportError:
#    pass
#try:
#    from . import dftd3
#except ImportError:
#    pass
#try:
#    from . import dmrgscf
#except ImportError:
#    pass
#try:
#    from . import icmpspt
#except ImportError:
#    pass
#try:
#    from . import shciscf
#except ImportError:
#    pass
#try:
#    from . import cornell_shci
#except ImportError:
#    pass

from . import pbc
from .pbc import __all__
del __all__



# pyscf/__config__.py

import os, sys

#
# All parameters initialized before loading pyscf_conf.py will be overwritten
# by the dynamic importing procedure.
#

DEBUG = False

MAX_MEMORY = int(os.environ.get('PYSCF_MAX_MEMORY', 4000)) # MB
TMPDIR = os.environ.get('TMPDIR', '.')
TMPDIR = os.environ.get('PYSCF_TMPDIR', TMPDIR)

VERBOSE = 3  # default logger level (logger.NOTE)
UNIT = 'angstrom'

#
# Loading pyscf_conf.py and overwriting above parameters
#
for conf_file in (os.environ.get('PYSCF_CONFIG_FILE', None),
                  os.path.join(os.path.abspath('.'), '.pyscf_conf.py'),
                  os.path.join(os.environ.get('HOME', '.'), '.pyscf_conf.py')):
    if conf_file is not None and os.path.isfile(conf_file):
        break
else:
    conf_file = None

if conf_file is not None:
    if sys.version_info < (3,0):
        import imp
        imp.load_source('pyscf.__config__', conf_file)
        del(imp)
    else:
        from importlib import machinery
        machinery.SourceFileLoader('pyscf.__config__', conf_file).load_module()
        del(machinery)
del(os, sys)

#
# All parameters initialized after loading pyscf_conf.py will be kept in the
# program.
#


# pyscf/__init__.py

# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

'''
*****************************************************
PySCF Python-based simulations of chemistry framework
*****************************************************

How to use
----------
There are two ways to access the documentation: the docstrings come with
the code, and an online program reference, available from
http://www.sunqm.net/pyscf/index.html

We recommend the enhanced Python interpreter `IPython <http://ipython.org>`_
and the web-based Python IDE `Ipython notebook <http://ipython.org/notebook.html>`_
to try out the package::

    >>> from pyscf import gto, scf
    >>> mol = gto.M(atom='H 0 0 0; H 0 0 1.2', basis='cc-pvdz')
    >>> mol.apply(scf.RHF).run()
    converged SCF energy = -1.06111199785749
    -1.06111199786

'''

__version__ = '2.0.1'

import os
import sys
# Avoid too many threads being created in OMP loops.
# See issue https://github.com/pyscf/pyscf/issues/317
if 'OPENBLAS_NUM_THREADS' not in os.environ:
    os.environ['OPENBLAS_NUM_THREADS'] = '1'
if 'MKL_NUM_THREADS' not in os.environ:
    os.environ['MKL_NUM_THREADS'] = '1'

# Load modules which are developed as plugins of the namespace package
PYSCF_EXT_PATH = os.getenv('PYSCF_EXT_PATH')
if PYSCF_EXT_PATH:
    for p in PYSCF_EXT_PATH.split(':'):
        if os.path.isdir(p):
            submodules = os.listdir(p)
            if 'pyscf' in submodules:
                # PYSCF_EXT_PATH points to the root directory of a submodule
                __path__.append(os.path.join(p, 'pyscf'))
            else:
                # Load all modules in PYSCF_EXT_PATH if it's a folder that
                # contains all extended modules
                for f in submodules:
                    __path__.append(os.path.join(p, f, 'pyscf'))
                del f
        elif os.path.exists(p):
            # Load all moduels defined inside the file PYSCF_EXT_PATH
            with open(p, 'r') as f:
                __path__.extend(f.read().splitlines())
            del f
    del p
elif '/site-packages/' in __file__ or '/dist-packages/' in __file__:
    # If pyscf has been installed in the standard runtime path (typically
    # under site-packages), and plugins are installed with the pip editable
    # mode, load namespace plugins. In this case, it is likely all modules are
    # managed by pip/conda or their virtual environments. It is safe to search
    # namespace plugins according to the old style of PEP 420.
    __path__ = __import__('pkgutil').extend_path(__path__, __name__)
else:
    # We need a better way to load plugins if pyscf is imported by the
    # environment variable PYTHONPATH. Current treatment may mix installed
    # plugins (e.g.  through pip install) with the developing plugins which
    # were accidentally placed under PYTHONPATH. When PYTHONPATH mechanism is
    # taken, an explicit list of extended paths (using environment
    # PYSCF_EXT_PATH) is recommended.
    __path__ = __import__('pkgutil').extend_path(__path__, __name__)
    if not all('/site-packages/' in p for p in __path__[1:]):
        sys.stderr.write('pyscf plugins found in \n%s\n'
                         'When PYTHONPATH is set, it is recommended to load '
                         'these plugins through the environment variable '
                         'PYSCF_EXT_PATH' % '\n'.join(__path__[1:]))

from distutils.version import LooseVersion
import numpy
if LooseVersion(numpy.__version__) <= '1.8.0':
    raise SystemError("You're using an old version of Numpy (%s). "
                      "It is recommended to upgrade numpy to 1.8.0 or newer. \n"
                      "You still can use all features of PySCF with the old numpy by removing this warning msg. "
                      "Some modules (DFT, CC, MRPT) might be affected because of the bug in old numpy." %
                      numpy.__version__)
elif '1.16.2' <= LooseVersion(numpy.__version__) < '1.18':
    #sys.stderr.write('Numpy 1.16 has memory leak bug  '
    #                 'https://github.com/numpy/numpy/issues/13808\n'
    #                 'It is recommended to downgrade to numpy 1.15 or older\n')
    import ctypes
    from numpy.core import _internal
    def _get_void_ptr(arr):
        simple_arr = numpy.asarray(_internal._unsafe_first_element_pointer(arr))
        c_arr = (ctypes.c_char * 0).from_buffer(simple_arr)
        return ctypes.cast(ctypes.byref(c_arr), ctypes.c_void_p)
    # patch _get_void_ptr as a workaround to numpy issue #13808
    _internal._get_void_ptr = _get_void_ptr

from pyscf import __config__
from pyscf import lib
from pyscf import gto
from pyscf import scf
from pyscf import ao2mo

# Whether to enable debug mode. When this flag is set, some modules may run
# extra debug code.
DEBUG = __config__.DEBUG

def M(**kwargs):
    '''Main driver to create Molecule object (mol) or Material crystal object (cell)'''
    from pyscf import __all__
    if kwargs.get('a') is not None:  # a is crystal lattice parameter
        return __all__.pbc.gto.M(**kwargs)
    else:  # Molecule
        return gto.M(**kwargs)

del os, sys, LooseVersion


# pyscf/post_scf.py

from . import mp, ci, cc, mcscf, fci, tdscf
try:
    from . import doci
except ImportError:
    pass


# pyscf/pbc/mp/test/test_dm.py

#!/usr/bin/env python
# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
from functools import reduce
import numpy as np
from pyscf.pbc import gto, scf, mp

cell = gto.Cell()
cell.atom = """
H  0.0 0.0 0.0
F  0.9 0.0 0.0
"""
cell.basis = 'sto-3g'
cell.a = [[2.82, 0, 0], [0, 2.82, 0], [0, 0, 2.82]]
cell.dimension = 1
cell.output = '/dev/null'
cell.build()

nk = [2,1,1]
kpts = cell.make_kpts(nk)
nkpts = len(kpts)
kmf = scf.KRHF(cell, kpts = kpts, exxdiv=None).density_fit()
kmf.kernel()

def tearDownModule():
    global cell, kmf
    del cell, kmf

class KnownValues(unittest.TestCase):
    def test_kmp2_contract_eri_dm(self):
        kmp2 = mp.KMP2(kmf)
        kmp2.kernel()
        e_tot = kmp2.e_tot

        mo_coeff = kmf.mo_coeff
        hcore = kmf.get_hcore()
        for k in range(nkpts):
            hcore[k] = reduce(np.dot, (mo_coeff[k].T.conj(), hcore[k], mo_coeff[k]))

        dm1 = kmp2.make_rdm1()
        dm2 = kmp2.make_rdm2()
        e = 0
        for k in range(nkpts):
            e += np.einsum('pq,qp', dm1[k], hcore[k]).real / nkpts
        ao2mo = kmp2._scf.with_df.ao2mo
        idx = 0
        for kp in range(nkpts):
            for kq in range(nkpts):
                for kr in range(nkpts):
                    ks = kmp2.khelper.kconserv[kp,kq,kr]
                    mop = mo_coeff[kp]
                    moq = mo_coeff[kq]
                    mor = mo_coeff[kr]
                    mos = mo_coeff[ks]
                    eri = ao2mo((mop,moq,mor,mos), 
                          (kpts[kp], kpts[kq], kpts[kr], kpts[ks]),
                          compact=False).reshape(mop.shape[-1],moq.shape[-1],
                          mor.shape[-1],mos.shape[-1]) / nkpts
                    e += np.einsum('pqrs,pqrs',dm2[idx], eri).real * 0.5 / nkpts
                    idx += 1
        e += cell.energy_nuc()
        self.assertAlmostEqual(e, e_tot, 9)

if __name__ == "__main__":
    print("Full Tests for kmp2 rdm")
    unittest.main()


# pyscf/pbc/mp/test/test_kpoint.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
import numpy as np

from pyscf.pbc import gto as pbcgto
from pyscf.pbc import scf as pbcscf
import pyscf.pbc.mp
import pyscf.pbc.mp.kmp2


cell = pbcgto.Cell()
cell.unit = 'B'
L = 7
cell.atom.extend([['Be', (L/2.,  L/2., L/2.)]])
cell.a = 7 * np.identity(3)
cell.a[1,0] = 5.0

cell.basis = 'gth-szv'
cell.pseudo = 'gth-pade-q2'
cell.mesh = [12]*3
cell.verbose = 5
cell.output = '/dev/null'
cell.build()

def build_h_cell():
    # Returns FCC H cell.
    cell = pbcgto.Cell()
    cell.atom = [['H', (0.000000000, 0.000000000, 0.000000000)],
                 ['H', (0.000000000, 0.500000000, 0.250000000)],
                 ['H', (0.500000000, 0.500000000, 0.500000000)],
                 ['H', (0.500000000, 0.000000000, 0.750000000)]]
    cell.unit = 'Bohr'
    cell.a = [[1.,0.,0.],[0.,1.,0],[0,0,2.2]]
    cell.verbose = 7
    cell.spin = 0
    cell.charge = 0
    cell.basis = [[0, [1.0, 1]],]
    cell.pseudo = 'gth-pade'
    cell.output = '/dev/null'
    cell.max_memory = 1000
    for i in range(len(cell.atom)):
        cell.atom[i][1] = tuple(np.dot(np.array(cell.atom[i][1]),np.array(cell.a)))
    cell.build()
    return cell

def run_kcell(cell, nk):
    abs_kpts = cell.make_kpts(nk, wrap_around=True)
    kmf = pbcscf.KRHF(cell, abs_kpts)
    kmf.conv_tol = 1e-12
    ekpt = kmf.scf()
    mp = pyscf.pbc.mp.kmp2.KMP2(kmf).run()
    return ekpt, mp.e_corr

def run_kcell_complex(cell, nk):
    abs_kpts = cell.make_kpts(nk, wrap_around=True)
    kmf = pbcscf.KRHF(cell, abs_kpts)
    kmf.conv_tol = 1e-12
    ekpt = kmf.scf()
    kmf.mo_coeff = [kmf.mo_coeff[i].astype(np.complex128) for i in range(np.prod(nk))]
    mp = pyscf.pbc.mp.kmp2.KMP2(kmf).run()
    return ekpt, mp.e_corr

class KnownValues(unittest.TestCase):
    def test_111(self):
        nk = (1, 1, 1)
        escf, emp = run_kcell(cell,nk)
        self.assertAlmostEqual(escf, -1.2061049658473704, 9)
        self.assertAlmostEqual(emp, -5.44597932944397e-06, 9)
        escf, emp = run_kcell_complex(cell,nk)
        self.assertAlmostEqual(emp, -5.44597932944397e-06, 9)

    def test_311_high_cost(self):
        nk = (3, 1, 1)
        escf, emp = run_kcell(cell,nk)
        self.assertAlmostEqual(escf, -1.0585001200928885, 9)
        self.assertAlmostEqual(emp, -7.9832274354253814e-06, 9)

    def test_h4_fcc_k2_frozen(self):
        '''Metallic hydrogen fcc lattice with frozen lowest lying occupied
        and highest lying virtual orbitals.  Checks versus a corresponding
        supercell calculation.

        NOTE: different versions of the davidson may converge to a different
        solution for the k-point IP/EA eom.  If you're getting the wrong
        root, check to see if it's contained in the supercell set of
        eigenvalues.'''
        cell = build_h_cell()

        nmp = [2, 1, 1]

        kmf = pbcscf.KRHF(cell)
        kmf.kpts = cell.make_kpts(nmp, scaled_center=[0.0,0.0,0.0])
        e = kmf.kernel()

        frozen = [[0, 3], []]
        mymp = pyscf.pbc.mp.kmp2.KMP2(kmf, frozen=frozen)
        ekmp2, _ = mymp.kernel()
        self.assertAlmostEqual(ekmp2, -0.022416773725207319, 6)
        self.assertAlmostEqual(mymp.e_tot, 2.155470531550681, 6)

        # Start of supercell calculations
        from pyscf.pbc.tools.pbc import super_cell
        supcell = super_cell(cell, nmp)
        supcell.build()
        mf = pbcscf.KRHF(supcell)
        e = mf.kernel()

        mysmp = pyscf.pbc.mp.kmp2.KMP2(mf, frozen=[0, 7])
        emp2, _ = mysmp.kernel()
        emp2 /= np.prod(nmp)
        self.assertAlmostEqual(emp2, -0.022416773725207319, 6)

    def test_h4_fcc_k2_frozen_df_nocc(self):
        '''Metallic hydrogen fcc, test different nocc at k points.'''
        cell = build_h_cell()

        nmp = [2, 1, 1]

        kmf = pbcscf.KRHF(cell).density_fit()
        kmf.kpts = cell.make_kpts(nmp, scaled_center=[0.0,0.0,0.0])
        e = kmf.kernel()

        frozen = [[0, 3], []]
        mymp = pyscf.pbc.mp.kmp2.KMP2(kmf, frozen=frozen)
        ekmp2, _ = mymp.kernel()
        self.assertAlmostEqual(ekmp2, -0.016333989667540873, 6)
        self.assertAlmostEqual(mymp.e_tot, 2.329841282521279, 6)

    def test_rdm1(self):
        cell = pbcgto.Cell()
        cell.atom = '''Al 0 0 0'''
        cell.basis = 'gth-szv'
        cell.pseudo = 'gth-pade'
        cell.a = '''
        2.47332919 0 1.42797728
        0.82444306 2.33187713 1.42797728
        0, 0, 2.85595455
        '''
        cell.unit = 'angstrom'
        cell.spin = 1
        cell.build()
        cell.verbose = 4
        cell.incore_anyway = True

        abs_kpts = cell.make_kpts((2, 1, 1), scaled_center=(.1, .2, .3))
        kmf = pbcscf.KRHF(cell, abs_kpts)
        kmf.conv_tol = 1e-12
        kmf.kernel()
        mp = pyscf.pbc.mp.kmp2.KMP2(kmf)
        mp.kernel(with_t2=True)
        self.assertAlmostEqual(mp.e_corr, -0.00162057921874043, 6)
        dm = mp.make_rdm1()
        np.testing.assert_allclose(np.trace(dm[0]) + np.trace(dm[1]), 6)
        for kdm in dm:
            np.testing.assert_allclose(kdm, kdm.conj().T)

    def test_kmp2_with_cderi(self):
        nk = (1, 1, 1)
        abs_kpts = cell.make_kpts(nk, wrap_around=True)
        kmf = pbcscf.KRHF(cell, abs_kpts).density_fit()
        kmf.conv_tol = 1e-12
        ekpt = kmf.scf()
        kmf2 = pbcscf.KRHF(cell, abs_kpts).density_fit()
        kmf2.conv_tol = 1e-12
        kmf2.with_df._cderi = kmf.with_df._cderi
        ekpt2 = kmf2.scf()
        mp = pyscf.pbc.mp.kmp2.KMP2(kmf2).run()        

        self.assertAlmostEqual(ekpt2, -1.2053666821021261, 9)
        self.assertAlmostEqual(mp.e_corr, -6.9881475423322723e-06, 9)


if __name__ == '__main__':
    print("Full kpoint test")
    unittest.main()



# pyscf/pbc/mp/test/test_mask.py

#!/usr/bin/env python

import unittest
import numpy as np

import pyscf.pbc.mp.kmp2
from pyscf.pbc.mp.kmp2 import get_nocc, get_nmo, get_frozen_mask

class fake_mp:
    def __init__(self, frozen, mo_occ, nkpts):
        self._nocc = None
        self._nmo = None
        self.frozen = frozen
        self.mo_occ = mo_occ
        self.nkpts = nkpts

    get_nocc = get_nocc
    get_nmo = get_nmo

class KnownValues(unittest.TestCase):
    def test_no_frozen(self):
        mp = fake_mp(frozen=None, mo_occ=[np.array([2, 2, 2, 0, 0]),], nkpts=1)
        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, 3)
        self.assertAlmostEqual(nmo, 5)

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertListEqual(nocc, [3])
        self.assertListEqual(nmo, [5])

    def test_frozen_int(self):
        mp = fake_mp(frozen=1, mo_occ=[np.array([2, 2, 2, 0, 0]), np.array([2, 2, 0, 0, 0])], nkpts=2)
        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, 2)
        self.assertAlmostEqual(nmo, 5)  # 2 occupied, 3 virtual

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertListEqual(nocc, [2, 1])
        self.assertListEqual(nmo, [4, 4])

    def test_frozen_list1(self):
        mp = fake_mp(frozen=[1,], mo_occ=[np.array([2, 2, 2, 0, 0]), np.array([2, 2, 0, 0, 0])], nkpts=2)
        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, 2)
        self.assertAlmostEqual(nmo, 5)  # 2 occupied, 3 virtual

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertListEqual(nocc, [2, 1])
        self.assertListEqual(nmo, [4, 4])

    def test_frozen_list2(self):
        # Freeze virtual not contained in set
        mp = fake_mp(frozen=[4, 5], mo_occ=[np.array([2, 2, 2, 0, 0]), np.array([2, 2, 0, 0, 0])], nkpts=2)
        self.assertRaises(RuntimeError, get_nocc, mp)
        self.assertRaises(RuntimeError, get_nmo, mp)  # Fails because it pads by calling get_nocc

    def test_frozen_repeated_orbital(self):
        mp = fake_mp(frozen=[[1, 1], [0]], mo_occ=[np.array([2, 2, 2, 0, 0]), np.array([2, 2, 0, 0, 0])], nkpts=2)
        self.assertRaises(RuntimeError, get_nocc, mp)
        self.assertRaises(RuntimeError, get_nmo, mp)  # Fails because it pads by calling get_nocc

    def test_frozen_kpt_list1(self):
        mp = fake_mp(frozen=[[0, 1,], [0]], mo_occ=[np.array([2, 2, 2, 0, 0]), np.array([2, 2, 0, 0, 0])], nkpts=2)
        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, 1)
        self.assertAlmostEqual(nmo, 4)  # 1 occupied, 3 virtual

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertListEqual(nocc, [1, 1])
        self.assertListEqual(nmo, [3, 4])

    def test_frozen_kpt_list2(self):
        mp = fake_mp(frozen=[[0,1],[],[0]], mo_occ=[np.array([2, 2, 2, 0, 0])] * 3, nkpts=3)
        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, 3)
        self.assertAlmostEqual(nmo, 5)  # 2nd k-point has 3 occupied and 2 virtual orbitals

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertListEqual(nocc, [1, 3, 2])
        self.assertListEqual(nmo, [3, 5, 4])

    def test_frozen_kpt_list3(self):
        mp = fake_mp(frozen=[[0,1,3],[3],[0]], mo_occ=[np.array([2, 2, 2, 0, 0])] * 3, nkpts=3)
        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, 3)
        self.assertAlmostEqual(nmo, 5)  # 2nd k-point has 3 occupied and 2 virtual orbitals

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertListEqual(nocc, [1, 3, 2])
        self.assertListEqual(nmo, [2, 4, 4])

if __name__ == '__main__':
    print("Full mask test")
    unittest.main()



# pyscf/pbc/mp/test/test_mask_uhf.py

#!/usr/bin/env python

import unittest
import numpy as np

import pyscf.pbc.mp.kump2
from pyscf.pbc.mp.kump2 import get_nocc, get_nmo, get_frozen_mask, padding_k_idx

class fake_ump:
    def __init__(self, frozen, mo_occ, nkpts):
        self._nocc = None
        self._nmo = None
        self.frozen = frozen
        self.mo_occ = mo_occ
        self.nkpts = nkpts

    get_nocc = get_nocc
    get_nmo = get_nmo

    @property
    def nocc(self):
        return self.get_nocc()
    @property
    def nmo(self):
        return self.get_nmo()

class KnownValues(unittest.TestCase):
    def test_no_frozen(self):
        mo_occ = [[np.array([1,1,0,0,0]),],
                  [np.array([1,1,1,0,0]),]]
        mp = fake_ump(frozen=None, mo_occ=mo_occ, nkpts=1)
        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, (2,3))
        self.assertAlmostEqual(nmo, (5,5))

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertAlmostEqual(nocc, ([2],[3]))
        self.assertAlmostEqual(nmo, ([5],[5]))

    def test_frozen_int(self):
        mo_occ = [[np.array([1,1,0,0,0]),],
                  [np.array([1,1,1,0,0]),]]
        mp = fake_ump(frozen=1, mo_occ=mo_occ, nkpts=1)
        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, (1,2))
        self.assertAlmostEqual(nmo, (4,4))  # 2 occupied, 3 virtual

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertAlmostEqual(nocc, ([1],[2]))
        self.assertAlmostEqual(nmo, ([4],[4]))

    def test_frozen_list1(self):
        mo_occ = [[np.array([1,1,0,0,0]),],
                  [np.array([1,1,1,0,0]),]]
        mp = fake_ump(frozen=([1,],[1,]), mo_occ=mo_occ, nkpts=1)
        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, (1,2))
        self.assertAlmostEqual(nmo, (4,4))

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertAlmostEqual(nocc, ([1],[2]))
        self.assertAlmostEqual(nmo, ([4],[4]))

    def test_frozen_list2(self):
        # Freeze virtual not contained in set
        mo_occ = [[np.array([1,1,0,0,0]),],
                  [np.array([1,1,1,0,0]),]]
        mp = fake_ump(frozen=([5,], [5,]), mo_occ=mo_occ, nkpts=1)
        self.assertRaises(RuntimeError, get_nocc, mp)
        self.assertRaises(RuntimeError, get_nmo, mp)  # Fails because it pads by calling get_nocc

    def test_frozen_list3(self):
        # Freeze all occupied alpha orbitals
        mo_occ = [[np.array([1,1,0,0,0]),],
                  [np.array([1,1,1,0,0]),]]
        mp = fake_ump(frozen=([0,1], [0,1]), mo_occ=mo_occ, nkpts=1)
        self.assertRaises(AssertionError, get_nocc, mp)  # No occupied alpha
        self.assertRaises(AssertionError, get_nmo, mp)  # Fails because calls get_nocc

    def test_frozen_list4(self):
        frozen = ([0,2], [0,1,3,4])
        mo_occ = [[np.array([1,1,0,0,0]),],
                  [np.array([1,1,1,0,0]),]]
        mp = fake_ump(frozen=frozen, mo_occ=mo_occ, nkpts=1)

        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, (1,1))
        self.assertAlmostEqual(nmo, (3,1))

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertAlmostEqual(nocc, ([1],[1]))
        self.assertAlmostEqual(nmo, ([3],[1]))

    def test_frozen_repeated_orbital(self):
        mo_occ = [[np.array([1,1,0,0,0]),],
                  [np.array([1,1,1,0,0]),]]
        mp = fake_ump(frozen=([0,1,1], [0,1]), mo_occ=mo_occ, nkpts=1)
        self.assertRaises(RuntimeError, get_nocc, mp)
        self.assertRaises(RuntimeError, get_nmo, mp)  # Fails because it pads by calling get_nocc

    def test_frozen_kpt_list1(self):
        frozen = ([[2], [0,]],
                  [[0,1,3,4], [0,2]])
        mo_occ = [[np.array([1, 1, 1, 0, 0]),np.array([1, 1, 0, 0, 0]),],
                  [np.array([1, 0, 0, 0, 0]),np.array([1, 1, 1, 0, 0]),]]
        mp = fake_ump(frozen=frozen, mo_occ=mo_occ, nkpts=2)

        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, (2, 1))
        self.assertAlmostEqual(nmo, (5, 3))  # Take the max occ and max vir from kpoints -
                                             # spin0: 2 occ (from first kpoint), 3 vir (from second kpoint)
                                             # spin1: 1 occ (from second kpoint),2 vir (from second kpoint)

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertAlmostEqual(nocc, ([2, 1], [0, 1]))
        self.assertAlmostEqual(nmo, ([4, 4], [1, 3]))

        #  Using `O` and `V` as occupied and virtual and `-` as a padded/frozen orbital,
        #  we will have the following orbital list
        #
        #        alpha spin      |  beta spin
        #  k1 :  O O | - V V     |  - | - V
        #  k2 :  O - | V V V     |  O | V V
        split_idx = padding_k_idx(mp, kind='split')
        outa = [[y for x in split_idx[0][idx] for y in x] for idx in range(2)]  # Flatten the list for ease-of-use
        outb = [[y for x in split_idx[1][idx] for y in x] for idx in range(2)]
        expecteda = [[0, 1, 0], [1, 2, 0, 1, 2]]  # [occ_idx, vir_idx] for alpha
        expectedb = [[0], [1, 0, 1]]              # [occ_idx, vir_idx] for beta
        self.assertAlmostEqual(outa, expecteda)
        self.assertAlmostEqual(outb, expectedb)

        outa, outb = padding_k_idx(mp, kind='joint')
        outa = [y for x in outa for y in x]
        outb = [y for x in outb for y in x]
        expecteda = [0,1,3,4,0,2,3,4]  # [occa_idx_k1, occa_idx_k2]
        expectedb = [2,0,1,2]
        self.assertAlmostEqual(outa, expecteda)
        self.assertAlmostEqual(outb, expectedb)

    def test_frozen_kpt_list2(self):
        frozen = ([[0,1], [0,]],
                  [[0], [0,1]])
        mo_occ = [[np.array([1, 1, 1, 0, 0]),np.array([1, 1, 0, 0, 0]),],
                  [np.array([1, 1, 0, 0, 0]),np.array([1, 1, 1, 0, 0]),]]
        mp = fake_ump(frozen=frozen, mo_occ=mo_occ, nkpts=2)

        nocc = get_nocc(mp)
        nmo = get_nmo(mp)
        self.assertAlmostEqual(nocc, (1, 1))
        self.assertAlmostEqual(nmo, (4, 4))  # We deleted the first occ from each of the alpha/beta, so
                                             # we have 1 occupied and 3 virtuals.

        nocc = get_nocc(mp, per_kpoint=True)
        nmo = get_nmo(mp, per_kpoint=True)
        self.assertAlmostEqual(nocc, ([1, 1], [1, 1]))
        self.assertAlmostEqual(nmo, ([3, 4], [4, 3]))

        #  Using `O` and `V` as occupied and virtual and `-` as a padded/frozen orbital,
        #  we will have the following orbital list
        #
        #        alpha spin  |  beta spin
        #  k1 :  O | - V V   |  O | V V V
        #  k2 :  O | V V V   |  O | - V V
        split_idx = padding_k_idx(mp, kind='split')
        outa = [[y for x in split_idx[0][idx] for y in x] for idx in range(2)]  # Flatten the list for ease-of-use
        outb = [[y for x in split_idx[1][idx] for y in x] for idx in range(2)]
        expecteda = [[0, 0], [1, 2, 0, 1, 2]]  # [occ_idx, vir_idx] for alpha
        expectedb = [[0, 0], [0, 1, 2, 1, 2]]  # [occ_idx, vir_idx] for beta
        self.assertAlmostEqual(outa, expecteda)
        self.assertAlmostEqual(outb, expectedb)

        outa, outb = padding_k_idx(mp, kind='joint')
        outa = [y for x in outa for y in x]
        outb = [y for x in outb for y in x]
        expecteda = [0,2,3,0,1,2,3]  # [occa_idx_k1, occa_idx_k2]
        expectedb = [0,1,2,3,0,2,3]
        self.assertAlmostEqual(outa, expecteda)
        self.assertAlmostEqual(outb, expectedb)


if __name__ == '__main__':
    print("Full mask test")
    unittest.main()



# pyscf/pbc/mp/test/test_padding.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import unittest
import numpy as np

from pyscf.pbc import gto as pbcgto
from pyscf.pbc import scf as pbcscf
import pyscf.pbc.mp
import pyscf.pbc.mp.kmp2
from pyscf.scf.addons import remove_linear_dep_


a0 = 1.78339987 * 0.96  # squeezed
atom = "C 0 0 0; C %.10f %.10f %.10f" % (a0*0.5, a0*0.5, a0*0.5)
a = np.asarray(
    [[0., a0, a0],
    [a0, 0., a0],
    [a0, a0, 0.]])
basis = "gth-dzvp"
pseudo = "gth-pade"
ke_cutoff = 100
cell = pbcgto.Cell(atom=atom, a=a, basis=basis, pseudo=pseudo,
                   ke_cutoff=ke_cutoff)
cell.verbose = 4
cell.output = '/dev/null'
cell.build()

def tearDownModule():
    global cell
    cell.stdout.close()
    del cell

def run_kcell(cell, nk):
    abs_kpts = cell.make_kpts(nk, wrap_around=True)
    kmf = pbcscf.KRHF(cell, abs_kpts)
    kmf = remove_linear_dep_(kmf, threshold=1e-5, lindep=1e-6)
    kmf.conv_tol = 1e-12
    ekpt = kmf.scf()
    mp = pyscf.pbc.mp.kmp2.KMP2(kmf).run()
    return ekpt, mp.e_corr

class KnownValues(unittest.TestCase):
    def test_221(self):
        cell1 = cell.copy()
        cell1.basis = 'gth-dzv'
        cell1.build()

        nk = [2, 2, 1]
        abs_kpts = cell1.make_kpts(nk, wrap_around=True)
        kmf = pbcscf.KRHF(cell1, abs_kpts)
        kmf = remove_linear_dep_(kmf, threshold=1e-4, lindep=1e-6)
        ekpt = kmf.scf()
        mp = pyscf.pbc.mp.kmp2.KMP2(kmf).run()

        self.assertAlmostEqual(ekpt,      -10.835614361742607, 8)
        self.assertAlmostEqual(mp.e_corr, -0.1522294774708119, 8)

    def test_222_high_cost(self):
        nk = (2, 2, 2)
        escf, emp = run_kcell(cell,nk)
        self.assertAlmostEqual(escf, -11.0152342492995, 9)
        self.assertAlmostEqual(emp, -0.233433093787577, 9)


if __name__ == '__main__':
    print("padding test")
    unittest.main()


# pyscf/pbc/__all__.py

from . import gto
from . import scf
from . import lib
from . import mp, ci, cc, tdscf
from . import ao2mo
from . import df
#from . import grad
#from . import lo
#from . import prop
from . import tools
from . import x2c
try:
    from . import dft
except (ImportError, IOError):
    pass


# pyscf/pbc/__init__.py

# Copyright 2014-2020 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# If extension plugins are installed in pyscf, search and load the pbc
# submodule in all plugins if applicable
if len(__import__('pyscf').__path__) > 1:
    __path__ = __import__('pkgutil').extend_path(__path__, __name__)

from pyscf.pbc import gto
from pyscf.pbc import scf
#from pyscf.pbc import tools

DEBUG = False

M = gto.M


# pyscf/pbc/mp/__init__.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pyscf.pbc import scf
from pyscf.pbc.mp import mp2
from pyscf.pbc.mp import kmp2

def RMP2(mf, frozen=None, mo_coeff=None, mo_occ=None):
    mf = scf.addons.convert_to_rhf(mf)
    return mp2.RMP2(mf, frozen, mo_coeff, mo_occ)

MP2 = RMP2

def UMP2(mf, frozen=None, mo_coeff=None, mo_occ=None):
    mf = scf.addons.convert_to_uhf(mf)
    return mp2.UMP2(mf, frozen, mo_coeff, mo_occ)

def GMP2(mf, frozen=None, mo_coeff=None, mo_occ=None):
    mf = scf.addons.convert_to_ghf(mf)
    return mp2.GMP2(mf, frozen, mo_coeff, mo_occ)

def KRMP2(mf, frozen=None, mo_coeff=None, mo_occ=None):
    return kmp2.KRMP2(mf, frozen, mo_coeff, mo_occ)

KMP2 = KRMP2


# pyscf/pbc/mp/kmp2.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: Timothy Berkelbach <tim.berkelbach@gmail.com>
#         James McClain <jdmcclain47@gmail.com>
#


'''
kpoint-adapted and spin-adapted MP2
t2[i,j,a,b] = <ij|ab> / D_ij^ab

t2 and eris are never stored in full, only a partial
eri of size (nkpts,nocc,nocc,nvir,nvir)
'''

import numpy as np
from scipy.linalg import block_diag
import h5py

from pyscf import lib
from pyscf.lib import logger, einsum
from pyscf.mp import mp2
from pyscf.pbc import df
from pyscf.pbc.lib import kpts_helper
from pyscf.lib.parameters import LARGE_DENOM
from pyscf import __config__

WITH_T2 = getattr(__config__, 'mp_mp2_with_t2', True)

def kernel(mp, mo_energy, mo_coeff, verbose=logger.NOTE, with_t2=WITH_T2):
    """Computes k-point RMP2 energy.

    Args:
        mp (KMP2): an instance of KMP2
        mo_energy (list): a list of numpy.ndarray. Each array contains MO energies of
                          shape (Nmo,) for one kpt
        mo_coeff (list): a list of numpy.ndarray. Each array contains MO coefficients
                         of shape (Nao, Nmo) for one kpt
        verbose (int, optional): level of verbosity. Defaults to logger.NOTE (=3).
        with_t2 (bool, optional): whether to compute t2 amplitudes. Defaults to WITH_T2 (=True).

    Returns:
        KMP2 energy and t2 amplitudes (=None if with_t2 is False)
    """
    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.new_logger(mp, verbose)

    mp.dump_flags()
    nmo = mp.nmo
    nocc = mp.nocc
    nvir = nmo - nocc
    nkpts = mp.nkpts

    with_df_ints = mp.with_df_ints and isinstance(mp._scf.with_df, df.GDF)

    mem_avail = mp.max_memory - lib.current_memory()[0]
    mem_usage = (nkpts * (nocc * nvir)**2) * 16 / 1e6
    if with_df_ints:
        mydf = mp._scf.with_df
        if mydf.auxcell is None:
            # Calculate naux based on precomputed GDF integrals
            naux = mydf.get_naoaux()
        else:
            naux = mydf.auxcell.nao_nr()

        mem_usage += (nkpts**2 * naux * nocc * nvir) * 16 / 1e6
    if with_t2:
        mem_usage += (nkpts**3 * (nocc * nvir)**2) * 16 / 1e6
    if mem_usage > mem_avail:
        raise MemoryError('Insufficient memory! MP2 memory usage %d MB (currently available %d MB)'
                          % (mem_usage, mem_avail))

    eia = np.zeros((nocc,nvir))
    eijab = np.zeros((nocc,nocc,nvir,nvir))

    fao2mo = mp._scf.with_df.ao2mo
    kconserv = mp.khelper.kconserv
    emp2 = 0.
    oovv_ij = np.zeros((nkpts,nocc,nocc,nvir,nvir), dtype=mo_coeff[0].dtype)

    mo_e_o = [mo_energy[k][:nocc] for k in range(nkpts)]
    mo_e_v = [mo_energy[k][nocc:] for k in range(nkpts)]

    # Get location of non-zero/padded elements in occupied and virtual space
    nonzero_opadding, nonzero_vpadding = padding_k_idx(mp, kind="split")

    if with_t2:
        t2 = np.zeros((nkpts, nkpts, nkpts, nocc, nocc, nvir, nvir), dtype=complex)
    else:
        t2 = None

    # Build 3-index DF tensor Lov
    if with_df_ints:
        Lov = _init_mp_df_eris(mp)

    for ki in range(nkpts):
        for kj in range(nkpts):
            for ka in range(nkpts):
                kb = kconserv[ki,ka,kj]
                # (ia|jb)
                if with_df_ints:
                    oovv_ij[ka] = (1./nkpts) * einsum("Lia,Ljb->iajb", Lov[ki, ka], Lov[kj, kb]).transpose(0,2,1,3)
                else:
                    orbo_i = mo_coeff[ki][:,:nocc]
                    orbo_j = mo_coeff[kj][:,:nocc]
                    orbv_a = mo_coeff[ka][:,nocc:]
                    orbv_b = mo_coeff[kb][:,nocc:]
                    oovv_ij[ka] = fao2mo((orbo_i,orbv_a,orbo_j,orbv_b),
                                         (mp.kpts[ki],mp.kpts[ka],mp.kpts[kj],mp.kpts[kb]),
                                         compact=False).reshape(nocc,nvir,nocc,nvir).transpose(0,2,1,3) / nkpts
            for ka in range(nkpts):
                kb = kconserv[ki,ka,kj]

                # Remove zero/padded elements from denominator
                eia = LARGE_DENOM * np.ones((nocc, nvir), dtype=mo_energy[0].dtype)
                n0_ovp_ia = np.ix_(nonzero_opadding[ki], nonzero_vpadding[ka])
                eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]

                ejb = LARGE_DENOM * np.ones((nocc, nvir), dtype=mo_energy[0].dtype)
                n0_ovp_jb = np.ix_(nonzero_opadding[kj], nonzero_vpadding[kb])
                ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb]

                eijab = lib.direct_sum('ia,jb->ijab',eia,ejb)
                t2_ijab = np.conj(oovv_ij[ka]/eijab)
                if with_t2:
                    t2[ki, kj, ka] = t2_ijab
                woovv = 2*oovv_ij[ka] - oovv_ij[kb].transpose(0,1,3,2)
                emp2 += einsum('ijab,ijab', t2_ijab, woovv).real

    log.timer("KMP2", *cput0)

    emp2 /= nkpts

    return emp2, t2


def _init_mp_df_eris(mp):
    """Compute 3-center electron repulsion integrals, i.e. (L|ov),
    where `L` denotes DF auxiliary basis functions and `o` and `v` occupied and virtual
    canonical crystalline orbitals. Note that `o` and `v` contain kpt indices `ko` and `kv`,
    and the third kpt index `kL` is determined by the conservation of momentum.

    Arguments:
        mp (KMP2) -- A KMP2 instance

    Returns:
        Lov (numpy.ndarray) -- 3-center DF ints, with shape (nkpts, nkpts, naux, nocc, nvir)
    """
    from pyscf.pbc.df import df
    from pyscf.ao2mo import _ao2mo
    from pyscf.pbc.lib.kpts_helper import gamma_point

    log = logger.Logger(mp.stdout, mp.verbose)

    if mp._scf.with_df._cderi is None:
        mp._scf.with_df.build()

    cell = mp._scf.cell
    if cell.dimension == 2:
        # 2D ERIs are not positive definite. The 3-index tensors are stored in
        # two part. One corresponds to the positive part and one corresponds
        # to the negative part. The negative part is not considered in the
        # DF-driven CCSD implementation.
        raise NotImplementedError

    nocc = mp.nocc
    nmo = mp.nmo
    nvir = nmo - nocc
    nao = cell.nao_nr()

    mo_coeff = _add_padding(mp, mp.mo_coeff, mp.mo_energy)[0]
    kpts = mp.kpts
    nkpts = len(kpts)
    if gamma_point(kpts):
        dtype = np.double
    else:
        dtype = np.complex128
    dtype = np.result_type(dtype, *mo_coeff)
    Lov = np.empty((nkpts, nkpts), dtype=object)

    cput0 = (logger.process_clock(), logger.perf_counter())

    bra_start = 0
    bra_end = nocc
    ket_start = nmo+nocc
    ket_end = ket_start + nvir
    with h5py.File(mp._scf.with_df._cderi, 'r') as f:
        kptij_lst = f['j3c-kptij'][:]
        tao = []
        ao_loc = None
        for ki, kpti in enumerate(kpts):
            for kj, kptj in enumerate(kpts):
                kpti_kptj = np.array((kpti, kptj))
                Lpq_ao = np.asarray(df._getitem(f, 'j3c', kpti_kptj, kptij_lst))

                mo = np.hstack((mo_coeff[ki], mo_coeff[kj]))
                mo = np.asarray(mo, dtype=dtype, order='F')
                if dtype == np.double:
                    out = _ao2mo.nr_e2(Lpq_ao, mo, (bra_start, bra_end, ket_start, ket_end), aosym='s2')
                else:
                    #Note: Lpq.shape[0] != naux if linear dependency is found in auxbasis
                    if Lpq_ao[0].size != nao**2:  # aosym = 's2'
                        Lpq_ao = lib.unpack_tril(Lpq_ao).astype(np.complex128)
                    out = _ao2mo.r_e2(Lpq_ao, mo, (bra_start, bra_end, ket_start, ket_end), tao, ao_loc)
                Lov[ki, kj] = out.reshape(-1, nocc, nvir)

    log.timer_debug1("transforming DF-MP2 integrals", *cput0)

    return Lov


def _padding_k_idx(nmo, nocc, kind="split"):
    """A convention used for padding vectors, matrices and tensors in case when occupation numbers depend on the
    k-point index.
    Args:
        nmo (Iterable): k-dependent orbital number;
        nocc (Iterable): k-dependent occupation numbers;
        kind (str): either "split" (occupied and virtual spaces are split) or "joint" (occupied and virtual spaces are
        the joint;

    Returns:
        Two lists corresponding to the occupied and virtual spaces for kind="split". Each list contains integer arrays
        with indexes pointing to actual non-zero entries in the padded vector/matrix/tensor. If kind="joint", a single
        list of arrays is returned corresponding to the entire MO space.
    """
    if kind not in ("split", "joint"):
        raise ValueError("The 'kind' argument must be one of 'split', 'joint'")

    if kind == "split":
        indexes_o = []
        indexes_v = []
    else:
        indexes = []

    nocc = np.array(nocc)
    nmo = np.array(nmo)
    nvirt = nmo - nocc
    dense_o = np.amax(nocc)
    dense_v = np.amax(nvirt)
    dense_nmo = dense_o + dense_v

    for k_o, k_nmo in zip(nocc, nmo):
        k_v = k_nmo - k_o
        if kind == "split":
            indexes_o.append(np.arange(k_o))
            indexes_v.append(np.arange(dense_v - k_v, dense_v))
        else:
            indexes.append(np.concatenate((
                np.arange(k_o),
                np.arange(dense_nmo - k_v, dense_nmo),
            )))

    if kind == "split":
        return indexes_o, indexes_v

    else:
        return indexes


def padding_k_idx(mp, kind="split"):
    """A convention used for padding vectors, matrices and tensors in case when occupation numbers depend on the
    k-point index.

    This implementation stores k-dependent Fock and other matrix in dense arrays with additional dimensions
    corresponding to k-point indexes. In case when the occupation numbers depend on the k-point index (i.e. a metal) or
    when some k-points have more Bloch basis functions than others the corresponding data structure has to be padded
    with entries that are not used (fictitious occupied and virtual degrees of freedom). Current convention stores these
    states at the Fermi level as shown in the following example.

    +----+--------+--------+--------+
    |    |  k=0   |  k=1   |  k=2   |
    |    +--------+--------+--------+
    |    | nocc=2 | nocc=3 | nocc=2 |
    |    | nvir=4 | nvir=3 | nvir=3 |
    +====+========+========+========+
    | v3 |  k0v3  |  k1v2  |  k2v2  |
    +----+--------+--------+--------+
    | v2 |  k0v2  |  k1v1  |  k2v1  |
    +----+--------+--------+--------+
    | v1 |  k0v1  |  k1v0  |  k2v0  |
    +----+--------+--------+--------+
    | v0 |  k0v0  |        |        |
    +====+========+========+========+
    |          Fermi level          |
    +====+========+========+========+
    | o2 |        |  k1o2  |        |
    +----+--------+--------+--------+
    | o1 |  k0o1  |  k1o1  |  k2o1  |
    +----+--------+--------+--------+
    | o0 |  k0o0  |  k1o0  |  k2o0  |
    +----+--------+--------+--------+

    In the above example, `get_nmo(mp, per_kpoint=True) == (6, 6, 5)`, `get_nocc(mp, per_kpoint) == (2, 3, 2)`. The
    resulting dense `get_nmo(mp) == 7` and `get_nocc(mp) == 3` correspond to padded dimensions. This function will
    return the following indexes corresponding to the filled entries of the above table:

    >>> padding_k_idx(mp, kind="split")
    ([(0, 1), (0, 1, 2), (0, 1)], [(0, 1, 2, 3), (1, 2, 3), (1, 2, 3)])

    >>> padding_k_idx(mp, kind="joint")
    [(0, 1, 3, 4, 5, 6), (0, 1, 2, 4, 5, 6), (0, 1, 4, 5, 6)]

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        kind (str): either "split" (occupied and virtual spaces are split) or "joint" (occupied and virtual spaces are
        the joint;

    Returns:
        Two lists corresponding to the occupied and virtual spaces for kind="split". Each list contains integer arrays
        with indexes pointing to actual non-zero entries in the padded vector/matrix/tensor. If kind="joint", a single
        list of arrays is returned corresponding to the entire MO space.
    """
    return _padding_k_idx(mp.get_nmo(per_kpoint=True), mp.get_nocc(per_kpoint=True), kind=kind)


def padded_mo_energy(mp, mo_energy):
    """
    Pads energies of active MOs.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        mo_energy (ndarray): original non-padded molecular energies;

    Returns:
        Padded molecular energies.
    """
    frozen_mask = get_frozen_mask(mp)
    padding_convention = padding_k_idx(mp, kind="joint")
    nkpts = mp.nkpts

    result = np.zeros((nkpts, mp.nmo), dtype=mo_energy[0].dtype)
    for k in range(nkpts):
        result[np.ix_([k], padding_convention[k])] = mo_energy[k][frozen_mask[k]]

    return result


def padded_mo_coeff(mp, mo_coeff):
    """
    Pads coefficients of active MOs.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        mo_coeff (ndarray): original non-padded molecular coefficients;

    Returns:
        Padded molecular coefficients.
    """
    frozen_mask = get_frozen_mask(mp)
    padding_convention = padding_k_idx(mp, kind="joint")
    nkpts = mp.nkpts

    result = np.zeros((nkpts, mo_coeff[0].shape[0], mp.nmo), dtype=mo_coeff[0].dtype)
    for k in range(nkpts):
        result[np.ix_([k], np.arange(result.shape[1]), padding_convention[k])] = mo_coeff[k][:, frozen_mask[k]]

    return result


def _frozen_sanity_check(frozen, mo_occ, kpt_idx):
    '''Performs a few sanity checks on the frozen array and mo_occ.

    Specific tests include checking for duplicates within the frozen array.

    Args:
        frozen (array_like of int): The orbital indices that will be frozen.
        mo_occ (:obj:`ndarray` of int): The occupuation number for each orbital
            resulting from a mean-field-like calculation.
        kpt_idx (int): The k-point that `mo_occ` and `frozen` belong to.

    '''
    frozen = np.array(frozen)
    nocc = np.count_nonzero(mo_occ > 0)

    assert nocc, 'No occupied orbitals?\n\nnocc = %s\nmo_occ = %s' % (nocc, mo_occ)
    all_frozen_unique = (len(frozen) - len(np.unique(frozen))) == 0
    if not all_frozen_unique:
        raise RuntimeError('Frozen orbital list contains duplicates!\n\nkpt_idx %s\n'
                           'frozen %s' % (kpt_idx, frozen))
    if len(frozen) > 0 and np.max(frozen) > len(mo_occ) - 1:
        raise RuntimeError('Freezing orbital not in MO list!\n\nkpt_idx %s\n'
                           'frozen %s\nmax orbital idx %s' % (kpt_idx, frozen, len(mo_occ) - 1))


def get_nocc(mp, per_kpoint=False):
    '''Number of occupied orbitals for k-point calculations.

    Number of occupied orbitals for use in a calculation with k-points, taking into
    account frozen orbitals.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        per_kpoint (bool, optional): True returns the number of occupied
            orbitals at each k-point.  False gives the max of this list.

    Returns:
        nocc (int, list of int): Number of occupied orbitals. For return type, see description of arg
            `per_kpoint`.

    '''
    for i, moocc in enumerate(mp.mo_occ):
        if np.any(moocc % 1 != 0):
            raise RuntimeError("Fractional occupation numbers encountered @ kp={:d}: {}. This may have been caused by "
                               "smearing of occupation numbers in the mean-field calculation. If so, consider "
                               "executing mf.smearing_method = False; mf.mo_occ = mf.get_occ() prior to calling "
                               "this".format(i, moocc))
    if mp._nocc is not None:
        return mp._nocc
    elif mp.frozen is None:
        nocc = [np.count_nonzero(mp.mo_occ[ikpt]) for ikpt in range(mp.nkpts)]
    elif isinstance(mp.frozen, (int, np.integer)):
        nocc = [(np.count_nonzero(mp.mo_occ[ikpt]) - mp.frozen) for ikpt in range(mp.nkpts)]
    elif isinstance(mp.frozen[0], (int, np.integer)):
        [_frozen_sanity_check(mp.frozen, mp.mo_occ[ikpt], ikpt) for ikpt in range(mp.nkpts)]
        nocc = []
        for ikpt in range(mp.nkpts):
            max_occ_idx = np.max(np.where(mp.mo_occ[ikpt] > 0))
            frozen_nocc = np.sum(np.array(mp.frozen) <= max_occ_idx)
            nocc.append(np.count_nonzero(mp.mo_occ[ikpt]) - frozen_nocc)
    elif isinstance(mp.frozen[0], (list, np.ndarray)):
        nkpts = len(mp.frozen)
        if nkpts != mp.nkpts:
            raise RuntimeError('Frozen list has a different number of k-points (length) than passed in mean-field/'
                               'correlated calculation.  \n\nCalculation nkpts = %d, frozen list = %s '
                               '(length = %d)' % (mp.nkpts, mp.frozen, nkpts))
        [_frozen_sanity_check(frozen, mo_occ, ikpt) for ikpt, frozen, mo_occ in zip(range(nkpts), mp.frozen, mp.mo_occ)]

        nocc = []
        for ikpt, frozen in enumerate(mp.frozen):
            max_occ_idx = np.max(np.where(mp.mo_occ[ikpt] > 0))
            frozen_nocc = np.sum(np.array(frozen) <= max_occ_idx)
            nocc.append(np.count_nonzero(mp.mo_occ[ikpt]) - frozen_nocc)
    else:
        raise NotImplementedError

    assert any(np.array(nocc) > 0), ('Must have occupied orbitals! \n\nnocc %s\nfrozen %s\nmo_occ %s' %
           (nocc, mp.frozen, mp.mo_occ))

    if not per_kpoint:
        nocc = np.amax(nocc)

    return nocc


def get_nmo(mp, per_kpoint=False):
    '''Number of orbitals for k-point calculations.

    Number of orbitals for use in a calculation with k-points, taking into account
    frozen orbitals.

    Note:
        If `per_kpoint` is False, then the number of orbitals here is equal to max(nocc) + max(nvir),
        where each max is done over all k-points.  Otherwise the number of orbitals is returned
        as a list of number of orbitals at each k-point.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        per_kpoint (bool, optional): True returns the number of orbitals at each k-point.
            For a description of False, see Note.

    Returns:
        nmo (int, list of int): Number of orbitals. For return type, see description of arg
            `per_kpoint`.

    '''
    if mp._nmo is not None:
        return mp._nmo

    if mp.frozen is None:
        nmo = [len(mp.mo_occ[ikpt]) for ikpt in range(mp.nkpts)]
    elif isinstance(mp.frozen, (int, np.integer)):
        nmo = [len(mp.mo_occ[ikpt]) - mp.frozen for ikpt in range(mp.nkpts)]
    elif isinstance(mp.frozen[0], (int, np.integer)):
        [_frozen_sanity_check(mp.frozen, mp.mo_occ[ikpt], ikpt) for ikpt in range(mp.nkpts)]
        nmo = [len(mp.mo_occ[ikpt]) - len(mp.frozen) for ikpt in range(mp.nkpts)]
    elif isinstance(mp.frozen, (list, np.ndarray)):
        nkpts = len(mp.frozen)
        if nkpts != mp.nkpts:
            raise RuntimeError('Frozen list has a different number of k-points (length) than passed in mean-field/'
                               'correlated calculation.  \n\nCalculation nkpts = %d, frozen list = %s '
                               '(length = %d)' % (mp.nkpts, mp.frozen, nkpts))
        [_frozen_sanity_check(fro, mo_occ, ikpt) for ikpt, fro, mo_occ in zip(range(nkpts), mp.frozen, mp.mo_occ)]

        nmo = [len(mp.mo_occ[ikpt]) - len(mp.frozen[ikpt]) for ikpt in range(nkpts)]
    else:
        raise NotImplementedError

    assert all(np.array(nmo) > 0), ('Must have a positive number of orbitals!\n\nnmo %s\nfrozen %s\nmo_occ %s' %
           (nmo, mp.frozen, mp.mo_occ))

    if not per_kpoint:
        # Depending on whether there are more occupied bands, we want to make sure that
        # nmo has enough room for max(nocc) + max(nvir) number of orbitals for occupied
        # and virtual space
        nocc = mp.get_nocc(per_kpoint=True)
        nmo = np.max(nocc) + np.max(np.array(nmo) - np.array(nocc))

    return nmo


def get_frozen_mask(mp):
    '''Boolean mask for orbitals in k-point post-HF method.

    Creates a boolean mask to remove frozen orbitals and keep other orbitals for post-HF
    calculations.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.

    Returns:
        moidx (list of :obj:`ndarray` of `np.bool`): Boolean mask of orbitals to include.

    '''
    moidx = [np.ones(x.size, dtype=np.bool) for x in mp.mo_occ]
    if mp.frozen is None:
        pass
    elif isinstance(mp.frozen, (int, np.integer)):
        for idx in moidx:
            idx[:mp.frozen] = False
    elif isinstance(mp.frozen[0], (int, np.integer)):
        frozen = list(mp.frozen)
        for idx in moidx:
            idx[frozen] = False
    elif isinstance(mp.frozen[0], (list, np.ndarray)):
        nkpts = len(mp.frozen)
        if nkpts != mp.nkpts:
            raise RuntimeError('Frozen list has a different number of k-points (length) than passed in mean-field/'
                               'correlated calculation.  \n\nCalculation nkpts = %d, frozen list = %s '
                               '(length = %d)' % (mp.nkpts, mp.frozen, nkpts))
        [_frozen_sanity_check(fro, mo_occ, ikpt) for ikpt, fro, mo_occ in zip(range(nkpts), mp.frozen, mp.mo_occ)]
        for ikpt, kpt_occ in enumerate(moidx):
            kpt_occ[mp.frozen[ikpt]] = False
    else:
        raise NotImplementedError

    return moidx


def _add_padding(mp, mo_coeff, mo_energy):
    nmo = mp.nmo

    # Check if these are padded mo coefficients and energies
    if not np.all([x.shape[1] == nmo for x in mo_coeff]):
        mo_coeff = padded_mo_coeff(mp, mo_coeff)

    if not np.all([x.shape[0] == nmo for x in mo_energy]):
        mo_energy = padded_mo_energy(mp, mo_energy)
    return mo_coeff, mo_energy


def make_rdm1(mp, t2=None, kind="compact"):
    r"""
    Spin-traced one-particle density matrix in the MO basis representation.
    The occupied-virtual orbital response is not included.

    dm1[p,q] = <q_alpha^\dagger p_alpha> + <q_beta^\dagger p_beta>

    The convention of 1-pdm is based on McWeeney's book, Eq (5.4.20).
    The contraction between 1-particle Hamiltonian and rdm1 is
    E = einsum('pq,qp', h1, rdm1)

    Args:
        mp (KMP2): a KMP2 kernel object;
        t2 (ndarray): a t2 MP2 tensor;
        kind (str): either 'compact' or 'padded' - defines behavior for k-dependent MO basis sizes;

    Returns:
        A k-dependent single-particle density matrix.
    """
    if kind not in ("compact", "padded"):
        raise ValueError("The 'kind' argument should be either 'compact' or 'padded'")
    d_imds = _gamma1_intermediates(mp, t2=t2)
    result = []
    padding_idxs = padding_k_idx(mp, kind="joint")
    for (oo, vv), idxs in zip(zip(*d_imds), padding_idxs):
        oo += np.eye(*oo.shape)
        d = block_diag(oo, vv)
        d += d.conj().T
        if kind == "padded":
            result.append(d)
        else:
            result.append(d[np.ix_(idxs, idxs)])
    return result


def make_rdm2(mp, t2=None, kind="compact"):
    r'''
    Spin-traced two-particle density matrix in MO basis

    .. math::

        dm2[p,q,r,s] = \sum_{\sigma,\tau} <p_\sigma^\dagger r_\tau^\dagger s_\tau q_\sigma>

    Note the contraction between ERIs (in Chemist's notation) and rdm2 is
    E = einsum('pqrs,pqrs', eri, rdm2)
    '''
    if kind not in ("compact", "padded"):
        raise ValueError("The 'kind' argument should be either 'compact' or 'padded'")
    if t2 is None: t2 = mp.t2
    dm1 = mp.make_rdm1(t2, "padded")
    nmo = mp.nmo
    nocc = mp.nocc
    nkpts = mp.nkpts
    dtype = t2.dtype

    dm2 = np.zeros((nkpts,nkpts,nkpts,nmo,nmo,nmo,nmo),dtype=dtype)
    for ki in range(nkpts):
        for kj in range(nkpts):
            for ka in range(nkpts):
                kb = mp.khelper.kconserv[ki, ka, kj]
                dovov = t2[ki, kj, ka].transpose(0,2,1,3) * 2 - t2[kj, ki, ka].transpose(1,2,0,3)
                dovov *= 2
                dm2[ki,ka,kj,:nocc,nocc:,:nocc,nocc:] = dovov
                dm2[ka,ki,kb,nocc:,:nocc,nocc:,:nocc] = dovov.transpose(1,0,3,2).conj()

    occidx = padding_k_idx(mp, kind="split")[0]
    for ki in range(nkpts):
        for i in occidx[ki]:
            dm1[ki][i,i] -= 2

    for ki in range(nkpts):
        for kp in range(nkpts):
            for i in occidx[ki]:
                dm2[ki,ki,kp,i,i,:,:] += dm1[kp].T * 2
                dm2[kp,kp,ki,:,:,i,i] += dm1[kp].T * 2
                dm2[kp,ki,ki,:,i,i,:] -= dm1[kp].T
                dm2[ki,kp,kp,i,:,:,i] -= dm1[kp]

    for ki in range(nkpts):
        for kj in range(nkpts):
            for i in occidx[ki]:
                for j in occidx[kj]:
                    dm2[ki,ki,kj,i,i,j,j] += 4
                    dm2[ki,kj,kj,i,j,j,i] -= 2

    if kind == "padded":
        return dm2
    else:
        idx = padding_k_idx(mp, kind="joint")
        result = []
        for kp in range(nkpts):
            for kq in range(nkpts):
                for kr in range(nkpts):
                    ks = mp.khelper.kconserv[kp, kq, kr]
                    result.append(dm2[kp,kq,kr][np.ix_(idx[kp],idx[kq],idx[kr],idx[ks])])
        return result


def _gamma1_intermediates(mp, t2=None):
    # Memory optimization should be here
    if t2 is None:
        t2 = mp.t2
    if t2 is None:
        raise NotImplementedError("Run kmp2.kernel with `with_t2=True`")
    nmo = mp.nmo
    nocc = mp.nocc
    nvir = nmo - nocc
    nkpts = mp.nkpts
    dtype = t2.dtype

    dm1occ = np.zeros((nkpts, nocc, nocc), dtype=dtype)
    dm1vir = np.zeros((nkpts, nvir, nvir), dtype=dtype)

    for ki in range(nkpts):
        for kj in range(nkpts):
            for ka in range(nkpts):
                kb = mp.khelper.kconserv[ki, ka, kj]

                dm1vir[kb] += einsum('ijax,ijay->yx', t2[ki][kj][ka].conj(), t2[ki][kj][ka]) * 2 -\
                              einsum('ijax,ijya->yx', t2[ki][kj][ka].conj(), t2[ki][kj][kb])
                dm1occ[kj] += einsum('ixab,iyab->xy', t2[ki][kj][ka].conj(), t2[ki][kj][ka]) * 2 -\
                              einsum('ixab,iyba->xy', t2[ki][kj][ka].conj(), t2[ki][kj][kb])
    return -dm1occ, dm1vir


class KMP2(mp2.MP2):
    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):

        if mo_coeff is None: mo_coeff = mf.mo_coeff
        if mo_occ is None: mo_occ = mf.mo_occ

        self.mol = mf.mol
        self._scf = mf
        self.verbose = self.mol.verbose
        self.stdout = self.mol.stdout
        self.max_memory = mf.max_memory

        self.frozen = frozen
        if isinstance(self._scf.with_df, df.GDF):
            self.with_df_ints = True
        else:
            self.with_df_ints = False

##################################################
# don't modify the following attributes, they are not input options
        self.kpts = mf.kpts
        self.mo_energy = mf.mo_energy
        self.nkpts = len(self.kpts)
        self.khelper = kpts_helper.KptsHelper(mf.cell, mf.kpts)
        self.mo_coeff = mo_coeff
        self.mo_occ = mo_occ
        self._nocc = None
        self._nmo = None
        self.e_corr = None
        self.e_hf = None
        self.t2 = None
        self._keys = set(self.__dict__.keys())

    get_nocc = get_nocc
    get_nmo = get_nmo
    get_frozen_mask = get_frozen_mask
    make_rdm1 = make_rdm1
    make_rdm2 = make_rdm2

    def dump_flags(self):
        logger.info(self, "")
        logger.info(self, "******** %s ********", self.__class__)
        logger.info(self, "nkpts = %d", self.nkpts)
        logger.info(self, "nocc = %d", self.nocc)
        logger.info(self, "nmo = %d", self.nmo)
        logger.info(self, "with_df_ints = %s", self.with_df_ints)

        if self.frozen is not None:
            logger.info(self, "frozen orbitals = %s", self.frozen)
        logger.info(
            self,
            "max_memory %d MB (current use %d MB)",
            self.max_memory,
            lib.current_memory()[0],
        )
        return self

    def kernel(self, mo_energy=None, mo_coeff=None, with_t2=WITH_T2):
        if mo_energy is None:
            mo_energy = self.mo_energy
        if mo_coeff is None:
            mo_coeff = self.mo_coeff
        if mo_energy is None or mo_coeff is None:
            log = logger.Logger(self.stdout, self.verbose)
            log.warn('mo_coeff, mo_energy are not given.\n'
                     'You may need to call mf.kernel() to generate them.')
            raise RuntimeError

        mo_coeff, mo_energy = _add_padding(self, mo_coeff, mo_energy)

        # TODO: compute e_hf for non-canonical SCF
        self.e_hf = self._scf.e_tot

        self.e_corr, self.t2 = \
                kernel(self, mo_energy, mo_coeff, verbose=self.verbose, with_t2=with_t2)
        logger.log(self, 'KMP2 energy = %.15g', self.e_corr)
        return self.e_corr, self.t2
KRMP2 = KMP2


from pyscf.pbc import scf
scf.khf.KRHF.MP2 = lib.class_as_method(KRMP2)
scf.kghf.KGHF.MP2 = None
scf.krohf.KROHF.MP2 = None


if __name__ == '__main__':
    from pyscf.pbc import gto, scf, mp

    cell = gto.Cell()
    cell.atom='''
    C 0.000000000000   0.000000000000   0.000000000000
    C 1.685068664391   1.685068664391   1.685068664391
    '''
    cell.basis = 'gth-szv'
    cell.pseudo = 'gth-pade'
    cell.a = '''
    0.000000000, 3.370137329, 3.370137329
    3.370137329, 0.000000000, 3.370137329
    3.370137329, 3.370137329, 0.000000000'''
    cell.unit = 'B'
    cell.verbose = 5
    cell.build()

    # Running HF and MP2 with 1x1x2 Monkhorst-Pack k-point mesh
    kmf = scf.KRHF(cell, kpts=cell.make_kpts([1,1,2]), exxdiv=None)
    ehf = kmf.kernel()

    mymp = mp.KMP2(kmf)
    emp2, t2 = mymp.kernel()
    print(emp2 - -0.204721432828996)



# pyscf/pbc/mp/kump2.py

#!/usr/bin/env python
# Copyright 2014-2021 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
# Author: James McClain <jdmcclain47@gmail.com>
#


'''
kpoint-adapted unrestricted MP2
t2[i,j,a,b] = <ij|ab> / D_ij^ab

t2 and eris are never stored in full, only a partial
eri of size (nkpts,nocc,nocc,nvir,nvir)
'''


import numpy as np

from pyscf import lib
from pyscf.lib import logger
from pyscf.pbc.mp import kmp2
from pyscf.pbc.lib import kpts_helper
from pyscf.pbc.mp.kmp2 import _frozen_sanity_check
from pyscf.lib.parameters import LARGE_DENOM

def kernel(mp, mo_energy, mo_coeff, verbose=logger.NOTE):
    raise NotImplementedError

def padding_k_idx(mp, kind="split"):
    """For a description, see `padding_k_idx` in kmp2.py.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        kind (str): either "split" (occupied and virtual spaces are split) or "joint" (occupied and virtual spaces are
        the joint;

    Returns:
        Two lists corresponding to the occupied and virtual spaces for kind="split". Each list contains integer arrays
        with indexes pointing to actual non-zero entries in the padded vector/matrix/tensor. If kind="joint", a single
        list of arrays is returned corresponding to the entire MO space.
    """
    if kind not in ("split", "joint"):
        raise ValueError("The 'kind' argument must be one of 'split', 'joint'")

    if kind == "split":
        indexes_oa = []
        indexes_va = []
        indexes_ob = []
        indexes_vb = []
    else:
        indexesa = []
        indexesb = []

    dense_oa, dense_ob = mp.nocc
    dense_nmoa, dense_nmob = mp.nmo
    dense_va = dense_nmoa - dense_oa
    dense_vb = dense_nmob - dense_ob

    nocca_per_kpt, noccb_per_kpt = np.asarray(get_nocc(mp, per_kpoint=True))
    nmoa_per_kpt, nmob_per_kpt = np.asarray(get_nmo(mp, per_kpoint=True))

    # alpha spin
    for k_oa, k_nmoa in zip(nocca_per_kpt, nmoa_per_kpt):
        k_va = k_nmoa - k_oa
        if kind == "split":
            indexes_oa.append(np.arange(k_oa))
            indexes_va.append(np.arange(dense_va - k_va, dense_va))
        else:
            indexesa.append(np.concatenate((
                np.arange(k_oa),
                np.arange(dense_nmoa - k_va, dense_nmoa),
            )))

    # beta spin
    for k_ob, k_nmob in zip(noccb_per_kpt, nmob_per_kpt):
        k_vb = k_nmob - k_ob
        if kind == "split":
            indexes_ob.append(np.arange(k_ob))
            indexes_vb.append(np.arange(dense_vb - k_vb, dense_vb))
        else:
            indexesb.append(np.concatenate((
                np.arange(k_ob),
                np.arange(dense_nmob - k_vb, dense_nmob),
            )))

    if kind == "split":
        return [indexes_oa, indexes_va], [indexes_ob, indexes_vb]
    else:
        return indexesa, indexesb


def padded_mo_energy(mp, mo_energy):
    """
    Pads energies of active MOs.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        mo_energy (ndarray): original non-padded molecular energies;

    Returns:
        Padded molecular energies.
    """
    frozen_mask = get_frozen_mask(mp)
    padding_convention = padding_k_idx(mp, kind="joint")
    nkpts = mp.nkpts

    result = (np.zeros((nkpts, mp.nmo), dtype=mo_energy[0][0].dtype),
              np.zeros((nkpts, mp.nmo), dtype=mo_energy[0][0].dtype))
    for spin in [0, 1]:
        for k in range(nkpts):
            result[spin][np.ix_([k], padding_convention[k])] = mo_energy[spin][k][frozen_mask[k]]

    return result


def padded_mo_coeff(mp, mo_coeff):
    """
    Pads coefficients of active MOs.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        mo_coeff (ndarray): original non-padded molecular coefficients;

    Returns:
        Padded molecular coefficients.
    """
    frozen_mask = get_frozen_mask(mp)
    padding_convention = padding_k_idx(mp, kind="joint")
    nkpts = mp.nkpts

    result = (np.zeros((nkpts, mo_coeff[0][0].shape[0], mp.nmo[0]), dtype=mo_coeff[0][0].dtype),
              np.zeros((nkpts, mo_coeff[1][0].shape[0], mp.nmo[1]), dtype=mo_coeff[0][0].dtype))
    for spin in [0, 1]:
        for k in range(nkpts):
            result[spin][np.ix_([k], np.arange(result[spin].shape[1]), padding_convention[spin][k])] = \
                mo_coeff[spin][k][:, frozen_mask[spin][k]]

    return result


def _is_arraylike(x):
    return isinstance(x, (tuple, list, np.ndarray))


def get_nocc(mp, per_kpoint=False):
    '''Number of occupied orbitals for k-point calculations.

    Number of occupied orbitals for use in a calculation with k-points, taking into
    account frozen orbitals.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        per_kpoint (bool, optional): True returns the number of occupied
            orbitals at each k-point.  False gives the max of this list.

    Returns:
        nocc (int, list of int): Number of occupied orbitals. For return type, see description of arg
            `per_kpoint`.

    Notes:
        For specifying frozen orbitals inside mp, the following options are accepted:

            +=========================+========================================+===============================+
            | Argument (Example)      | Argument Meaning                       | Example Meaning               |
            +=========================+========================================+===============================+
            | int (1)                 | Freeze the same number of orbitals     | Freeze one (lowest) orbital   |
            |                         | regardless of spin and/or kpt          | for all kpts and spin cases   |
            +-------------------------+----------------------------------------+-------------------------------+
            | 2-tuple of list of int  | inner list: List of orbitals indices   | Freeze the orbitals [0,4] for |
            | ([0, 4], [0, 5, 6])     |   to freeze at all kpts                | spin0, and orbitals [0,5,6]   |
            |                         | outer list: Spin index                 | for spin1 at all kpts         |
            +-------------------------+----------------------------------------+-------------------------------+
            | list(2) of list of list | inner list: list of orbital indices to | Freeze orbital 0 for spin0 at |
            | ([[0,],[]],             |   freeze at each kpt for given spin    | kpt0, and freeze orbital 0,1  |
            |  [[0,1],[4]])           | outer list: spin index                 | for spin1 at kpt0 and orbital |
            |                         |                                        | 4 at kpt1                     |
            +-------------------------+----------------------------------------+-------------------------------+

    '''
    for spin in [0,1]:
        for i, moocc in enumerate(mp.mo_occ[spin]):
            if np.any(moocc % 1 != 0):
                raise RuntimeError(
                    "Fractional occupation numbers encountered @ kp={:d}: {}.  "
                    "This may have been caused by smearing of occupation numbers "
                    "in the mean-field calculation. If so, consider executing "
                    "mf.smearing_method = False; mf.mo_occ = mf.get_occ() prior "
                    "to calling this".format(i, moocc))
    if mp._nocc is not None:
        return mp._nocc

    elif mp.frozen is None:
        nocc = [[np.count_nonzero(mp.mo_occ[0][k] > 0) for k in range(mp.nkpts)],
                [np.count_nonzero(mp.mo_occ[1][k] > 0) for k in range(mp.nkpts)]]

    elif isinstance(mp.frozen, (int, np.integer)):
        nocc = [0]*2
        for spin in [0,1]:
            nocc[spin] = [(np.count_nonzero(mp.mo_occ[spin][k] > 0) - mp.frozen) for k in range(mp.nkpts)]

    elif (_is_arraylike(mp.frozen[0]) and
          isinstance(mp.frozen[0][0], (int, np.integer))):  # case example: ([0, 4], [0, 5, 6])
        nocc = [0]*2
        assert(len(mp.frozen) == 2)
        for spin in [0,1]:
            [_frozen_sanity_check(mp.frozen[spin], mp.mo_occ[spin][ikpt], ikpt) for ikpt in range(mp.nkpts)]
            nocc_spin = []
            for ikpt in range(mp.nkpts):
                max_occ_idx = np.max(np.where(mp.mo_occ[spin][ikpt] > 0))
                frozen_nocc = np.sum(np.array(mp.frozen[spin]) <= max_occ_idx)
                nocc_spin.append(np.count_nonzero(mp.mo_occ[spin][ikpt]) - frozen_nocc)
            nocc[spin] = nocc_spin

    elif (_is_arraylike(mp.frozen[0]) and
          isinstance(mp.frozen[0][0], (list, np.ndarray))):  # case example: ([[0,],[]], [[0,1],[4]])
        assert(len(mp.frozen) == 2)
        for spin in [0,1]:
            nkpts = len(mp.frozen[spin])
            if nkpts != mp.nkpts:
                raise RuntimeError('Frozen list has a different number of k-points (length) than passed in'
                                   'mean-field/correlated calculation.  \n\nCalculation nkpts = %d, frozen'
                                   'list = %s (length = %d)' % (mp.nkpts, mp.frozen, nkpts))
        nocc = [0]*2
        for spin in [0,1]:
            [_frozen_sanity_check(frozen, mo_occ, ikpt)
             for ikpt, frozen, mo_occ in zip(range(nkpts), mp.frozen[spin], mp.mo_occ[spin])]
            nocc_spin = []
            for ikpt, frozen in enumerate(mp.frozen[spin]):
                max_occ_idx = np.max(np.where(mp.mo_occ[spin][ikpt] > 0))
                frozen_nocc = np.sum(np.array(frozen) <= max_occ_idx)
                nocc_spin.append(np.count_nonzero(mp.mo_occ[spin][ikpt]) - frozen_nocc)
            nocc[spin] = nocc_spin
    else:
        raise NotImplementedError('No known conversion for frozen %s' % mp.frozen)

    for spin in [0,1]:
        assert any(np.array(nocc[spin]) > 0), (
            'Must have occupied orbitals (spin=%d)! \n\nnocc %s\nfrozen %s\nmo_occ %s' %
            (spin, nocc, mp.frozen, mp.mo_occ))

    nocca, noccb = nocc
    if not per_kpoint:
        nocca = np.amax(nocca)
        noccb = np.amax(noccb)

    return nocca, noccb


def get_nmo(mp, per_kpoint=False):
    '''Number of orbitals for k-point calculations.

    Number of orbitals for use in a calculation with k-points, taking into account
    frozen orbitals.

    Note:
        If `per_kpoint` is False, then the number of orbitals here is equal to max(nocc) + max(nvir),
        where each max is done over all k-points.  Otherwise the number of orbitals is returned
        as a list of number of orbitals at each k-point.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.
        per_kpoint (bool, optional): True returns the number of orbitals at each k-point.
            For a description of False, see Note.

    Returns:
        nmo (int, list of int): Number of orbitals. For return type, see description of arg
            `per_kpoint`.

    '''
    if mp._nmo is not None:
        return mp._nmo

    nmo = [0, 0]
    if isinstance(mp.frozen, (int, np.integer)):
        for spin in [0,1]:
            nmo[spin] = [len(mp.mo_occ[spin][k]) - mp.frozen for k in range(mp.nkpts)]

    elif mp.frozen is None:
        nmo = [[len(mp.mo_occ[0][k]) for k in range(mp.nkpts)],
               [len(mp.mo_occ[1][k]) for k in range(mp.nkpts)]]

    elif (_is_arraylike(mp.frozen[0]) and
          isinstance(mp.frozen[0][0], (int, np.integer))):  # case example: ([0, 4], [0, 5, 6])
        assert(len(mp.frozen) == 2)
        for spin in [0,1]:
            [_frozen_sanity_check(mp.frozen[spin], mp.mo_occ[spin][ikpt], ikpt) for ikpt in range(mp.nkpts)]
            nmo[spin] = [len(mp.mo_occ[spin][ikpt]) - len(mp.frozen[spin]) for ikpt in range(mp.nkpts)]

    elif (_is_arraylike(mp.frozen[0]) and
          isinstance(mp.frozen[0][0], (list, np.ndarray))):  # case example: ([[0,],[]], [[0,1],[4]])
        assert(len(mp.frozen) == 2)
        for spin in [0,1]:
            nkpts = len(mp.frozen[spin])
            if nkpts != mp.nkpts:
                raise RuntimeError('Frozen list has a different number of k-points (length) than passed in'
                                   'mean-field/correlated calculation.  \n\nCalculation nkpts = %d, frozen'
                                   'list = %s (length = %d)' % (mp.nkpts, mp.frozen, nkpts))
        for spin in [0,1]:
            [_frozen_sanity_check(mp.frozen[spin][ikpt], mp.mo_occ[spin][ikpt], ikpt) for ikpt in range(mp.nkpts)]
            nmo[spin] = [len(mp.mo_occ[spin][ikpt]) - len(mp.frozen[spin][ikpt]) for ikpt in range(nkpts)]

    else:
        raise NotImplementedError('No known conversion for frozen %s' % mp.frozen)

    for spin in [0,1]:
        assert all(np.array(nmo[spin]) > 0), (
            'Must have a positive number of orbitals! (spin=%d)'
            '\n\nnmo %s\nfrozen %s\nmo_occ %s' % (spin, nmo, mp.frozen, mp.mo_occ))

    nmoa, nmob = nmo
    if not per_kpoint:
        # Depending on whether there are more occupied bands, we want to make sure    that
        # nmo has enough room for max(nocc) + max(nvir) number of orbitals for        occupied
        # and virtual space
        nocca, noccb = mp.get_nocc(per_kpoint=True)
        nmoa = np.amax(nocca) + np.max(np.array(nmoa) - np.array(nocca))
        nmob = np.amax(noccb) + np.max(np.array(nmob) - np.array(noccb))

    return nmoa, nmob


def get_frozen_mask(mp):
    '''Boolean mask for orbitals in k-point post-HF method.

    Creates a boolean mask to remove frozen orbitals and keep other orbitals for post-HF
    calculations.

    Args:
        mp (:class:`MP2`): An instantiation of an SCF or post-Hartree-Fock object.

    Returns:
        moidx (list of :obj:`ndarray` of `np.bool`): Boolean mask of orbitals to include.

    '''
    moidx = [[np.ones(x.size, dtype=np.bool) for x in mp.mo_occ[s]] for s in [0,1]]

    if mp.frozen is None:
        pass

    elif isinstance(mp.frozen, (int, np.integer)):
        for spin in [0,1]:
            for idx in moidx[spin]:
                idx[:mp.frozen] = False

    elif (_is_arraylike(mp.frozen[0]) and
          isinstance(mp.frozen[0][0], (int, np.integer))):  # case example: ([0, 4], [0, 5, 6])
        assert(len(mp.frozen) == 2)
        for spin in [0,1]:
            [_frozen_sanity_check(mp.frozen[spin], mp.mo_occ[spin][ikpt], ikpt) for ikpt in range(mp.nkpts)]
            for ikpt, kpt_occ in enumerate(moidx[spin]):
                kpt_occ[mp.frozen[spin]] = False

    elif (_is_arraylike(mp.frozen[0]) and
          isinstance(mp.frozen[0][0], (list, np.ndarray))):  # case example: ([[0,],[]], [[0,1],[4]])
        assert(len(mp.frozen) == 2)
        for spin in [0,1]:
            nkpts = len(mp.frozen[spin])
            if nkpts != mp.nkpts:
                raise RuntimeError('Frozen list has a different number of k-points (length) than passed in'
                                   'mean-field/correlated calculation.  \n\nCalculation nkpts = %d, frozen'
                                   'list = %s (length = %d)' % (mp.nkpts, mp.frozen, nkpts))
        for spin in [0,1]:
            [_frozen_sanity_check(mp.frozen[spin][ikpt], mp.mo_occ[spin][ikpt], ikpt) for ikpt in range(mp.nkpts)]
            for ikpt, kpt_occ in enumerate(moidx[spin]):
                kpt_occ[mp.frozen[spin][ikpt]] = False
    else:
        raise NotImplementedError('No known conversion for frozen %s' % mp.frozen)

    return moidx


def _add_padding(mp, mo_coeff, mo_energy):
    raise NotImplementedError("Implementation needs to be checked first")
    nmo = mp.nmo

    # Check if these are padded mo coefficients and energies
    if not np.all([x.shape[0] == nmo for x in mo_coeff]):
        mo_coeff = padded_mo_coeff(mp, mo_coeff)

    if not np.all([x.shape[0] == nmo for x in mo_energy]):
        mo_energy = padded_mo_energy(mp, mo_energy)
    return mo_coeff, mo_energy


class KUMP2(kmp2.KMP2):
    get_nocc = get_nocc
    get_nmo = get_nmo
    get_frozen_mask = get_frozen_mask

    def kernel(self, mo_energy=None, mo_coeff=None):
        raise NotImplementedError
        if mo_energy is None:
            mo_energy = self.mo_energy
        if mo_coeff is None:
            mo_coeff = self.mo_coeff
        if mo_energy is None or mo_coeff is None:
            log = logger.Logger(self.stdout, self.verbose)
            log.warn('mo_coeff, mo_energy are not given.\n'
                     'You may need to call mf.kernel() to generate them.')
            raise RuntimeError

        mo_coeff, mo_energy = _add_padding(self, mo_coeff, mo_energy)

        self.e_corr, self.t2 = \
                kernel(self, mo_energy, mo_coeff, verbose=self.verbose)
        logger.log(self, 'KMP2 energy = %.15g', self.e_corr)
        return self.e_corr, self.t2


from pyscf.pbc import scf
scf.kuhf.KUHF.MP2 = lib.class_as_method(KUMP2)


# pyscf/pbc/mp/mp2.py

#!/usr/bin/env python
# Copyright 2014-2018 The PySCF Developers. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pyscf import lib
from pyscf.mp import mp2
from pyscf.mp import ump2
from pyscf.mp import gmp2

class RMP2(mp2.RMP2):
    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
        if abs(mf.kpt).max() > 1e-9:
            raise NotImplementedError
        from pyscf.pbc.df.df_ao2mo import warn_pbc2d_eri
        warn_pbc2d_eri(mf)
        mp2.RMP2.__init__(self, mf, frozen, mo_coeff, mo_occ)

    def ao2mo(self, mo_coeff=None):
        ao2mofn = _gen_ao2mofn(self._scf)
        eris = mp2._make_eris(self, mo_coeff, ao2mofn, self.verbose)
        return eris

class UMP2(ump2.UMP2):
    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
        if abs(mf.kpt).max() > 1e-9:
            raise NotImplementedError
        from pyscf.pbc.df.df_ao2mo import warn_pbc2d_eri
        warn_pbc2d_eri(mf)
        ump2.UMP2.__init__(self, mf, frozen, mo_coeff, mo_occ)

    def ao2mo(self, mo_coeff=None):
        ao2mofn = _gen_ao2mofn(self._scf)
        eris = ump2._make_eris(self, mo_coeff, ao2mofn, self.verbose)
        return eris

class GMP2(gmp2.GMP2):
    def __init__(self, mf, frozen=None, mo_coeff=None, mo_occ=None):
        from pyscf.pbc.df.df_ao2mo import warn_pbc2d_eri
        warn_pbc2d_eri(mf)
        gmp2.GMP2.__init__(self, mf, frozen, mo_coeff, mo_occ)

    def ao2mo(self, mo_coeff=None):
        with_df = self._scf.with_df
        kpt = self._scf.kpt
        def ao2mofn(mos):
            orbo, orbv, orbo, orbv = mos
            nocc = orbo.shape[1]
            nvir = orbv.shape[1]
            nao = orbo.shape[0]
            orboa = orbo[:nao//2]
            orbob = orbo[nao//2:]
            orbva = orbv[:nao//2]
            orbvb = orbv[nao//2:]
            orbspin = getattr(orbo, 'orbspin', None)
            if orbspin is None:
                eri  = with_df.ao2mo((orboa,orbva,orboa,orbva), kpt)
                eri += with_df.ao2mo((orbob,orbvb,orbob,orbvb), kpt)
                eri1 = with_df.ao2mo((orboa,orbva,orbob,orbvb), kpt)
                eri += eri1
                eri += eri1.transpose(2,3,0,1)
                eri = eri.reshape(nocc,nvir,nocc,nvir)
            else:
                co = orboa + orbob
                cv = orbva + orbvb
                eri = with_df.ao2mo((co,cv,co,cv), kpt).reshape(nocc,nvir,nocc,nvir)
                sym_forbid = (orbspin[:nocc,None] != orbspin[nocc:])
                eri[sym_forbid,:,:] = 0
                eri[:,:,sym_forbid] = 0
            return eri
        return gmp2._make_eris_incore(self, mo_coeff, ao2mofn, self.verbose)

def _gen_ao2mofn(mf):
    with_df = mf.with_df
    kpt = mf.kpt
    def ao2mofn(mo_coeff):
        return with_df.ao2mo(mo_coeff, kpt, compact=False)
    return ao2mofn


from pyscf.pbc import scf
scf.hf.RHF.MP2 = lib.class_as_method(RMP2)
scf.uhf.UHF.MP2 = lib.class_as_method(UMP2)
scf.ghf.GHF.MP2 = lib.class_as_method(GMP2)
scf.rohf.ROHF.MP2 = None



# pyscf/pbc/mp/kmp2_stagger.py
'''
k-point spin-restricted periodic MP2 calculation using the staggered mesh method
Author: Xin Xing (xxing@berkeley.edu)
Reference: Staggered Mesh Method for Correlation Energy Calculations of Solids: Second-Order Møller–Plesset Perturbation Theory
           J. Chem. Theory Comput. 2021, 17, 8, 4733-4745
'''

import numpy as np
from pyscf import lib
from pyscf.lib import logger
from pyscf.pbc.lib import kpts_helper
from pyscf.lib.parameters import LARGE_DENOM
from pyscf import __config__
from pyscf.pbc.mp.kmp2 import padding_k_idx, _add_padding
from pyscf.pbc.mp import kmp2
from pyscf.pbc.tools.pbc import get_monkhorst_pack_size
from pyscf.pbc.scf.khf import get_occ
from pyscf.pbc import df, dft, scf
import h5py
from pyscf.pbc.lib.kpts_helper import unique

#   Minimum images of k-points to the first BZ
def minimum_image(cell, kpts): 
    """
    Compute the minimum image of k-points in 'kpts' in the first Brillouin zone
    
    Arguments: 
        cell -- a cell instance
        kpts -- a list of k-points 
    
    Returns: 
        kpts_bz -- a list of k-point in the first Brillouin zone
    """
    tmp_kpt = cell.get_scaled_kpts(kpts)
    tmp_kpt = tmp_kpt - np.floor(tmp_kpt)
    tmp_kpt[tmp_kpt > 0.5 - 1e-8] -= 1
    kpts_bz = cell.get_abs_kpts(tmp_kpt)
    return kpts_bz

#   Kernel function for computing the MP2 energy
def kernel(mp, mo_energy, mo_coeff, verbose=logger.NOTE):
    """Computes k-point RMP2 energy using the staggered mesh method

    Args:
        mp (KMP2_stagger): an instance of KMP2_stagger
        mo_energy (list): a list of numpy.ndarray. Each array contains MO energies of
                          shape (Nmo,) for one kpt
        mo_coeff (list): a list of numpy.ndarray. Each array contains MO coefficients
                         of shape (Nao, Nmo) for one kpt
        verbose (int, optional): level of verbosity. Defaults to logger.NOTE (=3).

    Returns:
        KMP2 energy
    """
    cput0 = (logger.process_clock(), logger.perf_counter())
    log = logger.new_logger(mp, verbose)

    nmo = mp.nmo
    nocc = mp.nocc
    nvir = nmo - nocc
    nkpts = mp.nkpts
    nkpts_ov = mp.nkpts_ov
    kconserv = mp.khelper.kconserv

    mem_avail = mp.max_memory - lib.current_memory()[0]
    mem_usage = (nkpts_ov * (nocc * nvir)**2) * 16 / 1e6
    if mp.with_df_ints:
        naux = mp._scf.with_df.auxcell.nao_nr()
        mem_usage += (nkpts_ov**2 * naux * nocc * nvir) * 16 / 1e6
    if mem_usage > mem_avail:
        raise MemoryError('Insufficient memory! MP2 memory usage %d MB (currently available %d MB)'
                          % (mem_usage, mem_avail))
    
    if mp.with_df_ints: 
        Lov = _init_mp_df_eris_stagger(mp)
    else: 
        with_df = df.FFTDF(mp.cell, mp.kpts)
        fao2mo = with_df.ao2mo

    

    #   info of occupied and virtual meshes
    nkpts_ov     = mp.nkpts_ov          #   number of points on the occupied/virtual mesh
    kpts_idx_vir = mp.kpts_idx_vir 
    kpts_idx_occ = mp.kpts_idx_occ

    eia   = np.zeros((nocc,nvir))
    eijab = np.zeros((nocc,nocc,nvir,nvir))

    emp2     = 0.0
    oovv_ij  = np.zeros((nkpts_ov,nocc,nocc,nvir,nvir), dtype=mo_coeff[0].dtype)

    mo_e_o = [mo_energy[k][:nocc] for k in range(nkpts)]
    mo_e_v = [mo_energy[k][nocc:] for k in range(nkpts)]

    #   get location of non-zero/padded elements in occupied and virtual space
    nonzero_opadding, nonzero_vpadding = padding_k_idx(mp, kind="split")

    for iki, ki in enumerate(kpts_idx_occ):
        for ikj, kj in enumerate(kpts_idx_occ):
            for ika, ka in enumerate(kpts_idx_vir):
                kb = kconserv[ki,ka,kj]
                ikb = np.where(kpts_idx_vir == kb)[0][0]

                if mp.with_df_ints:
                    oovv_ij[ika] = (1./nkpts_ov) * lib.einsum("Lia,Ljb->iajb", Lov[iki, ika], Lov[ikj, ikb]).transpose(0,2,1,3)
                else:
                    orbo_i = mo_coeff[ki][:,:nocc]
                    orbo_j = mo_coeff[kj][:,:nocc]
                    orbv_a = mo_coeff[ka][:,nocc:]
                    orbv_b = mo_coeff[kb][:,nocc:]
                    oovv_ij[ika] = fao2mo((orbo_i,orbv_a,orbo_j,orbv_b),
                                    (mp.kpts[ki],mp.kpts[ka],mp.kpts[kj],mp.kpts[kb]),
                                compact=False).reshape(nocc,nvir,nocc,nvir).transpose(0,2,1,3) / nkpts_ov
            for ika, ka in enumerate(kpts_idx_vir):
                kb = kconserv[ki,ka,kj]
                ikb = np.where(kpts_idx_vir == kb)[0][0]

                #   remove zero/padded elements from denominator
                eia = LARGE_DENOM * np.ones((nocc, nvir), dtype=mo_energy[0].dtype)
                n0_ovp_ia = np.ix_(nonzero_opadding[ki], nonzero_vpadding[ka][:nvir])
                eia[n0_ovp_ia] = (mo_e_o[ki][:,None] - mo_e_v[ka])[n0_ovp_ia]

                ejb = LARGE_DENOM * np.ones((nocc, nvir), dtype=mo_energy[0].dtype)
                n0_ovp_jb = np.ix_(nonzero_opadding[kj], nonzero_vpadding[kb][:nvir])
                ejb[n0_ovp_jb] = (mo_e_o[kj][:,None] - mo_e_v[kb])[n0_ovp_jb]

                #   energy calculation
                eijab = lib.direct_sum('ia,jb->ijab',eia,ejb)
                t2_ijab = np.conj(oovv_ij[ika]/eijab)
                woovv = 2*oovv_ij[ika] - oovv_ij[ikb].transpose(0,1,3,2)
                emp2 += np.einsum('ijab,ijab', t2_ijab, woovv).real

    log.timer("KMP2_stagger", *cput0)
    emp2 /= nkpts_ov
    return emp2


def _init_mp_df_eris_stagger(mp):
    """
    Compute 3-center electron repulsion integrals, i.e. (L|ov),
    where `L` denotes DF auxiliary basis functions and `o` and `v` occupied and virtual
    canonical crystalline orbitals that lie on two staggered uniform meshes. 
    Note that `o` and `v` contain kpt indices `ko` and `kv`,
    and the third kpt index `kL` is determined by the conservation of momentum.

    Arguments:
        mp (KMP2_stagger) -- A KMP2_stagger instance

    Returns:
        Lov (numpy.ndarray) -- 3-center DF ints, with shape (nkpts_ov, nkpts_ov, naux, nocc, nvir)
    """

    from pyscf.pbc.df import df
    from pyscf.ao2mo import _ao2mo

    log = logger.Logger(mp.stdout, mp.verbose)

    cell = mp._scf.cell
    if cell.dimension == 2:
        # 2D ERIs are not positive definite. The 3-index tensors are stored in
        # two part. One corresponds to the positive part and one corresponds
        # to the negative part. The negative part is not considered in the
        # DF-driven CCSD implementation.
        raise NotImplementedError

    if mp._scf.with_df._cderi is not None and mp.flag_submesh is True:
        #   When using submeshes for staggered mesh calculation, the 3c DF tensor from mean-field 
        #   calculation can be used directly. 
        with_df = mp._scf.with_df

    else:
        with_df = df.GDF(mp.cell, mp.kpts)

        #   Remove duplicate kpts
        uniq_idx = unique(mp.kpts)[1]
        kpts = np.asarray(mp.kpts)[uniq_idx]

        #   Construct 3-center GDF (L|ia) with ki and ka on two staggered meshes
        with_df.auxcell = df.make_modrho_basis(with_df.cell, with_df.auxbasis,
                                         with_df.exp_to_discard)
        kptij_lst = [(kpti, kpta) for kpti in mp.kpts_occ for kpta in mp.kpts_vir]
        kptij_lst = np.asarray(kptij_lst)

        if isinstance(with_df._cderi_to_save, str):
            cderi = with_df._cderi_to_save
        else:
            cderi = with_df._cderi_to_save.name

        with_df._cderi = cderi
        t1 = (logger.process_clock(), logger.perf_counter())
        with_df._make_j3c(with_df.cell, with_df.auxcell, kptij_lst, cderi)
        t1 = logger.timer_debug1(with_df, 'j3c', *t1)

    nocc = mp.nocc
    nmo = mp.nmo
    nvir = nmo - nocc
    nao = cell.nao_nr()

    #   info of occupied and virtual meshes
    nkpts_ov     = mp.nkpts_ov          #   number of points on the occupied/virtual mesh
    kpts_idx_vir = mp.kpts_idx_vir 
    kpts_idx_occ = mp.kpts_idx_occ
    kpts = mp.kpts


    mo_coeff = _add_padding(mp, mp.mo_coeff, mp.mo_energy)[0]
    dtype = np.result_type(np.complex128, *mo_coeff)

    Lov = np.empty((nkpts_ov, nkpts_ov), dtype=object)

    cput0 = (logger.process_clock(), logger.perf_counter())
    bra_start = 0
    bra_end = nocc
    ket_start = nmo+nocc
    ket_end = ket_start + nvir
    with h5py.File(with_df._cderi, 'r') as f:
        kptij_lst = f['j3c-kptij'][:]
        tao = []
        ao_loc = None

        for iki, ki in enumerate(kpts_idx_occ):
            for ika, ka in enumerate(kpts_idx_vir): 
                kpti = kpts[ki]
                kpta = kpts[ka]
                kpti_kpta = np.array( (kpti, kpta) )
                Lpq_ao = np.asarray(df._getitem(f, 'j3c', kpti_kpta, kptij_lst))

                mo = np.hstack((mo_coeff[ki], mo_coeff[ka]))
                mo = np.asarray(mo, dtype=dtype, order='F')
                
                #Note: Lpq.shape[0] != naux if linear dependency is found in auxbasis
                if Lpq_ao[0].size != nao**2:  # aosym = 's2'
                    Lpq_ao = lib.unpack_tril(Lpq_ao).astype(np.complex128)
                out = _ao2mo.r_e2(Lpq_ao, mo, (bra_start, bra_end, ket_start, ket_end), tao, ao_loc)
                Lov[iki, ika] = out.reshape(-1, nocc, nvir)

    log.timer_debug1("transforming DF-MP2 integrals", *cput0)

    return Lov


#   Class for staggered mesh MP2
class KMP2_stagger(kmp2.KMP2):
    def __init__(self, mf, frozen=None, flag_submesh=False):
        self.cell = mf.cell
        self._scf = mf
        self.verbose = self.cell.verbose
        self.stdout = self.cell.stdout
        self.max_memory = self.cell.max_memory

        #   MP2 energy 
        self.e_corr = None

        #   nocc and nmo variables
        self._nocc = None
        self._nmo = None

        #   Construction of orbitals and orbital energies on two staggered meshes
        nks = get_monkhorst_pack_size(mf.cell, mf.kpts)  #   mesh size
        self.flag_submesh = flag_submesh
        if flag_submesh is True: 
            #   Choice 1: Use two staggered submeshes of mf.kpts of half size along each dimension
            if np.any( nks % 2) : 
                print("The k-point mesh in the input mf has odd size in certain dimension. Cannot use its submeshes for staggered mesh calculation!")
                log = logger.Logger(self.stdout, self.verbose)
                log.warn('The k-point mesh in the input mf has odd size in certain dimension. \n'
                        'Cannot use its submeshes for staggered mesh calculation!\n'
                        'You may set "flag_submesh" to False')
                raise RuntimeError

            nks_ov = nks // 2
            shift = mf.kpts[ np.argmin( np.sum(mf.kpts**2, axis=1) )]   #   shift of the original k-point mesh
            half_shift = mf.cell.get_abs_kpts( [0.5/n for n in nks_ov] )
            kpts_vir = shift + mf.cell.make_kpts(nks_ov, with_gamma_point=True)
            kpts_occ = kpts_vir + half_shift

            self.kpts = mf.kpts
            self.mo_energy = mf.mo_energy
            self.mo_coeff = mf.mo_coeff
            self.mo_occ = get_occ(mf, mo_energy_kpts=mf.mo_energy)

        else: 
            #   Choice 2: Apply non-SCF calculation to compute orbitals and orbital energies on a 
            #             staggered mesh of the same size as mf.kpts. 
            if mf.cell.dimension < 3:
                # The non-SCF calculation of orbitals/orbital energies is currently only valid 
                # for 3D systems in PySCF. For example, when using get_band with a reference HF 
                # system of mesh 1*1*k to compute orbitals and orbital energies on a 1*1*2k mesh, 
                # the obtained mo_energy and  mo_coeff are not in good agreement with those computed 
                # from an SCF-HF calculation with mesh 1*1*2k.           
                raise NotImplementedError
            
            half_shift = mf.cell.get_abs_kpts( [0.5/n for n in nks] )
            kpts_vir = mf.kpts
            kpts_occ = kpts_vir + half_shift
            kpts = np.concatenate( (kpts_occ, kpts_vir), axis=0)
            if isinstance(mf, scf.KRHF):
                with lib.temporary_env(mf, exxdiv='vcut_sph', with_df=df.FFTDF(mf.cell, mf.kpts)):            
                    mo_energy, mo_coeff = mf.get_bands( kpts )
                    mo_occ = get_occ(mf, mo_energy_kpts=mo_energy)

            elif isinstance(mf, dft.KRKS):   
                with lib.temporary_env(mf, with_df=df.FFTDF(mf.cell, mf.kpts)):        
                    mo_energy, mo_coeff = mf.get_bands( kpts )
                    mo_occ = get_occ(mf, mo_energy_kpts=mo_energy)
                    
            self.kpts = kpts
            self.mo_energy = mo_energy
            self.mo_coeff = mo_coeff
            self.mo_occ = mo_occ
        
        if isinstance(self._scf.with_df, df.GDF):
            self.with_df_ints = True
        else: 
            self.with_df_ints = False
        
        #   Basic info
        self.nkpts = len(self.kpts)
        self.frozen = frozen
        self.khelper = kpts_helper.KptsHelper(self.cell, self.kpts)

        #   Mesh info for the staggered mesh method
        self.kpts_vir = kpts_vir
        self.kpts_occ = kpts_occ
        self.nkpts_ov = len(self.kpts_vir)

        #   Map kpts, kpts_vir, and kpts_occ to the first BZ for better processing
        kpts_bz = minimum_image(self.cell, self.kpts)
        kpts_occ_bz = minimum_image(self.cell, self.kpts_occ)
        kpts_vir_bz = minimum_image(self.cell, self.kpts_vir)

        #   indices of virtual k-points in self.kpts
        self.kpts_idx_vir = [ np.asarray(np.sum((kpts_bz - kvir)**2, axis=-1) < 1e-10).nonzero()[0] for kvir in kpts_vir_bz]
        self.kpts_idx_vir = np.concatenate(self.kpts_idx_vir).astype(int)

        #   indices of occupied k-points in self.kpts
        self.kpts_idx_occ = [ np.asarray(np.sum((kpts_bz - kocc)**2, axis=-1) < 1e-10).nonzero()[0] for kocc in kpts_occ_bz]
        self.kpts_idx_occ = np.concatenate(self.kpts_idx_occ).astype(int)

        if (
            len(self.kpts_idx_vir) != self.nkpts_ov or len(np.unique(self.kpts_idx_vir)) != self.nkpts_ov or
            len(self.kpts_idx_occ) != self.nkpts_ov or len(np.unique(self.kpts_idx_occ)) != self.nkpts_ov
            ):
            log = logger.Logger(self.stdout, self.verbose)
            log.warn('Cannot locate the provided occupied/virtual submeshes in the large k-point mesh')            
            raise RuntimeError

    def dump_flags(self):
        log = logger.Logger(self.stdout, self.verbose)
        log.info('')
        log.info('******** %s ********', self.__class__)
        log.info('method = %s', self.__class__.__name__)
        nocc = self.nocc
        nvir = self.nmo - nocc
        log.info('Staggerd mesh method for MP2 nocc = %d, nvir = %d', nocc, nvir)    
        
        nks_ov = get_monkhorst_pack_size(self.cell, self.kpts_vir)  #   mesh size
        if self.flag_submesh is True: 
            log.info('Two %d*%d*%d-sized submeshes of mf.kpts are used', nks_ov[0], nks_ov[1], nks_ov[2])
        else: 
            log.info('Two %d*%d*%d-sized meshes are used based on non-SCF calculation', nks_ov[0], nks_ov[1], nks_ov[2])
        
        if self.frozen is not None:
            log.info('frozen orbitals = %s', str(self.frozen))
        log.info('KMP2 energy = %.15g' % (self.e_corr))
        return self

    def kernel(self):
        if self.mo_energy is None or self.mo_coeff is None:
            log = logger.Logger(self.stdout, self.verbose)
            log.warn('mo_coeff, mo_energy are not provided.')
            raise RuntimeError

        mo_coeff, mo_energy = _add_padding(self, self.mo_coeff, self.mo_energy)

        self.e_corr = kernel(self, mo_energy, mo_coeff)
        self.dump_flags()        
        return self.e_corr


if __name__ == '__main__':
    from pyscf.pbc import df, gto, scf, mp
    
    '''
    Example calculation on hydrogen dimer
    '''
    cell = gto.Cell()
    cell.pseudo = 'gth-pade'
    cell.basis = 'gth-szv'
    cell.ke_cutoff=100
    cell.atom='''
        H 3.00   3.00   2.10
        H 3.00   3.00   3.90
        '''
    cell.a = '''
        6.0   0.0   0.0
        0.0   6.0   0.0
        0.0   0.0   6.0
        '''
    cell.unit = 'B'
    cell.verbose = 4
    cell.build()


    #   HF calculation using FFTDF
    nks_mf = [2,2,2] 
    kpts = cell.make_kpts(nks_mf, with_gamma_point=True)
    kmf = scf.KRHF(cell, kpts, exxdiv='ewald')
    ehf = kmf.kernel()

    #   staggered mesh KMP2 calculation using two submeshes of size [1,1,1] in kmf.kpts
    kmp = KMP2_stagger(kmf, flag_submesh=True)
    emp2 = kmp.kernel()    
    assert((abs(emp2 - -0.0160902544091997))<1e-5)

    #   staggered mesh KMP2 calculation using two meshes of size [2,2,2], one of them is kmf.kpts
    kmp = KMP2_stagger(kmf, flag_submesh=False)
    emp2 = kmp.kernel()
    assert((abs(emp2 -  -0.0140289970302513))<1e-5)

    #   standard KMP2 calculation
    kmp = mp.KMP2(kmf)
    emp2, _ = kmp.kernel()
    assert((abs(emp2 - -0.0143904878990777))<1e-5)


    #   HF calculation using GDF
    nks_mf = [2,2,2] 
    kpts = cell.make_kpts(nks_mf, with_gamma_point=True)
    kmf = scf.KRHF(cell, kpts, exxdiv='ewald')
    gdf = df.GDF(cell, kpts).build()
    kmf.with_df = gdf
    ehf = kmf.kernel()

    #   staggered mesh KMP2 calculation using two submeshes of size [1,1,1] in kmf.kpts
    kmp = KMP2_stagger(kmf, flag_submesh=True)
    emp2 = kmp.kernel()    
    assert((abs(emp2 - -0.0158364523431071))<1e-5)

    #   staggered mesh KMP2 calculation using two meshes of size [2,2,2], one of them is kmf.kpts
    kmp = KMP2_stagger(kmf, flag_submesh=False)
    emp2 = kmp.kernel()
    assert((abs(emp2 -  -0.0140280303691396))<1e-5)

    #   standard KMP2 calculation
    kmp = mp.KMP2(kmf)
    emp2, _ = kmp.kernel()
    assert((abs(emp2 - -0.0141829343769316))<1e-5)


    '''
    Example calculation on diamond
    '''

    cell = gto.Cell()
    cell.pseudo = 'gth-pade'
    cell.basis = 'gth-szv'
    cell.ke_cutoff=100
    cell.atom='''
        C     0.      0.      0.
        C     1.26349729, 0.7294805 , 0.51582061
        '''
    cell.a = '''
        2.52699457, 0.        , 0.    
        1.26349729, 2.18844149, 0.  
        1.26349729, 0.7294805 , 2.06328243
        '''
    cell.unit = 'angstrom'
    cell.verbose = 4
    cell.build()


    #   HF calculation using FFTDF
    nks_mf = [2,2,2] 
    kpts = cell.make_kpts(nks_mf, with_gamma_point=True)
    kmf = scf.KRHF(cell, kpts, exxdiv='ewald')
    ehf = kmf.kernel()

    #   staggered mesh KMP2 calculation using two submeshes of size [1,1,1] in kmf.kpts
    kmp = KMP2_stagger(kmf, flag_submesh=True)
    emp2 = kmp.kernel()    
    assert((abs(emp2 - -0.156289981810986))<1e-5)

    #   staggered mesh KMP2 calculation using two meshes of size [2,2,2], one of them is kmf.kpts
    kmp = KMP2_stagger(kmf, flag_submesh=False)
    emp2 = kmp.kernel()
    assert((abs(emp2 -  -0.105454107635884))<1e-5)

    #   standard KMP2 calculation
    kmp = mp.KMP2(kmf)
    emp2, _ = kmp.kernel()
    assert((abs(emp2 - -0.095517731535516))<1e-5)



    #   HF calculation using GDF
    nks_mf = [2,2,2] 
    kpts = cell.make_kpts(nks_mf, with_gamma_point=True)
    kmf = scf.KRHF(cell, kpts, exxdiv='ewald')
    gdf = df.GDF(cell, kpts).build()
    kmf.with_df = gdf
    ehf = kmf.kernel()

    #   staggered mesh KMP2 calculation using two submeshes of size [1,1,1] in kmf.kpts
    kmp = KMP2_stagger(kmf, flag_submesh=True)
    emp2 = kmp.kernel()    
    assert((abs(emp2 - -0.154923152683604))<1e-5)

    #   staggered mesh KMP2 calculation using two meshes of size [2,2,2], one of them is kmf.kpts
    kmp = KMP2_stagger(kmf, frozen=1, flag_submesh=False)
    emp2 = kmp.kernel()
    assert((abs(emp2 -  -0.105421948003715))<1e-5)

    #   standard KMP2 calculation
    kmp = mp.KMP2(kmf, frozen=1)
    emp2, _ = kmp.kernel()
    assert((abs(emp2 - -0.0952009565805345))<1e-5)
