Merge remote-tracking branch 'origin/main' into ml

.github/workflows/ci.yml

@@ -11,15 +11,13 @@ jobs:
         python-version: ["3.8", "3.9", "3.10", "3.11"]
     environment: ci
     steps:
-      - uses: actions/checkout@v2
+      - uses: actions/checkout@v4
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: install pyscf
         run: ./.github/workflows/install_pyscf.sh
-      - name: install pyscfadlib
-        run: ./.github/workflows/install_pyscfadlib.sh
       - name: test
         run: ./.github/workflows/run_test.sh
       - name: Upload coverage to codecov

.github/workflows/install_pyscf.sh

@@ -3,24 +3,12 @@ python -m pip install --upgrade pip
 python -m pip cache purge
 pip install wheel
 pip install numpy
-pip install 'scipy<1.11'
+pip install 'scipy<1.12'
 pip install h5py
 pip install jaxlib
 pip install jax
 pip install pytest
 pip install pytest-cov
 
-#pyscf
-git clone https://github.com/fishjojo/pyscf.git
-cd pyscf; git checkout ad; cd ..
-
-if [ "$RUNNER_OS" == "Linux" ]; then
-    os='linux'
-elif [ "$RUNNER_OS" == "macOS" ]; then
-    os='macos'
-else
-    echo "$RUNNER_OS not supported"
-    exit 1
-fi
-
-./.github/workflows/build_pyscf_"$os".sh
+pip install 'pyscf==2.3'
+pip install 'pyscfadlib==0.1.4'

.github/workflows/lint.sh

@@ -2,8 +2,5 @@
 export OMP_NUM_THREADS=1
 export PYTHONPATH=$(pwd):$(pwd)/pyscf:$PYTHONPATH
 echo "pyscfad = True" >> $HOME/.pyscf_conf.py
-echo "pyscf_numpy_backend = 'jax'" >> $HOME/.pyscf_conf.py
-echo "pyscf_scipy_linalg_backend = 'pyscfad'" >> $HOME/.pyscf_conf.py
-echo "pyscf_scipy_backend = 'jax'" >> $HOME/.pyscf_conf.py
 
 pylint pyscfad

.github/workflows/pylint.yml

@@ -7,13 +7,13 @@ jobs:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v2
-      - name: Set up Python 3.7
+      - name: Set up Python 3.11
         uses: actions/setup-python@v2
         with:
-          python-version: 3.7
+          python-version: 3.11
       - name: install pyscf
         run: ./.github/workflows/install_pyscf.sh
       - name: install pylint
-        run: pip install pylint
+        run: pip install 'pylint==2.17.7'
       - name: style check
         run: ./.github/workflows/lint.sh

.github/workflows/run_test.sh

@@ -1,15 +1,4 @@
 #!/usr/bin/env bash
 export OMP_NUM_THREADS=1
-export PYTHONPATH=$(pwd):$(pwd)/pyscf:$PYTHONPATH
-# preload MKL
-#if [[ -n "${MKLROOT}" ]]; then
-#  export LD_PRELOAD=$MKLROOT/lib/intel64/libmkl_avx2.so:$MKLROOT/lib/intel64/libmkl_sequential.so:$MKLROOT/lib/intel64/libmkl_core.so
-#fi
-echo "pyscfad = True" >> $HOME/.pyscf_conf.py
-echo "pyscf_numpy_backend = 'jax'" >> $HOME/.pyscf_conf.py
-echo "pyscf_scipy_linalg_backend = 'pyscfad'" >> $HOME/.pyscf_conf.py
-echo "pyscf_scipy_backend = 'jax'" >> $HOME/.pyscf_conf.py
-#echo "pyscfad_scf_implicit_diff = True" >> $HOME/.pyscf_conf.py
-#echo "pyscfad_ccsd_implicit_diff = True" >> $HOME/.pyscf_conf.py
 
 pytest ./pyscfad --cov-report xml --cov=. --verbosity=1 --durations=10

.pylintrc

@@ -159,6 +159,7 @@ disable=abstract-method,
         deprecated-method,
         arguments-out-of-order,
         consider-using-in,
+        invalid-unary-operand-type,
 
 
 [REPORTS]

LICENSE

@@ -1,6 +1,6 @@
 MIT License
 
-Copyright (c) 2021-2023 Xing Zhang
+Copyright (c) 2021-2024 Xing Zhang
 
 Permission is hereby granted, free of charge, to any person obtaining a copy
 of this software and associated documentation files (the "Software"), to deal

MANIFEST.in

@@ -1,11 +1,6 @@
 include MANIFEST.in
 include README.md setup.py LICENSE
 
