1from contextlib import contextmanager 2from ctypes import (c_bool, c_int, c_int32, c_int64, c_double, c_char_p, 3 c_char, POINTER, Structure, c_void_p, create_string_buffer) 4import sys 5 6import numpy as np 7from numpy.ctypeslib import as_array 8 9from . import _dll 10from .error import _error_handler 11import openmc.lib 12 13 14class _SourceSite(Structure): 15 _fields_ = [('r', c_double*3), 16 ('u', c_double*3), 17 ('E', c_double), 18 ('wgt', c_double), 19 ('delayed_group', c_int), 20 ('surf_id', c_int), 21 ('particle', c_int), 22 ('parent_id', c_int64), 23 ('progeny_id', c_int64)] 24 25 26# Define input type for numpy arrays that will be passed into C++ functions 27# Must be an int or double array, with single dimension that is contiguous 28_array_1d_int = np.ctypeslib.ndpointer(dtype=np.int32, ndim=1, 29 flags='CONTIGUOUS') 30_array_1d_dble = np.ctypeslib.ndpointer(dtype=np.double, ndim=1, 31 flags='CONTIGUOUS') 32 33_dll.openmc_calculate_volumes.restype = c_int 34_dll.openmc_calculate_volumes.errcheck = _error_handler 35_dll.openmc_cmfd_reweight.argtypes = c_bool, _array_1d_dble 36_dll.openmc_cmfd_reweight.restype = None 37_dll.openmc_finalize.restype = c_int 38_dll.openmc_finalize.errcheck = _error_handler 39_dll.openmc_find_cell.argtypes = [POINTER(c_double*3), POINTER(c_int32), 40 POINTER(c_int32)] 41_dll.openmc_find_cell.restype = c_int 42_dll.openmc_find_cell.errcheck = _error_handler 43_dll.openmc_hard_reset.restype = c_int 44_dll.openmc_hard_reset.errcheck = _error_handler 45_dll.openmc_init.argtypes = [c_int, POINTER(POINTER(c_char)), c_void_p] 46_dll.openmc_init.restype = c_int 47_dll.openmc_init.errcheck = _error_handler 48_dll.openmc_get_keff.argtypes = [POINTER(c_double*2)] 49_dll.openmc_get_keff.restype = c_int 50_dll.openmc_get_keff.errcheck = _error_handler 51_dll.openmc_initialize_mesh_egrid.argtypes = [ 52 c_int, _array_1d_int, c_double 53] 54_dll.openmc_initialize_mesh_egrid.restype = None 55_init_linsolver_argtypes = [_array_1d_int, c_int, _array_1d_int, c_int, c_int, 56 c_double, _array_1d_int, c_bool] 57_dll.openmc_initialize_linsolver.argtypes = _init_linsolver_argtypes 58_dll.openmc_initialize_linsolver.restype = None 59_dll.openmc_is_statepoint_batch.restype = c_bool 60_dll.openmc_master.restype = c_bool 61_dll.openmc_next_batch.argtypes = [POINTER(c_int)] 62_dll.openmc_next_batch.restype = c_int 63_dll.openmc_next_batch.errcheck = _error_handler 64_dll.openmc_plot_geometry.restype = c_int 65_dll.openmc_plot_geometry.restype = _error_handler 66_dll.openmc_run.restype = c_int 67_dll.openmc_run.errcheck = _error_handler 68_dll.openmc_reset.restype = c_int 69_dll.openmc_reset.errcheck = _error_handler 70_dll.openmc_reset_timers.restype = c_int 71_dll.openmc_reset_timers.errcheck = _error_handler 72_run_linsolver_argtypes = [_array_1d_dble, _array_1d_dble, _array_1d_dble, 73 c_double] 74_dll.openmc_run_linsolver.argtypes = _run_linsolver_argtypes 75_dll.openmc_run_linsolver.restype = c_int 76_dll.openmc_source_bank.argtypes = [POINTER(POINTER(_SourceSite)), POINTER(c_int64)] 77_dll.openmc_source_bank.restype = c_int 78_dll.openmc_source_bank.errcheck = _error_handler 79_dll.openmc_simulation_init.restype = c_int 80_dll.openmc_simulation_init.errcheck = _error_handler 81_dll.openmc_simulation_finalize.restype = c_int 82_dll.openmc_simulation_finalize.errcheck = _error_handler 83_dll.openmc_statepoint_write.argtypes = [c_char_p, POINTER(c_bool)] 84_dll.openmc_statepoint_write.restype = c_int 85_dll.openmc_statepoint_write.errcheck = _error_handler 86_dll.openmc_global_bounding_box.argtypes = [POINTER(c_double), 87 POINTER(c_double)] 88_dll.openmc_global_bounding_box.restype = c_int 89_dll.openmc_global_bounding_box.errcheck = _error_handler 90 91 92def global_bounding_box(): 93 """Calculate a global bounding box for the model""" 94 inf = sys.float_info.max 95 llc = np.zeros(3) 96 urc = np.zeros(3) 97 _dll.openmc_global_bounding_box(llc.ctypes.data_as(POINTER(c_double)), 98 urc.ctypes.data_as(POINTER(c_double))) 99 llc[llc == inf] = np.inf 100 urc[urc == inf] = np.inf 101 llc[llc == -inf] = -np.inf 102 urc[urc == -inf] = -np.inf 103 104 return llc, urc 105 106def calculate_volumes(): 107 """Run stochastic volume calculation""" 108 _dll.openmc_calculate_volumes() 109 110 111def current_batch(): 112 """Return the current batch of the simulation. 113 114 Returns 115 ------- 116 int 117 Current batch of the simulation 118 119 """ 120 return c_int.in_dll(_dll, 'current_batch').value 121 122 123def finalize(): 124 """Finalize simulation and free memory""" 125 _dll.openmc_finalize() 126 127 128def find_cell(xyz): 129 """Find the cell at a given point 130 131 Parameters 132 ---------- 133 xyz : iterable of float 134 Cartesian coordinates of position 135 136 Returns 137 ------- 138 openmc.lib.Cell 139 Cell containing the point 140 int 141 If the cell at the given point is repeated in the geometry, this 142 indicates which instance it is, i.e., 0 would be the first instance. 143 144 """ 145 index = c_int32() 146 instance = c_int32() 147 _dll.openmc_find_cell((c_double*3)(*xyz), index, instance) 148 return openmc.lib.Cell(index=index.value), instance.value 149 150 151def find_material(xyz): 152 """Find the material at a given point 153 154 Parameters 155 ---------- 156 xyz : iterable of float 157 Cartesian coordinates of position 158 159 Returns 160 ------- 161 openmc.lib.Material or None 162 Material containing the point, or None is no material is found 163 164 """ 165 index = c_int32() 166 instance = c_int32() 167 _dll.openmc_find_cell((c_double*3)(*xyz), index, instance) 168 169 mats = openmc.lib.Cell(index=index.value).fill 170 if isinstance(mats, (openmc.lib.Material, type(None))): 171 return mats 172 else: 173 return mats[instance.value] 174 175 176def hard_reset(): 177 """Reset tallies, timers, and pseudo-random number generator state.""" 178 _dll.openmc_hard_reset() 179 180 181def init(args=None, intracomm=None): 182 """Initialize OpenMC 183 184 Parameters 185 ---------- 186 args : list of str 187 Command-line arguments 188 intracomm : mpi4py.MPI.Intracomm or None 189 MPI intracommunicator 190 191 """ 192 if args is not None: 193 args = ['openmc'] + list(args) 194 else: 195 args = ['openmc'] 196 197 argc = len(args) 198 # Create the argv array. Note that it is actually expected to be of 199 # length argc + 1 with the final item being a null pointer. 200 argv = (POINTER(c_char) * (argc + 1))() 201 for i, arg in enumerate(args): 202 argv[i] = create_string_buffer(arg.encode()) 203 204 if intracomm is not None: 205 # If an mpi4py communicator was passed, convert it to void* to be passed 206 # to openmc_init 207 try: 208 from mpi4py import MPI 209 except ImportError: 210 intracomm = None 211 else: 212 address = MPI._addressof(intracomm) 213 intracomm = c_void_p(address) 214 215 _dll.openmc_init(argc, argv, intracomm) 216 217 218def is_statepoint_batch(): 219 """Return whether statepoint will be written in current batch or not. 220 221 Returns 222 ------- 223 bool 224 Whether is statepoint batch or not 225 226 """ 227 return _dll.openmc_is_statepoint_batch() 228 229 230def iter_batches(): 231 """Iterator over batches. 232 233 This function returns a generator-iterator that allows Python code to be run 234 between batches in an OpenMC simulation. It should be used in conjunction 235 with :func:`openmc.lib.simulation_init` and 236 :func:`openmc.lib.simulation_finalize`. For example: 237 238 .. code-block:: Python 239 240 with openmc.lib.run_in_memory(): 241 openmc.lib.simulation_init() 242 for _ in openmc.lib.iter_batches(): 243 # Look at convergence of tallies, for example 244 ... 245 openmc.lib.simulation_finalize() 246 247 See Also 248 -------- 249 openmc.lib.next_batch 250 251 """ 252 while True: 253 # Run next batch 254 status = next_batch() 255 256 # Provide opportunity for user to perform action between batches 257 yield 258 259 # End the iteration 260 if status != 0: 261 break 262 263 264def keff(): 265 """Return the calculated k-eigenvalue and its standard deviation. 266 267 Returns 268 ------- 269 tuple 270 Mean k-eigenvalue and standard deviation of the mean 271 272 """ 273 k = (c_double*2)() 274 _dll.openmc_get_keff(k) 275 return tuple(k) 276 277 278def master(): 279 """Return whether processor is master processor or not. 280 281 Returns 282 ------- 283 bool 284 Whether is master processor or not 285 286 """ 287 return _dll.openmc_master() 288 289 290def next_batch(): 291 """Run next batch. 292 293 Returns 294 ------- 295 int 296 Status after running a batch (0=normal, 1=reached maximum number of 297 batches, 2=tally triggers reached) 298 299 """ 300 status = c_int() 301 _dll.openmc_next_batch(status) 302 return status.value 303 304 305def plot_geometry(): 306 """Plot geometry""" 307 _dll.openmc_plot_geometry() 308 309 310def reset(): 311 """Reset tally results""" 312 _dll.openmc_reset() 313 314 315def reset_timers(): 316 """Reset timers.""" 317 _dll.openmc_reset_timers() 318 319 320def run(): 321 """Run simulation""" 322 _dll.openmc_run() 323 324 325def simulation_init(): 326 """Initialize simulation""" 327 _dll.openmc_simulation_init() 328 329 330def simulation_finalize(): 331 """Finalize simulation""" 332 _dll.openmc_simulation_finalize() 333 334 335def source_bank(): 336 """Return source bank as NumPy array 337 338 Returns 339 ------- 340 numpy.ndarray 341 Source sites 342 343 """ 344 # Get pointer to source bank 345 ptr = POINTER(_SourceSite)() 346 n = c_int64() 347 _dll.openmc_source_bank(ptr, n) 348 349 try: 350 # Convert to numpy array with appropriate datatype 351 bank_dtype = np.dtype(_SourceSite) 352 return as_array(ptr, (n.value,)).view(bank_dtype) 353 354 except ValueError as err: 355 # If a known numpy error was raised (github.com/numpy/numpy/issues 356 # /14214), re-raise with a more helpful error message. 357 if len(err.args) == 0: 358 raise err 359 if err.args[0].startswith('invalid shape in fixed-type tuple'): 360 raise ValueError('The source bank is too large to access via ' 361 'openmc.lib with this version of numpy. Use a different ' 362 'version of numpy or reduce the bank size (fewer particles ' 363 'per MPI process) so that it is smaller than 2 GB.') from err 364 else: 365 raise err 366 367 368def statepoint_write(filename=None, write_source=True): 369 """Write a statepoint file. 370 371 Parameters 372 ---------- 373 filename : str or None 374 Path to the statepoint to write. If None is passed, a default name that 375 contains the current batch will be written. 376 write_source : bool 377 Whether or not to include the source bank in the statepoint. 378 379 """ 380 if filename is not None: 381 filename = c_char_p(filename.encode()) 382 _dll.openmc_statepoint_write(filename, c_bool(write_source)) 383 384 385@contextmanager 386def run_in_memory(**kwargs): 387 """Provides context manager for calling OpenMC shared library functions. 388 389 This function is intended to be used in a 'with' statement and ensures that 390 OpenMC is properly initialized/finalized. At the completion of the 'with' 391 block, all memory that was allocated during the block is freed. For 392 example:: 393 394 with openmc.lib.run_in_memory(): 395 for i in range(n_iters): 396 openmc.lib.reset() 397 do_stuff() 398 openmc.lib.run() 399 400 Parameters 401 ---------- 402 **kwargs 403 All keyword arguments are passed to :func:`init`. 404 405 """ 406 init(**kwargs) 407 try: 408 yield 409 finally: 410 finalize() 411 412 413class _DLLGlobal: 414 """Data descriptor that exposes global variables from libopenmc.""" 415 def __init__(self, ctype, name): 416 self.ctype = ctype 417 self.name = name 418 419 def __get__(self, instance, owner): 420 return self.ctype.in_dll(_dll, self.name).value 421 422 def __set__(self, instance, value): 423 self.ctype.in_dll(_dll, self.name).value = value 424 425 426class _FortranObject: 427 def __repr__(self): 428 return "{}[{}]".format(type(self).__name__, self._index) 429 430 431class _FortranObjectWithID(_FortranObject): 432 def __init__(self, uid=None, new=True, index=None): 433 # Creating the object has already been handled by __new__. In the 434 # initializer, all we do is make sure that the object returned has an ID 435 # assigned. If the array index of the object is out of bounds, an 436 # OutOfBoundsError will be raised here by virtue of referencing self.id 437 self.id 438