from typing import Literal, Union
import numpy as np
from enum import IntEnum
L_MAP = {
"s": 0,
"p": 1,
"d": 2,
"f": 3,
"g": 4,
"h": 5,
}
L_MAP_INV = {key: L for L, key in L_MAP.items()}
Ls = Literal["s", "p", "d", "f", "g", "h"]
L_Inp = Union[int, Ls]
[docs]
def cart_size(L: int) -> int:
return (L + 2) * (L + 1) // 2
[docs]
def sph_size(L: int) -> int:
return 2 * L + 1
[docs]
def get_l(l_inp: L_Inp) -> int:
"""Convert shell label to angular moment quantum number l."""
try:
l = L_MAP[l_inp.lower()]
except (KeyError, AttributeError):
l = int(l_inp)
return l
[docs]
class BFType(IntEnum):
CARTESIAN = 1
PURE_SPHERICAL = 2
[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 cca_order(l):
"""Same as canonical_order()."""
inds = [
(l, 0, 0),
]
a = l
b = c = 0
for _ in range(((l + 1) * (l + 2) // 2) - 1):
if c < l - a:
b -= 1
c += 1
else:
a -= 1
c = 0
b = l - a
inds.append((a, b, c))
return inds
[docs]
def lm_iter(L_max):
for l in range(L_max + 1):
for m in range(-l, l + 1):
yield l, m
[docs]
def cart2ang_moms(cart):
"""Translate 'xxx' to (3, 0, 0) etc."""
ang_moms = np.zeros(3, dtype=int)
for char in cart:
# x: 120, y: 121, z: 122
ang_moms[ord(char.lower()) - 120] += 1
return tuple(ang_moms)
[docs]
def permut_for_order(cart_orders, pysis_order_func=cca_order):
"""Create permutation matrices to reorder Cartesian integrals."""
Ps = dict()
for cart_order in cart_orders:
lengths = [len(_) for _ in cart_order]
l0 = lengths[0]
assert all([length == l0 for length in lengths])
other_order = [cart2ang_moms(_) for _ in cart_order]
pysis_order = pysis_order_func(l0)
num = len(pysis_order)
P = np.zeros((num, num), dtype=int)
for i, po in enumerate(other_order):
ind = pysis_order.index(po)
P[i, ind] = 1
Ps[l0] = P
return Ps
[docs]
def symmetric_orthogonalization(mat, S, thresh=1e-6):
w, v = np.linalg.eigh(mat.T.dot(S).dot(mat))
mask = w >= thresh
w_isqrt = np.sqrt(w[mask])
S_isqrt = (v[:, mask] / w_isqrt).dot(v[:, mask].T)
return mat.dot(S_isqrt)