diff --git a/src/somd2/runner/_runner.py b/src/somd2/runner/_runner.py index e51ff1c..fe0c24e 100644 --- a/src/somd2/runner/_runner.py +++ b/src/somd2/runner/_runner.py @@ -432,29 +432,94 @@ 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 + _logger.debug("Found Cl- ion in system.") + 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 + _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) - 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 + _logger.debug("Found Na+ ion in system.") + 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 + _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) - 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 +527,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 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)