Source code for openmc.lib.filter

from collections.abc import Mapping
from ctypes import c_int, c_int32, c_double, c_char_p, POINTER, \
    create_string_buffer, c_size_t
from weakref import WeakValueDictionary

import numpy as np
from numpy.ctypeslib import as_array

from openmc.exceptions import AllocationError, InvalidIDError
from openmc.data.function import INTERPOLATION_SCHEME
from openmc import ParticleType
from . import _dll
from .core import _FortranObjectWithID
from .error import _error_handler
from .material import Material
from .mesh import _get_mesh


__all__ = [
    'Filter', 'AzimuthalFilter', 'CellFilter', 'CellbornFilter', 'CellfromFilter',
    'CellInstanceFilter', 'CollisionFilter', 'DistribcellFilter', 'DelayedGroupFilter',
    'EnergyFilter', 'EnergyoutFilter', 'EnergyFunctionFilter', 'LegendreFilter',
    'MaterialFilter', 'MeshFilter', 'MeshSurfaceFilter', 'MuFilter', 'ParticleFilter',
    'PolarFilter', 'SphericalHarmonicsFilter', 'SpatialLegendreFilter', 'SurfaceFilter',
    'UniverseFilter', 'ZernikeFilter', 'ZernikeRadialFilter', 'filters'
]

# Tally functions
_dll.openmc_cell_filter_get_bins.argtypes = [
    c_int32, POINTER(POINTER(c_int32)), POINTER(c_int32)]
_dll.openmc_cell_filter_get_bins.restype = c_int
_dll.openmc_cell_filter_get_bins.errcheck = _error_handler
_dll.openmc_energy_filter_get_bins.argtypes = [
    c_int32, POINTER(POINTER(c_double)), POINTER(c_size_t)]
_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_size_t, POINTER(c_double)]
_dll.openmc_energy_filter_set_bins.restype = c_int
_dll.openmc_energy_filter_set_bins.errcheck = _error_handler
_dll.openmc_energyfunc_filter_set_data.restype = c_int
_dll.openmc_energyfunc_filter_set_data.errcheck = _error_handler
_dll.openmc_energyfunc_filter_set_data.argtypes = [
    c_int32, c_size_t, POINTER(c_double), POINTER(c_double)]
_dll.openmc_energyfunc_filter_get_energy.resttpe = c_int
_dll.openmc_energyfunc_filter_get_energy.errcheck = _error_handler
_dll.openmc_energyfunc_filter_get_energy.argtypes = [
    c_int32, POINTER(c_size_t), POINTER(POINTER(c_double))]
_dll.openmc_energyfunc_filter_get_y.resttpe = c_int
_dll.openmc_energyfunc_filter_get_y.errcheck = _error_handler
_dll.openmc_energyfunc_filter_get_y.argtypes = [
    c_int32, POINTER(c_size_t), POINTER(POINTER(c_double))]
_dll.openmc_energyfunc_filter_get_interpolation.resttpe = c_int
_dll.openmc_energyfunc_filter_get_interpolation.errcheck = _error_handler
_dll.openmc_energyfunc_filter_get_interpolation.argtypes = [c_int32, POINTER(c_int)]
_dll.openmc_energyfunc_filter_set_interpolation.resttpe = c_int
_dll.openmc_energyfunc_filter_set_interpolation.errcheck = _error_handler
_dll.openmc_energyfunc_filter_set_interpolation.argtypes = [c_int32, c_char_p]
_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_num_bins.argtypes = [c_int32, POINTER(c_int)]
_dll.openmc_filter_get_num_bins.restype = c_int
_dll.openmc_filter_get_num_bins.errchck = _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_get_filter_index.argtypes = [c_int32, POINTER(c_int32)]
_dll.openmc_get_filter_index.restype = c_int
_dll.openmc_get_filter_index.errcheck = _error_handler
_dll.openmc_legendre_filter_get_order.argtypes = [c_int32, POINTER(c_int)]
_dll.openmc_legendre_filter_get_order.restype = c_int
_dll.openmc_legendre_filter_get_order.errcheck = _error_handler
_dll.openmc_legendre_filter_set_order.argtypes = [c_int32, c_int]
_dll.openmc_legendre_filter_set_order.restype = c_int
_dll.openmc_legendre_filter_set_order.errcheck = _error_handler
_dll.openmc_material_filter_get_bins.argtypes = [
    c_int32, POINTER(POINTER(c_int32)), POINTER(c_size_t)]
