From 1a520c9f4d598f96b4e6712fefa1d25b7d4da404 Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Fri, 4 Oct 2024 03:37:51 +0200
Subject: [PATCH] Adding material.get_element_atom_densities  (#3103)

Co-authored-by: Paul Romano <paul.k.romano@gmail.com>
---
 openmc/data/data.py               | 20 +++++++++-----
 openmc/material.py                | 43 +++++++++++++++++++++++++++++++
 tests/unit_tests/test_material.py | 22 ++++++++++++++++
 3 files changed, 78 insertions(+), 7 deletions(-)

diff --git a/openmc/data/data.py b/openmc/data/data.py
index d94d6aaaa39..408adf2e429 100644
--- a/openmc/data/data.py
+++ b/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 = []
diff --git a/openmc/material.py b/openmc/material.py
index 5b958c75a68..f550fd64900 100644
--- a/openmc/material.py
+++ b/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
diff --git a/tests/unit_tests/test_material.py b/tests/unit_tests/test_material.py
index 44c0730fbd2..94ba82571be 100644
--- a/tests/unit_tests/test_material.py
+++ b/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)