Repository: CORE-GATECH-GROUP/serpent-tools
Pull request #391 (2020-03-29 23:33:15)

API Export depletion data to pandas DataFrame
This PR introduces `DepletedMaterial.toDataFrame` that can produce a nice data frame of quantities stored on the instance. The columns are isotope names or ZAI identifiers, while the rows are index by time. Users can request specific isotopes by name or by ZAI, not both. Additionally, the index of the data frame can be set to points in time, burnup, or burnup step.

I've added a pytest-exclusive test file for this method as well. 


Patch:
diff --git a/serpentTools/data/plut_xs0.m b/serpentTools/data/plut_xs0.m
index f7f5610..bceb792 100644
--- a/serpentTools/data/plut_xs0.m
+++ b/serpentTools/data/plut_xs0.m
@@ -191,24 +191,6 @@ i94239_03c_xs = [
 5.87774E+00 2.27284E-03 2.72848E+00 2.41217E+00 2.27284E-03 9.49498E-02 1.42231E-01 1.91357E-02 2.38983E-02 6.89030E-03 0.00000E+00 8.01346E-03 0.00000E+00 0.00000E+00 0.00000E+00 0.00000E+00 0.00000E+00 0.00000E+00 4.33261E-03 0.00000E+00 5.77543E-03 0.00000E+00 0.00000E+00 7.03469E-02 1.68538E-03 5.14204E-03 2.52079E-03 6.96677E-03 6.93165E-03 1.49274E-02 6.96715E-03 7.46602E-03 5.97098E-03 7.96117E-03 1.07498E-02 1.34154E-02 3.67423E-04 4.65750E-03 5.32589E-03 4.56426E-03 4.43752E-03 1.66154E-03 1.81176E-03 1.87937E-03 1.40625E-03 1.56047E-03 2.24766E-01 1.60937E-02 0.00000E+00 
 ];
 
-i82206_09c_bra_mt = [
- 102 % (n,gamma)
-];
-
-i82206_09c_bra_f = [
-2.17000E-02 
-2.17000E-02 
-2.17000E-02 
-2.17000E-02 
-2.17000E-02 
-2.17000E-02 
-2.17000E-02 
-2.17000E-02 
-2.17000E-02 
-2.17000E-02 
-];
-
-
 i94239_03c_nu = [
 2.88002 0.00645 
 2.87172 0.00645 
diff --git a/serpentTools/objects/materials.py b/serpentTools/objects/materials.py
index 444fd47..5e1610b 100644
--- a/serpentTools/objects/materials.py
+++ b/serpentTools/objects/materials.py
@@ -399,3 +399,79 @@ class DepletedMaterial(DepletedMaterialBase):
                         xlabel=xlabel, ylabel=ylabel,
                         title=title, legend=legend)
         return ax
+
+    def toDataFrame(self, quantity, names=None, zai=None, index="days"):
+        """Create a pandas DataFrame for a property of interest
+
+        If ``names`` and ``zai`` are not provided, then the isotope
+        names will be used as the columns.
+
+        Parameters
+        ----------
+        quantity : str
+            Either a key in :attr:`data` or the string name of an attribute
+            like ``photonProdRate`` for :attr:`photonProdRate`
+        names : sequence of str, optional
+            Specific isotope names to obtain. Not compatible with ``zai``
+        zai : sequence of int, optional
+            Specific isotope zai to obtain. Not compatible with ``names``
+        index : {"days", "burnup", "step"}, optional
+            What array to use for the index or rows of the DataFrame.
+            Defaults to using :attr:`days`, but ``"burnup"`` can be passed
+            to use :attr:`burnup`, if it is present. The value of ``"step"``
+            will simply create a basic index that starts at 0 and increments
+            by one per row
+
+        Returns
+        -------
+        pandas.DataFrame
+            2D tabulated representation of the requested array. Columns
+            reflect isotopes, rows represent points in time
+
+        """
+        import pandas
+
+        if quantity in {"burnup", "volume"}:
+            raise ValueError("{} does not reflect 2D isotopic data".format(
+                quantity))
+
+        if names is not None and zai is not None:
+            raise ValueError("Cannot pass both isotope names and zai")
+
+        if index == "days":
+            dfIndex = pandas.Index(self.days, name="Time [d]")
+        elif index == "burnup":
+            bu = self.burnup
+            if bu is None:
+                raise AttributeError(
+                    "Burnup not set on material {}".format(self.name))
+            dfIndex = pandas.Index(bu, name="Burnup [MWd/kgU]")
+        elif index == "step":
+            dfIndex = pandas.Index(range(len(self.days)), name="Step")
+        else:
+            raise ValueError(
+                "Index must be days, burnup, or step, not {}".format(index))
+
+        data = self.data.get(quantity)
+
+        if data is None:
+            data = getattr(self, quantity, None)
+        if data is None:
+            raise AttributeError("Quantity {} not present as key in data "
+                                 "or as attribute".format(quantity))
+
+        if names is None:
+            if zai is None:
+                columns = pandas.Index(self.names, name="Isotope")
+                isoslice = slice(None)
+            else:
+                isoslice = self._getRowIndices("zai", zai)
+                columns = pandas.Index(
+                    [self.zai[x] for x in isoslice], name="Isotope ZAI")
+        else:
+            isoslice = self._getRowIndices("names", names)
+            columns = pandas.Index(
+                [self.names[x] for x in isoslice], name="Isotopes")
+
+        return pandas.DataFrame(
+            data[isoslice].T.copy(), index=dfIndex, columns=columns)
diff --git a/serpentTools/parsers/xsplot.py b/serpentTools/parsers/xsplot.py
index f550fb3..6768df4 100644
--- a/serpentTools/parsers/xsplot.py
+++ b/serpentTools/parsers/xsplot.py
@@ -1,8 +1,8 @@
 """Parser responsible for reading the ``*_xsplot.m`` files"""
 from numpy import array, float64
 