_dll.openmc_material_filter_get_bins.restype = c_int
_dll.openmc_material_filter_get_bins.errcheck = _error_handler
_dll.openmc_material_filter_set_bins.argtypes = [c_int32, c_size_t, POINTER(c_int32)]
_dll.openmc_material_filter_set_bins.restype = c_int
_dll.openmc_material_filter_set_bins.errcheck = _error_handler
_dll.openmc_mesh_filter_get_mesh.argtypes = [c_int32, POINTER(c_int32)]
_dll.openmc_mesh_filter_get_mesh.restype = c_int
_dll.openmc_mesh_filter_get_mesh.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
_dll.openmc_meshsurface_filter_get_mesh.argtypes = [c_int32, POINTER(c_int32)]
_dll.openmc_meshsurface_filter_get_mesh.restype = c_int
_dll.openmc_meshsurface_filter_get_mesh.errcheck = _error_handler
_dll.openmc_meshsurface_filter_set_mesh.argtypes = [c_int32, c_int32]
_dll.openmc_meshsurface_filter_set_mesh.restype = c_int
_dll.openmc_meshsurface_filter_set_mesh.errcheck = _error_handler
_dll.openmc_mesh_filter_get_translation.argtypes = [c_int32, POINTER(c_double*3)]
_dll.openmc_mesh_filter_get_translation.restype = c_int
_dll.openmc_mesh_filter_get_translation.errcheck = _error_handler
_dll.openmc_mesh_filter_set_translation.argtypes = [c_int32, POINTER(c_double*3)]
_dll.openmc_mesh_filter_set_translation.restype = c_int
_dll.openmc_mesh_filter_set_translation.errcheck = _error_handler
_dll.openmc_new_filter.argtypes = [c_char_p, POINTER(c_int32)]
_dll.openmc_new_filter.restype = c_int
_dll.openmc_new_filter.errcheck = _error_handler
_dll.openmc_spatial_legendre_filter_get_order.argtypes = [c_int32, POINTER(c_int)]
_dll.openmc_spatial_legendre_filter_get_order.restype = c_int
_dll.openmc_spatial_legendre_filter_get_order.errcheck = _error_handler
_dll.openmc_spatial_legendre_filter_set_order.argtypes = [c_int32, c_int]
_dll.openmc_spatial_legendre_filter_set_order.restype = c_int
_dll.openmc_spatial_legendre_filter_set_order.errcheck = _error_handler
_dll.openmc_sphharm_filter_get_order.argtypes = [c_int32, POINTER(c_int)]
_dll.openmc_sphharm_filter_get_order.restype = c_int
_dll.openmc_sphharm_filter_get_order.errcheck = _error_handler
_dll.openmc_sphharm_filter_set_order.argtypes = [c_int32, c_int]
_dll.openmc_sphharm_filter_set_order.restype = c_int
_dll.openmc_sphharm_filter_set_order.errcheck = _error_handler
_dll.openmc_zernike_filter_get_order.argtypes = [c_int32, POINTER(c_int)]
_dll.openmc_zernike_filter_get_order.restype = c_int
_dll.openmc_zernike_filter_get_order.errcheck = _error_handler
_dll.openmc_zernike_filter_set_order.argtypes = [c_int32, c_int]
_dll.openmc_zernike_filter_set_order.restype = c_int
_dll.openmc_zernike_filter_set_order.errcheck = _error_handler
_dll.tally_filters_size.restype = c_size_t

class Filter(_FortranObjectWithID):
    __instances = WeakValueDictionary()

    def __new__(cls, obj=None, uid=None, new=True, index=None):
        mapping = filters
        if index is None:
            if new:
                # Determine ID to assign
                if uid is None:
                    uid = max(mapping, default=0) + 1
                else:
                    if uid in mapping:
                        raise AllocationError('A filter with ID={} has already '
                                              'been allocated.'.format(uid))

                # Set the filter type -- note that the filter_type attribute
                # only exists on subclasses!
                index = c_int32()
                _dll.openmc_new_filter(cls.filter_type.encode(), index)
                index = index.value
            else:
                index = mapping[uid]._index

        if index not in cls.__instances:
            instance = super().__new__(cls)
            instance._index = index
            if uid is not None:
                instance.id = uid
            cls.__instances[index] = instance

        return cls.__instances[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)

    @property
    def n_bins(self):
        n = c_int()
        _dll.openmc_filter_get_num_bins(self._index, n)
        return n.value


