From 3122d0798fe501e090276e5f3942930474a8733f Mon Sep 17 00:00:00 2001
From: Will Boyd <wbinventor@gmail.com>
Date: Wed, 1 Mar 2017 16:16:55 -0500
Subject: [PATCH] Began adding code to count material instances

---
 openmc/geometry.py | 74 +++++++++++++++++++++++++++++++++++++++++++++-
 openmc/material.py | 14 +++++++++
 2 files changed, 87 insertions(+), 1 deletion(-)

diff --git a/openmc/geometry.py b/openmc/geometry.py
index a02f76cc8a..1dc8f6a4b0 100644
--- a/openmc/geometry.py
+++ b/openmc/geometry.py
@@ -497,6 +497,7 @@ class Geometry(object):
                                 univ = latt.universes[k][m]
                                 self.count_cell_instances(univ)
 
+                # FIXME: Figure out how to iterate over lattices
                 elif isinstance(latt, openmc.HexLattice):
 
                     # 3D Lattices
@@ -522,4 +523,75 @@ class Geometry(object):
                                     continue
                                 else:
                                     univ = latt.universes[j][k]
-                                    self.count_cell_instances(univ)
\ No newline at end of file
+                                    self.count_cell_instances(univ)
+
+
+    # FIXME:
+    def count_material_instances(self, universe=None):
+        if universe is None:
+            universe = self.root_universe
+
+        if universe is None:
+            raise RuntimeError(
+                'Unable to count material instances without a root universe')
+
+        # (Re-)initialize all material instances to 0
+        for material in self.get_all_materials().values():
+            material.num_instances = 0
+
+        for cell in universe.cells.values():
+
+            # If material-filled, we are finished with all levels
+            if cell.fill_type == 'material':
+                cell.fill.num_instances += 1
+
+            # If universe-filled, recursively count materials in filling universe
+            elif cell.fill_type == 'universe':
+                self.count_material_instances(cell.fill)
+
+            # If lattice-filled, recursively call for all universes in lattice
+            elif cell.fill_type == 'lattice':
+                latt = cell.fill
+
+                if isinstance(latt, openmc.RectLattice):
+                    # 3D Lattices
+                    if latt.ndim == 3:
+                        for j in range(latt.shape[2]):
+                            for k in range(latt.shape[1]):
+                                for m in range(latt.shape[0]):
+                                    univ = latt.universes[j][k][m]
+                                    self.count_material_instances(univ)
+                    # 2D Lattices
+                    else:
+                        for k in range(latt.shape[1]):
+                            for m in range(latt.shape[0]):
+                                univ = latt.universes[k][m]
+                                self.count_material_instances(univ)
+
+                # FIXME: Figure out how to iterate over hexagonal lattices
+                elif isinstance(latt, openmc.HexLattice):
+
+                    # 3D Lattices
+                    if latt.ndim == 3:
+                        for m in range(latt.num_axial):
+                            for k in range(2 * latt.num_rings - 1):
+                                for j in range(2 * latt.num_rings - 1):
+                                    if j + k < latt.num_rings + 1:
+                                        continue
+                                    elif j + k > 3 * latt.num_rings - 1:
+                                        continue
+                                    else:
+                                        univ = latt.universes[j][k][m]
+                                        self.count_material_instances(univ)
+
+                    # 2D Lattices
+                    else:
+                        for k in range(2 * latt.num_rings - 1):
+                            for j in range(2 * latt.num_rings - 1):
+                                if j + k < latt.num_rings + 1:
+                                    continue
+                                elif j + k > 3 * latt.num_rings - 1:
+                                    continue
+                                else:
+                                    univ = latt.universes[j][k]
+                                    self.count_material_instances(univ)
\ No newline at end of file
diff --git a/openmc/material.py b/openmc/material.py
index e6f974b5d5..c0a1f59254 100644
--- a/openmc/material.py
+++ b/openmc/material.py
@@ -75,6 +75,10 @@ class Material(object):
         Volume of the material in cm^3. This can either be set manually or
         calculated in a stochastic volume calculation and added via the
         :meth:`Material.add_volume_information` method.
+    num_instances : int
+        The number of instances of this cell throughout the geometry. This
+        property is initialized by calling the
+        :meth:`Geometry.count_cell_instances()` method.
 
     """
 
@@ -88,6 +92,7 @@ class Material(object):
         self._depletable = False
         self._volume = None
         self._atoms = {}
+        self._num_instances = None
 
         # A list of tuples (nuclide, percent, percent type)
         self._nuclides = []
@@ -195,6 +200,10 @@ class Material(object):
     def depletable(self):
         return self._depletable
 
+    @property
+    def num_instances(self):
+        return self._num_instances
+
     @property
     def elements(self):
         return self._elements
@@ -275,6 +284,11 @@ class Material(object):
             cv.check_type('material volume', volume, Real)
         self._volume = volume
 
+    @num_instances.setter
+    def num_instances(self, num_instances):
+        cv.check_type('num_instances', num_instances, Integral)
+        self._num_instances = num_instances
+
     @classmethod
     def from_hdf5(cls, group):
         """Create material from HDF5 group