<commit_before>#!/usr/bin/env python
#
# Author: <NAME> <<EMAIL>>

from functools import reduce
import numpy
import scipy.linalg
import pyscf.lib.parameters as param
from pyscf import gto

def lowdin(s):
    ''' new basis is |mu> c^{lowdin}_{mu i} '''
    e, v = scipy.linalg.eigh(s)
    idx = e > 1e-15
    return numpy.dot(v[:,idx]/numpy.sqrt(e[idx]), v[:,idx].T.conj())

def schmidt(s):
    c = numpy.linalg.cholesky(s)
    return scipy.linalg.solve_triangular(c, numpy.eye(c.shape[1]), lower=True,
                                         overwrite_b=False).T.conj()

def vec_lowdin(c, s=1):
    ''' lowdin orth for the metric c.T*s*c and get x, then c*x'''
    #u, w, vh = numpy.linalg.svd(c)
    #return numpy.dot(u, vh)
    # svd is slower than eigh
    return numpy.dot(c, lowdin(reduce(numpy.dot, (c.T,s,c))))

def vec_schmidt(c, s=1):
    ''' schmidt orth for the metric c.T*s*c and get x, then c*x'''
    if isinstance(s, numpy.ndarray):
        return numpy.dot(c, schmidt(reduce(numpy.dot, (c.T,s,c))))
    else:
        return numpy.linalg.qr(c)[0]

def weight_orth(s, weight):
    ''' new basis is |mu> c_{mu i}, c = w[(wsw)^{-1/2}]'''
    s1 = weight[:,None] * s * weight
    c = lowdin(s1)
    return weight[:,None] * c


def pre_orth_ao(mol, method='ANO'):
    '''Localized GTOs for each atom.  Possible localized methods include
    the ANO/MINAO projected basis or fraction-averaged RHF'''
    if method.upper() in ('ANO', 'MINAO'):
# Use ANO/MINAO basis to define the strongly occupied set
        return pre_orth_project_ano(mol, method)
    else:
        return pre_orth_ao_atm_scf(mol)

def pre_orth_project_ano(mol, basname):
    import pyscf.scf
    def search_atm_l(atm, l):
        idx = []
        p0 = 0
        for ib, b in enumerate(atm._bas):
            l0 = atm.bas_angular(ib)
            nctr = atm.bas_nctr(ib)
            if l0 == l:
                idx.extend(range(p0,p0+(2*l+1)*nctr))
            p0 += (2*l0+1) * nctr
        return idx

    aos = {}
    for symb in mol._basis.keys():
        rawsymb = gto.mole._rm_digit(symb)
        stdsymb = param.ELEMENTS[gto.mole._ELEMENTDIC[rawsymb.upper()]][0]
        atm = gto.Mole()
        atm._atm, atm._bas, atm._env = \
                atm.make_env([[stdsymb,(0,0,0)]], {stdsymb:mol._basis[symb]}, [])

        if 'GHOST' in symb.upper():
            aos[symb] = numpy.eye(atm.nao_nr())
            continue

        s0 = atm.intor_symmetric('cint1e_ovlp_sph')

        basis_add = gto.basis.load(basname, stdsymb)
        atmp = gto.Mole()
        atmp._atm, atmp._bas, atmp._env = \
                atmp.make_env([[stdsymb,(0,0,0)]], {stdsymb:basis_add}, [])
        ano = pyscf.scf.addons.project_mo_nr2nr(atmp, 1, atm)
        rm_ano = numpy.eye(ano.shape[0]) - reduce(numpy.dot, (ano, ano.T, s0))
        c = rm_ano.copy()
        for l in range(4):
            idx  = numpy.array(search_atm_l(atm, l))
            idxp = numpy.array(search_atm_l(atmp, l))
            if len(idx) > len(idxp) > 0:
# For angular l, first place the projected ANO, then the rest AOs.
                sdiag = reduce(numpy.dot, (rm_ano[:,idx].T, s0, rm_ano[:,idx])).diagonal()
                c[:,idx[:len(idxp)]] = ano[:,idxp]
                c[:,idx[len(idxp):]] = rm_ano[:,numpy.argsort(sdiag)[len(idxp):]]
            elif len(idxp) >= len(idx) > 0:
                c[:,idx] = ano[:,idxp[:len(idx)]]
        aos[symb] = c

    nao = mol.nao_nr()
    c = numpy.zeros((nao,nao))
    p0 = 0
    for ia in range(mol.natm):
        symb = mol.atom_symbol(ia)
        if symb in mol._basis:
            ano = aos[symb]
        else:<commit_msg>Fix bug in meta-lowdin, when ANO basis is smaller than the input basis<commit_after>#!/usr/bin/env python
#
# Author: <NAME> <<EMAIL>>

from functools import reduce
import numpy
import scipy.linalg
import pyscf.lib.parameters as param
from pyscf import gto

