From 55fb5de01b79b568fdb5f7407d10d0c998695135 Mon Sep 17 00:00:00 2001
From: Amelia Trainer <ameliajo@mit.edu>
Date: Fri, 5 Mar 2021 21:07:43 +0000
Subject: [PATCH] Got the python interface for the collision filter tentatively
 working

---
 include/openmc/tallies/filter_collision.h |   2 +-
 openmc/filter.py                          | 124 +++++++++++++++++++++-
 2 files changed, 124 insertions(+), 2 deletions(-)

diff --git a/include/openmc/tallies/filter_collision.h b/include/openmc/tallies/filter_collision.h
index 4773610416..b9b6d65e2a 100644
--- a/include/openmc/tallies/filter_collision.h
+++ b/include/openmc/tallies/filter_collision.h
@@ -23,7 +23,7 @@ public:
   //----------------------------------------------------------------------------
   // Methods
 
-  std::string type() const override { return "numCollisions"; }
+  std::string type() const override { return "collision"; }
 
   void from_xml(pugi::xml_node node) override;
 
diff --git a/openmc/filter.py b/openmc/filter.py
index 416698470b..5272c9494b 100644
--- a/openmc/filter.py
+++ b/openmc/filter.py
@@ -22,7 +22,8 @@ _FILTER_TYPES = (
     'universe', 'material', 'cell', 'cellborn', 'surface', 'mesh', 'energy',
     'energyout', 'mu', 'polar', 'azimuthal', 'distribcell', 'delayedgroup',
     'energyfunction', 'cellfrom', 'legendre', 'spatiallegendre',
-    'sphericalharmonics', 'zernike', 'zernikeradial', 'particle', 'cellinstance'
+    'sphericalharmonics', 'zernike', 'zernikeradial', 'particle', 'cellinstance',
+    'collision'
 )
 
 _CURRENT_NAMES = (
@@ -964,6 +965,127 @@ class MeshSurfaceFilter(MeshFilter):
         return pd.concat([df, pd.DataFrame(filter_dict)])
 
 
+class CollisionFilter(Filter):
+    """Bins tally events based on the number of collisions.
+
+    Parameters
+    ----------
+    values : Iterable of Int
+        A list or iterable of the number of collisions, as integer values. 
+        The events whose post-scattering collision number equals one of 
+        the provided values will be counted. 
+    filter_id : int
+        Unique identifier for the filter
+
+    Attributes
+    ----------
+    values : numpy.ndarray
+        An array of integer values representing the number of collisions events
+        by which to filter
+    id : int
+        Unique identifier for the filter
+    bins : numpy.ndarray
+        An array of integer values representing the number of collisions events
+        by which to filter
+    num_bins : Integral
+        The number of filter bins
+
+    """
+
+    def __init__(self, values, filter_id=None):
+        self.values = np.asarray(values)
+        self.bins = self.values
+        self.id = filter_id
+
+
+    def __repr__(self):
+        string = type(self).__name__ + '\n'
+        string += '{: <16}=\t{}\n'.format('\tValues', self.values)
+        string += '{: <16}=\t{}\n'.format('\tID', self.id)
+        return string
+
+    @Filter.bins.setter
+    def bins(self, bins):
+        Filter.bins.__set__(self, np.asarray(bins))
+
+    def check_bins(self, bins):
+        for x in bins:
+            # Values should be integers
+            cv.check_type('filter value', x, Integral)
+            cv.check_greater_than('filter value', x, 0, equality=True)
+
+    def get_bin_index(self, filter_bin):
+        i = np.where(self.bins[:, 1] == filter_bin[1])[0]
+        if len(i) == 0:
+            msg = 'Unable to get the bin index for Filter since "{0}" ' \
+                  'is not one of the bins'.format(filter_bin)
+            raise ValueError(msg)
+        else:
+            return i[0]
+
+    def get_pandas_dataframe(self, data_size, stride, **kwargs):
+        """Builds a Pandas DataFrame for the Filter's bins.
+
+        This method constructs a Pandas DataFrame object for the filter with
+        columns annotated by filter bin information. This is a helper method for
+        :meth:`Tally.get_pandas_dataframe`.
+
+        Parameters
+        ----------
+        data_size : int
+            The total number of bins in the tally corresponding to this filter
+        stride : int
+            Stride in memory for the filter
+
+        Returns
+        -------
+        pandas.DataFrame
+            A Pandas DataFrame with one column of the lower energy bound and one
+            column of upper energy bound for each filter bin.  The number of
+            rows in the DataFrame is the same as the total number of bins in the
+            corresponding tally, with the filter bin appropriately tiled to map
+            to the corresponding tally bins.
+
+        See also
+        --------
+        Tally.get_pandas_dataframe(), CrossFilter.get_pandas_dataframe()
+
+        """
+        # Initialize Pandas DataFrame
+        df = pd.DataFrame()
+
+        # Extract the lower and upper energy bounds, then repeat and tile
+        # them as necessary to account for other filters.
+        bins = np.repeat(self.bins[:], stride)
+        tile_factor = data_size // len(bins)
+        bins = np.tile(bins, tile_factor)
+
+        # Add the new energy columns to the DataFrame.
+        if hasattr(self, 'units'):
+            units = ' [{}]'.format(self.units)
+        else:
+            units = ''
+
+        df.loc[:, self.short_name.lower() + ' collisions' ] = bins
+
+        return df
+
+    def to_xml_element(self):
+        """Return XML Element representing the Filter.
+
+        Returns
+        -------
+        element : xml.etree.ElementTree.Element
+            XML element containing filter data
+
+        """
+        element = super().to_xml_element()
+        element[0].text = ' '.join(str(x) for x in self.values)
+        return element
+
+
+
+
 class RealFilter(Filter):
     """Tally modifier that describes phase-space and other characteristics