dedent before ast.parse for non-toplevel functions

Reviewed-on: #2

pyproject.toml

@@ -4,10 +4,12 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "mdevaluate"
-version = "24.01"
+version = "24.02"
 dependencies = [
     "mdanalysis",
     "pandas",
     "dask",
     "pathos",
+    "tables",
+    "pyedr"
 ]

requirements.txt

@@ -4,3 +4,4 @@ dask
 pathos
 tables
 pytest
+pyedr

src/mdevaluate/__init__.py

@@ -16,7 +16,7 @@ from . import reader
 from . import system
 from . import utils
 from . import extra
-from .logging import logger
+from .logging_util import logger
 
 
 def open(

src/mdevaluate/autosave.py

@@ -5,7 +5,7 @@ from typing import Optional, Callable, Iterable
 
 import numpy as np
 from .checksum import checksum
-from .logging import logger
+from .logging_util import logger
 
 autosave_directory: Optional[str] = None
 load_autosave_data = False

src/mdevaluate/checksum.py

@@ -1,9 +1,14 @@
 import functools
 import hashlib
-from .logging import logger
+from .logging_util import logger
 from types import ModuleType, FunctionType
 import inspect
 from typing import Iterable
+import ast
+import io
+import tokenize
+import re
+import textwrap
 
 import numpy as np
 
@@ -28,16 +33,43 @@ def version(version_nr: int, calls: Iterable = ()):
     return decorator
 
 
-def strip_comments(s: str):
+def strip_comments(source: str) -> str:
-    """Strips comment lines and docstring from Python source string."""
+    """Removes docstrings, comments, and irrelevant whitespace from Python source code."""
-    o = ""
+
-    in_docstring = False
+    # Step 1: Remove docstrings using AST
-    for l in s.split("\n"):
+    def remove_docstrings(node):
-        if l.strip().startswith(("#", '"', "'")) or in_docstring:
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Module)):
-            in_docstring = l.strip().startswith(('"""', "'''")) + in_docstring == 1
+            if (doc := ast.get_docstring(node, clean=False)):
+                first_stmt = node.body[0]
+                if isinstance(first_stmt, ast.Expr) and isinstance(first_stmt.value, ast.Constant):
+                    node.body.pop(0)  # Remove the docstring entirely
+        for child in ast.iter_child_nodes(node):
+            remove_docstrings(child)
+
+    tree = ast.parse(textwrap.dedent(source))
+    remove_docstrings(tree)
+    code_without_docstrings = ast.unparse(tree)
+
+    # Step 2: Remove comments using tokenize
+    tokens = tokenize.generate_tokens(io.StringIO(code_without_docstrings).readline)
+    result = []
+    last_lineno = -1
+    last_col = 0
+
+    for toknum, tokval, (srow, scol), (erow, ecol), line in tokens:
+        if toknum == tokenize.COMMENT:
             continue
-        o += l + "\n"
+        if srow > last_lineno:
-    return o
+            last_col = 0
+        if scol > last_col:
+            result.append(" " * (scol - last_col))
+        result.append(tokval)
+        last_lineno, last_col = erow, ecol
+
+    code_no_comments = ''.join(result)
+
+    # Step 3: Remove empty lines (whitespace-only or truly blank)
+    return "\n".join([line for line in code_no_comments.splitlines() if line.strip() != ""])
 
 
 def checksum(*args, csum=None):
@@ -73,7 +105,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/coordinates.py

@@ -1,6 +1,6 @@
 from functools import partial, wraps
 from copy import copy
-from .logging import logger
+from .logging_util import logger
 from typing import Optional, Callable, List, Tuple
 
 import numpy as np
@@ -218,7 +218,7 @@ class Coordinates:
             self.get_frame.clear_cache()
 
     def __iter__(self):
-        for i in range(len(self))[self._slice]:
+        for i in range(len(self.frames))[self._slice]:
             yield self[i]
 
     @singledispatchmethod
@@ -232,7 +232,7 @@ class Coordinates:
         return sliced
 
     def __len__(self):
