added a nested parameter to 2 methods in the Solution class

CHANGELOG.md

@@ -5,6 +5,16 @@ 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).
 
+## Unreleased
+
+### Added
+
+- `get_components_by_element`: A new keyword argument `nested` was added to methods `get_components_by_element`
+  and `get_el_amt_dict` of the `Solution` class. It defaults to `False` (no change from prior behavior), but
+  can be set to `True` to return a 2-level dictionary, with the element symbol as the key at the top level, and
+  the valence (float, or "unk" for unknown) as the key at the second level. This should make it easier for future
+  code to calculate the total amount of a given element regardless of its valence. (#284, @vineetbansal)
+
 ## [1.3.2] - 2025-09-15
 
 ### Fixed

src/pyEQL/solution.py

@@ -669,7 +669,7 @@ def viscosity_kinematic(self) -> Quantity:
         else:
             # TODO - fall back to the Jones-Dole model! There are currently no eyring parameters in the database!
             # proceed with the coefficients equal to zero and log a warning
-            self.logger.warning(f"Appropriate viscosity coefficients werenot found. Viscosity will be approximate.")
+            self.logger.warning("Appropriate viscosity coefficients were not found. Viscosity will be approximate.")
             G_123 = G_23 = 0
             x_cat = 0
 
@@ -1107,13 +1107,31 @@ def get_amount(self, solute: str, units: str = "mol/L") -> Quantity:
 
         raise ValueError(f"Unsupported unit {units} specified for get_amount")
 
-    def get_components_by_element(self) -> dict[str, list]:
+    def get_components_by_element(
+        self, nested: bool = False
+    ) -> dict[str, list[str]] | dict[str, dict[float | str, list[str]]]:
         """
         Return a list of all species associated with a given element.
 
-        Elements (keys) are suffixed with their oxidation state in parentheses, e.g.,
+        Args:
+            nested : bool
+                Whether to return a nested dictionary of <element>
+                to <valence> => <list of species> mapping. False by default.
+
+        Returns:
+            A mapping of element to a list of species in the solution.
+
+            If nested is False (default), elements (keys) are suffixed with
+            their oxidation state in parentheses, e.g.,
+
+            {"Na(1.0)":["Na[+1]", "NaOH(aq)"]}
 
-        {"Na(1.0)":["Na[+1]", "NaOH(aq)"]}
+            If nested is True, the dictionary is nested, e.g.,
+
+            {"Na": [{1:["Na[+1]", "NaOH(aq)"]}]}.
+
+            Note that the valence may be a string, assuming the value "unk"
+            denoting an unknown oxidation state.
 
         Species associated with each element are sorted in descending order of the amount
         present (i.e., the first species listed is the most abundant).
@@ -1132,20 +1150,41 @@ def get_components_by_element(self) -> dict[str, list]:
                 except (TypeError, IndexError):
                     self.logger.error(f"No oxidation state found for element {el}. Assigning '{UNKNOWN_OXI_STATE}'")
                     oxi_state = UNKNOWN_OXI_STATE
-                key = f"{el}({oxi_state})"
-                if d.get(key):
-                    d[key].append(s)
+                if d.get(el):
+                    if d[el].get(oxi_state):
+                        d[el][oxi_state].append(s)
+                    else:
+                        d[el][oxi_state] = [s]
                 else:
-                    d[key] = [s]
+                    d[el] = {oxi_state: [s]}
 
-        return d
+        if nested:
+            return d
+        return {f"{el}({val})": species for el, val_dict in d.items() for val, species in val_dict.items()}
 
-    def get_el_amt_dict(self):
+    def get_el_amt_dict(self, nested: bool = False) -> dict[str, float] | dict[str, dict[float | str, float]]:
         """
         Return a dict of Element: amount in mol.
 
-        Elements (keys) are suffixed with their oxidation state in parentheses,
-        e.g. "Fe(2.0)", "Cl(-1.0)".
+        Args:
+            nested : bool
+                Whether to return a nested dictionary of <element>
+                to <valence> => amount mapping. False by default.
+
+        Returns:
+            A mapping of element to its amount in moles in the solution.
+
+            If nested is False (default), elements (keys) are suffixed with
+            their oxidation state in parentheses, e.g.,
+
+            {"Fe(2.0)": 0.354, "Cl(-1.0)": 0.708}
+
+            If nested is True, the dictionary is nested, e.g.,
+
+            {"Fe": {2.0: 0.354}, "Cl": {-1.0: 0.708}}.}
+
+            Note that the valence may be a string, assuming the value "unk"
+            denoting an unknown oxidation state.
         """
         d = {}
         for s, mol in self.components.items():
@@ -1162,13 +1201,17 @@ def get_el_amt_dict(self):
                 except (TypeError, IndexError):
                     self.logger.error(f"No oxidation state found for element {el}. Assigning '{UNKNOWN_OXI_STATE}'")
                     oxi_state = UNKNOWN_OXI_STATE
-                key = f"{el}({oxi_state})"
-                if d.get(key):
-                    d[key] += stoich * mol
+                if d.get(el):
+                    if d[el].get(oxi_state):
+                        d[el][oxi_state] += stoich * mol
+                    else:
+                        d[el][oxi_state] = stoich * mol
                 else:
-                    d[key] = stoich * mol
+                    d[el] = {oxi_state: stoich * mol}
 
-        return d
+        if nested:
+            return d
+        return {f"{el}({val})": amount for el, val_dict in d.items() for val, amount in val_dict.items()}
 
     def get_total_amount(self, element: str, units: str) -> Quantity:
         """

