Solution: fix pH charge balancing bug and possible water ion imbalance

CHANGELOG.md

@@ -5,6 +5,20 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [1.1.3] - 2024-07-28
+
+### Fixed
+
+- `Solution`: Fix a bug in which setting `balance_charge` to `pH` could result in
+  negative concentration errors when charge balancing or after `equilibrate` was
+  called. `Solution` now correctly enforces the ion product of water (Kw=1e-14)
+  whenever adjusting the amounts of H+ or OH- for charge balancing.
+
+### Added
+
+- `Solution._adjust_charge_balance`: Added a privat helper method to consolidate charge
+  balancing code used in `__init__` and `equilibrate`.
+
 ## [1.1.2] - 2024-07-28
 
 ### Fixed

src/pyEQL/engines.py

@@ -687,20 +687,8 @@ def equilibrate(self, solution: "Solution") -> None:
             )
 
         # re-adjust charge balance for any missing species
-        # note that if balance_charge is set, it will have been passed to PHREEQC, so we only need to adjust
-        # for any missing species here.
-        charge_adjust = 0
-        for s in missing_species:
-            charge_adjust += -1 * solution.get_amount(s, "eq").magnitude
-        if charge_adjust != 0:
-            logger.warning(
-                "After equilibration, the charge balance of the solution was not electroneutral."
-                f" {charge_adjust} eq of charge were added via {solution._cb_species}"
-            )
-
-            if solution.balance_charge is not None:
-                z = solution.get_property(solution._cb_species, "charge")
-                solution.add_amount(solution._cb_species, f"{charge_adjust/z} mol")
+        # note that if balance_charge is set, it will have been passed to PHREEQC, so the only reason to re-adjust charge balance here is to account for any missing species.
+        solution._adjust_charge_balance()
 
         # rescale the solvent mass to ensure the total mass of solution does not change
         # this is important because PHREEQC and the pyEQL database may use slightly different molecular

src/pyEQL/solution.py

@@ -35,6 +35,7 @@
 EQUIV_WT_CACO3 = ureg.Quantity(100.09 / 2, "g/mol")
 # string to denote unknown oxidation states
 UNKNOWN_OXI_STATE = "unk"
+K_W = 1e-14  # ion product of water at 25 degC
 
 
 class Solution(MSONable):
@@ -242,7 +243,7 @@
 
         # set the pH with H+ and OH-
         self.add_solute("H+", str(10 ** (-1 * pH)) + "mol/L")
-        self.add_solute("OH-", str(10 ** (-1 * (14 - pH))) + "mol/L")
+        self.add_solute("OH-", str(K_W / (10 ** (-1 * pH))) + "mol/L")
 
         # populate the other solutes
         self._solutes = solutes
@@ -288,17 +289,7 @@
                 )
 
         # adjust charge balance, if necessary
-        if not np.isclose(cb, 0, atol=1e-8) and self.balance_charge is not None:
-            balanced = False
-            self.logger.info(
-                f"Solution is not electroneutral (C.B. = {cb} eq/L). Adding {self._cb_species} to compensate."
-            )
-            z = self.get_property(self._cb_species, "charge")
-            self.components[self._cb_species] += -1 * cb / z * self.volume.to("L").magnitude
-            if np.isclose(self.charge_balance, 0, atol=1e-8):
-                balanced = True
-            if not balanced:
-                warnings.warn(f"Unable to balance charge using species {self._cb_species}")
+        self._adjust_charge_balance()
 
     @property
     def mass(self) -> Quantity:
@@ -2293,6 +2284,57 @@
 
         return distance.to("nm")
 
+    def _adjust_charge_balance(self, atol=1e-8) -> None:
+        """Helper method to adjust the charge balance of the Solution."""
+        cb = self.charge_balance
+        if not np.isclose(cb, 0, atol=atol):
+            self.logger.info(f"Solution is not electroneutral (C.B. = {cb} eq/L).")
+            if self.balance_charge is None:
+                # Nothing to do.
+                self.logger.info("balance_charge is None, so no charge balancing will be performed.")
+                return
+
+            self.logger.info(
+                f"Solution is not electroneutral (C.B. = {cb} eq/L). Adjusting {self._cb_species} to compensate."
+            )
+
+            if self.balance_charge == "pH":
+                # the charge imbalance associated with the H+ / OH- system can be expressed
+                # as ([H+] - [OH-]) or ([H+] - K_W/[H+]). If we adjust H+, we also have to
+                # adjust OH- to maintain water equilibrium.
+                C_hplus = self.get_amount("H+", "mol/L").magnitude
+                start_imbalance = C_hplus - K_W / C_hplus
+                new_imbalance = start_imbalance - cb
+                # calculate the new concentration of H+ that will balance the charge
+                # solve H^2 - new_imbalance H - K_W  = 0, so a=1, b=-new_imbalance, c=-K_W
+                # check b^2 - 4ac; are there any real roots?
+                if new_imbalance**2 - 4 * 1 * K_W < 0:
+                    self.logger.error("Cannot balance charge by adjusting pH. The imbalance is too large.")
+                    return
+                new_hplus = max(
+                    [
+                        (new_imbalance + np.sqrt(new_imbalance**2 + 4 * 1 * K_W)) / 2,
+                        (new_imbalance - np.sqrt(new_imbalance**2 + 4 * 1 * K_W)) / 2,
+                    ]
+                )
+                self.set_amount("H+", f"{new_hplus} mol/L")
+                self.set_amount("OH-", f"{K_W/new_hplus} mol/L")
+                assert np.isclose(self.charge_balance, 0, atol=atol), f"{self.charge_balance}"
+                return
+
+            z = self.get_property(self._cb_species, "charge")
+            try:
+                self.add_amount(self._cb_species, f"{-1*cb/z} mol")
+                return
+            except ValueError:
+                # if the concentration is negative, it must mean there is not enough present.
+                # remove everything that's present and log an error.
+                self.components[self._cb_species] = 0
+                self.logger.error(
+                    f"There is not enough {self._cb_species} present to balance the charge. Try a different species."
+                )
+                return
+
     def _update_volume(self):
         """Recalculate the solution volume based on composition."""
         self._volume = self._get_solvent_volume() + self._get_solute_volume()

