Repository: CORE-GATECH-GROUP/serpent-tools
Pull request #276 (2019-02-12 17:21:17)

API Easier construction of BranchCollector
Three changes that help create BranchCollector objects:

1. Knowing the perturbations at call to `collect` is not necessary. If we assume that all the branches have the same number of perturbations, then we can determine this information from the BranchReader. This allows the user to simply do `BranchCollector.collect()` and then set the perturbations later.
2. class method `BranchCollector.fromFile` can be used to create a collector and process all group constants, e.g.
```
>>> from serpentTools import BranchCollector
>>> col = BranchCollector.fromFile('myfile.coe')
```
passing the perturbations to this method is also allowed
3. Pass a file path to the instantiation method. This will create a BranchingReader and read the file with that.
```
>>> from serpentTools import BranchCollector
>>> col = BranchCollector('myfile.coe')
```


Patch:
diff --git a/serpentTools/__init__.py b/serpentTools/__init__.py
index c27b59c..6be9854 100644
--- a/serpentTools/__init__.py
+++ b/serpentTools/__init__.py
@@ -3,6 +3,7 @@ from serpentTools.parsers import *
 from serpentTools.data import *
 from serpentTools.samplers import *
 from serpentTools.seed import *
+from serpentTools.xs import *
 from ._version import get_versions
 __version__ = get_versions()['version']
 del get_versions
diff --git a/serpentTools/xs.py b/serpentTools/xs.py
index ad053e1..255e5fb 100644
--- a/serpentTools/xs.py
+++ b/serpentTools/xs.py
@@ -10,6 +10,8 @@ from six import iteritems
 from six.moves import range
 from numpy import empty, nan, array, ndarray
 
+from serpentTools import BranchingReader
+
 __all__ = [
     'BranchCollector',
 ]
@@ -196,8 +198,10 @@ class BranchCollector(object):
 
     Parameters
     ----------
-    reader: :class:`serpentTools.parsers.branching.BranchingReader`
-        Read data from the passed reader
+    source: str or :class:`~serpentTools.parsers.branching.BranchingReader`
+        Coefficient file to be read, or a
+        :class:`~serpentTools.parsers.branching.BranchingReader`
+        that has read the file of interest
 
     Attributes
     ----------
@@ -222,9 +226,15 @@ class BranchCollector(object):
         '_perturbations', 'univIndex', '_burnups', '_states', '_shapes',
     )
 
-    def __init__(self, reader):
+    def __init__(self, source):
         warn("This is an experimental feature, subject to change.",
              UserWarning)
+        if isinstance(source, BranchingReader):
+            reader = source
+        else:
+            # assume file path or readable
+            reader = BranchingReader(source)
+            reader.read()
         self.filePath = reader.filePath
         self._branches = reader.branches
         self.xsTables = {}
@@ -432,23 +442,43 @@ class BranchCollector(object):
                 .format(self._burnups.size, value.size))
         self._burnups = value
 
-    def collect(self, perturbations):
+    def collect(self, perturbations=None):
         """
         Parse the contents of the file and collect cross sections
 
         Parameters
         ----------
-        perturbations: tuple
+        perturbations: tuple or None
             Tuple where each entry is a state that is perturbed across
             the analysis, e.g. ``("Tfuel", "RhoCool", "CR")``. These
             must appear in the same order as they are ordered in the
-            coefficient file.
+            coefficient file. If ``None``, then the number of
+            perturbations will be determined from the coefficient
+            file. This is used to set :attr:`pertubations` and
+            can be adjusted later
         """
-        if (isinstance(perturbations, str)
-           or not isinstance(perturbations, Iterable)):
-            self._perturbations = perturbations,
+        # get number of perturbations from number of
+        # items in keys of reader branches
+        for key in self._branches:
+            break
+        if isinstance(key, str):
+            nReaderPerts = 1
+        else:
+            nReaderPerts = len(key)
+
+        if perturbations is None:
+            perturbations = tuple(
+                ['p' + str(ii) for ii in range(nReaderPerts)])
+        elif (isinstance(perturbations, str)
+                or not isinstance(perturbations, Iterable)):
+            perturbations = perturbations,
         else:
-            self._perturbations = tuple(perturbations)
+            perturbations = tuple(perturbations)
+
+        assert len(perturbations) == nReaderPerts, "{} vs {}".format(
+            len(perturbations), nReaderPerts)
+
+        self._perturbations = perturbations
         sampleBranchKey = self._getBranchStates()
         sampleUniv = self._getUnivsBurnups(sampleBranchKey)
         xsSizes = self._getXsSizes(sampleUniv)
@@ -546,3 +576,29 @@ class BranchCollector(object):
         for xsKey, xsMat in iteritems(self.xsTables):
             for univIndex, univID in enumerate(self.univIndex):
                 self.universes[univID][xsKey] = xsMat[univIndex]
+
+    @classmethod
+    def fromFile(cls, filePath, perturbations=None):
+        """
+        Create a :class:`BranchCollector` from the contents of the file
+
+        Parameters
+        ----------
+        filePath: str
+            Location of coefficient file to be read
+        perturbations: None or iterable
+            Ordering of perturbation types in coefficient file.
+            If ``None``, the number of perturbations will be inferred
+            from file. Otherwise, the number of perturbations must
+            match those in the file. This value can be changed
+            after the fact using :attr:`perturbations`,
+            with insight gained from :attr:`states`
+
+        Returns
+        -------
+        :class:`BranchCollector` object that has processed the contents
+        of the file.
+        """
+        collector = BranchCollector(filePath)
+        collector.collect(perturbations)
+        return collector


Test patch:
diff --git a/serpentTools/tests/test_xs.py b/serpentTools/tests/test_xs.py
index aa27f11..ee3e31c 100644
--- a/serpentTools/tests/test_xs.py
+++ b/serpentTools/tests/test_xs.py
@@ -68,10 +68,12 @@ Group constant: {g}
 Universe: {u}
 Burnup: {b}""".strip()
 
+    ORIG_PERTS = ("BOR", "TFU")
+
     def collect(self):
         """Return the collector after collecting group constants"""
         collector = BranchCollector(self.reader)
-        collector.collect(("BOR", "TFU"))
+        collector.collect(self.ORIG_PERTS)
         return collector
 
     def getGroupConstant(self, pertVals, univID, gcKey, burnup, msg):
@@ -198,6 +200,31 @@ class BranchUnivTester(CollectedHelper):
         return xsTable[pertIndex + (burnupIndex, )]
 
 
+class UnknownPertCollectorTester(BranchCollectorTester):
+    """Helper for creating BranchCollectors without perturbations"""
+
+    def collect(self):
+        collector = BranchCollector(self.reader)
+        collector.collect()
+        return collector
+
+    def setUp(self):
+        BranchCollectorTester.setUp(self)
+        self.assertTrue(
+            len(self.collector.perturbations) == len(self.ORIG_PERTS),
+            msg="Failed to infer correct perturbations"
+        )
+
+
+class CollectorFromFileTester(UnknownPertCollectorTester):
+    """Test the collection of the data using the fromFile method"""
+
+    def collect(self):
+        """Overwrite to use the fromFile method"""
+        collector = BranchCollector.fromFile(self.reader.filePath)
+        return collector
+
+
 del CollectedHelper