#!/usr/bin/env python3
# [1] https://doi.org/10.1063/1.450106
# Efficient recursive computation of molecular integrals over Cartesian
# Gaussian functions
# Obara, Saika, 1986
# [2] https://doi.org/10.1002/9781119019572
# Molecular Electronic-Structure Theory
# Helgaker, Jørgensen, Olsen
# [3] https://doi.org/10.1021/acs.jctc.7b00788
# LIBRETA: Computerized Optimization and Code Synthesis for
# Electron Repulsion Integral Evaluation
# Jun Zhang
# [4] https://doi.org/10.1039/B413539C
# Efficient evaluation of three-center two-electron integrals over Gaussian functions
# Ahlrichs, 2004
# [5] EVALUATING MANY-ELECTRON MOLECULAR INTEGRALS FOR QUANTUM CHEMISTRY
# James Christopher Womack, PhD Thesis
# [6] https://doi.org/10.1063/1.4983393
# Efficient evaluation of three-center Coulomb integrals
# Samu, Kállay, 2017
# [7] https://arxiv.org/pdf/2210.03192.pdf
# Memory-Efficient Recursive Evaluation of 3-Center Gaussian Integrals
# Asadchev, Valeev, 2022
import argparse
from datetime import datetime
import functools
import itertools as it
import os
from pathlib import Path
import random
import string
import sys
import textwrap
import time
from jinja2 import Template
import numpy as np
from sympy import (
Array,
cse,
exp,
Expr,
flatten,
Function,
IndexedBase,
Matrix,
permutedims,
pi,
sqrt,
Symbol,
symbols,
tensorcontraction as tc,
tensorproduct as tp,
)
from sympy.codegen.ast import Assignment
from sympy.printing.numpy import NumPyPrinter
from sympy.printing.c import C99CodePrinter
# from pysisyphus.wavefunction.cart2sph import cart2sph_coeffs
try:
from pysisyphus.config import L_MAX, L_AUX_MAX
except ModuleNotFoundError:
L_MAX = 4
L_AUX_MAX = 5
L_MAP = {
0: "s",
1: "p",
2: "d",
3: "f",
4: "g",
5: "h",
}
KEYS = (
"cgto",
"ovlp",
"dpm",
"dqpm",
"qpm",
"kin",
"coul",
"3c2e",
"3c2e_sph",
)
ONE_THRESH = 1e-14
[docs]
def make_py_func(repls, reduced, args=None, name=None, doc_str="", return_array=True):
if args is None:
args = list()
# Generate random name, if no name was supplied
if name is None:
name = "func_" + "".join(
[random.choice(string.ascii_letters) for i in range(8)]
)
# This allows using the 'boys' function without producing an error
print_settings = {
"allow_unknown_functions": True,
}
print_func = NumPyPrinter(print_settings).doprint
assignments = [Assignment(lhs, rhs) for lhs, rhs in repls]
py_lines = [print_func(as_) for as_ in assignments]
return_val = print_func(reduced)
try:
n_return_vals = len(reduced)
except TypeError:
n_return_vals = 1
tpl = Template(
"""
def {{ name }}({{ args }}):
{% if doc_str %}
\"\"\"{{ doc_str }}\"\"\"
{% endif %}
{% for line in py_lines %}
{{ line }}
{% endfor %}
{% if return_array %}
# {{ n_return_vals }} item(s)
return numpy.array({{ return_val }})
{% else %}
return {{ return_val }}
{% endif %}
""",
trim_blocks=True,
lstrip_blocks=True,
)
rendered = textwrap.dedent(
tpl.render(
name=name,
args=args,
py_lines=py_lines,
return_val=return_val,
n_return_vals=n_return_vals,
return_array=return_array,
doc_str=doc_str,
)
).strip()
return rendered
[docs]
def make_c_func(repls, reduced, args=None, name=None, doc_str=""):
if args is None:
args = list()
# Generate random name, if no name was supplied
if name is None:
name = "func_" + "".join(
[random.choice(string.ascii_letters) for i in range(8)]
)
arg_strs = list()
for arg in args:
if arg.islower():
arg_str = f"double {arg}"
elif arg.isupper():
arg_str = f"double {arg}[3]"
else:
raise Exception
arg_strs.append(arg_str)
args_str = ", ".join(arg_strs)
# This allows using the 'boys' function without producing an error
print_settings = {
"allow_unknown_functions": True,
# Without disabling contract some expressions will raise ValueError.
"contract": False,
}
print_func = C99CodePrinter(print_settings).doprint
assignments = [Assignment(lhs, rhs) for lhs, rhs in repls]
repl_lines = [print_func(as_) for as_ in assignments]
res_lines = [print_func(red) for red in reduced]
res_len = len(reduced)
signature = f"double * {name}({args_str})"
tpl = Template(
"""
{{ signature }} {
{% if doc_str %}
/* {{ doc_str }} */
{% endif %}
static double {{ res_name }}[{{ res_len}}];
{% for line in repl_lines %}
const double {{ line }}
{% endfor %}
{% for rhs in res_lines %}
{{ res_name }}[{{ loop.index0}}] = {{ rhs }};
{% endfor %}
return {{ res_name }};
}
""",
trim_blocks=True,
lstrip_blocks=True,
)
rendered = textwrap.dedent(
tpl.render(
signature=signature,
res_name="results",
res_len=res_len,
repl_lines=repl_lines,
res_lines=res_lines, # c_lines=c_lines,
reduced=reduced,
doc_str=doc_str,
args_str=args_str,
)
).strip()
return rendered, signature
[docs]
def canonical_order(L):
inds = list()
for i in range(L + 1):
l = L - i
for n in range(i + 1):
m = i - n
inds.append((l, m, n))
return inds
[docs]
def shell_iter(Ls):
"""Iterator over cartesian product of L values in Ls."""
return it.product(*[canonical_order(L) for L in Ls])
[docs]
class CartGTO3d(Function):
"""3D Cartesian Gaussian function; not normalized."""
[docs]
@classmethod
@functools.cache
def eval(cls, i, j, k, a, Xa, Ya, Za):
Xa2 = Xa**2
Ya2 = Ya**2
Za2 = Za**2
return (Xa**i) * (Ya**j) * (Za**k) * exp(-a * (Xa2 + Ya2 + Za2))
[docs]
class CartGTOShell(Function):
[docs]
@classmethod
def eval(cls, La_tot, a, Xa, Ya, Za):
exprs = [CartGTO3d(*La, a, Xa, Ya, Za) for La, in shell_iter((La_tot,))]
# print(CartGTO3d.eval.cache_info())
return exprs
[docs]
class TwoCenter1d(Expr):
def __new__(self, a, A, b, B, C=None):
self.a = a
self.A = A
self.b = b
self.B = B
self.C = C
return super().__new__(self, a, A, b, B, C)
@property
def p(self):
"""Total exponent p."""
return self.a + self.b
@property
def mu(self):
"""Reduced exponent mu."""
return self.a * self.b / self.p
@property
def AB(self):
"""Relative coordinate/Gaussian separation X_AB."""
return self.A - self.B
@property
def P(self):
"""Center-of-charge coordinate."""
return (self.a * self.A + self.b * self.B) / self.p
@property
def PA(self):
"""Relative coordinate/Gaussian separation X_PA."""
return self.P - self.A
@property
def PB(self):
"""Relative coordinate/Gaussian separation X_PB."""
return self.P - self.B
@property
def PC(self):
"""Relative coordinate/Gaussian separation X_PC."""
return self.P - self.C
@property
def K(self):
return exp(-self.mu * self.AB**2)
[docs]
class Multipole1d(TwoCenter1d):
"""1d multipole-moment integral of order 'e', between primitive 1d Gaussians
Ga = G_i(a, r, A) and Gb = G_j(b, r, B) with Cartesian quantum number i and j,
exponents a and b, centered at A (B). The origin of the multipole expansion is
at C.
"""
[docs]
@functools.cache
def eval(self, i, j, e):
ang_moms = (i, j, e)
if any([_ < 0 for _ in ang_moms]):
return 0
recur = self.eval
def vrr(i, j, e, X):
return X * recur(i, j, e) + 1 / (2 * self.p) * (
i * recur(i - 1, j, e) + j * recur(i, j - 1, e) + e * recur(i, j, e - 1)
)
# Base case
if all([_ == 0 for _ in ang_moms]):
return sqrt(pi / self.p) * self.K
# Decrement i
elif i.is_positive:
return vrr(i - 1, j, e, self.PA)
# Decrement j
elif j.is_positive:
return vrr(i, j - 1, e, self.PB)
elif e.is_positive:
return vrr(i, j, e - 1, self.PC)
[docs]
class Multipole3d(Function):
[docs]
@classmethod
def eval(cls, La, Lb, a, b, A, B, Le, C):
x, y, z = [
Multipole1d(a, A[i], b, B[i], C[i]).eval(La[i], Lb[i], Le[i])
for i in range(3)
]
return x * y * z
[docs]
class Multipole3dShell(Function):
[docs]
@classmethod
def eval(cls, La_tot, Lb_tot, a, b, A, B, Le_tot=0, C=(0.0, 0.0, 0.0)):
exprs = [
Multipole3d(La, Lb, a, b, A, B, Le, C)
for Le, La, Lb in shell_iter((Le_tot, La_tot, Lb_tot))
]
# print(Multipole1d.eval.cache_info())
return exprs
[docs]
class DiagQuadrupole3dShell(Function):
[docs]
@classmethod
def eval(cls, La_tot, Lb_tot, a, b, A, B, C=(0.0, 0.0, 0.0)):
exprs = list()
for Le in ((2, 0, 0), (0, 2, 0), (0, 0, 2)):
for La, Lb in shell_iter((La_tot, Lb_tot)):
exprs.append(Multipole3d(La, Lb, a, b, A, B, Le, C))
# print(DiagQuadrupole3dShell.eval.cache_info())
return exprs
[docs]
class Overlap1d(TwoCenter1d):
[docs]
def eval(self, i, j):
return Multipole1d(self.a, self.A, self.b, self.B).eval(i, j, 0)
[docs]
class Overlap3dShell(Function):
"""Whole shell of 3d overlap integrals for a given L pair
over primitive gaussians."""
[docs]
@classmethod
def eval(cls, La_tot, Lb_tot, a, b, A, B):
exprs = Multipole3dShell(La_tot, Lb_tot, a, b, A, B)
# print(Multipole1d.eval.cache_info())
return exprs
[docs]
class Kinetic1d(TwoCenter1d):
[docs]
@functools.cache
def eval(self, i, j):
if i < 0 or j < 0:
return 0
recur = self.eval
def recur_rel(i, j, X):
return X * recur(i, j) + 1 / (2 * self.p) * (
i * recur(i - 1, j) + j * recur(i, j - 1)
)
def recur_ovlp(i, j):
return Overlap1d(self.a, self.A, self.b, self.B).eval(i, j)
# Base case
if i == 0 and j == 0:
return (
self.a - 2 * self.a**2 * (self.PA**2 + 1 / (2 * self.p))
) * recur_ovlp(i, j)
# Decrement i
elif i > 0:
# Eq. (9.3.41)
return recur_rel(i - 1, j, self.PA) + self.b / self.p * (
2 * self.a * recur_ovlp(i, j) - i * recur_ovlp(i - 2, j)
)
# Decrement j
elif j > 0:
# Eq. (9.3.41)
return recur_rel(i, j - 1, self.PB) + self.a / self.p * (
2 * self.b * recur_ovlp(i, j) - j * recur_ovlp(i, j - 2)
)
[docs]
class Kinetic3d(Function):
[docs]
@classmethod
def eval(cls, La, Lb, a, b, A, B):
Tx, Ty, Tz = [Kinetic1d(a, A[i], b, B[i]).eval(La[i], Lb[i]) for i in range(3)]
Sx, Sy, Sz = [Overlap1d(a, A[i], b, B[i]).eval(La[i], Lb[i]) for i in range(3)]
return Tx * Sy * Sz + Sx * Ty * Sz + Sx * Sy * Tz
[docs]
class Kinetic3dShell(Function):
[docs]
@classmethod
def eval(cls, La_tot, Lb_tot, a, b, A, B):
exprs = [
Kinetic3d(La, Lb, a, b, A, B) for La, Lb in shell_iter((La_tot, Lb_tot))
]
return exprs
# Placeholder for the Boys-function. The actual Boys-function will be
# imported in the generated python module.
boys = Function("boys")
[docs]
class Coulomb(TwoCenter1d):
"""Nucleus at C."""
[docs]
@functools.cache
def eval(self, i, k, m, j, l, n, N):
ang_moms = (i, k, m, j, l, n)
if any([am < 0 for am in ang_moms]):
return 0
def recur(N, *inds):
"""Simple wrapper to pass all required arguments."""
return self.eval(*inds, N) # , a, b, A, B, C)
def decr(to_decr, decr_ind):
one = np.zeros(3, dtype=int)
one[decr_ind] = 1
def Ni(inds):
return inds[decr_ind]
bra = np.array((i, k, m), dtype=int)
ket = np.array((j, l, n), dtype=int)
if to_decr == "bra":
X = self.PA
bra[decr_ind] -= 1
else:
X = self.PB
ket[decr_ind] -= 1
bra_decr = bra - one
ket_decr = ket - one
return (
X[decr_ind] * recur(N, *bra, *ket)
- self.PC[decr_ind] * recur(N + 1, *bra, *ket)
+ 1
/ (2 * self.p)
* Ni(bra)
* (recur(N, *bra_decr, *ket) - recur(N + 1, *bra_decr, *ket))
+ 1
/ (2 * self.p)
* Ni(ket)
* (recur(N, *bra, *ket_decr) - recur(N + 1, *bra, *ket_decr))
)
# Base case
if all([am == 0 for am in ang_moms]):
K = exp(-self.mu * self.AB.dot(self.AB))
return 2 * pi / self.p * K * boys(N, self.p * self.PC.dot(self.PC))
elif i > 0:
return decr("bra", 0)
elif j > 0:
return decr("ket", 0)
elif k > 0:
return decr("bra", 1)
elif l > 0:
return decr("ket", 1)
elif m > 0:
return decr("bra", 2)
elif n > 0:
return decr("ket", 2)
[docs]
class CoulombShell(Function):
[docs]
@classmethod
def eval(cls, La_tot, Lb_tot, a, b, A, B, C=(0.0, 0.0, 0.0)):
exprs = [
# Coulomb(*La, *Lb, 0, a, b, A, B, C)
Coulomb(a, A, b, B, C).eval(*La, *Lb, 0)
for La, Lb in shell_iter((La_tot, Lb_tot))
]
# print(Coulomb.eval.cache_info())
return exprs
[docs]
class TwoCenterTwoElectron(Function):
[docs]
@classmethod
@functools.cache
def eval(cls, ia, ja, ka, ib, jb, kb, N, a, b, A, B):
ang_moms = np.array((ia, ja, ka, ib, jb, kb), dtype=int)
ang_moms2d = ang_moms.reshape(-1, 3)
if any([am < 0 for am in ang_moms]):
return 0
p = a + b
P = (a * A + b * B) / p
mu = (a * b) / p
def recur(N, *inds):
return cls(*inds, N, a, b, A, B)
def vrr(bra_or_ket, cart_ind):
assert bra_or_ket in ("bra", "ket")
if bra_or_ket == "bra":
ind1 = 0
ind2 = 1
else:
ind1 = 1
ind2 = 0
exps = (a, b)
X = (P - A, P - B)[ind1][cart_ind]
l1 = ang_moms2d[ind1, cart_ind] - 1
exp1 = exps[ind1]
decr1 = ang_moms2d.copy()
decr1[ind1, cart_ind] -= 1
decr11 = decr1.copy()
decr11[ind1, cart_ind] -= 1
l2 = ang_moms2d[ind2, cart_ind]
exp2 = exps[ind2]
decr12 = ang_moms2d.copy()
decr12[ind1, cart_ind] -= 1
decr12[ind2, cart_ind] -= 1
decr1 = decr1.flatten()
decr11 = decr11.flatten()
decr12 = decr12.flatten()
return (
X * recur(N + 1, *decr1)
+ l1
/ (2 * exp1)
* (recur(N, *decr11) - exp2 / p * recur(N + 1, *decr11))
+ l2 / (2 * p) * recur(N + 1, *decr12)
)
# vrr_bra = functools.partial(vrr, bra_or_ket="bra")
vrr_bra = functools.partial(vrr, "bra")
vrr_ket = functools.partial(vrr, "ket")
# Base case
if (ang_moms == 0).all():
AB = A - B
return 2 * pi**2.5 / sqrt(p) / (a * b) * boys(N, mu * AB.dot(AB))
elif ia > 0:
return vrr_bra(0)
elif ja > 0:
return vrr_bra(1)
elif ka > 0:
return vrr_bra(2)
elif ib > 0:
return vrr_ket(0)
elif jb > 0:
return vrr_ket(1)
elif kb > 0:
return vrr_ket(2)
[docs]
class TwoCenterTwoElectronShell(Function):
[docs]
@classmethod
def eval(cls, La_tot, Lb_tot, a, b, A, B):
exprs = [
TwoCenterTwoElectron(*La, *Lb, 0, a, b, A, B)
for La, Lb in shell_iter((La_tot, Lb_tot))
]
# print(TwoCenterTwoElectron.eval.cache_info())
return exprs
[docs]
class ThreeCenterTwoElectronBase(Function):
"""
https://pubs.rsc.org/en/content/articlelanding/2004/CP/b413539c
There is an error in the base case (00|0). One must divide by
sqrt(eta + gamma), not multiply.
"""
[docs]
@classmethod
@functools.cache
def eval(cls, ia, ja, ka, ib, jb, kb, ic, jc, kc, N, a, b, c, A, B, C):
ang_moms = np.array((ia, ja, ka, ib, jb, kb, ic, jc, kc), dtype=int)
ang_moms2d = ang_moms.reshape(-1, 3)
if any([am < 0 for am in ang_moms]):
return 0
p = a + b
P = (a * A + b * B) / p
mu = (a * b) / p
X_PC = P - C
rho = p * c / (p + c)
def recur(N, *inds):
"""Simple wrapper to pass all required arguments.
Here we don't use ThreeCenterTwoElectronBase, but the derived classes,
that provide the 'aux_vrr' attribute, to distinguish between the different
vertical recursion relations."""
return cls(*inds, N, a, b, c, A, B, C)
def recur_hrr(cart_ind):
"""Horizontal recursion relation to transfer angular momentum in bra.
(a, b+1_i|c) = (a + 1_i,b|c) + X_AB (ab|c)
"""
assert N == 0
incr_ang_moms = ang_moms2d.copy()
incr_ang_moms[0, cart_ind] += 1 # Increase in left orbital
incr_ang_moms[1, cart_ind] -= 1 # Decrease in right orbital
decr_ang_moms = ang_moms2d.copy()
decr_ang_moms[1, cart_ind] -= 1 # Decrease in right orbital
incr_ang_moms = incr_ang_moms.flatten()
decr_ang_moms = decr_ang_moms.flatten()
AB_dir = (A - B)[cart_ind]
return recur(N, *incr_ang_moms) + AB_dir * recur(N, *decr_ang_moms)
def recur_vrr(cart_ind):
assert (ib, jb, kb) == (0, 0, 0)
assert (ic, jc, kc) == (0, 0, 0)
decr_a = ang_moms2d.copy()
decr_a[0, cart_ind] -= 1 # Decrease in bra left orbital
decr_aa = decr_a.copy()
decr_aa[0, cart_ind] -= 1 # Decrease in bra left orbital again
PA_dir = (P - A)[cart_ind]
PC_dir = (P - C)[cart_ind]
ai = (ia, ja, ka)[cart_ind] - 1
_2p = 2 * p
decr_a = decr_a.flatten()
decr_aa = decr_aa.flatten()
return (
PA_dir * recur(N, *decr_a)
- rho / p * PC_dir * recur(N + 1, *decr_a)
+ ai / _2p * (recur(N, *decr_aa) - rho / p * recur(N + 1, *decr_aa))
)
def recur_vrr_aux(cart_ind):
decr_c = ang_moms2d.copy()
decr_c[2, cart_ind] -= 1
decr_cc = decr_c.copy()
decr_cc[2, cart_ind] -= 1
decr_ac = decr_c.copy()
decr_ac[0, cart_ind] -= 1
decr_bc = decr_c.copy()
decr_bc[1, cart_ind] -= 1
PC_dir = (P - C)[cart_ind]
la = (ia, ja, ka)[cart_ind]
lb = (ib, jb, kb)[cart_ind]
lc = (ic, jc, kc)[cart_ind] - 1
decr_c = decr_c.flatten()
decr_cc = decr_cc.flatten()
decr_ac = decr_ac.flatten()
decr_bc = decr_bc.flatten()
return (
p / (p + c) * PC_dir * recur(N + 1, *decr_c)
+ lc
/ (2 * c)
* (recur(N, *decr_cc) - p / (p + c) * recur(N + 1, *decr_cc))
+ la / (2 * (p + c)) * recur(N + 1, *decr_ac)
+ lb / (2 * (p + c)) * recur(N + 1, *decr_bc)
)
def recur_vrr_aux_sph(cart_ind):
assert (ib, jb, kb) == (0, 0, 0)
decr_c = ang_moms2d.copy()
decr_c[2, cart_ind] -= 1
decr_ac = decr_c.copy()
decr_ac[0, cart_ind] -= 1
La = (ia, ja, ka)[cart_ind]
PC_dir = (P - C)[cart_ind]
return (
rho
/ c
* (
PC_dir * recur(N + 1, *decr_c.flatten())
+ La / (2 * p) * recur(N + 1, *decr_ac.flatten())
)
)
recur_vrr_aux_funcs = {
"cart": recur_vrr_aux,
"sph": recur_vrr_aux_sph,
}
recur_vrr_aux_func = recur_vrr_aux_funcs[cls.aux_vrr]
# Base case
if (ang_moms == 0).all():
X_AB = A - B
r2_PC = X_PC.dot(X_PC)
r2_AB = X_AB.dot(X_AB)
chi = rho * r2_PC
K = exp(-mu * r2_AB)
return 2 * pi**2.5 / sqrt(p + c) / (p * c) * K * boys(N, chi)
elif ib > 0:
return recur_hrr(0)
elif jb > 0:
return recur_hrr(1)
elif kb > 0:
return recur_hrr(2)
elif ic > 0:
return recur_vrr_aux_func(0)
elif jc > 0:
return recur_vrr_aux_func(1)
elif kc > 0:
return recur_vrr_aux_func(2)
elif ia > 0:
return recur_vrr(0)
elif ja > 0:
return recur_vrr(1)
elif ka > 0:
return recur_vrr(2)
[docs]
class ThreeCenterTwoElectron(ThreeCenterTwoElectronBase):
aux_vrr = "cart"
[docs]
class ThreeCenterTwoElectronSph(ThreeCenterTwoElectronBase):
aux_vrr = "sph"
[docs]
class ThreeCenterTwoElectronShell(Function):
[docs]
@classmethod
def eval(cls, La_tot, Lb_tot, Lc_tot, a, b, c, A, B, C):
exprs = [
ThreeCenterTwoElectron(*La, *Lb, *Lc, 0, a, b, c, A, B, C)
for La, Lb, Lc in shell_iter((La_tot, Lb_tot, Lc_tot))
]
# print(ThreeCenterTwoElectron.eval.cache_info())
return exprs
[docs]
class ThreeCenterTwoElectronSphShell(Function):
[docs]
@classmethod
def eval(cls, La_tot, Lb_tot, Lc_tot, a, b, c, A, B, C):
exprs = [
ThreeCenterTwoElectronSph(*La, *Lb, *Lc, 0, a, b, c, A, B, C)
for La, Lb, Lc in shell_iter((La_tot, Lb_tot, Lc_tot))
]
# print(ThreeCenterTwoElectron.eval.cache_info())
return exprs
[docs]
class ERI(Function):
"""Variables named according to libreta paper [3]."""
[docs]
@classmethod
def eval(
cls, ia, ja, ka, ib, jb, kb, ic, jc, kc, id_, jd, kd, N, a, b, c, d, A, B, C, D
):
ang_moms = np.array(
(ia, ja, ka, ib, jb, kb, ic, jc, kc, id_, jd, kd), dtype=int
)
ang_moms2d = ang_moms.reshape(-1, 3)
if any([am < 0 for am in ang_moms]):
return 0
AB = A - B
xi = a + b
P = (a * A + b * B) / xi # Eq. (5)
CD = C - D
zeta = c + d
Q = (c * C + d * D) / zeta # Eq. (6)
theta = xi * zeta / (xi + zeta) # Eq. (7)
def recur(N, *ang_moms):
"""Simple wrapper to pass all required arguments."""
return ERI(*ang_moms, N, a, b, c, d, A, B, C, D)
def recur_hrr(bra_or_ket, cart_ind):
"""Horizontal recursion relation to transfer angmoms in bra or ket.
(a, b+1_x|cd) = (a + 1_x,b|cd) + X_AB (ab|cd)
(ab|c, d + 1_x) = (ab|c + 1_x, d) + X_CD (ab|cd)
"""
if bra_or_ket == "bra":
XYZ = AB
incr_ind = 0
else:
XYZ = CD
incr_ind = 2
decr_ind = incr_ind + 1
incr_ang_moms = ang_moms2d.copy()
incr_ang_moms[incr_ind, cart_ind] += 1 # Increase in left orbital
incr_ang_moms[decr_ind, cart_ind] -= 1 # Decrease in right orbital
incr_ang_moms = incr_ang_moms.flatten()
decr_ang_moms = ang_moms2d.copy()
decr_ang_moms[decr_ind, cart_ind] -= 1
decr_ang_moms = decr_ang_moms.flatten()
return recur(N, *incr_ang_moms) + XYZ[cart_ind] * recur(N, *decr_ang_moms)
def recur_vrr(bra_or_ket, cart_ind):
assert (ib, jb, kb) == (0, 0, 0)
assert (id_, jd, kd) == (0, 0, 0)
if bra_or_ket == "bra":
XYZ = P - A
decr_ind = 0
decr_also_ind = 2
exp1, exp2 = zeta, xi
am1 = (ia, ja, ka)[cart_ind] - 1
am2 = (ic, jc, kc)[cart_ind]
sign = -1
else:
XYZ = Q - C
decr_ind = 2
decr_also_ind = 0
exp1, exp2 = xi, zeta
am1 = (ic, jc, kc)[cart_ind] - 1
am2 = (ia, ja, ka)[cart_ind]
sign = 1
decr_ang_moms = ang_moms2d.copy()
decr_ang_moms[decr_ind, cart_ind] -= 1
decr_ang_moms = decr_ang_moms.flatten()
decr2_ang_moms = ang_moms2d.copy()
decr2_ang_moms[decr_ind, cart_ind] -= 2 # Further decrease angular momentum
decr2_ang_moms = decr2_ang_moms.flatten()
decr_also_ang_moms = ang_moms2d.copy()
decr_also_ang_moms[decr_also_ind, cart_ind] -= 1
decr_also_ang_moms = decr_also_ang_moms.flatten()
denom = xi + zeta
quot = exp1 / denom
PQ = P - Q
return (
XYZ[cart_ind] * recur(N, *decr_ang_moms)
+ sign * quot * PQ[cart_ind] * recur(N + 1, *decr_ang_moms)
+ am1
/ (2 * exp2)
* (recur(N, *decr2_ang_moms) - quot * recur(N + 1, *decr2_ang_moms))
+ am2 / (2 * denom) * recur(N + 1, *decr_also_ang_moms)
)
# Base case
if all([am == 0 for am in ang_moms]):
RAB2 = AB.dot(AB)
RCD2 = CD.dot(CD)
Kab = exp(-(a * b) / xi * RAB2)
Kcd = exp(-(c * d) / zeta * RCD2)
PQ = P - Q
RPQ2 = PQ.dot(PQ)
return (
2
* pi ** (5 / 2)
/ (xi * zeta + sqrt(xi + zeta))
* Kab
* Kcd
* boys(N, theta * RPQ2)
) # Eq. (8) [ss|ss]^N
#
# Horizontal recursion relation to reduce angular momentum.
#
# Bra, basis func. with exp. b, centered at B with Lb_tot = (ib + jb + kb)
elif ib > 0:
return recur_hrr("bra", 0)
elif jb > 0:
return recur_hrr("bra", 1)
elif kb > 0:
return recur_hrr("bra", 2)
# Ket, basis func. with exp. d, centered at D with Ld_tot = (id_ + jd + kd)
elif id_ > 0:
return recur_hrr("ket", 0)
elif jd > 0:
return recur_hrr("ket", 1)
elif kd > 0:
return recur_hrr("ket", 2)
#
# Vertical recursion relation to reduce angular momentum.
#
# Bra, basis func. with exp. a, centered at A with La_tot = (ia + ja + ka)
elif ia > 0:
return recur_vrr("bra", 0)
elif ja > 0:
return recur_vrr("bra", 1)
elif ka > 0:
return recur_vrr("bra", 2)
# Ket, basis func. with exp. c, centered at C with Lc_tot = (ic + jc + kc)
elif ic > 0:
return recur_vrr("ket", 0)
elif jc > 0:
return recur_vrr("ket", 1)
elif kc > 0:
return recur_vrr("ket", 1)
[docs]
class ERIShell(Function):
[docs]
@classmethod
def eval(cls, La_tot, Lb_tot, Lc_tot, Ld_tot, a, b, c, d, A, B, C, D):
exprs = [
ERI(*La, *Lb, *Lc, *Ld, 0, a, b, c, d, A, B, C, D)
for La, Lb, Lc, Ld in shell_iter((La_tot, Lb_tot, Lc_tot, Ld_tot))
]
# print(ERI.eval.cache_info())
return exprs
[docs]
class SpinOrbit(Function):
"""1-electron spin-orbit interaction integral for Cartesian direction μ."""
[docs]
@classmethod
@functools.cache
def eval(cls, i, k, m, j, l, n, N, mu, a, b, A, B, C):
ang_moms = (i, k, m, j, l, n)
if any([am < 0 for am in ang_moms]):
return 0
p = a + b
P = (a * A + b * B) / p
X_PC = P - C
def coulomb(N, *inds):
return Coulomb(*inds, N, a, b, A, B, C)
def recur(N, *inds):
"""Simple wrapper to pass all required arguments."""
return SpinOrbit(*inds, N, mu, a, b, A, B, C)
def decr(to_decr, decr_ind):
one = np.zeros(3, dtype=int)
one[decr_ind] = 1
def Ni(inds):
return inds[decr_ind]
bra = np.array((i, k, m), dtype=int)
ket = np.array((j, l, n), dtype=int)
if to_decr == "bra":
X = P - A
bra[decr_ind] -= 1
else:
X = P - B
ket[decr_ind] -= 1
bra_decr = bra - one
ket_decr = ket - one
expr = (
# The initial part of this recursion is the same as for the 1-electron
# Coulomb integrals.
X[decr_ind] * recur(N, *bra, *ket)
- X_PC[decr_ind] * recur(N + 1, *bra, *ket)
+ 1
/ (2 * p)
* Ni(bra)
* (recur(N, *bra_decr, *ket) - recur(N + 1, *bra_decr, *ket))
+ 1
/ (2 * p)
* Ni(ket)
* (recur(N, *bra, *ket_decr) - recur(N + 1, *bra, *ket_decr))
# From here on, the recursion differs.
)
if to_decr == "bra":
X = B - C
exponent = b
else:
X = A - C
exponent = a
# Second to line in Eq. (A34) in [1]
expr += 2 * exponent * Matrix(one).cross(X)[mu] * coulomb(N + 1, *bra, *ket)
# Last line in Eq. (A34) in [1]
for k_ in range(
3
): # Use k_ instead of k, as this is already an angular momentum
one_k = np.zeros(3, dtype=int)
one_k[k_] = 1
one_ind = np.cross(one, one_k)[mu]
if to_decr == "bra":
expr += ket[k_] * one_ind * coulomb(N + 1, *bra, *(ket - one_k))
else:
expr += bra[k_] * one_ind * coulomb(N + 1, *(bra - one_k), *ket)
return expr
# Base case
if all([am == 0 for am in ang_moms]):
AC = A - C
BC = B - C
ACxBC = AC.cross(BC)
return 4 * p * ACxBC[m] * coulomb(N + 1, 0, 0, 0, 0, 0, 0)
elif i > 0:
return decr("bra", 0)
elif j > 0:
return decr("ket", 0)
elif k > 0:
return decr("bra", 1)
elif l > 0:
return decr("ket", 1)
elif m > 0:
return decr("bra", 2)
elif n > 0:
return decr("ket", 2)
[docs]
class SpinOrbitShell(Function):
[docs]
@classmethod
def eval(cls, La_tot, Lb_tot, a, b, A, B, C=(0.0, 0.0, 0.0)):
exprs = [
SpinOrbit(*La, *Lb, 0, mu, a, b, A, B, C)
for (La, Lb), mu in it.product(shell_iter((La_tot, Lb_tot)), range(3))
]
# print(SpinOrbit.eval.cache_info())
return exprs
[docs]
def get_center(i):
symbs = [Symbol(str(i) + ind, real=True) for ind in ("x", "y", "z")]
return Matrix([*symbs]).T # Return column vector
[docs]
def get_centers(i):
symbs = [Symbol(str(i) + ind, real=True) for ind in ("x", "y", "z")]
return Matrix([*symbs]).T # Return column vector
[docs]
def get_map(i, center_i):
array = IndexedBase(i, shape=3)
array_map = dict(zip(center_i, array))
return array, array_map
[docs]
def cart2spherical(L_tots, exprs):
assert len(L_tots) > 0
# Coefficient matrices for Cartesian-to-spherical conversion
coeffs = [Array(CART2SPH[L]) for L in L_tots]
cart_shape = [(l + 1) * (l + 2) // 2 for l in L_tots]
cart = Array(exprs).reshape(*cart_shape)
sph = tc(tp(coeffs[0], cart), (1, 2))
if len(L_tots) == 2:
sph = tc(tp(sph, coeffs[1].transpose()), (1, 2))
elif len(L_tots) == 3:
_, Cb, Cc = coeffs
sph = tc(tp(permutedims(sph, (0, 2, 1)), Cb.transpose()), (2, 3))
sph = tc(tp(permutedims(sph, (0, 2, 1)), Cc.transpose()), (2, 3))
else:
raise Exception(
"Cartesian -> spherical transformation for 4-center integrals "
"is not implemented!"
)
# Cartesian-to-spherical transformation introduces quite a number of
# multiplications by 1.0, which are uneccessary. Here, we try to drop
# some of them by replacing numbers very close to +1.0 with 1.
sph = sph.replace(lambda n: n.is_Number and (abs(n - 1) <= ONE_THRESH), lambda n: 1)
# TODO: maybe something along the lines
# sph = map(lambda expr: expr.evalf(), flatten(sph))
# is faster?
return flatten(sph)
[docs]
def gen_integral_exprs(
int_func,
L_maxs,
kind,
maps=None,
sph=False,
):
if maps is None:
maps = list()
ranges = [range(L + 1) for L in L_maxs]
for L_tots in it.product(*ranges):
time_str = time.strftime("%H:%M:%S")
start = datetime.now()
print(f"{time_str} - Generating {L_tots} {kind}")
sys.stdout.flush()
# Generate actual list of expressions.
exprs = int_func(*L_tots)
print("\t... generated expressions")
sys.stdout.flush()
# if sph:
# exprs = cart2spherical(L_tots, exprs)
# print("\t... did Cartesian -> Spherical conversion")
# sys.stdout.flush()
# Common subexpression elimination
repls, reduced = cse(list(exprs), order="none")
print("\t... did common subexpression elimination")
# Replacement expressions, used to form the reduced expressions.
for i, (lhs, rhs) in enumerate(repls):
rhs = rhs.evalf()
# Replace occurences of Ax, Ay, Az, ... with A[0], A[1], A[2], ...
rhs = functools.reduce(lambda rhs, map_: rhs.xreplace(map_), maps, rhs)
repls[i] = (lhs, rhs)
# Reduced expression, i.e., the final integrals/expressions.
for i, red in enumerate(reduced):
red = red.evalf()
reduced[i] = functools.reduce(
lambda red, map_: red.xreplace(map_), maps, red
)
# Carry out Cartesian-to-spherical transformation, if requested.
if sph:
reduced = cart2spherical(L_tots, reduced)
print("\t... did Cartesian -> Spherical conversion")
sys.stdout.flush()
dur = datetime.now() - start
print(f"\t... finished in {str(dur)} h")
sys.stdout.flush()
yield (repls, reduced), L_tots
[docs]
def render_py_funcs(exprs_Ls, args, base_name, doc_func, add_imports=None, comment=""):
if add_imports is None:
add_imports = ()
add_imports_str = "\n".join(add_imports)
if add_imports:
add_imports_str += "\n\n"
args = ", ".join((map(str, args)))
py_funcs = list()
for (repls, reduced), L_tots in exprs_Ls:
doc_str = doc_func(L_tots)
doc_str += "\n\nGenerated code; DO NOT modify by hand!"
name = base_name + "_" + "".join(str(l) for l in L_tots)
print(f"Rendering '{name}' ... ", end="")
start = time.time()
py_func = make_py_func(repls, reduced, args=args, name=name, doc_str=doc_str)
dur = time.time() - start
print(f"finished in {dur: >8.2f} s")
py_funcs.append(py_func)
py_funcs_joined = "\n\n".join(py_funcs)
if comment != "":
comment = f'"""\n{comment}\n"""\n\n'
np_tpl = Template(
"import numpy\n\n{{ add_imports }}{{ comment }}{{ py_funcs }}",
trim_blocks=True,
lstrip_blocks=True,
)
np_rendered = np_tpl.render(
comment=comment, add_imports=add_imports_str, py_funcs=py_funcs_joined
)
np_rendered = textwrap.dedent(np_rendered)
try:
from black import format_str, FileMode
np_rendered = format_str(np_rendered, mode=FileMode(line_length=90))
except ModuleNotFoundError:
print("Skipped formatting with black, as it is not installed!")
return np_rendered
[docs]
def render_c_funcs(exprs_Ls, args, base_name, doc_func, add_imports=None, comment=""):
if add_imports is not None:
raise Exception("Implement me!")
arg_strs = [str(arg) for arg in args]
funcs = list()
signatures = list()
for (repls, reduced), L_tots in exprs_Ls:
doc_str = doc_func(L_tots)
doc_str += "\n\n\t\tGenerated code; DO NOT modify by hand!"
doc_str = textwrap.dedent(doc_str)
name = base_name + "_" + "".join(str(l) for l in L_tots)
func, signature = make_c_func(
repls, reduced, args=arg_strs, name=name, doc_str=doc_str
)
funcs.append(func)
signatures.append(signature)
funcs_joined = "\n\n".join(funcs)
if comment != "":
comment = f"/*{comment}*/"
# Render C files
c_tpl = Template(
"#include <math.h>\n\n{{ comment }}\n\n{{ funcs }}",
trim_blocks=True,
lstrip_blocks=True,
)
c_rendered = c_tpl.render(comment=comment, funcs=funcs_joined)
c_rendered = textwrap.dedent(c_rendered)
# Render simple header file
h_rendered = "\n".join([f"{sig};" for sig in signatures])
return c_rendered, h_rendered
[docs]
def write_file(out_dir, name, rendered):
out_name = out_dir / name
with open(out_name, "w") as handle:
handle.write(rendered)
print(f"Wrote '{out_name}'.")
[docs]
def write_render(
ints_Ls,
args,
name,
doc_func,
out_dir,
comment="",
py_kwargs=None,
c=True,
c_kwargs=None,
):
if py_kwargs is None:
py_kwargs = {}
if c_kwargs is None:
c_kwargs = {}
ints_Ls = list(ints_Ls)
# Python
py_rendered = render_py_funcs(
ints_Ls, args, name, doc_func, comment=comment, **py_kwargs
)
write_file(out_dir, f"{name}.py", py_rendered)
# C
if c:
c_rendered, h_rendered = render_c_funcs(
ints_Ls,
args,
name,
doc_func,
comment=comment,
**c_kwargs,
)
write_file(out_dir, f"{name}.c", c_rendered)
write_file(out_dir, f"{name}.h", h_rendered)
[docs]
def parse_args(args):
parser = argparse.ArgumentParser()
parser.add_argument(
"--lmax",
default=L_MAX,
type=int,
help="Generate 1e-integrals up to this maximum angular momentum.",
)
parser.add_argument(
"--lauxmax",
default=L_AUX_MAX,
type=int,
help="Maximum angular moment for integrals using auxiliary functions.",
)
parser.add_argument(
"--write",
action="store_true",
help="Write out generated integrals to the current directory, potentially overwriting the present modules.",
)
parser.add_argument(
"--out-dir",
default="devel_ints",
help="Directory, where the generated integrals are written.",
)
keys_str = f"({', '.join(KEYS)})"
parser.add_argument(
"--keys",
nargs="+",
help=f"Generate only certain expressions. Possible keys are: {keys_str}. "
"If not given, all expressions are generated.",
)
return parser.parse_args()
[docs]
def run():
args = parse_args(sys.argv[1:])
l_max = args.lmax
l_aux_max = args.lauxmax
out_dir = Path(args.out_dir if not args.write else ".")
keys = args.keys
if keys is None:
keys = list()
try:
os.mkdir(out_dir)
except FileExistsError:
pass
try:
global CART2SPH
CART2SPH = cart2sph_coeffs(max(l_max, l_aux_max), zero_small=True)
except NameError:
print("cart2sph_coeffs import is deactivated or pysisyphus is not installed.")
# Cartesian basis function centers A and B.
center_A = get_center("A")
center_B = get_center("B")
center_C = get_center("C")
center_D = get_center("D")
# center_R = get_center("R")
Xa, Ya, Za = symbols("Xa Ya Za")
# Orbital exponents a, b, c, d.
a, b, c, d = symbols("a b c d", real=True)
# These maps will be used to convert {Ax, Ay, ...} to array quantities
# in the generated code. This way an iterable/np.ndarray can be used as
# function argument instead of (Ax, Ay, Az, Bx, By, Bz).
A, A_map = get_map("A", center_A)
B, B_map = get_map("B", center_B)
C, C_map = get_map("C", center_C)
D, D_map = get_map("D", center_D)
boys_import = ("from pysisyphus.wavefunction.ints.boys import boys",)
#################
# Cartesian GTO #
#################
def cart_gto():
def cart_gto_doc_func(L_tot):
(La_tot,) = L_tot
shell_a = L_MAP[La_tot]
return (
f"3D Cartesian {shell_a}-Gaussian shell.\n\n"
"Exponent a, centered at A, evaluated at (Xa, Ya, Za) + A."
)
# This code can evaluate multiple points at a time
cart_gto_Ls = gen_integral_exprs(
lambda La_tot: CartGTOShell(La_tot, a, Xa, Ya, Za),
(l_max,),
"cart_gto",
)
cart_gto_rendered = render_py_funcs(
cart_gto_Ls, (a, Xa, Ya, Za), "cart_gto3d", cart_gto_doc_func
)
write_file(out_dir, "gto3d.py", cart_gto_rendered)
print()
#####################
# Overlap integrals #
#####################
def overlap():
def ovlp_doc_func(L_tots):
La_tot, Lb_tot = L_tots
shell_a = L_MAP[La_tot]
shell_b = L_MAP[Lb_tot]
return f"Cartesian 3D ({shell_a}{shell_b}) overlap integral."
ovlp_ints_Ls = gen_integral_exprs(
lambda La_tot, Lb_tot: Overlap3dShell(
La_tot, Lb_tot, a, b, center_A, center_B
),
(l_max, l_max),
"overlap",
(A_map, B_map),
)
write_render(
ovlp_ints_Ls, (a, A, b, B), "ovlp3d", ovlp_doc_func, out_dir, c=True
)
print()
###########################
# Dipole moment integrals #
###########################
def dipole():
def dipole_doc_func(L_tots):
La_tot, Lb_tot = L_tots
shell_a = L_MAP[La_tot]
shell_b = L_MAP[Lb_tot]
return (
f"Cartesian 3D ({shell_a}{shell_b}) dipole moment integrals.\n"
"The origin is at C."
)
dipole_comment = """
Dipole integrals are given in the order:
for bf_a in basis_functions_a:
for bf_b in basis_functions_b:
for cart_dir in (x, y, z):
dipole_integrals(bf_a, bf_b, cart_dir)
So for <s_a|μ|s_b> it will be:
<s_a|x|s_b>
<s_a|y|s_b>
<s_a|z|s_b>
"""
dipole_ints_Ls = gen_integral_exprs(
lambda La_tot, Lb_tot: Multipole3dShell(
La_tot, Lb_tot, a, b, center_A, center_B, 1, center_C
),
(l_max, l_max),
"dipole moment",
(A_map, B_map, C_map),
)
write_render(
dipole_ints_Ls,
(a, A, b, B, C),
"dipole3d",
dipole_doc_func,
out_dir,
comment=dipole_comment,
c=True,
)
print()
###########################################
# Diagonal of quadrupole moment integrals #
###########################################
def diag_quadrupole():
def diag_quadrupole_doc_func(L_tots):
La_tot, Lb_tot = L_tots
shell_a = L_MAP[La_tot]
shell_b = L_MAP[Lb_tot]
return (
f"Cartesian 3D ({shell_a}{shell_b}) quadrupole moment integrals\n"
"for operators x², y² and z². The origin is at C."
)
diag_quadrupole_comment = """
Diagonal of the quadrupole moment matrix with operators x², y², z².
for rr in (xx, yy, zz):
for bf_a in basis_functions_a:
for bf_b in basis_functions_b:
quadrupole_integrals(bf_a, bf_b, rr)
"""
diag_quadrupole_ints_Ls = gen_integral_exprs(
lambda La_tot, Lb_tot: DiagQuadrupole3dShell(
La_tot, Lb_tot, a, b, center_A, center_B, center_C
),
(l_max, l_max),
"diag quadrupole moment",
(A_map, B_map, C_map),
)
write_render(
diag_quadrupole_ints_Ls,
(a, A, b, B, C),
"diag_quadrupole3d",
diag_quadrupole_doc_func,
out_dir,
comment=diag_quadrupole_comment,
c=True,
)
print()
###############################
# Quadrupole moment integrals #
###############################
def quadrupole():
def quadrupole_doc_func(L_tots):
La_tot, Lb_tot = L_tots
shell_a = L_MAP[La_tot]
shell_b = L_MAP[Lb_tot]
return (
f"Cartesian 3D ({shell_a}{shell_b}) quadrupole moment integrals.\n"
"The origin is at C."
)
quadrupole_comment = """
Quadrupole integrals contain the upper triangular part of the symmetric
3x3 quadrupole matrix.
/ xx xy xz \\
| yy yz |
\ zz /
"""
quadrupole_ints_Ls = gen_integral_exprs(
lambda La_tot, Lb_tot: Multipole3dShell(
La_tot, Lb_tot, a, b, center_A, center_B, 2, center_C
),
(l_max, l_max),
"quadrupole moment",
(A_map, B_map, C_map),
)
write_render(
quadrupole_ints_Ls,
(a, A, b, B, C),
"quadrupole3d",
quadrupole_doc_func,
out_dir,
comment=quadrupole_comment,
c=True,
)
print()
############################
# Kinetic energy integrals #
############################
def kinetic():
def kinetic_doc_func(L_tots):
La_tot, Lb_tot = L_tots
shell_a = L_MAP[La_tot]
shell_b = L_MAP[Lb_tot]
return f"Cartesian 3D ({shell_a}{shell_b}) kinetic energy integral."
kinetic_ints_Ls = gen_integral_exprs(
lambda La_tot, Lb_tot: Kinetic3dShell(
La_tot, Lb_tot, a, b, center_A, center_B
),
(l_max, l_max),
"kinetic",
(A_map, B_map),
)
write_render(
kinetic_ints_Ls,
(a, A, b, B),
"kinetic3d",
kinetic_doc_func,
out_dir,
c=True,
)
print()
#########################
# 1el Coulomb Integrals #
#########################
def coulomb():
def coulomb_doc_func(L_tots):
La_tot, Lb_tot = L_tots
shell_a = L_MAP[La_tot]
shell_b = L_MAP[Lb_tot]
return f"Cartesian ({shell_a}{shell_b}) 1-electron Coulomb integral."
coulomb_ints_Ls = gen_integral_exprs(
lambda La_tot, Lb_tot: CoulombShell(
La_tot, Lb_tot, a, b, center_A, center_B, center_C
),
(l_max, l_max),
"coulomb3d",
(A_map, B_map, C_map),
)
coulomb_rendered = render_py_funcs(
coulomb_ints_Ls,
(a, A, b, B, C),
"coulomb3d",
coulomb_doc_func,
add_imports=boys_import,
)
# TODO: handle add_args and Boys function in C
write_file(out_dir, "coulomb3d.py", coulomb_rendered)
print()
###############################################
# Two-center two-electron repulsion integrals #
###############################################
def _2center2electron():
def _2center2el_doc_func(L_tots):
La_tot, Lb_tot = L_tots
shell_a = L_MAP[La_tot]
shell_b = L_MAP[Lb_tot]
return (
f"Cartesian ({shell_a}|{shell_b}) "
"two-center two-electron repulsion integral."
)
_2center2el_ints_Ls = gen_integral_exprs(
lambda La_tot, Lb_tot: TwoCenterTwoElectronShell(
La_tot,
Lb_tot,
a,
b,
center_A,
center_B,
),
(l_aux_max, l_aux_max),
"_2center2el3d",
(A_map, B_map),
)
write_render(
_2center2el_ints_Ls,
(a, A, b, B),
"_2center2el3d",
_2center2el_doc_func,
out_dir,
c=False,
py_kwargs={"add_imports": boys_import},
)
print()
#################################################
# Three-center two-electron repulsion integrals #
#################################################
def _3center2electron():
def _3center2el_doc_func(L_tots):
La_tot, Lb_tot, Lc_tot = L_tots
shell_a = L_MAP[La_tot]
shell_b = L_MAP[Lb_tot]
shell_c = L_MAP[Lc_tot]
return (
f"Cartesian ({shell_a}{shell_b}|{shell_c}) "
"three-center two-electron repulsion integral."
)
_3center2el_ints_Ls = gen_integral_exprs(
lambda La_tot, Lb_tot, Lc_tot: ThreeCenterTwoElectronShell(
La_tot, Lb_tot, Lc_tot, a, b, c, center_A, center_B, center_C
),
(l_max, l_max, l_aux_max),
"_3center2el3d",
(A_map, B_map, C_map),
)
write_render(
_3center2el_ints_Ls,
(a, A, b, B, c, C),
"_3center2el3d",
_3center2el_doc_func,
out_dir,
c=False,
py_kwargs={"add_imports": boys_import},
)
print()
def _3center2electron_sph():
def _3center2el_doc_func(L_tots):
La_tot, Lb_tot, Lc_tot = L_tots
shell_a = L_MAP[La_tot]
shell_b = L_MAP[Lb_tot]
shell_c = L_MAP[Lc_tot]
return (
f"Cartesian ({shell_a}{shell_b}|{shell_c}) "
"three-center two-electron repulsion integral for conversion to "
"spherical harmonics.\nUses Ahlrichs' vertical recursion relation, "
"that leaves out some terms, that cancel\nwhen convertig to "
"spherical harmonics."
)
_3center2el_ints_Ls = gen_integral_exprs(
lambda La_tot, Lb_tot, Lc_tot: ThreeCenterTwoElectronSphShell(
La_tot, Lb_tot, Lc_tot, a, b, c, center_A, center_B, center_C
),
(l_max, l_max, l_aux_max),
"_3center2el3d_sph",
(A_map, B_map, C_map),
sph=False,
)
write_render(
_3center2el_ints_Ls,
(a, A, b, B, c, C),
"_3center2el3d_sph",
_3center2el_doc_func,
out_dir,
c=False,
py_kwargs={"add_imports": boys_import},
)
print()
####################################
# Spin-orbit interaction integrals #
####################################
# def so1el_doc_func(L_tots):
# La_tot, Lb_tot = L_tots
# shell_a = L_MAP[La_tot]
# shell_b = L_MAP[Lb_tot]
# return f"Cartesian ({shell_a}{shell_b}) 1-electron spin-orbit-interaction integral."
# I think this is still faulty!
# so1el_ints_Ls = gen_integral_exprs(
# lambda La_tot, Lb_tot: SpinOrbitShell(
# La_tot, Lb_tot, a, b, center_A, center_B, center_C
# ),
# (l_max, l_max),
# "so1el",
# (A_map, B_map, C_map),
# )
# so1el_rendered = render_py_funcs(
# so1el_ints_Ls,
# (a, A, b, B, C),
# "so1el",
# so1el_doc_func,
# add_imports=boys_import,
# )
# write_file(out_dir, "so1el.py", so1el_rendered)
# print()
#########################
# 2el Coulomb Integrals #
#########################
# def eri_doc_func(L_tots):
# La_tot, Lb_tot, Lc_tot, Ld_tot = L_tots
# shell_a = L_MAP[La_tot]
# shell_b = L_MAP[Lb_tot]
# shell_c = L_MAP[Lc_tot]
# shell_d = L_MAP[Ld_tot]
# return (
# f"Cartesian [{shell_a}{shell_b}|{shell_c}{shell_d}] "
# "2-electron electron repulsion integral."
# )
# I think this is still faulty!
# eri_ints_Ls = gen_integral_exprs(
# lambda La_tot, Lb_tot, Lc_tot, Ld_tot: ERIShell(
# La_tot,
# Lb_tot,
# Lc_tot,
# Ld_tot,
# a,
# b,
# c,
# d,
# center_A,
# center_B,
# center_C,
# center_D,
# ),
# # (l_max, l_max, l_max, l_max),
# (2, 2, 2, 2), # Stop at [dd|dd] for now
# "eri",
# (A_map, B_map, C_map, D_map),
# )
# eri_rendered = render_py_funcs(
# eri_ints_Ls,
# (a, A, b, B, c, C, d, D),
# "eri",
# eri_doc_func,
# add_imports=boys_import,
# )
# write_file(out_dir, "eri.py", eri_rendered)
# print()
funcs = {
"cgto": cart_gto, # Cartesian Gaussian-type-orbital for density evaluation
"ovlp": overlap, # Overlap integrals
"dpm": dipole, # Linear moment (dipole) integrals
"dqpm": diag_quadrupole, # Diagonal part of the quadrupole tensor
"qpm": quadrupole, # Quadratic moment (quadrupole) integrals
"kin": kinetic, # Kinetic energy integrals
"coul": coulomb, # 1-electron Coulomb integrals
# Integrals for density fitting
"2c2e": _2center2electron, # 2-center-2-electron integrals for DF
"3c2e": _3center2electron, # 3-center-2-electron integrals for density fitting
"3c2e_sph": _3center2electron_sph, # 3c2el for spherical transformation for DF
}
if len(keys) == 0:
keys = funcs.keys()
for key in keys:
funcs[key]()
if __name__ == "__main__":
run()