diff --git a/openmc/capi/cell.py b/openmc/capi/cell.py index fe3382589..e14df3f4a 100644 --- a/openmc/capi/cell.py +++ b/openmc/capi/cell.py @@ -7,12 +7,15 @@ from numpy.ctypeslib import as_array from . import _dll from .core import _FortranObjectWithID -from .error import _error_handler +from .error import _error_handler, AllocationError, InvalidIDError from .material import Material __all__ = ['Cell', 'cells'] # Cell functions +_dll.openmc_extend_cells.argtypes = [c_int32, POINTER(c_int32), POINTER(c_int32)] +_dll.openmc_extend_cells.restype = c_int +_dll.openmc_extend_cells.errcheck = _error_handler _dll.openmc_cell_get_id.argtypes = [c_int32, POINTER(c_int32)] _dll.openmc_cell_get_id.restype = c_int _dll.openmc_cell_get_id.errcheck = _error_handler @@ -51,11 +54,35 @@ class Cell(_FortranObjectWithID): """ __instances = WeakValueDictionary() - def __new__(cls, *args): - if args not in cls.__instances: - instance = super(Cell, self).__new__(cls) - cls.__instances[args] = instance - return cls.__instances[args] + def __new__(cls, uid=None, new=True, index=None): + mapping = cells + if index is None: + if new: + # Determine ID to assign + if uid is None: + try: + uid = max(mapping) + 1 + except ValueError: + uid = 1 + else: + if uid in mapping: + raise AllocationError('A cell with ID={} has already ' + 'been allocated.'.format(uid)) + + index = c_int32() + _dll.openmc_extend_cells(1, index, None) + index = index.value + else: + index = mapping[uid]._index + + if index not in cls.__instances: + instance = super(Cell, cls).__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): @@ -113,11 +140,11 @@ class _CellMapping(Mapping): except (AllocationError, InvalidIDError) as e: # __contains__ expects a KeyError to work correctly raise KeyError(str(e)) - return Cell(index.value) + return Cell(index=index.value) def __iter__(self): for i in range(len(self)): - yield Cell(i + 1).id + yield Cell(index=i + 1).id def __len__(self): return c_int32.in_dll(_dll, 'n_cells').value diff --git a/src/api.F90 b/src/api.F90 index 933faaa09..d0231c9a8 100644 --- a/src/api.F90 +++ b/src/api.F90 @@ -40,6 +40,7 @@ module openmc_api public :: openmc_energy_filter_get_bins public :: openmc_energy_filter_set_bins public :: openmc_extend_filters + public :: openmc_extend_cells public :: openmc_extend_materials public :: openmc_extend_tallies public :: openmc_filter_get_id diff --git a/src/geometry_header.F90 b/src/geometry_header.F90 index dd2cad7f1..63b1bbe95 100644 --- a/src/geometry_header.F90 +++ b/src/geometry_header.F90 @@ -427,6 +427,38 @@ contains ! C API FUNCTIONS !=============================================================================== + function openmc_extend_cells(n, index_start, index_end) result(err) bind(C) + ! Extend the cells array by n elements + integer(C_INT32_T), value, intent(in) :: n + integer(C_INT32_T), optional, intent(out) :: index_start + integer(C_INT32_T), optional, intent(out) :: index_end + integer(C_INT) :: err + + type(Cell), allocatable :: temp(:) ! temporary cells array + + if (n_cells == 0) then + ! Allocate cells array + allocate(cells(n)) + else + ! Allocate cells array with increased size + allocate(temp(n_cells + n)) + + ! Copy original cells to temporary array + temp(1:n_cells) = cells + + ! Move allocation from temporary array + call move_alloc(FROM=temp, TO=cells) + end if + + ! Return indices in cells array + if (present(index_start)) index_start = n_cells + 1 + if (present(index_end)) index_end = n_cells + n + n_cells = n_cells + n + + err = 0 + end function openmc_extend_cells + + function openmc_get_cell_index(id, index) result(err) bind(C) ! Return the index in the cells array of a cell with a given ID integer(C_INT32_T), value :: id