Deal with lmp.numpy.extract_atoms('f') returning an array larger than nlocal

.github/workflows/pytests.yml

@@ -21,9 +21,9 @@ jobs:
       use-mpi: True
 
     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 }}
 
@@ -39,7 +39,7 @@ jobs:
         run: conda install pip
 
       - name: Install pytest requirements from pip (specific pymatnext dependencies will be automatically installed when it is)
-        run: pip install wheel setuptools flake8 pytest pytest-cov
+        run: pip install wheel setuptools ruff pytest pytest-cov
 
       - name: Install latest ASE from gitlab
         run: pip install git+https://gitlab.com/ase/ase.git
@@ -59,12 +59,10 @@ jobs:
           conda install -c conda-forge mpi4py openmpi pytest-mpi
           pip install mpipool
 
-      - name: Lint with flake8
+      - name: Lint with ruff
         run: |
           # stop the build if there are Python syntax errors or undefined names
-          flake8 pymatnext/ --count --select=E9,F63,F7,F82 --show-source --statistics
-          # exit-zero treats all errors as warnings.
-          flake8 pymatnext/ --count --exit-zero --max-complexity=10 --max-line-length=120 --statistics
+          ruff check pymatnext
 
       - name: Test with pytest - plain
         if: env.coverage-on-version != matrix.python-version
@@ -95,7 +93,7 @@ jobs:
           mpirun -n 2 pytest --cov=pymatnext --cov-report term --cov-config=tests/.coveragerc --cov-report term-missing --cov-report term:skip-covered --with-mpi -k mpi --cov-append
 
       - name: 'Upload Coverage Data'
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v4
         if: env.coverage-on-version == matrix.python-version
         with:
           name: coverage-html-${{ matrix.python-version }}

pymatnext/analysis/tools/__init__.py

@@ -1 +1 @@
-from .xrd_pbc import analysis as xrd_pbc
+from .xrd_pbc import analysis as xrd_pbc # noqa: F401

pymatnext/analysis/tools/xrd_pbc.py

@@ -1,4 +1,7 @@
 import numpy as np
