optimize order parameter evaluation, allow hamiltonians to be evaluated on single states

Author: svandenhaute

psiflow/data.py

@@ -812,7 +812,7 @@ def _batch_frames(
 @join_app
 @typeguard.typechecked
 def batch_apply(
-    func: Callable,
+    funcs: list[Callable],
     batch_size: int,
     length: int,
     inputs: list = [],
@@ -821,6 +821,9 @@ def batch_apply(
     nbatches = math.ceil(length / batch_size)
     batches = [psiflow.context().new_file("data_", ".xyz") for _ in range(nbatches)]
     future = batch_frames(batch_size, inputs=[inputs[0]], outputs=batches)
-    evaluated = [func(Dataset(None, extxyz=e)) for e in future.outputs]
-    f = join_frames(inputs=[e.extxyz for e in evaluated], outputs=[outputs[0]])
+    datasets = [Dataset(None, extxyz=e) for e in future.outputs]
+    for func in funcs:
+        datasets = [func(d) for d in datasets]
+    # evaluated = [func(Dataset(None, extxyz=e)) for e in future.outputs]
+    f = join_frames(inputs=[d.extxyz for d in datasets], outputs=[outputs[0]])
     return f

psiflow/hamiltonians/_plumed.py

@@ -42,13 +42,14 @@ def try_manual_plumed_linking() -> str:
 def remove_comments_printflush(plumed_input: str) -> str:
     new_input = []
     for line in list(plumed_input.split("\n")):
-        if line.strip().startswith("#"):
+        pre_comment = line.strip().split("#")[0].strip()
+        if len(pre_comment) == 0:
             continue
-        if line.strip().startswith("PRINT"):
+        if pre_comment.startswith("PRINT"):
             continue
-        if line.strip().startswith("FLUSH"):
+        if pre_comment.startswith("FLUSH"):
             continue
-        new_input.append(line)
+        new_input.append(pre_comment)
     return "\n".join(new_input)
 
 

psiflow/hamiltonians/hamiltonian.py

@@ -1,11 +1,12 @@
 from __future__ import annotations  # necessary for type-guarding class methods
 
 import logging
-from typing import Callable, Optional
+from typing import Callable, Optional, Union
 
 import typeguard
 from parsl.app.app import python_app
 from parsl.app.futures import DataFuture
+from parsl.dataflow.futures import AppFuture
 from parsl.data_provider.files import File
 
 import psiflow
@@ -22,16 +23,20 @@ def evaluate_function(
     outputs: list = [],
     parsl_resource_specification: dict = {},
     **parameters,  # dict values can be futures, so app must wait for those
-) -> None:
+) -> Optional[Geometry]:
     import numpy as np
     from ase import Atoms
 
     from psiflow.data import _read_frames, _write_frames
     from psiflow.geometry import NullState
 
     assert len(inputs) >= 1
-    assert len(outputs) == 1
-    states = _read_frames(inputs=[inputs[0]])
+    if isinstance(inputs[0], Geometry):
+        assert len(outputs) == 0
+        states = [inputs[0]]
+    else:
+        assert len(outputs) == 1
+        states = _read_frames(inputs=[inputs[0]])
     calculators, index_mapping = load_calculators(states, inputs[1], **parameters)
     for i, state in enumerate(states):
         if state == NullState:
@@ -54,25 +59,41 @@ def evaluate_function(
                 print(e)
                 stress = np.zeros((3, 3))
             state.stress = stress
-    _write_frames(*states, outputs=[outputs[0]])
+    if isinstance(inputs[0], Geometry):
+        return states[0]
+    else:
+        _write_frames(*states, outputs=[outputs[0]])
 
 
 @typeguard.typechecked
 @psiflow.serializable  # otherwise MixtureHamiltonian.hamiltonians is not serialized
 class Hamiltonian:
     external: Optional[psiflow._DataFuture]
 
-    def evaluate(self, dataset: Dataset, batch_size: Optional[int] = 100) -> Dataset:
-        future = batch_apply(
-            self.single_evaluate,
-            batch_size,
-            dataset.length(),
-            inputs=[dataset.extxyz],
-            outputs=[
-                psiflow.context().new_file("data_", ".xyz")
-            ],  # join_app needs outputs kwarg here!
-        )
-        return Dataset(None, future.outputs[0])
+    def evaluate(
+        self,
+        arg: Union[Dataset, Geometry, AppFuture[Geometry]],
+        batch_size: Optional[int] = 100,
+    ) -> Union[AppFuture, Dataset]:
+        if isinstance(arg, Dataset):
+            future = batch_apply(
+                [self.single_evaluate],
+                batch_size,
+                arg.length(),
+                inputs=[arg.extxyz],
+                outputs=[
+                    psiflow.context().new_file("data_", ".xyz")
+                ],  # join_app needs outputs kwarg here!
+            )
+            return Dataset(None, future.outputs[0])
+        else:
+            future = self.evaluate_app(
+                self.load_calculators,
+                inputs=[arg, self.external],
+                outputs=[],
+                **self.parameters,
+            )
+            return future
 
     # mostly for internal use
     def single_evaluate(self, dataset: Dataset) -> Dataset:

psiflow/sampling/order.py

@@ -1,20 +1,22 @@
 from __future__ import annotations  # necessary for type-guarding class methods
 
-from typing import Union
+from functools import partial
+from typing import Union, Optional
 
 import typeguard
 from ase.units import kJ, mol
 from parsl.app.app import python_app
 from parsl.dataflow.futures import AppFuture
 
 import psiflow
-from psiflow.data import Dataset
+from psiflow.data import Dataset, batch_apply
 from psiflow.geometry import Geometry
 from psiflow.hamiltonians._plumed import PlumedHamiltonian
 from psiflow.hamiltonians.hamiltonian import Hamiltonian
 
 
-def _insert_in_state(
+@typeguard.typechecked
+def insert_in_state(
     state: Geometry,
     name: str,
 ) -> Geometry:
@@ -24,7 +26,32 @@ def _insert_in_state(
     return state
 
 
-insert_in_state = python_app(_insert_in_state, executors=["default_threads"])
+@typeguard.typechecked
+def _insert(
+    state_or_states: Union[Geometry, list[Geometry]],
+    name: str,
+) -> Union[list[Geometry], Geometry]:
+    if not isinstance(state_or_states, list):
+        return insert_in_state(state_or_states, name)
+    else:
+        for state in state_or_states:
+            insert_in_state(state, name)  # modify list in place
+        return state_or_states
+
+
+insert = python_app(_insert, executors=["default_threads"])
+
+
+@typeguard.typechecked
+def insert_in_dataset(
+    data: Dataset,
+    name: str,
+) -> Dataset:
+    geometries = insert(
+        data.geometries(),
+        name,
+    )
+    return Dataset(geometries)
 
 
 @typeguard.typechecked
@@ -51,11 +78,29 @@ def __init__(self, name: str, hamiltonian: Hamiltonian):
         super().__init__(name)
         self.hamiltonian = hamiltonian
 
-    def evaluate(self, state: Union[Geometry, AppFuture]) -> AppFuture:
-        return insert_in_state(
-            self.hamiltonian.evaluate(Dataset([state]))[0],
-            self.name,
-        )
+    def evaluate(
+        self,
+        arg: Union[Dataset, Geometry, AppFuture[Geometry]],
+        batch_size: Optional[int] = 100,
+    ) -> Union[Dataset, AppFuture]:
+        if isinstance(arg, Dataset):
+            # avoid batching the dataset twice:
+            # apply hamiltonian in batched sense and put insert afterwards
+            funcs = [
+                self.hamiltonian.single_evaluate,
+                partial(insert_in_dataset, name=self.name),
+            ]
+            future = batch_apply(
+                funcs,
+                batch_size,
+                arg.length(),
+                inputs=[arg.extxyz],
+                outputs=[psiflow.context().new_file("data_", ".xyz")],
+            )
+            return Dataset(None, future.outputs[0])
+        else:
+            state = self.hamiltonian.evaluate(arg)
+            return insert(state, self.name)
 
     def __eq__(self, other):
         if type(other) is not HamiltonianOrderParameter:

psiflow/sampling/walker.py

@@ -187,6 +187,7 @@ def quench(walkers: list[Walker], dataset: Dataset) -> None:
     coefficients = []
     for walker in walkers:
         c = all_hamiltonians.get_coefficients(1.0 * walker.hamiltonian)
+        assert c is not None
         coefficients.append(c)
     coefficients = np.array(coefficients)
 

tests/test_hamiltonian.py

@@ -32,9 +32,9 @@ def test_get_filename_hills():
 RESTART
 UNITS LENGTH=A ENERGY=kj/mol TIME=fs
 CV: VOLUME
-CV0: CV
+CV0: CV #lkasdjf
 METAD ARG=CV0 SIGMA=100 HEIGHT=2 PACE=50 LABEL=metad FILE=test_hills sdld
-METADD ARG=CV SIGMA=100 HEIGHT=2 PACE=50 LABEL=metad sdld
+METADD ARG=CV SIGMA=100 HEIGHT=2 PACE=50 LABEL=metad sdld #fjalsdkfj
 PRINT ARG=CV,metad.bias STRIDE=10 FILE=COLVAR
 FLUSH STRIDE=10
 """
@@ -62,6 +62,10 @@ def test_einstein(dataset, dataset_h2):
     for i in range(1, 10):
         assert evaluated[i].result().energy > 0.0
         assert not np.allclose(evaluated[i].result().stress, 0.0)
+        assert np.allclose(
+            evaluated[i].result().energy,
+            hamiltonian.evaluate(evaluated[i]).result().energy,
+        )
 
     # test evaluation with NullState in data
     data = hamiltonian.evaluate(dataset[:5] + Dataset([NullState]) + dataset[5:10])

tests/test_sampling.py

@@ -450,6 +450,13 @@ def test_order_parameter(dataset):
     assert state.energy is None
     assert np.allclose(CV, np.linalg.det(dataset[3].result().cell))
 
+    # test batch evaluation of order parameter
+    data = order.evaluate(dataset[:10], batch_size=5)
+    volumes = data.get("CV").result()
+    for i in range(10):
+        volume = np.linalg.det(dataset[i].result().cell)
+        assert np.allclose(volume, volumes[i])
+
 
 def test_walker_serialization(dataset, tmp_path):
     einstein = EinsteinCrystal(dataset[0], force_constant=0.1)