Implement spherical harmonics expansion filter

CMakeLists.txt

@@ -426,6 +426,7 @@ set(LIBOPENMC_FORTRAN_SRC
   src/tallies/tally_filter_meshsurface.F90
   src/tallies/tally_filter_mu.F90
   src/tallies/tally_filter_polar.F90
+  src/tallies/tally_filter_sph_harm.F90
   src/tallies/tally_filter_surface.F90
   src/tallies/tally_filter_universe.F90
   src/tallies/tally_header.F90

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_spherical_harmonics import *
 from openmc.trigger import *
 from openmc.tally_derivative import *
 from openmc.tallies import *

openmc/filter_spherical_harmonics.py

@@ -0,0 +1,150 @@
+from numbers import Integral
+from xml.etree import ElementTree as ET
+
+import numpy as np
+import pandas as pd
+
+import openmc.checkvalue as cv
+from . import Filter
+
+
+class SphericalHarmonicsFilter(Filter):
+    r"""Score spherical harmonic expansion moments up to specified order.
+
+    Parameters
+    ----------
+    order : int
+        Maximum spherical harmonics order
+    filter_id : int or None
+        Unique identifier for the filter
+
+    Attributes
+    ----------
+    order : int
+        Maximum spherical harmonics order
+    id : int
+        Unique identifier for the filter
+    cosine : {'scatter', 'particle'}
+        How to handle the cosine term.
+    num_bins : int
+        The number of filter bins
+
+    """
+
+    def __init__(self, order, filter_id=None):
+        self.order = order
+        self.id = filter_id
+        self._cosine = 'particle'
+
+    def __hash__(self):
+        string = type(self).__name__ + '\n'
+        string += '{: <16}=\t{}\n'.format('\tOrder', self.order)
+        string += '{: <16}=\t{}\n'.format('\tCosine', self.cosine)
+        return hash(string)
+
+    def __repr__(self):
+        string = type(self).__name__ + '\n'
+        string += '{: <16}=\t{}\n'.format('\tOrder', self.order)
+        string += '{: <16}=\t{}\n'.format('\tCosine', self.cosine)
+        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('spherical harmonics order', order, Integral)
+        cv.check_greater_than('spherical harmonics order', order, 0, equality=True)
+        self._order = order
+
+    @property
+    def cosine(self):
+        return self._cosine
+
+    @cosine.setter
+    def cosine(self, cosine):
+        cv.check_value('Spherical harmonics cosine treatment', cosine,
+                       ('scatter', 'particle'))
+        self._cosine = cosine
+
+    @property
+    def num_bins(self):
+        return (self._order + 1)**2
+
+    @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 '))
+
+        out = cls(group['order'].value, filter_id)
+        out.cosine = group['cosine'].value.decode()
+
+        return out
+
+    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
+        stride : int
+            Stride in memory for the filter
+
+        Returns
+        -------
+        pandas.DataFrame
+            A Pandas DataFrame with a column that is filled with strings
+            indicating spherical harmonics 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 = []
+        for n in range(self.order + 1):
+            bins.extend('Y{},{}'.format(n, m) for m in range(-n, n + 1))
+        bins = np.array(bins)
+
+        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 spherical harmonics filter data
+
+        """
+        element = ET.Element('filter')
+        element.set('id', str(self.id))
+        element.set('type', self.short_name.lower())
+        element.set('cosine', self.cosine)
+
+        subelement = ET.SubElement(element, 'order')
+        subelement.text = str(self.order)
+
+        return element

src/constants.F90

@@ -358,7 +358,7 @@ module constants
   integer, parameter :: NO_BIN_FOUND = -1
 
   ! Tally filter and map types
-  integer, parameter :: N_FILTER_TYPES = 17
+  integer, parameter :: N_FILTER_TYPES = 18
   integer, parameter :: &
        FILTER_UNIVERSE       = 1,  &
        FILTER_MATERIAL       = 2,  &
@@ -376,7 +376,8 @@ module constants
        FILTER_ENERGYFUNCTION = 14, &
        FILTER_CELLFROM       = 15, &
        FILTER_MESHSURFACE    = 16, &
-       FILTER_LEGENDRE       = 17
+       FILTER_LEGENDRE       = 17, &
+       FILTER_SPH_HARMONICS  = 18
 
   ! Mesh types
   integer, parameter :: &

src/tallies/tally_filter.F90

@@ -23,6 +23,7 @@ module tally_filter
   use tally_filter_meshsurface
   use tally_filter_mu
   use tally_filter_polar
+  use tally_filter_sph_harm
   use tally_filter_surface
   use tally_filter_universe
 
@@ -77,6 +78,8 @@ contains
           type_ = 'mu'
         type is (PolarFilter)
           type_ = 'polar'
+        type is (SphericalHarmonicsFilter)
+          type_ = 'sphericalharmonics'
         type is (SurfaceFilter)
           type_ = 'surface'
         type is (UniverseFilter)
@@ -150,6 +153,8 @@ contains
           allocate(MuFilter :: filters(index) % obj)
         case ('polar')
           allocate(PolarFilter :: filters(index) % obj)
+        case ('sphericalharmonics')
+          allocate(SphericalHarmonicsFilter :: filters(index) % obj)
         case ('surface')
           allocate(SurfaceFilter :: filters(index) % obj)
         case ('universe')

src/tallies/tally_filter_legendre.F90

@@ -75,7 +75,7 @@ contains
     integer,             intent(in) :: bin
     character(MAX_LINE_LEN)         :: label
 
-    label = "Legendre expansion order " // trim(to_str(bin - 1))
+    label = "Legendre expansion, P" // trim(to_str(bin - 1))
   end function text_label
 
 !===============================================================================

src/tallies/tally_filter_sph_harm.F90

@@ -0,0 +1,135 @@
+module tally_filter_sph_harm
+
+  use, intrinsic :: ISO_C_BINDING
+
+  use hdf5, only: HID_T
+
+  use constants
+  use error
+  use hdf5_interface
+  use math,                only: calc_pn, calc_rn
+  use particle_header,     only: Particle
+  use string,              only: to_str, to_lower
+  use tally_filter_header
+  use xml_interface
+
+  implicit none
+  private
+
+  integer, parameter :: COSINE_SCATTER = 1
+  integer, parameter :: COSINE_PARTICLE = 2
+
+!===============================================================================
+! LEGENDREFILTER gives Legendre moments of the change in scattering angle
+!===============================================================================
+
+  type, public, extends(TallyFilter) :: SphericalHarmonicsFilter
+    integer :: order
+    integer :: cosine
+  contains
+    procedure :: from_xml
+    procedure :: get_all_bins
+    procedure :: to_statepoint
+    procedure :: text_label
+  end type SphericalHarmonicsFilter
+
+contains
+
+!===============================================================================
+! SphericalHarmonicsFilter methods
+!===============================================================================
+
+  subroutine from_xml(this, node)
+    class(SphericalHarmonicsFilter), intent(inout) :: this
+    type(XMLNode), intent(in) :: node
+
+    character(MAX_WORD_LEN) :: temp_str
+
+    ! Get specified order
+    call get_node_value(node, "order", this % order)
+    this % n_bins = (this % order + 1)**2
+
+    ! Determine how cosine term is to be treated
+    if (check_for_node(node, "cosine")) then
+      call get_node_value(node, "cosine", temp_str)
+      select case (to_lower(temp_str))
+      case ('scatter')
+        this % cosine = COSINE_SCATTER
+      case ('particle')
+        this % cosine = COSINE_PARTICLE
+      end select
+    else
+      this % cosine = COSINE_PARTICLE
+    end if
+  end subroutine from_xml
+
+  subroutine get_all_bins(this, p, estimator, match)
+    class(SphericalHarmonicsFilter), intent(in)  :: this
+    type(Particle),      intent(in)  :: p
+    integer,             intent(in)  :: estimator
+    type(TallyFilterMatch),   intent(inout) :: match
+
+    integer :: i, j, n
+    integer :: num_nm
+    real(8) :: wgt
+    real(8) :: rn(2*this % order + 1)
+
+    ! TODO: Use recursive formula to calculate higher orders
+    j = 0
+    do n = 0, this % order
+      ! Determine cosine term for scatter expansion if necessary
+      if (this % cosine == COSINE_SCATTER) then
+        wgt = calc_pn(n, p % mu)
+      else
+        wgt = ONE
+      end if
+
+      ! Calculate n-th order spherical harmonics for (u,v,w)
+      num_nm = 2*n + 1
+      rn(1:num_nm) = calc_rn(n, p % last_uvw)
+
+      ! Append matching (bin,weight) for each moment
+      do i = 1, num_nm
+        j = j + 1
+        call match % bins % push_back(j)
+        call match % weights % push_back(wgt * rn(i))
+      end do
+    end do
+  end subroutine get_all_bins
+
+  subroutine to_statepoint(this, filter_group)
+    class(SphericalHarmonicsFilter), intent(in) :: this
+    integer(HID_T),      intent(in) :: filter_group
+
+    call write_dataset(filter_group, "type", "sphericalharmonics")
+    call write_dataset(filter_group, "n_bins", this % n_bins)
+    call write_dataset(filter_group, "order", this % order)
+    if (this % cosine == COSINE_SCATTER) then
+      call write_dataset(filter_group, "cosine", "scatter")
+    else
+      call write_dataset(filter_group, "cosine", "particle")
+    end if
+  end subroutine to_statepoint
+
+  function text_label(this, bin) result(label)
+    class(SphericalHarmonicsFilter), intent(in) :: this
+    integer,             intent(in) :: bin
+    character(MAX_LINE_LEN)         :: label
+
+    integer :: n, m
+
+    do n = 0, this % order
+      if (bin <= (n + 1)**2) then
+        m = (bin - n**2 - 1) - n
+        label = "Spherical harmonic expansion, Y" // trim(to_str(n)) // &
+             "," // trim(to_str(m))
+        exit
+      end if
+    end do
+  end function text_label
+
+!===============================================================================
+!                               C API FUNCTIONS
+!===============================================================================
+
+end module tally_filter_sph_harm

src/tallies/tally_header.F90

@@ -357,6 +357,8 @@ contains
       type is (LegendreFilter)
         j = FILTER_LEGENDRE
         this % estimator = ESTIMATOR_ANALOG
+      type is (SphericalHarmonicsFilter)
+        j = FILTER_SPH_HARMONICS
       end select
       this % find_filter(j) = i
     end do