From 7af56d8db22e552596989d691e44f6084081876f Mon Sep 17 00:00:00 2001
From: Paul Romano <paul.k.romano@gmail.com>
Date: Tue, 30 Aug 2022 21:17:50 -0500
Subject: [PATCH] Add get_material method on StepResult

---
 openmc/deplete/stepresult.py | 27 +++++++++++++++++++++++++++
 1 file changed, 27 insertions(+)

diff --git a/openmc/deplete/stepresult.py b/openmc/deplete/stepresult.py
index b26dfb44c4..41086baf31 100644
--- a/openmc/deplete/stepresult.py
+++ b/openmc/deplete/stepresult.py
@@ -6,6 +6,7 @@ timestep.
 
 from collections import OrderedDict
 import copy
+import warnings
 
 import h5py
 import numpy as np
@@ -198,6 +199,32 @@ class StepResult:
         new.rates = [r[ranges] for r in self.rates]
         return new
 
+    def get_material(self, mat_id):
+        """Return material object for given depleted composition
+
+        Parameters
+        ----------
+        mat_id : str
+            Material ID as a string
+
+        Returns
+        -------
+        openmc.Material
+            Equivalent material
+        """
+        with warnings.catch_warnings():
+            warnings.simplefilter('ignore', openmc.IDWarning)
+            material = openmc.Material(material_id=int(mat_id))
+        vol = self.volume[mat_id]
+        for nuc, _ in sorted(self.index_nuc.items(), key=lambda x: x[1]):
+            atoms = self[0, mat_id, nuc]
+            if atoms < 0.0:
+                continue
+            atom_per_bcm = atoms / vol * 1e-24
+            material.add_nuclide(nuc, atom_per_bcm)
+        material.volume = vol
+        return material
+
     def export_to_hdf5(self, filename, step):
         """Export results to an HDF5 file