Repository: CORE-GATECH-GROUP/serpent-tools
Pull request #290 (2019-03-06 16:28:32)

API ResultsReader to MATLAB
Works to address #282 by adding a `toMatlab` method for exporting content of the `ResultsReader` to matlab. 
The only thing that should be noted is the group constants are stored not as 2D matrices, as they would if you loaded the .mat file directly. Instead, the are constucted be `burnup, universe index, <groups>, value/unc`. Here, `groups` means either number of macro/micro energy groups, or a sub-2D matrix in the case of scattering


Patch:
diff --git a/serpentTools/parsers/results.py b/serpentTools/parsers/results.py
index 366959a..77166af 100644
--- a/serpentTools/parsers/results.py
+++ b/serpentTools/parsers/results.py
@@ -2,7 +2,7 @@
 from collections import OrderedDict
 from six import iteritems
 
-from numpy import array, vstack
+from numpy import array, vstack, empty
 from cycler import cycler
 from matplotlib import rcParams
 from matplotlib.pyplot import gca
@@ -29,11 +29,13 @@ from serpentTools.utils import (
     formatPlot,
     placeLegend,
     magicPlotDocDecorator,
+    deconvertVariableName,
 )
 from serpentTools.messages import (
     warning, SerpentToolsException,
     info,
 )
