<commit_before>#!/usr/bin/env python

import sys
import copy
import numpy
from pyscf import lib
from pyscf.fci import cistring
from pyscf import symm

def large_ci(ci, norb, nelec, tol=.1, return_strs=True):
    '''Search for the largest CI coefficients
    '''
    neleca, nelecb = _unpack(nelec)
    addra, addrb = numpy.where(abs(ci) > tol)
    if addra.size == 0: # No large CI coefficient > tol
        addra, addrb = numpy.unravel_index(numpy.argmax(abs(ci)), ci.shape)
        addra = numpy.asarray([addra])
        addrb = numpy.asarray([addrb])
    strsa = cistring.addrs2str(norb, neleca, addra)
    strsb = cistring.addrs2str(norb, nelecb, addrb)
    if return_strs:
        strsa = [bin(x) for x in strsa]
        strsb = [bin(x) for x in strsb]
        return list(zip(ci[addra,addrb], strsa, strsb))
    else:
        occslsta = cistring._strs2occslst(strsa, norb)
        occslstb = cistring._strs2occslst(strsb, norb)
        return list(zip(ci[addra,addrb], occslsta, occslstb))

def initguess_triplet(norb, nelec, binstring):
    '''Generate a triplet initial guess for FCI solver
    '''
    neleca, nelecb = _unpack(nelec)
    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    addr = cistring.str2addr(norb, neleca, int(binstring,2))
    ci0 = numpy.zeros((na,nb))
    ci0[addr,0] = numpy.sqrt(.5)
    ci0[0,addr] =-numpy.sqrt(.5)
    return ci0

def symm_initguess(norb, nelec, orbsym, wfnsym=0, irrep_nelec=None):
    '''Generate CI wavefunction initial guess which has the given symmetry.

    Args:
        norb : int
            Number of orbitals.
        nelec : int or 2-item list
            Number of electrons, or 2-item list for (alpha, beta) electrons
        orbsym : list of int
            The irrep ID for each orbital.

    Kwags:
        wfnsym : int
            The irrep ID of target symmetry
        irrep_nelec : dict
            Freeze occupancy for certain irreps

    Returns:
        CI coefficients 2D array which has the target symmetry.
    '''
    neleca, nelecb = _unpack(nelec)
    orbsym = numpy.asarray(orbsym)
    if not isinstance(orbsym[0], numpy.number):
        raise RuntimeError('TODO: convert irrep symbol to irrep id')

    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    ci1 = numpy.zeros((na,nb))

