From c2f08f5ae641b4f02b2eaca2429650f0f7402da9 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 13 Aug 2024 09:38:53 +0100 Subject: [PATCH 1/3] More robust search for existing ions. --- src/somd2/runner/_runner.py | 93 +++++++++++++++++++++++++++++++------ 1 file changed, 78 insertions(+), 15 deletions(-) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index e51ff1c..cf8bb2c 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -432,29 +432,86 @@ def _create_alchemical_ions(system, charge_diff): # Create the ions. for water in waters: - # Create an ion. This is to adjust the charge of the perturbed state - # to match that of the reference. + # Flag to indicate whether we need to reverse the alchemical ion + # perturbation, i.e. ion to water, rather than water to ion. + is_reverse = False + + # Create an ion to keep the charge constant throughout the + # perturbation. if charge_diff > 0: # Try to find a free chlorine ion so that we match parameters. try: - ion = system["element Cl"][0].molecule() - assert ion.num_atoms() == 1 + has_ion = False + ions = system["element Cl"].molecules() + for ion in ions: + if ion.num_atoms() == 1: + has_ion = True + break + + # If there isn't an ion, then try searching for a free sodium ion. + if not has_ion: + ions = system["element Na"].molecules() + for ion in ions: + if ion.num_atoms() == 1: + has_ion = True + is_reverse = True + break + + # If not found, create one using a template. + if not has_ion: + ion = _createChlorineIon( + water["element O"].coordinates(), model + ) + # If not found, create one using a template. except: ion = _createChlorineIon(water["element O"].coordinates(), model) - ion_str = "Cl-" + + # Create the ion string. + if is_reverse: + ion_str = "Na+" + else: + ion_str = "Cl-" + else: # Try to find a free sodium ion so that we match parameters. try: - ion = system["element Na"][0].molecule() - assert ion.num_atoms() == 1 + has_ion = False + ions = system["element Na"].molecules() + for ion in ions: + if ion.num_atoms() == 1: + has_ion = True + break + + # If there isn't an ion, then try searching for a free chlorine ion. + if not has_ion: + ions = system["element Cl"].molecules() + for ion in ions: + if ion.num_atoms() == 1: + has_ion = True + is_reverse = True + break + + # If not found, create one using a template. + if not has_ion: + ion = _createSodiumIon(water["element O"].coordinates(), model) + # If not found, create one using a template. except: ion = _createSodiumIon(water["element O"].coordinates(), model) - ion_str = "Na+" - # Create a perturbable molecule: water --> ion. - merged = _morph.merge(water, ion, map={"as_new_molecule": False}) + # Create the ion string. + if is_reverse: + ion_str = "Cl-" + else: + ion_str = "Na+" + + # Create an alchemical ion: ion --> water. + if is_reverse: + merged = _morph.merge(ion, water, map={"as_new_molecule": False}) + # Create an alchemical ion: water --> ion. + else: + merged = _morph.merge(water, ion, map={"as_new_molecule": False}) # Update the system. system.update(merged) @@ -462,11 +519,17 @@ def _create_alchemical_ions(system, charge_diff): # Get the index of the perturbed water. index = numbers.index(water.number()) - # Log that the water was perturbed. - _logger.info( - f"Water at molecule index {index} will be perturbed to " - f"{ion_str} to keep charge constant." - ) + # Log that we are adding an alchemical ion. + if is_reverse: + _logger.info( + f"Water at molecule index {index} will be perturbed from a " + f"{ion_str} ion to keep charge constant." + ) + else: + _logger.info( + f"Water at molecule index {index} will be perturbed to a " + f"{ion_str} ion to keep charge constant." + ) return system From c06057777b6993bfca0a66ec05543f9ab1914692 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 13 Aug 2024 14:27:59 +0100 Subject: [PATCH 2/3] Add alchemical ion test system with existing ions. --- tests/conftest.py | 7 +++++++ tests/runner/test_alchemical_ions.py | 6 ++++-- 2 files changed, 11 insertions(+), 2 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index d723056..b1b1968 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -14,3 +14,10 @@ def ethane_methanol_hmr(): mols = sr.load(sr.expand(sr.tutorial_url, "merged_molecule_hmr.s3")) mols = sr.morph.link_to_reference(mols) return mols + + +@pytest.fixture(scope="session") +def ethane_methanol_ions(): + mols = sr.load(sr.expand(sr.tutorial_url, "merged_molecule_ions.s3")) + mols = sr.morph.link_to_reference(mols) + return mols diff --git a/tests/runner/test_alchemical_ions.py b/tests/runner/test_alchemical_ions.py index 47e1b07..427262e 100644 --- a/tests/runner/test_alchemical_ions.py +++ b/tests/runner/test_alchemical_ions.py @@ -1,13 +1,15 @@ import math +import pytest from somd2.runner import Runner -def test_alchemical_ions(ethane_methanol): +@pytest.mark.parametrize("mols", ["ethane_methanol", "ethane_methanol_ions"]) +def test_alchemical_ions(mols, request): """Ensure that alchemical ions are added correctly.""" # Clone the system. - mols = ethane_methanol.clone() + mols = request.getfixturevalue(mols).clone() # Add 10 Cl- ions. new_mols = Runner._create_alchemical_ions(mols, 10) From 971421c72743d0c3e32fbbfa49778a1b1b54e143 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Tue, 13 Aug 2024 15:42:58 +0100 Subject: [PATCH 3/3] Add debug statement to show if ion was from system or template. --- src/somd2/runner/_runner.py | 8 ++++++++ 1 file changed, 8 insertions(+) diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index cf8bb2c..fe0c24e 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -446,6 +446,7 @@ def _create_alchemical_ions(system, charge_diff): for ion in ions: if ion.num_atoms() == 1: has_ion = True + _logger.debug("Found Cl- ion in system.") break # If there isn't an ion, then try searching for a free sodium ion. @@ -455,16 +456,19 @@ def _create_alchemical_ions(system, charge_diff): if ion.num_atoms() == 1: has_ion = True is_reverse = True + _logger.debug("Found Na+ ion in system.") break # If not found, create one using a template. if not has_ion: + _logger.debug(f"Creating Cl- ion from {model} water template.") ion = _createChlorineIon( water["element O"].coordinates(), model ) # If not found, create one using a template. except: + _logger.debug(f"Creating Cl- ion from {model} water template.") ion = _createChlorineIon(water["element O"].coordinates(), model) # Create the ion string. @@ -481,6 +485,7 @@ def _create_alchemical_ions(system, charge_diff): for ion in ions: if ion.num_atoms() == 1: has_ion = True + _logger.debug("Found Na+ ion in system.") break # If there isn't an ion, then try searching for a free chlorine ion. @@ -490,14 +495,17 @@ def _create_alchemical_ions(system, charge_diff): if ion.num_atoms() == 1: has_ion = True is_reverse = True + _logger.debug("Found Cl- ion in system.") break # If not found, create one using a template. if not has_ion: + _logger.debug(f"Creating Na+ ion from {model} water template.") ion = _createSodiumIon(water["element O"].coordinates(), model) # If not found, create one using a template. except: + _logger.debug(f"Creating Na+ ion from {model} water template.") ion = _createSodiumIon(water["element O"].coordinates(), model) # Create the ion string.