diff --git a/examples/jax/01_slater.py b/examples/jax/01_slater.py
new file mode 100644
index 00000000..145a46d9
--- /dev/null
+++ b/examples/jax/01_slater.py
@@ -0,0 +1,102 @@
+import pyscf.pbc.gto as gto
+import pyscf.gto as molgto
+import pyscf.pbc.dft as dft
+import numpy
+import pandas as pd
+
+
+def make_cell_basis():
+ """
+ Here we make an uncontracted basis for the cell.
+ Our JAX implementation requires this for efficient evaluation,
+ and often for solids the uncontracted basis is much better because the
+ others
+ """
+ cell = gto.Cell()
+ cell.atom = '''C 0. 0. 0.
+ C 0.8917 0.8917 0.8917
+ C 1.7834 1.7834 0.
+ C 2.6751 2.6751 0.8917
+ C 1.7834 0. 1.7834
+ C 2.6751 0.8917 2.6751
+ C 0. 1.7834 1.7834
+ C 0.8917 2.6751 2.6751'''
+ cell.basis = 'unc-ccecp-ccpvtz'
+ cell.pseudo = 'ccecp' #These are high accuracy ECP's for QMC.
+ cell.cart = True # also important for JAX efficiency.
+ cell.a = numpy.eye(3)*3.5668
+ cell.build()
+ return cell
+
+
+
+def run_dft(cell, chkfile):
+ mf = dft.RKS(cell)
+ mf.xc = 'lda,vwn'
+ mf = mf.multigrid_numint()
+ mf.chkfile = chkfile
+ mf.kernel()
+ return mf.e_tot
+
+
+def generate_etb_set(cell, alpha0=0.2, l_polarization=2):
+ """
+ Generate an even tempered basis which is selected to best reproduce the
+ basis in cell.
+
+ You can use this as written for the first 3 rows and all sp elements.
+
+ cell is a cell object with a given basis
+ alpha0 is the longest range you would like to go. typically 0.2 or 0.1
+ l_polarization is how many angular momentum functions you'd like to allow
+ """
+ #print(cell._basis)
+ new_basis ={}
+ for atomname, basis in cell._basis.items():
+ maxl_contract = 1 + l_polarization
+ # If you are doing a F-electron or 4d or 5d element then you
+ # need to update this.
+ if atomname in ['Sc', 'Ti', 'V', 'Cr', 'Mn', 'Fe', 'Co', 'Ni', 'Cu']:
+ maxl_contract = 2+l_polarization
+
+ # in this loop find the minimum and maximum exponents
+ # for each angular momentum channel.
+ maxexp = numpy.zeros(maxl_contract+1)
+ minexp = numpy.ones(maxl_contract+1)*1000
+ for element in basis:
+ #print(element)
+ l = element[0]
+ exponents = numpy.max([e[0] for e in element[1:]])
+ #print(exponents)
+ if l <= maxl_contract:
+ maxexp[l] = numpy.max([exponents, maxexp[l]])
+ minexp[l] = numpy.min([exponents, minexp[l]])
+
+ # Truncate the minimum exponent to alpha0
+ minexp[minexp < alpha0] = alpha0
+ # Now
+ etbs = []
+ for l, maxe in enumerate(maxexp):
+ n = int(numpy.log2(maxe/minexp[l]))+2
+ etbs.append((l, n, minexp[l], 2))
+ new_basis[atomname] = molgto.etbs(etbs)
+ newcell = cell.copy()
+ newcell.basis = new_basis
+ newcell.build()
+ return newcell
+
+
+if __name__ == "__main__":
+ cell_orig = make_cell_basis()
+ cell_etb = generate_etb_set(cell_orig, alpha0=0.2, l_polarization=2)
+ cell_exp2 = cell_orig.copy()
+ cell_exp2.exp_to_discard = 0.2
+ cell_exp2.build()
+
+ # Sometimes an even tempered basis is better than the uncontract and expand, and
+ # sometimes not.
+ print("etb basis", run_dft(cell_etb, 'etb0.2.hdf5'))
+ print("uncontracted + exp_to_discard", run_dft(cell_exp2, 'exp_to_discard0.2.hdf5'))
+
+ # You should probably also try decreasing alpha0, depending on the material.
+
diff --git a/examples/jax/02_check_jax.py b/examples/jax/02_check_jax.py
new file mode 100644
index 00000000..afec981c
--- /dev/null
+++ b/examples/jax/02_check_jax.py
@@ -0,0 +1,86 @@
+import os
+# here we are desperately trying to avoid multithreading
+# to get a good read on the single-threaded performance
+# you may or many not want this depending on your use case
+os.environ["OPENBLAS_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+os.environ["NUMEXPR_NUM_THREADS"] = "1"
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["XLA_FLAGS"] = (
+ "--xla_cpu_multi_thread_eigen=false intra_op_parallelism_threads=1 inter_op_parallelism_threads=1"
+)
+os.environ["JAX_NUM_CLIENTS"] = "1"
+os.environ["NPROC"] = "1"
+
+import pyqmc.api as pyq
+import jax
+import time
+import numpy as np
+import pandas as pd
+
+# you may or may not want to set this.
+jax.config.update('jax_platform_name', 'cpu')
+# you almost always want 64-bit math.
+jax.config.update("jax_enable_x64", True)
+print(jax.devices())
+
+
+def check_value(configs, wf_jax, wf_pyscf):
+ """
+ print out timing and difference in values for recompute()
+ """
+ vals_jax = wf_jax.recompute(configs)
+ jax.block_until_ready(vals_jax)
+ start = time.perf_counter()
+ vals_jax = wf_jax.recompute(configs)
+ jax.block_until_ready(vals_jax)
+ jax_time = time.perf_counter()
+ vals_pyscf = wf_pyscf.recompute(configs)
+ pyscf = time.perf_counter()
+ print(f"JAX time {jax_time - start} s PYSCF time {pyscf - jax_time} s "
+ f" Difference {np.mean(np.abs(vals_jax[0] - vals_pyscf[0]))}")
+
+
+def check_energy(configs, wf_jax, wf_pyscf) -> pd.DataFrame:
+ """
+ Check the various energies.
+ Note that between old and new ECPS they should only agree on average
+
+ """
+ enacc = {'old':pyq.EnergyAccumulator(cell, use_old_ecp=True),
+ 'new': pyq.EnergyAccumulator(cell, use_old_ecp=False) }
+ wfs = {'jax': wf_jax, 'pyscf': wf_pyscf}
+ data = []
+ for ecp in enacc.keys():
+ for wf in wfs.keys():
+ if wf == 'jax': #force compile
+ enacc[ecp](configs, wfs[wf])
+ start = time.perf_counter()
+ en = enacc[ecp](configs, wfs[wf])
+ end = time.perf_counter()
+ data.append({ 'time': end - start,
+ 'wf':wf,
+ 'ecp':ecp,
+ 'ecp_en': np.mean(en['ecp']),
+ 'grad2': np.mean(en['grad2']),
+ 'ke': np.mean(en['ke']),
+ 'total':np.mean(en['total']),
+ })
+ return pd.DataFrame(data)
+
+
+if __name__ == "__main__":
+ # we found that etb0.2 was the lowest DFT energy
+ cell, mf = pyq.recover_pyscf("etb0.2.hdf5")
+ wf_jax, _ = pyq.generate_slater(cell, mf, jax=True, nimages=2)
+ wf_pyscf, _ = pyq.generate_slater(cell, mf, jax=False, eval_gto_precision=1e-16 )
+
+ for nconfig in [10, 100, 1000]:
+ print(f"###### nconfig = {nconfig}")
+ configs = pyq.initial_guess(cell, nconfig)
+ check_value(configs, wf_jax, wf_pyscf)
+ df = check_energy(configs, wf_jax, wf_pyscf)
+ print(df)
+
+
+
diff --git a/examples/trial_wf/slater_geminal_etb_optimize.ipynb b/examples/trial_wf/slater_geminal_etb_optimize.ipynb
new file mode 100644
index 00000000..bc10aa30
--- /dev/null
+++ b/examples/trial_wf/slater_geminal_etb_optimize.ipynb
@@ -0,0 +1,518 @@
+{
+ "cells": [
+ {
+ "cell_type": "code",
+ "id": "initial_id",
+ "metadata": {
+ "collapsed": true,
+ "ExecuteTime": {
+ "end_time": "2025-12-08T13:43:15.296055Z",
+ "start_time": "2025-12-08T13:43:15.218371Z"
+ }
+ },
+ "source": [
+ "import os\n",
+ "import h5py\n",
+ "import pandas as pd\n",
+ "import matplotlib.pyplot as plt\n",
+ "import seaborn as sns\n",
+ "alpha_s = 0.2\n",
+ "alpha_p = 0.2\n",
+ "\n",
+ "df = []\n",
+ "ccvtz_energy = -16.8069\n",
+ "for n_s in range(0, 6):\n",
+ " for n_p in range(6):\n",
+ " for n_d in range(0, 4):\n",
+ " hdf_file = f\"data/slater_geminal_etb_{n_s}_{alpha_s}_{n_p}_{alpha_p}_{n_d}.hdf5\"\n",
+ " if os.path.exists(hdf_file):\n",
+ " print(f\"File {hdf_file} exists\")\n",
+ " with h5py.File(hdf_file, \"r\") as f:\n",
+ " energy = f[\"energy\"][()][-1]\n",
+ " error = f['energy_error'][()][-1]\n",
+ " if 'time' in f.keys():\n",
+ " time = f[\"time\"][()]\n",
+ " else:\n",
+ " time = 0\n",
+ " df.append({\n",
+ " \"n_s\": n_s,\n",
+ " \"n_p\": n_p,\n",
+ " \"n_d\":n_d,\n",
+ " \"energy\": energy,\n",
+ " 'error': error,\n",
+ " \"time\": time,\n",
+ " })\n",
+ "\n",
+ "\n",
+ "df = pd.DataFrame(df)\n",
+ "df\n"
+ ],
+ "outputs": [
+ {
+ "name": "stdout",
+ "output_type": "stream",
+ "text": [
+ "File data/slater_geminal_etb_0_0.2_0_0.2_1.hdf5 exists\n",
+ "File data/slater_geminal_etb_0_0.2_0_0.2_2.hdf5 exists\n",
+ "File data/slater_geminal_etb_0_0.2_1_0.2_0.hdf5 exists\n",
+ "File data/slater_geminal_etb_0_0.2_1_0.2_1.hdf5 exists\n",
+ "File data/slater_geminal_etb_0_0.2_1_0.2_2.hdf5 exists\n",
+ "File data/slater_geminal_etb_0_0.2_2_0.2_0.hdf5 exists\n",
+ "File data/slater_geminal_etb_0_0.2_2_0.2_1.hdf5 exists\n",
+ "File data/slater_geminal_etb_0_0.2_2_0.2_2.hdf5 exists\n",
+ "File data/slater_geminal_etb_1_0.2_0_0.2_0.hdf5 exists\n",
+ "File data/slater_geminal_etb_1_0.2_0_0.2_1.hdf5 exists\n",
+ "File data/slater_geminal_etb_1_0.2_0_0.2_2.hdf5 exists\n",
+ "File data/slater_geminal_etb_1_0.2_1_0.2_0.hdf5 exists\n",
+ "File data/slater_geminal_etb_1_0.2_1_0.2_1.hdf5 exists\n",
+ "File data/slater_geminal_etb_1_0.2_1_0.2_2.hdf5 exists\n",
+ "File data/slater_geminal_etb_1_0.2_2_0.2_0.hdf5 exists\n",
+ "File data/slater_geminal_etb_1_0.2_2_0.2_1.hdf5 exists\n",
+ "File data/slater_geminal_etb_1_0.2_2_0.2_2.hdf5 exists\n",
+ "File data/slater_geminal_etb_2_0.2_0_0.2_0.hdf5 exists\n",
+ "File data/slater_geminal_etb_2_0.2_0_0.2_1.hdf5 exists\n",
+ "File data/slater_geminal_etb_2_0.2_0_0.2_2.hdf5 exists\n",
+ "File data/slater_geminal_etb_2_0.2_1_0.2_0.hdf5 exists\n",
+ "File data/slater_geminal_etb_2_0.2_1_0.2_1.hdf5 exists\n",
+ "File data/slater_geminal_etb_2_0.2_1_0.2_2.hdf5 exists\n",
+ "File data/slater_geminal_etb_2_0.2_2_0.2_0.hdf5 exists\n",
+ "File data/slater_geminal_etb_2_0.2_2_0.2_1.hdf5 exists\n",
+ "File data/slater_geminal_etb_2_0.2_2_0.2_2.hdf5 exists\n"
+ ]
+ },
+ {
+ "data": {
+ "text/plain": [
+ " n_s n_p n_d energy error time\n",
+ "0 0 0 1 -16.733491 0.001864 10207.751945\n",
+ "1 0 0 2 -16.741963 0.001899 21817.270212\n",
+ "2 0 1 0 -16.781000 0.001554 7519.828752\n",
+ "3 0 1 1 -16.803994 0.001541 16429.181557\n",
+ "4 0 1 2 -16.782736 0.001365 32886.353608\n",
+ "5 0 2 0 -16.778963 0.001531 12022.304522\n",
+ "6 0 2 1 -16.794903 0.001064 25297.375293\n",
+ "7 0 2 2 -16.812330 0.001490 0.000000\n",
+ "8 1 0 0 -16.730423 0.001145 6190.928719\n",
+ "9 1 0 1 -16.750052 0.000903 12034.494479\n",
+ "10 1 0 2 -16.755024 0.002155 25005.136854\n",
+ "11 1 1 0 -16.801767 0.000844 8815.614302\n",
+ "12 1 1 1 -16.816471 0.001163 18745.997369\n",
+ "13 1 1 2 -16.816585 0.001368 0.000000\n",
+ "14 1 2 0 -16.803039 0.001425 13620.759366\n",
+ "15 1 2 1 -16.812902 0.001359 0.000000\n",
+ "16 1 2 2 -16.817182 0.001203 0.000000\n",
+ "17 2 0 0 -16.766942 0.001670 7466.260457\n",
+ "18 2 0 1 -16.779606 0.001174 13328.825939\n",
+ "19 2 0 2 -16.792576 0.000622 0.000000\n",
+ "20 2 1 0 -16.792028 0.001631 9854.208552\n",
+ "21 2 1 1 -16.811469 0.001701 0.000000\n",
+ "22 2 1 2 -16.819245 0.001173 0.000000\n",
+ "23 2 2 0 -16.802373 0.000936 0.000000\n",
+ "24 2 2 1 -16.759757 0.000936 0.000000\n",
+ "25 2 2 2 -16.630725 0.001892 0.000000"
+ ],
+ "text/html": [
+ "\n",
+ "\n",
+ "
\n",
+ " \n",
+ " \n",
+ " | \n",
+ " n_s | \n",
+ " n_p | \n",
+ " n_d | \n",
+ " energy | \n",
+ " error | \n",
+ " time | \n",
+ "
\n",
+ " \n",
+ " \n",
+ " \n",
+ " | 0 | \n",
+ " 0 | \n",
+ " 0 | \n",
+ " 1 | \n",
+ " -16.733491 | \n",
+ " 0.001864 | \n",
+ " 10207.751945 | \n",
+ "
\n",
+ " \n",
+ " | 1 | \n",
+ " 0 | \n",
+ " 0 | \n",
+ " 2 | \n",
+ " -16.741963 | \n",
+ " 0.001899 | \n",
+ " 21817.270212 | \n",
+ "
\n",
+ " \n",
+ " | 2 | \n",
+ " 0 | \n",
+ " 1 | \n",
+ " 0 | \n",
+ " -16.781000 | \n",
+ " 0.001554 | \n",
+ " 7519.828752 | \n",
+ "
\n",
+ " \n",
+ " | 3 | \n",
+ " 0 | \n",
+ " 1 | \n",
+ " 1 | \n",
+ " -16.803994 | \n",
+ " 0.001541 | \n",
+ " 16429.181557 | \n",
+ "
\n",
+ " \n",
+ " | 4 | \n",
+ " 0 | \n",
+ " 1 | \n",
+ " 2 | \n",
+ " -16.782736 | \n",
+ " 0.001365 | \n",
+ " 32886.353608 | \n",
+ "
\n",
+ " \n",
+ " | 5 | \n",
+ " 0 | \n",
+ " 2 | \n",
+ " 0 | \n",
+ " -16.778963 | \n",
+ " 0.001531 | \n",
+ " 12022.304522 | \n",
+ "
\n",
+ " \n",
+ " | 6 | \n",
+ " 0 | \n",
+ " 2 | \n",
+ " 1 | \n",
+ " -16.794903 | \n",
+ " 0.001064 | \n",
+ " 25297.375293 | \n",
+ "
\n",
+ " \n",
+ " | 7 | \n",
+ " 0 | \n",
+ " 2 | \n",
+ " 2 | \n",
+ " -16.812330 | \n",
+ " 0.001490 | \n",
+ " 0.000000 | \n",
+ "
\n",
+ " \n",
+ " | 8 | \n",
+ " 1 | \n",
+ " 0 | \n",
+ " 0 | \n",
+ " -16.730423 | \n",
+ " 0.001145 | \n",
+ " 6190.928719 | \n",
+ "
\n",
+ " \n",
+ " | 9 | \n",
+ " 1 | \n",
+ " 0 | \n",
+ " 1 | \n",
+ " -16.750052 | \n",
+ " 0.000903 | \n",
+ " 12034.494479 | \n",
+ "
\n",
+ " \n",
+ " | 10 | \n",
+ " 1 | \n",
+ " 0 | \n",
+ " 2 | \n",
+ " -16.755024 | \n",
+ " 0.002155 | \n",
+ " 25005.136854 | \n",
+ "
\n",
+ " \n",
+ " | 11 | \n",
+ " 1 | \n",
+ " 1 | \n",
+ " 0 | \n",
+ " -16.801767 | \n",
+ " 0.000844 | \n",
+ " 8815.614302 | \n",
+ "
\n",
+ " \n",
+ " | 12 | \n",
+ " 1 | \n",
+ " 1 | \n",
+ " 1 | \n",
+ " -16.816471 | \n",
+ " 0.001163 | \n",
+ " 18745.997369 | \n",
+ "
\n",
+ " \n",
+ " | 13 | \n",
+ " 1 | \n",
+ " 1 | \n",
+ " 2 | \n",
+ " -16.816585 | \n",
+ " 0.001368 | \n",
+ " 0.000000 | \n",
+ "
\n",
+ " \n",
+ " | 14 | \n",
+ " 1 | \n",
+ " 2 | \n",
+ " 0 | \n",
+ " -16.803039 | \n",
+ " 0.001425 | \n",
+ " 13620.759366 | \n",
+ "
\n",
+ " \n",
+ " | 15 | \n",
+ " 1 | \n",
+ " 2 | \n",
+ " 1 | \n",
+ " -16.812902 | \n",
+ " 0.001359 | \n",
+ " 0.000000 | \n",
+ "
\n",
+ " \n",
+ " | 16 | \n",
+ " 1 | \n",
+ " 2 | \n",
+ " 2 | \n",
+ " -16.817182 | \n",
+ " 0.001203 | \n",
+ " 0.000000 | \n",
+ "
\n",
+ " \n",
+ " | 17 | \n",
+ " 2 | \n",
+ " 0 | \n",
+ " 0 | \n",
+ " -16.766942 | \n",
+ " 0.001670 | \n",
+ " 7466.260457 | \n",
+ "
\n",
+ " \n",
+ " | 18 | \n",
+ " 2 | \n",
+ " 0 | \n",
+ " 1 | \n",
+ " -16.779606 | \n",
+ " 0.001174 | \n",
+ " 13328.825939 | \n",
+ "
\n",
+ " \n",
+ " | 19 | \n",
+ " 2 | \n",
+ " 0 | \n",
+ " 2 | \n",
+ " -16.792576 | \n",
+ " 0.000622 | \n",
+ " 0.000000 | \n",
+ "
\n",
+ " \n",
+ " | 20 | \n",
+ " 2 | \n",
+ " 1 | \n",
+ " 0 | \n",
+ " -16.792028 | \n",
+ " 0.001631 | \n",
+ " 9854.208552 | \n",
+ "
\n",
+ " \n",
+ " | 21 | \n",
+ " 2 | \n",
+ " 1 | \n",
+ " 1 | \n",
+ " -16.811469 | \n",
+ " 0.001701 | \n",
+ " 0.000000 | \n",
+ "
\n",
+ " \n",
+ " | 22 | \n",
+ " 2 | \n",
+ " 1 | \n",
+ " 2 | \n",
+ " -16.819245 | \n",
+ " 0.001173 | \n",
+ " 0.000000 | \n",
+ "
\n",
+ " \n",
+ " | 23 | \n",
+ " 2 | \n",
+ " 2 | \n",
+ " 0 | \n",
+ " -16.802373 | \n",
+ " 0.000936 | \n",
+ " 0.000000 | \n",
+ "
\n",
+ " \n",
+ " | 24 | \n",
+ " 2 | \n",
+ " 2 | \n",
+ " 1 | \n",
+ " -16.759757 | \n",
+ " 0.000936 | \n",
+ " 0.000000 | \n",
+ "
\n",
+ " \n",
+ " | 25 | \n",
+ " 2 | \n",
+ " 2 | \n",
+ " 2 | \n",
+ " -16.630725 | \n",
+ " 0.001892 | \n",
+ " 0.000000 | \n",
+ "
\n",
+ " \n",
+ "
\n",
+ "