Fix PyAPI HexLattice.get_unique_universes()

closes #413

openmc/checkvalue.py

@@ -1,3 +1,5 @@
+from collections import Iterable
+
 def check_type(name, value, expected_type, expected_iter_type=None):
     """Ensure that an object is of an expected type. Optionally, if the object is
     iterable, check that each element is of a particular type.
@@ -30,6 +32,82 @@ def check_type(name, value, expected_type, expected_iter_type=None):
                 raise ValueError(msg)
 
 
+def check_iterable_type(name, value, expected_type, min_depth=1, max_depth=1):
+    """Ensure that an object is an iterable containing an expected type.
+
+    Parameters
+    ----------
+    name : str
+        Description of value being checked
+    value : Iterable
+        Iterable, possibly of other iterables, that should ultimately contain
+        the expected type
+    expected_type : type
+        type that the iterable should contain
+    min_depth : int
+        The minimum number of layers of nested iterables there should be before
+        reaching the ultimately contained items
+    max_depth : int
+        The maximum number of layers of nested iterables ...
+    """
+    # Initialize the tree at the very first item.
+    tree = [value]
+    index = [0]
+
+    # Traverse the tree.
+    while index[0] != len(tree[0]):
+        # If we are done with this level of the tree, go to the next branch on
+        # the level above this one.
+        if index[-1] == len(tree[-1]):
+            del index[-1]
+            del tree[-1]
+            index[-1] += 1
+            continue
+
+        # Get a string representation of the current index in case we raise an
+        # exception.
+        form = '[' + '{:d}, '*(len(index)-1) + '{:d}]'
+        ind_str = form.format(*index)
+
+        # What is the current item we are looking at?
+        current_item = tree[-1][index[-1]]
+
+        # If this item is of the expected type, then we've reached the bottom
+        # level of this branch.
+        if isinstance(current_item, expected_type):
+            # Is this deep enough?
+            if len(tree) < min_depth:
+                msg = 'Error setting {0}: The item at {1} does not meet the ' \
+                      'minimum depth of {2}'.format(name, ind_str, min_depth)
+                raise ValueError(msg)
+
+            # This item is okay.  Move on to the next item.
+            index[-1] += 1
+
+        # If this item is not of the expected type, then it's either an error or
+        # another level of the tree that we need to pursue deeper.
+        else:
+            if isinstance(current_item, Iterable):
+                # The tree goes deeper here, let's explore it.
+                tree.append(current_item)
+                index.append(0)
+
+                # But first, have we exceeded the max depth?
+                if len(tree) > max_depth:
+                    msg = 'Error setting {0}: Found an iterable at {1}, items '\
+                          'in that iterable excceed the maximum depth of {2}' \
+                          .format(name, ind_str, max_depth)
+                    raise ValueError(msg)
+
+            else:
+                # This item is completely unexected.
+                msg = "Error setting {0}: Items must be of type '{1}', but " \
+                      "item at {2} is of type '{3}'"\
+                      .format(name, expected_type.__name__, ind_str,
+                              type(current_item).__name__)
+                raise ValueError(msg)
+
+
 def check_length(name, value, length_min, length_max=None):
     """Ensure that a sized object has length within a given range.
 

openmc/universe.py

@@ -7,7 +7,8 @@ import sys
 import numpy as np
 
 import openmc
-from openmc.checkvalue import check_type, check_length, check_greater_than
+from openmc.checkvalue import check_type, check_iterable_type, check_length, \
+                              check_greater_than
 
 if sys.version_info[0] >= 3:
     basestring = str
@@ -687,8 +688,11 @@ class Lattice(object):
 
     @universes.setter
     def universes(self, universes):
-        check_type('lattice universes', universes, Iterable)
-        self._universes = np.asarray(universes, dtype=Universe)
+        #check_type('lattice universes', universes, Iterable)
+        check_iterable_type('lattice universes', universes, Universe,
+                            min_depth=2, max_depth=3)
+        self._universes = universes
+        #self._universes = np.asarray(universes, dtype=Universe)
 
     def get_unique_universes(self):
         """Determine all unique universes in the lattice
@@ -701,13 +705,24 @@ class Lattice(object):
 
         """
 
-        unique_universes = np.unique(self._universes.ravel())
-        universes = {}
+        univs = dict()
+        for k in range(len(self._universes)):
+            for j in range(len(self._universes[k])):
+                if isinstance(self._universes[k][j], Universe):
+                    u = self._universes[k][j]
+                    if u._id not in univs:
+                        univs[u._id] = u
+                else:
+                    for i in range(len(self._universes[k][j])):
+                        u = self._universes[k][j][i]
+                        assert isinstance(u, Universe)
+                        if u._id not in univs:
+                            univs[u._id] = u
 
-        for universe in unique_universes:
-            universes[universe._id] = universe
+        if self._outer._id not in univs:
+            univs[self._outer._id] = self._outer
 
-        return universes
+        return univs
 
     def get_all_nuclides(self):
         """Return all nuclides contained in the lattice
@@ -1091,14 +1106,14 @@ class HexLattice(Lattice):
 
         # Set the number of axial positions.
         if n_dims == 3:
-            self.num_axial = self._universes.shape[0]
+            self.num_axial = len(self._universes)
         else:
             self._num_axial = None
 
         # Set the number of rings and make sure this number is consistent for
         # all axial positions.
         if n_dims == 3:
-            self.num_rings = len(self._universes[0])
+            self.num_rings = len(self._universes)
             for rings in self._universes:
                 if len(rings) != self._num_rings:
                     msg = 'HexLattice ID={0:d} has an inconsistent number of ' \
@@ -1106,7 +1121,7 @@ class HexLattice(Lattice):
                     raise ValueError(msg)
 
         else:
-            self.num_rings = self._universes.shape[0]
+            self.num_rings = len(self._universes)
 
         # Make sure there are the correct number of elements in each ring.
         if n_dims == 3: