Sampling overhaul

.gitignore

@@ -3,6 +3,7 @@ __pycache__/
 *.py[cod]
 *$py.class
 *.swp
+*.pyc
 
 # C extensions
 *.so

psiflow/__init__.py

@@ -34,41 +34,3 @@ def resolve_and_check(path: Path) -> Path:
 context = ExecutionContextLoader.context
 wait = ExecutionContextLoader.wait
 
-
-# TODO: EXECUTION
-#  - max_runtime is in seconds, max_simulation_time (and others) is in minutes
-#  - ExecutionDefinition gpu argument?
-#  - ExecutionDefinition wrap_in_timeout functionality
-#  - ExecutionDefinition wrap_in_srun functionality? Actually for MD this is more iffy right
-#  - ExecutionDefinition why properties?
-#  - ExecutionDefinition centralize wq_resources
-#  - configuration file with all options
-#  - timeout -s 9 or -s 15?
-#  - executor keys are hardcoded strings..
-#  - define a 'format_env_variables' util
-#  - update GPAW + ORCA + Default containers
-#    include s-dftd3 in modelevaluation + install s-dftd3 with openmp
-#  - what with mem_per_node / mem_per_worker
-#  - always GPU for training?
-#  - currently reference mpi_args have to be tuple according to typeguard..
-#  - cores_per_block has to be specified even when exclusive..?
-#  - can we do something with WQ priority?
-#  - see chatgpt convo for process memory limits and such
-#  - make /tmp for app workdirs an option?
-#  - what is scaling_cores_per_worker in WQ
-#  - can we clean up psiflow_internal slightly?
-#  -
-# TODO: REFERENCE
-#  - reference MPI args not really checked
-#  - mpi flags are very finicky across implementations --> use ENV args?
-#    OMPI_MCA_orte_report_bindings=1 I_MPI_DEBUG=4
-#  - commands ends with 'exit 0' - what if we do not want to exit yet?
-#  - some actual logging?
-#  - safe_compute_dataset functionality?
-#  - ReferenceDummy is a bit sloppy
-#  -
-# TODO: MISC
-#  - think about test efficiency
-#  - some imports take a very long time
-#  - fix serialisation
-#  -

psiflow/execution.py

@@ -10,7 +10,7 @@
 from threading import Lock
 
 # see https://stackoverflow.com/questions/59904631/python-class-constants-in-dataclasses
-from typing import Any, Optional, Union, ClassVar, Protocol
+from typing import Any, Optional, Union, ClassVar, Protocol, Iterable
 
 import parsl
 import psutil
@@ -155,7 +155,7 @@ def create_executor(self, path: Path) -> ParslExecutor:
 
             # only ModelEvaluation / ModelTraining / default_htex run in containers
             if not isinstance(self, ReferenceEvaluation) and self.container:
-                prepend = self.container.launch_command()
+                prepend = self.container.launch_command(self.gpu)
                 worker_executable = f"{prepend} work_queue_worker"
             else:
                 worker_executable = "work_queue_worker"
@@ -261,27 +261,37 @@ def client_command(self):
         command_list = ["psiflow-client"]
         return " ".join(command_list)
 
-    def get_client_args(
-        self,
-        hamiltonian_name: str,
-        nwalkers: int,
-        motion: str,
-    ) -> list[str]:
-        if "MACE" in hamiltonian_name:
-            if motion in ["minimize", "vibrations"]:
-                dtype = "float64"
-            else:
-                dtype = "float32"
-            nclients = min(nwalkers, self.max_workers)
-            if self.gpu:
-                template = "--dtype={} --device=cuda:{}"
-                args = [template.format(dtype, i) for i in range(nclients)]
-            else:
-                template = "--dtype={} --device=cpu"
-                args = [template.format(dtype) for i in range(nclients)]
-            return args
+    # def get_client_args(
+    #     self,
+    #     hamiltonian_name: str,
+    #     nwalkers: int,
+    #     motion: str,
+    # ) -> list[str]:
+    #     # TODO: redo this
+    #     if "MACE" in hamiltonian_name:
+    #         if motion in ["minimize", "vibrations"]:
+    #             dtype = "float64"
+    #         else:
+    #             dtype = "float32"
+    #         nclients = min(nwalkers, self.max_workers)
+    #         if self.gpu:
+    #             template = "--dtype={} --device=cuda:{}"
+    #             args = [template.format(dtype, i) for i in range(nclients)]
+    #         else:
+    #             template = "--dtype={} --device=cpu"
+    #             args = [template.format(dtype) for i in range(nclients)]
+    #         return args
+    #     else:
+    #         return [""]
+
+    def get_driver_devices(self, nwalkers: int) -> list[dict]:
+        # assumes driver is GPU capable
+        # TODO: what if only 1 gpu is available?
+        nclients = min(nwalkers, self.max_workers)
+        if self.gpu:
+            return [{'device': f'cuda:{i}'} for i in range(nclients)]
         else:
