diff --git a/openmc/material.py b/openmc/material.py index d9d4f241d3..ac8e73afe6 100644 --- a/openmc/material.py +++ b/openmc/material.py @@ -795,16 +795,30 @@ 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] + + if element: + matching_nuclides = [] + for nuclide in self._nuclides: + if re.split(r'(\d+)', nuclide.name)[0] == element: + matching_nuclides.append(nuclide.name) + return list(set(matching_nuclides)) + + else: + return list(set([x.name for x in self._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..8879119a1f 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() == ['Li7', 'Li6'] + assert mat.get_nuclides(element='Be') == [] + + mat.add_element('Be', 1.0) + assert mat.get_nuclides() == ['Li7', 'Li6', 'Be9'] + assert mat.get_nuclides(element='Be') == ['Be9'] + + def test_get_elements(): # test that zero elements exist on creation m = openmc.Material()