Add several filter-related functions to C API and bindings

openmc/capi/filter.py

@@ -1,12 +1,27 @@
-from ctypes import c_int, c_int32, c_double, c_char_p, POINTER
+from collections import Mapping
+from ctypes import c_int, c_int32, c_double, c_char_p, POINTER, \
+    create_string_buffer
+from weakref import WeakValueDictionary
+
+import numpy as np
+from numpy.ctypeslib import as_array
 
 from . import _dll
 from .error import _error_handler
 
 
-__all__ = []
+__all__ = ['FilterView', 'AzimuthalFilterView', 'CellFilterView',
+           'CellbornFilterView', 'CellfromFilterView', 'DistribcellFilterView',
+           'DelayedGroupFilterView', 'EnergyFilterView', 'EnergyoutFilterView',
+           'EnergyFunctionFilterView', 'MaterialFilterView', 'MeshFilterView',
+           'MuFilterView', 'PolarFilterView', 'SurfaceFilterView',
+           'UniverseFilterView', 'filters']
 
 # Tally functions
+_dll.openmc_energy_filter_get_bins.argtypes = [
+    c_int32, POINTER(POINTER(c_double)), POINTER(c_int32)]
+_dll.openmc_energy_filter_get_bins.restype = c_int
+_dll.openmc_energy_filter_get_bins.errcheck = _error_handler
 _dll.openmc_energy_filter_set_bins.argtypes = [c_int32, c_int32, POINTER(c_double)]
 _dll.openmc_energy_filter_set_bins.restype = c_int
 _dll.openmc_energy_filter_set_bins.errcheck = _error_handler
@@ -16,12 +31,157 @@ _dll.openmc_extend_filters.errcheck = _error_handler
 _dll.openmc_filter_get_id.argtypes = [c_int32, POINTER(c_int32)]
 _dll.openmc_filter_get_id.restype = c_int
 _dll.openmc_filter_get_id.errcheck = _error_handler
+_dll.openmc_filter_get_type.argtypes = [c_int32, c_char_p]
+_dll.openmc_filter_get_type.restype = c_int
+_dll.openmc_filter_get_type.errcheck = _error_handler
 _dll.openmc_filter_set_id.argtypes = [c_int32, c_int32]
 _dll.openmc_filter_set_id.restype = c_int
 _dll.openmc_filter_set_id.errcheck = _error_handler
 _dll.openmc_filter_set_type.argtypes = [c_int32, c_char_p]
 _dll.openmc_filter_set_type.restype = c_int
 _dll.openmc_filter_set_type.errcheck = _error_handler
+_dll.openmc_get_filter.argtypes = [c_int32, POINTER(c_int32)]
+_dll.openmc_get_filter.restype = c_int
+_dll.openmc_get_filter.errcheck = _error_handler
 _dll.openmc_mesh_filter_set_mesh.argtypes = [c_int32, c_int32]
 _dll.openmc_mesh_filter_set_mesh.restype = c_int
 _dll.openmc_mesh_filter_set_mesh.errcheck = _error_handler
+
+
+class FilterView(object):
+    __instances = WeakValueDictionary()
+
+    def __new__(cls, *args):
+        if args not in cls.__instances:
+            instance = super().__new__(cls)
+            cls.__instances[args] = instance
+        return cls.__instances[args]
+
+    def __init__(self, index):
+        self._index = index
+
+    @property
+    def id(self):
+        filter_id = c_int32()
+        _dll.openmc_filter_get_id(self._index, filter_id)
+        return filter_id.value
+
+    @id.setter
+    def id(self, filter_id):
+        _dll.openmc_filter_set_id(self._index, filter_id)
+
+
+class EnergyFilterView(FilterView):
+    @property
+    def bins(self):
+        energies = POINTER(c_double)()
+        n = c_int32()
+        _dll.openmc_energy_filter_get_bins(self._index, energies, n)
+        return as_array(energies, (n.value,))
+
+    @bins.setter
+    def bins(self, bins):
+        # Get numpy array as a double*
+        energies = np.asarray(bins)
+        energies_p = energies.ctypes.data_as(POINTER(c_double))
+
+        _dll.openmc_energy_filter_set_bins(
+            self._index, len(energies), energies_p)
+
+
+class EnergyoutFilterView(FilterView):
+    pass
+
+
+class AzimuthalFilterView(FilterView):
+    pass
+
+
+class CellFilterView(FilterView):
+    pass
+
+
+class CellbornFilterView(FilterView):
+    pass
+
+
+class CellfromFilterView(FilterView):
+    pass
+
+
+class DelayedGroupFilterView(FilterView):
+    pass
+
+
+class DistribcellFilterView(FilterView):
+    pass
+
+
+class EnergyFunctionFilterView(FilterView):
+    pass
+
+
+class MaterialFilterView(FilterView):
+    pass
+
+
+class MeshFilterView(FilterView):
+    pass
+
+
+class MuFilterView(FilterView):
+    pass
+
+
+class PolarFilterView(FilterView):
+    pass
+
+
+class SurfaceFilterView(FilterView):
+    pass
+
+
+class UniverseFilterView(FilterView):
+    pass
+
+
+_filter_type_map = {
+    'azimuthal': AzimuthalFilterView,
+    'cell': CellFilterView,
+    'cellborn': CellbornFilterView,
+    'cellfrom': CellfromFilterView,
+    'delayedgroup': DelayedGroupFilterView,
+    'distribcell': DistribcellFilterView,
+    'energy': EnergyFilterView,
+    'energyout': EnergyoutFilterView,
+    'energyfunction': EnergyFunctionFilterView,
+    'material': MaterialFilterView,
+    'mesh': MeshFilterView,
+    'mu': MuFilterView,
+    'polar': PolarFilterView,
+    'surface': SurfaceFilterView,
+    'universe': UniverseFilterView,
+}
+
+
+def _get_filter(index):
+    filter_type = create_string_buffer(20)
+    _dll.openmc_filter_get_type(index, filter_type)
+    filter_type = filter_type.value.decode()
+    return _filter_type_map[filter_type](index)
+
+
+class _FilterMapping(Mapping):
+    def __getitem__(self, key):
+        index = c_int32()
+        _dll.openmc_get_filter(key, index)
+        return _get_filter(index.value)
+
+    def __iter__(self):
+        for i in range(len(self)):
+            yield TallyView(i + 1).id
+
+    def __len__(self):
+        return c_int32.in_dll(_dll, 'n_filters').value
+
+filters = _FilterMapping()

openmc/capi/tally.py

@@ -6,6 +6,7 @@ from numpy.ctypeslib import as_array
 
 from . import _dll, NuclideView
 from .error import _error_handler
+from .filter import _get_filter
 
 
 __all__ = ['TallyView', 'tallies']
@@ -20,7 +21,8 @@ _dll.openmc_extend_tallies.errcheck = _error_handler
 _dll.openmc_tally_get_id.argtypes = [c_int32, POINTER(c_int32)]
 _dll.openmc_tally_get_id.restype = c_int
 _dll.openmc_tally_get_id.errcheck = _error_handler