-            return [""]
+            return [{'device': 'cpu'} for _ in range(nclients)]
 
     def wq_resources(self, nwalkers):
         if self.use_threadpool:
@@ -433,7 +443,7 @@ class ReferenceSpec(Protocol):
     name: ClassVar[str]
     reference_args: ClassVar[tuple[str, ...]]
     mpi_command: str
-    mpi_args: tuple[str, ...]
+    mpi_args: Iterable[str]
     executable: str
 
     def launch_command(self) -> str:

psiflow/free_energy/phonons.py

@@ -1,31 +1,28 @@
-from __future__ import annotations  # necessary for type-guarding class methods
-
 import xml.etree.ElementTree as ET
 from typing import Optional, Union
 
 import numpy as np
 import parsl
-import typeguard
 from ase.units import Bohr, Ha, J, _c, _hplanck, _k, kB, second
 from parsl.app.app import bash_app, python_app
 from parsl.dataflow.futures import AppFuture
 
 import psiflow
 from psiflow.data import Dataset
 from psiflow.geometry import Geometry, mass_weight
-from psiflow.hamiltonians import Hamiltonian, MixtureHamiltonian
+from psiflow.hamiltonians import Hamiltonian, MACEHamiltonian
 from psiflow.sampling.sampling import (
     setup_sockets,
-    make_force_xml,
-    make_start_command,
-    make_client_command
+    make_server_command,
+    make_driver_commands,
+    make_wait_for_sockets_command,
 )
+from psiflow.sampling.optimize import setup_forces, export_env_command
 from psiflow.utils.apps import multiply
 from psiflow.utils.io import load_numpy, save_xml
 from psiflow.utils import TMP_COMMAND, CD_COMMAND
 
 
-@typeguard.typechecked
 def _compute_frequencies(hessian: np.ndarray, geometry: Geometry) -> np.ndarray:
     assert hessian.shape[0] == hessian.shape[1]
     assert len(geometry) * 3 == hessian.shape[0]
@@ -35,7 +32,6 @@ def _compute_frequencies(hessian: np.ndarray, geometry: Geometry) -> np.ndarray:
 compute_frequencies = python_app(_compute_frequencies, executors=["default_threads"])
 
 
-@typeguard.typechecked
 def _harmonic_free_energy(
     frequencies: Union[float, np.ndarray],
     temperature: float,
@@ -65,9 +61,7 @@ def _harmonic_free_energy(
 harmonic_free_energy = python_app(_harmonic_free_energy, executors=["default_threads"])
 
 
-@typeguard.typechecked
 def setup_motion(
-    mode: str,
     asr: str,
     pos_shift: float,
     energy_shift: float,
@@ -91,43 +85,40 @@ def setup_motion(
 
 
 def _execute_ipi(
-    hamiltonian_names: list[str],
-    client_args: list[list[str]],
+    driver_kwargs: list[dict],
     command_server: str,
-    command_client: str,
-    stdout: str = "",
-    stderr: str = "",
+    env_vars: dict = {},
+    stdout: str = parsl.AUTO_LOGNAME,
+    stderr: str = parsl.AUTO_LOGNAME,
     inputs: list = [],
     outputs: list = [],
     parsl_resource_specification: Optional[dict] = None,
 ) -> str:
-    command_start = make_start_command(command_server, inputs[0], inputs[1])
-    commands_client = []
-    for i, name in enumerate(hamiltonian_names):
-        args = client_args[i]
-        assert len(args) == 1  # only have one client per hamiltonian
-        for arg in args:
-            commands_client += make_client_command(command_client, name, inputs[2 + i], inputs[1], arg),
-
-    command_end = f'{command_server} --cleanup'
-    command_copy = f'cp i-pi.output_full.hess {outputs[0]}'
+    file_xml, file_xyz_in, *files_in = inputs
+    command_start = make_server_command(
+        command_server, file_xml, file_xyz_in, outputs[0], [], []
+    )
+    command_wait = make_wait_for_sockets_command(
+        set(d["address"] for d in driver_kwargs)
+    )
+    commands_driver = make_driver_commands(driver_kwargs, file_xyz_in, files_in)
 
     command_list = [
         TMP_COMMAND,
         CD_COMMAND,
+        export_env_command(env_vars),
         command_start,
-        *commands_client,
+        command_wait,
+        *commands_driver,
         "wait",
-        command_end,
-        command_copy,
+        f"cp i-pi.output_full.hess {outputs[0]}",
     ]
     return "\n".join(command_list)
 
 
 execute_ipi = bash_app(_execute_ipi, executors=["ModelEvaluation"])
 
 
-@typeguard.typechecked
 def compute_harmonic(
     state: Union[Geometry, AppFuture],
     hamiltonian: Hamiltonian,
@@ -136,26 +127,23 @@ def compute_harmonic(
     pos_shift: float = 0.01,
     energy_shift: float = 0.00095,
 ) -> AppFuture:
-    hamiltonian: MixtureHamiltonian = 1 * hamiltonian
-    names = hamiltonian.get_named_components()
-    sockets = setup_sockets(names)
-    forces = make_force_xml(hamiltonian, names)
+    components, force_xml = setup_forces(hamiltonian)
+    sockets = setup_sockets(components)
 
     initialize = ET.Element("initialize", nbeads="1")
     start = ET.Element("file", mode="ase", cell_units="angstrom")
     start.text = " start_0.xyz "
     initialize.append(start)
-    motion = setup_motion(mode, asr, pos_shift, energy_shift)
+    motion = setup_motion(asr, pos_shift, energy_shift)
+    if mode != "fd":
+        raise NotImplementedError
 
     system = ET.Element("system")
     system.append(initialize)
     system.append(motion)
-    system.append(forces)
-
-    # output = setup_output(keep_trajectory)
+    system.append(force_xml)
 
     simulation = ET.Element("simulation", mode="static")
-    # simulation.append(output)
     for socket in sockets:
         simulation.append(socket)
     simulation.append(system)
@@ -169,33 +157,22 @@ def compute_harmonic(
         simulation,
         outputs=[context.new_file("input_", ".xml")],
     ).outputs[0]
-    inputs = [
-        input_future,
-        Dataset([state]).extxyz,
-    ]
-    inputs += hamiltonian.serialize(dtype="float64")
-
-    client_args = []
-    for name in names:
-        args = definition.get_client_args(name, 1, "vibrations")
-        client_args.append(args)
-    outputs = [
-        context.new_file("hess_", ".txt"),
-    ]
+    inputs = [input_future, Dataset([state]).extxyz]
 
-    command_server = definition.server_command()
-    command_client = definition.client_command()
-    resources = definition.wq_resources(1)
+    driver_kwargs = []
+    for i, comp in enumerate(components):
+        inputs.append(comp.hamiltonian.serialize_function(dtype="float64"))
+        kwargs = {"idx": i, "address": comp.address}
+        if isinstance(comp.hamiltonian, MACEHamiltonian):
+            kwargs |= definition.get_driver_devices(1)[0]
+        driver_kwargs.append(kwargs)
 
     result = execute_ipi(
-        names,
-        client_args,
-        command_server,
-        command_client,
-        stdout=parsl.AUTO_LOGNAME,
-        stderr=parsl.AUTO_LOGNAME,
+        driver_kwargs,
+        definition.server_command(),
+        env_vars=definition.env_vars,
         inputs=inputs,
-        outputs=outputs,
-        parsl_resource_specification=resources,
+        outputs=[context.new_file("hess_", ".txt")],
+        parsl_resource_specification=definition.wq_resources(1),
     )
-    return multiply(load_numpy(inputs=[result.outputs[0]]), Ha / Bohr**2)
+    return multiply(load_numpy(inputs=[result.outputs[0]]), Ha / Bohr**2)

psiflow/geometry.py

@@ -304,9 +304,16 @@ def from_string(cls, s: str, natoms: Optional[int] = None) -> Optional[Geometry]
             if key.startswith("order_"):
                 order[key.replace("order_", "")] = value
 
+        # TODO: this needs a better solution
+        cell = comment_dict.pop("Lattice", np.zeros((3, 3)))
+        if isinstance(cell, str):
+            # most likely an array of NaNs
+            cell = np.array(cell.split()).reshape((3, 3))
+        cell = cell.T # transposed!
+        pbc = sum(comment_dict.pop("pbc", [])) == 3
         geometry = cls(
             per_atom=per_atom,
-            cell=comment_dict.pop("Lattice", np.zeros((3, 3))).T,  # transposed!
+            cell=cell if pbc else np.zeros((3, 3)),
             energy=comment_dict.pop("energy", None),
             stress=comment_dict.pop("stress", None),
             delta=comment_dict.pop("delta", None),