From b10e54e7ec1ba7c657a53fa7a6a63dc8c4466f3f Mon Sep 17 00:00:00 2001 From: Koki Shimada Date: Wed, 1 Apr 2026 04:50:58 +0000 Subject: [PATCH 1/2] Use set operations for unique atom detection in HybridTopologyFactory Replace linear-scan loops with set difference operations in _set_mappings for finding unique atoms and environment atom maps. Also use dict comprehension for _env_old_to_new_map construction. --- .../openmm_rfe/_rfe_utils/relative.py | 25 +++++++++---------- 1 file changed, 12 insertions(+), 13 deletions(-) diff --git a/src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py b/src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py index 7297689b1..8fefabf11 100644 --- a/src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py +++ b/src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py @@ -437,22 +437,21 @@ def _set_mappings(self, old_to_new_map, core_old_to_new_map): self._new_to_hybrid_map = {} # Get unique atoms - # old system first - self._unique_old_atoms = [] - for particle_idx in range(self._old_system.getNumParticles()): - if particle_idx not in self._old_to_new_map.keys(): - self._unique_old_atoms.append(particle_idx) + old_mapped = self._old_to_new_map.keys() + self._unique_old_atoms = sorted( + set(range(self._old_system.getNumParticles())) - old_mapped + ) - self._unique_new_atoms = [] - for particle_idx in range(self._new_system.getNumParticles()): - if particle_idx not in self._new_to_old_map.keys(): - self._unique_new_atoms.append(particle_idx) + new_mapped = self._new_to_old_map.keys() + self._unique_new_atoms = sorted( + set(range(self._new_system.getNumParticles())) - new_mapped + ) # Get env atoms (i.e. atoms mapped not in core) - self._env_old_to_new_map = {} - for key, value in old_to_new_map.items(): - if key not in self._core_old_to_new_map.keys(): - self._env_old_to_new_map[key] = value + self._env_old_to_new_map = { + k: v for k, v in old_to_new_map.items() + if k not in self._core_old_to_new_map + } self._env_new_to_old_map = self._invert_dict(self._env_old_to_new_map) From bfd99e44af4b2571fadfc735c45d5962b9bf9beb Mon Sep 17 00:00:00 2001 From: Koki Shimada Date: Wed, 1 Apr 2026 05:52:51 +0000 Subject: [PATCH 2/2] Remove redundant _find_bond_parameters call in _handle_harmonic_bonds Reuse the result already stored in new_bond_parameters instead of calling _find_bond_parameters a second time with the same arguments. --- src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py b/src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py index 8fefabf11..6b9ada0ba 100644 --- a/src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py +++ b/src/openfe/protocols/openmm_rfe/_rfe_utils/relative.py @@ -1220,9 +1220,7 @@ def _handle_harmonic_bonds(self): r0_new = r0_old k_new = 0.0*unit.kilojoule_per_mole/unit.angstrom**2 else: - # TODO - why is this being recalculated? - [index1, index2, r0_new, k_new] = self._find_bond_parameters( - new_system_bond_force, index1_new, index2_new) + [_, _, r0_new, k_new] = new_bond_parameters self._hybrid_system_forces['core_bond_force'].addBond( index1_hybrid, index2_hybrid, [r0_old, k_old, r0_new, k_new])