API Easier construction of BranchCollector

serpentTools/__init__.py

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

serpentTools/tests/test_xs.py

@@ -68,10 +68,12 @@ class CollectedHelper(_BranchTestHelper):
 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 @@ def _getGroupConstant(self, pertVals, univID, gcKey, burnup):
         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
 
 

serpentTools/xs.py

@@ -10,6 +10,8 @@
 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 @@ def burnups(self, value):
                 .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 @@ def _populateUniverses(self):
         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