-from serpentTools.messages import error, warning
 from serpentTools.engines import KeywordParser
+from serpentTools.messages import warning, info, debug, error
 from serpentTools.parsers.base import BaseReader
 from serpentTools.objects.xsdata import XSData
 
@@ -53,44 +53,51 @@ class XSPlotReader(BaseReader):
 
     def _read(self):
         """Read through the depletion file and store requested data."""
+        info('Preparing to read {}'.format(self.filePath))
         keys = ['E', r'i\d{4,5}', r'm\w']
         separators = ['\n', '];', '\r\n']
 
         with KeywordParser(self.filePath, keys, separators) as parser:
             for chunk in parser.yieldChunks():
 
-                lead = chunk[0].strip()
-                data = chunk[1:]
-                if lead.startswith("E = ["):
+                if chunk[0][:5] == 'E = [':
                     # The energy grid
-                    self.metadata['egrid'] = array(data, dtype=float64)
+                    self.metadata['egrid'] = array(chunk[1:], dtype=float64)
 
-                elif lead.endswith('majorant_xs = ['):
+                elif chunk[0][:15] == 'majorant_xs = [':
                     # L-inf norm on all XS on all materials
-                    self.metadata['majorant_xs'] = array(data,
+                    self.metadata['majorant_xs'] = array(chunk[1:],
                                                          dtype=float64)
 
-                elif lead.endswith('_mt = ['):
-                    xsname = lead[:lead.index("_mt")]
-                    isiso = lead[0] == 'i'
+                elif chunk[0][-7:] == 'mt = [\n':
+                    debug('found mt specification')
+                    xsname = chunk[0][:-8]
+                    isiso = True if chunk[0][0] == 'i' else False
                     self.xsections[xsname] = XSData(xsname, self.metadata,
                                                     isIso=isiso)
                     self.xsections[xsname].setMTs(chunk)
 
-                elif lead.endswith('_xs = ['):
-                    xsname = lead[:lead.index("_xs")]
+                elif chunk[0][-7:] == 'xs = [\n':
+                    debug('found xs specification')
+                    xsname = chunk[0][:-8]
                     self.xsections[xsname].setData(chunk)
 
-                elif lead.endswith('_nu = ['):
-                    xsname = lead[:lead.index("_nu")]
+                elif chunk[0][-7:] == 'nu = [\n':
+                    debug('found nu specification')
+                    xsname = chunk[0][:-8]
                     self.xsections[xsname].setNuData(chunk)
 
-                elif lead.endswith("bra_f = ["):
-                    xsname = lead[:lead.index("_f")]
-                    self.xsections[xsname].setData(chunk)
+                elif 'bra_f' in chunk[0]:
+                    warning("There is this weird 'bra_f' XS. these seem to be"
+                            " constant. recording to metadata instead.")
+                    self.metadata[xsname].setData(chunk)
 
                 else:
-                    raise ValueError("Unidentifiable entry\n{}".format(chunk))
+                    print(chunk)
+                    error('Unidentifiable entry {}'.format(chunk[0]))
+
+        info('Done reading xsplot file')
+        debug('  found {} xs listings'.format(len(self.xsections)))
 
     def _precheck(self):
         """do a quick scan to ensure this looks like a xsplot file."""


Test patch:
diff --git a/tests/test_pt_depletion.py b/tests/test_pt_depletion.py
new file mode 100644
index 0000000..650f778
--- /dev/null
+++ b/tests/test_pt_depletion.py
@@ -0,0 +1,58 @@
+import sys
+import numpy
+import pytest
+import serpentTools
+
+
+@pytest.fixture(scope="module")
+def referenceDepReader():
+    return serpentTools.readDataFile("ref_dep.m")
+
+
+@pytest.fixture(scope="module")
+def referenceMaterial(referenceDepReader):
+    return referenceDepReader.materials["fuel"]
+
+
+@pytest.mark.skipif(sys.version_info < (3, 6), reason="pandas requires 3.6+")
+@pytest.mark.parametrize("quantity", ["adens", "ingTox", "inhTox"])
+def test_toDataFrame(referenceMaterial, quantity):
+    adensFrame = referenceMaterial.toDataFrame(quantity)
+    assert numpy.array_equal(
+        adensFrame.columns.values, referenceMaterial.names
+    )
+    assert numpy.array_equal(adensFrame.index.values, referenceMaterial.days)
+    assert numpy.array_equal(
+        adensFrame.values, referenceMaterial.data[quantity].T
+    )
+
+    reference = referenceMaterial.getValues(
+        "days", quantity, names=["U235", "Xe135", "Pu239"]
+    )
+    frame0 = referenceMaterial.toDataFrame(
+        quantity, names=["U235", "Xe135", "Pu239"], index="burnup"
+    )
+    frame1 = referenceMaterial.toDataFrame(
+        quantity, zai=[922350, 541350, 942390], index="step"
+    )
+
+    assert numpy.array_equal(frame0.values, reference.T)
+    assert numpy.array_equal(frame0.values, frame1.values)
+    assert numpy.array_equal(frame0.index.values, referenceMaterial.burnup)
+    assert numpy.array_equal(frame0.columns.values, ["U235", "Xe135", "Pu239"])
+    assert numpy.array_equal(frame1.columns.values, [922350, 541350, 942390])
+    assert numpy.array_equal(
+        frame1.index.values, range(len(referenceMaterial.days))
+    )
+
+
+@pytest.mark.skipif(sys.version_info < (3, 6), reason="pandas requires 3.6+")
+def test_toDataFrameErrors(referenceMaterial):
+    with pytest.raises(ValueError, match=".*2D"):
+        referenceMaterial.toDataFrame("burnup")
+
+    with pytest.raises(ValueError, match="Cannot pass both"):
+        referenceMaterial.toDataFrame("adens", names=["U235"], zai=[922350])
+
+    with pytest.raises(AttributeError, match=".*bad"):
+        referenceMaterial.toDataFrame("bad")
diff --git a/tests/test_xsplot.py b/tests/test_xsplot.py
index b87a39c..57b0ac0 100644
--- a/tests/test_xsplot.py
+++ b/tests/test_xsplot.py
@@ -1,6 +1,6 @@
 """Test the xsplot reader."""
 import unittest
-from numpy import ndarray, array, newaxis
+from numpy import ndarray, array
 
 from numpy.testing import assert_array_equal
 
@@ -302,8 +302,8 @@ class XSPlotTester(unittest.TestCase):
             },
             'i8016_03c': {
                 'MT': [
-                    1, 101, 2, 102, 107, 51, 52, 53, 54, 22, 55, 91,
-                    103, 104, 56, 57, 28, 108, 105, 23, 16],
+                    1, 101, 2, 102, 107, 51, 52, 53, 54, 22, 55,
+                    3, 104, 56, 57, 28, 108, 105, 23, 16],
                 'MTdescrip': [
                     ' Total', ' Sum of absorption',
                     ' elastic scattering', ' (n,gamma)',
@@ -314,14 +314,11 @@ class XSPlotTester(unittest.TestCase):
                     ' inelastic scatt. to 4. excited state',
                     ' (n,nalpha)',
                     ' inelastic scatt. to 5. excited state',
-                    ' inelastic scattering to continuum',
-                    ' (n,p)',
+                    ' inelastic scattering to continuum', ' (n,p)',
                     ' (n,d)',
                     ' inelastic scatt. to 6. excited state',
                     ' inelastic scatt. to 7. excited state',
-                    ' (n,np)',
-                    ' (n,2alpha)',
-                    ' (n,t)',
+                    ' (n,np)', ' (n,2alpha)', ' (n,t)',
                     ' (n,n3alpha)',
                     ' (n,2n)',
                 ],
