From 87d96bc75f0093c5d8abd3d71207f9a817fd815f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Thu, 23 Jan 2025 11:56:24 -0500
Subject: [PATCH 01/11] Optional multiple comps in modeller generation

---
 openfe/protocols/openmm_afe/base.py           |  2 +-
 .../protocols/openmm_md/plain_md_methods.py   |  2 +-
 .../protocols/openmm_rfe/equil_rfe_methods.py |  2 +-
 .../protocols/openmm_utils/system_creation.py | 57 ++++++++++---------
 .../openmm_utils/system_validation.py         |  1 +
 .../test_openmm_equil_rfe_protocols.py        |  4 +-
 6 files changed, 37 insertions(+), 31 deletions(-)

diff --git a/openfe/protocols/openmm_afe/base.py b/openfe/protocols/openmm_afe/base.py
index 4e201d220..7f121267f 100644
--- a/openfe/protocols/openmm_afe/base.py
+++ b/openfe/protocols/openmm_afe/base.py
@@ -452,7 +452,7 @@ def _get_modeller(
 
             # get OpenMM modeller + dictionary of resids for each component
             system_modeller, comp_resids = system_creation.get_omm_modeller(
-                protein_comp=protein_component,
+                protein_comps=protein_component,
                 solvent_comp=solvent_component,
                 small_mols=smc_components,
                 omm_forcefield=system_generator.forcefield,
diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py
index 2770580a9..e848fda13 100644
--- a/openfe/protocols/openmm_md/plain_md_methods.py
+++ b/openfe/protocols/openmm_md/plain_md_methods.py
@@ -591,7 +591,7 @@ def run(self, *, dry=False, verbose=True,
 
             # c. get OpenMM Modeller + a resids dictionary for each component
             stateA_modeller, comp_resids = system_creation.get_omm_modeller(
-                protein_comp=protein_comp,
+                protein_comps=protein_comp,
                 solvent_comp=solvent_comp,
                 small_mols=smc_components,
                 omm_forcefield=system_generator.forcefield,
diff --git a/openfe/protocols/openmm_rfe/equil_rfe_methods.py b/openfe/protocols/openmm_rfe/equil_rfe_methods.py
index 545dc02f5..a5c2959f3 100644
--- a/openfe/protocols/openmm_rfe/equil_rfe_methods.py
+++ b/openfe/protocols/openmm_rfe/equil_rfe_methods.py
@@ -783,7 +783,7 @@ def run(self, *, dry=False, verbose=True,
 
             # c. get OpenMM Modeller + a dictionary of resids for each component
             stateA_modeller, comp_resids = system_creation.get_omm_modeller(
-                protein_comp=protein_comp,
+                protein_comps=protein_comp,
                 solvent_comp=solvent_comp,
                 small_mols=dict(chain(off_small_mols['stateA'],
                                       off_small_mols['both'])),
diff --git a/openfe/protocols/openmm_utils/system_creation.py b/openfe/protocols/openmm_utils/system_creation.py
index 77ae8274d..227c4908f 100644
--- a/openfe/protocols/openmm_utils/system_creation.py
+++ b/openfe/protocols/openmm_utils/system_creation.py
@@ -11,7 +11,7 @@
 from openff.toolkit import Molecule as OFFMol
 from openff.units.openmm import to_openmm, ensure_quantity
 from openmmforcefields.generators import SystemGenerator
-from typing import Optional
+from typing import Optional, Iterable
 from pathlib import Path
 from gufe.settings import OpenMMSystemGeneratorFFSettings, ThermoSettings
 from gufe import (
@@ -134,9 +134,9 @@ def get_system_generator(
 
 
 def get_omm_modeller(
-    protein_comp: Optional[ProteinComponent],
+    protein_comps: Iterable[ProteinComponent] | ProteinComponent,
     solvent_comp: Optional[SolventComponent],
-    small_mols: dict[SmallMoleculeComponent, OFFMol],
+    small_mols: Iterable[SmallMoleculeComponent] | SmallMoleculeComponent,
     omm_forcefield : app.ForceField,
     solvent_settings : OpenMMSolvationSettings
 ) -> ModellerReturn:
@@ -146,11 +146,11 @@ def get_omm_modeller(
 
     Parameters
     ----------
-    protein_comp : Optional[ProteinComponent]
+    protein_comps : Iterable[ProteinComponent] or ProteinComponent
       Protein Component, if it exists.
-    solvent_comp : Optional[ProteinCompoinent]
+    solvent_comp : Optional[SolventComponent]
       Solvent Component, if it exists.
-    small_mols : dict
+    small_mols : Iterable[SmallMoleculeComponent] or SmallMoleculeComponent
       Small molecules to add.
     omm_forcefield : app.ForceField
       ForceField object for system.
@@ -188,28 +188,33 @@ def _add_small_mol(comp,
     # Create empty modeller
     system_modeller = app.Modeller(app.Topology(), [])
 
-    # If there's a protein in the system, we add it first to the Modeller
-    if protein_comp is not None:
-        system_modeller.add(protein_comp.to_openmm_topology(),
-                            protein_comp.to_openmm_positions())
-        # add missing virtual particles (from crystal waters)
-        system_modeller.addExtraParticles(omm_forcefield)
-        component_resids[protein_comp] = np.array(
-          [r.index for r in system_modeller.topology.residues()]
-        )
-        # if we solvate temporarily rename water molecules to 'WAT'
-        # see openmm issue #4103
-        if solvent_comp is not None:
-            for r in system_modeller.topology.residues():
-                if r.name == 'HOH':
-                    r.name = 'WAT'
+    # We first add all the protein components, if any
+    if protein_comps:
+        protein_comps = iter(protein_comps)  # Support both input iterable or not
+        for protein_comp in protein_comps:
+            system_modeller.add(protein_comp.to_openmm_topology(),
+                                protein_comp.to_openmm_positions())
+            # add missing virtual particles (from crystal waters)
+            system_modeller.addExtraParticles(omm_forcefield)
+            component_resids[protein_comp] = np.array(
+              [r.index for r in system_modeller.topology.residues()]
+            )
+            # if we solvate temporarily rename water molecules to 'WAT'
+            # see openmm issue #4103
+            if solvent_comp is not None:
+                for r in system_modeller.topology.residues():
+                    if r.name == 'HOH':
+                        r.name = 'WAT'
 
     # Now loop through small mols
-    for comp, mol in small_mols.items():
-        _add_small_mol(comp, mol, system_modeller, component_resids)
-
-    # Add solvent if neeeded
-    if solvent_comp is not None:
+    if small_mols:
+        small_mols = iter(small_mols)  # Support both input iterable or not
+        for small_mol_comp in small_mols:
+            _add_small_mol(small_mol_comp, small_mol_comp.to_openff(), system_modeller,
+                           component_resids)
+
+    # Add solvent if needed
+    if solvent_comp:
         conc = solvent_comp.ion_concentration
         pos = solvent_comp.positive_ion
         neg = solvent_comp.negative_ion
diff --git a/openfe/protocols/openmm_utils/system_validation.py b/openfe/protocols/openmm_utils/system_validation.py
index f1710d93e..fadae7ff9 100644
--- a/openfe/protocols/openmm_utils/system_validation.py
+++ b/openfe/protocols/openmm_utils/system_validation.py
@@ -37,6 +37,7 @@ def get_alchemical_components(
     ValueError
       If there are any duplicate components in states A or B.
     """
+    # TODO: We probably want to rewrite this using sets. As in _check_states_compatibility in feflow.protocols.ProteinMutationProtocol
     matched_components: dict[Component, Component]  = {}
     alchemical_components: dict[str, list[Component]] = {
         'stateA': [], 'stateB': [],
diff --git a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py
index 41065c177..f0e1d3919 100644
--- a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py
+++ b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py
@@ -1697,7 +1697,7 @@ def benzene_solvent_openmm_system(benzene_modifications):
     )
 
     modeller, _ = system_creation.get_omm_modeller(
-        protein_comp=None,
+        protein_comps=None,
         solvent_comp=openfe.SolventComponent(),
         small_mols={smc: offmol},
         omm_forcefield=system_generator.forcefield,
@@ -1738,7 +1738,7 @@ def benzene_tip4p_solvent_openmm_system(benzene_modifications):
     )
 
     modeller, _ = system_creation.get_omm_modeller(
-        protein_comp=None,
+        protein_comps=None,
         solvent_comp=openfe.SolventComponent(),
         small_mols={smc: offmol},
         omm_forcefield=system_generator.forcefield,

From 1608c1cb385358dbee119f2332910eef4c7e7abe Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Wed, 29 Jan 2025 19:07:14 -0500
Subject: [PATCH 02/11] Specifying units in docstring

---
 openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
index 41e83238c..979eab96c 100644
--- a/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
+++ b/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
@@ -677,7 +677,7 @@ def set_and_check_new_positions(mapping, old_topology, new_topology,
         state.
     tolerance : float
         Warning threshold for deviations along any dimension (x,y,z) in mapped
-        atoms between the "old" and "new" positions. Default 1.0.
+        atoms between the "old" and "new" positions, in Angstroms. Default 1.0.
     """
     # Get the positions in Angstrom as raw numpy arrays
     old_pos_array = old_positions.value_in_unit(omm_unit.angstrom)

From 7cd99194b04d458dd1168e8f68753908193149a8 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Tue, 25 Mar 2025 17:10:22 -0400
Subject: [PATCH 03/11] Accept multiple solvent components in arg. Homogenizing
 accepted component types.

---
 openfe/protocols/openmm_afe/base.py           |  2 +-
 .../protocols/openmm_md/plain_md_methods.py   |  2 +-
 .../protocols/openmm_rfe/equil_rfe_methods.py |  2 +-
 .../protocols/openmm_utils/system_creation.py | 22 +++++++++++--------
 .../test_openmm_equil_rfe_protocols.py        |  4 ++--
 5 files changed, 18 insertions(+), 14 deletions(-)

diff --git a/openfe/protocols/openmm_afe/base.py b/openfe/protocols/openmm_afe/base.py
index 18ea699ab..472793956 100644
--- a/openfe/protocols/openmm_afe/base.py
+++ b/openfe/protocols/openmm_afe/base.py
@@ -459,7 +459,7 @@ def _get_modeller(
             # get OpenMM modeller + dictionary of resids for each component
             system_modeller, comp_resids = system_creation.get_omm_modeller(
                 protein_comps=protein_component,
-                solvent_comp=solvent_component,
+                solvent_comps=solvent_component,
                 small_mols=smc_components,
                 omm_forcefield=system_generator.forcefield,
                 solvent_settings=solvation_settings,
diff --git a/openfe/protocols/openmm_md/plain_md_methods.py b/openfe/protocols/openmm_md/plain_md_methods.py
index fed5bf10d..24cb35423 100644
--- a/openfe/protocols/openmm_md/plain_md_methods.py
+++ b/openfe/protocols/openmm_md/plain_md_methods.py
@@ -604,7 +604,7 @@ def run(self, *, dry=False, verbose=True,
             # c. get OpenMM Modeller + a resids dictionary for each component
             stateA_modeller, comp_resids = system_creation.get_omm_modeller(
                 protein_comps=protein_comp,
-                solvent_comp=solvent_comp,
+                solvent_comps=solvent_comp,
                 small_mols=smc_components,
                 omm_forcefield=system_generator.forcefield,
                 solvent_settings=solvation_settings,
diff --git a/openfe/protocols/openmm_rfe/equil_rfe_methods.py b/openfe/protocols/openmm_rfe/equil_rfe_methods.py
index d33d2ad9c..2bb7ad167 100644
--- a/openfe/protocols/openmm_rfe/equil_rfe_methods.py
+++ b/openfe/protocols/openmm_rfe/equil_rfe_methods.py
@@ -797,7 +797,7 @@ def run(self, *, dry=False, verbose=True,
             # c. get OpenMM Modeller + a dictionary of resids for each component
             stateA_modeller, comp_resids = system_creation.get_omm_modeller(
                 protein_comps=protein_comp,
-                solvent_comp=solvent_comp,
+                solvent_comps=solvent_comp,
                 small_mols=dict(chain(off_small_mols['stateA'],
                                       off_small_mols['both'])),
                 omm_forcefield=system_generator.forcefield,
diff --git a/openfe/protocols/openmm_utils/system_creation.py b/openfe/protocols/openmm_utils/system_creation.py
index 05ee8822c..fad6a4566 100644
--- a/openfe/protocols/openmm_utils/system_creation.py
+++ b/openfe/protocols/openmm_utils/system_creation.py
@@ -8,7 +8,6 @@
 import numpy.typing as npt
 from openmm import app, MonteCarloBarostat
 from openmm import unit as omm_unit
-from openff.toolkit import Molecule as OFFMol
 from openff.units.openmm import to_openmm, ensure_quantity
 from openmmforcefields.generators import SystemGenerator
 from typing import Optional, Iterable
@@ -134,9 +133,9 @@ def get_system_generator(
 
 
 def get_omm_modeller(
-    protein_comps: Iterable[ProteinComponent] | ProteinComponent,
-    solvent_comp: Optional[SolventComponent],
-    small_mols: Iterable[SmallMoleculeComponent] | SmallMoleculeComponent,
+    protein_comps: Optional[Iterable[ProteinComponent] | ProteinComponent],
+    solvent_comps: Optional[Iterable[SolventComponent] | SolventComponent],
+    small_mols: Optional[Iterable[SmallMoleculeComponent] | SmallMoleculeComponent],
     omm_forcefield : app.ForceField,
     solvent_settings : OpenMMSolvationSettings
 ) -> ModellerReturn:
@@ -146,11 +145,11 @@ def get_omm_modeller(
 
     Parameters
     ----------
-    protein_comps : Iterable[ProteinComponent] or ProteinComponent
+    protein_comps : Optional[Iterable[ProteinComponent] or ProteinComponent]
       Protein Component, if it exists.
-    solvent_comp : Optional[SolventComponent]
+    solvent_comps : Optional[Iterable[SolventComponent] or SolventComponent]
       Solvent Component, if it exists.
-    small_mols : Iterable[SmallMoleculeComponent] or SmallMoleculeComponent
+    small_mols : Optional[Iterable[SmallMoleculeComponent] or SmallMoleculeComponent]
       Small molecules to add.
     omm_forcefield : app.ForceField
       ForceField object for system.
@@ -201,7 +200,7 @@ def _add_small_mol(comp,
             )
             # if we solvate temporarily rename water molecules to 'WAT'
             # see openmm issue #4103
-            if solvent_comp is not None:
+            if solvent_comps is not None:
                 for r in system_modeller.topology.residues():
                     if r.name == 'HOH':
                         r.name = 'WAT'
@@ -214,7 +213,12 @@ def _add_small_mol(comp,
                            component_resids)
 
     # Add solvent if neeeded
-    if solvent_comp is not None:
+    if solvent_comps:
+        # TODO: Support multiple solvent components? Is there a use case for it?
+        # Error out when we have more than one solvent component in the states/systems
+        if len(set(solvent_comps)) > 1:
+            raise ValueError("More than one solvent component found in systems. Only one supported.")
+        solvent_comp = set(solvent_comps).pop()  # Get the first (and only) solvent component
         # Do unit conversions if necessary
         solvent_padding = None
         box_size = None
diff --git a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py
index c3918b178..a9b7ef353 100644
--- a/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py
+++ b/openfe/tests/protocols/test_openmm_equil_rfe_protocols.py
@@ -1718,7 +1718,7 @@ def benzene_solvent_openmm_system(benzene_modifications):
 
     modeller, _ = system_creation.get_omm_modeller(
         protein_comps=None,
-        solvent_comp=openfe.SolventComponent(),
+        solvent_comps=openfe.SolventComponent(),
         small_mols={smc: offmol},
         omm_forcefield=system_generator.forcefield,
         solvent_settings=settings.solvation_settings,
@@ -1759,7 +1759,7 @@ def benzene_tip4p_solvent_openmm_system(benzene_modifications):
 
     modeller, _ = system_creation.get_omm_modeller(
         protein_comps=None,
-        solvent_comp=openfe.SolventComponent(),
+        solvent_comps=openfe.SolventComponent(),
         small_mols={smc: offmol},
         omm_forcefield=system_generator.forcefield,
         solvent_settings=settings.solvation_settings,

From 7c6dc2ccabbe30d5ec1ed47d52bf82adc70203ce Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Tue, 25 Mar 2025 17:11:01 -0400
Subject: [PATCH 04/11] We can handle multiple residue ids

---
 openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py | 5 +----
 1 file changed, 1 insertion(+), 4 deletions(-)

diff --git a/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
index 979eab96c..cdf207705 100644
--- a/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
+++ b/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
@@ -390,10 +390,6 @@ def _get_indices(topology, resids):
     residue_name : str
         Name of the residue to get the indices for.
     """
-    # TODO: remove, this shouldn't be necessary anymore
-    if len(resids) > 1:
-        raise ValueError("multiple residues were found")
-
     # create list of openmm residues
     top_res = [r for r in topology.residues() if r.index in resids]
 
@@ -502,6 +498,7 @@ def pick_H(i, j, x, y) -> int:
 
             to_del.append(pick_H(i, j, x, y))
 
+    # to_del = set(to_del) # Making a set to avoid repeated indices (TODO: TEST!!)
     for idx in to_del:
         del no_const_old_to_new_atom_map[idx]
 

From a0d3aea718070ecb455a0e53631555dda785458f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Fri, 11 Apr 2025 17:00:22 -0400
Subject: [PATCH 05/11] Make sure we handle single solvent_comp as well (force
 iterable)

---
 openfe/protocols/openmm_utils/system_creation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/openfe/protocols/openmm_utils/system_creation.py b/openfe/protocols/openmm_utils/system_creation.py
index fad6a4566..252ed7544 100644
--- a/openfe/protocols/openmm_utils/system_creation.py
+++ b/openfe/protocols/openmm_utils/system_creation.py
@@ -216,7 +216,7 @@ def _add_small_mol(comp,
     if solvent_comps:
         # TODO: Support multiple solvent components? Is there a use case for it?
         # Error out when we have more than one solvent component in the states/systems
-        if len(set(solvent_comps)) > 1:
+        if len(set(iter(solvent_comps))) > 1:
             raise ValueError("More than one solvent component found in systems. Only one supported.")
         solvent_comp = set(solvent_comps).pop()  # Get the first (and only) solvent component
         # Do unit conversions if necessary

From 542b99ac0e9199d7930299be375f61e33d5104ad Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Fri, 11 Apr 2025 17:15:59 -0400
Subject: [PATCH 06/11] reworking handling logic. Should accept both single and
 multiple comps.

---
 openfe/protocols/openmm_utils/system_creation.py | 16 +++++++++++++---
 1 file changed, 13 insertions(+), 3 deletions(-)

diff --git a/openfe/protocols/openmm_utils/system_creation.py b/openfe/protocols/openmm_utils/system_creation.py
index 252ed7544..082bf2956 100644
--- a/openfe/protocols/openmm_utils/system_creation.py
+++ b/openfe/protocols/openmm_utils/system_creation.py
@@ -189,7 +189,10 @@ def _add_small_mol(comp,
 
     # We first add all the protein components, if any
     if protein_comps:
-        protein_comps = iter(protein_comps)  # Support both input iterable or not
+        try:
+            protein_comps = iter(protein_comps)
+        except TypeError:
+            protein_comps = {protein_comps}  # make it a set/iterable with the comp
         for protein_comp in protein_comps:
             system_modeller.add(protein_comp.to_openmm_topology(),
                                 protein_comp.to_openmm_positions())
@@ -207,16 +210,23 @@ def _add_small_mol(comp,
 
     # Now loop through small mols
     if small_mols:
-        small_mols = iter(small_mols)  # Support both input iterable or not
+        try:
+            small_mols = iter(small_mols)
+        except:
+            small_mols = {small_mols} # make it a set/iterable with the comp
         for small_mol_comp in small_mols:
             _add_small_mol(small_mol_comp, small_mol_comp.to_openff(), system_modeller,
                            component_resids)
 
     # Add solvent if neeeded
     if solvent_comps:
+        try:
+            solvent_comps = iter(solvent_comps)
+        except TypeError:
+            solvent_comps = {solvent_comps}  # make it a set/iterable with the comp
         # TODO: Support multiple solvent components? Is there a use case for it?
         # Error out when we have more than one solvent component in the states/systems
-        if len(set(iter(solvent_comps))) > 1:
+        if len(set(solvent_comps)) > 1:
             raise ValueError("More than one solvent component found in systems. Only one supported.")
         solvent_comp = set(solvent_comps).pop()  # Get the first (and only) solvent component
         # Do unit conversions if necessary

From 3916cd35a091d7d5f4a493b4afc31b842f206e33 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Tue, 22 Apr 2025 17:51:01 -0400
Subject: [PATCH 07/11] Using list to simplify processing

---
 openfe/protocols/openmm_utils/system_creation.py | 15 ++++++++-------
 1 file changed, 8 insertions(+), 7 deletions(-)

diff --git a/openfe/protocols/openmm_utils/system_creation.py b/openfe/protocols/openmm_utils/system_creation.py
index 082bf2956..080aab745 100644
--- a/openfe/protocols/openmm_utils/system_creation.py
+++ b/openfe/protocols/openmm_utils/system_creation.py
@@ -212,23 +212,24 @@ def _add_small_mol(comp,
     if small_mols:
         try:
             small_mols = iter(small_mols)
-        except:
-            small_mols = {small_mols} # make it a set/iterable with the comp
+        except TypeError:
+            small_mols = {small_mols}  # make it a set/iterable with the comp
         for small_mol_comp in small_mols:
             _add_small_mol(small_mol_comp, small_mol_comp.to_openff(), system_modeller,
                            component_resids)
 
     # Add solvent if neeeded
     if solvent_comps:
+        # Making it a list to make our life easier -- TODO: Maybe there's a better type for this
         try:
-            solvent_comps = iter(solvent_comps)
+            solvent_comps = list(set(solvent_comps))  # if given iterable
         except TypeError:
-            solvent_comps = {solvent_comps}  # make it a set/iterable with the comp
+            solvent_comps = [solvent_comps]  # if not iterable, given single obj
         # TODO: Support multiple solvent components? Is there a use case for it?
-        # Error out when we have more than one solvent component in the states/systems
-        if len(set(solvent_comps)) > 1:
+        # Error out when we iter(have more than one solvent component in the states/systems
+        if len(solvent_comps) > 1:
             raise ValueError("More than one solvent component found in systems. Only one supported.")
-        solvent_comp = set(solvent_comps).pop()  # Get the first (and only) solvent component
+        solvent_comp = solvent_comps[0]  # Get the first (and only?) solvent component
         # Do unit conversions if necessary
         solvent_padding = None
         box_size = None

From 47e73cdee37880c8d04856078c848a59b60c3f3b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Wed, 18 Jun 2025 15:49:09 -0400
Subject: [PATCH 08/11] Raising error when solvent component is None. Remove
 unused import.

---
 .../protocols/openmm_rfe/equil_rfe_methods.py | 23 ++++++++++++-------
 .../openmm_utils/system_validation.py         |  1 -
 2 files changed, 15 insertions(+), 9 deletions(-)

diff --git a/openfe/protocols/openmm_rfe/equil_rfe_methods.py b/openfe/protocols/openmm_rfe/equil_rfe_methods.py
index 2bb7ad167..1b2e7de7c 100644
--- a/openfe/protocols/openmm_rfe/equil_rfe_methods.py
+++ b/openfe/protocols/openmm_rfe/equil_rfe_methods.py
@@ -154,14 +154,21 @@ def _get_alchemical_charge_difference(
                           "correction has been requested. Unfortunately "
                           "only absolute differences of 1 are supported.")
                 raise ValueError(errmsg)
-
-            ion = {-1: solvent_component.positive_ion,
-                   1: solvent_component.negative_ion}[difference]
-            wmsg = (f"A charge difference of {difference} is observed "
-                    "between the end states. This will be addressed by "
-                    f"transforming a water into a {ion} ion")
-            logger.warning(wmsg)
-            warnings.warn(wmsg)
+            if solvent_component:
+                ion = {-1: solvent_component.positive_ion,
+                       1: solvent_component.negative_ion}[difference]
+                wmsg = (f"A charge difference of {difference} is observed "
+                        "between the end states. This will be addressed by "
+                        f"transforming a water into a {ion} ion")
+                logger.warning(wmsg)
+                warnings.warn(wmsg)
+            else:
+                errmsg = (f"A charge difference of {difference} is observed "
+                          "between the end states. This would be addressed by "
+                          f"transforming a water into an ion. However, "
+                          "No solvent component was specified when trying to"
+                          "apply the charge correction.")
+                raise ValueError(errmsg)
         else:
             wmsg = (f"A charge difference of {difference} is observed "
                     "between the end states. No charge correction has "
diff --git a/openfe/protocols/openmm_utils/system_validation.py b/openfe/protocols/openmm_utils/system_validation.py
index fadae7ff9..9eefc9d92 100644
--- a/openfe/protocols/openmm_utils/system_validation.py
+++ b/openfe/protocols/openmm_utils/system_validation.py
@@ -5,7 +5,6 @@
 Protocols.
 """
 from typing import Optional, Tuple
-from openff.toolkit import Molecule as OFFMol
 from gufe import (
     Component, ChemicalSystem, SolventComponent, ProteinComponent,
     SmallMoleculeComponent

From 4641fb9da4df0f4929111b13edbf798eb1446f64 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Tue, 23 Sep 2025 00:22:52 -0400
Subject: [PATCH 09/11] Using set instead. Avoid repeated constraint indices

---
 .../protocols/openmm_rfe/_rfe_utils/topologyhelpers.py   | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
index cdf207705..dd4fdf357 100644
--- a/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
+++ b/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
@@ -470,7 +470,7 @@ def pick_H(i, j, x, y) -> int:
     # there are two reasons constraints would invalidate a mapping entry
     # 1) length of constraint changed (but both constrained)
     # 2) constraint removed to harmonic bond (only one constrained)
-    to_del = []
+    to_del = set()
     for (i, j), l_old in old_constraints.items():
         x, y = old_to_new_atom_map[i], old_to_new_atom_map[j]
 
@@ -481,12 +481,12 @@ def pick_H(i, j, x, y) -> int:
                 l_new = new_constraints.pop((y, x))
             except KeyError:
                 # type 2) constraint doesn't exist in new system
-                to_del.append(pick_H(i, j, x, y))
+                to_del.add(pick_H(i, j, x, y))
                 continue
 
         # type 1) constraint length changed
         if l_old != l_new:
-            to_del.append(pick_H(i, j, x, y))
+            to_del.add(pick_H(i, j, x, y))
 
     # iterate over new_constraints (we were .popping items out)
     # (if any left these are type 2))
@@ -496,9 +496,8 @@ def pick_H(i, j, x, y) -> int:
         for x, y in new_constraints:
             i, j = new_to_old[x], new_to_old[y]
 
-            to_del.append(pick_H(i, j, x, y))
+            to_del.add(pick_H(i, j, x, y))
 
-    # to_del = set(to_del) # Making a set to avoid repeated indices (TODO: TEST!!)
     for idx in to_del:
         del no_const_old_to_new_atom_map[idx]
 

From 77079537e90041b48d4937d18e82abaff36ee8c7 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Tue, 23 Sep 2025 23:28:43 -0400
Subject: [PATCH 10/11] Revert "Using set instead. Avoid repeated constraint
 indices"

This reverts commit 4641fb9da4df0f4929111b13edbf798eb1446f64.
---
 .../protocols/openmm_rfe/_rfe_utils/topologyhelpers.py   | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py b/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
index dd4fdf357..cdf207705 100644
--- a/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
+++ b/openfe/protocols/openmm_rfe/_rfe_utils/topologyhelpers.py
@@ -470,7 +470,7 @@ def pick_H(i, j, x, y) -> int:
     # there are two reasons constraints would invalidate a mapping entry
     # 1) length of constraint changed (but both constrained)
     # 2) constraint removed to harmonic bond (only one constrained)
-    to_del = set()
+    to_del = []
     for (i, j), l_old in old_constraints.items():
         x, y = old_to_new_atom_map[i], old_to_new_atom_map[j]
 
@@ -481,12 +481,12 @@ def pick_H(i, j, x, y) -> int:
                 l_new = new_constraints.pop((y, x))
             except KeyError:
                 # type 2) constraint doesn't exist in new system
-                to_del.add(pick_H(i, j, x, y))
+                to_del.append(pick_H(i, j, x, y))
                 continue
 
         # type 1) constraint length changed
         if l_old != l_new:
-            to_del.add(pick_H(i, j, x, y))
+            to_del.append(pick_H(i, j, x, y))
 
     # iterate over new_constraints (we were .popping items out)
     # (if any left these are type 2))
@@ -496,8 +496,9 @@ def pick_H(i, j, x, y) -> int:
         for x, y in new_constraints:
             i, j = new_to_old[x], new_to_old[y]
 
-            to_del.add(pick_H(i, j, x, y))
+            to_del.append(pick_H(i, j, x, y))
 
+    # to_del = set(to_del) # Making a set to avoid repeated indices (TODO: TEST!!)
     for idx in to_del:
         del no_const_old_to_new_atom_map[idx]
 

From 018721f8adef3c8e86d4d14b04abd942139ae55e Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Iv=C3=A1n=20Pulido?=
 <2949729+ijpulidos@users.noreply.github.com>
Date: Tue, 23 Sep 2025 23:55:34 -0400
Subject: [PATCH 11/11] hot fix type annotations

---
 openfe/protocols/openmm_utils/omm_settings.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/openfe/protocols/openmm_utils/omm_settings.py b/openfe/protocols/openmm_utils/omm_settings.py
index 8e2e5c28b..5f51d78b3 100644
--- a/openfe/protocols/openmm_utils/omm_settings.py
+++ b/openfe/protocols/openmm_utils/omm_settings.py
@@ -342,17 +342,17 @@ class Config:
     """Size of the simulation timestep. Default 4.0 * unit.femtosecond."""
     langevin_collision_rate: FloatQuantity['1/picosecond'] = 1.0 / unit.picosecond
     """Collision frequency. Default 1.0 / unit.pisecond."""
-    reassign_velocities = False
+    reassign_velocities: bool = False
     """
     If ``True``, velocities are reassigned from the Maxwell-Boltzmann
     distribution at the beginning of each MC move. Default ``False``.
     """
-    n_restart_attempts = 20
+    n_restart_attempts: int = 20
     """
     Number of attempts to restart from Context if there are NaNs in the
     energies after integration. Default 20.
     """
-    constraint_tolerance = 1e-06
+    constraint_tolerance: float = 1e-06
     """Tolerance for the constraint solver. Default 1e-6."""
     barostat_frequency: FloatQuantity['timestep'] = 25.0 * unit.timestep  # todo: IntQuantity
     """