Source code for wavefunction_analysis.embedding.mol_lo_tools

import scipy, numpy
import pyscf

[docs] def lo_weight_on_ao(mol : pyscf.gto.Mole = None, coeff_ao_lo : numpy.ndarray = None, ovlp_ao : numpy.ndarray = None): """A function to calculate the weight of local orbitals on atomic orbitals. Parameters: mol : pyscf.gto.Mole The molecule object. coeff_ao_lo : numpy.ndarray The coefficient matrix to transform atomic orbitals to local orbitals. ovlp_ao : numpy.ndarray The overlap matrix of atomic orbitals. Returns: w2_ao_lo : numpy.ndarray The weight of local orbitals on atomic orbitals. """ if ovlp_ao is None: assert mol is not None ovlp_ao = mol.intor_symmetric('int1e_ovlp') assert ovlp_ao is not None nao = ovlp_ao.shape[0] nlo = coeff_ao_lo.shape[1] assert coeff_ao_lo.shape == (nao, nlo) assert ovlp_ao.shape == (nao, nao) ovlp_ao_sqrt = scipy.linalg.sqrtm(ovlp_ao) w_ao_lo = numpy.dot(ovlp_ao_sqrt, coeff_ao_lo) w2_ao_lo = numpy.einsum('np,np->np', w_ao_lo, w_ao_lo) return w2_ao_lo
[docs] def lo_weight_on_atom(mol : pyscf.gto.Mole, coeff_ao_lo : numpy.ndarray = None, ovlp_ao : numpy.ndarray = None): """A function to calculate the weight of local orbitals on atomic orbitals. Parameters: mol : pyscf.gto.Mole The molecule object. coeff_ao_lo : numpy.ndarray The coefficient matrix to transform atomic orbitals to local orbitals. ovlp_ao : numpy.ndarray The overlap matrix of atomic orbitals. Returns: w2_atm_lo : numpy.ndarray The weight of local orbitals on each atom. """ if ovlp_ao is None: ovlp_ao = mol.intor_symmetric('int1e_ovlp') assert ovlp_ao is not None natm = mol.natm nao = mol.nao nlo = coeff_ao_lo.shape[1] assert coeff_ao_lo.shape == (nao, nlo) assert ovlp_ao.shape == (nao, nao) w2_ao_lo = lo_weight_on_ao(mol, coeff_ao_lo, ovlp_ao) w2_atm_lo = numpy.zeros((natm, nlo)) ao_slice_by_atom = mol.aoslice_by_atom() for iatm in range(natm): ao_slice_iatm = ao_slice_by_atom[iatm, 2:4] w2_atm_lo[iatm] = numpy.sum(w2_ao_lo[ao_slice_iatm[0]:ao_slice_iatm[1]], axis=0) return w2_atm_lo
[docs] def lo_weight_on_frag(frag_atms_list : list, mol : pyscf.gto.Mole, coeff_ao_lo : numpy.ndarray = None, ovlp_ao : numpy.ndarray = None): """A function to calculate the weight of local orbitals on fragments. Parameters: frag_atms_list : list A list of atoms in each fragment. mol : pyscf.gto.Mole The molecule object. coeff_ao_lo : numpy.ndarray The coefficient matrix to transform atomic orbitals to local orbitals. ovlp_ao : numpy.ndarray The overlap matrix of atomic orbitals. Returns: w2_frag_lo : numpy.ndarray The weight of local orbitals on each fragment. """ if ovlp_ao is None: ovlp_ao = mol.intor_symmetric('int1e_ovlp') assert ovlp_ao is not None nfrag = len(frag_atms_list) natm = mol.natm nao = mol.nao nlo = coeff_ao_lo.shape[1] w2_atm_lo = lo_weight_on_atom(mol, coeff_ao_lo, ovlp_ao) w2_frag_lo = numpy.zeros((nfrag, nlo)) assert coeff_ao_lo.shape == (nao, nlo) assert ovlp_ao.shape == (nao, nao) assert w2_atm_lo.shape == (natm, nlo) for ifrag, frag_atms in enumerate(frag_atms_list): w2_frag_lo[ifrag] = numpy.sum(w2_atm_lo[frag_atms], axis=0) return w2_frag_lo
[docs] def partition_lo_to_atms(mol : pyscf.gto.Mole = None, coeff_ao_lo : numpy.ndarray = None, w2_atm_lo : numpy.ndarray = None, min_weight : float = 0.9,): """A function to partition local orbitals to atoms. Parameters: mol : pyscf.gto.Mole The molecule object. coeff_ao_lo : numpy.ndarray The coefficient matrix to transform atomic orbitals to local orbitals. w2_atm_lo : numpy.ndarray The weight of local orbitals on each atom. min_weight : float The minimum weight of local orbitals on each atom. If the max weight is smaller than min_weight, the local orbital will not be assigned to the atom. Returns: lo_idx_on_atm_list : list The list of local orbital indices on each atom. """ if w2_atm_lo is None: w2_atm_lo = lo_weight_on_atom(mol, coeff_ao_lo, None) natm = mol.natm nlo = coeff_ao_lo.shape[1] assert w2_atm_lo.shape == (natm, nlo) lo_idx_on_atm_list = [[] for iatm in range(natm)] for lo in range(nlo): iatm = numpy.argmax(w2_atm_lo[:, lo]) if w2_atm_lo[iatm, lo] > min_weight: lo_idx_on_atm_list[iatm].append(lo) #print("LO %3d is assigned to atom %3d with weight % 6.4f" % (lo, iatm, w2_atm_lo[iatm, lo])) else: print('Warning: The weight of local orbital %d on atom %d is % 6.4f, which is smaller than the tolerance % 6.4f!' % (lo, iatm, w2_atm_lo[iatm, lo], tol)) return lo_idx_on_atm_list
[docs] def partition_lo_to_frags(frag_atms_list : list, mol : pyscf.gto.Mole = None, coeff_ao_lo : numpy.ndarray = None, w2_frag_lo : numpy.ndarray = None, min_weight : float = 0.9,): """A function to partition local orbitals to fragments. Parameters: frag_atms_list : list The list of atoms in each fragment. mol : pyscf.gto.Mole The molecule object. coeff_ao_lo : numpy.ndarray The coefficient matrix to transform atomic orbitals to local orbitals. w2_frag_lo : numpy.ndarray The weight of local orbitals on each atom. min_weight : float The minimum weight of local orbitals on each atom. If the max weight is smaller than min_weight, the local orbital will not be assigned to the atom. Returns: lo_idx_on_frag_list : list The list of local orbital indices on each fragment. """ if w2_frag_lo is None: w2_frag_lo = lo_weight_on_frag(frag_atms_list, mol, coeff_ao_lo, None) nfrag = len(frag_atms_list) nlo = coeff_ao_lo.shape[1] assert w2_frag_lo.shape == (nfrag, nlo) lo_idx_on_frag_list = [[] for ifrag in range(nfrag)] for lo in range(nlo): ifrag = numpy.argmax(w2_frag_lo[:, lo]) if w2_frag_lo[ifrag, lo] > min_weight: lo_idx_on_frag_list[ifrag].append(lo) #print("LO %3d is assigned to fragment %3d, weight = %6.4f" % (lo, ifrag, w2_frag_lo[ifrag, lo])) else: print('Warning: The weight of local orbital %d on fragment %d is % 6.4f, which is smaller than the tolerance % 6.4f!' % (lo, ifrag, w2_frag_lo[ifrag, lo], min_weight)) return lo_idx_on_frag_list
partition_lo_to_imps = partition_lo_to_frags