-recursive-include pyscfadlib/thirdparty *.so *.dylib
-include pyscfadlib/*.so pyscfadlib/*.dylib pyscfadlib/config.h.in
-
-# source code
-recursive-include pyscfadlib *.c *.h CMakeLists.txt
+include pyscfad/geomopt/log.ini
 
 global-exclude *.py[cod]
-prune pyscfadlib/build

README.md

@@ -8,73 +8,35 @@ PySCF with Auto-differentiation
 Installation
 ------------
 
----
-* To install the latest release, use the following commands:
+* To install the latest release, run:
 ```
-# install cmake
-pip install cmake
-
-# install OpenMP runtime used with clang
-# On Linux:
-sudo apt update
-sudo apt install libomp-dev
-
-# On OSX:
-brew install libomp
-
-# install pyscf
-pip install 'pyscf @ git+https://github.com/fishjojo/pyscf.git@ad#egg=pyscf' 
-pip install 'pyscf-properties @ git+https://github.com/fishjojo/properties.git@ad' 
-
-# install pyscfad
 pip install pyscfad
 ```
 
----
-* To install the development version, use the following command instead:
+* To install the development version, run:
 ```
 pip install git+https://github.com/fishjojo/pyscfad.git
 ```
 
-* The dependencies can be installed via a predefined conda environment
-```
-conda env create -f environment.yml
-conda activate pyscfad_env
-```
-
-* Alternatively, the dependencies can be installed from source
-```
-pip install numpy scipy h5py
-pip install jax jaxlib jaxopt
-
-# install pyscf
-cd $HOME; git clone https://github.com/fishjojo/pyscf.git
-cd pyscf; git checkout ad 
-cd pyscf/lib; mkdir build 
-cd build; cmake ..; make
-
-export PYTHONPATH=$HOME/pyscf:$PYTHONPATH
-```
-
----
-* One can also run PySCFAD inside a docker container:
-```
-docker pull fishjojo/pyscfad:latest
-docker run -rm -t -i fishjojo/pyscfad:latest /bin/bash
-```
+`pyscfad` depends on `pyscf==2.3.0`.
+Other dependencies include
+`numpy>=1.17`, `scipy<1.12`, `h5py`,
+`pyscfadlib`,
+`jax>=0.3.25` and `jaxlib>=0.3.25`.
 
 Running examples
 ----------------
 
 * In order to perform AD calculations, 
 the following lines need to be added to 
-the PySCF configure file($HOME/.pyscf\_conf.py)
+the PySCF configure file ($HOME/.pyscf\_conf.py)
 ```
 pyscfad = True
-pyscf_numpy_backend = 'jax'
-pyscf_scipy_linalg_backend = 'pyscfad'
-pyscf_scipy_backend = 'jax'
-# The followings are optional
-pyscfad_scf_implicit_diff = True
-pyscfad_ccsd_implicit_diff = True
 ```
+
+Citing PySCFAD
+--------------
+The following paper should be cited in publications utilizing the PySCFAD program package:
+
+[Differentiable quantum chemistry with PySCF for molecules and materials at the mean-field level and beyond](https://doi.org/10.1063/5.0118200), 
+X. Zhang, G. K.-L. Chan, *J. Chem. Phys.*, **157**, 204801 (2022)

examples/cc/00-simple.py

@@ -1,5 +1,3 @@
-import jax
-from pyscfad import gto, scf, cc
 '''
 Reference nuclear gradient
 [[0.0, 0.0, -1.15101379e-01]
@@ -8,10 +6,19 @@
 Note that without implicit differentiation turned on,
 the gradient will be initial guess dependent.
 '''
+import jax
+from pyscfad import gto, scf, cc
+from pyscfad import config
+
+# implicit differentiation of SCF iterations
+config.update('pyscfad_scf_implicit_diff', True)
+# implicit differentiation of CC iterations
+config.update('pyscfad_ccsd_implicit_diff', True)
+
 mol = gto.Mole()
 mol.atom = 'H 0. 0. 0.; F 0. 0. 1.1'
 mol.basis = 'ccpvdz'
-mol.verbose = 3
+mol.verbose = 4
 mol.incore_anyway = True
 mol.build()
 

examples/dft/00-simple.py

@@ -1,12 +1,11 @@
-import pyscf
 from pyscfad import gto, dft
 
 """
 Analytic nuclear gradient for RKS computed by auto-differentiation
 """
 
 mol = gto.Mole()
-mol.atom = 'H 0 0 0; H 0 0 0.74'  # in Angstrom
+mol.atom = 'H 0 0 0; H 0 0 0.74'
 mol.basis = '631g'
 mol.verbose=5
 mol.build()

examples/gw/01-rpa.py

@@ -1,12 +1,3 @@
-import jax
-from pyscf import df as pyscf_df
-from pyscfad import gto, dft, scf, df
-from pyscfad.gw import rpa
-from pyscfad import config
-
-config.update('pyscfad_scf_implicit_diff', True)
-config.update('pyscfad_moleintor_opt', True)
-
 '''
 RPA e_tot, e_hf, e_corr =  -76.26428191794197 -75.95645187758402 -0.30783004035795963
 
@@ -39,6 +30,18 @@
 
 '''
 
+import jax
+from pyscf import df as pyscf_df
+from pyscfad import gto, dft, scf, df
+from pyscfad.gw import rpa
+from pyscfad import config
+
+config.update('pyscfad_scf_implicit_diff', True)
+# Using optimized C implementation for gradients calculations.
+# This requires the `pyscfadlib` package, which can be installed with
+# `pip install pyscfadlib`
+#config.update('pyscfad_moleintor_opt', True)
+
 mol = gto.Mole()
 mol.verbose = 4
 mol.atom = [

examples/gw/03-rpa-geomopt-with-pyberny.py

@@ -0,0 +1,68 @@
+'''
+Geometry optimization using RPA gradients calculated with PySCFAD and PyBerny
+as a molecular geometry optimizer
+'''
+import jax
+from pyscf import df as pyscf_df
+from pyscfad import gto, dft, df
+from pyscfad.gw import rpa
+from pyscf.geomopt.berny_solver import optimize, to_berny_geom
+from berny import Berny, geomlib
+import warnings
+warnings.simplefilter("ignore")
+
+def energy_(mol, with_df):
+    mf = dft.RKS(mol)
+    mf.xc = 'pbe'
+    mf.kernel(dm0=None)
+
+    mymp = rpa.RPA(mf)
+    mymp.with_df = with_df
+    mymp.kernel()
+    return mymp.e_tot
+
+def solver(geom, val_and_grad):
+    mol = gto.Mole()
+    mol.verbose = 0
+    mol.atom = geom
+    mol.basis = 'ccpvdz'
+    mol.build(trace_exp=False, trace_ctr_coeff=False)
+
+    auxbasis = pyscf_df.addons.make_auxbasis(mol, mp2fit=True)
+    auxmol = df.addons.make_auxmol(mol, auxbasis)
+    with_df = df.DF(mol, auxmol=auxmol)
+
+    e_tot, jac = val_and_grad(mol, with_df)
+
+    return e_tot, jac[0].coords + jac[1].mol.coords + jac[1].auxmol.coords
+
+
+# fictitious molecule to initialize PyBerny optimizer
+mol_ = gto.Mole()
+mol_.build(atom = 'H 0 0 0; H 0 0 1.', basis = 'sto-3g')
+geom = to_berny_geom(mol_)
+optimizer = Berny(geom) # initialize geometry optimizer
+
+val_and_grad = jax.value_and_grad(energy_, (0,1))
+
+for iter_, geom in enumerate(optimizer):
+    energy, gradients = solver(list(geom), val_and_grad)
+    optimizer.send((energy, gradients))
+    print(f'iter={iter_+1}   energy={energy:.10f}')
+
+print('\nOptimized geometry:')
+print(geom.coords)
+
+'''
+iter=1   energy=-1.1691629669
+iter=2   energy=-1.1857669916
+iter=3   energy=-1.1780975967
+iter=4   energy=-1.1915441270
+iter=5   energy=-1.1919771635
+iter=6   energy=-1.1920713876
+iter=7   energy=-1.1920721970
+
+Optimized geometry:
+[[0.         0.         0.11749648]
+ [0.         0.         0.88250352]]
+'''

examples/mp/00-simple.py

@@ -1,15 +1,18 @@
+"""
+Analytic nuclear gradient for MP2 computed by auto-differentiation
+"""
 import jax
 import pyscf
 from pyscfad import gto, scf, mp
+from pyscfad import config
 
-"""
-Analytic nuclear gradient for MP2 computed by auto-differentiation
-"""
+# implicit differentiation of SCF iterations
+config.update('pyscfad_scf_implicit_diff', True)
 
 mol = gto.Mole()
-mol.atom = 'H 0 0 0; H 0 0 0.74'  # in Angstrom
+mol.atom = 'H 0 0 0; H 0 0 0.74'
 mol.basis = '631g'
-mol.verbose=5
+mol.verbose=4
 mol.build()
 
 def mp2(mol, dm0=None):

examples/mp/10-oomp2.py

@@ -1,6 +1,6 @@
 from scipy.optimize import minimize
 from jax import value_and_grad
-from pyscf import numpy as np
+from jax import numpy as np
 from pyscfad import util
 from pyscfad.tools import rotate_mo1
 from pyscfad import gto, scf, mp