Reworking main for loop of get_filter_indices based on comments from @paulromano

This commit is contained in:
Patrick Shriwise 2022-02-16 08:39:18 -06:00
parent ccec9ee7bf
commit e2c432484e

View file

@ -1092,29 +1092,16 @@ class Tally(IDManagerMixin):
for i, self_filter in enumerate(self.filters):
# If a user-requested Filter, get the user-requested bins
for j, test_filter in enumerate(filters):
if type(self_filter) is openmc.EnergyFunctionFilter:
indices = [self_filter.get_bin_index(None)]
break
if type(self_filter) is test_filter:
bins = filter_bins[j]
indices = np.array([self_filter.get_bin_index(b) for b in bins])
break
else:
# If not a user-requested Filter, get all bins
if isinstance(self_filter, openmc.DistribcellFilter):
# Create list of cell instance IDs for distribcell Filters
bins = list(range(self_filter.num_bins))
indices = np.arange(self_filter.num_bins)
elif isinstance(self_filter, openmc.EnergyFunctionFilter):
# EnergyFunctionFilters don't have bins so just add a None
bins = [None]
else:
# Create list of IDs for bins for all other filter types
bins = self_filter.bins
# Add indices for each bin in this Filter to the list
if type(self_filter) in filters:
indices = np.array([self_filter.get_bin_index(b) for b in bins])
else:
indices = np.arange(len(bins))
filter_indices.append(indices)
# Account for stride in each of the previous filters