Add alchemical groups (#15)

* Add alchemical groups

* Update environment

* Update README

* Fix typos in README

* Fix replacement of name in modification_kwargs

* Fix setting of MLPotential in strategies

* Update tests

* Update hooks and formatting

* Update environments

* Update action version in workflow

* Bump Python version from 3.10 to 3.12

* Correct env name

* Update pre/post/skp dependencies of MLInterpolationModification

* Update dependencies for ml mods

* Update order for ml mods

* Keep edge/node ordering

* Bump to version 0.2.2

* Add Tang-Toennies to LJSoftCore

* Updates to OpenFF strategy

* Generalise CT

Author: JMorado

.github/workflows/main.yaml

@@ -15,7 +15,7 @@ jobs:
     strategy:
       matrix:
         os: [ubuntu-latest]
-        python-version: ["3.10"]
+        python-version: ["3.12"]
 
     runs-on: ${{ matrix.os }}
     steps:
@@ -28,15 +28,15 @@ jobs:
           python-version: ${{ matrix.python-version }}
 
       - name: Pre-commit hooks
-        uses: pre-commit/action@v3.0.0
+        uses: pre-commit/action@v3.0.1
 
   build:
     needs: pre-commit
 
     strategy:
       matrix:
         os: [ubuntu-latest] #, macos-latest, windows-latest]
-        python-version: ["3.10"]
+        python-version: ["3.12"]
 
     defaults:
       run:
@@ -48,7 +48,7 @@ jobs:
         uses: actions/checkout@v4
 
       - name: Set up environment
-        uses: mamba-org/setup-micromamba@v1
+        uses: mamba-org/setup-micromamba@v2
         with:
           environment-file: environment.yaml
           environment-name: gha-test-env
@@ -94,7 +94,7 @@ jobs:
   #   strategy:
   #     matrix:
   #         os: [ubuntu-latest]
-  #         python-version: ["3.10"]
+  #         python-version: ["3.12"]
 
   #   runs-on: ${{ matrix.os }}
   #   steps:

.pre-commit-config.yaml

@@ -1,41 +1,13 @@
-# Pre-commit hooks for Python code
-# Last revision by: Joao Morado
-# Last revision date: 8.01.2023
-# See https://pre-commit.com for more information
-# See https://pre-commit.com/hooks.html for more hooks
-
 repos:
--   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v3.2.0
-    hooks:
-    -   id: trailing-whitespace
-    -   id: end-of-file-fixer
-      #-   id: check-added-large-files
--   repo: https://github.com/psf/black
-    rev: 23.1.0
-    hooks:
-    -   id: black
-    -   id: black-jupyter
--   repo: https://github.com/keewis/blackdoc
-    rev: v0.3.8
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.12.8 
     hooks:
-    -   id: blackdoc
-      #-   id: blackdoc-autoupdate-black
--   repo: https://github.com/PyCQA/isort
-    rev: 5.12.0
-    hooks:
-    -   id: isort
-        args: ["--profile", "black", "--filter-files"]
--   repo: https://github.com/PyCQA/flake8
-    rev: 6.0.0
-    hooks:
-    -   id: flake8
-        additional_dependencies: [flake8-docstrings]
-        args: [--max-line-length=127, --exit-zero]
-        verbose: True
--   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.1.1
+      - id: ruff
+        args: ["--fix"] 
+
+  - repo: https://github.com/pre-commit/mirrors-mypy
+    rev: v1.17.1  
     hooks:
-    -   id: mypy
-        args: [--no-strict-optional, --ignore-missing-imports, --namespace-packages, --explicit-package-bases]
-        additional_dependencies: ["types-PyYAML"]
+      - id: mypy
+        verbose: true
+        entry: bash -c 'mypy "$@" || true' --

README.md

@@ -11,12 +11,10 @@ A package to run hybrid ML/MM free energy simulations.
 1. [Installation](#installation)
 2. [Alchemical Modifications](#alchemical-modifications)
 3. [Running a Multistate Equilibrium Free Energy Simulation](#running-a-multistate-equilibrium-free-energy-simulation)
+   1. [Using Multiple Alchemical Groups](#using-multiple-alchemical-groups)
 4. [Dynamics and EMLE settings](#dynamics-and-emle-settings)
-
     1. [Sire Strategy](#sire-strategy)
     2. [OpenFF Strategy](#openff-strategy)
-
-
 5. [Log Level](#log-level)
 
 ## Installation
@@ -98,6 +96,78 @@ U_kln = fes.run_single_state(1000, 1000, 6)
 np.save("U_kln_mm_sol_6.npy", np.asarray(U_kln))
 ```
 
+### Using Multiple Alchemical Groups
+
+For more complex transformations, you can define multiple alchemical groups that can be transformed independently or simultaneously. This is particularly useful when you want to apply different transformations to different regions of your system or transform multiple ligands separately.
+
+To use multiple alchemical groups, specify the group name as a suffix after a colon in the lambda schedule:
+
+```python
+from fes_ml.fes import FES
+import numpy as np
+
+# Define lambda schedule for multiple alchemical groups
+lambda_schedule = {
+    # Group 1: Turn off LJ and charges for ligand 1
+    "LJSoftCore:ligand1": [1.0, 0.8, 0.6, 0.4, 0.2, 0.0, 0.0, 0.0],
+    "ChargeScaling:ligand1": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 0.5, 0.0],
+
+    # Group 2: Turn off LJ and charges for ligand 2
+    "LJSoftCore:ligand2": [1.0, 1.0, 0.8, 0.6, 0.4, 0.2, 0.0, 0.0],
+    "ChargeScaling:ligand2": [1.0, 1.0, 1.0, 1.0, 1.0, 0.5, 0.33, 0.0],
+
+    # Group 3: Interpolate between MM and ML for the entire system
+    "MLInterpolation:system": [0.0, 0.0, 0.0, 0.2, 0.4, 0.6, 0.8, 1.0]
+}
+
+# Define atom indices for each alchemical group
+ligand1_atoms = [1, 2, 3, 4, 5]       # Atoms belonging to first ligand
+ligand2_atoms = [20, 21, 22, 23, 24]  # Atoms belonging to second ligand
+system_atoms = list(range(1, 50))     # All atoms for ML/MM interpolation
+
+# Define per-group alchemical atoms
+modifications_kwargs = {
+    "LJSoftCore:ligand1": {
+        "alchemical_atoms": ligand1_atoms
+    },
+    "ChargeScaling:ligand1": {
+        "alchemical_atoms": ligand1_atoms
+    },
+    "LJSoftCore:ligand2": {
+        "alchemical_atoms": ligand2_atoms
+    },
+    "ChargeScaling:ligand2": {
+        "alchemical_atoms": ligand2_atoms
+    },
+    "MLInterpolation:system": {
+        "alchemical_atoms": system_atoms
+    }
+}
+```
+
+#### Multiple Instances of the Same Modification Type
+
+You can also use multiple instances of the same modification type for the same group of atoms. For example, to interpolate between two sets of `CustomLJ` parameters:
+
+```python
+lambda_schedule = {
+    "LJSoftCore:openff1": [1.0, 0.8, 0.6, 0.4, 0.2, 0.0],
+    "LJSoftCore:openff2": [0.0, 0.2, 0.4, 0.6, 0.8, 1.0],
+    "CustomLJ:openff1": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+    "CustomLJ:openff2": [1.0, 1.0, 1.0, 1.0, 1.0, 1.0],
+}
+
+# Define different LJ parameters for each region
+modifications_kwargs = {
+    "CustomLJ:openff1": {
+        "lj_offxml": "openff_unconstrained-1.0.0.offxml",
+    },
+    "CustomLJ:openff2": {
+        "lj_offxml": "openff_unconstrained-2.0.0.offxml",
+    }
+}
+```
+
 ## Dynamics and EMLE settings
 
 In fes-ml, the default strategy for creating OpenMM systems is through Sire. Additionally, fes-ml offers the OpenFF strategy. You can select the desired creation strategy, either `'sire'` or `'openff'`, using the `strategy_name` argument when calling the `fes.create_alchemical_states` method to create the alchemical systems. Most other simulation configurations can also be set by passing additional arguments to this method. For details on customization, refer to the definitions of the `SireCreationStrategy` and `OpenFFCreationStrategy` classes.

analysis/analyse.py

@@ -6,29 +6,26 @@ channels:
 
 dependencies:
   - ambertools
-  - ase
   - compilers
-  - cudatoolkit<11.9
-  - deepmd-kit
+  - cudatoolkit=11.8
   - eigen
   - loguru
   - openmm>=8.1
   - openmm-torch
+  - openmm-ml
   - openmmforcefields
   - openff-toolkit
   - openff-interchange
   - nnpops
   - pip
   - pybind11
-  - pytorch
+  - pytorch=*=*cuda*
   - python
   - pyyaml
   - sire
   - torchani
   - pygit2
-  - xtb-python
   - pip:
     - git+https://github.com/chemle/emle-engine.git
-    - git+https://github.com/openmm/openmm-ml.git
     - coloredlogs
     - mace-torch

environment.yaml

@@ -125,7 +125,7 @@ def water_atom(
 
         # Select atoms for which you want to calculate RDF
         water = universe.select_atoms(f"type O and not {ligand_selection}")
-        ligand = universe.select_atoms(ligand_selection)
+        # ligand = universe.select_atoms(ligand_selection)
         atoms = [
             universe.select_atoms(f"index {atom.index} and {ligand_selection}")
             for atom in universe.select_atoms(ligand_selection)
@@ -327,10 +327,10 @@ def __init__(self) -> None:
     def _plot_stdout_with_time_base(txt_file, save_location, data_name, prop):
         df = pd.read_csv(txt_file, sep=",")
         data = df[f"{data_name}"]
-        steps = df[f'#"Step"']
+        steps = df['#"Step"']
 
         plt.plot(steps, data, "r-", linewidth=2, color="maroon")
-        plt.xlabel(f"Steps")
+        plt.xlabel("Steps")
         plt.ylabel(f"{data_name}")
         plt.title(f"Plot of {prop} for {txt_file}")
         # plt.savefig(f"{save_location}")
@@ -378,12 +378,12 @@ def plot_energy_all_windows(
             txt_file = f"{folder}/{base_name}{w}.txt"
             df = pd.read_csv(txt_file, sep=",")
             data = df[f"{data_name}"]
-            steps = df[f'#"Step"']
+            steps = df['#"Step"']
             x_axis = [r for r in range(1, len(steps) + 1, 1)]
 
             plt.plot(x_axis, data, linewidth=2, alpha=0.3, label=label, color=colors[w])
             legend_labels.append(label)
-            plt.xlabel(f"Steps")
+            plt.xlabel("Steps")
             plt.ylabel(f"{data_name}")
             plt.title(f"Plot of {prop} for {txt_file}")
             plt.legend(bbox_to_anchor=(1, 0.5), loc="center left")

environment_rascal.yaml

@@ -9,7 +9,7 @@
     out = f"OUTPUT_{i}"
     try:
         f = glob.glob(out + "/*npy")[0]
-    except:
+    except Exception:
         break
     U_kln.append(np.load(f)[:, frames_disc::step])
 

examples/additional_scripts/analysis.py

@@ -3,7 +3,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 from openmm import unit
-from pymbar import MBAR, timeseries
+from pymbar import MBAR
 
 if len(sys.argv) != 3 or sys.argv[1] in ["-h", "--help"]:
     print("Usage: python analyse.py <file_name> <temperature>")

examples/openff_strategy/mts_benchmark/ml/agg_output.py

@@ -9,7 +9,7 @@
     out = f"OUTPUT_{i}"
     try:
         f = glob.glob(out + "/*npy")[0]
-    except:
+    except Exception:
         break
     U_kln.append(np.load(f)[:, frames_disc::step])