@@ -736,29 +733,6 @@ class XSPlotTester(unittest.TestCase):
                     [2.93887e-01, 1.36424e-01, 1.13642e-04,
                      1.20609e-01, 5.96505e-01, 1.28477e-02]]),
             },
-            "i82206_09c_bra": {
-                "MT": [102],
-                "MTdescrip": [" (n,gamma)"],
-                "hasNuData": False,
-                "isIso": True,
-                'metadata': {
-                    'egrid': array(
-                        [1.00000e-08, 1.03891e-07, 1.07934e-06,
-                         1.12135e-05, 1.16498e-04, 1.21032e-03,
-                         1.25742e-02, 1.30635e-01, 1.35719e+00,
-                         1.41000e+01]),
-                    'majorant_xs': array(
-                        [78.4253, 36.1666, 2.54417, 13.0654,
-                         4.27811, 0.822536, 0.781066, 0.598564,
-                         0.34175, 0.293887]),
-                },
-                "name": "i82206_09c_bra",
-                "xsdata": array([
-                    2.17000E-02, 2.17000E-02, 2.17000E-02, 2.17000E-02,
-                    2.17000E-02, 2.17000E-02, 2.17000E-02, 2.17000E-02,
-                    2.17000E-02, 2.17000E-02]
-                )[:, newaxis],  # Promote to 2D
-            },
         }
 
         # newly obtained results: