From 085125ef866f62cb4f0486385d5c299fde550219 Mon Sep 17 00:00:00 2001
From: Jonathan Shimwell <mail@jshimwell.com>
Date: Tue, 23 Aug 2022 16:30:53 +0100
Subject: [PATCH] added optional nuclide to atom densities

---
 openmc/material.py                | 21 ++++++++++++++-------
 tests/unit_tests/test_material.py |  7 +++++++
 2 files changed, 21 insertions(+), 7 deletions(-)

diff --git a/openmc/material.py b/openmc/material.py
index 167eaba84a..16b5923d7f 100644
--- a/openmc/material.py
+++ b/openmc/material.py
@@ -825,9 +825,15 @@ class Material(IDManagerMixin):
 
         return nuclides
 
-    def get_nuclide_atom_densities(self):
-        """Returns all nuclides in the material and their atomic densities in
-        units of atom/b-cm
+    def get_nuclide_atom_densities(self, nuclide: Optional[str] = None):
+        """Returns one or all nuclides in the material and their atomic
+        densities in units of atom/b-cm
+
+        Parameters
+        ----------
+        nuclides : str, optional
+            Nuclide for which atom density is desired. If not specified, the
+            atom density for the all the material is given.
 
         .. versionchanged:: 0.13.1
             The values in the dictionary were changed from a tuple containing
@@ -862,10 +868,11 @@ class Material(IDManagerMixin):
         nuc_densities = []
         nuc_density_types = []
 
-        for nuclide in self.nuclides:
-            nucs.append(nuclide.name)
-            nuc_densities.append(nuclide.percent)
-            nuc_density_types.append(nuclide.percent_type)
+        for loop_nuclide in self.nuclides:
+            if nuclide is None or nuclide == loop_nuclide.name:
+                nucs.append(loop_nuclide.name)
+                nuc_densities.append(loop_nuclide.percent)
+                nuc_density_types.append(loop_nuclide.percent_type)
 
         nucs = np.array(nucs)
         nuc_densities = np.array(nuc_densities)
diff --git a/tests/unit_tests/test_material.py b/tests/unit_tests/test_material.py
index 451c5308a4..1e5b8a797e 100644
--- a/tests/unit_tests/test_material.py
+++ b/tests/unit_tests/test_material.py
@@ -339,6 +339,13 @@ def test_get_nuclide_atom_densities(uo2):
         assert density > 0
 
 
+def test_get_specific_nuclide_atom_densities(uo2):
+    print(uo2)
+    nuc = uo2.get_nuclide_atom_densities(nuclide='O16')
+    assert list(nuc.keys()) == ['O16']
+    assert list(nuc.values())[0] > 0
+
+
 def test_get_nuclide_atoms():
     mat = openmc.Material()
     mat.add_nuclide('Li6', 1.0)