Refactor runner and add support for replica exchange.

src/somd2/app/run.py

@@ -39,7 +39,7 @@ def cli():
 
     from somd2 import _logger
     from somd2.config import Config
-    from somd2.runner import Runner
+    from somd2.runner import Runner, RepexRunner
 
     from somd2.io import yaml_to_dict
 
@@ -80,7 +80,10 @@ def cli():
     config = Config(**args)
 
     # Instantiate a Runner object to run the simulation.
-    runner = Runner(system, config)
+    if config.replica_exchange:
+        runner = RepexRunner(system, config)
+    else:
+        runner = Runner(system, config)
 
     # Run the simulation.
     runner.run()

src/somd2/config/_config.py

@@ -119,7 +119,8 @@ def __init__(
         platform="auto",
         max_threads=None,
         max_gpus=None,
-        run_parallel=True,
+        oversubscription_factor=1,
+        replica_exchange=False,
         output_directory="output",
         restart=False,
         write_config=True,
@@ -196,17 +197,16 @@ def __init__(
             Lennard-Jones interaction.
 
         restraints: sire.mm._MM.Restraints
-            A single set of restraints, or a list of
-            sets of restraints that will be applied to
-            the atoms during the simulation.
+            A single set of restraints, or a list of sets of restraints that
+            will be applied to the atoms during the simulation.
 
         constraint: str
             Constraint type to use for non-perturbable molecules.
 
         perturbable_constraint: str
             Constraint type to use for perturbable molecules. If None, then
-            this will be set according to what is chosen for the
-            non-perturbable constraint.
+            this will be set according to what is chosen for the non-perturbable
+            constraint.
 
         include_constrained_energies: bool
             Whether to include constrained energies in the potential.
@@ -244,7 +244,8 @@ def __init__(
             Whether to use constraints during equilibration.
 
         energy_frequency: str
-            Frequency at which to output energy data.
+            Frequency at which to output energy data. If running using 'replica_exchange',
+            then this will also be the frequency at which replica swaps are attempted.
 
         save_trajectories: bool
             Whether to save trajectory files
@@ -270,8 +271,12 @@ def __init__(
             Maximum number of GPUs to use for simulation (Default None, uses all available.)
             Does nothing if platform is set to CPU.
 
-        run_parallel: bool
-            Whether to run simulation in parallel.
+        oversubscription_factor: int
+            Factor by which to oversubscribe jobs on GPUs during replica exchange simulations.
+
+        replica_exchange: bool
+            Whether to run replica exchange simulation. Currently this can only be used when
+            GPU resources are available.
 
         output_directory: str
             Path to a directory to store output files.
@@ -349,7 +354,8 @@ def __init__(
         self.platform = platform
         self.max_threads = max_threads
         self.max_gpus = max_gpus
-        self.run_parallel = run_parallel
+        self.oversubscription_factor = oversubscription_factor
+        self.replica_exchange = replica_exchange
         self.restart = restart
         self.somd1_compatibility = somd1_compatibility
         self.pert_file = pert_file
@@ -1203,9 +1209,12 @@ def max_gpus(self, max_gpus):
         ):
             if "CUDA_VISIBLE_DEVICES" in _os.environ:
                 self._max_gpus = len(_os.environ["CUDA_VISIBLE_DEVICES"].split(","))
+            elif "OPENCL_VISIBLE_DEVICES" in _os.environ:
+                self._max_gpus = len(_os.environ["OPENCL_VISIBLE_DEVICES"].split(","))
+            elif "HIP_VISIBLE_DEVICES" in _os.environ:
+                self._max_gpus = len(_os.environ["HIP_VISIBLE_DEVICES"].split(","))
             else:
                 self._max_gpus = 0
-
         else:
             try:
                 self._max_gpus = int(max_gpus)
@@ -1217,14 +1226,31 @@ def max_gpus(self, max_gpus):
                 )
 
     @property
-    def run_parallel(self):
-        return self._run_parallel
-
-    @run_parallel.setter
-    def run_parallel(self, run_parallel):
-        if not isinstance(run_parallel, bool):
-            raise ValueError("'run_parallel' must be of type 'bool'")
-        self._run_parallel = run_parallel
+    def oversubscription_factor(self):
+        return self._oversubscription_factor
+
+    @oversubscription_factor.setter
+    def oversubscription_factor(self, oversubscription_factor):
+        if not isinstance(oversubscription_factor, int):
+            try:
+                oversubscription_factor = int(oversubscription_factor)
+            except:
+                raise ValueError("'oversubscription_factor' must be of type 'int'")
+
+        if oversubscription_factor < 1:
+            raise ValueError("'oversubscription_factor' must be greater than 1")
+
+        self._oversubscription_factor = oversubscription_factor
+
+    @property
+    def replica_exchange(self):
+        return self._replica_exchange
+
+    @replica_exchange.setter
+    def replica_exchange(self, replica_exchange):
+        if not isinstance(replica_exchange, bool):
+            raise ValueError("'replica_exchange' must be of type 'bool'")
+        self._replica_exchange = replica_exchange
 
     @property
     def restart(self):

src/somd2/io/_io.py

@@ -50,11 +50,13 @@ def dataframe_to_parquet(df, metadata, filepath=None, filename=None):
 
     metadata: dict
         Dictionary containing metadata to be saved with the dataframe.
-        Currently just temperature and lambda value.
 
     filepath: str or pathlib.PosixPath
         The of the parent directory in to which the parquet file will be saved.
         If None, save to current working directory.
+
+    filename: str
+        The name of the parquet file to be saved. If None, a default name will be used.
     """
 
     if filepath is None:
@@ -81,7 +83,7 @@ def dataframe_to_parquet(df, metadata, filepath=None, filename=None):
     return filepath / filename
 
 
-def dict_to_yaml(data_dict, path, filename="config.yaml"):
+def dict_to_yaml(data_dict, filename="config.yaml", path=None):
     """
     Write a dictionary to a YAML file.
 
@@ -91,16 +93,20 @@ def dict_to_yaml(data_dict, path, filename="config.yaml"):
     data_dict: dict
         The dictionary to be written to a YAML file.
 
-    path: str or pathlib.PosixPath
-        The path to the YAML file to be written.
-
     filename: str
         The name of the YAML file to be written (default 'config.yaml').
+
+    path: str or pathlib.PosixPath
+        The path to the YAML file to be written.
     """
     import yaml as _yaml
 
-    try:
+    if path is None:
+        path = _Path(filename)
+    else:
         path = _Path(path) / filename
+
+    try:
         # Ensure the parent directory for the file exists
         path.parent.mkdir(parents=True, exist_ok=True)
         # Open the file in write mode and write the dictionary as YAML

src/somd2/runner/__init__.py

@@ -20,3 +20,4 @@
 #####################################################################
 
 from ._runner import *
+from ._repex import *