Merge pull request #153 from KingsburyLab/bugfix

Solution: fix auto charge balance bug

Author: rkingsbury

CHANGELOG.md

@@ -5,6 +5,14 @@ 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.2] - 2024-07-28
+
+### Fixed
+
+- `Solution`: Fix a bug in which setting `balance_charge` to `auto` when the initial
+  composition was electroneutral would cause errors and/or improper charge balancing
+  after `equilibrate` was called.
+
 ## [1.1.1] - 2024-07-27
 
 ### Fixed

src/pyEQL/engines.py

@@ -229,10 +229,7 @@ def _setup_ppsol(self, solution: "Solution") -> None:
                 d[key] = str(mol / solv_mass)
 
             # tell PHREEQC which species to use for charge balance
-            if (
-                solution.balance_charge is not None
-                and solution.balance_charge in solution.get_components_by_element()[el]
-            ):
+            if solution.balance_charge is not None and solution._cb_species in solution.get_components_by_element()[el]:
                 d[key] += " charge"
 
         # create the PHREEQC solution object
@@ -698,18 +695,12 @@ def equilibrate(self, solution: "Solution") -> None:
         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.balance_charge}"
+                f" {charge_adjust} eq of charge were added via {solution._cb_species}"
             )
 
-            if solution.balance_charge is None:
-                pass
-            elif solution.balance_charge == "pH":
-                solution.components["H+"] += charge_adjust
-            elif solution.balance_charge == "pE":
-                raise NotImplementedError
-            else:
-                z = solution.get_property(solution.balance_charge, "charge")
-                solution.add_amount(solution.balance_charge, f"{charge_adjust/z} mol")
+            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")
 
         # 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

@@ -261,39 +261,44 @@ def __init__(
             for item in self._solutes:
                 self.add_solute(*item)
 
-        # adjust the charge balance, if necessary
+        # determine the species that will be used for charge balancing, when needed.
+        # this is necessary to do even if the composition is already electroneutral,
+        # because the appropriate species also needs to be passed to equilibrate
+        # to keep from distorting the charge balance.
         cb = self.charge_balance
+        if self.balance_charge is None:
+            self._cb_species = None
+        elif self.balance_charge == "pH":
+            self._cb_species = "H[+1]"
+        elif self.balance_charge == "pE":
+            raise NotImplementedError("Balancing charge via redox (pE) is not yet implemented!")
+        elif self.balance_charge == "auto":
+            # add the most abundant ion of the opposite charge
+            if cb <= 0:
+                self._cb_species = max(self.cations, key=self.cations.get)
+            elif cb > 0:
+                self._cb_species = max(self.anions, key=self.anions.get)
+        else:
+            ions = set().union(*[self.cations, self.anions])  # all ions
+            self._cb_species = self.balance_charge
+            if self._cb_species not in ions:
+                raise ValueError(
+                    f"Charge balancing species {self._cb_species} was not found in the solution!. "
+                    f"Species {ions} were found."
+                )
+
+        # 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 {balance_charge} to compensate."
+                f"Solution is not electroneutral (C.B. = {cb} eq/L). Adding {self._cb_species} to compensate."
             )
-            if self.balance_charge == "pH":
-                self.components["H+"] += (
-                    -1 * cb * self.volume.to("L").magnitude
-                )  # if C.B. is negative, we need to add cations. H+ is 1 eq/mol
+            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
-            elif self.balance_charge == "pE":
-                raise NotImplementedError("Balancing charge via redox (pE) is not yet implemented!")
-            else:
-                ions = set().union(*[self.cations, self.anions])  # all ions
-                if self.balance_charge == "auto":
-                    # add the most abundant ion of the opposite charge
-                    if cb <= 0:
-                        self.balance_charge = max(self.cations, key=self.cations.get)
-                    elif cb > 0:
-                        self.balance_charge = max(self.anions, key=self.anions.get)
-                if self.balance_charge not in ions:
-                    raise ValueError(
-                        f"Charge balancing species {self.balance_charge} was not found in the solution!. "
-                        f"Species {ions} were found."
-                    )
-                z = self.get_property(self.balance_charge, "charge")
-                self.components[self.balance_charge] += -1 * cb / z * self.volume.to("L").magnitude
-                balanced = True
-
             if not balanced:
-                warnings.warn(f"Unable to balance charge using species {self.balance_charge}")
+                warnings.warn(f"Unable to balance charge using species {self._cb_species}")
 
     @property
     def mass(self) -> Quantity:

tests/test_solution.py

@@ -268,16 +268,35 @@ def test_charge_balance(s3, s5, s5_pH, s6, s6_Ca):
         volume="1 L",
         balance_charge="auto",
     )
-    assert s.balance_charge == "Na[+1]"
+    assert s.balance_charge == "auto"
+    assert s._cb_species == "Na[+1]"
     assert np.isclose(s.charge_balance, 0, atol=1e-8)
     s.equilibrate()
-    assert s.balance_charge == "Na[+1]"
+    assert s.balance_charge == "auto"
+    assert s._cb_species == "Na[+1]"
+    assert np.isclose(s.charge_balance, 0, atol=1e-8)
 
     s = Solution({"Na+": "2 mM", "Cl-": "1 mM"}, balance_charge="auto")
-    assert s.balance_charge == "Cl[-1]"
+    assert s.balance_charge == "auto"
+    assert s._cb_species == "Cl[-1]"
     assert np.isclose(s.charge_balance, 0, atol=1e-8)
     s.equilibrate()
-    assert s.balance_charge == "Cl[-1]"
+    assert s.balance_charge == "auto"
+    assert s._cb_species == "Cl[-1]"
+    assert np.isclose(s.charge_balance, 0, atol=1e-8)
+
+    # check "auto" with an electroneutral solution
+    s = Solution({"Na+": "2 mM", "Cl-": "2 mM"}, balance_charge="auto")
+    assert s.balance_charge == "auto"
+    assert s._cb_species == "Na[+1]"
+    assert np.isclose(s.charge_balance, 0, atol=1e-8)
+    s.equilibrate()
+    assert s.balance_charge == "auto"
+    assert s._cb_species == "Na[+1]"
+    assert np.isclose(s.charge_balance, 0, atol=1e-8)
+
+    with pytest.raises(ValueError, match=r"Charge balancing species Zr\[\+4\] was not found"):
+        s = Solution({"Na+": "2 mM", "Cl-": "2 mM"}, balance_charge="Zr[+4]")
 
 
 def test_alkalinity_hardness(s3, s5, s6):