Repository: CORE-GATECH-GROUP/serpent-tools
Pull request #257 (2018-12-21 14:08:15)

API DepletionReader export to Matlab file
Utilizes [`scipy.io.savemat`](https://docs.scipy.org/doc/scipy/reference/generated/scipy.io.savemat.html?highlight=savemat#scipy.io.savemat) to take contents of the reader and export them to a binary `.mat` file. This can be loaded into MATLAB, and in many cases is faster and more successful than reading in the text file.
The method supports converting variables back to their original SERPENT_STYLE names or not, and optionally saving metadata as well.

Added unit tests that cover three cases:

- Reconvert variable names back and write metadata
- Reconvert variable names but don't write metadata
- Don't convert variable names but do write metadata

- [x] PR fits the [project scope](http://serpent-tools.readthedocs.io/en/latest/develop/contributing.html#project-scope)
- [x] Code to be committed is [DRY](https://en.wikipedia.org/wiki/Don%27t_repeat_yourself) and matches the [code style](http://serpent-tools.readthedocs.io/en/latest/develop/codestyle.html#code-style) of the project


Patch:
diff --git a/docs/api/depletion.rst b/docs/api/depletion.rst
index 83b09f5..b2d35da 100644
--- a/docs/api/depletion.rst
+++ b/docs/api/depletion.rst
@@ -4,12 +4,6 @@
 Depletion Reader
 ================
 
-.. warning::
-
-    Does not support depleted materials with underscores,
-    i.e. ``fuel_1`` will not be matched with the current methods.
-    Follow up on, or claim, the issue on GitHub: :issue:`58`
-
 .. autoclass:: serpentTools.parsers.depletion.DepletionReader
     :special-members: __getitem__
 
diff --git a/docs/changelog.rst b/docs/changelog.rst
index e62847f..2d2a76a 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -13,7 +13,11 @@ Changelog
 
 Next
 ====
-* :meth:`serpentTools.settings.rc.loadYaml` uses ``safe_load``
+* :pull:`256` :meth:`serpentTools.settings.rc.loadYaml` uses ``safe_load``
+* :pull:`257` |depletionReader| now can utilize 
+  :meth:`~serpentTools.parsers.depletion.DepletionReader.saveAsMatlab` for
+  exporting data to a binary ``.mat`` file
+
 .. _v0.6.0:
 
 0.6.0
diff --git a/serpentTools/parsers/depletion.py b/serpentTools/parsers/depletion.py
index 5614d93..269f229 100644
--- a/serpentTools/parsers/depletion.py
+++ b/serpentTools/parsers/depletion.py
@@ -12,7 +12,7 @@ from ._collections import DEPLETION_VARIABLES
 from serpentTools.engines import KeywordParser
 from serpentTools.parsers.base import MaterialReader
 from serpentTools.objects.materials import DepletedMaterial
-
+from serpentTools.utils.core import deconvertVariableName
 from serpentTools.messages import (
     warning, debug, error, SerpentToolsException,
 )
@@ -354,6 +354,65 @@ class DepletionReader(DepPlotMixin, MaterialReader):
             lower, upper, 'burnup')
         return similar
 
+    def saveAsMatlab(self, fileP, reconvert=True, metadata=True,
+                     **savematKwargs):
+        """
+        Write a binary MATLAB file from the contents of this reader
+
+
+        .. note::
+
+            Vectors will be stored as either column or row
+            matrices, e.g. ``1xN`` or ``Nx1``. This can be
+            controlled by passing ``oned_as`` to be either
+            ``'row'`` or ``'column'``
+
+        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, e.g.
+            ``MAT_TOTAL_ING_TOX``. Otherwise, maintain the ``mixedCase``
+            style, ``total_ingTox``.
+        metadata: bool or str or list of strings
+            If this evaluates to true, then write all metadata to the
+            file as well.
+        savematKwargs:
+            Additional keyword arguments to pass to
+            :func:`scipy.io.savemat`
+
+        Raises
+        ------
+        ImportError:
+            If :term:`scipy` is not installed
+
+        See Also
+        --------
+        :func:`scipy.io.savemat`
+        """
+        from scipy.io import savemat
+
+        # set these here to reduce number of if/elses
+
+        if reconvert:
+            converter = deconvert
+        else:
+            converter = prepToMatlab
+
+        if metadata:
+            data = {k.upper() if reconvert else k: v
+                    for k, v in iteritems(self.metadata)}
+        else:
+            data = {}
+
+        for matName, material in iteritems(self.materials):
+            for varName, varData in iteritems(material.data):
+                data[converter(matName, varName)] = varData
+
+        savemat(fileP, data, **savematKwargs)
+
 
 def getMaterialNameAndVariable(mlabName):
     for serpVar in DEPLETION_VARIABLES:
@@ -374,3 +433,17 @@ def getMatlabVarName(chunk):
     if isinstance(chunk, list):
         return getMatlabVarName(chunk[0])
     return chunk[:chunk.index(' ')]
+
+
+_MAT_FMT_ORIG = "MAT_{}_{}"
+_MAT_FMT_CC = "{}_{}"
+
+
+def deconvert(material, variable):
+    """Restore the original name as present in the depletion file"""
+    return _MAT_FMT_ORIG.format(material, deconvertVariableName(variable))
+
+
+def prepToMatlab(material, variable):
+    """Create the name of a single variable for MATLAB"""
+    return _MAT_FMT_CC.format(material, variable)
diff --git a/serpentTools/utils/core.py b/serpentTools/utils/core.py
index 3a19ea7..e56468e 100644
--- a/serpentTools/utils/core.py
+++ b/serpentTools/utils/core.py
@@ -153,6 +153,17 @@ def convertVariableName(variable):
                                      for item in lowerSplits[1:]])
 
 