-        return len(self.frames)
+        return len(self.frames[self._slice])
 
     def __checksum__(self):
         return checksum(self.frames, self.atom_filter, self._slice, self.mode)
@@ -261,6 +261,7 @@ class CoordinatesMap:
         self.frames = self.coordinates.frames
         self.atom_subset = self.coordinates.atom_subset
         self.function = function
+        self._slice = slice(None)
         if isinstance(function, partial):
             self._description = self.function.func.__name__
         else:
@@ -645,7 +646,7 @@ def next_neighbors(
         return distances_new, indices_new
     else:
         atoms_pbc, atoms_pbc_index = pbc_points(
-            query_atoms, box, thickness=distance_upper_bound + 0.1, index=True, **kwargs
+            atoms, box, thickness=distance_upper_bound + 0.1, index=True, **kwargs
         )
         tree = KDTree(atoms_pbc)
         distances, indices = tree.query(
@@ -696,7 +697,7 @@ def number_of_neighbors(
         atoms = atoms % np.diag(box)
         tree = KDTree(atoms, boxsize=np.diag(box))
     else:
-        atoms_pbc = pbc_points(query_atoms, box, thickness=r_max + 0.1, **kwargs)
+        atoms_pbc = pbc_points(atoms, box, thickness=r_max + 0.1, **kwargs)
         tree = KDTree(atoms_pbc)
 
     num_of_neighbors = tree.query_ball_point(query_atoms, r_max, return_length=True)

src/mdevaluate/correlation.py

@@ -2,7 +2,7 @@ from typing import Callable, Optional
 
 import numpy as np
 from numpy.typing import ArrayLike
-from scipy.special import legendre
+from scipy.special import legendre, jn
 import dask.array as darray
 from functools import partial
 from scipy.spatial import KDTree
@@ -183,6 +183,12 @@ def msd(
         displacements = displacements_without_drift(start_frame, end_frame, trajectory)
     if axis == "all":
         return (displacements**2).sum(axis=1).mean()
+    elif axis == "xy" or axis == "yx":
+        return (displacements[:, [0, 1]]**2).sum(axis=1).mean()
+    elif axis == "xz" or axis == "zx":
+        return (displacements[:, [0, 2]]**2).sum(axis=1).mean()
+    elif axis == "yz" or axis == "zy":
+        return (displacements[:, [1, 2]]**2).sum(axis=1).mean()
     elif axis == "x":
         return (displacements[:, 0] ** 2).mean()
     elif axis == "y":
@@ -211,6 +217,15 @@ def isf(
     if axis == "all":
         distance = (displacements**2).sum(axis=1) ** 0.5
         return np.sinc(distance * q / np.pi).mean()
+    elif axis == "xy" or axis == "yx":
+        distance = (displacements[:, [0, 1]]**2).sum(axis=1) ** 0.5
+        return np.real(jn(0, distance * q)).mean()
+    elif axis == "xz" or axis == "zx":
+        distance = (displacements[:, [0, 2]]**2).sum(axis=1) ** 0.5
+        return np.real(jn(0, distance * q)).mean()
+    elif axis == "yz" or axis == "zy":
+        distance = (displacements[:, [1, 2]]**2).sum(axis=1) ** 0.5
+        return np.real(jn(0, distance * q)).mean()
     elif axis == "x":
         distance = np.abs(displacements[:, 0])
         return np.mean(np.cos(np.abs(q * distance)))
@@ -262,6 +277,12 @@ def van_hove_self(
         vectors = displacements_without_drift(start_frame, end_frame, trajectory)
     if axis == "all":
         delta_r = (vectors**2).sum(axis=1) ** 0.5
+    elif axis == "xy" or axis == "yx":
+        delta_r = (vectors[:, [0, 1]]**2).sum(axis=1) ** 0.5
+    elif axis == "xz" or axis == "zx":
+        delta_r = (vectors[:, [0, 2]]**2).sum(axis=1) ** 0.5
+    elif axis == "yz" or axis == "zy":
+        delta_r = (vectors[:, [1, 2]]**2).sum(axis=1) ** 0.5
     elif axis == "x":
         delta_r = np.abs(vectors[:, 0])
     elif axis == "y":
@@ -410,9 +431,9 @@ def non_gaussian_parameter(
     trajectory: Coordinates = None,
     axis: str = "all",
 ) -> float:
-    """
+    r"""
     Calculate the non-Gaussian parameter.
-    ..math:
+    .. math:
       \alpha_2 (t) =
         \frac{3}{5}\frac{\langle r_i^4(t)\rangle}{\langle r_i^2(t)\rangle^2} - 1
     """
@@ -423,6 +444,15 @@ def non_gaussian_parameter(
     if axis == "all":
         r = (vectors**2).sum(axis=1)
         dimensions = 3
+    elif axis == "xy" or axis == "yx":
+        r = (vectors[:, [0, 1]]**2).sum(axis=1)
+        dimensions = 2
+    elif axis == "xz" or axis == "zx":
+        r = (vectors[:, [0, 2]]**2).sum(axis=1)
+        dimensions = 2
+    elif axis == "yz" or axis == "zy":
+        r = (vectors[:, [1, 2]]**2).sum(axis=1)
+        dimensions = 2
     elif axis == "x":
         r = vectors[:, 0] ** 2
         dimensions = 1

src/mdevaluate/distribution.py

@@ -147,9 +147,9 @@ def rdf(
         distances = [d for dist in distances for d in dist]
 
     hist, bins = np.histogram(distances, bins=bins, range=(0, bins[-1]), density=False)
-    hist = hist / len(atoms_a)
+    hist = hist / len(atoms_b)
     hist = hist / (4 / 3 * np.pi * bins[1:] ** 3 - 4 / 3 * np.pi * bins[:-1] ** 3)
-    n = len(atoms_b) / np.prod(np.diag(atoms_b.box))
+    n = len(atoms_a) / np.prod(np.diag(atoms_a.box))
     hist = hist / n
 
     return hist
@@ -306,7 +306,7 @@ def next_neighbor_distribution(
     )[1]
     resname_nn = reference.residue_names[nn]
     count_nn = (resname_nn == atoms.residue_names.reshape(-1, 1)).sum(axis=1)
-    return np.histogram(count_nn, bins=bins, normed=normed)[0]
+    return np.histogram(count_nn, bins=bins, density=normed)[0]
 
 
 def hbonds(

src/mdevaluate/extra/chill.py

@@ -11,7 +11,7 @@ from mdevaluate.coordinates import CoordinateFrame, Coordinates
 from mdevaluate.pbc import pbc_points
 
 
-def a_ij(atoms: ArrayLike, N: int = 4, l: int = 3) -> tuple[NDArray, NDArray]:
+def calc_aij(atoms: ArrayLike, N: int = 4, l: int = 3) -> tuple[NDArray, NDArray]:
     tree = KDTree(atoms)
 
     dist, indices = tree.query(atoms, N + 1)
@@ -84,18 +84,18 @@ def count_ice_types(iceTypes: NDArray) -> NDArray:
 
 
 def selector_ice(
-    start_frame: CoordinateFrame,
+    oxygen_atoms_water: CoordinateFrame,
-    traj: Coordinates,
     chosen_ice_types: ArrayLike,
     combined: bool = True,
+    next_neighbor_distance: float = 0.35,
 ) -> NDArray:
-    oxygen = traj.subset(atom_name="OW")
+    atoms = oxygen_atoms_water
-    atoms = oxygen[start_frame.step]
+    atoms_PBC = pbc_points(atoms, thickness=next_neighbor_distance * 2.2)
-    atoms = atoms % np.diag(atoms.box)
+    aij, indices = calc_aij(atoms_PBC)
-    atoms_PBC = pbc_points(atoms, thickness=1)
-    aij, indices = a_ij(atoms_PBC)
     tree = KDTree(atoms_PBC)
-    neighbors = tree.query_ball_point(atoms_PBC, 0.35, return_length=True)
+    neighbors = tree.query_ball_point(
+        atoms_PBC, next_neighbor_distance, return_length=True
+    ) - 1
     index_SOL = atoms_PBC.tolist().index(atoms[0].tolist())
     index_SOL = np.arange(index_SOL, index_SOL + len(atoms))
     ice_Types = classify_ice(aij, indices, neighbors, index_SOL)
@@ -117,9 +117,9 @@ def selector_ice(
 def ice_types(trajectory: Coordinates, segments: int = 10000) -> pd.DataFrame:
     def ice_types_distribution(frame: CoordinateFrame, selector: Callable) -> NDArray:
         atoms_PBC = pbc_points(frame, thickness=1)
-        aij, indices = a_ij(atoms_PBC)
+        aij, indices = calc_aij(atoms_PBC)
         tree = KDTree(atoms_PBC)
-        neighbors = tree.query_ball_point(atoms_PBC, 0.35, return_length=True)
+        neighbors = tree.query_ball_point(atoms_PBC, 0.35, return_length=True) - 1
         index = selector(frame, atoms_PBC)
         ice_types_data = classify_ice(aij, indices, neighbors, index)
         ice_parts_data = count_ice_types(ice_types_data)
@@ -161,8 +161,8 @@ def ice_types(trajectory: Coordinates, segments: int = 10000) -> pd.DataFrame:
 def ice_clusters_traj(
     traj: Coordinates, segments: int = 10000, skip: float = 0.1
 ) -> list:
-    def ice_clusters(frame: CoordinateFrame, traj: Coordinates) -> Tuple[float, list]:
+    def ice_clusters(frame: CoordinateFrame) -> Tuple[float, list]:
-        selection = selector_ice(frame, traj, [0, 1, 2])
+        selection = selector_ice(frame, [0, 1, 2])
         if len(selection) == 0:
             return frame.time, []
         else:
@@ -194,6 +194,6 @@ def ice_clusters_traj(
         np.int_(np.linspace(len(traj) * skip, len(traj) - 1, num=segments))
     )
     all_clusters = [
-        ice_clusters(traj[frame_index], traj) for frame_index in frame_indices
+        ice_clusters(traj[frame_index]) for frame_index in frame_indices
     ]
     return all_clusters

src/mdevaluate/extra/free_energy_landscape.py

@@ -1,9 +1,9 @@
 from functools import partial
+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
@@ -11,6 +11,56 @@ import multiprocessing as mp
 from ..coordinates import Coordinates
 
 
+def _pbc_points_reduced(
+    coordinates: ArrayLike,
+    pore_geometry: str,
+    box: Optional[NDArray] = None,
+    thickness: Optional[float] = None,
+) -> tuple[NDArray, NDArray]:
+    if box is None:
+        box = coordinates.box
+    if pore_geometry == "cylindrical":
+        grid = np.array([[i, j, k] for k in [-1, 0, 1] for j in [0] for i in [0]])
+        indices = np.tile(np.arange(len(coordinates)), 3)
+    elif pore_geometry == "slit":
+        grid = np.array(
+            [[i, j, k] for k in [0] for j in [1, 0, -1] for i in [-1, 0, 1]]
+        )
+        indices = np.tile(np.arange(len(coordinates)), 9)
+    else:
+        raise ValueError(
+            f"pore_geometry is {pore_geometry}, should either be "
+            f"'cylindrical' or 'slit'"
+        )
+    coordinates_pbc = np.concatenate([coordinates + v @ box for v in grid], axis=0)
+    size = np.diag(box)
+
+    if thickness is not None:
+        mask = np.all(coordinates_pbc > -thickness, axis=1)
+        coordinates_pbc = coordinates_pbc[mask]
+        indices = indices[mask]
+        mask = np.all(coordinates_pbc < size + thickness, axis=1)
+        coordinates_pbc = coordinates_pbc[mask]
+        indices = indices[mask]
+
+    return coordinates_pbc, indices
+
+
+def _build_tree(points, box, r_max, pore_geometry):
+    if np.all(np.diag(np.diag(box)) == box):
+        tree = KDTree(points % box, boxsize=box)
+        points_pbc_index = None
+    else:
+        points_pbc, points_pbc_index = _pbc_points_reduced(
+            points,
+            pore_geometry,
+            box,
+            thickness=r_max + 0.01,
+        )
+        tree = KDTree(points_pbc)
+    return tree, points_pbc_index
+
+
 def occupation_matrix(
     trajectory: Coordinates,
     edge_length: float = 0.05,
@@ -28,11 +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 = [
-        np.arange(len(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
@@ -72,19 +118,19 @@ def _calc_histogram(
 
 
 def find_maxima(
-    occupation_df: pd.DataFrame, box: ArrayLike, edge_length: float = 0.05
+    occupation_df: pd.DataFrame, box: ArrayLike, radius: float, pore_geometry: str
 ) -> pd.DataFrame:
     maxima_df = occupation_df.copy()
     maxima_df["maxima"] = None
     points = np.array(maxima_df[["x", "y", "z"]])
-    tree = KDTree(points, boxsize=box)
+    tree, points_pbc_index = _build_tree(points, box, radius, pore_geometry)
-    all_neighbors = tree.query_ball_point(
-        points, edge_length * 3 ** (1 / 2) + edge_length / 100
-    )
     for i in range(len(maxima_df)):
         if maxima_df.loc[i, "maxima"] is not None:
             continue
-        neighbors = np.array(all_neighbors[i])
+        maxima_pos = maxima_df.loc[i, ["x", "y", "z"]]
+        neighbors = np.array(tree.query_ball_point(maxima_pos, radius))
+        if points_pbc_index is not None:
+            neighbors = points_pbc_index[neighbors]
         neighbors = neighbors[neighbors != i]
         if len(neighbors) == 0:
             maxima_df.loc[i, "maxima"] = True
@@ -104,23 +150,39 @@ def _calc_energies(
     maxima_df: pd.DataFrame,
     bins: ArrayLike,
     box: NDArray,
+    pore_geometry: str,
     T: float,
+    nodes: int = 8,
 ) -> NDArray:
     points = np.array(maxima_df[["x", "y", "z"]])
-    tree = KDTree(points, boxsize=box)
+    tree, points_pbc_index = _build_tree(points, box, bins[-1], pore_geometry)
     maxima = maxima_df.loc[maxima_indices, ["x", "y", "z"]]
     maxima_occupations = np.array(maxima_df.loc[maxima_indices, "occupation"])
     num_of_neighbors = np.max(
         tree.query_ball_point(maxima, bins[-1], return_length=True)
     )
-    distances, indices = tree.query(
+    split_maxima = []
-        maxima, k=num_of_neighbors, distance_upper_bound=bins[-1]
+    for i in range(0, len(maxima), 1000):
-    )
+        split_maxima.append(maxima[i : i + 1000])
+
+    distances = []
+    indices = []
+    for maxima in split_maxima:
+        distances_step, indices_step = tree.query(
+            maxima, k=num_of_neighbors, distance_upper_bound=bins[-1], workers=nodes
+        )
+        distances.append(distances_step)
+        indices.append(indices_step)
+    distances = np.concatenate(distances)
+    indices = np.concatenate(indices)
     all_energy_hist = []
     all_occupied_bins_hist = []
     if distances.ndim == 1:
         current_distances = distances[1:][distances[1:] <= bins[-1]]
-        current_indices = indices[1:][distances[1:] <= bins[-1]]
+        if points_pbc_index is None:
+            current_indices = indices[1:][distances[1:] <= bins[-1]]
+        else:
+            current_indices = points_pbc_index[indices[1:][distances[1:] <= bins[-1]]]
         energy = (
             -np.log(maxima_df.loc[current_indices, "occupation"] / maxima_occupations)
             * T
@@ -131,8 +193,12 @@ def _calc_energies(
         return result
     for i, maxima_occupation in enumerate(maxima_occupations):
         current_distances = distances[i, 1:][distances[i, 1:] <= bins[-1]]
-        current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]]
+        if points_pbc_index is None:
-
+            current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]]
+        else:
+            current_indices = points_pbc_index[
+                indices[i, 1:][distances[i, 1:] <= bins[-1]]
+            ]
         energy = (
             -np.log(maxima_df.loc[current_indices, "occupation"] / maxima_occupation)
             * T
@@ -168,9 +234,12 @@ def distance_resolved_energies(
     distance_bins: ArrayLike,
     r_bins: ArrayLike,
     box: NDArray,
+    pore_geometry: str,
     temperature: float,
+    nodes: int = 8,
 ) -> pd.DataFrame:
     results = []
+    distances = []
     for i in range(len(distance_bins) - 1):
         maxima_indices = np.array(
             maxima_df.index[
@@ -179,11 +248,22 @@ def distance_resolved_energies(
                 * (maxima_df["maxima"] == True)
             ]
         )
-        results.append(
+        try:
-            _calc_energies(maxima_indices, maxima_df, r_bins, box, temperature)
+            results.append(
-        )
+                _calc_energies(
+                    maxima_indices,
+                    maxima_df,
+                    r_bins,
+                    box,
+                    pore_geometry,
+                    temperature,
+                    nodes,
+                )
+            )
+            distances.append((distance_bins[i] + distance_bins[i + 1]) / 2)
+        except ValueError:
+            pass
 
-    distances = (distance_bins[:-1] + distance_bins[1:]) / 2
     radii = (r_bins[:-1] + r_bins[1:]) / 2
     d = np.array([d for d in distances for r in radii])
     r = np.array([r for d in distances for r in radii])
@@ -192,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 = []
@@ -201,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/extra/water.py

@@ -265,7 +265,7 @@ def LSI(
     )
     distributions = np.array(
         [
-            LSI_distribution(trajectory[frame_index], trajectory, bins, selector=None)
+            LSI_distribution(trajectory[frame_index], bins, selector=None)
             for frame_index in frame_indices
         ]
     )

src/mdevaluate/pbc.py

@@ -7,7 +7,7 @@ from numpy.typing import ArrayLike, NDArray
 
 from itertools import product
 
-from .logging import logger
+from .logging_util import logger
 
 if TYPE_CHECKING:
     from mdevaluate.coordinates import CoordinateFrame
@@ -156,7 +156,7 @@ def nojump(frame: CoordinateFrame, usecache: bool = True) -> CoordinateFrame:
                     [m[i0 : abstep + 1].sum(axis=0) for m in reader.nojump_matrices]
                 ).T
             )
-            * frame.box.diagonal()
+            @ frame.box
         )
 
         reader._nojump_cache[abstep] = delta
@@ -173,7 +173,7 @@ def nojump(frame: CoordinateFrame, usecache: bool = True) -> CoordinateFrame:
                     ]
                 ).T
             )
-            * frame.box.diagonal()
+            @ frame.box
         )
     return frame - delta
 

src/mdevaluate/reader.py

@@ -19,13 +19,13 @@ import MDAnalysis
 from scipy import sparse
 
 from .checksum import checksum
-from .logging import logger
+from .logging_util import logger
 from . import atoms
 from .coordinates import Coordinates
 
 CSR_ATTRS = ("data", "indices", "indptr")
 NOJUMP_MAGIC = 2016
-Group_RE = re.compile("\[ ([-+\w]+) \]")
+Group_RE = re.compile(r"\[ ([-+\w]+) \]")
 
 
 class NojumpError(Exception):
@@ -152,7 +152,7 @@ def nojump_load_filename(reader: BaseReader):
         )
         if os.path.exists(full_path_fallback):
             return full_path_fallback
-    if os.path.exists(fname) or is_writeable(directory):
+    if os.path.exists(full_path) or is_writeable(directory):
         return full_path
     else:
         user_data_dir = os.path.join("/data/", os.environ["HOME"].split("/")[-1])
@@ -187,19 +187,33 @@ def nojump_save_filename(reader: BaseReader):
 
 def parse_jumps(trajectory: Coordinates):
     prev = trajectory[0].whole
-    box = prev.box.diagonal()
+    box = prev.box
     SparseData = namedtuple("SparseData", ["data", "row", "col"])
     jump_data = (
         SparseData(data=array("b"), row=array("l"), col=array("l")),
         SparseData(data=array("b"), row=array("l"), col=array("l")),
         SparseData(data=array("b"), row=array("l"), col=array("l")),
     )
-
     for i, curr in enumerate(trajectory):
         if i % 500 == 0:
             logger.debug("Parse jumps Step: %d", i)
-        delta = ((curr - prev) / box).round().astype(np.int8)
+        r3 = np.subtract(curr, prev)
+        delta_z = np.array(np.rint(np.divide(r3[:, 2], box[2][2])), dtype=np.int8)
+        r2 = np.subtract(
+            r3,
+            (np.rint(np.divide(r3[:, 2], box[2][2])))[:, np.newaxis]
+            * box[2][np.newaxis, :],
+            )
+        delta_y = np.array(np.rint(np.divide(r2[:, 1], box[1][1])), dtype=np.int8)
+        r1 = np.subtract(
+            r2,
+            (np.rint(np.divide(r2[:, 1], box[1][1])))[:, np.newaxis]
+            * box[1][np.newaxis, :],
+            )
+        delta_x = np.array(np.rint(np.divide(r1[:, 0], box[0][0])), dtype=np.int8)
+        delta = np.array([delta_x, delta_y, delta_z]).T
         prev = curr
+        box = prev.box
         for d in range(3):
             (col,) = np.where(delta[:, d] != 0)
             jump_data[d].col.extend(col)

src/mdevaluate/utils.py

@@ -14,7 +14,7 @@ from scipy.ndimage import uniform_filter1d
 from scipy.interpolate import interp1d
 from scipy.optimize import curve_fit
 
-from .logging import logger
+from .logging_util import logger
 from .functions import kww, kww_1e
 
 

test/test_checksum.py

@@ -5,11 +5,11 @@ import numpy as np
 
 def test_checksum():
     salt = checksum.SALT
-    checksum.SALT = ''
+    checksum.SALT = ""
     assert checksum.checksum(1) == 304942582444936629325699363757435820077590259883
-    assert checksum.checksum('42') == checksum.checksum(42)
+    assert checksum.checksum("42") == checksum.checksum(42)
     cs1 = checksum.checksum(999)
-    checksum.SALT = '999'
+    checksum.SALT = "999"
     assert cs1 != checksum.checksum(999)
 
     a = np.array([1, 2, 3])
@@ -19,7 +19,6 @@ def test_checksum():
 
 
 def test_version():
-
     @checksum.version(1)
     def f1():
         pass

test/test_coordinates.py

@@ -7,7 +7,7 @@ from mdevaluate import coordinates
 
 @pytest.fixture
 def trajectory(request):
-    return mdevaluate.open(os.path.join(os.path.dirname(__file__), 'data/water'))
+    return mdevaluate.open(os.path.join(os.path.dirname(__file__), "data/water"))
 
 
 def test_coordinates_getitem(trajectory):

test/test_free_energy_landscape.py

@@ -13,42 +13,24 @@ def trajectory(request):
 
 
 def test_get_fel(trajectory):
-    test_array = np.array(
+    test_array = np.array([210., 214., 209., 192., 200., 193., 230., 218., 266.])
-        [
-            174.46253634,
-            174.60905476,
-            178.57658092,
-            182.43001192,
-            180.57916378,
-            176.49886217,
-            178.96018547,
-            181.13561782,
-            178.31026314,
-            176.08903996,
-            180.71215345,
-            181.59703135,
-            180.34329368,
-            187.02474488,
-            197.99167477,
-            214.05788031,
-            245.58571282,
-            287.52457507,
-            331.53492965,
-        ]
-    )
 
     OW = trajectory.subset(atom_name="OW")
-
+    box = trajectory[0].box
-    box = np.diag(trajectory[0].box)
+    box_voxels = (np.diag(box) // [0.05, 0.05, 0.05] + [1, 1, 1]) * [0.05, 0.05, 0.05]
-    box_voxels = (box // [0.05, 0.05, 0.05] + [1, 1, 1]) * [0.05, 0.05, 0.05]
+    occupation_matrix = fel.occupation_matrix(OW, skip=0, segments=10)
-    occupation_matrix = fel.occupation_matrix(OW, skip=0, segments=1000)
+    radius_maxima = 0.05 * 3 ** (1 / 2) + 0.05 / 100
-    maxima_matrix = fel.find_maxima(occupation_matrix, box=box_voxels, edge_length=0.05)
+    maxima_matrix = fel.find_maxima(
-    maxima_matrix = fel.add_distances(maxima_matrix, "cylindrical", box / 2)
+        occupation_matrix,
-    r_bins = np.arange(0, 2, 0.02)
+        box=box_voxels,
-    distance_bins = np.arange(0.05, 2.05, 0.1)
+        radius=radius_maxima,
+        pore_geometry="cylindrical",
+    )
+    maxima_matrix = fel.add_distances(maxima_matrix, "cylindrical", np.diag(box) / 2)
+    r_bins = np.arange(0, 0.5, 0.02)
+    distance_bins = np.arange(1.8, 1.9, 0.01)
     energy_df = fel.distance_resolved_energies(
-        maxima_matrix, distance_bins, r_bins, box, 225
+        maxima_matrix, distance_bins, r_bins, box, "cylindrical", 225
     )
     result = fel.find_energy_maxima(energy_df, r_min=0.05, r_max=0.15)
-
     assert (np.round(np.array(result["energy"])) == np.round(test_array)).all()

test/test_pbc.py

@@ -9,6 +9,6 @@ def test_pbc_diff():
     y = np.random.rand(10, 3)
     box = np.ones((3,))
 
-    assert (pbc.pbc_diff(x, x, box) == approx(0))
+    assert pbc.pbc_diff(x, x, box) == approx(0)
-    dxy = (pbc.pbc_diff(x, y, box)**2).sum(axis=1)**0.5
+    dxy = (pbc.pbc_diff(x, y, box) ** 2).sum(axis=1) ** 0.5
     assert (dxy <= 0.75**0.5).all()

test/test_utils.py

@@ -8,7 +8,7 @@ from mdevaluate import utils
 @pytest.fixture
 def logdata(request):
     xdata = np.logspace(-1, 3, 50)
-    ydata = np.exp(- (xdata)**0.7)
+    ydata = np.exp(-((xdata) ** 0.7))
     return xdata, ydata
 
 
@@ -18,16 +18,16 @@ def test_filon_fourier_transformation(logdata):
     xdata_zero = copy(xdata)
     xdata_zero[0] = 0
     _, filon = utils.filon_fourier_transformation(xdata_zero, ydata)
-    assert not np.isnan(filon).any(), 'There are NaN values in the filon result!'
+    assert not np.isnan(filon).any(), "There are NaN values in the filon result!"
 
     freqs = np.logspace(-4, 1)
     filon_freqs, filon_imag = utils.filon_fourier_transformation(
-        xdata, xdata, frequencies=freqs, derivative='linear', imag=True
+        xdata, xdata, frequencies=freqs, derivative="linear", imag=True
-        )
+    )
 
     assert (freqs == filon_freqs).all()
 
     freqs, filon_real = utils.filon_fourier_transformation(
-        xdata, xdata, frequencies=freqs, derivative='linear', imag=False
+        xdata, xdata, frequencies=freqs, derivative="linear", imag=False
-        )
+    )
     assert np.isclose(filon_imag.real, filon_real).all()