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