refactor xtb

Author: fishjojo

pyscfad/backend/_jax/lax/linalg.py

@@ -105,7 +105,7 @@ def _eigh_gen_jvp_rule(primals, tangents, *, lower, itype, deg_thresh):
                   precision=lax.Precision.HIGHEST)
 
     if type(at) is ad_util.Zero:
-        vt_at_v = lax.zeros_like_array(a)
+        vt_at_v = jnp.zeros_like(a)
     else:
         vt_at_v = dot(_H(v), dot(at, v))
 

pyscfad/ml/scf/hf_pad.py

@@ -1,4 +1,4 @@
-# Copyright 2021-2025 The PySCFAD Authors
+# Copyright 2025-2026 The PySCFAD Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -14,224 +14,26 @@
 
 import numpy
 from pyscfad import numpy as np
-from pyscfad.lib import logger
 from pyscfad.scf import hf_lite as hf
 from pyscfad.scipy.linalg import eigh
-from pyscfad.ops import stop_grad
-
-from functools import partial
-from jax import custom_jvp
-from jax.lax import while_loop
-
-def _fermi_entropy(occ):
-    occ = occ / 2.0
-    occ_safe = np.where(np.logical_or(occ < 1e-10, occ > 1 - 1e-10), 0.5, occ)
-    ent_term = occ_safe * np.log(occ_safe) + (1 - occ_safe) * np.log(1 - occ_safe)
-
-    return -2 * np.where(
-        np.logical_or(occ < 1e-10, occ > 1 - 1e-10),
-        0.,
-        ent_term
-    ).sum()
-
-def _fermi_smearing_occ(mu, mo_energy, sigma, mo_mask):
-    de = (mo_energy - mu) / sigma
-    de_ = np.where(np.less(de, 40.), de, 0)
-    occ = np.where(np.less(de, 40.), 1. / (np.exp(de_) + 1.), 0.)
-    occ = np.where(mo_mask, occ, 0)
-    return occ
-
-@custom_jvp
-def _smearing_solve_mu(mo_es, nocc, sigma, mo_mask):
-    def cond_fun(value):
-        mu, nerr = value
-        return nerr**2 > 1e-12
-
-    def body_fun(value):
-        ''' One Halley step '''
-        mu, nerr = value
-        occ = _fermi_smearing_occ(mu, mo_es, sigma, mo_mask)
-        grad = occ * (1.-occ) / sigma
-        hess = grad * (1.-2*occ) / sigma
-        occ = np.sum(occ)
-        nerr = occ - nocc
-        grad = np.sum(grad)
-        hess = np.sum(hess)
-        dmu = -nerr * grad / (grad**2 - .5 * hess * nerr)
-        return mu + dmu, nerr
-
-    mu, nerr = while_loop(cond_fun, body_fun, (mo_es[nocc-1], 1e2))
-    return np.array([mu])
-
-@_smearing_solve_mu.defjvp
-def _smearing_solve_mu_jvp(primals, tangents):
-    mo_es, nocc, sigma, mo_mask = primals
-    dmo_e, _, _, _ = tangents
-
-    mu = _smearing_solve_mu(mo_es, nocc, sigma, mo_mask)
-    occ = _fermi_smearing_occ(mu, mo_es, sigma, mo_mask)
-    dndmu = occ * (1.-occ) / sigma
-    dndmu = np.where(np.abs(dndmu) < 1e-10, 0., dndmu)
-    return mu, np.dot(dndmu, dmo_e)[None] / (1e-13 + np.sum(dndmu))
-
-def _smearing_optimize(mo_es, nocc, sigma, mo_mask):
-    mu = _smearing_solve_mu(mo_es, nocc, sigma, mo_mask)
-    mo_occs = _fermi_smearing_occ(mu, mo_es, sigma, mo_mask)
-    return mu, mo_occs
-
-def make_mo_mask(mo_energy, mo_coeff, ao_mask):
-    mask_fake_ao = np.asarray(1 - ao_mask, dtype=bool)
-    mo_coeff_fake_ao = np.where(mask_fake_ao[:,None], mo_coeff, 0)
-    tmp = np.linalg.norm(mo_coeff_fake_ao, axis=0)
-    #mask = np.where(np.logical_and(mo_energy>1e8, tmp>1e-12), False, True)
-    mask = np.where(np.logical_and(abs(mo_energy)<1e-12, tmp>1e-12), False, True)
-    return mask
-
-def get_occ(mf, mo_energy=None, mo_coeff=None):
-    # NOTE assuming mo_energy is in ascending order
-    if mo_energy is None:
-        mo_energy = mf.mo_energy
-    if mo_coeff is None:
-        mo_coeff = mf.mo_coeff
-    nmo = mo_energy.size
-    #e_sort = mo_energy
-
-    mask = make_mo_mask(mo_energy, mo_coeff, mf.mol.ao_mask)
-    nocc = mf.tot_electrons // 2
-
-    if mf.sigma is not None and mf.sigma > 0:
-        mu, mo_occ = _smearing_optimize(mo_energy, stop_grad(nocc), stop_grad(mf.sigma), stop_grad(mask))
-        mo_occ *= 2
-    else:
-        pick = (np.cumsum(mask) <= nocc) & mask
-        mo_occ = np.where(pick, 2., 0.)
-
-        e_homo = np.max(np.where(pick, mo_energy, -np.inf))
-        e_lumo = np.min(np.where(mask & ~pick, mo_energy, np.inf))
-        if mf.verbose >= logger.DEBUG:
-            logger.debug(mf, "  HOMO = %.15g  LUMO = %.15g", e_homo, e_lumo)
-
-    if mf.verbose >= logger.DEBUG:
-        numpy.set_printoptions(threshold=nmo)
-        logger.debug(mf, "  mo_energy =\n%s", mo_energy)
-        numpy.set_printoptions(threshold=1000)
-    return mo_occ
-
-def get_homo_lumo_energy(mf, mo_energy=None, mo_coeff=None):
-    if mf.sigma is not None and mf.sigma > 0:
-        raise NotImplementedError
-
-    if mo_energy is None:
-        mo_energy = mf.mo_energy
-    if mo_coeff is None:
-        mo_coeff = mf.mo_coeff
-
-    mask = make_mo_mask(mo_energy, mo_coeff, mf.mol.ao_mask)
-    nocc = mf.tot_electrons // 2
-    pick = (np.cumsum(mask) <= nocc) & mask
-
-    e_homo = np.max(np.where(pick, mo_energy, -np.inf))
-    e_lumo = np.min(np.where(mask & ~pick, mo_energy, np.inf))
-    if mf.verbose >= logger.DEBUG:
-        logger.debug(mf, "  HOMO = %.15g  LUMO = %.15g", e_homo, e_lumo)
-    return e_homo, e_lumo
-
-def make_rdm1(mo_coeff, mo_occ, **kwargs):
-    dm = (mo_coeff * mo_occ) @ mo_coeff.conj().T
-    return dm
-
-def get_grad(mo_coeff, mo_occ, fock_ao):
-    fock_mo = mo_coeff.conj().T @ fock_ao @ mo_coeff
-    mask = np.where(mo_occ > 0, True, False)
-    g = 2 * fock_mo * ((1-mask[:,None]) * mask[None,:])
-    return g.ravel()
 
 class SCF(hf.SCF):
