"""
Contains the physical baseline model (e.g. CASSCF).
This module provides functionality to calculate a basline solution using pySCF and functions to
"""
import logging
import jax.numpy as jnp
import pyscf
import pyscf.ci
import pyscf.mcscf
import numpy as np
from deeperwin.configuration import PhysicalConfig, CASSCFConfig
from deeperwin.utils import get_el_ion_distance_matrix
from scipy.optimize import curve_fit
logger = logging.getLogger("dpe")
#################################################################################
############################ Orbital functions ##################################
#################################################################################
[docs]def eval_gaussian_orbital(el_ion_diff, el_ion_dist, ind_nuc, alpha, weights, angular_momenta):
"""
Evaluate a single orbital on multiple points in space r.
Each orbital is evaluated at N points in 3D space.
Each orbital consist of k gaussians centered around a nucleus at position R.
Args:
el_ion_diff: Position where orbital is evaluated, relative to nucleus, shape [3]
el_ion_dist: norm of difference, i.e. distance between electron position and the nucleus; shape []
alpha: Gaussian decay coefficients [k]
weights: Linear expansion coeffs of orbitals [k]
angular_momenta: Integers describing the power of each cartesian coordinate, shape [3]
Returns:
Array of shape [N]
"""
r_sqr = jnp.expand_dims(el_ion_dist ** 2, axis=-1) # shape [1]
pre_fac = jnp.prod(el_ion_diff ** angular_momenta, axis=-1) # shape []
phi_gto = pre_fac * jnp.sum(weights * jnp.exp(-alpha * r_sqr), axis=-1)
return phi_gto
[docs]def eval_atomic_orbitals(el_ion_diff, el_ion_dist, orbitals, ao_cusp_params=None):
"""
Args:
el_ion_diff: shape [N_batch x n_el x N_ion x 3]
el_ion_dist: shape [N_batch x n_el x N_ion]
orbitals:
Returns:
"""
outputs = []
for i,ao in enumerate(orbitals):
ind_nuc = ao[0]
diff = el_ion_diff[..., ind_nuc, :]
dist = el_ion_dist[..., ind_nuc]
gto = eval_gaussian_orbital(diff, dist, *ao)
if (ao_cusp_params is not None) and (ao_cusp_params[i] is not None):
r_c = ao_cusp_params[i][0]
sto = _eval_cusp_atomic_orbital(dist, *ao_cusp_params[i])
weight_gto = 0.5 * (jnp.sign(dist - r_c) + 1)
weight_sto = 1.0 - weight_gto
gto = weight_gto * gto + weight_sto * sto
outputs.append(gto)
return jnp.stack(outputs, axis=-1)
def _eval_cusp_molecular_orbital(dist, r_c, offset, sign, poly):
dist = jnp.minimum(dist, r_c)
n = jnp.arange(5)
r_n = dist[..., jnp.newaxis, jnp.newaxis] ** n
p_r = jnp.sum(r_n * poly, axis=-1)
psi = offset + sign*jnp.exp(p_r)
return psi
def _eval_cusp_atomic_orbital(dist, r_c, offset, sign, poly):
dist = jnp.minimum(dist, r_c)
n = jnp.arange(5)
r_n = dist[..., jnp.newaxis] ** n
p_r = jnp.dot(r_n, poly)
psi = offset + sign*jnp.exp(p_r)
return psi
[docs]def evaluate_molecular_orbitals(el_ion_diff, el_ion_dist, atomic_orbitals, mo_coeff, ao_cusp_params=None, mo_cusp_params=None):
aos = eval_atomic_orbitals(el_ion_diff, el_ion_dist, atomic_orbitals, ao_cusp_params)
mos = aos @ mo_coeff
if mo_cusp_params is not None:
r_c, offset, sign, poly, coeff_1s = mo_cusp_params
sto_mask = jnp.heaviside(r_c - el_ion_dist, 0.0)[..., jnp.newaxis]
sto = _eval_cusp_molecular_orbital(el_ion_dist, r_c, offset, sign, poly) - jnp.sum(aos[..., np.newaxis, :, np.newaxis] * coeff_1s, axis=-2)
sto = jnp.sum(sto*sto_mask, axis=-2) # sum over all ions, masking only the once that are within r_c of an electron
mos = mos + sto # subtract GTO-part of 1s orbitals and replace by STO-part
return mos
#################################################################################
####################### Compile-time helper functions ###########################
#################################################################################
[docs]def build_pyscf_molecule(R, Z, charge=0, spin=0, basis_set='6-311G'):
"""
Args:
R: ion positions, shape [Nx3]
Z: ion charges, shape [N]
charge: integer
spin: integer; number of excess spin-up electrons
basis_set (str): basis set identifier
Returns:
pyscf.Molecule
"""
molecule = pyscf.gto.Mole()
molecule.atom = [[Z_, tuple(R_)] for R_, Z_ in zip(R, Z)]
molecule.unit = "bohr"
molecule.basis = basis_set
molecule.cart = True
molecule.spin = spin # 2*n_up - n_down
molecule.charge = charge
molecule.output = "/dev/null"
molecule.verbose = 0 # suppress output to console
molecule.max_memory = 10e3 # maximum memory in megabytes (i.e. 10e3 = 10GB)
molecule.build()
return molecule
def _get_gto_normalization_factor(alpha, angular_momenta):
l_tot = np.sum(angular_momenta)
fac_alpha = (2 * alpha / np.pi) ** (3 / 4) * (8 * alpha) ** (l_tot / 2)
fac = np.array([np.math.factorial(x) for x in angular_momenta])
fac2 = np.array([np.math.factorial(2 * x) for x in angular_momenta])
fac_exponent = np.sqrt(np.prod(fac) / np.prod(fac2))
factor = fac_alpha * fac_exponent
return factor
def _get_atomic_orbital_basis_functions(molecule):
"""
Args:
molecule: pyscf molecule object
Returns:
orbitals: tuple of tuple, each containing index of nucleus, nuclear coordinate, exponential decay constant alpha, weights and angular momentum exponents
"""
ao_labels = molecule.ao_labels(None)
n_basis = len(ao_labels)
ind_basis = 0
atomic_orbitals = []
for ind_nuc, (element, atom_pos) in enumerate(molecule._atom):
for gto_data in molecule._basis[element]:
l = gto_data[0]
gto_data = np.array(gto_data[1:])
alpha = gto_data[:, 0]
weights = gto_data[:, 1:]
for ind_contraction in range(weights.shape[1]):
n_orientations = [1, 3, 6][l]
for m in range(n_orientations): # 1,3,6, ... = number of orbitals per angular momentum
shape_string = ao_labels[ind_basis][3] # string of the form 'xxy' or 'zz'
angular_momenta = np.array([shape_string.count(x) for x in ['x', 'y', 'z']], dtype=int)
normalization = _get_gto_normalization_factor(alpha, angular_momenta)
atomic_orbitals.append([np.array(ind_nuc, int),
alpha,
weights[:, ind_contraction] * normalization,
angular_momenta])
ind_basis += 1
assert len(
atomic_orbitals) == n_basis, "Could not properly construct basis functions. You probably tried to use a valence-only basis (e.g. cc-pVDZ) for an all-electron calculation." # noqa
return atomic_orbitals
[docs]def build_pyscf_molecule_from_physical_config(physical_config: PhysicalConfig, basis_set):
R, Z = np.array(physical_config.R), np.array(physical_config.Z)
charge = sum(Z) - physical_config.n_electrons
spin = 2 * physical_config.n_up - physical_config.n_electrons
return build_pyscf_molecule(R, Z, charge, spin, basis_set)
[docs]def get_hartree_fock_solution(physical_config: PhysicalConfig, basis_set):
molecule = build_pyscf_molecule_from_physical_config(physical_config, basis_set)
atomic_orbitals = _get_atomic_orbital_basis_functions(molecule)
hf = pyscf.scf.HF(molecule)
hf.verbose = 0 # suppress output to console
hf.kernel()
return atomic_orbitals, hf
[docs]def get_baseline_solution(physical_config: PhysicalConfig, casscf_config: CASSCFConfig, n_dets: int):
n_el, n_up, R, Z = physical_config.get_basic_params()
n_dn = n_el - n_up
atomic_orbitals, hf = get_hartree_fock_solution(physical_config, casscf_config.basis_set)
if n_dets == 1 or casscf_config.only_hf:
# Only 1 determinant => no MRCI/CASSCF calculation required
mo_coeff = hf.mo_coeff
if len(mo_coeff) != 2:
mo_coeff = [mo_coeff, mo_coeff]
ci_weights = np.ones(1) if not casscf_config.only_hf else np.ones(32)
ind_orbitals = [np.arange(n_up)[None, :], np.arange(n_dn)[None, :]] if not casscf_config.only_hf else\
[np.tile(np.arange(n_up), n_dets).reshape((-1, n_up)), np.tile(np.arange(n_dn), n_dets).reshape((-1, n_dn))]
E_hf, E_casscf = hf.e_tot, np.nan
else:
# Run CASSCF to get excited determinants
casscf = pyscf.mcscf.UCASSCF(hf, physical_config.n_cas_orbitals, physical_config.n_cas_electrons)
casscf.kernel()
E_hf, E_casscf = hf.e_tot, casscf.e_tot
mo_coeff = list(casscf.mo_coeff) # tuple of spin_up, spin_down
ind_orbitals = _get_orbital_indices(casscf)
ci_weights = casscf.ci.flatten()
logger.debug(f"Total nuber of CASSCF-determinants before truncation: {len(ci_weights)}")
if n_dets < len(ci_weights):
ind_largest = np.argsort(np.abs(ci_weights))[::-1][:n_dets]
share_captured = np.sum(ci_weights[ind_largest]**2) / np.sum(ci_weights**2)
logger.debug(f"Share of CI-weights captured by {n_dets} dets: {share_captured:.3%}")
ci_weights = ci_weights[ind_largest]
ci_weights = jnp.array(ci_weights / np.sum(ci_weights ** 2))
ind_orbitals = ind_orbitals[0][ind_largest], ind_orbitals[1][ind_largest]
# Calculate cusp-correction-parameters for molecular orbitals
ao_cusp_params = None
mo_cusp_params = [None, None]
if casscf_config.cusps and (casscf_config.cusps.cusp_type == "mo"):
mo_cusp_params = [calculate_molecular_orbital_cusp_params(atomic_orbitals, mo_coeff[i], R, Z, casscf_config.cusps.r_cusp_el_ion_scale) for i in range(2)]
elif casscf_config.cusps and (casscf_config.cusps.cusp_type == 'ao'):
ao_cusp_params = [_calculate_ao_cusp_params(*ao, physical_config.Z[ao[0]]) for ao in atomic_orbitals]
return (atomic_orbitals, ao_cusp_params, mo_cusp_params, mo_coeff, ind_orbitals, ci_weights), (E_hf, E_casscf)
def _get_effective_charge(Z: int, n: int, s_lower_shell=0.85, s_same_shell=0.35):
"""
Calculates the approximate effective charge for an electron, using Slater's Rule of shielding.
Args:
Z: Nuclear charge of an atom
n: Principal quantum number of the elctron for which the effective charge is being calculated
Returns:
Z_eff (float): effective charge
"""
shielding = 0
for n_shell in range(1, n+1):
n_electrons_in_shell = 2 * n_shell ** 2
if n_shell == n:
n_el_in_lower_shells = sum([2*k**2 for k in range(1,n_shell)])
n_electrons_in_shell = min(Z - n_el_in_lower_shells, n_electrons_in_shell) - 1
if n_shell == n:
shielding += n_electrons_in_shell * s_same_shell
elif n_shell == (n-1):
shielding += n_electrons_in_shell * s_lower_shell
else:
shielding += n_electrons_in_shell
return max(Z - shielding, 1)
def _get_atomic_orbital_envelope_exponents(physical_config: PhysicalConfig, basis_set):
molecule = build_pyscf_molecule_from_physical_config(physical_config, basis_set)
atomic_orbitals = _get_atomic_orbital_basis_functions(molecule)
Z_ions = np.array(physical_config.Z)
n_ions = len(Z_ions)
n_ao = []
for ind_nuc in range(n_ions):
ao_nuc = [ao for ao in atomic_orbitals if ao[0] == ind_nuc]
nr_of_s_orbitals = len([ao for ao in ao_nuc if sum(ao[-1]) == 0])
nr_of_p_orbitals = len([ao for ao in ao_nuc if sum(ao[-1]) == 1])
nr_of_d_orbitals = len([ao for ao in ao_nuc if sum(ao[-1]) == 2])
assert (nr_of_s_orbitals + nr_of_p_orbitals + nr_of_d_orbitals) == len(ao_nuc), "Initialization currently does support d-orbitals"
n_ao.append(np.arange(1, nr_of_s_orbitals + 1))
n_ao.append(np.repeat(np.arange(2, nr_of_p_orbitals // 3 + 2), 3))
n_ao.append(np.repeat(np.arange(3, nr_of_d_orbitals // 6 + 3), 6))
n_ao = np.concatenate(n_ao)
Z_ao = np.array([Z_ions[ao[0]] for ao in atomic_orbitals])
Z_ao = np.array([_get_effective_charge(Z, n) for Z, n in zip(Z_ao, n_ao)])
alpha_ao = Z_ao / n_ao
alpha_ao = alpha_ao * 0.7 # rescale to better fit empirical values
return alpha_ao
[docs]def get_envelope_exponents_from_atomic_orbitals(physical_config: PhysicalConfig, basis_set='6-31G', pad_full_det=False):
"""
Calculate an initial guess for the envelope parameters using a baseline calculation.
Envelopes have the form coeff_I exp(-alpha_I * r).
Returns:
coeff_values: Tuple of len 2; for each spin has shape [n_ions x n_orb]
alpha_values: Tuple of len 2; for each spin has shape [n_ions x n_orb]
"""
n_up, n_dn, n_ions = physical_config.n_up, physical_config.n_dn, physical_config.n_ions
(atomic_orbitals, _, _, mo_coeff, _, _), _ = get_baseline_solution(physical_config,
CASSCFConfig(basis_set=basis_set, cusps=None),
n_dets=1)
mo_coeff = mo_coeff[0][:, :n_up], mo_coeff[1][:, :n_dn] # keep only occupied orbitals
alpha_ao = _get_atomic_orbital_envelope_exponents(physical_config, basis_set)
ind_nuc_ao = np.array([ao[0] for ao in atomic_orbitals], dtype=int)
coeff_values = [np.zeros([n_ions, n_up]), np.zeros([n_ions, n_dn])]
alpha_values = [np.zeros([n_ions, n_up]), np.zeros([n_ions, n_dn])]
for spin in range(2):
mo_weights = mo_coeff[spin] ** 2 / np.sum(mo_coeff[spin] ** 2, axis=0) # relative contribution of each ao
for ind_mo in range([n_up, n_dn][spin]):
for ind_nuc in range(n_ions):
index_nuc = (ind_nuc_ao == ind_nuc)
mo_weights_nuc = mo_weights[:, ind_mo][index_nuc]
coeff_values[spin][ind_nuc, ind_mo] = np.sum(mo_weights_nuc) + 1e-3
alpha_values[spin][ind_nuc, ind_mo] = (np.dot(mo_weights_nuc, alpha_ao[index_nuc]) + 1e-6) / (np.sum(mo_weights_nuc) + 1e-6)
for spin in range(2):
# undo softplus which will be applied in network
alpha_values[spin] = np.log(np.exp(alpha_values[spin] - 1))
if pad_full_det:
coeff_values[0] = np.concatenate([coeff_values[0], np.ones([n_ions, n_dn])], axis=1)
alpha_values[0] = np.concatenate([alpha_values[0], np.ones([n_ions, n_dn])], axis=1)
coeff_values[1] = np.concatenate([np.ones([n_ions, n_up]), coeff_values[1]], axis=1)
alpha_values[1] = np.concatenate([np.ones([n_ions, n_up]), alpha_values[1]], axis=1)
return coeff_values, alpha_values
def _int_to_spin_tuple(x):
if type(x) == int:
return (x,) * 2
else:
return x
[docs]def split_results_into_spins(hf):
if len(hf.mo_occ.shape) == 2:
return hf.mo_coeff, hf.mo_energy, hf.mo_occ
else:
return [hf.mo_coeff, hf.mo_coeff], [hf.mo_energy, hf.mo_energy], [hf.mo_occ / 2, hf.mo_occ / 2]
[docs]def get_orbital_type(atomic_orbital):
atom_ind, _, _, l = atomic_orbital
if sum(l) == 0:
orbital_type = "s"
elif sum(l) == 1:
if l[0]:
orbital_type = "px"
elif l[1]:
orbital_type = "py"
else:
orbital_type = "pz"
else:
orbital_type = "other"
return f"Nuc{atom_ind}: {orbital_type}"
[docs]def get_p_orbital_indices_per_atom(atomic_orbitals):
current_atom = -1
p_orb_indices = []
for i,ao in enumerate(atomic_orbitals):
if ao[0] != current_atom:
current_atom = ao[0]
p_orb_indices.append([None, None, None])
l = tuple(ao[-1])
if l == (1, 0, 0):
p_orb_indices[-1][0] = i
elif l == (0, 1, 0):
p_orb_indices[-1][1] = i
elif l == (0, 0, 1):
p_orb_indices[-1][2] = i
return p_orb_indices
def _get_orbitals_by_cas_type(casscf):
"""
Splits orbitals into fixed orbitals (that are either always occupied, or always unoccupied) and active orbitals (that are occupied in some determinants, and unoccupied in others).
Returns:
tuple containing
- **fixed_orbitals** (list of 2 lists): For each spin, a list of indices of fixed orbitals
- **active_orbitals** (list of 2 lists): For each spin, a list of indices of active orbitals
"""
n_core = _int_to_spin_tuple(casscf.ncore)
n_cas = _int_to_spin_tuple(casscf.ncas)
active_orbitals = [list(range(n_core[s], n_core[s] + n_cas[s])) for s in range(2)]
fixed_orbitals = [list(range(n_core[s])) for s in range(2)]
return fixed_orbitals, active_orbitals
def _get_orbital_indices(casscf):
"""
Parse the output of the pySCF CASSCF calculation to determine which orbitals are included in which determinant.
Returns:
(list of 2 np.arrays): First array for spin-up, second array for spin-down. Each array has shape [N_determinants x n_electrons_of_spin] and contains the indices of occupied orbitals in each determinant.
"""
fixed, active = _get_orbitals_by_cas_type(casscf)
nelcas = _int_to_spin_tuple(casscf.nelecas)
occ_up = pyscf.fci.cistring._gen_occslst(active[0], nelcas[0])
occ_down = pyscf.fci.cistring._gen_occslst(active[1], nelcas[1])
orbitals_up = []
orbitals_dn = []
for o_up in occ_up:
for o_dn in occ_down:
orbitals_up.append(fixed[0] + list(o_up))
orbitals_dn.append(fixed[1] + list(o_dn))
return [jnp.array(orbitals_up, dtype=int), jnp.array(orbitals_dn, dtype=int)]
##########################################################################################
##################################### Cusp functions #####################################
##########################################################################################
def _get_local_energy_of_cusp_orbital(r_c, offset, sign, poly, Z, phi_others=0.0):
r = np.linspace(1e-6, r_c, 100)[:, np.newaxis]
p_0 = poly[0] + poly[1] * r + poly[2] * r ** 2 + poly[3] * r ** 3 + poly[4] * r ** 4
p_1 = poly[1] + 2 * poly[2] * r + 3 * poly[3] * r ** 2 + 4 * poly[4] * r ** 3
p_2 = 2 * poly[2] + 6 * poly[3] * r + 12 * poly[4] * r ** 2
prefac = sign * np.exp(p_0) / (offset + sign*np.exp(p_0) + phi_others)
E_l = -prefac * (p_1/r +0.5 *p_2 + 0.5 * p_1**2) - Z/r
penalty = np.nanvar(E_l, axis=0)
# penalty = jnp.max(jnp.abs(E_l - E_l[-1]), axis=0)
return penalty
def _calculate_mo_cusp_params(phi_rc_0, phi_rc_1, phi_rc_2, phi_0, phi_0_others, Z, r_c):
if np.abs(phi_0) < 1e-6 and np.abs(phi_rc_0) < 1e-6:
return 0.0, 0.0, np.zeros(5)
n_cusp_trials = 500
phi_new_0 = phi_0 * (1.0 + np.concatenate([np.logspace(-2, 1, n_cusp_trials//2), -np.logspace(-2, 1, n_cusp_trials//2)]))
sign = jnp.sign(phi_new_0 - phi_rc_0)
offset = 2 * phi_rc_0 - phi_new_0 # = "C"
phi_rc_shifted = phi_rc_0 - offset # = "R(r_c)"
x1 = np.log(jnp.abs(phi_rc_shifted))
x2 = phi_rc_1 / phi_rc_shifted
x3 = phi_rc_2 / phi_rc_shifted
x4 = -Z * (phi_new_0 + phi_0_others) / (phi_new_0 - offset)
x5 = np.log(np.abs(phi_new_0 - offset))
rc2 = r_c*r_c
rc3 = rc2 * r_c
rc4 = rc2 * rc2
poly = np.array([
x5,
x4,
6*x1/rc2 - 3*x2/r_c + 0.5 * x3 -3*x4/r_c - 6*x5/rc2 - 0.5*x2*x2,
-8*x1/rc3 + 5*x2/rc2 - x3/r_c + 3*x4/rc2 + 8*x5/rc3 + x2*x2/r_c,
3*x1/rc4 - 2*x2/rc3 + 0.5*x3/rc2 - x4/rc3 - 3*x5/rc4 - 0.5*x2*x2/rc2
])
E_loc_cusp = _get_local_energy_of_cusp_orbital(r_c, offset, sign, poly, Z, phi_0_others)
ind_opt = np.nanargmin(E_loc_cusp)
return offset[ind_opt], sign[ind_opt], poly[:,ind_opt]
def _calculate_ao_cusp_params(ind_nuc, alpha, gto_coeffs, angular_momenta, Z):
if sum(angular_momenta) != 0:
return None
r_c = jnp.minimum(0.5, 1/Z)
phi_rc_0 = jnp.sum(gto_coeffs * np.exp(-alpha * r_c ** 2))
phi_rc_1 = jnp.sum(gto_coeffs * (-2 * alpha * r_c) * np.exp(-alpha * r_c ** 2))
phi_rc_2 = jnp.sum(gto_coeffs * (-2 * alpha + 4 * (alpha * r_c) ** 2) * np.exp(-alpha * r_c ** 2))
phi_0 = jnp.sum(gto_coeffs)
n_cusp_trials = 500
phi_new_0 = phi_0 * np.linspace(-1, 3, n_cusp_trials)
sign = jnp.sign(phi_new_0 - phi_rc_0)
offset = 2 * phi_rc_0 - phi_new_0
phi_shifted = phi_rc_0 - offset
p0 = jnp.log((phi_new_0 - offset) * sign)
p1 = -Z * (offset * sign * jnp.exp(-p0) + 1)
A = jnp.array([[r_c ** 2, r_c ** 3, r_c ** 4],
[2 * r_c, 3 * r_c ** 2, 4 * r_c ** 3],
[2, 6 * r_c, 12 * r_c ** 2]]
)
b = jnp.array([
jnp.log(phi_shifted * sign) - p0 - p1 * r_c,
phi_rc_1 / phi_shifted - p1,
phi_rc_2 / phi_shifted - (phi_rc_1 / phi_shifted) ** 2
])
poly = jnp.concatenate([jnp.array([p0, p1]), jnp.linalg.solve(A, b)])
E_loc_cusp = _get_local_energy_of_cusp_orbital(r_c, offset, sign, poly, Z, phi_others=0.0)
ind_opt = jnp.nanargmin(E_loc_cusp)
return r_c, offset[ind_opt], sign[ind_opt], poly[:,ind_opt]
[docs]def calculate_molecular_orbital_cusp_params(atomic_orbitals, mo_coeff, R, Z, r_cusp_scale):
n_molecular_orbitals, n_nuclei, n_atomic_orbitals = mo_coeff.shape[1], len(R), len(atomic_orbitals)
cusp_rc = np.minimum(r_cusp_scale / Z, 0.5)
cusp_offset = np.zeros([n_nuclei, n_molecular_orbitals])
cusp_sign = np.zeros([n_nuclei, n_molecular_orbitals])
cusp_poly = np.zeros([n_nuclei, n_molecular_orbitals, 5])
cusp_1s_coeffs = np.zeros([n_nuclei, n_atomic_orbitals, n_molecular_orbitals])
for nuc_idx in range(n_nuclei):
for mo_idx in range(n_molecular_orbitals):
diff, dist = get_el_ion_distance_matrix(jnp.array(R[nuc_idx]), R)
ao = eval_atomic_orbitals(diff, dist, atomic_orbitals)
is_centered_1s = np.array([(a[0] == nuc_idx) and (sum(a[3]) == 0) for a in atomic_orbitals])
phi_0_1s = (is_centered_1s * ao) @ mo_coeff[:, mo_idx]
phi_0_others = ((1-is_centered_1s) * ao) @ mo_coeff[:, mo_idx]
phi_rc_0, phi_rc_1, phi_rc_2 = 0.0, 0.0, 0.0
r_c = cusp_rc[nuc_idx]
for i, (_, alpha, weights, _) in enumerate(atomic_orbitals):
if is_centered_1s[i]:
phi_rc_0 += np.sum(weights * np.exp(-alpha * r_c ** 2)) * mo_coeff[i, mo_idx]
phi_rc_1 += np.sum(weights * (-2 * alpha * r_c) * np.exp(-alpha * r_c ** 2)) * mo_coeff[i, mo_idx]
phi_rc_2 += np.sum(weights * (-2 * alpha + 4 * (alpha * r_c) ** 2) * np.exp(-alpha * r_c ** 2)) * mo_coeff[i, mo_idx]
cusp_1s_coeffs[nuc_idx, :, mo_idx] = is_centered_1s * mo_coeff[:, mo_idx] # n_nuc x n_atomic_orbitals x n_molec_orbitals
cusp_offset[nuc_idx, mo_idx], cusp_sign[nuc_idx, mo_idx], cusp_poly[nuc_idx, mo_idx] = _calculate_mo_cusp_params(phi_rc_0,
phi_rc_1,
phi_rc_2,
phi_0_1s,
phi_0_others,
Z[nuc_idx],
r_c)
return cusp_rc, cusp_offset, cusp_sign, cusp_poly, cusp_1s_coeffs