dedent before ast.parse for non-toplevel functions

Reviewed-on: #2

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):
-    """Strips comment lines and docstring from Python source string."""
-    o = ""
-    in_docstring = False
-    for l in s.split("\n"):
-        if l.strip().startswith(("#", '"', "'")) or in_docstring:
-            in_docstring = l.strip().startswith(('"""', "'''")) + in_docstring == 1
+def strip_comments(source: str) -> str:
+    """Removes docstrings, comments, and irrelevant whitespace from Python source code."""
+
+    # Step 1: Remove docstrings using AST
+    def remove_docstrings(node):
+        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Module)):
+            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"
-    return o
+        if srow > last_lineno:
+            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

src/mdevaluate/correlation.py

@@ -431,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
     """

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
@@ -47,6 +46,21 @@ def _pbc_points_reduced(
     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,
@@ -64,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.arange(len(frame_indices))[i : i + size]
-        for i in range(0, len(frame_indices), size)
-    ]
+    indices = np.array_split(frame_indices, nodes)
     pool = mp.Pool(nodes)
     results = pool.map(
         partial(_calc_histogram, trajectory=trajectory, bins=bins), indices
@@ -113,23 +123,14 @@ def find_maxima(
     maxima_df = occupation_df.copy()
     maxima_df["maxima"] = None
     points = np.array(maxima_df[["x", "y", "z"]])
-    if np.all(np.diag(np.diag(box)) == box):
-        tree = KDTree(points, boxsize=box)
-        all_neighbors = tree.query_ball_point(points, radius)
-    else:
-        points_pbc, points_pbc_index = _pbc_points_reduced(
-            points,
-            pore_geometry,
-            box,
-            thickness=radius + 0.01,
-        )
-        tree = KDTree(points_pbc)
-        all_neighbors = tree.query_ball_point(points, radius)
-        all_neighbors = points_pbc_index[all_neighbors]
+    tree, points_pbc_index = _build_tree(points, box, radius, pore_geometry)
     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
@@ -154,16 +155,7 @@ def _calc_energies(
     nodes: int = 8,
 ) -> NDArray:
     points = np.array(maxima_df[["x", "y", "z"]])
-    if np.all(np.diag(np.diag(box)) == box):
-        tree = KDTree(points, boxsize=box)
-    else:
-        points_pbc, points_pbc_index = _pbc_points_reduced(
-            points,
-            pore_geometry,
-            box,
-            thickness=bins[-1] + 0.01,
-        )
-        tree = KDTree(points_pbc)
+    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(
@@ -187,7 +179,7 @@ def _calc_energies(
     all_occupied_bins_hist = []
     if distances.ndim == 1:
         current_distances = distances[1:][distances[1:] <= bins[-1]]
-        if np.all(np.diag(np.diag(box)) == box):
+        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]]]
@@ -201,7 +193,7 @@ def _calc_energies(
         return result
     for i, maxima_occupation in enumerate(maxima_occupations):
         current_distances = distances[i, 1:][distances[i, 1:] <= bins[-1]]
-        if np.all(np.diag(np.diag(box)) == box):
+        if points_pbc_index is None:
             current_indices = indices[i, 1:][distances[i, 1:] <= bins[-1]]
         else:
             current_indices = points_pbc_index[
@@ -280,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 = []
@@ -289,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)
-        energies.append(np.max(p3(np.linspace(r_min, r_max, 1000))))
+        p3 = Poly.fit(x[mask], y[mask], deg=degree)
+        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

@@ -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

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):

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_free_energy_landscape.py

@@ -13,19 +13,7 @@ def trajectory(request):
 
 
 def test_get_fel(trajectory):
-    test_array = np.array(
-        [
-            184.4909136,
-            233.79320471,
-            223.12003988,
-            228.49746397,
-            200.5626769,
-            212.82484221,
-            165.10818396,
-            170.74123681,
-            175.86672931,
-        ]
-    )
+    test_array = np.array([210., 214., 209., 192., 200., 193., 230., 218., 266.])
 
     OW = trajectory.subset(atom_name="OW")
     box = trajectory[0].box