########################
# pass 1: The fixed occs
    orbleft = numpy.ones(norb, dtype=bool)
    stra = numpy.zeros(norb, dtype=bool)
    strb = numpy.zeros(norb, dtype=bool)
    if irrep_nelec is not None:
        for k,n in irrep_nelec.items():
            orbleft[orbsym==k] = False
            if isinstance(n, (int, numpy.number)):
                idx = numpy.where(orbsym==k)[0][:n//2]
                stra[idx] = True
                strb[idx] = True
            else:
                na, nb = n
                stra[numpy.where(orbsym==k)[0][:na]] = True
                strb[numpy.where(orbsym==k)[0][:nb]] = True
                if (na-nb)%2:
                    wfnsym ^= k

    orbleft = numpy.where(orbleft)[0]
    neleca_left = neleca - stra.sum()
    nelecb_left = nelecb - strb.sum()
    spin = neleca_left - nelecb_left
    assert(neleca_left >= 0)
    assert(nelecb_left >= 0)
    assert(spin >= 0)

########################
# pass 2: search pattern
    def gen_str_iter(orb_list, nelec):
        if nelec == 1:
            for i in orb_list:
                yield [i]
        elif nelec >= len(orb_list):
            yield orb_list
        else:
            restorb = orb_list[1:]
            #yield from gen_str_iter(restorb, nelec)
            for x in gen_str_iter(restorb, nelec):
                yield x
            for x in gen_str_iter(restorb, nelec-1):
                yield [orb_list[0]] + x

# search for alpha and beta pattern which match to the required symmetry
    def query(target, nelec_atmost, spin, orbsym):
        norb = len(orbsym)
        for excite_level in range(1, nelec_atmost+1):
            for beta_only in gen_str_iter(range(norb), excite_level):
                alpha_allow = [i for i in range(norb) if i not in beta_only]
                alpha_orbsym = orbsym[alpha_allow]
                alpha_target = target
                for i in beta_only:
                    alpha_target ^= orbsym[i]
                alpha_only = symm.route(alpha_target, spin+excite_level, alpha_orbsym)
                if alpha_only:
                    alpha_only = [alpha_allow[i] for i in alpha_only]
                    return alpha_only, beta_only
        raise RuntimeError('No pattern found for wfn irrep %s over orbsym %s'
                           % (target, orbsym))

    if spin == 0:
        aonly = bonly = []
        if wfnsym != 0:
            aonly, bonly = query(wfnsym, neleca_left, spin, orbsym[orbleft])
    else:
        # 1. assume "nelecb_left" doubly occupied orbitals
        # search for alpha pattern which match to the required symmetry
        aonly, bonly = orbleft[symm.route(wfnsym, spin, orbsym[orbleft])], []
        # dcompose doubly occupied orbitals, search for alpha and beta pattern
        if len(aonly) != spin:
            aonly, bonly = query(wfnsym, neleca_left, spin, orbsym[orbleft])

    ndocc = neleca_left - len(aonly) # == nelecb_left - len(bonly)
    docc_allow = numpy.ones(len(orbleft), dtype=bool)
    docc_allow[aonly] = False
    docc_allow[bonly] = False
    docclst = orbleft[numpy.where(docc_allow)[0]][:ndocc]
    stra[docclst] = True
    strb[docclst] = True

    def find_addr_(stra, aonly, nelec):
        stra[orbleft[aonly]] = True
        return cistring.str2addr(norb, nelec, ('%i'*norb)%tuple(stra)[::-1])
    if bonly:
        if spin > 0:
            aonly, socc_only = aonly[:-spin], aonly[-spin:]
            stra[orbleft[socc_only]] = True
        stra1 = stra.copy()
        strb1 = strb.copy()

        addra = find_addr_(stra, aonly, neleca)
        addrb = find_addr_(strb, bonly, nelecb)
        addra1 = find_addr_(stra1, bonly, neleca)
        addrb1 = find_addr_(strb1, aonly, nelecb)
        ci1[addra,addrb] = ci1[addra1,addrb1] = numpy.sqrt(.5)
    else:
        addra = find_addr_(stra, aonly, neleca)
        addrb = find_addr_(strb, bonly, nelecb)
        ci1[addra,addrb] = 1

#    target = 0
#    for i,k in enumerate(stra):
#        if k:
#            target ^= orbsym[i]
#    for i,k in enumerate(strb):
#        if k:
#            target ^= orbsym[i]
#    print target
    return ci1


def _symmetrize_wfn(ci, strsa, strsb, orbsym, wfnsym=0):
    ci = ci.reshape(strsa.size,strsb.size)
    airreps = numpy.zeros(strsa.size, dtype=numpy.int32)
    birreps = numpy.zeros(strsb.size, dtype=numpy.int32)
    for i, ir in enumerate(orbsym):
        airreps[numpy.bitwise_and(strsa, 1<<i) > 0] ^= ir
        birreps[numpy.bitwise_and(strsb, 1<<i) > 0] ^= ir
    mask = (airreps.reshape(-1,1) ^ birreps) == wfnsym
    ci1 = numpy.zeros_like(ci)
    ci1[mask] = ci[mask]
    ci1 *= 1/numpy.linalg.norm(ci1)
    return ci1
def symmetrize_wfn(ci, norb, nelec, orbsym, wfnsym=0):
    '''Symmetrize the CI wavefunction by zeroing out the determinants which
    do not have the right symmetry.

    Args:
        ci : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        nelec : int or 2-item list
            Number of electrons, or 2-item list for (alpha, beta) electrons
        orbsym : list of int
            The irrep ID for each orbital.

    Kwags:
        wfnsym : int
            The irrep ID of target symmetry

    Returns:
        2D array which is the symmetrized CI coefficients
    '''
    neleca, nelecb = _unpack(nelec)
    strsa = numpy.asarray(cistring.gen_strings4orblist(range(norb), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(norb), nelecb))
    return _symmetrize_wfn(ci, strsa, strsb, orbsym, wfnsym)

def _guess_wfnsym(ci, strsa, strsb, orbsym):
    na = len(strsa)
    nb = len(strsb)
    if isinstance(ci, numpy.ndarray) and ci.ndim <= 2:
        assert(ci.size == na*nb)
        idx = numpy.argmax(ci)
    else:
        assert(ci[0].size == na*nb)
        idx = ci[0].argmax()
    stra = strsa[idx // nb]
    strb = strsb[idx % nb ]

    airrep = 0
    birrep = 0
    for i, ir in enumerate(orbsym):
        if (stra & (1<<i)):
            airrep ^= ir
        if (strb & (1<<i)):
            birrep ^= ir
    return airrep ^ birrep
def guess_wfnsym(ci, norb, nelec, orbsym):
    '''Guess the wavefunction symmetry based on the non-zero elements in the
    given CI coefficients.

    Args:
        ci : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        nelec : int or 2-item list
            Number of electrons, or 2-item list for (alpha, beta) electrons
        orbsym : list of int
            The irrep ID for each orbital.

    Returns:
        Irrep ID
    '''
    neleca, nelecb = _unpack(nelec)
    strsa = numpy.asarray(cistring.gen_strings4orblist(range(norb), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(norb), nelecb))
    return _guess_wfnsym(ci, strsa, strsb, orbsym)


def des_a(ci0, norb, neleca_nelecb, ap_id):
    r'''Construct (N-1)-electron wavefunction by removing an alpha electron from
    the N-electron wavefunction.

    ... math::

        |N-1\rangle = \hat{a}_p |N\rangle

    Args:
        ci0 : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        (neleca,nelecb) : (int,int)
            Number of (alpha, beta) electrons of the input CI function
        ap_id : int
            Orbital index (0-based), for the annihilation operator

    Returns:
        2D array, row for alpha strings and column for beta strings.  Note it
        has different number of rows to the input CI coefficients
    '''
    neleca, nelecb = neleca_nelecb
    if ci0.ndim == 1:
        ci0 = ci0.reshape(cistring.num_strings(norb, neleca),
                          cistring.num_strings(norb, nelecb))
    if neleca <= 0:
        return numpy.zeros((0, ci0.shape[1]))
    des_index = cistring.gen_des_str_index(range(norb), neleca)
    na_ci1 = cistring.num_strings(norb, neleca-1)
    ci1 = numpy.zeros((na_ci1, ci0.shape[1]))

    entry_has_ap = (des_index[:,:,1] == ap_id)
    addr_ci0 = numpy.any(entry_has_ap, axis=1)
    addr_ci1 = des_index[entry_has_ap,2]
    sign = des_index[entry_has_ap,3]
    #print(addr_ci0)
    #print(addr_ci1)
    ci1[addr_ci1] = sign.reshape(-1,1) * ci0[addr_ci0]
    return ci1

def des_b(ci0, norb, neleca_nelecb, ap_id):
    r'''Construct (N-1)-electron wavefunction by removing a beta electron from
    N-electron wavefunction.

    Args:
        ci0 : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        (neleca,nelecb) : (int,int)
            Number of (alpha, beta) electrons of the input CI function
        ap_id : int
            Orbital index (0-based), for the annihilation operator

    Returns:
        2D array, row for alpha strings and column for beta strings. Note it
        has different number of columns to the input CI coefficients.
    '''
    neleca, nelecb = neleca_nelecb
    if ci0.ndim == 1:
        ci0 = ci0.reshape(cistring.num_strings(norb, neleca),
                          cistring.num_strings(norb, nelecb))
    if nelecb <= 0:
        return numpy.zeros((ci0.shape[0], 0))
    des_index = cistring.gen_des_str_index(range(norb), nelecb)
    nb_ci1 = cistring.num_strings(norb, nelecb-1)
    ci1 = numpy.zeros((ci0.shape[0], nb_ci1))

    entry_has_ap = (des_index[:,:,1] == ap_id)
    addr_ci0 = numpy.any(entry_has_ap, axis=1)
    addr_ci1 = des_index[entry_has_ap,2]
    sign = des_index[entry_has_ap,3]
    # This sign prefactor accounts for interchange of operators with alpha and beta spins
    if neleca % 2 == 1:
        sign *= -1
    ci1[:,addr_ci1] = ci0[:,addr_ci0] * sign
    return ci1

def cre_a(ci0, norb, neleca_nelecb, ap_id):
    r'''Construct (N+1)-electron wavefunction by adding an alpha electron in
    the N-electron wavefunction.

    ... math::

        |N+1\rangle = \hat{a}^+_p |N\rangle

    Args:
        ci0 : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        (neleca,nelecb) : (int,int)
            Number of (alpha, beta) electrons of the input CI function
        ap_id : int
            Orbital index (0-based), for the creation operator

    Returns:
        2D array, row for alpha strings and column for beta strings. Note it
        has different number of rows to the input CI coefficients.
    '''
    neleca, nelecb = neleca_nelecb
    if ci0.ndim == 1:
        ci0 = ci0.reshape(cistring.num_strings(norb, neleca),
                          cistring.num_strings(norb, nelecb))
    if neleca >= norb:
        return numpy.zeros((0, ci0.shape[1]))
    cre_index = cistring.gen_cre_str_index(range(norb), neleca)
    na_ci1 = cistring.num_strings(norb, neleca+1)
    ci1 = numpy.zeros((na_ci1, ci0.shape[1]))

    entry_has_ap = (cre_index[:,:,0] == ap_id)
    addr_ci0 = numpy.any(entry_has_ap, axis=1)
    addr_ci1 = cre_index[entry_has_ap,2]
    sign = cre_index[entry_has_ap,3]
    ci1[addr_ci1] = sign.reshape(-1,1) * ci0[addr_ci0]
    return ci1

# construct (N+1)-electron wavefunction by adding a beta electron to
# N-electron wavefunction:
def cre_b(ci0, norb, neleca_nelecb, ap_id):
    r'''Construct (N+1)-electron wavefunction by adding a beta electron in
    the N-electron wavefunction.

    Args:
        ci0 : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        (neleca,nelecb) : (int,int)
            Number of (alpha, beta) electrons of the input CI function
        ap_id : int
            Orbital index (0-based), for the creation operator

    Returns:
        2D array, row for alpha strings and column for beta strings. Note it
        has different number of columns to the input CI coefficients.
    '''
    neleca, nelecb = neleca_nelecb
    if ci0.ndim == 1:
        ci0 = ci0.reshape(cistring.num_strings(norb, neleca),
                          cistring.num_strings(norb, nelecb))
    if nelecb >= norb:
        return numpy.zeros((ci0.shape[0], 0))
    cre_index = cistring.gen_cre_str_index(range(norb), nelecb)
    nb_ci1 = cistring.num_strings(norb, nelecb+1)
    ci1 = numpy.zeros((ci0.shape[0], nb_ci1))

    entry_has_ap = (cre_index[:,:,0] == ap_id)
    addr_ci0 = numpy.any(entry_has_ap, axis=1)
    addr_ci1 = cre_index[entry_has_ap,2]
    sign = cre_index[entry_has_ap,3]
    # This sign prefactor accounts for interchange of operators with alpha and beta spins
    if neleca % 2 == 1:
        sign *= -1
    ci1[:,addr_ci1] = ci0[:,addr_ci0] * sign
    return ci1


def energy(h1e, eri, fcivec, norb, nelec, link_index=None):
    '''Compute the FCI electronic energy for given Hamiltonian and FCI vector.
    '''
    from pyscf.fci import direct_spin1
    h2e = direct_spin1.absorb_h1e(h1e, eri, norb, nelec, .5)
    ci1 = direct_spin1.contract_2e(h2e, fcivec, norb, nelec, link_index)
    return numpy.dot(fcivec.reshape(-1), ci1.reshape(-1))


def reorder(ci, nelec, orbidxa, orbidxb=None):
    '''Reorder the CI coefficients, to adapt the reordered orbitals (The relation
    of the reordered orbitals and original orbitals is  new = old[idx]).  Eg.

    The orbital ordering indices ``orbidx = [2,0,1]`` indicates the map
    old orbital  a b c  ->   new orbital  c a b.  The strings are reordered as<commit_msg>Update fci addons to handle negative particle numbers<commit_after>#!/usr/bin/env python

import sys
import copy
import numpy
from pyscf import lib
from pyscf.fci import cistring
from pyscf import symm

def large_ci(ci, norb, nelec, tol=.1, return_strs=True):
    '''Search for the largest CI coefficients
    '''
    neleca, nelecb = _unpack(nelec)
    addra, addrb = numpy.where(abs(ci) > tol)
    if addra.size == 0: # No large CI coefficient > tol
        addra, addrb = numpy.unravel_index(numpy.argmax(abs(ci)), ci.shape)
        addra = numpy.asarray([addra])
        addrb = numpy.asarray([addrb])
    strsa = cistring.addrs2str(norb, neleca, addra)
    strsb = cistring.addrs2str(norb, nelecb, addrb)
    if return_strs:
        strsa = [bin(x) for x in strsa]
        strsb = [bin(x) for x in strsb]
        return list(zip(ci[addra,addrb], strsa, strsb))
    else:
        occslsta = cistring._strs2occslst(strsa, norb)
        occslstb = cistring._strs2occslst(strsb, norb)
        return list(zip(ci[addra,addrb], occslsta, occslstb))

def initguess_triplet(norb, nelec, binstring):
    '''Generate a triplet initial guess for FCI solver
    '''
    neleca, nelecb = _unpack(nelec)
    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    addr = cistring.str2addr(norb, neleca, int(binstring,2))
    ci0 = numpy.zeros((na,nb))
    ci0[addr,0] = numpy.sqrt(.5)
    ci0[0,addr] =-numpy.sqrt(.5)
    return ci0

def symm_initguess(norb, nelec, orbsym, wfnsym=0, irrep_nelec=None):
    '''Generate CI wavefunction initial guess which has the given symmetry.

    Args:
        norb : int
            Number of orbitals.
        nelec : int or 2-item list
            Number of electrons, or 2-item list for (alpha, beta) electrons
        orbsym : list of int
            The irrep ID for each orbital.

    Kwags:
        wfnsym : int
            The irrep ID of target symmetry
        irrep_nelec : dict
            Freeze occupancy for certain irreps

    Returns:
        CI coefficients 2D array which has the target symmetry.
    '''
    neleca, nelecb = _unpack(nelec)
    orbsym = numpy.asarray(orbsym)
    if not isinstance(orbsym[0], numpy.number):
        raise RuntimeError('TODO: convert irrep symbol to irrep id')

    na = cistring.num_strings(norb, neleca)
    nb = cistring.num_strings(norb, nelecb)
    ci1 = numpy.zeros((na,nb))

########################
# pass 1: The fixed occs
    orbleft = numpy.ones(norb, dtype=bool)
    stra = numpy.zeros(norb, dtype=bool)
    strb = numpy.zeros(norb, dtype=bool)
    if irrep_nelec is not None:
        for k,n in irrep_nelec.items():
            orbleft[orbsym==k] = False
            if isinstance(n, (int, numpy.number)):
                idx = numpy.where(orbsym==k)[0][:n//2]
                stra[idx] = True
                strb[idx] = True
            else:
                na, nb = n
                stra[numpy.where(orbsym==k)[0][:na]] = True
                strb[numpy.where(orbsym==k)[0][:nb]] = True
                if (na-nb)%2:
                    wfnsym ^= k

    orbleft = numpy.where(orbleft)[0]
    neleca_left = neleca - stra.sum()
    nelecb_left = nelecb - strb.sum()
    spin = neleca_left - nelecb_left
    assert(neleca_left >= 0)
    assert(nelecb_left >= 0)
    assert(spin >= 0)

########################
# pass 2: search pattern
    def gen_str_iter(orb_list, nelec):
        if nelec == 1:
            for i in orb_list:
                yield [i]
        elif nelec >= len(orb_list):
            yield orb_list
        else:
            restorb = orb_list[1:]
            #yield from gen_str_iter(restorb, nelec)
            for x in gen_str_iter(restorb, nelec):
                yield x
            for x in gen_str_iter(restorb, nelec-1):
                yield [orb_list[0]] + x

# search for alpha and beta pattern which match to the required symmetry
    def query(target, nelec_atmost, spin, orbsym):
        norb = len(orbsym)
        for excite_level in range(1, nelec_atmost+1):
            for beta_only in gen_str_iter(range(norb), excite_level):
                alpha_allow = [i for i in range(norb) if i not in beta_only]
                alpha_orbsym = orbsym[alpha_allow]
                alpha_target = target
                for i in beta_only:
                    alpha_target ^= orbsym[i]
                alpha_only = symm.route(alpha_target, spin+excite_level, alpha_orbsym)
                if alpha_only:
                    alpha_only = [alpha_allow[i] for i in alpha_only]
                    return alpha_only, beta_only
        raise RuntimeError('No pattern found for wfn irrep %s over orbsym %s'
                           % (target, orbsym))

    if spin == 0:
        aonly = bonly = []
        if wfnsym != 0:
            aonly, bonly = query(wfnsym, neleca_left, spin, orbsym[orbleft])
    else:
        # 1. assume "nelecb_left" doubly occupied orbitals
        # search for alpha pattern which match to the required symmetry
        aonly, bonly = orbleft[symm.route(wfnsym, spin, orbsym[orbleft])], []
        # dcompose doubly occupied orbitals, search for alpha and beta pattern
        if len(aonly) != spin:
            aonly, bonly = query(wfnsym, neleca_left, spin, orbsym[orbleft])

    ndocc = neleca_left - len(aonly) # == nelecb_left - len(bonly)
    docc_allow = numpy.ones(len(orbleft), dtype=bool)
    docc_allow[aonly] = False
    docc_allow[bonly] = False
    docclst = orbleft[numpy.where(docc_allow)[0]][:ndocc]
    stra[docclst] = True
    strb[docclst] = True

    def find_addr_(stra, aonly, nelec):
        stra[orbleft[aonly]] = True
        return cistring.str2addr(norb, nelec, ('%i'*norb)%tuple(stra)[::-1])
    if bonly:
        if spin > 0:
            aonly, socc_only = aonly[:-spin], aonly[-spin:]
            stra[orbleft[socc_only]] = True
        stra1 = stra.copy()
        strb1 = strb.copy()

        addra = find_addr_(stra, aonly, neleca)
        addrb = find_addr_(strb, bonly, nelecb)
        addra1 = find_addr_(stra1, bonly, neleca)
        addrb1 = find_addr_(strb1, aonly, nelecb)
        ci1[addra,addrb] = ci1[addra1,addrb1] = numpy.sqrt(.5)
    else:
        addra = find_addr_(stra, aonly, neleca)
        addrb = find_addr_(strb, bonly, nelecb)
        ci1[addra,addrb] = 1

#    target = 0
#    for i,k in enumerate(stra):
#        if k:
#            target ^= orbsym[i]
#    for i,k in enumerate(strb):
#        if k:
#            target ^= orbsym[i]
#    print target
    return ci1


def _symmetrize_wfn(ci, strsa, strsb, orbsym, wfnsym=0):
    ci = ci.reshape(strsa.size,strsb.size)
    airreps = numpy.zeros(strsa.size, dtype=numpy.int32)
    birreps = numpy.zeros(strsb.size, dtype=numpy.int32)
    for i, ir in enumerate(orbsym):
        airreps[numpy.bitwise_and(strsa, 1<<i) > 0] ^= ir
        birreps[numpy.bitwise_and(strsb, 1<<i) > 0] ^= ir
    mask = (airreps.reshape(-1,1) ^ birreps) == wfnsym
    ci1 = numpy.zeros_like(ci)
    ci1[mask] = ci[mask]
    ci1 *= 1/numpy.linalg.norm(ci1)
    return ci1
def symmetrize_wfn(ci, norb, nelec, orbsym, wfnsym=0):
    '''Symmetrize the CI wavefunction by zeroing out the determinants which
    do not have the right symmetry.

    Args:
        ci : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        nelec : int or 2-item list
            Number of electrons, or 2-item list for (alpha, beta) electrons
        orbsym : list of int
            The irrep ID for each orbital.

    Kwags:
        wfnsym : int
            The irrep ID of target symmetry

    Returns:
        2D array which is the symmetrized CI coefficients
    '''
    neleca, nelecb = _unpack(nelec)
    strsa = numpy.asarray(cistring.gen_strings4orblist(range(norb), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(norb), nelecb))
    return _symmetrize_wfn(ci, strsa, strsb, orbsym, wfnsym)

def _guess_wfnsym(ci, strsa, strsb, orbsym):
    na = len(strsa)
    nb = len(strsb)
    if isinstance(ci, numpy.ndarray) and ci.ndim <= 2:
        assert(ci.size == na*nb)
        idx = numpy.argmax(ci)
    else:
        assert(ci[0].size == na*nb)
        idx = ci[0].argmax()
    stra = strsa[idx // nb]
    strb = strsb[idx % nb ]

    airrep = 0
    birrep = 0
    for i, ir in enumerate(orbsym):
        if (stra & (1<<i)):
            airrep ^= ir
        if (strb & (1<<i)):
            birrep ^= ir
    return airrep ^ birrep
def guess_wfnsym(ci, norb, nelec, orbsym):
    '''Guess the wavefunction symmetry based on the non-zero elements in the
    given CI coefficients.

    Args:
        ci : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        nelec : int or 2-item list
            Number of electrons, or 2-item list for (alpha, beta) electrons
        orbsym : list of int
            The irrep ID for each orbital.

    Returns:
        Irrep ID
    '''
    neleca, nelecb = _unpack(nelec)
    strsa = numpy.asarray(cistring.gen_strings4orblist(range(norb), neleca))
    strsb = numpy.asarray(cistring.gen_strings4orblist(range(norb), nelecb))
    return _guess_wfnsym(ci, strsa, strsb, orbsym)


def des_a(ci0, norb, neleca_nelecb, ap_id):
    r'''Construct (N-1)-electron wavefunction by removing an alpha electron from
    the N-electron wavefunction.

    ... math::

        |N-1\rangle = \hat{a}_p |N\rangle

    Args:
        ci0 : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        (neleca,nelecb) : (int,int)
            Number of (alpha, beta) electrons of the input CI function
        ap_id : int
            Orbital index (0-based), for the annihilation operator

    Returns:
        2D array, row for alpha strings and column for beta strings.  Note it
        has different number of rows to the input CI coefficients
    '''
    neleca, nelecb = neleca_nelecb
    if neleca <= 0:
        return numpy.zeros_like(ci0)
    if ci0.ndim == 1:
        ci0 = ci0.reshape(cistring.num_strings(norb, neleca),
                          cistring.num_strings(norb, nelecb))
    des_index = cistring.gen_des_str_index(range(norb), neleca)
    na_ci1 = cistring.num_strings(norb, neleca-1)
    ci1 = numpy.zeros((na_ci1, ci0.shape[1]))

    entry_has_ap = (des_index[:,:,1] == ap_id)
    addr_ci0 = numpy.any(entry_has_ap, axis=1)
    addr_ci1 = des_index[entry_has_ap,2]
    sign = des_index[entry_has_ap,3]
    #print(addr_ci0)
    #print(addr_ci1)
    ci1[addr_ci1] = sign.reshape(-1,1) * ci0[addr_ci0]
    return ci1

def des_b(ci0, norb, neleca_nelecb, ap_id):
    r'''Construct (N-1)-electron wavefunction by removing a beta electron from
    N-electron wavefunction.

    Args:
        ci0 : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        (neleca,nelecb) : (int,int)
            Number of (alpha, beta) electrons of the input CI function
        ap_id : int
            Orbital index (0-based), for the annihilation operator

    Returns:
        2D array, row for alpha strings and column for beta strings. Note it
        has different number of columns to the input CI coefficients.
    '''
    neleca, nelecb = neleca_nelecb
    if nelecb <= 0:
        return numpy.zeros_like(ci0)
    if ci0.ndim == 1:
        ci0 = ci0.reshape(cistring.num_strings(norb, neleca),
                          cistring.num_strings(norb, nelecb))
    des_index = cistring.gen_des_str_index(range(norb), nelecb)
    nb_ci1 = cistring.num_strings(norb, nelecb-1)
    ci1 = numpy.zeros((ci0.shape[0], nb_ci1))

    entry_has_ap = (des_index[:,:,1] == ap_id)
    addr_ci0 = numpy.any(entry_has_ap, axis=1)
    addr_ci1 = des_index[entry_has_ap,2]
    sign = des_index[entry_has_ap,3]
    # This sign prefactor accounts for interchange of operators with alpha and beta spins
    if neleca % 2 == 1:
        sign *= -1
    ci1[:,addr_ci1] = ci0[:,addr_ci0] * sign
    return ci1

def cre_a(ci0, norb, neleca_nelecb, ap_id):
    r'''Construct (N+1)-electron wavefunction by adding an alpha electron in
    the N-electron wavefunction.

    ... math::

        |N+1\rangle = \hat{a}^+_p |N\rangle

    Args:
        ci0 : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        (neleca,nelecb) : (int,int)
            Number of (alpha, beta) electrons of the input CI function
        ap_id : int
            Orbital index (0-based), for the creation operator

    Returns:
        2D array, row for alpha strings and column for beta strings. Note it
        has different number of rows to the input CI coefficients.
    '''
    neleca, nelecb = neleca_nelecb
    if neleca >= norb:
        return numpy.zeros_like(ci0)
    if ci0.ndim == 1:
        ci0 = ci0.reshape(cistring.num_strings(norb, neleca),
                          cistring.num_strings(norb, nelecb))
    cre_index = cistring.gen_cre_str_index(range(norb), neleca)
    na_ci1 = cistring.num_strings(norb, neleca+1)
    ci1 = numpy.zeros((na_ci1, ci0.shape[1]))

    entry_has_ap = (cre_index[:,:,0] == ap_id)
    addr_ci0 = numpy.any(entry_has_ap, axis=1)
    addr_ci1 = cre_index[entry_has_ap,2]
    sign = cre_index[entry_has_ap,3]
    ci1[addr_ci1] = sign.reshape(-1,1) * ci0[addr_ci0]
    return ci1

# construct (N+1)-electron wavefunction by adding a beta electron to
# N-electron wavefunction:
def cre_b(ci0, norb, neleca_nelecb, ap_id):
    r'''Construct (N+1)-electron wavefunction by adding a beta electron in
    the N-electron wavefunction.

    Args:
        ci0 : 2D array
            CI coefficients, row for alpha strings and column for beta strings.
        norb : int
            Number of orbitals.
        (neleca,nelecb) : (int,int)
            Number of (alpha, beta) electrons of the input CI function
        ap_id : int
            Orbital index (0-based), for the creation operator

    Returns:
        2D array, row for alpha strings and column for beta strings. Note it
        has different number of columns to the input CI coefficients.
    '''
    neleca, nelecb = neleca_nelecb
    if nelecb >= norb:
        return numpy.zeros_like(ci0)
    if ci0.ndim == 1:
        ci0 = ci0.reshape(cistring.num_strings(norb, neleca),
                          cistring.num_strings(norb, nelecb))
    cre_index = cistring.gen_cre_str_index(range(norb), nelecb)
    nb_ci1 = cistring.num_strings(norb, nelecb+1)
    ci1 = numpy.zeros((ci0.shape[0], nb_ci1))

    entry_has_ap = (cre_index[:,:,0] == ap_id)
    addr_ci0 = numpy.any(entry_has_ap, axis=1)
    addr_ci1 = cre_index[entry_has_ap,2]
    sign = cre_index[entry_has_ap,3]
    # This sign prefactor accounts for interchange of operators with alpha and beta spins
    if neleca % 2 == 1:
        sign *= -1
    ci1[:,addr_ci1] = ci0[:,addr_ci0] * sign
    return ci1


def energy(h1e, eri, fcivec, norb, nelec, link_index=None):
    '''Compute the FCI electronic energy for given Hamiltonian and FCI vector.
    '''
    from pyscf.fci import direct_spin1
    h2e = direct_spin1.absorb_h1e(h1e, eri, norb, nelec, .5)
    ci1 = direct_spin1.contract_2e(h2e, fcivec, norb, nelec, link_index)
    return numpy.dot(fcivec.reshape(-1), ci1.reshape(-1))


def reorder(ci, nelec, orbidxa, orbidxb=None):
    '''Reorder the CI coefficients, to adapt the reordered orbitals (The relation
    of the reordered orbitals and original orbitals is  new = old[idx]).  Eg.

    The orbital ordering indices ``orbidx = [2,0,1]`` indicates the map
    old orbital  a b c  ->   new orbital  c a b.  The strings are reordered as
{'fraction_non_alphanumeric': 0.10662951940733237, 'fraction_numerical': 0.01563984043563604, 'mean_word_length': 2.4371055495103375, 'pattern_counts': {'":': 0, '<': 25, '<?xml version=': 0, '>': 27, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1832943', 'n_tokens_mistral': 11508, 'n_tokens_neox': 11101, 'n_words': 2938}