From 0e8eaecf2faa223aa32d1b768bd9d9c025127632 Mon Sep 17 00:00:00 2001
From: Will Boyd <wbinventor@gmail.com>
Date: Fri, 11 Sep 2015 16:41:11 -0400
Subject: [PATCH] Multigroup Chi class is now working and verfied for 2-groups

---
 openmc/filter.py    |  5 +++--
 openmc/mgxs/mgxs.py | 11 ++++++++---
 openmc/tallies.py   | 22 ++++++++++++++++------
 3 files changed, 27 insertions(+), 11 deletions(-)

diff --git a/openmc/filter.py b/openmc/filter.py
index 52c927bed5..435d08c448 100644
--- a/openmc/filter.py
+++ b/openmc/filter.py
@@ -321,8 +321,7 @@ class Filter(object):
             # Filter bins for distribcell are the "IDs" of each unique placement
             # of the Cell in the Geometry (integers starting at 0)
             elif self.type == 'distribcell':
-                val = np.where(self.bins == filter_bin)[0][0]
-                filter_index = val
+                filter_index = filter_bin
 
             # Use ID for all other Filters (e.g., material, cell, etc.)
             else:
@@ -363,6 +362,8 @@ class Filter(object):
 
         elif self.type in ['energy', 'energyout']:
             bin = (self.bins[bin_index], self.bins[bin_index+1])
+        elif self.type == 'distribcell':
+            bin = (self.bins[0],)
         else:
             bin = (self.bins[bin_index],)
 
diff --git a/openmc/mgxs/mgxs.py b/openmc/mgxs/mgxs.py
index 39eb79160e..d92daa5d6b 100644
--- a/openmc/mgxs/mgxs.py
+++ b/openmc/mgxs/mgxs.py
@@ -1280,14 +1280,19 @@ class Chi(MultiGroupXS):
 
         sum_nu_fission_in = nu_fission_in.summation(filters=['energy'],
                                                     filter_bins=energy_bins)
+
+        # FIXME: CrossFilter for energy + energy messes up tally arithmetic
+        sum_nu_fission_in.remove_filter(sum_nu_fission_in.filters[-1])
+
         self._xs_tally = nu_fission_out / sum_nu_fission_in
 
         # Compute the total across all groups per subdomain
         if self.domain_type == 'distribcell':
-            subdomain_indices = self.get_subdomain_indices()
-            filter_bins = [(i,) for i in subdomain_indices]
+            domain_filter = self.tallies['nu-fission-in'].filters[0]
+            num_subdomains = domain_filter.num_bins
+            filter_bins = [((i,),) for i in range(num_subdomains)]
         else:
-            filter_bins = [(self.domain,)]
+            filter_bins = [((self.domain,),)]
 
         # Normalize chi to 1.0
         norm = self._xs_tally.summation(filters=[self.domain_type],
diff --git a/openmc/tallies.py b/openmc/tallies.py
index 827d98b613..c0fb1615f5 100644
--- a/openmc/tallies.py
+++ b/openmc/tallies.py
@@ -852,6 +852,9 @@ class Tally(object):
                         for k in range(filter.num_bins):
                             bins.append((filter.bins[k], filter.bins[k+1]))
 
+                    elif filter.type == 'distribcell':
+                        bins = np.arange(filter.num_bins)
+
                     # Create list of IDs for bins for all other Filter types
                     else:
                         bins = filter.bins
@@ -1529,12 +1532,19 @@ class Tally(object):
             stride *= filter.num_bins
 
         filters = [filter1.type, filter2.type]
-        filter1_bins = np.arange(filter.num_bins)
-        filter2_bins = np.arange(filter2.num_bins)
+        if filter1.type == 'distribcell':
+            filter1_bins = np.arange(filter.num_bins)
+        else:
+            filter1_bins = [(filter1.get_bin(i)) for i in range(filter1.num_bins)]
+
+        if filter1.type == 'distribcell':
+            filter2_bins = np.arange(filter2.num_bins)
+        else:
+            filter2_bins = [filter2.get_bin(i) for i in range(filter2.num_bins)]
 
         if self.sum is not None:
             for bin1, bin2 in itertools.product(filter1_bins, filter2_bins):
-                filter_bins = [(filter1.get_bin(bin1),), (filter2.get_bin(bin2),)]
+                filter_bins = [(bin1,), (bin2,)]
                 data = self.get_values(filters=filters,
                                        filter_bins=filter_bins, value='sum')
                 indices = swap_tally.get_filter_indices(filters, filter_bins)
@@ -1542,7 +1552,7 @@ class Tally(object):
 
         if self.sum_sq is not None:
             for bin1, bin2 in itertools.product(filter1_bins, filter2_bins):
-                filter_bins = [(filter1.get_bin(bin1),), (filter2.get_bin(bin2),)]
+                filter_bins = [(bin1,), (bin2,)]
                 data = self.get_values(filters=filters,
                                        filter_bins=filter_bins, value='sum_sq')
                 indices = swap_tally.get_filter_indices(filters, filter_bins)
@@ -1550,7 +1560,7 @@ class Tally(object):
 
         if self.sum is not None:
             for bin1, bin2 in itertools.product(filter1_bins, filter2_bins):
-                filter_bins = [(filter1.get_bin(bin1),), (filter2.get_bin(bin2),)]
+                filter_bins = [(bin1,), (bin2,)]
                 data = self.get_values(filters=filters,
                                        filter_bins=filter_bins, value='mean')
                 indices = swap_tally.get_filter_indices(filters, filter_bins)
@@ -1558,7 +1568,7 @@ class Tally(object):
 
         if self.sum is not None:
             for bin1, bin2 in itertools.product(filter1_bins, filter2_bins):
-                filter_bins = [(filter1.get_bin(bin1),), (filter2.get_bin(bin2),)]
+                filter_bins = [(bin1,), (bin2,)]
                 data = self.get_values(filters=filters,
                                        filter_bins=filter_bins, value='std_dev')
                 indices = swap_tally.get_filter_indices(filters, filter_bins)