Repository: CORE-GATECH-GROUP/serpent-tools

Issue:
[Feature] Auto-reshape scattering matrices for homogenized universes
Implement a setting that will automatically reshape scattering matrices for homogenized universes. This should be done in such a way that both the BranchingReader and ResultsReader honor this request.

Possibly a setting like `xs.reshapeScatter`?

Discussion:

Patch:
diff --git a/docs/defaultSettings.rst b/docs/defaultSettings.rst
index aa16202..8893098 100644
--- a/docs/defaultSettings.rst
+++ b/docs/defaultSettings.rst
@@ -206,6 +206,19 @@ If true, store the infinite medium cross sections.
   Type: bool
   
 
+.. _xs-reshapeScatter:
+
+---------------------
+``xs.reshapeScatter``
+---------------------
+
+If true, reshape the scattering matrices to square matrices. By default, these matrices are stored as vectors.
+::
+
+  Default: False
+  Type: bool
+  
+
 .. _xs-variableExtras:
 
 ---------------------
diff --git a/docs/welcome/changelog.rst b/docs/welcome/changelog.rst
index 8628472..62f0e1d 100644
--- a/docs/welcome/changelog.rst
+++ b/docs/welcome/changelog.rst
@@ -12,6 +12,9 @@ Next
 * :pull:`131` Updated variable groups between ``2.1.29`` and ``2.1.30`` - include
   poison cross section, kinetic parameters, six factor formula (2.1.30 exclusive),
   and minor differences
+* :pull:`141` - Setting :ref:`xs-reshapeScatter` can be used to reshape scatter
+  matrices on :py:class:`~serpentTools.objects.containers.HomogUniv` 
+  objects to square matrices
 
 .. _vDeprecated:
 
diff --git a/serpentTools/objects/containers.py b/serpentTools/objects/containers.py
index f9ec978..9626270 100644
--- a/serpentTools/objects/containers.py
+++ b/serpentTools/objects/containers.py
@@ -1,19 +1,22 @@
-""" Custom-built containers for storing data from serpent outputs
+""" 
+Custom-built containers for storing data from serpent outputs
 
 Contents
 --------
-:py:class:`~serpentTools.objects.containers.HomogUniv`
-:py:class:`~serpentTools.objects.containers.BranchContainer
-:py:class:`~serpentTools.objects.containers.DetectorBase`
-:py:class:`~serpentTools.objects.containers.Detector`
+* :py:class:`~serpentTools.objects.containers.HomogUniv`
+* :py:class:`~serpentTools.objects.containers.BranchContainer
+* :py:class:`~serpentTools.objects.containers.DetectorBase`
+* :py:class:`~serpentTools.objects.containers.Detector`
 
 """
 from collections import OrderedDict
+from itertools import product
 
 from matplotlib import pyplot
+from numpy import (array, arange, unique, log, divide, ones_like, hstack,
+                   ndarray)
 
-from numpy import array, arange, unique, log, divide, ones_like, hstack
-
+from serpentTools.settings import rc
 from serpentTools.plot import cartMeshPlot, plot, magicPlotDocDecorator
 from serpentTools.objects import NamedObject, convertVariableName
 from serpentTools.messages import warning, SerpentToolsException, debug
@@ -22,9 +25,16 @@ DET_COLS = ('value', 'energy', 'universe', 'cell', 'material', 'lattice',
             'reaction', 'zmesh', 'ymesh', 'xmesh', 'tally', 'error', 'scores')
 """Name of the columns of the data"""
 
+SCATTER_MATS = set()
+SCATTER_ORDERS = 8
+
+for xsSpectrum, xsType in product({'INF', 'B1'},
+                                  {'S', 'SP'}):
+    SCATTER_MATS.update({'{}_{}{}'.format(xsSpectrum, xsType, xx)
+                        for xx in range(SCATTER_ORDERS)})
 
 __all__ = ('DET_COLS', 'HomogUniv', 'BranchContainer', 'Detector',
-           'DetectorBase')
+           'DetectorBase', 'SCATTER_MATS', 'SCATTER_ORDERS')
 
 
 def isNonNeg(value):
@@ -69,6 +79,11 @@ class HomogUniv(NamedObject):
         Relative uncertainties for leakage-corrected group constants
     metadata: dict
         Other values that do not not conform to inf/b1 dictionaries
+    reshapedMats: bool
+        ``True`` if scattering matrices have been reshaped to square
+        matrices. Otherwise, these matrices are stored as vectors. 
+    numGroups: None or int
+        Number of energy groups
 
     Raises
     ------
@@ -98,6 +113,12 @@ class HomogUniv(NamedObject):
         self.b1Unc = {}
         self.infUnc = {}
         self.metadata = {}
+        self.__reshaped = rc['xs.reshapeScatter']
+        self.numGroups = None
+
+    @property
+    def reshaped(self):
+        return self.__reshaped
 
     def __str__(self):
         extras = []
@@ -113,8 +134,14 @@ class HomogUniv(NamedObject):
                                   extras or '')
     
     def addData(self, variableName, variableValue, uncertainty=False):
-        """
-        sets the value of the variable and, optionally, the associate s.d.
+        r"""
+        Sets the value of the variable and, optionally, the associate s.d.
+
+        .. versionadded:: 0.5.0
+
+            Reshapes scattering matrices according to setting 
+            `xs.reshapeScatter`. Matrices are of the form
+            :math:`S[i, j]=\Sigma_{s,i\rightarrow j}`
 
         .. warning::
 
@@ -127,8 +154,7 @@ class HomogUniv(NamedObject):
         variableValue:
             Variable Value
         uncertainty: bool
-            Set to ``True`` in order to retrieve the
-            uncertainty associated to the expected values
+            Set to ``True`` if this data is an uncertainty 
 
         Raises
         ------
@@ -136,12 +162,31 @@ class HomogUniv(NamedObject):
             If the uncertainty flag is not boolean
 
         """
-
-        # 1. Check the input type
-        variableName = convertVariableName(variableName)
         if not isinstance(uncertainty, bool):
             raise TypeError('The variable uncertainty has type {}, '
                             'should be boolean.'.format(type(uncertainty)))
+        if not isinstance(variableValue, ndarray):
+            debug("Converting {} from {} to array".format(
+                variableName, type(variableValue)))
+            variableValue = array(variableValue)
+        ng = self.numGroups
+        if self.__reshaped and variableName in SCATTER_MATS:
+            if ng is None:
+                warning("Number of groups is unknown at this time. "
+                        "Will not reshape variable {}"
+                        .format(variableName))
+            else:
+                variableValue = variableValue.reshape(ng, ng)
+        incomingGroups = variableValue.shape[0] 
+        if ng is None:
+            self.numGroups = incomingGroups 
+        elif incomingGroups != ng and variableName not in SCATTER_MATS:
+            warning("Variable {} appears to have different group structure. "
+                    "Current: {} vs. incoming: {}"
+                    .format(variableName, ng, incomingGroups))
+
+        variableName = convertVariableName(variableName)
+
         # 2. Pointer to the proper dictionary
         setter = self._lookup(variableName, uncertainty)
         # 3. Check if variable is already present. Then set the variable.
@@ -166,7 +211,7 @@ class HomogUniv(NamedObject):
         x:
             Variable Value
         dx:
-            Associated uncertainty
+            Associated uncertainty if ``uncertainty``
 
         Raises
         ------
diff --git a/serpentTools/parsers/branching.py b/serpentTools/parsers/branching.py
index fa771cb..ffda4bb 100644
--- a/serpentTools/parsers/branching.py
+++ b/serpentTools/parsers/branching.py
@@ -134,6 +134,10 @@ class BranchingReader(XSReader):
                 possibleEndOfFile=step == numVariables - 1)
             varName = splitList[0]
             varValues = [float(xx) for xx in splitList[2:]]
+            if not varValues:
+                debug("No data present for variable {}. Skipping"
+                      .format(varName))
+                continue
             if self._checkAddVariable(varName):
                 if self.settings['areUncsPresent']:
                     vals, uncs = splitItems(varValues)
diff --git a/serpentTools/settings.py b/serpentTools/settings.py
index a7ca265..745eea6 100644
--- a/serpentTools/settings.py
+++ b/serpentTools/settings.py
@@ -124,6 +124,12 @@ defaultSettings = {
         'description': 'If true, store the critical leakage cross sections.',
         'type': bool
     },
