Fix format and lint errors and relax requirements (#309)

* Allow python > 3.11
* Allow numpy >= 1.26
* Bump CI python versions (min version is 3.10)
* Fix lint / format errors introduced in #296 

I split the AFQMC driver into a base class and then implementations for
FT / ZT afqmc which allows for more customization of the factory
methods.

.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!
@@ -108,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: |
@@ -131,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" fqe
+          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

.gitignore

@@ -108,4 +108,6 @@ 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*$

dev/run_tests.py

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

examples/03-custom_observable/run_afqmc.py

@@ -15,6 +15,8 @@
 # Author: Fionn Malone <[email protected]>
 #
 
+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/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)

examples/08-custom_walker/run_afqmc.py

@@ -94,7 +94,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)

examples/16-ft_afqmc/run_afqmc.py

@@ -1,11 +1,12 @@
 import json
-import numpy
 
+import numpy
 from ueg import UEG
-from ipie.config import MPI
+
 from ipie.addons.thermal.qmc.calc import build_thermal_afqmc_driver
-from ipie.analysis.extraction import extract_observable
 from ipie.analysis.autocorr import reblock_by_autocorr
+from ipie.analysis.extraction import extract_observable
+from ipie.config import MPI
 
 comm = MPI.COMM_WORLD
 
@@ -65,7 +66,7 @@
 # 3. Run thermal AFQMC calculation.
 afqmc.run(verbose=verbose)
 afqmc.finalise()
-afqmc.estimators.compute_estimators(afqmc.hamiltonian, afqmc.trial, afqmc.walkers)
+afqmc.estimators.compute_estimators(hamiltonian=afqmc.hamiltonian, trial=afqmc.trial, walker_batch=afqmc.walkers)
 
 if comm.rank == 0:
     energy_data = extract_observable(afqmc.estimators.filename, "energy")

ipie/addons/free_projection/estimators/energy.py

@@ -30,7 +30,7 @@ def __init__(
     ):
         super().__init__(system, ham, trial, filename)
 
-    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(system, hamiltonian, walkers, trial)

ipie/addons/free_projection/estimators/tests/test_estimators.py

@@ -68,8 +68,12 @@ def test_estimator_handler_fp():
         handler["energy1"] = estim
         handler.json_string = ""
         handler.initialize(comm)
-        handler.compute_estimators(comm, system, ham, trial, walker_batch)
-        handler.compute_estimators(comm, system, ham, trial, walker_batch)
+        handler.compute_estimators(
+            system=system, hamiltonian=ham, trial=trial, walker_batch=walker_batch
+        )
+        handler.compute_estimators(
+            system=system, hamiltonian=ham, trial=trial, walker_batch=walker_batch
+        )
 
 
 if __name__ == "__main__":

ipie/addons/free_projection/qmc/fp_afqmc.py

