Merge develop into main (#310)

.github/workflows/ci.yml

@@ -13,11 +13,11 @@ jobs:
       - name: Setup MPI
         uses: mpi4py/setup-mpi@v1
         with:
-          mpi: openmpi 
+          mpi: openmpi
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -r requirements.txt 
+          pip install -r requirements.txt
           pip install -r dev/dev.txt
       - name: pylint
         run: |
@@ -32,7 +32,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: [ '3.8', '3.9', '3.10' ]
+        python-version: [ '3.10', '3.11', '3.12' ]
         mpi:
           - openmpi
     runs-on:  ubuntu-latest
@@ -53,7 +53,7 @@ jobs:
         pip install -r dev/dev.txt
         pip install pyblock
         # https://github.com/JoonhoLee-Group/ipie/issues/278
-        pip install "pyscf<=2.3.0"
+        pip install pyscf
     - name: Install package
       run: |
         # HACK FOR LEGACY CODE!
@@ -81,10 +81,21 @@ jobs:
           python -m pip install --upgrade pip
           pip install -r requirements.txt
           pip install pytest-xdist
+      - name: Install package
+        run: |
+          python -m pip install -e .
       - name: No mpi4py
         run: |
           python -c "from ipie.config import MPI; assert \"FakeComm\" in str(MPI.COMM_WORLD)"
           python -m pytest -n=auto
+      - name: wicks helper
+        run: |
+          sudo apt-get -y update && sudo apt-get -y install cmake
+          cc --version
+          root=$(pwd)
+          cd ipie/lib/wicks && mkdir build && cd build && cmake .. && make VERBOSE=1 && cd $root
+          # pytest ipie/trial_wavefunction/ -m wicks
+          # pytest ipie/lib/wicks -m wicks
   integration:
     strategy:
       fail-fast: false
@@ -97,11 +108,11 @@ jobs:
       - name: Setup MPI
         uses: mpi4py/setup-mpi@v1
         with:
-          mpi: openmpi 
+          mpi: openmpi
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -r requirements.txt 
+          pip install -r requirements.txt
           pip install -r dev/dev.txt
       - name: Install package
         run: |
@@ -120,14 +131,14 @@ jobs:
       - name: Install dependencies
         run: |
           python -m pip install --upgrade pip
-          pip install -r requirements.txt 
+          pip install -r requirements.txt
           pip install pyblock
           # https://github.com/JoonhoLee-Group/ipie/issues/278
-          pip install "pyscf<=2.3.0"
+          pip install pyscf fqe
       - name: Install package
         run: |
           python -m pip install -e .
       - name: Test Examples
         timeout-minutes: 10
         run: |
-          python dev/run_tests.py --examples
+          python dev/run_tests.py --examples

.github/workflows/publish.yml

@@ -1,4 +1,4 @@
-name: Publish ipie 
+name: Publish ipie
 
 on:
   push:
@@ -37,7 +37,7 @@ jobs:
         name: python-package-distributions
         path: dist/
     - name: Test the wheel
-      run: | 
+      run: |
         ls dist
         wheel=$(find dist/*whl)
         echo $wheel
@@ -48,9 +48,9 @@ jobs:
     - name: Setup MPI
       uses: mpi4py/setup-mpi@v1
       with:
-        mpi: openmpi 
+        mpi: openmpi
     - name: Test the wheel
-      run: | 
+      run: |
         pip install mpi4py
         python -c "import mpi4py"
         python -c "from ipie.config import MPI; assert \"mpi4py\" in str(MPI.COMM_WORLD)"

.gitignore

@@ -108,4 +108,7 @@ ipie/qmc/tests/reference_data/**/*.h5
 *.txt
 *wheels*
 *FCIDUMP*
-*out*
+*out*
+
+*.code-workspace
+

.pylintrc

@@ -23,8 +23,8 @@ output-format=colorized
 disable=all
 enable=
       no-member,
-    # renable once classes are tidied
-;     attribute-defined-outside-init,
+      # renable once classes are tidied
+      # attribute-defined-outside-init,
       consider-using-f-string,
       useless-object-inheritance,
       unused-variable,
@@ -68,4 +68,4 @@ enable=
       no-else-return,
 # https://github.com/pylint-dev/pylint/issues/2178
 # ignore urls
-ignore-long-lines=^\s*(# )?<?https?://\S+>?$|^\s*(\w*\s*=\s*)?(\"|\').*(\"|\'),?\s*$
+ignore-long-lines=^\s*(# )?<?https?://\S+>?$|^\s*(\w*\s*=\s*)?(\"|\').*(\"|\'),?\s*$

README.rst

@@ -22,6 +22,7 @@ ipie inherits a lot of QMC features from pauxy.
     :target: https://arxiv.org/abs/2209.04015
 
 Copyright by Joonho Lee (joonholee@g.harvard.edu)
+Our first release paper can be found at https://pubs.acs.org/doi/10.1021/acs.jctc.2c00934
 
 Features
 --------

dev/run_tests.py

@@ -1,4 +1,3 @@
-#!/usr/bin/env python3
 import argparse
 import glob
 import os

docs/source/conf.py

@@ -24,7 +24,7 @@
 
 import mock
 
-MOCK_MODULES = ["h5py","mpi4py"]
+MOCK_MODULES = ["h5py", "mpi4py"]
 for mod_name in MOCK_MODULES:
     sys.modules[mod_name] = mock.Mock()
 autodoc_mock_imports = MOCK_MODULES

docs/source/dev/index.md

@@ -65,3 +65,13 @@ export IPIE_USE_GPU=1; mpirun -np 1 pytest -m gpu -sv
 
 -   Note if running CPU test afterwards it may be necessary to clear the
     environment variable!
+
+
+## Releasing a package
+
+1. Change the version from 'X.Y.Z.dev0' to 'X.Y.Z'
+2. Merge develop into main. Do not squash the merge (enables sensible release notes.)
+3. Push tag 'vX.Y.Z' (no dev0).
+4. Check actions and ensure build and publish steps run.
+5. Create new branch and bump version to 'X.(Y+1).Z.dev0'.
+6. Send PR for this branch into develop.

examples/01-simple/scf.py

@@ -1,4 +1,3 @@
-
 # Copyright 2022 The ipie Developers. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

examples/03-custom_observable/run_afqmc.py

@@ -15,6 +15,8 @@
 # Author: Fionn Malone <fmalone@google.com>
 #
 
+from typing import Dict
+
 import numpy as np
 from pyscf import gto, scf
 
@@ -86,11 +88,11 @@ def __init__(self, ham):
         # Must specify that we're dealing with array valued estimator
         self.scalar_estimator = False
 
-    def compute_estimator(self, system, walker_batch, hamiltonian, trial_wavefunction):
-        trial_wavefunction.calc_greens_function(walker_batch, build_full=True)
-        numer = np.einsum("w,wii->i", walker_batch.weight, walker_batch.Ga + walker_batch.Gb)
+    def compute_estimator(self, system=None, walkers=None, hamiltonian=None, trial=None):
+        trial.calc_greens_function(walkers, build_full=True)
+        numer = np.einsum("w,wii->i", walkers.weight, walkers.Ga + walkers.Gb)
         self["DiagGNumer"] = numer
-        self["DiagGDenom"] = sum(walker_batch.weight)
+        self["DiagGDenom"] = sum(walkers.weight)
 
 
 afqmc = build_afqmc_driver(comm, nelec=mol.nelec)
@@ -127,25 +129,28 @@ def __init__(self, ham):
         # Must specify that we're dealing with array valued estimator
         self.scalar_estimator = False
 
-    def compute_estimator(self, system, walker_batch, hamiltonian, trial_wavefunction):
-        trial_wavefunction.calc_greens_function(walker_batch, build_full=True)
+    def compute_estimator(self, system=None, walkers=None, hamiltonian=None, trial=None):
+        trial.calc_greens_function(walkers, build_full=True)
         numer = np.array(
             [
-                np.einsum("w,wij->ij", walker_batch.weight, walker_batch.Ga),
-                np.einsum("w,wij->ij", walker_batch.weight, walker_batch.Gb),
+                np.einsum("w,wij->ij", walkers.weight, walkers.Ga),
+                np.einsum("w,wij->ij", walkers.weight, walkers.Gb),
             ]
         )
 
         # For multidimensional arrays we must flatten the data
         self["GNumer"] = numer.ravel()
-        self["GDenom"] = sum(walker_batch.weight)
+        self["GDenom"] = sum(walkers.weight)
 
 
 afqmc = build_afqmc_driver(comm, nelec=mol.nelec)
 # Let us override the number of blocks to keep it short
 afqmc.params.num_blocks = 20
 # We can now add this to the estimator handler object in the afqmc driver
-add_est = {"diagG": Diagonal1RDM(ham=afqmc.hamiltonian), "1RDM": Mixed1RDM(ham=afqmc.hamiltonian)}
+add_est: Dict[str, EstimatorBase] = {
+    "diagG": Diagonal1RDM(ham=afqmc.hamiltonian),
+    "1RDM": Mixed1RDM(ham=afqmc.hamiltonian),
+}
 afqmc.run(additional_estimators=add_est)
 # We can extract the qmc data as as a pandas data frame like so
 from ipie.analysis.extraction import extract_observable

examples/05-frozen_core/scf.py

@@ -1,4 +1,3 @@
-
 # Copyright 2022 The ipie Developers. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");

examples/06-gpu/chunked_chol.py

@@ -0,0 +1,23 @@
+from pyscf import cc, gto, scf
+from ipie.utils.mpi import make_splits_displacements
+import h5py
+import numpy as np
+import gc 
+
+
+mol = gto.M(
+    atom=[("H", 1.6 * i, 0, 0) for i in range(0, 4)],
+    basis="sto-6g",
+    verbose=4,
+    unit="Bohr",
+)
+mf = scf.UHF(mol)
+mf.chkfile = "scf.chk"
+mf.kernel()
+
+from ipie.utils.from_pyscf import gen_ipie_input_from_pyscf_chk
+gen_ipie_input_from_pyscf_chk(mf.chkfile, verbose=0)
+
+
+from ipie.utils.chunk_large_chol import split_cholesky
+split_cholesky('hamiltonian.h5', 4) # split the cholesky to 4 subfiles

examples/06-gpu/run_afqmc_chunked.py

@@ -0,0 +1,117 @@
+import sys
+import h5py
+import numpy
+
+from ipie.hamiltonians.generic_chunked import GenericRealCholChunked as HamGeneric
+from ipie.qmc.afqmc import AFQMC
+from ipie.systems.generic import Generic
+from ipie.trial_wavefunction.single_det import SingleDet
+from ipie.utils.mpi import MPIHandler
+
+import os
+from ipie.utils.backend import arraylib as xp
+from pyscf import gto
+
+try:
+    import cupy
+    from mpi4py import MPI
+except ImportError:
+    sys.exit(0)
+
+from chunked_chol import *
+
+mol = gto.M(
+    atom=[("H", 1.6 * i, 0, 0) for i in range(0, 4)],
+    basis="sto-6g",
+    verbose=4,
+    unit="Bohr",
+)
+
+from ipie.config import config
+config.update_option("use_gpu", True)
+
+gpu_number_per_node = 4
+nmembers = 4
+gpu_id = MPI.COMM_WORLD.rank % gpu_number_per_node
+xp.cuda.Device(gpu_id).use()
+
+comm = MPI.COMM_WORLD
+num_walkers = 1224 // comm.size
+nsteps = 25
+nblocks = 100
+timestep = 0.005
+rng_seed = None
+
+
+with h5py.File("hamiltonian.h5") as fa:
+    e0 = fa["e0"][()]
+    hcore = fa["hcore"][()]
+
+rank = comm.Get_rank()
+size = comm.Get_size()
+srank = rank % nmembers
+
+from ipie.utils.mpi import MPIHandler, make_splits_displacements
+handler = MPIHandler(nmembers=nmembers)
+
+from ipie.utils.pack_numba import pack_cholesky
+
+num_basis = hcore.shape[-1]
+with h5py.File(f"chol_{srank}.h5") as fa:
+    chol_chunk = fa["chol"][()]
+
+chunked_chols = chol_chunk.shape[-1]
+num_chol = handler.scomm.allreduce(chunked_chols, op=MPI.SUM)
+
+chol_chunk_view = chol_chunk.reshape((num_basis, num_basis, -1))
+cp_shape = (num_basis * (num_basis + 1) // 2, chol_chunk_view.shape[-1])
+chol_packed_chunk = numpy.zeros(cp_shape, dtype=chol_chunk_view.dtype)
+sym_idx = numpy.triu_indices(num_basis)
+pack_cholesky(sym_idx[0], sym_idx[1], chol_packed_chunk, chol_chunk_view)
+del chol_chunk_view
+
+split_size = make_splits_displacements(num_chol, nmembers)[0]
+assert chunked_chols == split_size[srank]
+
+with h5py.File("wavefunction.h5") as fa:
+    phi0a = fa["phi0_alpha"][()]
+    psiT = fa["psi_T_alpha"][()]
+
+
+num_basis = hcore.shape[-1]
+mol_nelec = mol.nelec
+system = Generic(nelec=mol_nelec)
+ham = HamGeneric(
+    numpy.array([hcore, hcore]),
+    None,
+    chol_chunk,
+    chol_packed_chunk,
+    e0, handler
+)
+ham.nchol = num_chol
+ham.handler = handler
+
+trial = SingleDet(numpy.hstack([psiT, psiT]), mol_nelec, num_basis, handler)
+trial.build()
+trial.half_rotate(ham)
+
+from ipie.walkers.uhf_walkers import UHFWalkers
+from ipie.utils.mpi import MPIHandler
+walkers = UHFWalkers(numpy.hstack([phi0a, phi0a]), system.nup, system.ndown, ham.nbasis, num_walkers, mpi_handler=handler)
+walkers.build(trial)
+
+afqmc = AFQMC.build(
+    mol_nelec,
+    ham,
+    trial,
+    walkers,
+    num_walkers,
+    rng_seed,
+    nsteps,
+    nblocks,
+    timestep,
+    mpi_handler=handler)
+
+
+afqmc.run()
+afqmc.finalise(verbose=True)

examples/07-custom_trial/run_afqmc.py

@@ -91,7 +91,7 @@ def __init__(
             trial=trial,
         )
 
-    def compute_estimator(self, system, walkers, hamiltonian, trial, istep=1):
+    def compute_estimator(self, system, walkers, hamiltonian, trial):
         trial.calc_greens_function(walkers)
         # Need to be able to dispatch here
         energy = local_energy_batch(system, hamiltonian, walkers, trial)