avoid unnecessary yaff imports

psiflow/walkers/dynamic.py

@@ -84,6 +84,26 @@ def molecular_dynamics_yaff(
     return " ".join(command_list)
 
 
+def parse_yaff_output(stdout):
+    temperatures = []
+    counter = 0
+    time = 0
+    for line in stdout.split("\n"):
+        if "VERLET" in line:
+            try:
+                _ = [float(s) for s in line.split()[1:]]
+            except ValueError:
+                continue
+            temperatures.append(float(line.split()[3]))
+            counter = int(line.split()[1])
+            time = float(line.split()[6])
+        else:
+            pass
+    if len(temperatures) == 0:
+        temperatures.append(-1)
+    return counter, np.mean(np.array(temperatures)), time
+
+
 @python_app(executors=["default_threads"])
 def molecular_dynamics_yaff_post(
     inputs: list[File] = [],
@@ -92,7 +112,7 @@ def molecular_dynamics_yaff_post(
     from ase.io import read
 
     from psiflow.data import FlowAtoms, NullState
-    from psiflow.walkers.utils import parse_yaff_output
+    from psiflow.walkers.dynamic import parse_yaff_output
 
     with open(inputs[1], "r") as f:
         stdout = f.read()

psiflow/walkers/random.py

@@ -3,6 +3,7 @@
 from collections import namedtuple
 from typing import Any, Union
 
+import numpy as np
 import typeguard
 from ase import Atoms
 from parsl.app.app import python_app
@@ -15,6 +16,55 @@
 Metadata = namedtuple("Metadata", ["state", "counter", "reset"])
 
 
+def apply_strain(strain, box0):
+    """Applies a strain tensor to a reference box
+
+    The resulting strained box matrix is obtained based on:
+
+        box = box0 @ sqrt(2 * strain + I)
+
+    where the second argument is computed based on a diagonalization of
+    2 * strain + I.
+
+    Parameters
+    ----------
+
+    strain : ndarray of shape (3, 3)
+        desired strain matrix
+
+    box0 : ndarray of shape (3, 3)
+        reference box matrix
+
+    """
+    assert np.allclose(strain, strain.T)
+    A = 2 * strain + np.eye(3)
+    values, vectors = np.linalg.eigh(A)
+    sqrtA = vectors @ np.sqrt(np.diag(values)) @ vectors.T
+    box = box0 @ sqrtA
+    return box
+
+
+def compute_strain(box, box0):
+    """Computes the strain of a given box with respect to a reference
+
+    The strain matrix is defined by the following expression
+
+        strain = 0.5 * (inv(box0) @ box @ box.T @ inv(box0).T - I)
+
+    Parameters
+    ----------
+
+    box : ndarray of shape (3, 3)
+        box matrix for which to compute the strain
+
+    box0 : ndarray of shape (3, 3)
+        reference box matrix
+
+    """
+    box0inv = np.linalg.inv(box0)
+    return 0.5 * (box0inv @ box @ box.T @ box0inv.T - np.eye(3))
+
+
 @typeguard.typechecked
 def random_perturbation(
     state: FlowAtoms,
@@ -24,7 +74,7 @@ def random_perturbation(
 
     import numpy as np
 
-    from psiflow.walkers.utils import apply_strain
+    from psiflow.walkers.random import apply_strain
 
     state = copy.deepcopy(state)
     np.random.seed(parameters["seed"])

psiflow/walkers/utils.py

@@ -198,75 +198,6 @@ def __call__(self, iterative):
             self.temperatures.append(iterative.temp)
 
 
-def apply_strain(strain, box0):
-    """Applies a strain tensor to a reference box
-
-    The resulting strained box matrix is obtained based on:
-
-        box = box0 @ sqrt(2 * strain + I)
-
-    where the second argument is computed based on a diagonalization of
-    2 * strain + I.
-
-    Parameters
-    ----------
-
-    strain : ndarray of shape (3, 3)
-        desired strain matrix
-
-    box0 : ndarray of shape (3, 3)
-        reference box matrix
-
-    """
-    assert np.allclose(strain, strain.T)
-    A = 2 * strain + np.eye(3)
-    values, vectors = np.linalg.eigh(A)
-    sqrtA = vectors @ np.sqrt(np.diag(values)) @ vectors.T
-    box = box0 @ sqrtA
-    return box
-
-
-def compute_strain(box, box0):
-    """Computes the strain of a given box with respect to a reference
-
-    The strain matrix is defined by the following expression
-
-        strain = 0.5 * (inv(box0) @ box @ box.T @ inv(box0).T - I)
-
-    Parameters
-    ----------
-
-    box : ndarray of shape (3, 3)
-        box matrix for which to compute the strain
-
-    box0 : ndarray of shape (3, 3)
-        reference box matrix
-
-    """
-    box0inv = np.linalg.inv(box0)
-    return 0.5 * (box0inv @ box @ box.T @ box0inv.T - np.eye(3))
-
-
-def parse_yaff_output(stdout):
-    temperatures = []
-    counter = 0
-    time = 0
-    for line in stdout.split("\n"):
-        if "VERLET" in line:
-            try:
-                _ = [float(s) for s in line.split()[1:]]
-            except ValueError:
-                continue
-            temperatures.append(float(line.split()[3]))
-            counter = int(line.split()[1])
-            time = float(line.split()[6])
-        else:
-            pass
-    if len(temperatures) == 0:
-        temperatures.append(-1)
-    return counter, np.mean(np.array(temperatures)), time
-
-
 def max_temperature(temperature: float, natoms: int, quantile: float) -> float:
     ndof = 3 * natoms
     return chi2.ppf(1 - quantile, ndof) * temperature / ndof

tests/test_walkers.py

@@ -19,8 +19,8 @@
     RandomWalker,
     load_walker,
 )
-from psiflow.walkers.dynamic import parse_openmm_output
-from psiflow.walkers.utils import get_velocities_at_temperature, parse_yaff_output
+from psiflow.walkers.dynamic import parse_openmm_output, parse_yaff_output
+from psiflow.walkers.utils import get_velocities_at_temperature
 
 
 def test_random_walker_multiply(dataset, tmp_path):