tests/test_solution.py

@@ -202,8 +202,8 @@ def test_init_engines():
 
 
 def test_component_subsets(s2):
-    assert s2.cations == {"Na[+1]": 8, "H[+1]": 2e-7}
-    assert s2.anions == {"Cl[-1]": 8, "OH[-1]": 2e-7}
+    assert list(s2.cations.keys()) == ["Na[+1]", "H[+1]"]
+    assert list(s2.anions.keys()) == ["Cl[-1]", "OH[-1]"]
     assert list(s2.neutrals.keys()) == ["H2O(aq)"]
 
 
@@ -285,6 +285,32 @@ def test_charge_balance(s3, s5, s5_pH, s6, s6_Ca):
     assert s._cb_species == "Cl[-1]"
     assert np.isclose(s.charge_balance, 0, atol=1e-8)
 
+    # check 'pH' when the solution needs to be made more POSITIVE
+    s = Solution({"Na+": "2 mM", "Cl-": "1 mM"}, balance_charge="pH", pH=4)
+    assert s.balance_charge == "pH"
+    assert s._cb_species == "H[+1]"
+    assert np.isclose(s.charge_balance, 0, atol=1e-8)
+    assert s.pH > 4
+    s.equilibrate()
+    assert s.balance_charge == "pH"
+    assert s._cb_species == "H[+1]"
+    assert np.isclose(s.charge_balance, 0, atol=1e-8)
+
+    # check 'pH' when the imbalance is extreme
+    s = Solution({"Na+": "2 mM", "Cl-": "1 M"}, balance_charge="pH", pH=4)
+    assert s.balance_charge == "pH"
+    assert s._cb_species == "H[+1]"
+    assert np.isclose(s.charge_balance, 0, atol=1e-8)
+    assert np.isclose(s.pH, 0, atol=0.1)
+    s.equilibrate()
+    assert s.balance_charge == "pH"
+    assert s._cb_species == "H[+1]"
+    assert np.isclose(s.charge_balance, 0, atol=1e-8)
+
+    # check warning when there isn't enough to balance
+    s = Solution({"Na+": "1 M", "K+": "2 mM", "Cl-": "2 mM"}, balance_charge="K+")
+    assert s.get_amount("K+", "mol/L") == 0
+
     # check "auto" with an electroneutral solution
     s = Solution({"Na+": "2 mM", "Cl-": "2 mM"}, balance_charge="auto")
     assert s.balance_charge == "auto"
@@ -405,16 +431,16 @@ def test_components_by_element(s1, s2):
     assert s1.get_components_by_element() == {
         "H(1.0)": [
             "H2O(aq)",
-            "H[+1]",
             "OH[-1]",
+            "H[+1]",
         ],
         "O(-2.0)": ["H2O(aq)", "OH[-1]"],
     }
     assert s2.get_components_by_element() == {
         "H(1.0)": [
             "H2O(aq)",
-            "H[+1]",
             "OH[-1]",
+            "H[+1]",
         ],
         "O(-2.0)": ["H2O(aq)", "OH[-1]"],
         "Na(1.0)": ["Na[+1]"],
@@ -498,8 +524,8 @@ def test_equilibrate(s1, s2, s5_pH):
     orig_solv_mass = s5_pH.solvent_mass.magnitude
     set(s5_pH.components.keys())
     s5_pH.equilibrate()
-    assert np.isclose(s5_pH.get_total_amount("Ca", "mol").magnitude, 0.001)
-    assert np.isclose(s5_pH.get_total_amount("C(4)", "mol").magnitude, 0.001)
+    assert np.isclose(s5_pH.get_total_amount("Ca", "mol").magnitude, 0.001, atol=1e-7)
+    assert np.isclose(s5_pH.get_total_amount("C(4)", "mol").magnitude, 0.001, atol=1e-7)
     # due to the large pH shift, water mass and density need not be perfectly conserved
     assert np.isclose(s5_pH.solvent_mass.magnitude, orig_solv_mass, atol=1e-3)
     assert np.isclose(s5_pH.density.magnitude, orig_density, atol=1e-3)