Add formula.element to accumulate across isotopes and ions

periodictable/formulas.py

@@ -13,7 +13,7 @@
 from pyparsing import (Literal, Optional, White, Regex,
                        ZeroOrMore, OneOrMore, Forward, StringEnd, Group)
 
-from .core import default_table, isatom, isisotope, change_table
+from .core import default_table, isatom, isisotope, change_table, Ion
 from .constants import avogadro_number
 from .util import require_keywords, cell_volume
 
@@ -323,6 +323,29 @@ def hill(self):
         """
         return formula(self.atoms)
 
+    @property
+    def elements(self):
+        """
+        { *element*: *count*, ... }
+
+        Composition of the molecule.  Referencing this attribute computes
+        the *count* as the total number of each element in the chemical
+        formula, summed across all subgroups. Isotopes and ions are considered
+        equivalent.
+        """
+        return count_elements(self)
+
+    @property
+    def elements_hill(self):
+        """
+        Formula
+
+        Convert the formula to a formula in Hill notation.  Carbon appears
+        first followed by hydrogen then the remaining elements in alphabetical
+        order.
+        """
+        return formula(self.elements)
+
     def natural_mass_ratio(self):
         """
         Natural mass to isotope mass ratio.
@@ -891,10 +914,34 @@ def _count_atoms(seq):
             partial = _count_atoms(fragment)
         else:
             partial = {fragment: 1}
-        for el, elcount in partial.items():
-            if el not in total:
-                total[el] = 0
-            total[el] += elcount*count
+        for atom, atom_count in partial.items():
+            if atom not in total:
+                total[atom] = 0
+            total[atom] += atom_count*count
+    return total
+
+def count_elements(compound, by_isotope=False):
+    """
+    Composition of the molecule.
+
+    Returns {*element*: *count*, ...} where the *count* is the total number
+    of each element in the chemical formula, summed across all isotopes and
+    ionization levels. If *by_isotope* is True, then sum the individual isotopes
+    separately across all ionization levels.
+    """
+    total = {}
+    # Note: could accumulate charge at the same time as counting elements.
+    for part, count in formula(compound).atoms.items():
+        # Resolve isotopes and ions to the underlying element. Four cases:
+        #    isotope with charge needs fragment.element.element
+        #    isotope without charge needs fragment.element
+        #    element with charge needs fragment.element
+        #    element without charge needs fragment
+        if isinstance(part, Ion):
+            part = part.element
+        if not by_isotope:
+            part = getattr(part, "element", part)
+        total[part] = count + total.get(part, 0)
     return total
 
 def _immutable(seq):
@@ -950,7 +997,7 @@ def _str_atoms(seq):
     """
     Convert formula structure to string.
     """
-    #print "str", seq
+    #print("str", seq)
     ret = ""
     for count, fragment in seq:
         if isatom(fragment):

test/test_formulas.py

@@ -2,8 +2,9 @@
 from copy import deepcopy
 from pickle import loads, dumps
 
-from periodictable import Ca, C, O, H, Fe, Ni, Si, D, Na, Cl, Co, Ti
+from periodictable import Ca, C, O, H, Fe, Ni, Si, D, Na, Cl, Co, Ti, S
 from periodictable import formula, mix_by_weight, mix_by_volume
+import periodictable.formulas
 
 def test():
     ikaite = formula()
@@ -43,6 +44,16 @@ def test():
     # Check atom count
     assert formula("Fe2O4+3H2O").atoms == {Fe: 2, O: 7, H: 6}
 
+    # Check element count. The formula includes element, charged element,
+    # isotope and charged isotope. The "3" in front forces recursion into a
+    # formula tree.
+    assert formula("3HDS{6+}O{2-}3O[16]{2-}").elements == {S: 3, O: 12, H: 6}
+    assert str(formula("HDS{6+}O{2-}3O[16]{2-}").elements_hill) == "H2O4S"
+    isotopes = periodictable.formulas.count_elements(
+        formula("HDS{6+}O{2-}3O[16]{2-}"), by_isotope=True
+    )
+    assert isotopes == {S: 1, O: 3, O[16]:1, H: 1, D: 1}
+
     # Check charge
     assert formula("P{5+}O{2-}4").charge == -3
     try: