mod: use numba multipole backend by default

currently, this only supports multipole integrals to arbitrary order,
but they are much faster than the python/numpy integrals and can be
evaluated in parallel. the ability to evaluate (overlap, dipole and
quadrupole) integrals will already bring us very far. Overlap integrals
are used for the Wavefunction sanity check, dipole integrals are used to
calculate (transition) dipole moments used in spectra plotting and
diabatization and quadrupole moment integrals are also used in
diabatization.

pysisyphus/wavefunction/backend_numba_multipole.py

@@ -202,7 +202,7 @@ def get_multipole_ints_cart_numba(
     shells_b,
     symmetric,
 ):
-    components = 2 * Le + 1
+    components = (Le + 2) * (Le + 1) // 2
 
     tot_size_a = 0
     for shell in shells_a:
@@ -214,7 +214,6 @@ def get_multipole_ints_cart_numba(
 
     # Allocate final integral array
     integrals = np.zeros((tot_size_a, tot_size_b, components))
-    shells_b = shells_a
     nshells_a = len(shells_a)
     nshells_b = len(shells_b)
 
@@ -261,8 +260,12 @@ def get_multipole_ints_cart(
         shellstructs_b,
         symmetric,
     )
-    if integrals.shape[2] == 1:
-        integrals = np.squeeze(integrals, axis=2)
+    # TODO: Drop the transpose and restructure the code to handle
+    # (shellsa, shellsb, multipoles), as this is more efficient.
+    # See issue #312 on Github.
+    integrals = np.transpose(integrals, axes=(2, 0, 1))
+    if integrals.shape[0] == 1:
+        integrals = np.squeeze(integrals, axis=0)
     return integrals
 
 

pysisyphus/wavefunction/shells.py

@@ -208,7 +208,7 @@ def __init__(
         shells: List[Shell],
         screen: bool = False,
         ordering: Ordering = "native",
-        backend: Union[str, IntegralBackend] = IntegralBackend.PYTHON,
+        backend: Union[str, IntegralBackend] = IntegralBackend.NUMBA_MULTIPOLE,
     ):
         self.shells = shells
         self.ordering = ordering

pysisyphus/wavefunction/wavefunction.py

@@ -496,6 +496,7 @@ def dipole_ints(
     ) -> NDArray[float]:
         if origin is None:
             origin = self.get_origin(kind=kind)
+        origin = np.array(origin)
 
         dipole_ints = {
             BFType.CARTESIAN: lambda *args: self.shells.get_dipole_ints_cart(*args),
@@ -508,6 +509,7 @@ def quadrupole_ints(
     ) -> NDArray[float]:
         if origin is None:
             origin = self.get_origin(kind=kind)
+        origin = np.array(origin)
 
         quadrupole_ints = {
             BFType.CARTESIAN: lambda *args: self.shells.get_quadrupole_ints_cart(*args),

tests/test_wavefunction/test_numba_backend.py

@@ -45,6 +45,7 @@ def py_wf():
     return Wavefunction.from_file(
         "lib:01_ch4_tzvp.json",
         shell_kwargs={
+            "backend": "python",
             "ordering": "pysis",
         },
     )

tests/test_wavefunction/test_shell.py

@@ -57,7 +57,7 @@ def _charge_center(mol):
     origin = _charge_center(mol)
     with mol.with_common_orig(origin):
         quad = mol.intor("int1e_rr").reshape(3, 3, nao, nao)
-    shells = Shells.from_pyscf_mol(mol)
+    shells = Shells.from_pyscf_mol(mol, backend="python")
     pysis_quad = shells.get_quadrupole_ints_sph(origin)
     np.testing.assert_allclose(pysis_quad, quad, atol=1e-14)
 
@@ -112,7 +112,7 @@ def test_2c2e(mol_auxmol):
     N_aux = 1 / np.diag(S_aux) ** 0.5
     NaNa = N_aux[:, None] * N_aux[None, :]
 
-    aux_shells = Shells.from_pyscf_mol(auxmol)
+    aux_shells = Shells.from_pyscf_mol(auxmol, backend="python")
     integrals = aux_shells.get_2c2el_ints_cart()
 
     np.testing.assert_allclose(integrals, int2c * NaNa)
@@ -122,7 +122,7 @@ def test_3c2e(mol_auxmol):
     mol, auxmol = mol_auxmol
     int3c = df.incore.aux_e2(mol, auxmol, "int3c2e_sph")
 
-    shells = Shells.from_pyscf_mol(mol)
+    shells = Shells.from_pyscf_mol(mol, backend="python")
     aux_shells = Shells.from_pyscf_mol(auxmol)
     integrals = shells.get_3c2el_ints_sph(aux_shells)
 

tests/test_wavefunction/test_wavefunction.py

@@ -175,7 +175,7 @@ def test_quadrupole_moment():
 
 def test_one_el_terms():
     fn = WF_LIB_DIR / "orca_ch4_sto3g.json"
-    wf = Wavefunction.from_orca_json(fn)
+    wf = Wavefunction.from_orca_json(fn, shell_kwargs={"backend": "python"})
     shells = wf.shells
     T = shells.T_sph
     V = shells.V_sph
@@ -188,7 +188,9 @@ def test_one_el_terms():
 
 
 def test_eval_esp():
-    wf = Wavefunction.from_file(WF_LIB_DIR / "orca_lih_ccpvdz.molden.input")
+    wf = Wavefunction.from_file(
+        WF_LIB_DIR / "orca_lih_ccpvdz.molden.input", shell_kwargs={"backend": "python"}
+    )
     coords3d = np.array(
         (
             (0.0, 0.0, 1.0),