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

Reviewed-on: #5

src/mdevaluate/__init__.py

@@ -16,7 +16,7 @@ from . import reader
 from . import system
 from . import utils
 from . import extra
-from .logging_util import logger
+from .logging 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_util import logger
+from .logging import logger
 
 autosave_directory: Optional[str] = None
 load_autosave_data = False

src/mdevaluate/checksum.py

@@ -1,14 +1,9 @@
 import functools
 import hashlib
-from .logging_util import logger
+from .logging 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
 
@@ -33,46 +28,19 @@ def version(version_nr: int, calls: Iterable = ()):
     return decorator
 
 
-def strip_comments(source: str) -> str:
+def strip_comments(s: str):
-    """Removes docstrings, comments, and irrelevant whitespace from Python source code."""
+    """Strips comment lines and docstring from Python source string."""
-
+    o = ""
-    # Step 1: Remove docstrings using AST
+    in_docstring = False
-    def remove_docstrings(node):
+    for l in s.split("\n"):
-        if isinstance(node, (ast.FunctionDef, ast.AsyncFunctionDef, ast.ClassDef, ast.Module)):
+        if l.strip().startswith(("#", '"', "'")) or in_docstring:
-            if (doc := ast.get_docstring(node, clean=False)):
+            in_docstring = l.strip().startswith(('"""', "'''")) + in_docstring == 1
-                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
-        if srow > last_lineno:
+        o += l + "\n"
-            last_col = 0
+    return o
-        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, _seen=None):
+def checksum(*args, csum=None):
     """
     Calculate a checksum of any object, by sha1 hash.
 
@@ -92,15 +60,7 @@ def checksum(*args, csum=None, _seen=None):
         csum = hashlib.sha1()
         csum.update(str(SALT).encode())
 
-    if _seen is None:
-        _seen = set()
-
     for arg in args:
-        obj_id = id(arg)
-        if obj_id in _seen:
-            continue
-        _seen.add(obj_id)
-
         if hasattr(arg, "__checksum__"):
             logger.debug("Checksum via __checksum__: %s", str(arg))
             csum.update(str(arg.__checksum__()).encode())
@@ -117,15 +77,15 @@ def checksum(*args, csum=None, _seen=None):
             for key in sorted(merged):  # deterministic ordering
                 v = merged[key]
                 if v is not arg:
-                    checksum(v, csum=csum, _seen=_seen)
+                    checksum(v, csum=csum)
         elif isinstance(arg, functools.partial):
             logger.debug("Checksum via partial for %s", str(arg))
-            checksum(arg.func, csum=csum, _seen=_seen)
+            checksum(arg.func, csum=csum)
             for x in arg.args:
-                checksum(x, csum=csum, _seen=_seen)
+                checksum(x, csum=csum)
             for k in sorted(arg.keywords.keys()):
                 csum.update(k.encode())
-                checksum(arg.keywords[k], csum=csum, _seen=_seen)
+                checksum(arg.keywords[k], csum=csum)
         elif isinstance(arg, np.ndarray):
             csum.update(arg.tobytes())
         else:

src/mdevaluate/coordinates.py

@@ -1,6 +1,6 @@
 from functools import partial, wraps
 from copy import copy
-from .logging_util import logger
+from .logging 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/pbc.py

@@ -7,7 +7,7 @@ from numpy.typing import ArrayLike, NDArray
 
 from itertools import product
 
-from .logging_util import logger
+from .logging import logger
 
 if TYPE_CHECKING:
     from mdevaluate.coordinates import CoordinateFrame
@@ -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
 
 

src/mdevaluate/reader.py

@@ -19,13 +19,13 @@ import MDAnalysis
 from scipy import sparse
 
 from .checksum import checksum
-from .logging_util import logger
+from .logging import logger
 from . import atoms
 from .coordinates import Coordinates
 
 CSR_ATTRS = ("data", "indices", "indptr")
 NOJUMP_MAGIC = 2016
-Group_RE = re.compile(r"\[ ([-+\w]+) \]")
+Group_RE = re.compile("\[ ([-+\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_util import logger
+from .logging import logger
 from .functions import kww, kww_1e