tests/test_solution.py

@@ -118,8 +118,8 @@ def test_empty_solution():
     assert np.isclose(s1.pE, 8.5)
     # it should contain H2O, H+, and OH- species
     assert set(s1.components.keys()) == {"H2O(aq)", "OH[-1]", "H[+1]"}
-    assert np.isclose(s1.density.to('kg/m**3').magnitude, 997.0479, atol=0.1)
-    assert np.isclose(s1.viscosity_kinematic.to('mm**2/s').magnitude, 0.8917, atol=1e-3) # 1 cSt = 1 mm2/s
+    assert np.isclose(s1.density.to("kg/m**3").magnitude, 997.0479, atol=0.1)
+    assert np.isclose(s1.viscosity_kinematic.to("mm**2/s").magnitude, 0.8917, atol=1e-3)  # 1 cSt = 1 mm2/s
     assert np.isclose(s1.viscosity_dynamic, s1.viscosity_kinematic * s1.density, atol=1e-8)
 
 
@@ -376,18 +376,28 @@ def test_get_el_amt_dict(s6):
     # scale volume to 8L
     s6 *= 8
     d = s6.get_el_amt_dict()
+    d_nested = s6.get_el_amt_dict(nested=True)
     for el, amt in zip(
         ["H(1.0)", "O(-2.0)", "Ca(2.0)", "Mg(2.0)", "Na(1.0)", "Ag(1.0)", "C(4.0)", "S(6.0)", "Br(-1.0)"],
         [water_mol * 2 * 8, (water_mol + 0.018 + 0.24) * 8, 0.008, 0.040, 0.08, 0.08, 0.048, 0.48, 0.16],
         strict=False,
     ):
         assert np.isclose(d[el], amt, atol=1e-3)
 
+        el_no_valence = el.split("(")[0]
+        valence = float(el.split("(")[1].split(")")[0])
+        assert np.isclose(d_nested[el_no_valence][valence], amt, atol=1e-3)
+
     s = Solution({"Fe+2": "1 mM", "Fe+3": "5 mM", "FeCl2": "1 mM", "FeCl3": "5 mM"})
     d = s.get_el_amt_dict()
+    d_nested = s.get_el_amt_dict(nested=True)
     for el, amt in zip(["Fe(2.0)", "Fe(3.0)", "Cl(-1.0)"], [0.002, 0.01, 0.002 + 0.015], strict=False):
         assert np.isclose(d[el], amt, atol=1e-3)
 
+        el_no_valence = el.split("(")[0]
+        valence = float(el.split("(")[1].split(")")[0])
+        assert np.isclose(d_nested[el_no_valence][valence], amt, atol=1e-3)
+
 
 def test_p(s2):
     assert np.isclose(s2.p("Na+"), -1 * np.log10(s2.get_activity("Na+")))
@@ -491,6 +501,76 @@ def test_components_by_element(s1, s2):
     }
 
 
+def test_components_by_element_nested(s1, s2):
+    assert s1.get_components_by_element(nested=True) == {
+        "H": {
+            1.0: ["H2O(aq)", "OH[-1]", "H[+1]"],
+        },
+        "O": {
+            -2.0: ["H2O(aq)", "OH[-1]"],
+        },
+    }
+
+    assert s2.get_components_by_element(nested=True) == {
+        "H": {
+            1.0: ["H2O(aq)", "OH[-1]", "H[+1]"],
+        },
+        "O": {
+            -2.0: ["H2O(aq)", "OH[-1]"],
+        },
+        "Na": {
+            1.0: ["Na[+1]"],
+        },
+        "Cl": {
+            -1.0: ["Cl[-1]"],
+        },
+    }
+
+    if platform.machine() == "arm64" and platform.system() == "Darwin":
+        pytest.skip(reason="arm64 not supported")
+
+    s2.equilibrate()
+
+    assert s2.get_components_by_element(nested=True) == {
+        "H": {
+            1.0: [
+                "H2O(aq)",
+                "OH[-1]",
+                "H[+1]",
+                "HCl(aq)",
+                "NaOH(aq)",
+                "HClO(aq)",
+                "HClO2(aq)",
+            ],
+            0.0: ["H2(aq)"],
+        },
+        "O": {
+            -2.0: [
+                "H2O(aq)",
+                "OH[-1]",
+                "NaOH(aq)",
+                "HClO(aq)",
+                "ClO[-1]",
+                "ClO2[-1]",
+                "ClO3[-1]",
+                "ClO4[-1]",
+                "HClO2(aq)",
+            ],
+            0.0: ["O2(aq)"],
+        },
+        "Na": {
+            1.0: ["Na[+1]", "NaCl(aq)", "NaOH(aq)"],
+        },
+        "Cl": {
+            -1.0: ["Cl[-1]", "NaCl(aq)", "HCl(aq)"],
+            1.0: ["HClO(aq)", "ClO[-1]"],
+            3.0: ["ClO2[-1]", "HClO2(aq)"],
+            5.0: ["ClO3[-1]"],
+            7.0: ["ClO4[-1]"],
+        },
+    }
+
+
 def test_get_total_amount(s2):
     assert np.isclose(s2.get_total_amount("Na(1)", "g").magnitude, 8 * 58, 44)
     assert np.isclose(s2.get_total_amount("Na", "mol").magnitude, 8)