-_dll.openmc_tally_get_filters.argtypes = [c_int32, POINTER(c_int32), POINTER(c_int)]
+_dll.openmc_tally_get_filters.argtypes = [
+    c_int32, POINTER(POINTER(c_int32)), POINTER(c_int)]
 _dll.openmc_tally_get_filters.restype = c_int
 _dll.openmc_tally_get_filters.errcheck = _error_handler
 _dll.openmc_tally_get_nuclides.argtypes = [
@@ -52,12 +54,14 @@ class TallyView(object):
     Parameters
     ----------
     index : int
-         Index in the `tallys` array.
+         Index in the `tallies` array.
 
     Attributes
     ----------
     id : int
         ID of the tally
+    filters : list
+        List of views to tally filters
     nuclides : list of str
         List of nuclides to score results for
     results : numpy.ndarray
@@ -81,6 +85,13 @@ class TallyView(object):
         _dll.openmc_tally_get_id(self._index, tally_id)
         return tally_id.value
 
+    @property
+    def filters(self):
+        filt_idx = POINTER(c_int32)()
+        n = c_int()
+        _dll.openmc_tally_get_filters(self._index, filt_idx, n)
+        return [_get_filter(filt_idx[i]) for i in range(n.value)]
+
     @property
     def nuclides(self):
         nucs = POINTER(c_int)()

src/api.F90

@@ -30,10 +30,12 @@ module openmc_api
   public :: openmc_calculate_volumes
   public :: openmc_cell_get_id
   public :: openmc_cell_set_temperature
+  public :: openmc_energy_filter_get_bins
   public :: openmc_energy_filter_set_bins
   public :: openmc_extend_filters
   public :: openmc_extend_tallies
   public :: openmc_filter_get_id
+  public :: openmc_filter_get_type
   public :: openmc_filter_set_id
   public :: openmc_filter_set_type
   public :: openmc_finalize

src/error.F90

@@ -31,6 +31,7 @@ module error
   integer(C_INT), public, bind(C) :: E_ARGUMENT_INVALID = -16
   integer(C_INT), public, bind(C) :: E_WRONG_TYPE = -17
   integer(C_INT), public, bind(C) :: E_FILTER_NOT_ALLOCATED = -18
+  integer(C_INT), public, bind(C) :: E_FILTER_INVALID_ID = -19
 
   ! Warning codes
   integer(C_INT), public, bind(C) :: W_BELOW_MIN_BOUND = 1

src/tallies/tally_filter.F90

@@ -1149,6 +1149,68 @@ contains
 !                               C API FUNCTIONS
 !===============================================================================
 
+  function openmc_filter_get_type(index, type) result(err) bind(C)
+    ! Get the type of a filter
+    integer(C_INT32_T), value, intent(in) :: index
+    character(kind=C_CHAR), intent(out) :: type(*)
+    integer(C_INT) :: err
+
+    integer :: i
+    character(20) :: type_
+
+    if (index >= 1 .and. index <= n_filters) then
+      if (allocated(filters(index) % obj)) then
+        ! Get type as a Fortran string
+        select type (f => filters(index) % obj)
+        type is (AzimuthalFilter)
+          type_ = 'azimuthal'
+        type is (CellFilter)
+          type_ = 'cell'
+        type is (CellbornFilter)
+          type_ = 'cellborn'
+        type is (CellfromFilter)
+          type_ = 'cellfrom'
+        type is (DelayedGroupFilter)
+          type_ = 'delayedgroup'
+        type is (DistribcellFilter)
+          type_ = 'distribcell'
+        type is (EnergyFilter)
+          type_ = 'energy'
+        type is (EnergyoutFilter)
+          type_ = 'energyout'
+        type is (EnergyFunctionFilter)
+          type_ = 'energyfunction'
+        type is (MaterialFilter)
+          type_ = 'material'
+        type is (MeshFilter)
+          type_ = 'mesh'
+        type is (MuFilter)
+          type_ = 'mu'
+        type is (PolarFilter)
+          type_ = 'polar'
+        type is (SurfaceFilter)
+          type_ = 'surface'
+        type is (UniverseFilter)
+          type_ = 'universe'
+        end select
+
+        ! Convert Fortran string to null-terminated C string. We assume the
+        ! caller has allocated a char array buffer
+        do i = 1, len_trim(type_)
+          type(i) = type_(i:i)
+        end do
+        type(len_trim(type_) + 1) = C_NULL_CHAR
+
+        err = 0
+      else
+        err = E_FILTER_NOT_ALLOCATED
+      end if
+    else
+      err = E_OUT_OF_BOUNDS
+    end if
+  end function openmc_filter_get_type
+
+
   function openmc_filter_set_type(index, type) result(err) bind(C)
     ! Set the type of a filter
     integer(C_INT32_T), value, intent(in) :: index

src/tallies/tally_filter_energy.F90

@@ -15,6 +15,7 @@ module tally_filter_energy
 
   implicit none
   private
+  public :: openmc_energy_filter_get_bins
   public :: openmc_energy_filter_set_bins
 
 !===============================================================================
@@ -165,6 +166,32 @@ contains
 !                               C API FUNCTIONS
 !===============================================================================
 
+  function openmc_energy_filter_get_bins(index, energies, n) result(err) bind(C)
+    ! Return the bounding energies for an energy filter
+    integer(C_INT32_T), value :: index
+    type(C_PTR), intent(out) :: energies
+    integer(C_INT32_T), intent(out) :: n
+    integer(C_INT) :: err
+
+    if (index >= 1 .and. index <= n_filters) then
+      if (allocated(filters(index) % obj)) then
+        select type (f => filters(index) % obj)
+        type is (EnergyFilter)
+          energies = C_LOC(f % bins)
+          n = size(f % bins)
+          err = 0
+        class default
+          err = E_WRONG_TYPE
+        end select
+      else
+        err = E_FILTER_NOT_ALLOCATED
+      end if
+    else
+      err = E_OUT_OF_BOUNDS
+    end if
+  end function openmc_energy_filter_get_bins
+
+
   function openmc_energy_filter_set_bins(index, n, energies) result(err) bind(C)
     ! Set the bounding energies for an energy filter
     integer(C_INT32_T), value, intent(in) :: index

src/tallies/tally_filter_header.F90

@@ -191,4 +191,23 @@ contains
     end if
   end function openmc_filter_set_id
 
+
+  function openmc_get_filter(id, index) result(err) bind(C)
+    ! Returns the index in the filters array of a filter with a given ID
+    integer(C_INT32_T), value :: id
+    integer(C_INT32_T), intent(out) :: index
+    integer(C_INT) :: err
+
+    if (allocated(filters)) then
+      if (filter_dict % has_key(id)) then
+        index = filter_dict % get_key(id)
+        err = 0
+      else
+        err = E_FILTER_INVALID_ID
+      end if
+    else
+      err = E_FILTER_NOT_ALLOCATED
+    end if
+  end function openmc_get_filter
+
 end module tally_filter_header