Allow user to specify counterion.

Author: lohedges

src/BioSimSpace/Sandpit/Exscientia/Solvent/_ions.py

@@ -40,12 +40,8 @@
     "rubidium": ("RB", 1),
     "cs": ("CS", 1),
     "cesium": ("CS", 1),
-    "f": ("F", -1),
-    "fluoride": ("F", -1),
     "cl": ("CL", -1),
     "chloride": ("CL", -1),
-    "br": ("BR", -1),
-    "bromide": ("BR", -1),
     "mg": ("MG", 2),
     "magnesium": ("MG", 2),
     "ca": ("CA", 2),
@@ -78,6 +74,7 @@ def addIons(
     ion_conc=0,
     is_neutral=True,
     preserved_waters=None,
+    counter_ion=None,
     work_dir=None,
     property_map={},
 ):
@@ -128,6 +125,15 @@ def addIons(
         that genion cannot select them for replacement. Useful for
         protecting crystallographic or binding-site water molecules.
 
+    counter_ion : str
+        The name of the ion to use for the opposite charge slot in the
+        genion call. By default genion uses Na\\ :sup:`+` as the counter
+        cation and Cl\\ :sup:`-` as the counter anion. Use this parameter
+        to override the counter-ion species, e.g. ``"br"`` when adding a
+        cation. The counter-ion must have the opposite charge sign to
+        ``ion``. Requires ``is_neutral=True`` (otherwise no counter-ions
+        are added and the parameter has no effect).
+
     work_dir : str
         The working directory for the process.
 
@@ -235,6 +241,35 @@ def addIons(
                     "'preserved_waters' items must be int indices or Molecule objects."
                 )
 
+    # Validate and resolve counter_ion.
+    counter_ion_name = None
+    counter_ion_charge = None
+    if counter_ion is not None:
+        if not isinstance(counter_ion, str):
+            raise TypeError("'counter_ion' must be of type 'str'")
+        counter_key = counter_ion.strip().lower()
+        if counter_key not in _ion_map:
+            raise ValueError(
+                f"Unsupported counter_ion '{counter_ion}'. Supported ions are: {ions()}"
+            )
+        counter_ion_name, counter_ion_charge = _ion_map[counter_key]
+        # The counter-ion must have the opposite charge sign to the requested ion.
+        if ion_charge > 0 and counter_ion_charge >= 0:
+            raise ValueError(
+                f"'counter_ion' must be negatively charged when 'ion' is positively "
+                f"charged. '{counter_ion}' has charge {counter_ion_charge:+d}."
+            )
+        if ion_charge < 0 and counter_ion_charge <= 0:
+            raise ValueError(
+                f"'counter_ion' must be positively charged when 'ion' is negatively "
+                f"charged. '{counter_ion}' has charge {counter_ion_charge:+d}."
+            )
+        if not is_neutral:
+            raise ValueError(
+                "'counter_ion' has no effect when 'is_neutral=False'. Set "
+                "'is_neutral=True' to enable counter-ion addition."
+            )
+
     if work_dir is not None and not isinstance(work_dir, str):
         raise TypeError("'work_dir' must be of type 'str'")
 
@@ -280,6 +315,8 @@ def addIons(
         num_ions,
         is_neutral,
         preserved_mols,
+        counter_ion_name,
+        counter_ion_charge,
         work_dir,
         property_map,
     )
@@ -292,6 +329,8 @@ def _add_ions(
     num_ions,
     is_neutral,
     preserved_mols=None,
+    counter_ion_name=None,
+    counter_ion_charge=None,
     work_dir=None,
     property_map={},
 ):
@@ -322,6 +361,13 @@ def _add_ions(
         between the non-water solute and the new water+ions in the result,
         mirroring the ordering used by ``_solvate`` for crystal waters.
 
+    counter_ion_name : str
+        The GROMACS residue name of the counter-ion (e.g. ``"BR"``), or
+        ``None`` to use genion's default (NA/CL).
+
+    counter_ion_charge : int
+        The integer charge of the counter-ion, or ``None``.
+
     work_dir : str
         The working directory for the process.
 
@@ -481,22 +527,29 @@ def _add_ions(
         # Build the genion command.
         # genion supports one positive ion type (-pname/-pq/-np) and one
         # negative ion type (-nname/-nq/-nn). We use the appropriate slot
-        # for the requested ion. When is_neutral=True, genion adds the
-        # default NA/CL as counter-ions to bring the total charge to zero.
+        # for the requested ion. When is_neutral=True, genion adds counter-ions
+        # to bring the total charge to zero; counter_ion_name overrides the
+        # default counter-ion species (NA for cations, CL for anions).
         if ion_charge > 0:
             command = (
                 "%s genion -s ions.tpr -o ions_out.gro -p system.top"
                 " -pname %s -pq %d" % (_gmx_exe, ion_name, ion_charge)
             )
             if num_ions > 0:
                 command += " -np %d" % num_ions
+            # Override the default Cl- counter-ion if requested.
+            if counter_ion_name is not None:
+                command += " -nname %s -nq %d" % (counter_ion_name, counter_ion_charge)
         else:
             command = (
                 "%s genion -s ions.tpr -o ions_out.gro -p system.top"
                 " -nname %s -nq %d" % (_gmx_exe, ion_name, ion_charge)
             )
             if num_ions > 0:
                 command += " -nn %d" % num_ions
+            # Override the default Na+ counter-ion if requested.
+            if counter_ion_name is not None:
+                command += " -pname %s -pq %d" % (counter_ion_name, counter_ion_charge)
 
         if is_neutral:
             command += " -neutral"

src/BioSimSpace/Solvent/_ions.py

@@ -40,12 +40,8 @@
     "rubidium": ("RB", 1),
     "cs": ("CS", 1),
     "cesium": ("CS", 1),
-    "f": ("F", -1),
-    "fluoride": ("F", -1),
     "cl": ("CL", -1),
     "chloride": ("CL", -1),
-    "br": ("BR", -1),
-    "bromide": ("BR", -1),
     "mg": ("MG", 2),
     "magnesium": ("MG", 2),
     "ca": ("CA", 2),
@@ -78,6 +74,7 @@ def addIons(
     ion_conc=0,
     is_neutral=True,
     preserved_waters=None,
+    counter_ion=None,
     work_dir=None,
     property_map={},
 ):
@@ -128,6 +125,15 @@ def addIons(
         that genion cannot select them for replacement. Useful for
         protecting crystallographic or binding-site water molecules.
 
+    counter_ion : str
+        The name of the ion to use for the opposite charge slot in the
+        genion call. By default genion uses Na\\ :sup:`+` as the counter
+        cation and Cl\\ :sup:`-` as the counter anion. Use this parameter
+        to override the counter-ion species, e.g. ``"br"`` when adding a
+        cation. The counter-ion must have the opposite charge sign to
+        ``ion``. Requires ``is_neutral=True`` (otherwise no counter-ions
+        are added and the parameter has no effect).
+
     work_dir : str
         The working directory for the process.
 
@@ -235,6 +241,35 @@ def addIons(
                     "'preserved_waters' items must be int indices or Molecule objects."
                 )
 
+    # Validate and resolve counter_ion.
+    counter_ion_name = None
+    counter_ion_charge = None
+    if counter_ion is not None:
+        if not isinstance(counter_ion, str):
+            raise TypeError("'counter_ion' must be of type 'str'")
+        counter_key = counter_ion.strip().lower()
+        if counter_key not in _ion_map:
+            raise ValueError(
+                f"Unsupported counter_ion '{counter_ion}'. Supported ions are: {ions()}"
+            )
+        counter_ion_name, counter_ion_charge = _ion_map[counter_key]
+        # The counter-ion must have the opposite charge sign to the requested ion.
+        if ion_charge > 0 and counter_ion_charge >= 0:
+            raise ValueError(
+                f"'counter_ion' must be negatively charged when 'ion' is positively "
+                f"charged. '{counter_ion}' has charge {counter_ion_charge:+d}."
+            )
+        if ion_charge < 0 and counter_ion_charge <= 0:
+            raise ValueError(
+                f"'counter_ion' must be positively charged when 'ion' is negatively "
+                f"charged. '{counter_ion}' has charge {counter_ion_charge:+d}."
+            )
+        if not is_neutral:
+            raise ValueError(
+                "'counter_ion' has no effect when 'is_neutral=False'. Set "
+                "'is_neutral=True' to enable counter-ion addition."
+            )
+
     if work_dir is not None and not isinstance(work_dir, str):
         raise TypeError("'work_dir' must be of type 'str'")
 
@@ -280,6 +315,8 @@ def addIons(
         num_ions,
         is_neutral,
         preserved_mols,
+        counter_ion_name,
+        counter_ion_charge,
         work_dir,
         property_map,
     )
@@ -292,6 +329,8 @@ def _add_ions(
     num_ions,
     is_neutral,
     preserved_mols=None,
+    counter_ion_name=None,
+    counter_ion_charge=None,
     work_dir=None,
     property_map={},
 ):
@@ -322,6 +361,13 @@ def _add_ions(
         between the non-water solute and the new water+ions in the result,
         mirroring the ordering used by ``_solvate`` for crystal waters.
 
+    counter_ion_name : str
+        The GROMACS residue name of the counter-ion (e.g. ``"BR"``), or
+        ``None`` to use genion's default (NA/CL).
+
+    counter_ion_charge : int
+        The integer charge of the counter-ion, or ``None``.
+
     work_dir : str
         The working directory for the process.
 
@@ -481,22 +527,29 @@ def _add_ions(
         # Build the genion command.
         # genion supports one positive ion type (-pname/-pq/-np) and one
         # negative ion type (-nname/-nq/-nn). We use the appropriate slot
-        # for the requested ion. When is_neutral=True, genion adds the
-        # default NA/CL as counter-ions to bring the total charge to zero.
+        # for the requested ion. When is_neutral=True, genion adds counter-ions
+        # to bring the total charge to zero; counter_ion_name overrides the
+        # default counter-ion species (NA for cations, CL for anions).
         if ion_charge > 0:
             command = (
                 "%s genion -s ions.tpr -o ions_out.gro -p system.top"
                 " -pname %s -pq %d" % (_gmx_exe, ion_name, ion_charge)
             )
             if num_ions > 0:
                 command += " -np %d" % num_ions
+            # Override the default Cl- counter-ion if requested.
+            if counter_ion_name is not None:
+                command += " -nname %s -nq %d" % (counter_ion_name, counter_ion_charge)
         else:
             command = (
                 "%s genion -s ions.tpr -o ions_out.gro -p system.top"
                 " -nname %s -nq %d" % (_gmx_exe, ion_name, ion_charge)
             )
             if num_ions > 0:
                 command += " -nn %d" % num_ions
+            # Override the default Na+ counter-ion if requested.
+            if counter_ion_name is not None:
+                command += " -pname %s -pq %d" % (counter_ion_name, counter_ion_charge)
 
         if is_neutral:
             command += " -neutral"

tests/Sandpit/Exscientia/Solvent/test_ions.py

@@ -38,7 +38,7 @@ def test_ions():
     assert isinstance(ion_list, list)
     assert len(ion_list) > 0
     # Check a representative selection of expected ions are present.
-    for ion in ["br", "ca", "cl", "cs", "f", "k", "li", "mg", "na", "rb", "zn"]:
+    for ion in ["ca", "cl", "cs", "k", "li", "mg", "na", "rb", "zn"]:
         assert ion in ion_list
 
 
@@ -163,3 +163,39 @@ def test_add_ions_with_existing_ions(solvated_system_with_ions):
 
     assert num_na_after == num_na_before
     assert num_cl_after == num_cl_before
+
+
+@pytest.mark.skipif(not has_gromacs, reason="Requires GROMACS to be installed")
+def test_add_ions_counter_ion(solvated_system):
+    """
+    Test that counter_ion overrides the default Na+ counter-ion.
+
+    Adding CL- with counter_ion="k" should result in K+ ions being added
+    for neutralisation instead of the default NA+.
+    """
+    num_ions = 2
+
+    result = BSS.Solvent.addIons(
+        solvated_system, "cl", num_ions=num_ions, is_neutral=True, counter_ion="k"
+    )
+
+    # Check that CL ions were added.
+    try:
+        cl_molecules = result.search("resname CL").molecules()
+    except Exception:
+        pytest.fail("No CL ions found in the result system.")
+    assert len(cl_molecules) == num_ions
+
+    # K+ counter-ions should be present (2 Cl- → -2 charge → 2 K+).
+    try:
+        k_molecules = result.search("resname K").molecules()
+    except Exception:
+        pytest.fail("No K counter-ions found in the result system.")
+    assert len(k_molecules) == num_ions
+
+    # The default NA counter-ion should NOT be present.
+    try:
+        na_count = len(result.search("resname NA").molecules())
+    except Exception:
+        na_count = 0
+    assert na_count == 0

tests/Solvent/test_ions.py

@@ -38,7 +38,7 @@ def test_ions():
     assert isinstance(ion_list, list)
     assert len(ion_list) > 0
     # Check a representative selection of expected ions are present.
-    for ion in ["br", "ca", "cl", "cs", "f", "k", "li", "mg", "na", "rb", "zn"]:
+    for ion in ["ca", "cl", "cs", "k", "li", "mg", "na", "rb", "zn"]:
         assert ion in ion_list
 
 
@@ -163,3 +163,39 @@ def test_add_ions_with_existing_ions(solvated_system_with_ions):
 
     assert num_na_after == num_na_before
     assert num_cl_after == num_cl_before
+
+
+@pytest.mark.skipif(not has_gromacs, reason="Requires GROMACS to be installed")
+def test_add_ions_counter_ion(solvated_system):
+    """
+    Test that counter_ion overrides the default Na+ counter-ion.
+
+    Adding CL- with counter_ion="k" should result in K+ ions being added
+    for neutralisation instead of the default NA+.
+    """
+    num_ions = 2
+
+    result = BSS.Solvent.addIons(
+        solvated_system, "cl", num_ions=num_ions, is_neutral=True, counter_ion="k"
+    )
+
+    # Check that CL ions were added.
+    try:
+        cl_molecules = result.search("resname CL").molecules()
+    except Exception:
+        pytest.fail("No CL ions found in the result system.")
+    assert len(cl_molecules) == num_ions
+
+    # K+ counter-ions should be present (2 Cl- → -2 charge → 2 K+).
+    try:
+        k_molecules = result.search("resname K").molecules()
+    except Exception:
+        pytest.fail("No K counter-ions found in the result system.")
+    assert len(k_molecules) == num_ions
+
+    # The default NA counter-ion should NOT be present.
+    try:
+        na_count = len(result.search("resname NA").molecules())
+    except Exception:
+        na_count = 0
+    assert na_count == 0