@@ -62,7 +62,7 @@ def build(
         trial_wavefunction,
         walkers=None,
         num_walkers: int = 100,
-        seed: int = None,
+        seed: Optional[int] = None,
         num_steps_per_block: int = 25,
         num_blocks: int = 100,
         timestep: float = 0.005,
@@ -282,7 +282,7 @@ def setup_estimators(
 
     def run(
         self,
-        psi=None,
+        walkers=None,
         estimator_filename="estimate.h5",
         verbose=True,
         additional_estimators: Optional[Dict[str, EstimatorBase]] = None,
@@ -300,8 +300,8 @@ def run(
         """
         self.setup_timers()
         tzero_setup = time.time()
-        if psi is not None:
-            self.walkers = psi
+        if walkers is not None:
+            self.walkers = walkers
         self.setup_timers()
         eshift = 0.0
         self.walkers.orthogonalise()
@@ -360,7 +360,10 @@ def run(
                 start = time.time()
                 if step % self.params.num_steps_per_block == 0:
                     self.estimators[block_number].compute_estimators(
-                        comm, self.system, self.hamiltonian, self.trial, self.walkers
+                        system=self.system,
+                        hamiltonian=self.hamiltonian,
+                        trial=self.trial,
+                        walker_batch=self.walkers,
                     )
                     self.estimators[block_number].print_block(
                         comm,

ipie/addons/thermal/analysis/extraction.py

@@ -19,6 +19,7 @@
 
 from ipie.utils.misc import get_from_dict
 
+
 def set_info(frame, md):
     ncols = len(frame.columns)
     system = md.get("system")
@@ -61,4 +62,3 @@ def set_info(frame, md):
         frame["cholesky_treshold"] = chol
 
     return list(frame.columns[ncols:])
-

ipie/addons/thermal/analysis/thermal_analysis.py

@@ -18,21 +18,16 @@
 
 #!/usr/bin/env python
 
-import sys
 import argparse
-
 import glob
+import sys
+
 import numpy
-import scipy.optimize
 import pandas as pd
-
-from ipie.analysis.extraction import (
-        extract_observable,
-        get_metadata, 
-        get_sys_param
-        )
+import scipy.optimize
 
 from ipie.addons.thermal.analysis.extraction import set_info
+from ipie.analysis.extraction import extract_observable, get_metadata
 
 
 def parse_args(args):
@@ -49,18 +44,10 @@ def parse_args(args):
         Command line arguments.
     """
 
-    parser = argparse.ArgumentParser(description = __doc__)
-    parser.add_argument('-c', '--chem-pot', dest='fit_chem_pot',
-                        action='store_true', default=False,
-                        help='Estimate optimal chemical potential')
-    parser.add_argument('-n', '--nav', dest='nav', type=float,
-                        help='Target electron density.')
-    parser.add_argument('-o', '--order', dest='order', type=int,
-                        default=3, help='Order polynomial to fit.')
-    parser.add_argument('-p', '--plot', dest='plot', action='store_true',
-                        help='Plot density vs. mu.')
-    parser.add_argument('-f', nargs='+', dest='filenames',
-                        help='Space-separated list of files to analyse.')
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "-f", nargs="+", dest="filenames", help="Space-separated list of files to analyse."
+    )
 
     options = parser.parse_args(args)
 
@@ -76,9 +63,9 @@ def analyse(files, block_idx=1):
     files = sorted(files)
 
     for f in files:
-        data_energy = extract_observable(f, name='energy', block_idx=block_idx)
-        data_nav = extract_observable(f, name='nav', block_idx=block_idx)
-        data = pd.concat([data_energy, data_nav['Nav']], axis=1)
+        data_energy = extract_observable(f, name="energy", block_idx=block_idx)
+        data_nav = extract_observable(f, name="nav", block_idx=block_idx)
+        data = pd.concat([data_energy, data_nav["Nav"]], axis=1)
         md = get_metadata(f)
         keys = set_info(data, md)
         sims.append(data[1:])
@@ -108,20 +95,6 @@ def nav_mu(mu, coeffs):
     return numpy.polyval(coeffs, mu)
 
 
-def find_chem_pot(data, target, vol, order=3, plot=False):
-    print(f"# System volume: {vol}.")
-    print(f"# Target number of electrons: {vol * target}.")
-    nav = data.Nav.values / vol
-    nav_error = data.Nav_error.values / vol
-    # Half filling special case where error bar is zero.
-    zeros = numpy.where(nav_error == 0)[0]
-    nav_error[zeros] = 1e-8
-    mus = data.mu.values
-    delta = nav - target
-    s = 0
-    e = len(delta)
-    rmin = None
-
 def main(args):
     """Run reblocking and data analysis on PAUXY output.
 
@@ -136,28 +109,16 @@ def main(args):
     """
 
     options = parse_args(args)
-    if '*' in options.filenames[0]:
+    if "*" in options.filenames[0]:
         files = glob.glob(options.filenames[0])
 
     else:
         files = options.filenames
 
     data = analyse(files)
 
-    if options.fit_chem_pot:
-        name = get_sys_param(files[0], 'name')
-        vol = 1.
-        mu = find_chem_pot(data, options.nav, vol,
-                           order=options.order, plot=options.plot)
-
-        if mu is not None:
-            print("# Optimal chemical potential found to be: {}.".format(mu))
-
-        else:
-            print("# Failed to find chemical potential.")
-
     print(data.to_string(index=False))
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main(sys.argv[1:])