-    def __init__(self, mol, **kwargs):
-        super().__init__(mol, **kwargs)
-        self.sigma = None
-
-    def get_occ(self, mo_energy=None, mo_coeff=None):
-        return get_occ(self, mo_energy=mo_energy, mo_coeff=mo_coeff)
-
-    @property
-    def tot_electrons(self):
-        return self.mol.tot_electrons()
-
-    def make_rdm1(self, mo_coeff=None, mo_occ=None, **kwargs):
-        if mo_coeff is None:
-            mo_coeff = self.mo_coeff
-        if mo_occ is None:
-            mo_occ = self.mo_occ
-        return make_rdm1(mo_coeff, mo_occ, **kwargs)
-
-    def get_grad(self, mo_coeff, mo_occ, fock=None):
-        if fock is None:
-            dm1 = self.make_rdm1(mo_coeff, mo_occ)
-            fock = self.get_hcore(self.mol) + self.get_veff(self.mol, dm1)
-        return get_grad(mo_coeff, mo_occ, fock)
+    # NOTE too large shift will give large errors
+    padding_level_shift = 1e6
 
     def _eigh(self, h, s):
         ao_mask = self.mol.ao_mask
         mask = np.asarray(1 - ao_mask, dtype=np.int32)
         s = s + np.diag(mask)
-        #h = np.where(np.outer(ao_mask, ao_mask), h, 0.)
-        #h1e_diag = np.diag(h)
-        #h1e_diag = np.where(ao_mask, h1e_diag, 1e10)
-        #h = np.fill_diagonal(h, h1e_diag, inplace=False)
+        h = np.where(np.outer(ao_mask, ao_mask), h, 0.)
+        h = h + np.diag(mask) * self.padding_level_shift
         return eigh(h, s)
 
-    get_homo_lumo_energy = get_homo_lumo_energy
+    def mo_mask(self, mo_energy=None, mo_coeff=None):
+        if mo_energy is None:
+            mo_energy = self.mo_energy
 
-SCFPad = SCF
+        thr = .99 * self.padding_level_shift
+        return np.where(np.greater(mo_energy, thr), False, True)
 
-
-if __name__ == "__main__":
-    import numpy
-    import jax
-    from pyscfad.ml.gto import make_basis_array, MolePad
-    #from pyscfad.xtb import basis as xtb_basis
-    #bfile = xtb_basis.get_basis_filename()
-
-    import os
-    from pyscf.gto import basis
-    bfile = os.path.dirname(basis.__file__) + "/sto-3g.dat"
-    basis = make_basis_array(bfile, max_number=8)
-
-    numbers = np.array([8, 1, 1, 0], dtype=np.int32)
-    coords = np.array(
-        [
-            [0.00000,  0.00000,  0.00000],
-            [1.43355,  0.00000, -0.95296],
-            [1.43355,  0.00000,  0.95296],
-            [0.00000,  0.00000,  0.00000],
-        ]
-    )
-
-    @jax.jit
-    def energy(numbers, coords):
-        mol = MolePad(numbers, coords, basis=basis, verbose=4)
-        dm0 = np.zeros((mol.nao, mol.nao))
-        for i in range(5):
-            dm0 = dm0.at[i,i].set(2.)
-        mf = SCFPad(mol)
-        ehf = mf.kernel(dm0=dm0)
-        return ehf
-
-    ehf = energy(numbers, coords)
-    print(ehf)
-    print(energy._cache_size())
-
-    numbers = np.array([7, 1, 1, 1], dtype=np.int32)
-    coords = np.array(
-        [
-            [-0.80650,       -1.00659,        0.02850],
-            [-0.50540,       -0.31299,        0.68220],
-            [ 0.00620,       -1.41579,       -0.38500],
-            [-1.32340,       -0.54779,       -0.69350],
-        ]
-    ) / 0.52917721067121
-
-    ehf = energy(numbers, coords)
-    print(ehf)
-    print(energy._cache_size())
+SCFPad = SCF

pyscfad/ml/xtb/__init__.py

@@ -1,4 +1,4 @@
-# Copyright 2021-2025 The PySCFAD Authors
+# Copyright 2025-2026 The PySCFAD Authors
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.