+from serpentTools.io import MatlabConverter
 
 
 MapStrVersions = {
@@ -678,3 +680,166 @@ class ResultsReader(XSReader):
             return OrderedDict([[item, '{}{}'.format(item, tail)]
                                 for item in y])
         return y
+
+    def toMatlab(self, fileP, reconvert=True, append=True,
+                 format='5', longNames=True, compress=True,
+                 oned='row'):
+        """
+        Write a binary MATLAB file from the contents of this reader
+
+        The group constant data will be presented as a multi-dimensional
+        array, rather than a stacked 2D matrix. The axis are ordered
+        ``burnup, universeIndex, group, value/uncertainty``
+
+        The ordering of the universes can be found in the ``'UNIVERSES'``
+        vector if ``reconvert==True``, otherwise ``'universes'``. Each
+        universe ID is present in this vector, ordered to their
+        position along the second axis in the matrix.
+
+        Parameters
+        ----------
+        fileP: str or file-like object
+            Name of the file to write
+        reconvert: bool
+            If this evaluates to true, convert values back into their
+            original form as they appear in the output file.
+        append: bool
+            If true and a file exists under ``output``, append to that file.
+            Otherwise the file will be overwritten
+        format: {'5', '4'}
+            Format of file to write. ``'5'`` for MATLAB 5 to 7.2, ``'4'`` for
+            MATLAB 4
+        longNames: bool
+            If true, allow variable names to reach 63 characters,
+            which works with MATLAB 7.6+. Otherwise, maximum length is
+            31 characters
+        compress: bool
+            If true, compress matrices on write
+        oned: {'row', 'col'}:
+            Write one-dimensional arrays as row vectors if
+            ``oned=='row'`` (default), or column vectors
+
+        Examples
+        --------
+
+        >>> import serpentTools
+        >>> from scipy.io import loadmat
+        >>> r = serpentTools.readDataFile('pwr_res.m')
+        # convert to original variable names
+        >>> r.toMatlab('pwr_res.mat', True)
+        >>> loaded = loadmat('pwr_res.mat')
+        >>> loaded['ABS_KEFF']
+        array([[0.991938, 0.00145 ],
+               [0.181729, 0.0024  ]])
+        >>> kinf = loaded['INF_KINF']
+        >>> kinf.shape
+        (2, 1, 1, 2)
+        >>> kinf[:, 0, 0, 0]
+        array([0.993385, 0.181451])
+        >>> tot = loaded['INF_TOT']
+        >>> tot.shape
+        (2, 1, 2, 2)
+        >>> tot[:, 0, :, 0]
+        array([[0.496553, 1.21388 ],
+               [0.481875, 1.30993 ]])
+        # use the universes key to identify ordering of universes
+        >>> loaded['UNIVERSES']
+        array([0])
+
+        Raises
+        ------
+        ImportError:
+            If :term:`scipy` is not installed
+
+        See Also
+        --------
+        * :func:`scipy.io.savemat`
+        """
+        converter = MatlabConverter(self, fileP)
+        return converter.convert(reconvert, append, format, longNames,
+                                 compress, oned)
+
+    def _gather_matlab(self, reconvert):
+        if reconvert:
+            varFunc = getSerpentCaseName
+            out = {
+                varFunc(key): value for key, value in iteritems(self.metadata)
+            }
+            out.update({
+                varFunc(key): value for key, value in iteritems(self.resdata)
+            })
+        else:
+            out = {}
+            varFunc = getMixedCaseName
+            out.update(self.metadata)
+            out.update(self.resdata)
+        out.update(self._gather_univdata(varFunc))
+        return out
+
+    def _gather_univdata(self, func):
+        numAppearances = {}
+
+        for key in self.universes:
+            if key[0] not in numAppearances:
+                numAppearances[key[0]] = 1
+                continue
+            numAppearances[key[0]] += 1
+        # check to make sure all universes appear an identical
+        # number of times
+        burnupVals = set(numAppearances.values())
+        if len(burnupVals) != 1:
+            raise SerpentToolsException(
+                "Universes appear a different number of times:\n{}"
+                .format(numAppearances))
+
+        shapeStart = burnupVals.pop(), len(numAppearances)
+        univOrder = tuple(sorted(numAppearances))
+
+        univData = {func('universes'): univOrder}
+
+        for univKey, univ in iteritems(self.universes):
+
+            # position in matrix
+            uIndex = univOrder.index(univKey[0])
+            bIndex = univKey[2]
+
+            for expName, uncName in zip(
+                    ('infExp', 'b1Exp', 'gc'), ('infUnc', 'b1Unc', 'gcUnc')):
+                expD = getattr(univ, expName)
+                uncD = getattr(univ, uncName)
+                gatherPairedUnivData(univ, uIndex, bIndex, shapeStart, func,
+                                     expD, uncD, univData)
+
+        return univData
+
+
+def getMixedCaseName(name):
+    """
+    Return the name of the variable used for exporting - camelCase
+    """
+    return name
+
+
+def getSerpentCaseName(name):
+    """
+    Return the name of the variable used for exporting - SERPENT_CASE
+    """
+    return deconvertVariableName(name)
+
+
+def gatherPairedUnivData(universe, uIndex, bIndex, shapeStart, convFunc,
+                         expD, uncD, univData):
+    """Helper function to update universe dictionary for exporting"""
+    for xsKey, xsVal in iteritems(expD):
+        outKey = convFunc(xsKey)
+        xsUnc = uncD[xsKey]
+        if outKey not in univData:
+            shape = xsVal.shape
+            if not shape:  # single value like infKinf
+                shape = 1,
+            data = empty((shapeStart + shape + (2, )))
+            univData[outKey] = data
+        else:
+            data = univData[outKey]
+        data[bIndex, uIndex, ..., 0] = xsVal
+        data[bIndex, uIndex, ..., 1] = xsUnc


Test patch:
diff --git a/serpentTools/tests/test_toMatlab.py b/serpentTools/tests/test_toMatlab.py
index bf51ef1..9aaf24f 100644
--- a/serpentTools/tests/test_toMatlab.py
+++ b/serpentTools/tests/test_toMatlab.py
@@ -7,7 +7,7 @@ from numpy import arange
 from numpy.testing import assert_array_equal
 from serpentTools.objects import Detector
 from serpentTools.tests.utils import MatlabTesterHelper
-from serpentTools.data import getFile
+from serpentTools.data import getFile, readDataFile
 from serpentTools.parsers import DepmtxReader
 
 
@@ -124,7 +124,103 @@ class ConvertedFullDepmtxMatlabTester(DepmtxMatlabHelper):
     SPARSE = False
 
 
-del Det2MatlabHelper, DepmtxMatlabHelper
+class ResultToMatlabHelper(MatlabTesterHelper):
+    """Class for testing the result reader matlab export"""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.reader = readDataFile("pwr_res.m")
+
+    def setUp(self):
+        MatlabTesterHelper.setUp(self)
+        from scipy.io import loadmat
+        self.stream = BytesIO()
+        if self.CONVERT:
+            from serpentTools.parsers.results import getSerpentCaseName
+            self.convertFunc = getSerpentCaseName
+        else:
+            from serpentTools.parsers.results import getMixedCaseName
+            self.convertFunc = getMixedCaseName
+        self.reader.toMatlab(self.stream, self.CONVERT, oned='column')
+        self.gathered = loadmat(self.stream)
+
+    def checkMaybe1DArrays(self, expValue, loaded, origKey):
+        """Check arrays when the gathered value is likely a 2D row vector"""
+        if loaded.shape != expValue.shape:
+            self.assertEqual(expValue.shape[0], loaded.shape[0],
+                             msg=origKey)
+            loaded = loaded[:, 0]
+        assert_array_equal(expValue, loaded, err_msg=origKey)
+
+    def test_gatheredMetadata(self):
+        """Test the writing of metadata to matlab"""
+        for origKey, expValue in iteritems(self.reader.metadata):
+            expKey = self.convertFunc(origKey)
+            self.assertTrue(expKey in self.gathered, msg=origKey)
+            actual = self.gathered[expKey]
+            if isinstance(expValue, str):
+                actValue = actual[0]
+                self.assertEqual(expValue, actValue, msg=origKey)
+            elif isinstance(expValue, (int, float)):
+                # stored as (1, 1) matrix
+                actValue = actual.flatten()[0]
+                self.assertEqual(expValue, actValue, msg=origKey)
+            else:  # try numpy array
+                self.checkMaybe1DArrays(expValue, actual, origKey)
+
+    def test_dumpedResdata(self):
+        """Test the writing of results data to matlab"""
+        for origKey, expValue in iteritems(self.reader.resdata):
+            expKey = self.convertFunc(origKey)
+            self.assertTrue(expKey in self.gathered, msg=origKey)
+            self.checkMaybe1DArrays(expValue, self.gathered[expKey], origKey)
+
+    def test_dumpedUnivData(self):
+        """Test the writing of universe data to matlab"""
+        univIxKey = self.convertFunc("universes")
+        burnIxKey = self.convertFunc("burnStep")
+        self.assertTrue(univIxKey in self.gathered)
+        self.assertTrue(burnIxKey in self.gathered)
+        univOrder = tuple(self.gathered[univIxKey])
+        burnOrder = tuple(self.gathered[burnIxKey].flatten())
+        for univKey, universe in iteritems(self.reader.universes):
+            self.assertTrue(univKey[0] in univOrder,
+                            msg="{} // {}".format(univKey, univOrder))
+            uIndex = univOrder.index(univKey[0])
+            actBurnIx = univKey[2]
+            self.assertTrue(actBurnIx in burnOrder,
+                            msg="{} // {}".format(univKey, burnOrder))
+            burnIndex = burnOrder.index(actBurnIx)
+            self.checkUnivContents(universe, uIndex, burnIndex,
+                                   universe.infExp, universe.infUnc)
+            self.checkUnivContents(universe, uIndex, burnIndex,
+                                   universe.b1Exp, universe.b1Unc)
+            self.checkUnivContents(universe, uIndex, burnIndex,
+                                   universe.gc, universe.gcUnc)
+
+    def checkUnivContents(self, universe, univIndex, burnIndex,
+                          expGcD, uncGcD):
+        """
+        Check the contents of a specific subset of group constant dictionaries
+        """
+        for origKey, expValue in iteritems(expGcD):
+            uncValue = uncGcD[origKey]
+            expKey = self.convertFunc(origKey)
+            self.assertTrue(expKey in self.gathered, msg=origKey)
+            actData = self.gathered[expKey][burnIndex, univIndex, ..., :]
+            assert_array_equal(expValue, actData[..., 0], err_msg=origKey)
+            assert_array_equal(uncValue, actData[..., 1], err_msg=origKey)
+
+
+class SerpentCaseResults2MatlabTester(ResultToMatlabHelper):
+    CONVERT = True
+
+
+class MixedCaseResults2MatlabTester(ResultToMatlabHelper):
+    CONVERT = False
+
+
+del Det2MatlabHelper, DepmtxMatlabHelper, ResultToMatlabHelper
 
 if __name__ == '__main__':
     from unittest import main