Small fixes (#62)

* fix dynamic polarizability

* update dft

* update Davidson solver

* update tdscf

* update doc

Author: fishjojo

.github/workflows/install_pyscf.sh

@@ -7,6 +7,6 @@ pip install pytest-cov
 pip install numpy
 pip install scipy
 pip install h5py
-pip install jaxlib
-pip install jax
-pip install 'pyscf>=2.3,<2.7'
+pip install 'jaxlib<=0.4.35'
+pip install 'jax<=0.4.35'
+pip install 'pyscf>=2.3'

README.md

@@ -31,7 +31,7 @@ The dependent C library `pyscfadlib` can be compiled from source following the i
 `numpy`, `scipy`,
 `pyscf>=2.3.0`,
 `pyscfadlib>=0.1.4`,
-`jax>=0.4.14`, and `jaxlib>=0.4.14`.
+`jax>=0.4.14,<=0.4.35`, and `jaxlib>=0.4.14,<=0.4.35`.
 
 Citing PySCFAD
 --------------

doc/source/getting_started/install.rst

@@ -82,9 +82,15 @@ Package                                                supported versions
 `numpy <https://numpy.org>`_                           >=1.17
 `scipy <https://scipy.org>`_                           >=1.7
 `h5py <https://www.h5py.org/>`_                        >=2.7
-`jax <https://jax.readthedocs.io/en/latest/>`_         >=0.4.14
-`jaxlib <https://pypi.org/project/jaxlib/>`_           >=0.4.14
+`jax <https://jax.readthedocs.io/en/latest/>`_         >=0.4.14,<=0.4.35
+`jaxlib <https://pypi.org/project/jaxlib/>`_           >=0.4.14,<=0.4.35
 `pyscf <https://pyscf.org/>`_                          >=2.3.0
 `pyscfadlib <https://pypi.org/project/pyscfadlib/>`_   >=0.1.4
 =====================================================  ==================
 
+.. note::
+
+    Since jax version 0.4.36, the tracing machinery has been modified
+    to eliminate data-dependent tracing, which conflicts with pyscfad's flexibility.
+    A comprehensive update to ensure compatibility with later jax versions may be
+    introduced in the future pyscfad 0.2 release.

examples/scf/10-polarizability.py

@@ -1,4 +1,3 @@
-import numpy
 import jax
 from pyscfad import numpy as np
 from pyscfad import gto, scf
@@ -19,7 +18,7 @@ def apply_E(E):
     mf.kernel()
     return mf.dip_moment(mol, mf.make_rdm1(), unit='AU', verbose=0)
 
-E0 = numpy.zeros((3))
+E0 = np.zeros((3))
 polar = jax.jacfwd(apply_E)(E0)
 print(polar)
 

examples/scf/11-born_charge.py

@@ -4,7 +4,7 @@
 import numpy
 import jax
 from pyscfad import gto, scf
-from pyscfad.lib import numpy as jnp
+from pyscfad import numpy as np
 
 mol = gto.Mole()
 mol.atom = '''H  ,  0.   0.   0.
@@ -17,7 +17,7 @@ def dip_moment(mol):
     ao_dip = mol.intor_symmetric('int1e_r', comp=3)
     h1 = mf.get_hcore()
     E = numpy.zeros((3))
-    mf.get_hcore = lambda *args, **kwargs: h1 + jnp.einsum('x,xij->ij', E, ao_dip)
+    mf.get_hcore = lambda *args, **kwargs: h1 + np.einsum('x,xij->ij', E, ao_dip)
     mf.kernel()
     return mf.dip_moment(mol, mf.make_rdm1(), unit='AU', verbose=0)
 

examples/scf/12-raman.py

@@ -10,7 +10,7 @@
 '''
 import jax
 from pyscfad import gto, scf
-from pyscfad.lib import numpy as np
+from pyscfad import numpy as np
 from pyscfad.prop.polarizability.rhf import Polarizability
 from pyscfad.prop.thermo import vib
 

examples/scf/21-ir.py

@@ -8,7 +8,7 @@
 '''
 import jax
 from pyscfad import gto, scf, cc
-from pyscfad.lib import numpy as np
+from pyscfad import numpy as np
 from pyscfad.prop.thermo import vib
 
 mol = gto.Mole()

pyscfad/dft/libxc.py

@@ -59,10 +59,13 @@ def _eval_xc_comp_jvp(xc_code, spin, relativity, deriv, omega, verbose,
             vrho1, vsigma1, vlapl1, vtau1 = _vxc_partial_deriv(rho, val, val1, 'MGGA')
             vrho_jvp = np.einsum('np,np->p', vrho1, rho_t)
             vsigma_jvp = np.einsum('np,np->p', vsigma1, rho_t)
-            vlapl_jvp = np.einsum('np,np->p', vlapl1, rho_t)
+            if vlapl1 is None:
+                vlapl_jvp = None
+            else:
+                vlapl_jvp = np.einsum('np,np->p', vlapl1, rho_t)
             vtau_jvp = np.einsum('np,np->p', vtau1, rho_t)
             vrho1 = vsigma1 = vlapl1 = vtau1 = None
-            jvp = np.vstack((vrho_jvp, vsigma_jvp, vlapl_jvp, vtau_jvp))
+            jvp = (vrho_jvp, vsigma_jvp, vlapl_jvp, vtau_jvp)
         else:
             raise NotImplementedError
     else:
@@ -92,7 +95,10 @@ def _exc_partial_deriv(rho, exc, vxc, xctype='LDA'):
         dsigma = vxc[1] / rho[0] * 2. * rho[1:4]
         exc1 = np.vstack((drho, dsigma))
         if xctype == 'MGGA':
-            dlap = vxc[2] / rho[0]
+            if vxc[2] is None:
+                dlap = np.zeros_like(rho[0])
+            else:
+                dlap = vxc[2] / rho[0]
             dtau = vxc[3] / rho[0]
             exc1 = np.vstack((exc1, dlap, dtau))
     else:
@@ -108,10 +114,21 @@ def _vxc_partial_deriv(rho, vxc, fxc, xctype='LDA'):
         vrho1 = np.vstack((fxc[0], fxc[1] * 2. * rho[1:4]))
         vsigma1 = np.vstack((fxc[1], fxc[2] * 2. * rho[1:4]))
         if xctype == 'MGGA':
-            vrho1 = np.vstack((vrho1, fxc[5], fxc[6]))
-            vsigma1 = np.vstack((vsigma1, fxc[8], fxc[9]))
-            vlapl1 = np.vstack((fxc[5], fxc[8] * 2. * rho[1:4], fxc[3], fxc[7]))
-            vtau1 = np.vstack((fxc[6], fxc[9] * 2. * rho[1:4], fxc[7], fxc[4]))
+            ZERO = np.zeros_like(rho[0])
+            if vxc[2] is None:
+                fxc3 = fxc5 = fxc7 = fxc8 = ZERO
+            else:
+                fxc3 = fxc[3]
+                fxc5 = fxc[5]
+                fxc7 = fxc[7]
+                fxc8 = fxc[8]
+            vrho1 = np.vstack((vrho1, fxc5, fxc[6]))
+            vsigma1 = np.vstack((vsigma1, fxc8, fxc[9]))
+            if vxc[2] is None:
+                vlapl1 = None
+            else:
+                vlapl1 = np.vstack((fxc5, fxc8 * 2. * rho[1:4], fxc3, fxc7))
+            vtau1 = np.vstack((fxc[6], fxc[9] * 2. * rho[1:4], fxc7, fxc[4]))
     else:
         raise KeyError
     return vrho1, vsigma1, vlapl1, vtau1