Merge pull request #120 from KingsburyLab/bugfix

Solution.get_total_amount: fix error on ppm unit

CHANGELOG.md

@@ -7,12 +7,19 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Added
+
+- `utils.interpret_units`: New method to "sanitize" environmental engineering units like ppm to strings that `pint`
+  can understand, e.g. ppm -> mg/L. This method is now used in `get_amount` and `get_total_amount` to ensure consistency
+  in how they process units.
+
 ### Changed
 
 - CI: `pre-commit autoupdate`
 
 ### Fixed
 
+- `Solution.get_total_amount`: Fixed an issue in which `ppm` units would fail.
 - `Solution`: Fixed an issue in which repeated calls to `equilibrate` when using `NativeEOS` or `PHREEQCEOS` would
   change the mass of the `Solution` slightly. This was attributed to the fact that `pyEQL` and `PHREEQC` use slightly
   different molecular weights for water.

src/pyEQL/solution.py

@@ -33,7 +33,7 @@
 from pyEQL.logging_system import logger
 from pyEQL.salt_ion_match import Salt
 from pyEQL.solute import Solute
-from pyEQL.utils import FormulaDict, create_water_substance, standardize_formula
+from pyEQL.utils import FormulaDict, create_water_substance, interpret_units, standardize_formula
 
 EQUIV_WT_CACO3 = ureg.Quantity(100.09 / 2, "g/mol")
 # string to denote unknown oxidation states
@@ -997,23 +997,17 @@ def get_amount(self, solute: str, units: str = "mol/L") -> Quantity:
             :meth:`get_osmolarity`
             :meth:`get_osmolality`
             :meth:`get_total_moles_solute`
+            :func:`pyEQL.utils.interpret_units`
         """
         z = 1
         # sanitized unit to be passed to pint
-        _units = units
         if "eq" in units:
             _units = units.replace("eq", "mol")
             z = self.get_property(solute, "charge")
             if z == 0:  # uncharged solutes have zero equiv concentration
                 return ureg.Quantity(0, _units)
-        elif units == "m":  # molal
-            _units = "mol/kg"
-        elif units == "ppm":
-            _units = "mg/L"
-        elif units == "ppb":
-            _units = "ug/L"
-        elif units == "ppt":
-            _units = "ng/L"
+        else:
+            _units = interpret_units(units)
 
         # retrieve the number of moles of solute and its molecular weight
         try:
@@ -1119,7 +1113,7 @@ def get_el_amt_dict(self):
 
         return d
 
-    def get_total_amount(self, element: str, units) -> Quantity:
+    def get_total_amount(self, element: str, units: str) -> Quantity:
         """
         Return the total amount of 'element' (across all solutes) in the solution.
 
@@ -1128,19 +1122,21 @@ def get_total_amount(self, element: str, units) -> Quantity:
                 oxidation state in parentheses, e.g., "Na(1.0)", "Fe(2.0)", or "O(0.0)". If no oxidation state
                 is given, the total concentration of the element (over all oxidation states) is returned.
             units : str
-                        Units desired for the output. Examples of valid units are
-                        'mol/L','mol/kg','mol', 'kg', and 'g/L'
+                Units desired for the output. Any unit understood by `get_amount` can be used. Examples of valid
+                units are 'mol/L','mol/kg','mol', 'kg', and 'g/L'.
 
         Returns:
             The total amount of the element in the solution, in the specified units
 
         See Also:
             :meth:`get_amount`
+            :func:`pyEQL.utils.interpret_units`
         """
         TOT: Quantity = 0
 
-        # standardize the element formula
+        # standardize the element formula and units
         el = str(Element(element.split("(")[0]))
+        units = interpret_units(units)
 
         # enumerate the species whose concentrations we need
         comp_by_element = self.get_components_by_element()

src/pyEQL/utils.py

@@ -15,6 +15,27 @@
 from pyEQL import ureg
 
 
+def interpret_units(unit: str) -> str:
+    """
+    Translate commonly used environmental units such as 'ppm' into strings that `pint` can understand.
+
+    Args:
+        unit: string representing the unit to translate
+
+    Returns: a unit that pint can understand
+    """
+    if unit == "m":  # molal
+        return "mol/kg"
+    if unit == "ppm":
+        return "mg/L"
+    if unit == "ppb":
+        return "ug/L"
+    if unit == "ppt":
+        return "ng/L"
+    # if all else fails, return the unit we were provided
+    return unit
+
+
 @lru_cache
 def standardize_formula(formula: str):
     """