+
+from ase.utils.xrdebye import wavelengths
+
 ASF_c = np.asarray([
     0.0, 
     0.001305, 
@@ -364,8 +367,6 @@
 
 last_params = {}
 
-from ase.utils.xrdebye import wavelengths
-
 def analysis(atoms=None, twotheta_range=[1.0, 180.0], twotheta_n=None, wavelength=wavelengths['CuKa1'], do_Lorentz=False, header=False):
     """do XRD analysis on one atomic config
     Parameters

pymatnext/analysis/utils.py

@@ -1,6 +1,4 @@
 import numpy as np
-from scipy import stats
-import warnings
 
 def calc_log_a(iters, n_walkers, n_cull, each_cull=False):
     if each_cull:
@@ -9,9 +7,10 @@ def calc_log_a(iters, n_walkers, n_cull, each_cull=False):
         # also assume that iters array increments by one for each cull (i.e. not exactly NS iters)
         # X_n = \prod_{i=0}^n \frac{N-i\%P}{N+1-i\%P}
         # \log X_n = \sum_{i=0}^n \log (N-i\%P) - \log(N+1-i\%P)
-        i_range_mod_n_cull = np.array(range(0, iters[-1] + 1)) % n_cull
-        log_X_n_term = np.log(n_walkers - i_range_mod_n_cull) - np.log(n_walkers + 1 - i_range_mod_n_cull)
-        log_X_n = np.cumsum(log_X_n_term)
+        ## using leading underscore to suppress ruff for this unused bit of code
+        _i_range_mod_n_cull = np.array(range(0, iters[-1] + 1)) % n_cull
+        _log_X_n_term = np.log(n_walkers - _i_range_mod_n_cull) - np.log(n_walkers + 1 - _i_range_mod_n_cull)
+        _log_X_n = np.cumsum(_log_X_n_term)
         # a(iter[i]) = X(iter[i-1]) - X(iter[i])
         #     = prod(0..iter[i-1]) (N-i%P)/(N+1-i%P) - prod(0..iter[i]) (N-i%P)/(N+1-i%P)
         #     = [ prod(0..iter[i-1]) (N-i%P)/(N+1-i%P) ] * (1 - prod(iter[i-1]+1..iter[i]) (N-i%P)/(N+1-i%P))
@@ -59,7 +58,7 @@ def calc_Z_terms(beta, log_a, Es, flat_V_prior=False, N_atoms=None, Vs=None):
     return (Z_term, log_shift)
 
 
-def analyse_T(T, Es, E_shift, Vs, extra_vals, log_a, flat_V_prior, N_atoms, kB, n_extra_DOF, p_entropy_min=5.0):
+def analyse_T(T, Es, E_shift, Vs, extra_vals, log_a, flat_V_prior, N_atoms, kB, n_extra_DOF, p_entropy_min=5.0, sum_f=np.sum):
     """Do an analysis at a single temperature
 
     Note that some parameters (e.g. percentiles used for low and high extent of the contributing iterations,
@@ -106,20 +105,20 @@ def analyse_T(T, Es, E_shift, Vs, extra_vals, log_a, flat_V_prior, N_atoms, kB,
     #     Z_term = Z_term_true * exp(-log_shift)
     # exp(-log_shift) constant factor doesn't matter for quantities that are calculated from sums
     # weighted with Z_term and normalized by Z_term_sum
-    Z_term_sum = np.sum(Z_term)
+    Z_term_sum = sum_f(Z_term)
 
-    U_pot = np.sum(Z_term * Es) / Z_term_sum
+    U_pot = sum_f(Z_term * Es) / Z_term_sum
 
     if N_atoms is not None:
-        N = np.sum(Z_term * N_atoms) / Z_term_sum
+        N = sum_f(Z_term * N_atoms) / Z_term_sum
         n_extra_DOF * N
 
     U = n_extra_DOF / (2.0 * beta) + U_pot + E_shift
 
     Cvp = n_extra_DOF * kB / 2.0 + kB * beta * beta * (sum(Z_term * Es**2) / Z_term_sum - U_pot**2)
 
     if Vs is not None:
-        V = np.sum(Z_term * Vs) / Z_term_sum
+        V = sum_f(Z_term * Vs) / Z_term_sum
         thermal_exp = -1.0 / V * kB * beta * beta * (sum(Z_term * Vs) * sum(Z_term * Es) / Z_term_sum - sum(Z_term * Vs * Es)) / Z_term_sum
     else:
         V = None
@@ -128,7 +127,7 @@ def analyse_T(T, Es, E_shift, Vs, extra_vals, log_a, flat_V_prior, N_atoms, kB,
     if extra_vals is not None and len(extra_vals) > 0:
         extra_vals_out = []
         for v in (extra_vals):
-            extra_vals_out.append(np.sum(Z_term * v, axis=-1) / Z_term_sum)
+            extra_vals_out.append(sum_f(Z_term * v, axis=-1) / Z_term_sum)
 
     # undo shift of Z_term
     log_Z = np.log(Z_term_sum) + log_shift
@@ -165,7 +164,7 @@ def analyse_T(T, Es, E_shift, Vs, extra_vals, log_a, flat_V_prior, N_atoms, kB,
 
     probabilities = Z_term / Z_term_sum
     probabilities = probabilities[np.where(probabilities > 0.0)]
-    p_entropy = -np.sum(probabilities * np.log(probabilities))
+    p_entropy = -sum_f(probabilities * np.log(probabilities))
 
     results_dict.update({'low_percentile_config': low_percentile_config,
                          'mode_config': mode_config,

pymatnext/cli/ns_analyse.py

@@ -15,17 +15,17 @@
 try:
     import ase.units
     GPa = ase.units.GPa
-except:
+except ModuleNotFoundError:
     GPa = None
 
 try:
     from mpi4py import MPI
-except:
+except ModuleNotFoundError:
     MPI = None
 
 try:
     from matplotlib.figure import Figure
-except:
+except ModuleNotFoundError:
     Figure = None
 
 def main():
@@ -62,9 +62,9 @@ def main():
         args.nT = 1
 
     if args.accurate_sum:
-        sum=math.fsum
+        sum_f = math.fsum
     else:
-        sum=np.sum
+        sum_f = np.sum
 
     if MPI is not None:
         comm_rank = MPI.COMM_WORLD.Get_rank()
@@ -102,10 +102,10 @@ def colname(colname_str):
         analysis_header = None
         with open(infile) as fin:
             # find header, either real data or previous analysis
-            for l in fin:
-                if l.startswith('#'):
+            for line in fin:
+                if line.startswith('#'):
                     try:
-                        header = json.loads(l[1:])
+                        header = json.loads(line[1:])
                         if isinstance(header, list):
                             analysis_header = header
                             header = None
@@ -126,12 +126,12 @@ def colname(colname_str):
 
             if analysis_header is None:
                 # read real sampled data
-                for l_i, l in enumerate(fin):
+                for l_i, line in enumerate(fin):
                     if args.line_end is not None and l_i >= args.line_end:
                         break
                     if l_i < args.line_skip or l_i % args.interval != 0:
                         continue
-                    f = l.split()
+                    f = line.split()
                     ## should we ignore certain kinds of malformed lines?
                     ## try:
                     it = int(f[0])
@@ -199,7 +199,8 @@ def colname(colname_str):
                 T = args.Tmin + i_T * args.dT
                 results_dict = utils.analyse_T(T, Es, E_min, Vs, vals, log_a, flat_V_prior, natoms,
                                                args.kB, header.get('n_extra_DOF_per_atom', 3),
-                                               p_entropy_min=args.probability_entropy_minimum)
+                                               p_entropy_min=args.probability_entropy_minimum,
+                                               sum_f=sum_f)
                 if item_keys is None:
                     item_keys = list(results_dict.keys())
                     try:
@@ -225,7 +226,8 @@ def colname(colname_str):
             try:
                 data = MPI.COMM_WORLD.gather(data, root = 0)
                 data = [item for sublist in data for item in sublist]
-            except:
+            except Exception as exc:
+                print("BOB got exception", exc)
                 pass
 
         else:
@@ -234,14 +236,14 @@ def colname(colname_str):
             T_max = args.Tmin + (args.nT - 1) * args.dT
             data = []
             with open(infile) as fin:
-                for l in fin:
-                    if l.startswith('#'):
+                for line in fin:
+                    if line.startswith('#'):
                         continue
-                    fields = l.strip().split()
+                    fields = line.strip().split()
                     for f_i, f in enumerate(fields):
                         try:
                             fields[f_i] = float(f)
-                        except:
+                        except ValueError:
                             fields[f_i] = f
 
                     if fields[0] >= args.Tmin and fields[0] <= T_max:

pymatnext/cli/ns_analyse_traj.py

@@ -5,7 +5,7 @@
 
 try:
     from matplotlib.figure import Figure
-except:
+except ModuleNotFoundError:
     Figure = None
 
 import sys
@@ -77,9 +77,9 @@ def main():
             raise RuntimeError("Need --samples_file OR --walkers and --cull")
 
     if args.accurate_sum:
-        sum=math.fsum
+        sum_f = math.fsum
     else:
-        sum=np.sum
+        sum_f = np.sum
 
     def traj_configs(trajfiles, ns_iters):
         ns_iters = ns_iters.split(":")
@@ -135,7 +135,7 @@ def traj_configs(trajfiles, ns_iters):
         analysis_mod_json = analysis_str.split(maxsplit=1)
         try:
             analysis_func = import_module(analysis_mod_json[0]).analysis
-        except:
+        except ModuleNotFoundError:
             analysis_func = import_module('pymatnext.analysis.tools.' + analysis_mod_json[0]).analysis
 
         analysis_args = []
@@ -159,7 +159,6 @@ def traj_configs(trajfiles, ns_iters):
     # check for cached
     found_cached = []
     if args.cache:
-        trajfile_a = [np.asarray(list(f)) for f in args.trajfile]
         not_all = False
         for i in range(len(args.trajfile[0])):
             if not all([f[i] == args.trajfile[0][i] for f in args.trajfile]):
@@ -191,7 +190,7 @@ def traj_configs(trajfiles, ns_iters):
         iterator = traj_configs(args.trajfile, args.ns_iters)
     else:
         iterator = tqdm(traj_configs(args.trajfile, args.ns_iters))
-        iterator.set_description(f"iter <skipping>")
+        iterator.set_description("iter <skipping>")
     for at in iterator:
         if all(at.numbers == 0):
             at.numbers = at.arrays['type']
@@ -256,7 +255,7 @@ def traj_configs(trajfiles, ns_iters):
         outfiles[analysis] = open(args.output + '.' + analysis + '.data', 'w')
     for T_i, T in enumerate(args.temperature):
         results_dict = utils.analyse_T(T, Es, E_min, Vs, extra_vals, log_a, args.flat_V_prior, natoms,
-                                       args.kB, 0, args.delta_P is not None and args.delta_P != 0.0)
+                                       args.kB, 0, args.delta_P is not None and args.delta_P != 0.0, sum_f=sum_f)
         for analysis, res in zip(args.analysis, results_dict['extra_vals']):
             outfiles[analysis].write(f'# T {T} analysis {analysis}\n')
             for v in res.T:

pymatnext/cli/sample.py

@@ -6,7 +6,6 @@
 import re
 
 import time
-import importlib
 import pprint
 import itertools
 import json
@@ -15,7 +14,6 @@
 from argparse import ArgumentParser
 
 import toml
-import numpy as np
 
 from pymatnext.ns import NS
 from pymatnext.params import check_fill_defaults
@@ -43,9 +41,9 @@ def init_MPI():
         # from https://stackoverflow.com/questions/49868333/fail-fast-with-mpi4py
         def mpiabort_excepthook(type, value, traceback_obj):
             sys.stderr.write(f"{MPI.COMM_WORLD.rank} Aborting because of exception {value}\n")
-            for l in traceback.format_tb(traceback_obj):
-                for ll in l.splitlines():
-                    sys.stderr.write(f"{MPI.COMM_WORLD.rank} {ll.rstrip()}\n")
+            for line in traceback.format_tb(traceback_obj):
+                for line_split in line.splitlines():
+                    sys.stderr.write(f"{MPI.COMM_WORLD.rank} {line_split.rstrip()}\n")
             sys.stderr.flush()
             MPI.COMM_WORLD.Abort()
             sys.__excepthook__(type, value, traceback_obj)
@@ -190,7 +188,7 @@ def sample(args, MPI, NS_comm, walker_comm):
             # truncate .NS_samples file
             f_samples = open(ns_file_name, "r+")
             # skip header
-            l = f_samples.readline()
+            _ = f_samples.readline()
             line_i = None
             while True:
                 line = f_samples.readline()

pymatnext/loop_exit/__init__.py

@@ -1,8 +1,5 @@
-from copy import deepcopy
 import importlib
 
-import numpy as np
-
 from pymatnext.params import check_fill_defaults
 from .loop_exit_params import param_defaults
 

pymatnext/ns.py

@@ -1,4 +1,3 @@
-from copy import deepcopy
 import importlib
 
 import re