Miscellaneous Hybrid topology factory optimization

.github/workflows/ci.yaml

@@ -23,12 +23,11 @@ defaults:
 jobs:
   tests:
     runs-on: ${{ matrix.os }}-latest
-    name: "💻-${{matrix.os }} 🐍-${{ matrix.python-version }} - pydantic ${{matrix.pydantic}}"
+    name: "💻-${{matrix.os }} 🐍-${{ matrix.python-version }}"
     strategy:
       fail-fast: false
       matrix:
         os: ["ubuntu"]
-        pydantic: ["1", "2"]
         python-version:
           - "3.10"
           - "3.11"
@@ -53,7 +52,6 @@ jobs:
           cache-downloads: true
           create-args: >-
             python=${{ matrix.python-version }}
-            pydantic=${{ matrix.pydantic }}
           init-shell: bash
 
       - name: "Install"
@@ -82,12 +80,12 @@ jobs:
           # Set the OFE_SLOW_TESTS to True if running a Cron job
           OFE_SLOW_TESTS: ${{ fromJSON('{"false":"false","true":"true"}')[github.event_name != 'pull_request'] }}
         run: |
-              pytest -n auto -v --cov=feflow --cov-report=xml --durations=10
+              pytest -n logical -v --cov=feflow --cov-report=xml --durations=10
 
       - name: codecov
-        if: ${{ github.repository == 'choderalab/feflow'
+        if: ${{ github.repository == 'OpenFreeEnergy/feflow'
                 && github.event_name == 'pull_request' }}
-        uses: codecov/codecov-action@v3
+        uses: codecov/codecov-action@v4
         with:
           token: ${{ secrets.CODECOV_TOKEN }}
           file: coverage.xml

devtools/conda-envs/test_env.yaml

@@ -4,16 +4,16 @@ channels:
   - openeye
 dependencies:
     # Base depends
-  - gufe >=0.9.5
+  - gufe ~=1.7.1
   - numpy
-  - openfe >=0.15  # TODO: Remove once we don't depend on openfe
+  - openfe ~=1.7.0   # TODO: Remove once we don't depend on openfe
   - openff-units
   - openmm
   - openmmforcefields >=0.14.1  # TODO: remove when upstream deps fix this
-  - pymbar <4
-  - pydantic >=1.10.17
+  - openmmtools >=0.23.0
+  - pymbar ~=4.0
+  - pydantic >=1.10.17, <3
   - python
-
     # Testing (optional deps)
   - espaloma_charge  # To us Espaloma FF in tests
   - openeye-toolkits

docs/tutorials/protein-mutations-neq.ipynb

@@ -0,0 +1,49 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "id": "47736a8f-7446-4600-b06d-4c8ffd05f66c",
+   "metadata": {},
+   "source": [
+    "TODO:\n",
+    "Add a tutorial here on how to use the protein-mutation capabilities in `FEFlow`.\n",
+    "\n",
+    "What should it have?\n",
+    "* How to mutate a protein from a PDB file with `pdbfixer`\n",
+    "* How to create the mapping between the WT and MUT\n",
+    "  * check mapping is correct (sanity check)\n",
+    "* Setting up the protocol\n",
+    "* Executing the protocol"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "fa00a848-79a2-46c3-a141-6e24ce0f7f4f",
+   "metadata": {},
+   "outputs": [],
+   "source": []
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.12.9"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

feflow/mapping/protein_mutation.py

@@ -0,0 +1,119 @@
+"""
+Module with mapping objects that are useful for the protein mutation protocol.
+"""
+
+import re
+from gufe import ProteinComponent
+from gufe.mapping import AtomMapping
+
+# TODO: Find a better place for this pattern string, maybe feflow.utils?
+regex_mut_str = r"^[a-zA-Z]{3}-?\d+-?[a-zA-Z]{3}$"
+
+
+class ProteinMutationMapping(AtomMapping):
+    """
+    Container for an atom mapping for single residue mutations in proteins.
+
+    This is a specialized version of :class:`.AtomMapping` for
+    :class:`.ProteinComponent` which stores the mapping as a dict of
+    integers.
+    """
+
+    componentA: ProteinComponent
+    componentB: ProteinComponent
+    _mutation_spec: re.Match[regex_mut_str]
+    _compA_to_compB: dict[int, int]
+
+    def __init__(
+        self,
+        componentA: ProteinComponent,
+        componentB: ProteinComponent,
+        componentA_to_componentB: dict[int, int],
+        mutation_spec: re.Match[regex_mut_str],
+    ):
+        """
+        Parameters
+        ----------
+        componentA, componentB : ProteinComponent
+          the protein molecules on either end of the mapping
+        componentA_to_componentB : dict[int, int]
+          correspondence of indices of atoms between wild type and mutation; the
+          keys are indices in componentA and the values are indices in
+          componentB.
+          These are checked that they are within the possible indices of the
+          respective components.
+        mutation_spec: str
+          Mutation specification string. E.g. "LYS42ALA".
+        """
+        super().__init__(componentA, componentB)
+
+        # validate compA_to_compB
+        nA = self.componentA.to_openmm_topology().getNumAtoms()
+        nB = self.componentB.to_openmm_topology().getNumAtoms()
+        for i, j in componentA_to_componentB.items():
+            if not (0 <= i < nA):
+                raise ValueError(
+                    f"Got invalid index for ComponentA ({i}); " f"must be 0 <= n < {nA}"
+                )
+            if not (0 <= j < nB):
+                raise ValueError(
+                    f"Got invalid index for ComponentB ({i}); " f"must be 0 <= n < {nB}"
+                )
+
+        self._compA_to_compB = componentA_to_componentB
+        self._mutation_spec = mutation_spec
+
+    @property
+    def componentA_to_componentB(self) -> dict[int, int]:
+        return dict(self._compA_to_compB)
+
+    @property
+    def componentB_to_componentA(self) -> dict[int, int]:
+        return {v: k for k, v in self._compA_to_compB.items()}
+
+    @property
+    def componentA_unique(self):
+        return (
+            i
+            for i in range(self.componentA.to_openmm_topology().getNumAtoms())
+            if i not in self._compA_to_compB
+        )
+
+    @property
+    def componentB_unique(self):
+        return (
+            i
+            for i in range(self.componentB.to_openmm_topology().getNumAtoms())
+            if i not in self._compA_to_compB.values()
+        )
+
+    @property
+    def mutation_spec(self):
+        """String with the specification of the mutation."""
+        return self._mutation_spec
+
+    @classmethod
+    def _defaults(cls):
+        return {}
+
+    def _to_dict(self) -> dict:
+        """Seralize to dictionary"""
+        return {
+            "componentA": self.componentA,
+            "componentB": self.componentB,
+            "componentA_to_componentB": self._compA_to_compB,
+            "mutation_spec": self.mutation_spec,
+        }
+
+    @classmethod
+    def _from_dict(cls, dct: dict):
+        """Deserialize from dictionary"""
+        mapping = dct["componentA_to_componentB"]
+        fixed = {int(k): int(v) for k, v in mapping.items()}
+
+        return cls(
+            componentA=dct["componentA"],
+            componentB=dct["componentB"],
+            componentA_to_componentB=fixed,
+            mutation_spec=dct["mutation_spec"],
+        )