have to wrap delta with np.array to make sure it is ndarray and result stays CoordinateFrame

Reviewed-on: #5

src/mdevaluate/checksum.py

@@ -73,7 +73,9 @@ def checksum(*args, csum=None):
         elif isinstance(arg, FunctionType):
             csum.update(strip_comments(inspect.getsource(arg)).encode())
             c = inspect.getclosurevars(arg)
-            for v in {**c.nonlocals, **c.globals}.values():
+            merged = {**c.nonlocals, **c.globals}
+            for key in sorted(merged):  # deterministic ordering
+                v = merged[key]
                 if v is not arg:
                     checksum(v, csum=csum)
         elif isinstance(arg, functools.partial):

src/mdevaluate/distribution.py

@@ -182,10 +182,10 @@ def tetrahedral_order(
     )
 
     # Connection vectors
-    neighbors_1 -= atoms
+    neighbors_1 = pbc_diff(neighbors_1, atoms, box=atoms.box)
-    neighbors_2 -= atoms
+    neighbors_2 = pbc_diff(neighbors_2, atoms, box=atoms.box)
-    neighbors_3 -= atoms
+    neighbors_3 = pbc_diff(neighbors_3, atoms, box=atoms.box)
-    neighbors_4 -= atoms
+    neighbors_4 = pbc_diff(neighbors_4, atoms, box=atoms.box)
 
     # Normed Connection vectors
     neighbors_1 /= np.linalg.norm(neighbors_1, axis=-1).reshape(-1, 1)

src/mdevaluate/extra/free_energy_landscape.py

@@ -4,7 +4,6 @@ from typing import Optional
 import numpy as np
 from numpy.typing import ArrayLike, NDArray
 from numpy.polynomial.polynomial import Polynomial as Poly
-import math
 from scipy.spatial import KDTree
 import pandas as pd
 import multiprocessing as mp
@@ -49,7 +48,7 @@ def _pbc_points_reduced(
 
 def _build_tree(points, box, r_max, pore_geometry):
     if np.all(np.diag(np.diag(box)) == box):
-        tree = KDTree(points, boxsize=box)
+        tree = KDTree(points % box, boxsize=box)
         points_pbc_index = None
     else:
         points_pbc, points_pbc_index = _pbc_points_reduced(
@@ -79,8 +78,7 @@ def occupation_matrix(
     z_bins = np.arange(0, box[2][2] + edge_length, edge_length)
     bins = [x_bins, y_bins, z_bins]
     # Trajectory is split for parallel computing
-    size = math.ceil(len(frame_indices) / nodes)
+    indices = np.array_split(frame_indices, nodes)
-    indices = [frame_indices[i : i + size] for i in range(0, len(frame_indices), size)]
     pool = mp.Pool(nodes)
     results = pool.map(
         partial(_calc_histogram, trajectory=trajectory, bins=bins), indices
@@ -274,7 +272,11 @@ def distance_resolved_energies(
 
 
 def find_energy_maxima(
-    energy_df: pd.DataFrame, r_min: float, r_max: float
+    energy_df: pd.DataFrame,
+    r_min: float,
+    r_max: float,
+    r_eval: float = None,
+    degree: int = 2,
 ) -> pd.DataFrame:
     distances = []
     energies = []
@@ -283,6 +285,9 @@ def find_energy_maxima(
         x = np.array(data_d["r"])
         y = np.array(data_d["energy"])
         mask = (x >= r_min) * (x <= r_max)
-        p3 = Poly.fit(x[mask], y[mask], deg=2)
+        p3 = Poly.fit(x[mask], y[mask], deg=degree)
-        energies.append(np.max(p3(np.linspace(r_min, r_max, 1000))))
+        if r_eval is None:
+            energies.append(np.max(p3(np.linspace(r_min, r_max, 1000))))
+        else:
+            energies.append(p3(r_eval))
     return pd.DataFrame({"d": distances, "energy": energies})

src/mdevaluate/pbc.py

@@ -149,32 +149,21 @@ def nojump(frame: CoordinateFrame, usecache: bool = True) -> CoordinateFrame:
             i0 = 0
             delta = 0
 
-        delta = (
+        delta = (delta
-            delta
+            + np.vstack(
-            + np.array(
+                [m[i0 : abstep + 1].sum(axis=0) for m in reader.nojump_matrices]
-                np.vstack(
+            ).T)
-                    [m[i0 : abstep + 1].sum(axis=0) for m in reader.nojump_matrices]
-                ).T
-            )
-            @ frame.box
-        )
 
         reader._nojump_cache[abstep] = delta
         while len(reader._nojump_cache) > NOJUMP_CACHESIZE:
             reader._nojump_cache.popitem(last=False)
-        delta = delta[selection, :]
     else:
-        delta = (
+        delta = np.vstack(
-            np.array(
+                [m[: frame.step + 1, selection].sum(axis=0) for m in reader.nojump_matrices]
-                np.vstack(
-                    [
-                        m[: frame.step + 1, selection].sum(axis=0)
-                        for m in reader.nojump_matrices
-                    ]
                 ).T
-            )
+    
-            @ frame.box
+    delta = delta[selection, :]
-        )
+    delta = np.array(delta @ frame.box)
     return frame - delta