Adding material.get_element_atom_densities (#3103)

Co-authored-by: Paul Romano <[email protected]>

openmc/data/data.py

@@ -549,7 +549,18 @@ def gnds_name(Z, A, m=0):
     return f'{ATOMIC_SYMBOL[Z]}{A}'
 
 
-def isotopes(element):
+
+def _get_element_symbol(element: str) -> str:
+    if len(element) > 2:
+        symbol = ELEMENT_SYMBOL.get(element.lower())
+        if symbol is None:
+            raise ValueError(f'Element name "{element}" not recognized')
+        return symbol
+    else:
+        return element
+
+
+def isotopes(element: str) -> list[tuple[str, float]]:
     """Return naturally occurring isotopes and their abundances
 
     .. versionadded:: 0.12.1
@@ -570,12 +581,7 @@ def isotopes(element):
         If the element name is not recognized
 
     """
-    # Convert name to symbol if needed
-    if len(element) > 2:
-        symbol = ELEMENT_SYMBOL.get(element.lower())
-        if symbol is None:
-            raise ValueError(f'Element name "{element}" not recognised')
-        element = symbol
+    element = _get_element_symbol(element)
 
     # Get the nuclides present in nature
     result = []

openmc/material.py

@@ -18,6 +18,7 @@
 from .mixin import IDManagerMixin
 from openmc.checkvalue import PathLike
 from openmc.stats import Univariate, Discrete, Mixture
+from openmc.data.data import _get_element_symbol
 
 
 # Units for density supported by OpenMC
@@ -1075,6 +1076,48 @@ def get_nuclide_atom_densities(self, nuclide: str | None = None) -> dict[str, fl
 
         return nuclides
 
+    def get_element_atom_densities(self, element: str | None = None) -> dict[str, float]:
+        """Returns one or all elements in the material and their atomic
+        densities in units of atom/b-cm
+
+        .. versionadded:: 0.15.1
+
+        Parameters
+        ----------
+        element : str, optional
+            Element for which atom density is desired. If not specified, the
+            atom density for each element in the material is given.
+
+        Returns
+        -------
+        elements : dict
+            Dictionary whose keys are element names and values are densities in
+            [atom/b-cm]
+
+        """
+        if element is not None:
+            element = _get_element_symbol(element)
+
+        nuc_densities = self.get_nuclide_atom_densities()
+
+        # Initialize an empty dictionary for summed values
+        densities = {}
+
+        # Accumulate densities for each nuclide
+        for nuclide, density in nuc_densities.items():
+            nuc_element = openmc.data.ATOMIC_SYMBOL[openmc.data.zam(nuclide)[0]]
+            if element is None or element == nuc_element:
+                if nuc_element not in densities:
+                    densities[nuc_element] = 0.0
+                densities[nuc_element] += float(density)
+
+        # If specific element was requested, make sure it is present
+        if element is not None and element not in densities:
+                raise ValueError(f'Element {element} not found in material.')
+
+        return densities
+
+
     def get_activity(self, units: str = 'Bq/cm3', by_nuclide: bool = False,
                      volume: float | None = None) -> dict[str, float] | float:
         """Returns the activity of the material or for each nuclide in the

tests/unit_tests/test_material.py

@@ -380,6 +380,28 @@ def test_get_nuclide_atom_densities_specific(uo2):
     assert all_nuc['O16'] == one_nuc['O16']
 
 
+def test_get_element_atom_densities(uo2):
+    for element, density in uo2.get_element_atom_densities().items():
+        assert element in ('U', 'O')
+        assert density > 0
+
+
+def test_get_element_atom_densities_specific(uo2):
+    one_nuc = uo2.get_element_atom_densities('O')
+    assert list(one_nuc.keys()) == ['O']
+    assert list(one_nuc.values())[0] > 0
+
+    one_nuc = uo2.get_element_atom_densities('uranium')
+    assert list(one_nuc.keys()) == ['U']
+    assert list(one_nuc.values())[0] > 0
+
+    with pytest.raises(ValueError, match='not found'):
+        uo2.get_element_atom_densities('Li')
+
+    with pytest.raises(ValueError, match='not recognized'):
+        uo2.get_element_atom_densities('proximium')
+
+
 def test_get_nuclide_atoms():
     mat = openmc.Material()
     mat.add_nuclide('Li6', 1.0)