Add HexLattice.show_indices staticmethod and clarify universes description

openmc/lattice.py

@@ -30,8 +30,8 @@ class Lattice(object):
         Unique identifier for the lattice
     name : str
         Name of the lattice
-    pitch : float
-        Pitch of the lattice in cm
+    pitch : Iterable of float
+        Pitch of the lattice in each direction in cm
     outer : openmc.Universe
         A universe to fill all space outside the lattice
     universes : Iterable of Iterable of openmc.Universe
@@ -259,8 +259,9 @@ class RectLattice(Lattice):
     lower_left : Iterable of float
         The coordinates of the lower-left corner of the lattice. If the lattice
         is two-dimensional, only the x- and y-coordinates are specified.
-    pitch : float
-        Pitch of the lattice in cm
+    pitch : Iterable of float
+        Pitch of the lattice in the x, y, and (if applicable) z directions in
+        cm.
     outer : openmc.Universe
         A universe to fill all space outside the lattice
     universes : Iterable of Iterable of openmc.Universe
@@ -512,13 +513,19 @@ class HexLattice(Lattice):
     center : Iterable of float
         Coordinates of the center of the lattice. If the lattice does not have
         axial sections then only the x- and y-coordinates are specified
-    pitch : float
-        Pitch of the lattice in cm
+    pitch : Iterable of float
+        Pitch of the lattice in cm. The first item in the iterable specifies the
+        pitch in the radial direction and, if the lattice is 3D, the second item
+        in the iterable specifies the pitch in the axial direction.
     outer : openmc.Universe
         A universe to fill all space outside the lattice
     universes : Iterable of Iterable of openmc.Universe
         A two- or three-dimensional list/array of universes filling each element
-        of the lattice
+        of the lattice. Each sub-list corresponds to one ring of universes and
+        should be ordered from outermost ring to innermost ring. The universes
+        within each sub-list are ordered from the "top" and proceed in a
+        clockwise fashion. The :meth:`HexLattice.show_indices` method can be
+        used to help figure out indices for this property.
 
     """
 
@@ -882,3 +889,107 @@ class HexLattice(Lattice):
         # Join the rows together and return the string.
         universe_ids = '\n'.join(rows)
         return universe_ids
+
+    @staticmethod
+    def show_indices(num_rings):
+        """Return a diagram of the hexagonal lattice layout with indices.
+
+        This method can be used to show the proper indices to be used when
+        setting the :attr:`HexLattice.universes` property. For example, running
+        this method with num_rings=3 will return the following diagram::
+
+                      (0, 0)
+                (0,11)      (0, 1)
+          (0,10)      (1, 0)      (0, 2)
+                (1, 5)      (1, 1)
+          (0, 9)      (2, 0)      (0, 3)
+                (1, 4)      (1, 2)
+          (0, 8)      (1, 3)      (0, 4)
+                (0, 7)      (0, 5)
+                      (0, 6)
+
+        Parameters
+        ----------
+        num_rings : int
+            Number of rings in the hexagonal lattice
+
+        Returns
+        -------
+        str
+            Diagram of the hexagonal lattice showing indices
+
+        """
+
+        # Find the largest string and count the number of digits so we can
+        # properly pad the output string later
+        largest_index = 6*(num_rings - 1)
+        n_digits_index = len(str(largest_index))
+        n_digits_ring = len(str(num_rings - 1))
+        str_form = '({{:{}}},{{:{}}})'.format(n_digits_ring, n_digits_index)
+        pad = ' '*(n_digits_index + n_digits_ring + 3)
+
+        # Initialize the list for each row.
+        rows = [[] for i in range(1 + 4 * (num_rings-1))]
+        middle = 2 * (num_rings - 1)
+
+        # Start with the degenerate first ring.
+        rows[middle] = [str_form.format(num_rings - 1, 0)]
+
+        # Add universes one ring at a time.
+        for r in range(1, num_rings):
+            # r_prime increments down while r increments up.
+            r_prime = num_rings - 1 - r
+            theta = 0
+            y = middle + 2*r
+
+            for i in range(r):
+                # Climb down the top-right.
+                rows[y].append(str_form.format(r_prime, theta))
+                y -= 1
+                theta += 1
+
+            for i in range(r):
+                # Climb down the right.
+                rows[y].append(str_form.format(r_prime, theta))
+                y -= 2
+                theta += 1
+
+            for i in range(r):
+                # Climb down the bottom-right.
+                rows[y].append(str_form.format(r_prime, theta))
+                y -= 1
+                theta += 1
+
+            for i in range(r):
+                # Climb up the bottom-left.
+                rows[y].insert(0, str_form.format(r_prime, theta))
+                y += 1
+                theta += 1
+
+            for i in range(r):
+                # Climb up the left.
+                rows[y].insert(0, str_form.format(r_prime, theta))
+                y += 2
+                theta += 1
+
+            for i in range(r):
+                # Climb up the top-left.
+                rows[y].insert(0, str_form.format(r_prime, theta))
+                y += 1
+                theta += 1
+
+        # Flip the rows and join each row into a single string.
+        rows = [pad.join(x) for x in rows[::-1]]
+
+        # Pad the beginning of the rows so they line up properly.
+        for y in range(num_rings - 1):
+            rows[y] = (num_rings - 1 - y)*pad + rows[y]
+            rows[-1 - y] = (num_rings - 1 - y)*pad + rows[-1 - y]
+
+        for y in range(num_rings % 2, num_rings, 2):
+            rows[middle + y] = pad + rows[middle + y]
+            if y != 0:
+                rows[middle - y] = pad + rows[middle - y]
+
+        # Join the rows together and return the string.
+        return '\n'.join(rows)