def lowdin(s):
    ''' new basis is |mu> c^{lowdin}_{mu i} '''
    e, v = scipy.linalg.eigh(s)
    idx = e > 1e-15
    return numpy.dot(v[:,idx]/numpy.sqrt(e[idx]), v[:,idx].T.conj())

def schmidt(s):
    c = numpy.linalg.cholesky(s)
    return scipy.linalg.solve_triangular(c, numpy.eye(c.shape[1]), lower=True,
                                         overwrite_b=False).T.conj()

def vec_lowdin(c, s=1):
    ''' lowdin orth for the metric c.T*s*c and get x, then c*x'''
    #u, w, vh = numpy.linalg.svd(c)
    #return numpy.dot(u, vh)
    # svd is slower than eigh
    return numpy.dot(c, lowdin(reduce(numpy.dot, (c.T,s,c))))

def vec_schmidt(c, s=1):
    ''' schmidt orth for the metric c.T*s*c and get x, then c*x'''
    if isinstance(s, numpy.ndarray):
        return numpy.dot(c, schmidt(reduce(numpy.dot, (c.T,s,c))))
    else:
        return numpy.linalg.qr(c)[0]

def weight_orth(s, weight):
    ''' new basis is |mu> c_{mu i}, c = w[(wsw)^{-1/2}]'''
    s1 = weight[:,None] * s * weight
    c = lowdin(s1)
    return weight[:,None] * c


def pre_orth_ao(mol, method='ANO'):
    '''Localized GTOs for each atom.  Possible localized methods include
    the ANO/MINAO projected basis or fraction-averaged RHF'''
    if method.upper() in ('ANO', 'MINAO'):
# Use ANO/MINAO basis to define the strongly occupied set
        return pre_orth_project_ano(mol, method)
    else:
        return pre_orth_ao_atm_scf(mol)

def pre_orth_project_ano(mol, basname):
    import pyscf.scf
    def search_atm_l(atm, l):
        idx = []
        p0 = 0
        for ib, b in enumerate(atm._bas):
            l0 = atm.bas_angular(ib)
            nctr = atm.bas_nctr(ib)
            if l0 == l:
                idx.extend(range(p0,p0+(2*l+1)*nctr))
            p0 += (2*l0+1) * nctr
        return idx

    aos = {}
    for symb in mol._basis.keys():
        rawsymb = gto.mole._rm_digit(symb)
        stdsymb = param.ELEMENTS[gto.mole._ELEMENTDIC[rawsymb.upper()]][0]
        atm = gto.Mole()
        atm._atm, atm._bas, atm._env = \
                atm.make_env([[stdsymb,(0,0,0)]], {stdsymb:mol._basis[symb]}, [])

        if 'GHOST' in symb.upper():
            aos[symb] = numpy.eye(atm.nao_nr())
            continue

        s0 = atm.intor_symmetric('cint1e_ovlp_sph')

        basis_add = gto.basis.load(basname, stdsymb)
        atmp = gto.Mole()
        atmp._atm, atmp._bas, atmp._env = \
                atmp.make_env([[stdsymb,(0,0,0)]], {stdsymb:basis_add}, [])
        ano = pyscf.scf.addons.project_mo_nr2nr(atmp, 1, atm)
        rm_ano = numpy.eye(ano.shape[0]) - reduce(numpy.dot, (ano, ano.T, s0))
        c = rm_ano.copy()
        for l in range(4):
            idx  = numpy.asarray(search_atm_l(atm, l))
            idxp = numpy.asarray(search_atm_l(atmp, l))
            if len(idx) > len(idxp) > 0:
# For angular l, first place the projected ANO, then the rest AOs.
                sdiag = reduce(numpy.dot, (rm_ano[:,idx].T, s0, rm_ano[:,idx])).diagonal()
                nleft = (len(idx) - len(idxp)) // (2*l+1)
                shell_average = numpy.einsum('ij->i', sdiag.reshape(-1,l*2+1))
                shell_rest = numpy.argsort(-shell_average)[:nleft]
                idx_rest = []
                for k in shell_rest:
                    idx_rest.extend(idx[k*(2*l+1):(k+1)*(2*l+1)])
                c[:,idx[:len(idxp)]] = ano[:,idxp]
                c[:,idx[len(idxp):]] = rm_ano[:,idx_rest]
            elif len(idxp) >= len(idx) > 0:  # More ANOs than the mol basis functions
                c[:,idx] = ano[:,idxp[:len(idx)]]
        aos[symb] = c

    nao = mol.nao_nr()
    c = numpy.zeros((nao,nao))
    p0 = 0
    for ia in range(mol.natm):
        symb = mol.atom_symbol(ia)
        if symb in mol._basis:
            ano = aos[symb]
        else:
{'fraction_non_alphanumeric': 0.16043898615103214, 'fraction_numerical': 0.012281160177684871, 'mean_word_length': 2.313852813852814, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 24, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2437268', 'n_tokens_mistral': 3064, 'n_tokens_neox': 2987, 'n_words': 629}