1from collections.abc import Iterable
2from io import StringIO
3from numbers import Real
4from warnings import warn
5
6import numpy as np
7
8import openmc.checkvalue as cv
9from openmc.mixin import EqualityMixin
10from openmc.stats import Univariate, Tabular, Uniform, Legendre
11from .function import INTERPOLATION_SCHEME
12from .data import EV_PER_MEV
13from .endf import get_head_record, get_cont_record, get_tab1_record, \
14    get_list_record, get_tab2_record
15
16
17class AngleDistribution(EqualityMixin):
18    """Angle distribution as a function of incoming energy
19
20    Parameters
21    ----------
22    energy : Iterable of float
23        Incoming energies in eV at which distributions exist
24    mu : Iterable of openmc.stats.Univariate
25        Distribution of scattering cosines corresponding to each incoming energy
26
27    Attributes
28    ----------
29    energy : Iterable of float
30        Incoming energies in eV at which distributions exist
31    mu : Iterable of openmc.stats.Univariate
32        Distribution of scattering cosines corresponding to each incoming energy
33
34    """
35
36    def __init__(self, energy, mu):
37        super().__init__()
38        self.energy = energy
39        self.mu = mu
40
41    @property
42    def energy(self):
43        return self._energy
44
45    @property
46    def mu(self):
47        return self._mu
48
49    @energy.setter
50    def energy(self, energy):
51        cv.check_type('angle distribution incoming energy', energy,
52                      Iterable, Real)
53        self._energy = energy
54
55    @mu.setter
56    def mu(self, mu):
57        cv.check_type('angle distribution scattering cosines', mu,
58                      Iterable, Univariate)
59        self._mu = mu
60
61    def to_hdf5(self, group):
62        """Write angle distribution to an HDF5 group
63
64        Parameters
65        ----------
66        group : h5py.Group
67            HDF5 group to write to
68
69        """
70
71        dset = group.create_dataset('energy', data=self.energy)
72
73        # Make sure all data is tabular
74        mu_tabular = [mu_i if isinstance(mu_i, Tabular) else
75                      mu_i.to_tabular() for mu_i in self.mu]
76
77        # Determine total number of (mu,p) pairs and create array
78        n_pairs = sum([len(mu_i.x) for mu_i in mu_tabular])
79        pairs = np.empty((3, n_pairs))
80
81        # Create array for offsets
82        offsets = np.empty(len(mu_tabular), dtype=int)
83        interpolation = np.empty(len(mu_tabular), dtype=int)
84        j = 0
85
86        # Populate offsets and pairs array
87        for i, mu_i in enumerate(mu_tabular):
88            n = len(mu_i.x)
89            offsets[i] = j
90            interpolation[i] = 1 if mu_i.interpolation == 'histogram' else 2
91            pairs[0, j:j+n] = mu_i.x
92            pairs[1, j:j+n] = mu_i.p
93            pairs[2, j:j+n] = mu_i.c
94            j += n
95
96        # Create dataset for distributions
97        dset = group.create_dataset('mu', data=pairs)
98
99        # Write interpolation as attribute
100        dset.attrs['offsets'] = offsets
101        dset.attrs['interpolation'] = interpolation
102
103    @classmethod
104    def from_hdf5(cls, group):
105        """Generate angular distribution from HDF5 data
106
107        Parameters
108        ----------
109        group : h5py.Group
110            HDF5 group to read from
111
112        Returns
113        -------
114        openmc.data.AngleDistribution
115            Angular distribution
116
117        """
118        energy = group['energy'][()]
119        data = group['mu']
120        offsets = data.attrs['offsets']
121        interpolation = data.attrs['interpolation']
122
123        mu = []
124        n_energy = len(energy)
125        for i in range(n_energy):
126            # Determine length of outgoing energy distribution and number of
127            # discrete lines
128            j = offsets[i]
129            if i < n_energy - 1:
130                n = offsets[i+1] - j
131            else:
132                n = data.shape[1] - j
133
134            interp = INTERPOLATION_SCHEME[interpolation[i]]
135            mu_i = Tabular(data[0, j:j+n], data[1, j:j+n], interp)
136            mu_i.c = data[2, j:j+n]
137
138            mu.append(mu_i)
139
140        return cls(energy, mu)
141
142    @classmethod
143    def from_ace(cls, ace, location_dist, location_start):
144        """Generate an angular distribution from ACE data
145
146        Parameters
147        ----------
148        ace : openmc.data.ace.Table
149            ACE table to read from
150        location_dist : int
151            Index in the XSS array corresponding to the start of a block,
152            e.g. JXS(9).
153        location_start : int
154            Index in the XSS array corresponding to the start of an angle
155            distribution array
156
157        Returns
158        -------
159        openmc.data.AngleDistribution
160            Angular distribution
161
162        """
163        # Set starting index for angle distribution
164        idx = location_dist + location_start - 1
165
166        # Number of energies at which angular distributions are tabulated
167        n_energies = int(ace.xss[idx])
168        idx += 1
169
170        # Incoming energy grid
171        energy = ace.xss[idx:idx + n_energies]*EV_PER_MEV
172        idx += n_energies
173
174        # Read locations for angular distributions
175        lc = ace.xss[idx:idx + n_energies].astype(int)
176        idx += n_energies
177
178        mu = []
179        for i in range(n_energies):
180            if lc[i] > 0:
181                # Equiprobable 32 bin distribution
182                n_bins = 32
183                idx = location_dist + abs(lc[i]) - 1
184                cos = ace.xss[idx:idx + n_bins + 1]
185                pdf = np.zeros(n_bins + 1)
186                pdf[:n_bins] = 1.0/(n_bins*np.diff(cos))
187                cdf = np.linspace(0.0, 1.0, n_bins + 1)
188
189                mu_i = Tabular(cos, pdf, 'histogram', ignore_negative=True)
190                mu_i.c = cdf
191            elif lc[i] < 0:
192                # Tabular angular distribution
193                idx = location_dist + abs(lc[i]) - 1
194                intt = int(ace.xss[idx])
195                n_points = int(ace.xss[idx + 1])
196                # Data is given as rows of (values, PDF, CDF)
197                data = ace.xss[idx + 2:idx + 2 + 3*n_points]
198                data.shape = (3, n_points)
199
200                mu_i = Tabular(data[0], data[1], INTERPOLATION_SCHEME[intt])
201                mu_i.c = data[2]
202            else:
203                # Isotropic angular distribution
204                mu_i = Uniform(-1., 1.)
205
206            mu.append(mu_i)
207
208        return cls(energy, mu)
209
210    @classmethod
211    def from_endf(cls, ev, mt):
212        """Generate an angular distribution from an ENDF evaluation
213
214        Parameters
215        ----------
216        ev : openmc.data.endf.Evaluation
217            ENDF evaluation
218        mt : int
219            The MT value of the reaction to get angular distributions for
220
221        Returns
222        -------
223        openmc.data.AngleDistribution
224            Angular distribution
225
226        """
227        file_obj = StringIO(ev.section[4, mt])
228
229        # Read HEAD record
230        items = get_head_record(file_obj)
231        lvt = items[2]
232        ltt = items[3]
233
234        # Read CONT record
235        items = get_cont_record(file_obj)
236        li = items[2]
237        nk = items[4]
238        center_of_mass = (items[3] == 2)
239
240        # Check for obsolete energy transformation matrix. If present, just skip
241        # it and keep reading
242        if lvt > 0:
243            warn('Obsolete energy transformation matrix in MF=4 angular '
244                 'distribution.')
245            for _ in range((nk + 5)//6):
246                file_obj.readline()
247
248        if ltt == 0 and li == 1:
249            # Purely isotropic
250            energy = np.array([0., ev.info['energy_max']])
251            mu = [Uniform(-1., 1.), Uniform(-1., 1.)]
252
253        elif ltt == 1 and li == 0:
254            # Legendre polynomial coefficients
255            params, tab2 = get_tab2_record(file_obj)
256            n_energy = params[5]
257
258            energy = np.zeros(n_energy)
259            mu = []
260            for i in range(n_energy):
261                items, al = get_list_record(file_obj)
262                temperature = items[0]
263                energy[i] = items[1]
264                coefficients = np.asarray([1.0] + al)
265                mu.append(Legendre(coefficients))
266
267        elif ltt == 2 and li == 0:
268            # Tabulated probability distribution
269            params, tab2 = get_tab2_record(file_obj)
270            n_energy = params[5]
271
272            energy = np.zeros(n_energy)
273            mu = []
274            for i in range(n_energy):
275                params, f = get_tab1_record(file_obj)
276                temperature = params[0]
277                energy[i] = params[1]
278                if f.n_regions > 1:
279                    raise NotImplementedError('Angular distribution with multiple '
280                                              'interpolation regions not supported.')
281                mu.append(Tabular(f.x, f.y, INTERPOLATION_SCHEME[f.interpolation[0]]))
282
283        elif ltt == 3 and li == 0:
284            # Legendre for low energies / tabulated for high energies
285            params, tab2 = get_tab2_record(file_obj)
286            n_energy_legendre = params[5]
287
288            energy_legendre = np.zeros(n_energy_legendre)
289            mu = []
290            for i in range(n_energy_legendre):
291                items, al = get_list_record(file_obj)
292                temperature = items[0]
293                energy_legendre[i] = items[1]
294                coefficients = np.asarray([1.0] + al)
295                mu.append(Legendre(coefficients))
296
297            params, tab2 = get_tab2_record(file_obj)
298            n_energy_tabulated = params[5]
299
300            energy_tabulated = np.zeros(n_energy_tabulated)
301            for i in range(n_energy_tabulated):
302                params, f = get_tab1_record(file_obj)
303                temperature = params[0]
304                energy_tabulated[i] = params[1]
305                if f.n_regions > 1:
306                    raise NotImplementedError('Angular distribution with multiple '
307                                              'interpolation regions not supported.')
308                mu.append(Tabular(f.x, f.y, INTERPOLATION_SCHEME[f.interpolation[0]]))
309
310            energy = np.concatenate((energy_legendre, energy_tabulated))
311
312        return AngleDistribution(energy, mu)
313