From 8941b11a3df626d4327c68505dc3ed6ebcbaab9c Mon Sep 17 00:00:00 2001 From: svdenhau Date: Thu, 2 Nov 2023 22:53:39 +0100 Subject: [PATCH] avoid unnecessary yaff imports --- psiflow/walkers/dynamic.py | 22 +++++++++++- psiflow/walkers/random.py | 52 +++++++++++++++++++++++++++- psiflow/walkers/utils.py | 69 -------------------------------------- tests/test_walkers.py | 4 +-- 4 files changed, 74 insertions(+), 73 deletions(-) diff --git a/psiflow/walkers/dynamic.py b/psiflow/walkers/dynamic.py index 5ce52f0..3eb6604 100644 --- a/psiflow/walkers/dynamic.py +++ b/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() diff --git a/psiflow/walkers/random.py b/psiflow/walkers/random.py index bef8e45..d32dae6 100644 --- a/psiflow/walkers/random.py +++ b/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"]) diff --git a/psiflow/walkers/utils.py b/psiflow/walkers/utils.py index b2b50aa..a5f0560 100644 --- a/psiflow/walkers/utils.py +++ b/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 diff --git a/tests/test_walkers.py b/tests/test_walkers.py index eb7c6a0..deb2c31 100644 --- a/tests/test_walkers.py +++ b/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):