Merge pull request #64 from OpenBioSim/feature_repex

Add support for HREX and REST2

Author: lohedges

README.md

@@ -59,11 +59,58 @@ This is a `bzip2` compressed file that will need to be extracted before use.
 
 #### Running SOMD2 using one or more GPUs
 
-In order to run using GPUs you will first need to set the relevant environment variable. For example, to run using 4 CUDA enabled GPUS set `CUDA_VISIBLE_DEVICES=0,1,2,3` (for openCL and HIP use `OPENCL_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` respectively). 
-
-By default SOMD2 will run using the CPU platform, however if the relevant environment variable has been set (as above) the new platform will be detected and set. In the case that this detection fails, or if there are multiple platforms available, the `--platform` option can be set (for example `--platform cuda`).
-
-By default, SOMD2 will automatically manage the distribution of lambda windows across all listed devices. In order to restrict the number of devices used the `--max_gpus` option can be set, for example setting `max_gpus=2` while `CUDA_VISIBLE_DEVICES` are set as above would restrict SOMD2 to using only GPUs 0 and 1.
+In order to run using GPUs you will first need to set the relevant environment
+variable. For example, to run using 4 CUDA enabled GPUS set `CUDA_VISIBLE_DEVICES=0,1,2,3`
+(for openCL and HIP use `OPENCL_VISIBLE_DEVICES` and `HIP_VISIBLE_DEVICES` respectively).
+
+By default `SOMD2` will run using the CPU platform, however if the relevant
+environment variable has been set (as above) the new platform will be detected
+and set. In the case that this detection fails, or if there are multiple platforms
+available, the `--platform` option can be set (for example `--platform cuda`).
+
+By default, `SOMD2` will automatically manage the distribution of lambda windows
+across all listed devices. In order to restrict the number of devices used
+the `--max_gpus` option can be set, for example setting `max_gpus=2` while
+`CUDA_VISIBLE_DEVICES` are set as above would restrict `SOMD2` to using only
+GPUs 0 and 1.
+
+## Replica exchange
+
+`SOMD2` supports Hamiltonian replica exchange (HREX) simulations, which can be
+enabled using the `--replica-exchange` option. Note that dynamics contexts will
+be created up-front for all replicas, so this can be memory intensive. As such,
+replica exchange is intended for use on multi-GPU nodes with a large amount of
+memory. For optimal performance, it is recommended that the number of replicas
+be a multiple of the number of GPUs. It is also possible to oversubscribe the
+GPUs, i.e. have more than one replica running on a GPU at a time. This can be
+controlled via the `--oversubscription-factor` option, e.g. a value of 2 would
+allow 2 replicas to run on each GPU at a time.
+
+The swap frequency for replica exchange is controlled by the `energy-frequency`
+option, i.e. we compute the energies for all replicas at this frequency, then
+attempt to mix the replicas. A larger value will improve performance, but may
+reduce the efficiency of the exchange.
+
+## REST2
+
+We also support Replica Exchange with Solute Scaling
+([REST2](https://pubs.acs.org/doi/10.1021/jp204407d)) simulations to facilitate sampling for perturbations
+involving conformational changes, e.g.  ring flips. This can be enabled
+using the `--rest2-scale` option, which specifies the "temperature" of the
+REST2 region relative to the rest of the system. By default, the REST2 region
+comprises _all_ atoms in perturbable molecules, but can be controlled via the
+`--rest2-selection` option. This should be a `Sire` selection string that specifies
+additional atoms of interest, i.e. those in regular, non-perturbable molecules.
+If the selection does contain atoms within perturbable molecules, then only
+those atoms within the perturbable molecules will be considered as part of the
+REST2 region, i.e. you can select a sub-set of atoms within a perturbable
+molecule to be scaled.
+
+By default, the REST2 schedule is a triangular function that starts and ends
+at 1.0, with a peak at the middle of the lambda schedule corresponding to
+the value of `--rest2-scale`. By passing multiple values for `--rest2-scale`, the
+user can fully control the schedule. When doing so, the number of values must
+match the number of lambda windows.
 
 ## Analysis
 
@@ -89,6 +136,19 @@ pmf2, overlap2 = BSS.FreeEnergy.Relative.analyse("output2")
 free_nrg = BSS.FreeEnergy.Relative.difference(pmf1, pmf2)
 ```
 
+## Truncated MBAR analysis
+
+When running HREX with a large number of replicas it can become computationally
+expensive to compute energies. (We need the energies of each replica at each
+lamdba value.) As a shortcut, it's possible to truncate the neighbourhood of
+windows for which we compute energies, then use a large null energy for the
+remaining windows. This can be controlled via the `--num-energy-neighbours` option.
+For example, setting this to 2 would compute energies for the current window and
+its two neighbours on either side. The value assigned to the remaining windows
+can be controlled via the `--null-energy` option. The number of neighbours should
+be chosen as a trade off between accuracy and computational cost. A value of around
+20% of the number of replicas has been found to be a good starting point.
+
 ## Note for SOMD1 users
 
 For existing users of `somd1`, it's possible to generate input for `somd2` by passing
@@ -124,3 +184,7 @@ somd2 somd1.bss --pert-file somd1.pert --somd1-compatibility
 ```
 
 (This only shows the limited options required. Others will take default values and can be set accordingly.)
+
+If you want to load an existing system from a perturbation file and use the
+new `somd2` ghost atom bonded-term modifications, then simply omit the
+`--somd1-compatibility` option.

environment.yaml

@@ -8,3 +8,4 @@ dependencies:
   - biosimspace
   - git
   - loguru
+  - numba

src/somd2/__init__.py

@@ -1,7 +1,7 @@
 ######################################################################
 # SOMD2: GPU accelerated alchemical free-energy engine.
 #
-# Copyright: 2023-2024
+# Copyright: 2023-2025
 #
 # Authors: The OpenBioSim Team <team@openbiosim.org>
 #
@@ -19,6 +19,17 @@
 # along with SOMD2. If not, see <http://www.gnu.org/licenses/>.
 #####################################################################
 
+# Make sure we used the mixed API so we can use BioSimSpace.
+try:
+    import sire as _sr
+
+    _sr.use_mixed_api(support_old_module_names=False)
+    _sr.convert.supported_formats()
+
+    del _sr
+except ImportError:
+    pass
+
 # Disable Sire progress bars until we work out the best way to handle
 # them for the SOMD2 runner, i.e. when running multiple dynamics objects
 # in parallel.

src/somd2/_utils/__init__.py

@@ -1,7 +1,7 @@
 ######################################################################
 # SOMD2: GPU accelerated alchemical free-energy engine.
 #
-# Copyright: 2023-2024
+# Copyright: 2023-2025
 #
 # Authors: The OpenBioSim Team <team@openbiosim.org>
 #

src/somd2/_utils/_ghosts.py

@@ -1,7 +1,7 @@
 ######################################################################
 # SOMD2: GPU accelerated alchemical free-energy engine.
 #
-# Copyright: 2023-2024
+# Copyright: 2023-2025
 #
 # Authors: The OpenBioSim Team <team@openbiosim.org>
 #
@@ -19,6 +19,8 @@
 # along with SOMD2. If not, see <http://www.gnu.org/licenses/>.
 #####################################################################
 
+__all__ = ["boresch"]
+
 from sire.system import System as _System
 from sire.legacy.System import System as _LegacySystem
 
@@ -32,7 +34,7 @@
 from . import _lam_sym
 
 
-def _boresch(system, k_hard=100, k_soft=5, optimise_angles=True):
+def boresch(system, k_hard=100, k_soft=5, optimise_angles=True):
     """
     Apply Boresch modifications to ghost atom bonded terms to avoid non-physical
     coupling between the ghost atoms and the physical region.

src/somd2/_utils/_somd1.py

@@ -1,7 +1,7 @@
 ######################################################################
 # SOMD2: GPU accelerated alchemical free-energy engine.
 #
-# Copyright: 2023-2024
+# Copyright: 2023-2025
 #
 # Authors: The OpenBioSim Team <team@openbiosim.org>
 #
@@ -19,14 +19,72 @@
 # along with SOMD2. If not, see <http://www.gnu.org/licenses/>.
 #####################################################################
 
+__all__ = ["apply_pert", "make_compatible", "reconstruct_system"]
+
 from sire.system import System as _System
 from sire.legacy.System import System as _LegacySystem
 
 import sire.legacy.MM as _SireMM
 import sire.legacy.Mol as _SireMol
 
 
-def _make_compatible(system):
+def apply_pert(system, pert_file):
+    """
+    Helper function to apply a perturbation to a reference system.
+
+    Parameters
+    ----------
+
+    system: sr.system.System
+        The reference system.
+
+    pert_file: str
+        Path to a stream file containing the perturbation to apply to the
+        reference system.
+
+    Returns
+    -------
+
+    system: sire.system.System
+        The perturbable system.
+    """
+
+    if not isinstance(system, _System):
+        raise TypeError("'system' must be of type 'sr.system.System'.")
+
+    if not isinstance(pert_file, str):
+        raise TypeError("'pert_file' must be of type 'str'.")
+
+    from sire import morph as _morph
+
+    # Get the non-water molecules in the system.
+    non_waters = system["not water"].molecules()
+
+    # Try to apply the perturbation to each non-water molecule.
+    is_pert = False
+    for mol in non_waters:
+        # Exclude ions.
+        if mol.num_atoms() > 1:
+            try:
+                pert_mol = _morph.create_from_pertfile(mol, pert_file)
+                is_pert = True
+                break
+            except:
+                pass
+
+    if not is_pert:
+        raise ValueError(f"Failed to apply the perturbation in '{pert_file}'.")
+
+    # Update the molecule.
+    system.update(pert_mol)
+
+    # Link to the reference state.
+    system = _morph.link_to_reference(system)
+
+    return system
+
+
+def make_compatible(system):
     """
     Makes a perturbation SOMD1 compatible.
 
@@ -547,54 +605,108 @@ def _make_compatible(system):
     return system
 
 
-def _apply_pert(system, pert_file):
+def reconstruct_system(system):
     """
-    Helper function to apply a perturbation to a reference system.
+    Reconstruct a perturbable system to its original state, i.e. extract the
+    end states for each perturbable molecule and re-merge them using the original
+    mapping. This removes any ghost atom modifications applied via a perturbation
+    file.
 
     Parameters
     ----------
 
-    system: sr.system.System
-        The reference system.
-
-    pert_file: str
-        Path to a stream file containing the perturbation to apply to the
-        reference system.
+    system : sire.system.System, sire.legacy.System.System
+        The system containing the molecules to be perturbed.
 
     Returns
     -------
 
-    system: sire.system.System
-        The perturbable system.
+    system : sire.system.System
+        The updated system.
     """
 
-    if not isinstance(system, _System):
-        raise TypeError("'system' must be of type 'sr.system.System'.")
-
-    if not isinstance(pert_file, str):
-        raise TypeError("'pert_file' must be of type 'str'.")
+    import BioSimSpace as _BSS
 
     from sire import morph as _morph
 
-    # Get the non-water molecules in the system.
-    non_waters = system["not water"].molecules()
+    # Check the system is a Sire system.
+    if not isinstance(system, (_System, _LegacySystem)):
+        raise TypeError(
+            "'system' must of type 'sire.system.System' or 'sire.legacy.System.System'"
+        )
 
-    # Try to apply the perturbation to each non-water molecule.
-    is_pert = False
-    for mol in non_waters:
-        # Exclude ions.
-        if mol.num_atoms() > 1:
-            try:
-                pert_mol = _morph.create_from_pertfile(mol, pert_file)
-                is_pert = True
-                break
-            except:
-                pass
+    # Extract the legacy system.
+    if isinstance(system, _LegacySystem):
+        system = _System(system)
 
-    if not is_pert:
-        raise ValueError(f"Failed to apply the perturbation in '{pert_file}'.")
+    # Clone the system.
+    system = system.clone()
 
-    # Update the molecule.
-    system.update(pert_mol)
+    # Search for perturbable molecules.
+    try:
+        pert_mols = system.molecules("property is_perturbable")
+    except KeyError:
+        raise KeyError("No perturbable molecules in the system")
+
+    # Store a dummy element for ghost atoms.
+    ghost = _SireMol.Element(0)
+
+    # Loop over all perturbable molecules.
+    for mol in pert_mols:
+
+        # Extract the end states.
+        ref = _morph.extract_reference(mol)
+        pert = _morph.extract_perturbed(mol)
+
+        # Find the indices for non-ghost atoms.
+        ref_idxs = []
+        for x, (a, e) in enumerate(
+            zip(ref.property("ambertype"), ref.property("element"))
+        ):
+            if a == "du" or e == ghost:
+                continue
+            else:
+                ref_idxs.append(x)
+        pert_idxs = []
+        for x, (a, e) in enumerate(
+            zip(pert.property("ambertype"), pert.property("element"))
+        ):
+            if a == "du" or e == ghost:
+                continue
+            else:
+                pert_idxs.append(x)
+
+        # Convert to BioSimSpace molecules and extract the non-ghost atoms.
+        ref = _BSS._SireWrappers.Molecule(ref).extract(ref_idxs)
+        pert = _BSS._SireWrappers.Molecule(pert).extract(pert_idxs)
+
+        # Work out the mapping.
+        idx0 = 0
+        idx1 = 0
+        mapping = {}
+        for x, atom in enumerate(mol.atoms()):
+            at0 = atom.property("ambertype0")
+            at1 = atom.property("ambertype1")
+
+            if at0 != "du" and at1 != "du":
+                mapping[idx0] = idx1
+
+            if at0 != "du":
+                idx0 += 1
+
+            if at1 != "du":
+                idx1 += 1
+
+        # Re-merge the molecules.
+        merged = _BSS.Align.merge(ref, pert, mapping=mapping, force=True)
+
+        # Give the molecule the same number as the original.
+        merged = merged._sire_object.edit().renumber(mol.number()).molecule().commit()
+
+        # Update the system.
+        system.update(merged)
+
+    # Link to the reference state.
+    system = _morph.link_to_reference(system)
 
     return system

src/somd2/app/__init__.py

@@ -1,7 +1,7 @@
 ######################################################################
 # SOMD2: GPU accelerated alchemical free-energy engine.
 #
-# Copyright: 2023-2024
+# Copyright: 2023-2025
 #
 # Authors: The OpenBioSim Team <team@openbiosim.org>
 #