fix scipy gmres dtype error

Author: fishjojo

pyscfad/_src/scipy/sparse/linalg.py

@@ -11,7 +11,7 @@ def _matvec(x):
         # NOTE result may not be writable
         # (required by scipy>=1.12), so make a copy
         return numpy.array(ops.to_numpy(Ax))
-    A = LinearOperator((b.size, b.size), matvec=_matvec)
+    A = LinearOperator((b.size, b.size), matvec=_matvec, dtype=b.dtype)
     return A
 
 def gmres(A_or_matvec, b, x0=None, *,

pyscfad/cc/test/conftest.py

@@ -17,3 +17,20 @@ def get_mol():
     yield mol
 
     config.reset()
+
+@pytest.fixture
+def get_opt_mol():
+    config.update("pyscfad_moleintor_opt", True)
+    config.update('pyscfad_scf_implicit_diff', True)
+    config.update('pyscfad_ccsd_implicit_diff', True)
+
+    mol = gto.Mole()
+    mol.atom = 'H 0. 0. 0.; F 0. 0. 1.1'
+    mol.basis = '631g'
+    mol.verbose = 0
+    mol.incore_anyway = True
+    mol.max_memory = 7000
+    mol.build(trace_exp=False, trace_ctr_coeff=False)
+    yield mol
+
+    config.reset()

pyscfad/cc/test/test_ccsd_t.py pyscfad/cc/test/test_ccsd_t_slow.pypyscfad/cc/test/test_ccsd_t.py renamed to pyscfad/cc/test/test_ccsd_t_slow.py

@@ -1,11 +1,6 @@
-import pytest
 import numpy
 import jax
-from pyscfad import gto, scf, cc
-from pyscfad import config
-
-config.update('pyscfad_scf_implicit_diff', True)
-config.update('pyscfad_ccsd_implicit_diff', True)
+from pyscfad import scf, cc
 
 def test_nuc_grad(get_mol):
     mol = get_mol
@@ -16,8 +11,7 @@ def energy(mol):
         mycc.kernel()
         et = mycc.ccsd_t()
         return mycc.e_tot + et
-    with jax.disable_jit():
-        g1 = jax.grad(energy)(mol).coords
+    g1 = jax.grad(energy)(mol).coords
     g0 = numpy.array([[0., 0., -8.60709468e-02],
                       [0., 0.,  8.60709468e-02]])
     assert(abs(g1-g0).max() < 1e-6)

pyscfad/cc/test/test_dcsd.py

@@ -1,4 +1,3 @@
-import pytest
 import numpy
 import jax
 from pyscfad import scf
@@ -12,8 +11,7 @@ def energy(mol):
         mycc = dfdcsd.RDCSD(mf)
         mycc.kernel()
         return mycc.e_tot
-    with jax.disable_jit():
-        g1 = jax.grad(energy)(mol).coords
+    g1 = jax.grad(energy)(mol).coords
     # finite difference
     g0 = numpy.array([[0., 0., -0.08500490828],
                       [0., 0.,  0.08500490828]])

pyscfad/cc/test/test_opt_dfccsd.py

@@ -0,0 +1,25 @@
+import numpy
+import jax
+from pyscfad import scf
+from pyscfad.cc import dfccsd
+
+def test_dfccsdt_nuc_grad(get_opt_mol):
+    mol = get_opt_mol
+    def energy(mol):
+        mf = scf.RHF(mol).density_fit()
+        mf.kernel()
+
+        mycc = dfccsd.RCCSD(mf)
+        eris = mycc.ao2mo()
+        mycc.kernel(eris=eris)
+        et = mycc.ccsd_t(eris=eris)
+        return mycc.e_tot + et
+
+    e, jac = jax.value_and_grad(energy)(mol)
+
+    e0 = -100.10156178822595
+    assert abs(e - e0) < 1e-6
+
+    g0 = numpy.array([[0., 0., -0.0860735932],
+                      [0., 0.,  0.0860735932]])
+    assert abs(jac.coords-g0).max() < 1e-6

pyscfad/cc/test/test_rccsd.py

@@ -1,4 +1,3 @@
-import pytest
 import numpy
 import jax
 from pyscfad import scf, cc
@@ -11,8 +10,7 @@ def energy(mol):
         mycc = cc.RCCSD(mf)
         mycc.kernel()
         return mycc.e_tot
-    with jax.disable_jit():
-        g1 = jax.grad(energy)(mol).coords
+    g1 = jax.grad(energy)(mol).coords
     g0 = numpy.array([[0., 0., -0.0873564848],
                       [0., 0.,  0.0873564848]])
     assert(abs(g1-g0).max() < 1e-6)
@@ -25,8 +23,7 @@ def energy(mol):
         mycc = cc.dfccsd.RCCSD(mf)
         mycc.kernel()
         return mycc.e_tot
-    with jax.disable_jit():
-        g1 = jax.grad(energy)(mol).coords
+    g1 = jax.grad(energy)(mol).coords
     # finite difference
     g0 = numpy.array([[0., 0., -0.0873569023],
                       [0., 0.,  0.0873569023]])

pyscfad/cc/test/test_rccsd_hess.py

@@ -1,12 +1,6 @@
-import pytest
 import numpy
 import jax
-from pyscf import lib
-from pyscfad import gto, scf, cc
-from pyscfad import config
-
-config.update('pyscfad_scf_implicit_diff', True)
-config.update('pyscfad_ccsd_implicit_diff', True)
+from pyscfad import scf, cc
 
 def test_nuc_hessian(get_mol):
     mol = get_mol
@@ -16,8 +10,7 @@ def energy(mol):
         mycc = cc.RCCSD(mf)
         mycc.kernel()
         return mycc.e_tot
-    with jax.disable_jit():
-        h1 = jax.jacrev(jax.jacrev(energy))(mol).coords.coords
+    h1 = jax.jacrev(jax.jacrev(energy))(mol).coords.coords
     h0 = numpy.array(
         [[[[ 4.20246014e-02, 0, 0],
            [-4.20246014e-02, 0, 0]],
@@ -33,5 +26,3 @@ def energy(mol):
            [0, 0,  1.53241780e-01]]]]
     )
     assert(abs(h1-h0).max() < 5e-5)
-    #f = lib.fp(h1)
-    #assert(abs(f - -0.18529155578401263) < 1e-6)