diff --git a/openmc/material.py b/openmc/material.py index d9d4f241d3..f137ec0cd7 100644 --- a/openmc/material.py +++ b/openmc/material.py @@ -795,16 +795,33 @@ class Material(IDManagerMixin): return sorted({re.split(r'(\d+)', i)[0] for i in self.get_nuclides()}) - def get_nuclides(self): - """Returns all nuclides in the material + def get_nuclides(self, element: Optional[str] = None): + """Returns a list of all nuclides in the material, if the element + argument is specified then just nuclides of that element are returned. + + Parameters + ---------- + element : str + Specifies the element to match when searching through the nuclides Returns ------- nuclides : list of str List of nuclide names - """ - return [x.name for x in self._nuclides] + + matching_nuclides = [] + if element: + for nuclide in self._nuclides: + if re.split(r'(\d+)', nuclide.name)[0] == element: + if nuclide.name not in matching_nuclides: + matching_nuclides.append(nuclide.name) + else: + for nuclide in self._nuclides: + if nuclide.name not in matching_nuclides: + matching_nuclides.append(nuclide.name) + + return matching_nuclides def get_nuclide_densities(self): """Returns all nuclides in the material and their densities diff --git a/tests/unit_tests/test_material.py b/tests/unit_tests/test_material.py index 8bb8817958..26ccbc9feb 100644 --- a/tests/unit_tests/test_material.py +++ b/tests/unit_tests/test_material.py @@ -299,6 +299,23 @@ def test_isotropic(): assert m2.isotropic == ['H1'] +def test_get_nuclides(): + mat = openmc.Material() + + mat.add_nuclide('Li6', 1.0) + assert mat.get_nuclides() == ['Li6'] + assert mat.get_nuclides(element='Li') == ['Li6'] + assert mat.get_nuclides(element='Be') == [] + + mat.add_element('Li', 1.0) + assert mat.get_nuclides() == ['Li6', 'Li7'] + assert mat.get_nuclides(element='Be') == [] + + mat.add_element('Be', 1.0) + assert mat.get_nuclides() == ['Li6', 'Li7', 'Be9'] + assert mat.get_nuclides(element='Be') == ['Be9'] + + def test_get_elements(): # test that zero elements exist on creation m = openmc.Material()