API ResultsReader to MATLAB

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 @@
     formatPlot,
     placeLegend,
     magicPlotDocDecorator,
+    deconvertVariableName,
 )
 from serpentTools.messages import (
     warning, SerpentToolsException,
     info,
 )
+from serpentTools.io import MatlabConverter
 
 
 MapStrVersions = {
@@ -678,3 +680,166 @@ def _expandPlotIterables(y, tail=''):
             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

serpentTools/tests/test_toMatlab.py

@@ -7,7 +7,7 @@
 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