1from collections.abc import Mapping
2from ctypes import c_int, c_int32, c_double, c_char_p, POINTER, c_size_t
3from weakref import WeakValueDictionary
4
5import numpy as np
6from numpy.ctypeslib import as_array
7
8from openmc.exceptions import AllocationError, InvalidIDError, OpenMCError
9from . import _dll, Nuclide
10from .core import _FortranObjectWithID
11from .error import _error_handler
12
13
14__all__ = ['Material', 'materials']
15
16# Material functions
17_dll.openmc_extend_materials.argtypes = [c_int32, POINTER(c_int32), POINTER(c_int32)]
18_dll.openmc_extend_materials.restype = c_int
19_dll.openmc_extend_materials.errcheck = _error_handler
20_dll.openmc_get_material_index.argtypes = [c_int32, POINTER(c_int32)]
21_dll.openmc_get_material_index.restype = c_int
22_dll.openmc_get_material_index.errcheck = _error_handler
23_dll.openmc_material_add_nuclide.argtypes = [
24    c_int32, c_char_p, c_double]
25_dll.openmc_material_add_nuclide.restype = c_int
26_dll.openmc_material_add_nuclide.errcheck = _error_handler
27_dll.openmc_material_get_id.argtypes = [c_int32, POINTER(c_int32)]
28_dll.openmc_material_get_id.restype = c_int
29_dll.openmc_material_get_id.errcheck = _error_handler
30_dll.openmc_material_get_densities.argtypes = [
31    c_int32, POINTER(POINTER(c_int)), POINTER(POINTER(c_double)),
32    POINTER(c_int)]
33_dll.openmc_material_get_densities.restype = c_int
34_dll.openmc_material_get_densities.errcheck = _error_handler
35_dll.openmc_material_get_density.argtypes = [c_int32, POINTER(c_double)]
36_dll.openmc_material_get_density.restype = c_int
37_dll.openmc_material_get_density.errcheck = _error_handler
38_dll.openmc_material_get_temperature.argtypes = [c_int32, POINTER(c_double)]
39_dll.openmc_material_get_temperature.restype = c_int
40_dll.openmc_material_get_temperature.errcheck = _error_handler
41_dll.openmc_material_get_volume.argtypes = [c_int32, POINTER(c_double)]
42_dll.openmc_material_get_volume.restype = c_int
43_dll.openmc_material_get_volume.errcheck = _error_handler
44_dll.openmc_material_set_density.argtypes = [c_int32, c_double, c_char_p]
45_dll.openmc_material_set_density.restype = c_int
46_dll.openmc_material_set_density.errcheck = _error_handler
47_dll.openmc_material_set_densities.argtypes = [
48    c_int32, c_int, POINTER(c_char_p), POINTER(c_double)]
49_dll.openmc_material_set_densities.restype = c_int
50_dll.openmc_material_set_densities.errcheck = _error_handler
51_dll.openmc_material_set_id.argtypes = [c_int32, c_int32]
52_dll.openmc_material_set_id.restype = c_int
53_dll.openmc_material_set_id.errcheck = _error_handler
54_dll.openmc_material_get_name.argtypes = [c_int32, POINTER(c_char_p)]
55_dll.openmc_material_get_name.restype = c_int
56_dll.openmc_material_get_name.errcheck = _error_handler
57_dll.openmc_material_set_name.argtypes = [c_int32, c_char_p]
58_dll.openmc_material_set_name.restype = c_int
59_dll.openmc_material_set_name.errcheck = _error_handler
60_dll.openmc_material_set_volume.argtypes = [c_int32, c_double]
61_dll.openmc_material_set_volume.restype = c_int
62_dll.openmc_material_set_volume.errcheck = _error_handler
63_dll.n_materials.argtypes = []
64_dll.n_materials.restype = c_size_t
65
66
67class Material(_FortranObjectWithID):
68    """Material stored internally.
69
70    This class exposes a material that is stored internally in the OpenMC
71    library. To obtain a view of a material with a given ID, use the
72    :data:`openmc.lib.materials` mapping.
73
74    Parameters
75    ----------
76    uid : int or None
77        Unique ID of the material
78    new : bool
79        When `index` is None, this argument controls whether a new object is
80        created or a view to an existing object is returned.
81    index : int or None
82         Index in the `materials` array.
83
84    Attributes
85    ----------
86    id : int
87        ID of the material
88    nuclides : list of str
89        List of nuclides in the material
90    densities : numpy.ndarray
91        Array of densities in atom/b-cm
92    name : str
93        Name of the material
94    temperature : float
95        Temperature of the material in [K]
96    volume : float
97        Volume of the material in [cm^3]
98
99    """
100    __instances = WeakValueDictionary()
101
102    def __new__(cls, uid=None, new=True, index=None):
103        mapping = materials
104        if index is None:
105            if new:
106                # Determine ID to assign
107                if uid is None:
108                    uid = max(mapping, default=0) + 1
109                else:
110                    if uid in mapping:
111                        raise AllocationError('A material with ID={} has already '
112                                              'been allocated.'.format(uid))
113
114                index = c_int32()
115                _dll.openmc_extend_materials(1, index, None)
116                index = index.value
117            else:
118                index = mapping[uid]._index
119        elif index == -1:
120            # Special value indicates void material
121            return None
122
123        if index not in cls.__instances:
124            instance = super(Material, cls).__new__(cls)
125            instance._index = index
126            if uid is not None:
127                instance.id = uid
128            cls.__instances[index] = instance
129
130        return cls.__instances[index]
131
132    @property
133    def id(self):
134        mat_id = c_int32()
135        _dll.openmc_material_get_id(self._index, mat_id)
136        return mat_id.value
137
138    @id.setter
139    def id(self, mat_id):
140        _dll.openmc_material_set_id(self._index, mat_id)
141
142    @property
143    def name(self):
144        name = c_char_p()
145        _dll.openmc_material_get_name(self._index, name)
146        return name.value.decode()
147
148    @name.setter
149    def name(self, name):
150        name_ptr = c_char_p(name.encode())
151        _dll.openmc_material_set_name(self._index, name_ptr)
152
153    @property
154    def temperature(self):
155        temperature = c_double()
156        _dll.openmc_material_get_temperature(self._index, temperature)
157        return temperature.value
158
159    @property
160    def volume(self):
161        volume = c_double()
162        try:
163            _dll.openmc_material_get_volume(self._index, volume)
164        except OpenMCError:
165            return None
166        return volume.value
167
168    @volume.setter
169    def volume(self, volume):
170        _dll.openmc_material_set_volume(self._index, volume)
171
172    @property
173    def nuclides(self):
174        return self._get_densities()[0]
175        return nuclides
176
177    @property
178    def density(self):
179      density = c_double()
180      try:
181          _dll.openmc_material_get_density(self._index, density)
182      except OpenMCError:
183          return None
184      return density.value
185
186    @property
187    def densities(self):
188        return self._get_densities()[1]
189
190    def _get_densities(self):
191        """Get atom densities in a material.
192
193        Returns
194        -------
195        list of string
196            List of nuclide names
197        numpy.ndarray
198            Array of densities in atom/b-cm
199
200        """
201        # Allocate memory for arguments that are written to
202        nuclides = POINTER(c_int)()
203        densities = POINTER(c_double)()
204        n = c_int()
205
206        # Get nuclide names and densities
207        _dll.openmc_material_get_densities(self._index, nuclides, densities, n)
208
209        # Convert to appropriate types and return
210        nuclide_list = [Nuclide(nuclides[i]).name for i in range(n.value)]
211        density_array = as_array(densities, (n.value,))
212        return nuclide_list, density_array
213
214    def add_nuclide(self, name, density):
215        """Add a nuclide to a material.
216
217        Parameters
218        ----------
219        name : str
220            Name of nuclide, e.g. 'U235'
221        density : float
222            Density in atom/b-cm
223
224        """
225        _dll.openmc_material_add_nuclide(self._index, name.encode(), density)
226
227    def set_density(self, density, units='atom/b-cm'):
228        """Set density of a material.
229
230        Parameters
231        ----------
232        density : float
233            Density
234        units : {'atom/b-cm', 'g/cm3'}
235            Units for density
236
237        """
238        _dll.openmc_material_set_density(self._index, density, units.encode())
239
240    def set_densities(self, nuclides, densities):
241        """Set the densities of a list of nuclides in a material
242
243        Parameters
244        ----------
245        nuclides : iterable of str
246            Nuclide names
247        densities : iterable of float
248            Corresponding densities in atom/b-cm
249
250        """
251        # Convert strings to an array of char*
252        nucs = (c_char_p * len(nuclides))()
253        nucs[:] = [x.encode() for x in nuclides]
254
255        # Get numpy array as a double*
256        d = np.asarray(densities)
257        dp = d.ctypes.data_as(POINTER(c_double))
258
259        _dll.openmc_material_set_densities(self._index, len(nuclides), nucs, dp)
260
261
262class _MaterialMapping(Mapping):
263    def __getitem__(self, key):
264        index = c_int32()
265        try:
266            _dll.openmc_get_material_index(key, index)
267        except (AllocationError, InvalidIDError) as e:
268            # __contains__ expects a KeyError to work correctly
269            raise KeyError(str(e))
270        return Material(index=index.value)
271
272    def __iter__(self):
273        for i in range(len(self)):
274            yield Material(index=i).id
275
276    def __len__(self):
277        return _dll.n_materials()
278
279    def __repr__(self):
280        return repr(dict(self))
281
282materials = _MaterialMapping()
283