[docs]class EnergyFilter(Filter): filter_type = 'energy' def __init__(self, bins=None, uid=None, new=True, index=None): super().__init__(uid, new, index) if bins is not None: self.bins = bins @property def bins(self): energies = POINTER(c_double)() n = c_size_t() _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 CollisionFilter(Filter): filter_type = 'collision' class EnergyoutFilter(EnergyFilter): filter_type = 'energyout' class AzimuthalFilter(Filter): filter_type = 'azimuthal' class CellFilter(Filter): filter_type = 'cell' @property def bins(self): cells = POINTER(c_int32)() n = c_int32() _dll.openmc_cell_filter_get_bins(self._index, cells, n) return as_array(cells, (n.value,)) class CellbornFilter(Filter): filter_type = 'cellborn' class CellfromFilter(Filter): filter_type = 'cellfrom' class CellInstanceFilter(Filter): filter_type = 'cellinstance' class DelayedGroupFilter(Filter): filter_type = 'delayedgroup' class DistribcellFilter(Filter): filter_type = 'distribcell' class EnergyFunctionFilter(Filter): filter_type = 'energyfunction' def __new__(cls, energy=None, y=None, uid=None, new=True, index=None): return super().__new__(cls, uid=uid, new=new, index=index) def __init__(self, energy=None, y=None, uid=None, new=True, index=None): if (energy is None) != (y is None): raise AttributeError("Need both energy and y or neither") super().__init__(uid, new, index) if energy is not None: self.set_data(energy, y) def set_data(self, energy, y): """Set the interpolation information for the filter Parameters ---------- energy : numpy.ndarray Independent variable for the interpolation y : numpy.ndarray Dependent variable for the interpolation interpolation : {'histogram', 'linear-linear', 'linear-log', 'log-linear', 'log-log', 'quadratic', 'cubic'} Interpolation scheme """ energy_array = np.asarray(energy) y_array = np.asarray(y) energy_p = energy_array.ctypes.data_as(POINTER(c_double)) y_p = y_array.ctypes.data_as(POINTER(c_double)) _dll.openmc_energyfunc_filter_set_data( self._index, len(energy_array), energy_p, y_p) @property def energy(self): return self._get_attr(_dll.openmc_energyfunc_filter_get_energy) @property def y(self): return self._get_attr(_dll.openmc_energyfunc_filter_get_y) @property def interpolation(self) -> str: interp = c_int() _dll.openmc_energyfunc_filter_get_interpolation(self._index, interp) return INTERPOLATION_SCHEME[interp.value] @interpolation.setter def interpolation(self, interp: str): interp_ptr = c_char_p(interp.encode()) _dll.openmc_energyfunc_filter_set_interpolation(self._index, interp_ptr) def _get_attr(self, cfunc): array_p = POINTER(c_double)() n = c_size_t() cfunc(self._index, n, array_p) return as_array(array_p, (n.value, )) class LegendreFilter(Filter): filter_type = 'legendre' def __init__(self, order=None, uid=None, new=True, index=None): super().__init__(uid, new, index) if order is not None: self.order = order @property def order(self): temp_order = c_int() _dll.openmc_legendre_filter_get_order(self._index, temp_order) return temp_order.value @order.setter def order(self, order): _dll.openmc_legendre_filter_set_order(self._index, order)
[docs]class MaterialFilter(Filter): filter_type = 'material' def __init__(self, bins=None, uid=None, new=True, index=None): super().__init__(uid, new, index) if bins is not None: self.bins = bins @property def bins(self): materials = POINTER(c_int32)() n = c_size_t() _dll.openmc_material_filter_get_bins(self._index, materials, n) return [Material(index=materials[i]) for i in range(n.value)] @bins.setter def bins(self, materials): # Get material indices as int32_t[] n = len(materials) bins = (c_int32*n)(*(m._index for m in materials)) _dll.openmc_material_filter_set_bins(self._index, n, bins)
[docs]class MeshFilter(Filter): filter_type = 'mesh' def __init__(self, mesh=None, uid=None, new=True, index=None): super().__init__(uid, new, index) if mesh is not None: self.mesh = mesh @property def mesh(self): index_mesh = c_int32() _dll.openmc_mesh_filter_get_mesh(self._index, index_mesh) return _get_mesh(index_mesh.value) @mesh.setter def mesh(self, mesh): _dll.openmc_mesh_filter_set_mesh(self._index, mesh._index) @property def translation(self): translation = (c_double*3)() _dll.openmc_mesh_filter_get_translation(self._index, translation) return tuple(translation) @translation.setter def translation(self, translation): _dll.openmc_mesh_filter_set_translation(self._index, (c_double*3)(*translation))
[docs]class MeshSurfaceFilter(Filter): filter_type = 'meshsurface' def __init__(self, mesh=None, uid=None, new=True, index=None): super().__init__(uid, new, index) if mesh is not None: self.mesh = mesh @property def mesh(self): index_mesh = c_int32() _dll.openmc_meshsurface_filter_get_mesh(self._index, index_mesh) return _get_mesh(index_mesh.value) @mesh.setter def mesh(self, mesh): _dll.openmc_meshsurface_filter_set_mesh(self._index, mesh._index) @property def translation(self): translation = (c_double*3)() _dll.openmc_mesh_filter_get_translation(self._index, translation) return tuple(translation) @translation.setter def translation(self, translation): _dll.openmc_mesh_filter_set_translation(self._index, (c_double*3)(*translation))
class MuFilter(Filter): filter_type = 'mu' class ParticleFilter(Filter): filter_type = 'particle' @property def bins(self): particle_i = np.zeros((self.n_bins,), dtype=c_int) _dll.openmc_particle_filter_get_bins( self._index, particle_i.ctypes.data_as(POINTER(c_int))) return [ParticleType(i) for i in particle_i] class PolarFilter(Filter): filter_type = 'polar' class SphericalHarmonicsFilter(Filter): filter_type = 'sphericalharmonics' def __init__(self, order=None, uid=None, new=True, index=None): super().__init__(uid, new, index) if order is not None: self.order = order @property def order(self): temp_order = c_int() _dll.openmc_sphharm_filter_get_order(self._index, temp_order) return temp_order.value @order.setter def order(self, order): _dll.openmc_sphharm_filter_set_order(self._index, order) class SpatialLegendreFilter(Filter): filter_type = 'spatiallegendre' def __init__(self, order=None, uid=None, new=True, index=None): super().__init__(uid, new, index) if order is not None: self.order = order @property def order(self): temp_order = c_int() _dll.openmc_spatial_legendre_filter_get_order(self._index, temp_order) return temp_order.value @order.setter def order(self, order): _dll.openmc_spatial_legendre_filter_set_order(self._index, order) class SurfaceFilter(Filter): filter_type = 'surface' class UniverseFilter(Filter): filter_type = 'universe' class ZernikeFilter(Filter): filter_type = 'zernike' def __init__(self, order=None, uid=None, new=True, index=None): super().__init__(uid, new, index) if order is not None: self.order = order @property def order(self): temp_order = c_int() _dll.openmc_zernike_filter_get_order(self._index, temp_order) return temp_order.value @order.setter def order(self, order): _dll.openmc_zernike_filter_set_order(self._index, order) class ZernikeRadialFilter(ZernikeFilter): filter_type = 'zernikeradial' _FILTER_TYPE_MAP = { 'azimuthal': AzimuthalFilter, 'cell': CellFilter, 'cellborn': CellbornFilter, 'cellfrom': CellfromFilter, 'cellinstance': CellInstanceFilter, 'delayedgroup': DelayedGroupFilter, 'distribcell': DistribcellFilter, 'energy': EnergyFilter, 'energyout': EnergyoutFilter, 'energyfunction': EnergyFunctionFilter, 'legendre': LegendreFilter, 'material': MaterialFilter, 'mesh': MeshFilter, 'meshsurface': MeshSurfaceFilter, 'mu': MuFilter, 'particle': ParticleFilter, 'polar': PolarFilter, 'sphericalharmonics': SphericalHarmonicsFilter, 'spatiallegendre': SpatialLegendreFilter, 'surface': SurfaceFilter, 'universe': UniverseFilter, 'zernike': ZernikeFilter, 'zernikeradial': ZernikeRadialFilter } 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=index) class _FilterMapping(Mapping): def __getitem__(self, key): index = c_int32() try: _dll.openmc_get_filter_index(key, index) except (AllocationError, InvalidIDError) as e: # __contains__ expects a KeyError to work correctly raise KeyError(str(e)) return _get_filter(index.value) def __iter__(self): for i in range(len(self)): yield _get_filter(i).id def __len__(self): return _dll.tally_filters_size() def __repr__(self): return repr(dict(self)) filters = _FilterMapping()