Fix autoaux

pyscf/df/addons.py

@@ -30,6 +30,9 @@
 ETB_BETA = getattr(__config__, 'df_addons_aug_dfbasis', 2.0)
 FIRST_ETB_ELEMENT = getattr(__config__, 'df_addons_aug_start_at', 36)  # 'Rb'
 
+LVAL_SCHEME = 'PySCF'  # or 'Gaussian', 'ORCA'
+L_INC = 1
+
 # Obtained from http://www.psicode.org/psi4manual/master/basissets_byfamily.html
 DEFAULT_AUXBASIS = {
     # AO basis       JK-fit                     MP2-fit
@@ -73,6 +76,69 @@ class load(ao2mo.load):
     def __init__(self, eri, dataname='j3c'):
         ao2mo.load.__init__(self, eri, dataname)
 
+def _aug_etb_element(nuc_charge, basis, beta):
+    # lval
+    if LVAL_SCHEME == 'Gaussian':
+        # Yang, JCP, 127, 074102
+        # 0: H - Be, 1: B - Ca, 2: Sc - Ba, 3: La -
+        if nuc_charge <= 2:
+            max_shells = 0
+        elif nuc_charge <= 18:
+            max_shells = 1
+        elif nuc_charge <= 54:
+            max_shells = 2
+        else:
+            max_shells = 3
+    elif LVAL_SCHEME == 'ORCA':
+        # Stoychev, JCTC, 13, 554
+        # 0: H - Be, 1: B - Ca, 2: Sc - Ba, 3: La -
+        conf = elements.NRSRHFS_CONFIGURATION[nuc_charge]
+        max_shells = max(3 - conf.count(0), 0)
+    else:
+        # PySCF original settings
+        # 1: H - Be, 2: B - Ca, 3: Sc - La, 4: Ce -
+        conf = elements.CONFIGURATION[nuc_charge]
+        max_shells = 4 - conf.count(0)
+
+    emin_by_l = [1e99] * 8
+    emax_by_l = [0] * 8
+    l_max = 0
+    for b in basis:
+        l = b[0]
+        l_max = max(l_max, l)
+
+        if isinstance(b[1], int):
+            e_c = numpy.array(b[2:])
+        else:
+            e_c = numpy.array(b[1:])
+        es = e_c[:,0]
+        cs = e_c[:,1:]
+        es = es[abs(cs).max(axis=1) > 1e-3]
+        emax_by_l[l] = max(es.max(), emax_by_l[l])
+        emin_by_l[l] = min(es.min(), emin_by_l[l])
+
+    l_max1 = l_max + 1
+    emin_by_l = numpy.array(emin_by_l[:l_max1])
+    emax_by_l = numpy.array(emax_by_l[:l_max1])
+
+    # Estimate the exponents ranges by geometric average
+    emax = numpy.sqrt(numpy.einsum('i,j->ij', emax_by_l, emax_by_l))
+    emin = numpy.sqrt(numpy.einsum('i,j->ij', emin_by_l, emin_by_l))
+    liljsum = numpy.arange(l_max1)[:,None] + numpy.arange(l_max1)
+    # Setting l_max_aux as JCTC, 13, 554
+    l_max_aux = max(max_shells*2, l_max + L_INC)
+    emax_by_l = [emax[liljsum==ll].max() for ll in range(l_max_aux+1)]
+    emin_by_l = [emin[liljsum==ll].min() for ll in range(l_max_aux+1)]
+    # Tune emin and emax
+    emin_by_l = numpy.array(emin_by_l) * 2
+    emax_by_l = numpy.array(emax_by_l) * 2
+
+    ns = numpy.log((emax_by_l+emin_by_l)/emin_by_l) / numpy.log(beta)
+    etb = []
+    for l, n in enumerate(numpy.ceil(ns).astype(int)):
+        if n > 0:
+            etb.append((l, n, emin_by_l[l], beta))
+    return etb
 
 def aug_etb_for_dfbasis(mol, dfbasis=DFBASIS, beta=ETB_BETA,
                         start_at=FIRST_ETB_ELEMENT):
@@ -88,46 +154,8 @@ def aug_etb_for_dfbasis(mol, dfbasis=DFBASIS, beta=ETB_BETA,
         if nuc_charge < nuc_start:
             newbasis[symb] = dfbasis
         else:
-            conf = elements.CONFIGURATION[nuc_charge]
-            max_shells = 4 - conf.count(0)
-            emin_by_l = [1e99] * 8
-            emax_by_l = [0] * 8
-            l_max = 0
-            for b in mol._basis[symb]:
-                l = b[0]
-                l_max = max(l_max, l)
-                if l >= max_shells+1:
-                    continue
-
-                if isinstance(b[1], int):
-                    e_c = numpy.array(b[2:])
-                else:
-                    e_c = numpy.array(b[1:])
-                es = e_c[:,0]
-                cs = e_c[:,1:]
-                es = es[abs(cs).max(axis=1) > 1e-3]
-                emax_by_l[l] = max(es.max(), emax_by_l[l])
-                emin_by_l[l] = min(es.min(), emin_by_l[l])
-
-            l_max1 = l_max + 1
-            emin_by_l = numpy.array(emin_by_l[:l_max1])
-            emax_by_l = numpy.array(emax_by_l[:l_max1])
-
-# Estimate the exponents ranges by geometric average
-            emax = numpy.sqrt(numpy.einsum('i,j->ij', emax_by_l, emax_by_l))
-            emin = numpy.sqrt(numpy.einsum('i,j->ij', emin_by_l, emin_by_l))
-            liljsum = numpy.arange(l_max1)[:,None] + numpy.arange(l_max1)
-            emax_by_l = [emax[liljsum==ll].max() for ll in range(l_max1*2-1)]
-            emin_by_l = [emin[liljsum==ll].min() for ll in range(l_max1*2-1)]
-            # Tune emin and emax
-            emin_by_l = numpy.array(emin_by_l) * 2
-            emax_by_l = numpy.array(emax_by_l) * 2
-
-            ns = numpy.log((emax_by_l+emin_by_l)/emin_by_l) / numpy.log(beta)
-            etb = []
-            for l, n in enumerate(numpy.ceil(ns).astype(int)):
-                if n > 0:
-                    etb.append((l, n, emin_by_l[l], beta))
+            basis = mol._basis[symb]
+            etb = _aug_etb_element(nuc_charge, basis, beta)
             if etb:
                 newbasis[symb] = gto.expand_etbs(etb)
             else:
@@ -233,7 +261,7 @@ def make_auxmol(mol, auxbasis=None):
             etb_abs = aug_etb_for_dfbasis(mol, start_at=0)
             for elem, bas in _basis.items():
                 if bas == 'etb':
-                    _basis[elem] = auxbasis_etb[elem]
+                    _basis[elem] = etb_abs[elem]
 
     pmol._basis = pmol.format_basis(_basis)
 

pyscf/df/autoaux.py

@@ -0,0 +1,118 @@
+#!/usr/bin/env python
+# Copyright 2024 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+'''
+The AutoAux algorithm by ORCA
+
+Ref:
+    JCTC, 13 (2016), 554
+'''
+
+from math import factorial
+import np as np
+from pyscf import gto
+
+F_LAUX = np.array([20, 4.0, 4.0, 3.5, 2.5, 2.0, 2.0])
+BETA_BIG = np.array([1.8, 2.0, 2.2, 2.2, 2.2, 2.3, 3.0, 3.0])
+BETA_SMALL = 1.8
+L_INC = 1
+
+def _primitive_emin_emax(basis):
+    l_max = max(b[0] for b in basis)
+    emin_by_l = [1e99] * (l_max+1)
+    emax_by_l = [0] * (l_max+1)
+    a_eff_by_l = [0] * (l_max+1)
+
+    for b in basis:
+        l = b[0]
+        if isinstance(b[1], int):
+            e_c = np.array(b[2:])
+        else:
+            e_c = np.array(b[1:])
+        es = e_c[:,0]
+        cs = e_c[:,1:]
+        cs = np.einsum('pi,p->pi', cs, gto.gto_norm(l, es))
+        emax_by_l[l] = max(es.max(), emax_by_l[l])
+        emin_by_l[l] = min(es.min(), emin_by_l[l])
+
+        ee = es[:,None] + es
+        r_ints = gto.gaussian_int(l*2+3, ee)
+        r_exp = np.einsum('pi,pq,qi->i', cs, r_ints, cs)
+
+        k = 2**(2*l+1) * (factorial(l+1))**2 / factorial(2*l+2)
+        a_eff_by_l[l] = max(2 * k**2 / (np.pi * r_exp.min()),
+                            a_eff_by_l)
+    return np.array(emax_by_l), np.array(emin_by_l), np.array(a_eff_by_l)
+
+def _auto_aux_element(Z, basis):
+    a_max_by_l, a_min_by_l, a_eff_by_l = _primitive_emin_emax(basis)
+    l_max1 = a_max_by_l.size
+
+    l_max = l_max1 - 1
+    if Z <= 2:
+        l_val = 0
+    elif Z <= 18:
+        l_val = 1
+    elif Z <= 54:
+        l_val = 2
+    else:
+        l_val = 3
+    l_max_aux = max(l_val*2, l_max+L_INC)
+
+    a_min_prim = a_min_by_l[:,None] + a_min_by_l
+    a_max_prim = a_max_by_l[:,None] + a_max_by_l
+    a_max_aux = a_eff_by_l[:,None] + a_eff_by_l
+
+    liljsum = np.arange(l_max1)[:,None] + np.arange(l_max1)
+    a_min_by_l = [a_min_prim[liljsum==ll].min() for ll in range(l_max_aux+1)]
+    a_max_by_l = [a_max_prim[liljsum==ll].max() for ll in range(l_max_aux+1)]
+    a_aux_by_l = [a_max_aux [liljsum==ll].max() for ll in range(l_max_aux+1)]
+
+    a_max_adjust = [max(F_LAUX[l] * a_aux_by_l[l], a_max_by_l[l])
+                    for l in range(l_val*2+1)]
+    a_max_adjust = a_max_adjust + a_aux_by_l[l_val*2+1 : l_max_aux+1]
+
+    emin = np.array(a_min_by_l)
+    emax = np.array(a_max_adjust) + 1e-3
+
+    ns_small = np.log(emax[:l_val*2+1] / emin[:l_val*2+1]) / np.log(BETA_SMALL)
+    etb = []
+    for l, n in enumerate(np.ceil(ns_small).astype(int)):
+        if n > 0:
+            etb.append((l, n, emin[l], BETA_SMALL))
+
+    if l_max_aux > l_val*2:
+        ns_big = (np.log(emax[l_val*2+1:] / emin[l_val*2+1:])
+                  / np.log(BETA_BIG[l_val*2+1:]))
+        for l, n in enumerate(np.ceil(ns_big).astype(int)):
+            if n > 0:
+                l = l + l_val*2+1
+                beta = BETA_BIG[l]
+                etb.append((l, n, emin[l], beta))
+    return etb
+
+def auto_aux(mol):
+    uniq_atoms = {a[0] for a in mol._atom}
+    newbasis = {}
+    for symb in uniq_atoms:
+        Z = gto.charge(symb)
+        basis = mol._basis[symb]
+        etb = _auto_aux_element(Z, basis)
+        if etb:
+            newbasis[symb] = etb
+        else:
+            raise RuntimeError(f'Failed to generate even-tempered auxbasis for {symb}')
+    return newbasis