+    'xs.reshapeScatter': {
+        'default': False,
+        'description': 'If true, reshape the scattering matrices to square matrices. '
+                       'By default, these matrices are stored as vectors.',
+       'type': bool
+    },
     'xs.variableGroups': {
         'default': [],
         'description': ('Name of variable groups from variables.yaml to be '
@@ -254,10 +260,6 @@ class UserSettingsLoader(dict):
         self.__originals = {}
         dict.__init__(self, self._defaultLoader.retrieveDefaults())
 
-    def __setitem__(self, key, value):
-        self._checkStoreOriginal(key)
-        self.setValue(key, value)
-
     def __enter__(self):
         self.__inside= True
         return self
@@ -268,10 +270,6 @@ class UserSettingsLoader(dict):
             self[key] = originalValue
         self.__originals= {}
 
-    def _checkStoreOriginal(self, key):
-        if self.__inside:
-            self.__originals[key] = self[key]
-
     def setValue(self, name, value):
         """Set the value of a specific setting.
 
@@ -290,6 +288,8 @@ class UserSettingsLoader(dict):
             If the value is not of the correct type
 
         """
+        if self.__inside:
+            self.__originals[name] = self[name]
         if name not in self:
             raise KeyError('Setting {} does not exist'.format(name))
         self._defaultLoader[name].validate(value)
@@ -299,6 +299,8 @@ class UserSettingsLoader(dict):
         dict.__setitem__(self, name, value)
         messages.debug('Updated setting {} to {}'.format(name, value))
 
+    __setitem__ = setValue
+
     def getReaderSettings(self, settingsPreffix):
         """Get all module-wide and reader-specific settings.
 


Test patch:
diff --git a/serpentTools/tests/test_container.py b/serpentTools/tests/test_container.py
index fbed4dc..25a1188 100644
--- a/serpentTools/tests/test_container.py
+++ b/serpentTools/tests/test_container.py
@@ -1,43 +1,43 @@
 """Test the container object. """
 
 import unittest
-from serpentTools.objects import containers
+
+from six import iteritems
+from numpy import array, arange
+from numpy.testing import assert_allclose
+
+from serpentTools.settings import rc
+from serpentTools.objects.containers import HomogUniv
 from serpentTools.parsers import DepletionReader
 
+NUM_GROUPS = 5
 
-class HomogenizedUniverseTester(unittest.TestCase):
-    """ Class to test the Homogenized Universe """
 
-    @classmethod
-    def setUpClass(cls):
-        cls.univ = containers.HomogUniv('dummy', 0, 0, 0)
-        cls.Exp = {}
-        cls.Unc = {}
-        # Data definition
-        cls.Exp = {'B1_1': 1, 'B1_2': [1, 2], 'B1_3': [1, 2, 3],
-                   'INF_1': 3, 'INF_2': [4, 5], 'INF_3': [6, 7, 8],
-                   'MACRO_E': (.1, .2, .3, .1, .2, .3), 'Test_1': 'ciao'}
-        cls.Unc = {'B1_1': 1e-1, 'B1_2': [1e-1, 2e-1], 'B1_3': [1e-1, 2e-1,
-                                                                3e-1],
-                   'INF_1': 3e-1, 'INF_2': [4e-1, 5e-1], 'INF_3': [6e-1, 7e-1,
-                                                                   8e-1],
-                   'MACRO_E': (.1e-1, .2e-1, .3e-1, .1e-1, .2e-1, .3e-1),
-                   'Test_1': 'addio'}
+class _HomogUnivTestHelper(unittest.TestCase):
+    """Class that runs the tests for the two sub-classes
+    
+    Subclasses will differ in how the ``mat`` data
+    is arranged. For one case, the ``mat`` will be a 
+    2D matrix.
+    """
 
+    def setUp(self):
+        self.univ, vec, mat = self.getParams()
+        # Data definition
+        rawData = {'B1_1': vec, 'B1_AS_LIST': list(range(NUM_GROUPS)),
+                  'INF_1': vec, 'INF_S0': mat}
+        meta = {'MACRO_E': vec}
         # Partial dictionaries
-        cls.b1Exp = {'b11': 1, 'b12': [1, 2], 'b13': [1, 2, 3]}
-        cls.b1Unc = {'b11': (1, 0.1), 'b12': ([1, 2], [.1, .2]), 'b13':
-            ([1, 2, 3], [.1, .2, .3])}
-        cls.infExp = {'inf1': 3, 'inf2': [4, 5], 'inf3': [6, 7, 8]}
-        cls.infUnc = {'inf1': (3, .3), 'inf2': ([4, 5], [.4, .5]), 'inf3':
-            ([6, 7, 8], [.6, .7, .8])}
-        cls.meta = {'macroE': (.1e-1, .2e-1, .3e-1, .1e-1, .2e-1, .3e-1),
-                    'test1': 'addio'}
+        self.b1Unc = self.b1Exp = {'b11': vec}
+        self.infUnc = self.infExp = {'inf1': vec, 'infS0': mat}
+        self.meta = {'macroE': vec}
+
         # Use addData
-        for kk in cls.Exp:
-            cls.univ.addData(kk, cls.Exp[kk], False)
-        for kk in cls.Unc:
-            cls.univ.addData(kk, cls.Unc[kk], True)
+        for key, value in iteritems(rawData):
+            self.univ.addData(key, value, uncertainty=False)
+            self.univ.addData(key, value, uncertainty=True)
+        for key, value in iteritems(meta):
+            self.univ.addData(key, value)
 
     def test_getB1Exp(self):
         """ Get Expected vales from B1 dictionary"""
@@ -45,16 +45,16 @@ class HomogenizedUniverseTester(unittest.TestCase):
         # Comparison
         for kk in self.univ.b1Exp:
             d[kk] = self.univ.get(kk, False)
-        self.assertDictEqual(self.b1Exp, d)
-
+        compareDictOfArrays(self.b1Exp, d, 'b1 values')
+        
     def test_getB1Unc(self):
         """ Get Expected vales and associated uncertainties from B1 dictionary
         """
         d = {}
         # Comparison
         for kk in self.univ.b1Exp:
-            d[kk] = self.univ.get(kk, True)
-        self.assertDictEqual(self.b1Unc, d)
+            d[kk] = self.univ.get(kk, True)[1]
+        compareDictOfArrays(self.b1Unc, d, 'b1 uncertainties')
 
     def test_getInfExp(self):
         """ Get Expected vales from Inf dictionary"""
@@ -62,7 +62,7 @@ class HomogenizedUniverseTester(unittest.TestCase):
         # Comparison
         for kk in self.univ.infExp:
             d[kk] = self.univ.get(kk, False)
-        self.assertDictEqual(self.infExp, d)
+        compareDictOfArrays(self.infExp, d, 'infinite values')
 
     def test_getInfUnc(self):
         """ Get Expected vales and associated uncertainties from Inf dictionary
@@ -70,8 +70,8 @@ class HomogenizedUniverseTester(unittest.TestCase):
         d = {}
         # Comparison
         for kk in self.univ.infUnc:
-            d[kk] = self.univ.get(kk, True)
-        self.assertDictEqual(self.infUnc, d)
+            d[kk] = self.univ.get(kk, True)[1]
+        compareDictOfArrays(self.infUnc, d, 'infinite uncertainties')
 
     def test_getMeta(self):
         """ Get metaData from corresponding dictionary"""
@@ -79,8 +79,64 @@ class HomogenizedUniverseTester(unittest.TestCase):
         # Comparison
         for kk in self.univ.metadata:
             d[kk] = self.univ.get(kk, False)
-        self.assertDictEqual(self.meta, d)
+        compareDictOfArrays(self.meta, d, 'metadata')
+
+    def test_getBothInf(self):
+        """
+        Verify that the value and the uncertainty are returned if the
+        flag is passed.
+        """
+        expected, uncertainties = {}, {}
+        for key in self.infExp.keys():
+            value, unc = self.univ.get(key, True)
+            expected[key] = value
+            uncertainties[key] = unc
+        compareDictOfArrays(self.infExp, expected, 'infinite values')
+        compareDictOfArrays(self.infUnc, uncertainties, 
+                            'infinite uncertainties')
+
+
+class VectoredHomogUnivTester(_HomogUnivTestHelper):
+    """Class for testing HomogUniv that does not reshape scatter matrices"""
+
+    def getParams(self):
+        univ, vec, mat = getParams()
+        self.assertFalse(univ.reshaped)
+        return univ, vec, mat
+
+
+class ReshapedHomogUnivTester(_HomogUnivTestHelper):
+    """Class for testing HomogUniv that does reshape scatter matrices"""
+
+    def getParams(self):
+        from serpentTools.settings import rc
+        with rc:
+            rc.setValue('xs.reshapeScatter', True)
+            univ, vec, mat = getParams()
+            univ.numGroups = NUM_GROUPS
+            self.assertTrue(univ.reshaped)
+        return univ, vec, mat.reshape(NUM_GROUPS, NUM_GROUPS)
+
+
+def getParams():
+    """Return the universe, vector, and matrix for testing."""
+    univ = HomogUniv(300, 0, 0, 0)
+    vec = arange(NUM_GROUPS)
+    mat = arange(NUM_GROUPS ** 2)
+    return univ, vec, mat
+
+
+def compareDictOfArrays(expected, actualDict, dataType):
+    for key, value in iteritems(expected):
+        actual = actualDict[key]
+        assert_allclose(value, actual, 
+                err_msg="Error in {} dictionary: key={}"
+                .format(dataType, key))
 
+del _HomogUnivTestHelper
 
 if __name__ == '__main__':
-    unittest.main()
+    from serpentTools import rc
+    with rc:
+        rc['verbosity'] = 'debug'
+        unittest.main()