+def deconvertVariableName(variable):
+    """Convert ``mixedCase`` to ``SERPENT_CASE``"""
+    out = ""
+    for char in variable:
+        if char.isupper():
+            out += '_' + char
+            continue
+        out += char.upper()
+    return out
+
+
 LEADER_TO_WIKI = "http://serpent.vtt.fi/mediawiki/index.php/"
 
 


Test patch:
diff --git a/serpentTools/tests/test_depletion.py b/serpentTools/tests/test_depletion.py
index 32d91a8..7063f1b 100644
--- a/serpentTools/tests/test_depletion.py
+++ b/serpentTools/tests/test_depletion.py
@@ -1,5 +1,5 @@
 """Test the depletion file."""
-from unittest import TestCase
+from unittest import TestCase, skipIf
 from os import remove
 
 from numpy import array
@@ -11,7 +11,9 @@ from serpentTools.data import getFile
 from serpentTools.settings import rc
 from serpentTools.parsers.depletion import (
     DepletionReader, getMaterialNameAndVariable, getMatlabVarName,
+    prepToMatlab, deconvert,
 )
+from serpentTools.utils import hasScipy
 from serpentTools.tests.utils import LoggerMixin
 
 
@@ -260,6 +262,132 @@ class DepletionUtilTester(TestCase):
         self.assertEqual(self.BAD_0, getMatlabVarName(self.VAR_CHUNK[0]))
 
 
+class DepMatlabExportHelper(TestCase):
+    """
+    Class that tests the saveAsMatlab method for DepletionReaders
+    """
+
+    MAT_FILE_TEMPLATE = "dep_c{}_m{}.mat"
+
+    @classmethod
+    def setUpClass(cls):
+        cls.reader = DepletionReader(DEP_FILE_PATH)
+        cls.reader.read()
+
+    def setUp(self):
+        self.MAT_FILE = self.MAT_FILE_TEMPLATE.format(
+            1 if self.RECONVERT else 0,
+            1 if self.SAVE_METADATA else 0,
+        )
+
+    @skipIf(not hasScipy(), "scipy not installed")
+    def test_saveMatlab(self):
+        """
+        Verify that the reader can write data to a .mat file
+        """
+        from scipy.io import loadmat
+        self.reader.saveAsMatlab(self.MAT_FILE, self.RECONVERT,
+                                 self.SAVE_METADATA)
+        loaded = loadmat(self.MAT_FILE)
+        self.check(loaded)
+        if self.SAVE_METADATA:
+            self.checkMetadata(loaded)
+        remove(self.MAT_FILE)
+
+    def check(self, loaded):
+        """Check the contents of the data loaded from the .mat file"""
+        for materialName, material in iteritems(self.reader.materials):
+            for variableName, data in iteritems(material.data):
+                expectedName = self.constructExpectedVarName(
+                    materialName, variableName)
+
+                # savemat saves vectors as row vectors by default
+
+                if len(data.shape) == 1:
+                    data = data.reshape(1, data.size)
+                assert expectedName in loaded, (
+                    "Variable {} from material {}"
+                    .format(materialName, variableName))
+                assert_equal(
+                    loaded[expectedName], data,
+                    err_msg="{} {}".format(materialName, variableName))
+
+    def checkMetadata(self, loaded):
+        """Check the contents of the metadata loaded from mat file"""
+        pass
+
+    def constructExpectedVarName(self, material, variable):
+        """
+        Return name of the array in the .mat file for a material and variable
+        """
+        raise NotImplementedError
+
+
+class DepMatlabExportConverter(object):
+    """Mixin class that converts variables back to orgiginal form"""
+
+    RECONVERT = True
+
+    @staticmethod
+    def constructExpectedVarName(material, variable):
+        return deconvert(material, variable)
+
+
+class DepMatlabExportNoConverter(object):
+    """Mixin class that does not convert to original form"""
+
+    RECONVERT = False
+
+    @staticmethod
+    def constructExpectedVarName(material, variable):
+        return prepToMatlab(material, variable)
+
+
+class DepMatlabExportMetadataTester(DepMatlabExportHelper):
+    """Class that sends metadata to the matlab file and checks"""
+
+    SAVE_METADATA = True
+
+    def checkMetadata(self, loaded):
+        for key, data in iteritems(self.reader.metadata):
+            expected = key.upper() if self.RECONVERT else key
+            assert expected in loaded, (
+                "Missing {} from metadata".format(expected))
+
+
+class DepMatlabExportWithoutMetadataTester(DepMatlabExportHelper):
+    """Class that does not send metadata to the matlab file and checks"""
+
+    SAVE_METADATA = False
+
+    def checkMetadata(self, loaded):
+        for key, data in iteritems(self.reader.metadata):
+            assert key not in loaded, (
+                "{} from metadata present and should not be".format(key))
+
+
+class DepMatlabExportConvWithMetadata(
+        DepMatlabExportConverter, DepMatlabExportMetadataTester):
+    """Class that saves a matlab file with converted names and metadata"""
+    pass
+
+
+class DepMatlabExportConvWithoutMetadata(
+        DepMatlabExportConverter, DepMatlabExportWithoutMetadataTester):
+    """Class that saves a matlab file with converted names and no metadata"""
+    pass
+
+
+class DepMatlabExportNoConvWithMetadata(
+        DepMatlabExportNoConverter, DepMatlabExportMetadataTester):
+    """Class that saves a matlab file without converted names and metadata"""
+    pass
+
+
+del DepMatlabExportMetadataTester, DepMatlabExportWithoutMetadataTester
+del DepMatlabExportHelper
+
+
 if __name__ == '__main__':
     from unittest import main
     main()