Fix spatial Legendre filter and add associated Python class

openmc/__init__.py

@@ -16,6 +16,7 @@ from openmc.universe import *
 from openmc.lattice import *
 from openmc.filter import *
 from openmc.filter_legendre import *
+from openmc.filter_spatial_legendre import *
 from openmc.filter_spherical_harmonics import *
 from openmc.trigger import *
 from openmc.tally_derivative import *

openmc/filter_spatial_legendre.py

@@ -0,0 +1,173 @@
+from numbers import Integral, Real
+from xml.etree import ElementTree as ET
+
+import numpy as np
+import pandas as pd
+
+import openmc.checkvalue as cv
+from . import Filter
+
+
+class SpatialLegendreFilter(Filter):
+    r"""Score Legendre expansion moments in space up to specified order.
+
+    This filter allows scores to be multiplied by Legendre polynomials of the
+    the particle's position along a particular axis, normalized to a given
+    range, up to a user-specified order.
+
+    Parameters
+    ----------
+    order : int
+        Maximum Legendre polynomial order
+    axis : {'x', 'y', 'z'}
+        Axis along which to take the expansion
+    minimum : float
+        Minimum value along selected axis
+    maximum : float
+        Maximum value along selected axis
+    filter_id : int or None
+        Unique identifier for the filter
+
+    Attributes
+    ----------
+    order : int
+        Maximum Legendre polynomial order
+    id : int
+        Unique identifier for the filter
+    num_bins : int
+        The number of filter bins
+
+    """
+
+    def __init__(self, order, axis, minimum, maximum, filter_id=None):
+        self.order = order
+        self.axis = axis
+        self.minimum = minimum
+        self.maximum = maximum
+        self.id = filter_id
+
+    def __hash__(self):
+        string = type(self).__name__ + '\n'
+        string += '{: <16}=\t{}\n'.format('\tOrder', self.order)
+        return hash(string)
+
+    def __repr__(self):
+        string = type(self).__name__ + '\n'
+        string += '{: <16}=\t{}\n'.format('\tOrder', self.order)
+        string += '{: <16}=\t{}\n'.format('\tID', self.id)
+        return string
+
+    @property
+    def order(self):
+        return self._order
+
+    @order.setter
+    def order(self, order):
+        cv.check_type('Legendre order', order, Integral)
+        cv.check_greater_than('Legendre order', order, 0, equality=True)
+        self._order = order
+
+    @property
+    def axis(self):
+        return self._axis
+
+    @axis.setter
+    def axis(self, axis):
+        cv.check_value('axis', axis, ('x', 'y', 'z'))
+        self._axis = axis
+
+    @property
+    def minimum(self):
+        return self._minimum
+
+    @minimum.setter
+    def minimum(self, minimum):
+        cv.check_type('minimum', minimum, Real)
+        self._minimum = minimum
+
+    @property
+    def maximum(self):
+        return self._maximum
+
+    @maximum.setter
+    def maximum(self, maximum):
+        cv.check_type('maximum', maximum, Real)
+        self._maximum = maximum
+
+    @property
+    def num_bins(self):
+        return self._order + 1
+
+    @classmethod
+    def from_hdf5(cls, group, **kwargs):
+        if group['type'].value.decode() != cls.short_name.lower():
+            raise ValueError("Expected HDF5 data for filter type '"
+                             + cls.short_name.lower() + "' but got '"
+                             + group['type'].value.decode() + " instead")
+
+        filter_id = int(group.name.split('/')[-1].lstrip('filter '))
+        order = group['order'].value
+        axis = group['axis'].value.decode()
+        min_, max_ = group['min'].value, group['max'].value
+
+        return cls(order, axis, min_, max_, filter_id)
+
+    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 : Integral
+            The total number of bins in the tally corresponding to this filter
+
+        Returns
+        -------
+        pandas.DataFrame
+            A Pandas DataFrame with a column that is filled with strings
+            indicating Legendre orders. The number of rows in the DataFrame is
+            the same as the total number of bins in the corresponding tally.
+
+        See also
+        --------
+        Tally.get_pandas_dataframe(), CrossFilter.get_pandas_dataframe()
+
+        """
+        # Initialize Pandas DataFrame
+        df = pd.DataFrame()
+
+        bins = np.array(['P{}'.format(i) for i in range(self.order + 1)])
+        filter_bins = np.repeat(bins, stride)
+        tile_factor = data_size // len(filter_bins)
+        filter_bins = np.tile(filter_bins, tile_factor)
+        df = pd.concat([df, pd.DataFrame(
+            {self.short_name.lower(): filter_bins})])
+
+        return df
+
+    def to_xml_element(self):
+        """Return XML Element representing the filter.
+
+        Returns
+        -------
+        element : xml.etree.ElementTree.Element
+            XML element containing Legendre filter data
+
+        """
+        element = ET.Element('filter')
+        element.set('id', str(self.id))
+        element.set('type', self.short_name.lower())
+
+        subelement = ET.SubElement(element, 'order')
+        subelement.text = str(self.order)
+        subelement = ET.SubElement(element, 'axis')
+        subelement.text = self.axis
+        subelement = ET.SubElement(element, 'min')
+        subelement.text = str(self.minimum)
+        subelement = ET.SubElement(element, 'max')
+        subelement.text = str(self.maximum)
+
+        return element

src/tallies/tally_filter_sptl_legendre.F90

@@ -9,19 +9,26 @@ module tally_filter_sptl_legendre
   use hdf5_interface
   use math,                only: calc_pn
   use particle_header,     only: Particle
-  use string,              only: to_str
+  use string,              only: to_str, to_lower
   use tally_filter_header
   use xml_interface
 
   implicit none
   private
 
+  integer, parameter :: AXIS_X = 1
+  integer, parameter :: AXIS_Y = 2
+  integer, parameter :: AXIS_Z = 3
+
 !===============================================================================
 ! SpatialLEGENDREFILTER gives Legendre moments
 !===============================================================================
 
   type, public, extends(TallyFilter) :: SpatialLegendreFilter
     integer :: order
+    integer :: axis
+    real(8) :: min
+    real(8) :: max
   contains
     procedure :: from_xml
     procedure :: get_all_bins
@@ -39,8 +46,22 @@ contains
     class(SpatialLegendreFilter), intent(inout) :: this
     type(XMLNode), intent(in) :: node
 
-    ! Get specified order
+    character(MAX_WORD_LEN) :: axis
+
+    ! Get attributes from XML
     call get_node_value(node, "order", this % order)
+    call get_node_value(node, "axis", axis)
+    select case (to_lower(axis))
+    case ('x')
+      this % axis = AXIS_X
+    case ('y')
+      this % axis = AXIS_Y
+    case ('z')
+      this % axis = AXIS_Z
+    end select
+    call get_node_value(node, "min", this % min)
+    call get_node_value(node, "max", this % max)
+
     this % n_bins = this % order + 1
   end subroutine from_xml
 
@@ -52,27 +73,40 @@ contains
 
     integer :: i
     real(8) :: wgt
-    real(8) :: leg_x
-    real(8) :: xmin
-    real(8) :: xmax
+    real(8) :: x       ! Position on specified axis
+    real(8) :: x_norm  ! Normalized position
 
-    ! TODO: Use recursive formula to calculate higher orders
-    do i = 0, this % order
-    !i is the order, second var is the value normalized between 0 and 1 given by
-    ! norm_x = 2*(leg_x-xmin)/(xmax-xmin)-1
-      wgt = 2.0*(leg_x - xmin)/(xmax - xmin) - 1
-      call match % bins % push_back(i + 1)
-      call match % weights % push_back(wgt)
-    end do
+    x = p % coord(1) % xyz(this % axis)
+    if (this % min <= x .and. x <= this % max) then
+      ! Calculate normalized position between min and max
+      x_norm = TWO*(x - this % min)/(this % max - this % min) - ONE
+
+      ! TODO: Use recursive formula to calculate higher orders
+      do i = 0, this % order
+        wgt = calc_pn(i, x_norm)
+        call match % bins % push_back(i + 1)
+        call match % weights % push_back(wgt)
+      end do
+    end if
   end subroutine get_all_bins
 
   subroutine to_statepoint(this, filter_group)
     class(SpatialLegendreFilter), intent(in) :: this
     integer(HID_T),      intent(in) :: filter_group
 
-    call write_dataset(filter_group, "type", "legendre")
+    call write_dataset(filter_group, "type", "spatiallegendre")
     call write_dataset(filter_group, "n_bins", this % n_bins)
     call write_dataset(filter_group, "order", this % order)
+    select case (this % axis)
+    case (AXIS_X)
+      call write_dataset(filter_group, 'axis', 'x')
+    case (AXIS_Y)
+      call write_dataset(filter_group, 'axis', 'y')
+    case (AXIS_Z)
+      call write_dataset(filter_group, 'axis', 'z')
+    end select
+    call write_dataset(filter_group, 'min', this % min)
+    call write_dataset(filter_group, 'max', this % max)
   end subroutine to_statepoint
 
   function text_label(this, bin) result(label)