Using gufe get typed components API

Author: ijpulidos

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/protocols/nonequilibrium_cycling.py

@@ -138,23 +138,16 @@ def _execute(self, ctx, *, protocol, state_a, state_b, mapping, **inputs):
         )
         from feflow.utils.hybrid_topology import HybridTopologyFactory
         from feflow.utils.charge import get_alchemical_charge_difference
-        from feflow.utils.misc import (
-            get_typed_components,
-            register_ff_parameters_template,
-        )
+        from feflow.utils.misc import register_ff_parameters_template
 
         # Check compatibility between states (same receptor and solvent)
         self._check_states_compatibility(state_a, state_b)
 
         # Get receptor components from systems if found (None otherwise)
-        solvent_comps = get_typed_components(
-            state_a, SolventComponent
-        )  # this returns a set
-        solvent_comp_a = (
-            solvent_comps.pop() if solvent_comps else None
-        )  # Get the first component if exists
-        protein_comps_a = get_typed_components(state_a, ProteinComponent)
-        small_mols_a = get_typed_components(state_a, SmallMoleculeComponent)
+        solvent_comps = state_a.get_components_of_type(SolventComponent)
+        solvent_comp_a = solvent_comps.pop() if solvent_comps else None # there must be at most one solvent comp
+        protein_comps_a = state_a.get_components_of_type(ProteinComponent)
+        small_mols_a = state_a.get_components_of_type(SmallMoleculeComponent)
 
         # Get alchemical components
         alchemical_comps = get_alchemical_components(state_a, state_b)
@@ -187,10 +180,9 @@ def _execute(self, ctx, *, protocol, state_a, state_b, mapping, **inputs):
         # Parameterizing small molecules
         self.logger.info("Parameterizing molecules")
         # Get small molecules from states
-        # TODO: Refactor if/when gufe provides the functionality https://github.com/OpenFreeEnergy/gufe/issues/251
-        state_a_small_mols = get_typed_components(state_a, SmallMoleculeComponent)
-        state_b_small_mols = get_typed_components(state_b, SmallMoleculeComponent)
-        all_small_mols = state_a_small_mols | state_b_small_mols
+        state_a_small_mols = state_a.get_components_of_type(SmallMoleculeComponent)
+        state_b_small_mols = state_b.get_components_of_type(SmallMoleculeComponent)
+        all_small_mols = set(state_a_small_mols) | set(state_b_small_mols)
 
         # Generate and register FF parameters in the system generator template
         all_openff_mols = [comp.to_openff() for comp in all_small_mols]

feflow/utils/misc.py

@@ -8,41 +8,6 @@
 import openmm.app
 
 
-# TODO: should this be a method for the gufe.ChemicalSystem class?
-def get_typed_components(
-    system: gufe.ChemicalSystem, comptype: type[gufe.Component]
-) -> set[gufe.Component]:
-    """
-    Retrieve all components of a specific type from a `gufe.ChemicalSystem`.
-
-    This function searches the components within the provided chemical system
-    and returns a list of all components matching the specified type.
-
-    Parameters
-    ----------
-    system : gufe.ChemicalSystem
-        The chemical system from which to extract components.
-    comptype : Type[gufe.Component]
-        The type of component to search for, such as `ProteinComponent`,
-        `SmallMoleculeComponent`, or `SolventComponent`.
-
-    Returns
-    -------
-    set[gufe.Component]
-        A set of unique components matching the specified type. If no components
-        of the given type are found, an empty set is returned.
-
-    """
-    if not issubclass(comptype, gufe.Component):
-        raise TypeError(
-            f"`comptype` must be a subclass of `gufe.Component`. Got: {comptype}"
-        )
-
-    ret_comps = {comp for comp in system.values() if isinstance(comp, comptype)}
-
-    return ret_comps
-
-
 def register_ff_parameters_template(system_generator, charge_settings, openff_mols):
     """
     Register force field parameters in the system generator using provided charge settings