diff --git a/CHANGELOG.md b/CHANGELOG.md index 214b7ef8..d596bd4e 100644 --- a/CHANGELOG.md +++ b/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 diff --git a/src/pyEQL/solution.py b/src/pyEQL/solution.py index 0fc84ef9..30b81fb5 100644 --- a/src/pyEQL/solution.py +++ b/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 + to => 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 + to => 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: """ diff --git a/tests/test_solution.py b/tests/test_solution.py index 9dede33c..30939603 100644 --- a/tests/test_solution.py +++ b/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,6 +376,7 @@ 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], @@ -383,11 +384,20 @@ def test_get_el_amt_dict(s6): ): 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)