Repository: CORE-GATECH-GROUP/serpent-tools (python)

Issue:
[ENH] Allow the DepletionReader to capture materials with underscores
Currently, the `DepletionReader` doesn't capture materials with underscores. This is due to the regular expressions used to match the material name and variable name. With variables like `ING_TOX`, the underscores throw off the engine.

Pull request description:
ENH Capture depleted material with underscores 
Closes #58 by allowing the DepletionReader to work with materials containing underscores. Had to move away from the regular expression parsing because some variables contain underscores, while some don't. Added a bunch of variables into `serpentTools/parsers/_collections.py`. These variables are used in the process of figuring out what to strip out from the variable name to obtain the actual name of the material.

Added some tests to ensure that we can indeed do this, as well as some tests for 
utilities that break up the material names.

- [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
- [x] New and/or changed features are [well documented](http://serpent-tools.readthedocs.io/en/latest/develop/documentation.html#documentation) and have adequate testing

## TODO
- [x] Add this PR to changelog
- [x] Find places where we say we can't do this and remove them

Patch:
diff --git a/docs/changelog.rst b/docs/changelog.rst
index 7a2a43f..8ce106c 100644
--- a/docs/changelog.rst
+++ b/docs/changelog.rst
@@ -32,7 +32,7 @@ Changelog
   |depmtxReader|
 * :pull:`252` - Better axis and colorbar labeling for |detector| mesh plots
 * :pull:`254` - Better plotting of single concentrations with |depmtxReader|
-
+* :pull:`255` - |depletionReader| can capture material with underscores now!
 
 Deprecations
 ------------
diff --git a/docs/examples/DepletionReader.rst b/docs/examples/DepletionReader.rst
index 9bd20b7..9427c2f 100644
--- a/docs/examples/DepletionReader.rst
+++ b/docs/examples/DepletionReader.rst
@@ -296,12 +296,6 @@ replacements
 .. image:: DepletionReader_files/DepletionReader_26_0.png
 
 
-Limitations
------------
-
-Currently, the |depReader| cannot catch materials with underscore in the name,
-due to variables like ``ING_TOX`` also containing an underscore.
-Issue `#58 <https://github.com/CORE-GATECH-GROUP/serpent-tools/issues/58>`_
 
 .. _depletion-settings:
 
diff --git a/examples/DepletionReader.ipynb b/examples/DepletionReader.ipynb
index 37a10b8..0fdcc00 100644
--- a/examples/DepletionReader.ipynb
+++ b/examples/DepletionReader.ipynb
@@ -458,14 +458,6 @@
     "         labelFmt=\"{mat}: {iso} // {zai}\", loglog=True);"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Limitations\n",
-    "Currently, the `DepletionReader` cannot catch materials with underscore in the name, due to variables like `ING_TOX` also containing underscore. This is currently an open issue [#58](https://github.com/CORE-GATECH-GROUP/serpent-tools/issues/58)"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
diff --git a/serpentTools/parsers/_collections.py b/serpentTools/parsers/_collections.py
index 0b93156..3eaaf89 100644
--- a/serpentTools/parsers/_collections.py
+++ b/serpentTools/parsers/_collections.py
@@ -95,3 +95,20 @@ RES_DATA_NO_UNCS = {
 Set containing keys for objects stored in :attr:`ResultsReader.resdata`
 that do not contain uncertainties.
 """
+
+DEPLETION_VARIABLES = {
+    'ING_TOX',
+    'INH_TOX',
+    'VOLUME',
+    'BURNUP',
+    'ADENS',
+    'MDENS',
+    'H',
+    'A',
+    'SF',
+    'GSRC',
+    'MASS',
+}
+"""
+Set of all variables stored for material in depletion output
+"""
diff --git a/serpentTools/parsers/depletion.py b/serpentTools/parsers/depletion.py
index 528929d..5614d93 100644
--- a/serpentTools/parsers/depletion.py
+++ b/serpentTools/parsers/depletion.py
@@ -8,7 +8,7 @@ from serpentTools.utils import (
     magicPlotDocDecorator, formatPlot, DEPLETION_PLOT_LABELS,
     convertVariableName, str2vec,
 )
-
+from ._collections import DEPLETION_VARIABLES
 from serpentTools.engines import KeywordParser
 from serpentTools.parsers.base import MaterialReader
 from serpentTools.objects.materials import DepletedMaterial
@@ -147,30 +147,15 @@ class DepletionReader(DepPlotMixin, MaterialReader):
 
     def __init__(self, filePath):
         MaterialReader.__init__(self, filePath, 'depletion')
-        self._matPatterns = self._makeMaterialRegexs()
-        self._matchMatNVar = r'[A-Z]{3}_([0-9a-zA-Z]*)_([A-Z]*_?[A-Z]*)'
-        # Captures material name and variable from string
-        #  MAT_fuel1_ADENS --> ('fuel1', 'ADENS')
-        #  MAT_fUeL1g_ING_TOX --> ('fUeL1', 'ING_TOX')
-
-        self._matchTotNVar = r'[A-Z]{3}_([A-Z]*_?[A-Z]*)'
-        # Captures variables for total block from string::
-        #  TOT_ADENS --> ('ADENS', )
-        #  ING_TOX --> ('ING_TOX', )
+        patterns = self.settings['materials'] or ['.*']
+        # match all materials if nothing given
+        self._matPatterns = [re.compile(mat) for mat in patterns]
         DepPlotMixin.__init__(self)
 
     def __getitem__(self, name):
         """Retrieve a material from :attr:`materials`."""
         return self.materials[name]
 
-    def _makeMaterialRegexs(self):
-        """Return the patterns by which to find the requested materials."""
-        patterns = self.settings['materials'] or ['.*']
-        # match all materials if nothing given
-        if any(['_' in pat for pat in patterns]):
-            warning('Materials with underscores are not supported.')
-        return [re.compile(mat) for mat in patterns]
-
     def _read(self):
         """Read through the depletion file and store requested data."""
         keys = ['MAT', 'TOT'] if self.settings['processTotal'] else ['MAT']
@@ -178,12 +163,12 @@ class DepletionReader(DepPlotMixin, MaterialReader):
         separators = ['\n', '];', '\r\n']
         with KeywordParser(self.filePath, keys, separators) as parser:
             for chunk in parser.yieldChunks():
-                if 'MAT' in chunk[0]:
-                    self._addMaterial(chunk)
-                elif 'TOT' in chunk[0]:
-                    self._addTotal(chunk)
-                else:
-                    self._addMetadata(chunk)
+                mvar = getMatlabVarName(chunk)
+                if mvar[:3] in ['MAT', 'TOT']:
+                    name, variable = getMaterialNameAndVariable(mvar)
+                    self._checkAddData(chunk, name, variable)
+                    continue
+                self._addMetadata(chunk)
         if 'days' in self.metadata:
             for mKey in self.materials:
                 self.materials[mKey].days = self.metadata['days']
@@ -210,23 +195,14 @@ class DepletionReader(DepPlotMixin, MaterialReader):
     def _cleanSingleLine(chunk):
         return chunk[0][chunk[0].index('[') + 1:chunk[0].index(']')]
 
-    def _addMaterial(self, chunk):
-        """Add data from a MAT chunk."""
-        name, variable = self._getGroupsFromChunk(self._matchMatNVar, chunk)
-        if any([re.match(pat, name) for pat in self._matPatterns]):
-            self._processChunk(chunk, name, variable)
-
-    def _addTotal(self, chunk):
-        """Add data from a TOT chunk"""
-        variable = self._getGroupsFromChunk(self._matchTotNVar, chunk)[0]
-        self._processChunk(chunk, 'total', variable)
-
-    def _getGroupsFromChunk(self, regex, chunk):
-        match = re.match(regex, chunk[0])
-        if match:
-            return match.groups()
-        raise Exception('{} not determine match from the following chunk:\n'
-                        '{}'.format(self, ''.join(chunk)))
+    def _checkAddData(self, chunk, name, variable):
+        if name == 'total':
+            if not self.settings['processTotal']:
+                return
+            return self._processChunk(chunk, name, variable)
+        for pattern in self._matPatterns:
+            if pattern.match(name):
+                return self._processChunk(chunk, name, variable)
 
     def _processChunk(self, chunk, name, variable):
         if (self.settings['materialVariables']
@@ -377,3 +353,24 @@ class DepletionReader(DepPlotMixin, MaterialReader):
             self.metadata['burnup'], other.metadata['burnup'],
             lower, upper, 'burnup')
         return similar
+
+
+def getMaterialNameAndVariable(mlabName):
+    for serpVar in DEPLETION_VARIABLES:
+        leng = len(serpVar)
+        if serpVar == mlabName[-len(serpVar):]:
+            matName = mlabName[4:-(leng + 1)]
+            if not matName:
+                if 'TOT' == mlabName[:3]:
+                    return 'total', serpVar
+                break
+            return matName, serpVar
+    raise ValueError(
+        "Could not find a match for {}. "
+        "Please alert developers".format(mlabName))
+
+
+def getMatlabVarName(chunk):
+    if isinstance(chunk, list):
+        return getMatlabVarName(chunk[0])
+    return chunk[:chunk.index(' ')]


Test patch:
diff --git a/serpentTools/tests/test_depletion.py b/serpentTools/tests/test_depletion.py
index baa35e4..32d91a8 100644
--- a/serpentTools/tests/test_depletion.py
+++ b/serpentTools/tests/test_depletion.py
@@ -1,43 +1,77 @@
 """Test the depletion file."""
-import unittest
+from unittest import TestCase
+from os import remove
 
 from numpy import array
 from numpy.testing import assert_equal
 
 from six import iteritems
 
-from serpentTools.parsers import read
 from serpentTools.data import getFile
 from serpentTools.settings import rc
-from serpentTools.parsers.depletion import DepletionReader
+from serpentTools.parsers.depletion import (
+    DepletionReader, getMaterialNameAndVariable, getMatlabVarName,
+)
+from serpentTools.tests.utils import LoggerMixin
 
 
-class _DepletionTestHelper(unittest.TestCase):
+DEP_FILE = 'ref_dep.m'
+DEP_FILE_PATH = getFile(DEP_FILE)
+
+FILE_WITH_UNDERSCORES = 'underscores_dep.m'
+ORIG_FUEL = "fuel"
+NEW_FUEL_NAME = "fuel_0"
+
+
+def setUpModule():
+    """
+    Set up the module
+
+    *. Create a new file with underscores
+    """
+    with open(DEP_FILE_PATH) as incoming:
+        with open(FILE_WITH_UNDERSCORES, 'w') as out:
+            for line in incoming:
+                out.write(line.replace(ORIG_FUEL, NEW_FUEL_NAME))
+
+
+def tearDownModule():
+    """
+    Tear down the module
+
+    *. Remove new file with underscores
+    """
+    remove(FILE_WITH_UNDERSCORES)
+
+
+class _DepletionTestHelper(TestCase):
     """Base class to setup the depletion reader and material tests."""
 
+    PROCESS_TOTAL = True
+    MATERIAL = ORIG_FUEL
+    FILE = DEP_FILE_PATH
+    EXPECTED_VARIABLES = ['BURNUP', 'ADENS', 'ING_TOX']
+
     @classmethod
     def setUpClass(cls):
-        cls.filePath = getFile('ref_dep.m')
-        cls.processTotal = True
-        cls.materials = ['fuel']
-        cls.expectedMaterials = set(cls.materials)
-        cls.expectedVariables = {'BURNUP', 'ADENS', 'ING_TOX'}
-        if cls.processTotal:
+        # Silence the debugger but we don't care about checking messages
+        logger = LoggerMixin()
+        logger.attach()
+
+        cls.expectedMaterials = {cls.MATERIAL, }
+        if cls.PROCESS_TOTAL:
             cls.expectedMaterials.add('total')
+        # Set some settings to control the output
         with rc as tempRC:
             tempRC['verbosity'] = 'debug'
-            tempRC['depletion.processTotal'] = True
-            tempRC['depletion.materials'] = ['fuel', ]
-            tempRC['depletion.materialVariables'] = [
-                'BURNUP', 'ADENS', 'ING_TOX']
-            tempRC['depletion.processTotal'] = cls.processTotal
-            tempRC['depletion.materials'] = cls.materials
+            tempRC['depletion.materials'] = [cls.MATERIAL, ]
+            tempRC['depletion.processTotal'] = cls.PROCESS_TOTAL
             tempRC['depletion.materialVariables'] = (
-                list(cls.expectedVariables)
+                cls.EXPECTED_VARIABLES
             )
-            cls.reader = DepletionReader(cls.filePath)
-        cls.reader.read()
-        read(cls.filePath, DepletionReader)
+            cls.reader = DepletionReader(cls.FILE)
+            cls.reader.read()
+        logger.detach()
 
 
 class DepletionTester(_DepletionTestHelper):
@@ -82,16 +116,14 @@ class DepletionTester(_DepletionTestHelper):
 class DepletedMaterialTester(_DepletionTestHelper):
     """Class that tests the functionality of the DepletedMaterial class"""
 
-    @classmethod
-    def setUpClass(cls):
-        _DepletionTestHelper.setUpClass()
-        cls.material = cls.reader.materials['fuel']
-        cls.requestedDays = [0.0, 0.5, 1.0, 5.0, 10.0, 25.0, 50.0]
-        cls.fuelBU = cls.material.burnup
+    def setUp(self):
+        self.material = self.reader.materials[self.MATERIAL]
+        self.requestedDays = [0.0, 0.5, 1.0, 5.0, 10.0, 25.0, 50.0]
+        self.fuelBU = self.material.burnup
 
     def test_materials(self):
         """Verify the materials are read in properly."""
-        self.assertIn('fuel', self.reader.materials)
+        self.assertIn(self.MATERIAL, self.reader.materials)
         expectedAdens = array([
             [0.00000E+00, 2.44791E-10, 1.07741E-09, 7.54422E-09, 1.54518E-08,
              2.45253E-08, 3.05523E-08, 3.98843E-08, 4.28827E-08, 4.37783E-08,
@@ -191,5 +223,43 @@ class DepletedMaterialTester(_DepletionTestHelper):
                            names=['Xe135', 'U235'])
 
 
+class UnderscoreDepMaterialTester(DepletedMaterialTester):
+    """Class that reads from a file with underscored fuel names"""
+
+    MATERIAL = NEW_FUEL_NAME
+    FILE = FILE_WITH_UNDERSCORES
+
+
+class DepletionUtilTester(TestCase):
+    """
+    Test case that tests some utilities used by the DepletionReader
+    """
+
+    BAD_0 = "MAT_FUEL_DOES_NOT_EXIST"
+    BAD_1 = "BAD_FUEL_ADENS"
+
+    VAR_CHUNK = ['{} = [\n'.format(BAD_0), ]
+
+    def test_getMaterialAndVariable(self):
+        """Test the function for extracting names and variables"""
+        valid = {
+            'MAT_fuel_ADENS': ('fuel', 'ADENS'),
+            'TOT_ING_TOX': ('total', 'ING_TOX'),
+            'MAT_FUEL_2_H': ('FUEL_2', 'H'),
+        }
+        for matVar, (expName, expVar) in iteritems(valid):
+            actName, actVar = getMaterialNameAndVariable(matVar)
+            self.assertEqual(expName, actName, msg=matVar)
+            self.assertEqual(expVar, actVar, msg=matVar)
+        with self.assertRaises(ValueError):
+            getMaterialNameAndVariable(self.BAD_0)
+            getMaterialNameAndVariable(self.BAD_1)
+
+    def test_getMatlabVarName(self):
+        self.assertEqual(self.BAD_0, getMatlabVarName(self.VAR_CHUNK))
+        self.assertEqual(self.BAD_0, getMatlabVarName(self.VAR_CHUNK[0]))
+
+
 if __name__ == '__main__':
-    unittest.main()
+    from unittest import main
+    main()