<commit_before>#
# Author: <NAME> <<EMAIL>>
#

'''
Generate symmetry adapted basis
'''

from functools import reduce
import numpy
import pyscf.lib
import pyscf.gto
from pyscf.symm import geom
from pyscf.symm import param

OP_PARITY_ODD = {
    'E'  : (0, 0, 0),
    'C2x': (0, 1, 1),
    'C2y': (1, 0, 1),
    'C2z': (1, 1, 0),
    'i'  : (1, 1, 1),
    'sx' : (1, 0, 0),
    'sy' : (0, 1, 0),
    'sz' : (0, 0, 1),
}

def tot_parity_odd(op, l, m):
    if op == 'E':
        return 0
    else:
        ox,oy,oz = OP_PARITY_ODD[op]
        gx,gy,gz = param.SPHERIC_GTO_PARITY_ODD[l][l+m]
        return (ox and gx)^(oy and gy)^(oz and gz)

def symm_adapted_basis(gpname, eql_atom_ids, atoms, basis_tab):
    if gpname in ('Dooh', 'Coov'):
        return linearmole_symm_adapted_basis(gpname, eql_atom_ids, atoms, basis_tab)
    opdic = geom.symm_ops(gpname)
    ops = [opdic[op] for op in param.OPERATOR_TABLE[gpname]]
    chartab = param.CHARACTER_TABLE[gpname]
    nirrep = chartab.__len__()
    nfn, basoff = _basis_offset_for_atoms(atoms, basis_tab)
    coords = numpy.around([a[1] for a in atoms], decimals=geom.PLACE-1)

    sodic = {}
    def add_so(ir, c):
        if ir in sodic:
            sodic[ir].append(c)
        else:
            sodic[ir] = [c]
    for atom_ids in eql_atom_ids:
        at0 = atoms[atom_ids[0]]
        symb = pyscf.gto.mole._symbol(at0[0])
        op_coords = numpy.around([numpy.dot(at0[1], op) for op in ops],
                                 decimals=geom.PLACE-1)
# num atoms in atom_ids smaller than num ops?
## the coord-ids generated from the ops sequence
#        id_coords = numpy.argsort(geom.argsort_coords(op_coords))
## which atom_id that each coord-id corresponds to
#        from_atom_id = numpy.argsort(geom.argsort_coords(coords[atom_ids]))
#        op_relate_atoms = atom_ids[from_atom_id[id_coords]]
        coords0 = coords[atom_ids]
        idx = []
        for c in op_coords:
            idx.append(numpy.argwhere(numpy.sum(abs(coords0-c),axis=1)
                                      <geom.GEOM_THRESHOLD)[0,0])
        op_relate_atoms = numpy.array(atom_ids)[idx]

        ib = 0
        if symb in basis_tab:
            bas0 = basis_tab[symb]
        else:
            bas0 = basis_tab[pyscf.gto.mole._rm_digit(symb)]
        for b in bas0:
            angl = b[0]
            for i in range(_num_contract(b)):
                for m in range(-angl,angl+1):
                    sign = [-1 if tot_parity_odd(op,angl,m) else 1
                            for op in param.OPERATOR_TABLE[gpname]]
                    c = numpy.zeros((nirrep,nfn))
                    for op_id, atm_id in enumerate(op_relate_atoms):<commit_msg>Fix symm.basis<commit_after>#
# Author: <NAME> <<EMAIL>>
#

'''
Generate symmetry adapted basis
'''

from functools import reduce
import numpy
import pyscf.lib
import pyscf.gto
from pyscf.symm import geom
from pyscf.symm import param

OP_PARITY_ODD = {
    'E'  : (0, 0, 0),
    'C2x': (0, 1, 1),
    'C2y': (1, 0, 1),
    'C2z': (1, 1, 0),
    'i'  : (1, 1, 1),
    'sx' : (1, 0, 0),
    'sy' : (0, 1, 0),
    'sz' : (0, 0, 1),
}

def tot_parity_odd(op, l, m):
    if op == 'E':
        return 0
    else:
        ox,oy,oz = OP_PARITY_ODD[op]
        gx,gy,gz = param.SPHERIC_GTO_PARITY_ODD[l][l+m]
        return (ox and gx)^(oy and gy)^(oz and gz)

def symm_adapted_basis(gpname, eql_atom_ids, atoms, basis_tab):
    if gpname in ('Dooh', 'Coov'):
        return linearmole_symm_adapted_basis(gpname, eql_atom_ids, atoms, basis_tab)
    opdic = geom.symm_ops(gpname)
    ops = [opdic[op] for op in param.OPERATOR_TABLE[gpname]]
    chartab = param.CHARACTER_TABLE[gpname]
    nirrep = chartab.__len__()
    nfn, basoff = _basis_offset_for_atoms(atoms, basis_tab)
    coords = numpy.around([a[1] for a in atoms], decimals=geom.PLACE-1)

    sodic = {}
    def add_so(ir, c):
        if ir in sodic:
            sodic[ir].append(c)
        else:
            sodic[ir] = [c]
    for atom_ids in eql_atom_ids:
        at0 = atoms[atom_ids[0]]
        symb = pyscf.gto.mole._symbol(at0[0])
        op_coords = numpy.around([numpy.dot(at0[1], op) for op in ops],
                                 decimals=geom.PLACE-1)
# num atoms in atom_ids smaller than num ops?
## the coord-ids generated from the ops sequence
#        id_coords = numpy.argsort(geom.argsort_coords(op_coords))
## which atom_id that each coord-id corresponds to
#        from_atom_id = numpy.argsort(geom.argsort_coords(coords[atom_ids]))
#        op_relate_atoms = atom_ids[from_atom_id[id_coords]]
        coords0 = coords[atom_ids]
        idx = []
        for c in op_coords:
            idx.append(numpy.argwhere(numpy.sum(abs(coords0-c),axis=1)
                                      <geom.TOLERANCE)[0,0])
        op_relate_atoms = numpy.array(atom_ids)[idx]

        ib = 0
        if symb in basis_tab:
            bas0 = basis_tab[symb]
        else:
            bas0 = basis_tab[pyscf.gto.mole._rm_digit(symb)]
        for b in bas0:
            angl = b[0]
            for i in range(_num_contract(b)):
                for m in range(-angl,angl+1):
                    sign = [-1 if tot_parity_odd(op,angl,m) else 1
                            for op in param.OPERATOR_TABLE[gpname]]
                    c = numpy.zeros((nirrep,nfn))
                    for op_id, atm_id in enumerate(op_relate_atoms):
{'fraction_non_alphanumeric': 0.12807606263982102, 'fraction_numerical': 0.018642803877703208, 'mean_word_length': 2.245614035087719, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 9, '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': '1945110', 'n_tokens_mistral': 2139, 'n_tokens_neox': 2058, 'n_words': 456}