simulations.py 57 KB
Newer Older
1
""" Module to setup a parametrized simulation, that work on probability distributions """
2
import os
3
import numpy as np
marcus's avatar
marcus committed
4
from astrotools import auger, coord, cosmic_rays, healpytools as hpt, nucleitools as nt
5

6
PATH = os.path.split(__file__)[0]
7

8

9
def set_fisher_smeared_sources(nside, sources, delta, source_fluxes=None):
MarCus's avatar
MarCus committed
10
    """
11
    Smears the source positions (optional fluxes) with a fisher distribution of width delta.
12
13

    :param nside: nside of the HEALPix pixelization (default: 64)
MarCus's avatar
MarCus committed
14
    :type nside: int
15
    :param sources: array of shape (3, n_sources) that point towards the center of the sources
16
    :param delta: float or array with same length as sources: width of the fisher distribution (in radians)
17
    :param source_fluxes: corresponding cosmic ray fluxes of the sources of shape (n_sources).
Marcus Wirtz's avatar
Marcus Wirtz committed
18
    :return: healpy map (with npix entries) for the smeared sources normalized to sum 1
MarCus's avatar
MarCus committed
19
    """
20
    npix = hpt.nside2npix(nside)
21
    nsrc = np.shape(sources)[1]
MarCus's avatar
MarCus committed
22
    eg_map = np.zeros(npix)
23
24
25
26
    if isinstance(delta, (int, float)):
        delta = np.repeat(delta, nsrc)
    if len(delta) != nsrc:
        raise ValueError("Number of deltas must be 1 or equal to number of sources")
27
    for i, v_src in enumerate(sources.T):
28
        eg_map_add = hpt.fisher_pdf(nside, *v_src, k=1. / delta[i] ** 2)
29
        if source_fluxes is not None:
30
31
            eg_map_add *= source_fluxes[i]
        eg_map += eg_map_add
MarCus's avatar
MarCus committed
32
    return eg_map / eg_map.sum()
33
34


35
36
37
38
39
40
41
42
43
44
45
46
47
def sample_from_m_distributions(weight_matrix, n):
    """
    Sample n events from m distributions each having k bins and defined in weight_matrix.

    :param weight_matrix: shape (m, k), hosting m different distributions each with k bins.
    :param n: Number of events drawn for each of the m distributions.
    :return indices: Indice array with values 0..(k-1), in shape (m, n)
    """
    weight_matrix /= weight_matrix.sum(axis=1, keepdims=True)
    s, r = weight_matrix.cumsum(axis=1), np.random.random(n)
    return (r[np.newaxis, np.newaxis] > s[:, :, np.newaxis]).sum(axis=1)


48
49
50
51
52
53
54
55
56
57
58
59
class BaseSimulation:
    """
    Base class where the classes ObservedBound() and SourceBound() inherit from.
    """

    def __init__(self, nsets, ncrs):
        nsets = int(nsets)
        ncrs = int(ncrs)
        self.shape = (nsets, ncrs)
        self.crs = cosmic_rays.CosmicRaysSets((nsets, ncrs))
        self.nsets = nsets
        self.ncrs = ncrs
60
        self.signal_label = None
61
62
63
64
65
66
67
68
69
70
71
72

    def set_xmax(self, z2a='double', model='EPOS-LHC'):
        """
        Calculate Xmax bei gumbel distribution for the simulated energies and charges.

        :param z2a: How the charge is converted to mass number ['double', 'empiric', 'stable', 'abundance']
        :param model: Hadronic interaction for gumbel distribution
        :return: no return
        """
        assert 'xmax' not in self.crs.keys(), "Try to re-assign xmax values!"

        if (not hasattr(self.crs, 'log10e')) or (not hasattr(self.crs, 'charge')):
73
74
            raise Exception("""Use function set_energy() and set_charges() before using function set_xmax.
                            If you Use SourceBound simulation attenuate() also has to be used additionally""")
75
76
77
78
79
80
81
        mass = getattr(nt.Charge2Mass(self.crs['charge']), z2a)()
        mass = np.hstack(mass) if isinstance(mass, np.ndarray) else mass
        xmax = auger.rand_xmax(np.hstack(self.crs['log10e']), mass, model=model)
        self.crs['xmax'] = np.reshape(xmax, self.shape)

        return self.crs['xmax']

82
83
84
85
86
87
88
89
90
91
92
    def apply_uncertainties(self, err_e=None, err_a=None, err_xmax=None):
        """
        Apply measurement uncertainties.

        :param err_e: relative uncertainty on the energy (typical 0.14)
        :param err_a: angular uncertainty in degree on the arrival directions (typical 0.5 degree)
        :param err_xmax: absolute uncertainty on the shower depth xmax (typical 15 g/cm^2)
        """
        if err_e is not None:
            self.crs['log10e'] += np.log10(1 + np.random.normal(scale=err_e, size=self.shape))

93
94
95
96
97
        if err_a is not None:
            vecs = coord.rand_fisher_vec(self.crs['vecs'], 1./np.deg2rad(err_a)**2)
            lon, lat = coord.vec2ang(vecs)
            self.crs['lon'] = lon.reshape(self.shape)
            self.crs['lat'] = lat.reshape(self.shape)
98
99
100
101

        if err_xmax is not None:
            self.crs['xmax'] += np.random.normal(err_xmax)

102
103
104
105
    def shuffle_events(self):
        """
        Independently shuffle the cosmic rays of each set.
        """
106
        # This function can be simplified in the future using np.take_along_axis()
107
        shuffle_ids = np.random.permutation(np.prod(self.shape)).reshape(self.shape)
108
        shuffle_ids = np.argsort(shuffle_ids, axis=1)
109
110
        sets_ids = np.repeat(np.arange(self.nsets), self.ncrs).reshape(self.shape)
        for _key in self.crs.shape_array.dtype.names:
111
            self.crs[_key] = self.crs[_key][sets_ids, shuffle_ids]
112
113
114
115
        self.signal_label = self.signal_label[sets_ids, shuffle_ids]
        if 'vecs' in self.crs.keys():
            sets_ids_3d = np.repeat(np.arange(3), np.prod(self.shape)).reshape((3,) + self.shape)
            self.crs['vecs'] = self.crs['vecs'][sets_ids_3d, sets_ids, np.stack([shuffle_ids] * 3)]
116

117
118

class ObservedBound(BaseSimulation):
MarCus's avatar
MarCus committed
119
    """
MarCus's avatar
MarCus committed
120
121
122
    Class to simulate cosmic ray arrival scenario by sources located at the sky, including energies, charges, smearing
    and galactic magnetic field effects.
    This is an observed bound simulation, thus energies and composition is set on Earth and might differ at sources.
MarCus's avatar
MarCus committed
123
    """
Martin Urban's avatar
Martin Urban committed
124

125
    def __init__(self, nside, nsets, ncrs):
MarCus's avatar
MarCus committed
126
        """
127
        Initialization of ObservedBound simulation.
128

129
        :param nside: nside of the HEALPix pixelization (default: 64)
130
        :param nsets: number of simulated cosmic ray sets
131
        :param ncrs: number of cosmic rays per set
MarCus's avatar
MarCus committed
132
        """
133
        BaseSimulation.__init__(self, nsets, ncrs)
134
        self.nside = nside
MarCus's avatar
MarCus committed
135
        self.npix = hpt.nside2npix(nside)
136

137
        self.crs['nside'] = nside
138
139
140
        self.sources = None
        self.source_fluxes = None

Martin Urban's avatar
Martin Urban committed
141
        self.rigidities = None
MarCus's avatar
MarCus committed
142
143
        self.rig_bins = None
        self.cr_map = None
144
        self.lensed = None
Marcus Wirtz's avatar
Marcus Wirtz committed
145
        self.exposure = {'map': None, 'a0': None, 'zmax': None}
marcus's avatar
marcus committed
146
        self.signal_idx = None
147

148
    def set_energy(self, log10e_min, log10e_max=None, energy_spectrum=None, **kwargs):
MarCus's avatar
MarCus committed
149
        """
150
        Setting the energies of the simulated cosmic ray set.
151

152
153
        :param log10e_min: Either minimum energy (in log10e) for AUGER setup or power-law
                           or numpy.array of energies in shape (nsets, ncrs)
154
155
        :type log10e_min: Union[np.ndarray, float]
        :param log10e_max: Maximum energy for AUGER setup
156
        :param energy_spectrum: model that is defined in below class EnergySpectrum
157
        :param kwargs: additional named keywords which will be passed to class EnergySpectrum
158
        :return: energies in log10e
MarCus's avatar
MarCus committed
159
        """
160
161
        assert 'log10e' not in self.crs.keys(), "Try to re-assign energies!"

162
163
164
165
        if isinstance(log10e_min, np.ndarray):
            if log10e_min.shape == self.shape:
                self.crs['log10e'] = log10e_min
            elif log10e_min.size == self.ncrs:
166
167
                print("Warning: the same energies have been used for all simulated sets (nsets).")
                self.crs['log10e'] = np.tile(log10e_min, self.nsets).reshape(self.shape)
168
            else:
169
                raise Exception("Shape of input energies not in format (nsets, ncrs).")
170
        elif isinstance(log10e_min, (float, np.float, int, np.int)):
171
172
173
174
175
176
177
            if energy_spectrum is not None:
                log10e = getattr(EnergySpectrum(self.shape, log10e_min, log10e_max), energy_spectrum)(**kwargs)
            else:
                if (log10e_min < 17.) or (log10e_min > 21.):
                    print("Warning: Specified parameter log10e_min below 17 or above 20.5.")
                log10e = auger.rand_energy_from_auger(self.shape, log10e_min=log10e_min, log10e_max=log10e_max)
            self.crs['log10e'] = log10e
178
179
180
        else:
            raise Exception("Input of emin could not be understood.")

181
182
        return self.crs['log10e']

183
    def set_charges(self, charge, **kwargs):
MarCus's avatar
MarCus committed
184
        """
185
        Setting the charges of the simulated cosmic ray set.
186

187
        :param charge: Either charge number that is used or numpy.array of charges in shape (nsets, ncrs) or keyword
Martin Urban's avatar
Martin Urban committed
188
        :type: charge: Union[np.ndarray, str, float]
189
        :return: charges
MarCus's avatar
MarCus committed
190
        """
191
192
        assert 'charge' not in self.crs.keys(), "Try to re-assign charges!"

193
194
        if isinstance(charge, np.ndarray):
            if charge.shape == self.shape:
195
                self.crs['charge'] = charge
196
            elif charge.size == self.ncrs:
197
                self.crs['charge'] = np.tile(charge, self.nsets).reshape(self.shape)
198
            else:
199
                raise Exception("Shape of input charges not in format (nsets, ncrs).")
200
        elif isinstance(charge, (float, np.float, int, np.int)):
marcus's avatar
marcus committed
201
            self.crs['charge'] = charge
202
        elif isinstance(charge, str):
MarCus's avatar
MarCus committed
203
204
            if not hasattr(self.crs, 'log10e'):
                raise Exception("Use function set_energy() before accessing a composition model.")
205
            self.crs['charge'] = getattr(CompositionModel(self.shape, self.crs['log10e']), charge.lower())(**kwargs)
206
207
208
        elif isinstance(charge, dict):
            z = np.array([nt.ELEMENT_CHARGE[key] for key in charge])
            self.crs['charge'] = np.random.choice(z, self.shape, p=[charge[key] for key in charge])
209
210
211
        else:
            raise Exception("Input of charge could not be understood.")

212
213
        return self.crs['charge']

214
215
216
217
218
219
220
    def set_xmax(self, z2a='double', model='EPOS-LHC'):
        """
        Calculate Xmax bei gumbel distribution for the simulated energies and charges.
        For more information refer to BaseSimulation.set_xmax().
        """
        return BaseSimulation.set_xmax(self, z2a, model)

221
    def set_sources(self, sources, fluxes=None):
MarCus's avatar
MarCus committed
222
        """
223
        Define source position and optional weights (cosmic ray luminosity).
224

225
226
        :param sources: array of shape (3, n_sources) that point towards the center of the sources or integer for
                        number of random sources or keyword ['sbg']
227
        :param fluxes: corresponding cosmic ray fluxes of the sources of shape (n_sources).
228
        :return: no return
MarCus's avatar
MarCus committed
229
        """
230
        if isinstance(sources, np.ndarray):
231
232
233
234
            if (len(np.shape(sources)) == 1) and len(sources) == 3:
                sources = np.reshape(sources, (3, 1))
            assert len(np.shape(sources)) == 2
            assert np.shape(sources)[0] == 3
235
            self.sources = sources
Marcus Wirtz's avatar
Marcus Wirtz committed
236
237
238
            if fluxes is not None:
                assert fluxes.size == len(sources.T)
                self.source_fluxes = fluxes
239
        elif isinstance(sources, (int, np.int)):
MarCus's avatar
MarCus committed
240
            src_pix = np.random.randint(0, self.npix, sources)
241
            self.sources = np.array(hpt.pix2vec(self.nside, src_pix))
Marcus Wirtz's avatar
Marcus Wirtz committed
242
243
244
            if fluxes is not None:
                assert fluxes.size == sources
                self.source_fluxes = fluxes
245
        elif isinstance(sources, str):
Marcus Wirtz's avatar
Marcus Wirtz committed
246
            self.sources, self.source_fluxes = getattr(SourceScenario(), sources.lower())()[:2]
247
248
249
        else:
            raise Exception("Source scenario not understood.")

250
    def set_rigidity_bins(self, lens_or_bins, cover_rigidity=True):
MarCus's avatar
MarCus committed
251
        """
252
        Defines the bin centers of the rigidities.
253
254

        :param lens_or_bins: Either the binning of the lens (left bin borders) or the lens itself
255
        :return: no return
MarCus's avatar
MarCus committed
256
257
        """
        if self.rig_bins is None:
258
            if 'log10e' not in self.crs.keys():
259
                raise Exception("Cannot define rigidity bins without energies specified.")
260
            if 'charge' not in self.crs.keys():
MarCus's avatar
MarCus committed
261
                print("Warning: Energy dependent deflection instead of rigidity dependent (set_charges to avoid)")
262
263

            if isinstance(lens_or_bins, np.ndarray):
Martin Urban's avatar
Martin Urban committed
264
                bins = lens_or_bins  # type: np.array
Marcus Wirtz's avatar
Marcus Wirtz committed
265
                dbins = bins[1] - bins[0]
266
            else:
Marcus Wirtz's avatar
Marcus Wirtz committed
267
                bins = np.array(lens_or_bins.log10r_mins)
Marcus Wirtz's avatar
Marcus Wirtz committed
268
                dbins = lens_or_bins.dlog10r
269
270
            rigidities = self.crs['log10e']
            if 'charge' in self.crs.keys():
271
                rigidities = rigidities - np.log10(self.crs['charge'])
272
273
            if cover_rigidity:
                assert np.all(np.min(rigidities) >= np.min(bins - dbins / 2.)), "Rigidities not covered by lens!"
274
            idx = np.digitize(rigidities, bins) - 1
Marcus Wirtz's avatar
Marcus Wirtz committed
275
            rigs = (bins + dbins / 2.)[idx]
276
            rigs = rigs.reshape(self.shape)
PaulaBigalke's avatar
PaulaBigalke committed
277
            self.rigidities = rigs
MarCus's avatar
MarCus committed
278
            self.rig_bins = np.unique(rigs)
279

MarCus's avatar
MarCus committed
280
        return self.rig_bins
281

282
    def smear_sources(self, delta, dynamic=None):
MarCus's avatar
MarCus committed
283
        """
284
        Smears the source positions with a fisher distribution of width delta (optional dynamic smearing).
285

286
287
        :param delta: either the constant width of the fisher distribution or dynamic (delta / R[10 EV]), in radians
        :param dynamic: if True, function applies dynamic smearing (delta / R[EV]) with value delta at 10 EV rigidity
288
        :return: no return
MarCus's avatar
MarCus committed
289
        """
290
291
292
        if self.sources is None:
            raise Exception("Cannot smear sources without positions.")

293
        if (dynamic is None) or (dynamic is False):
MarCus's avatar
MarCus committed
294
            shape = (1, self.npix)
295
            eg_map = np.reshape(set_fisher_smeared_sources(self.nside, self.sources, delta, self.source_fluxes), shape)
296
        else:
MarCus's avatar
MarCus committed
297
            if self.rig_bins is None:
298
                raise Exception("Cannot dynamically smear sources without rigidity bins (use set_rigidity_bins()).")
MarCus's avatar
MarCus committed
299
            eg_map = np.zeros((self.rig_bins.size, self.npix))
MarCus's avatar
MarCus committed
300
            for i, rig in enumerate(self.rig_bins):
Martin Urban's avatar
Martin Urban committed
301
                delta_temp = delta / 10 ** (rig - 19.)
302
                eg_map[i] = set_fisher_smeared_sources(self.nside, self.sources, delta_temp, self.source_fluxes)
MarCus's avatar
MarCus committed
303
        self.cr_map = eg_map
304

MarCus's avatar
MarCus committed
305
    def lensing_map(self, lens, cache=None):
MarCus's avatar
MarCus committed
306
        """
307
        Apply a galactic magnetic field to the extragalactic map.
308

Martin Urban's avatar
Martin Urban committed
309
        :param lens: Instance of astrotools.gamale.Lens class, for the galactic magnetic field
MarCus's avatar
MarCus committed
310
        :param cache: Caches all the loaded lens parts (increases speed, but may consume a lot of memory!)
311
        :return: no return
MarCus's avatar
MarCus committed
312
        """
313
314
        if self.lensed:
            print("Warning: Cosmic Ray maps were already lensed before.")
315

MarCus's avatar
MarCus committed
316
        if self.rig_bins is None:
317
318
            self.set_rigidity_bins(lens)

MarCus's avatar
MarCus committed
319
        if self.cr_map is None:
320
            self._fix_point_source()
MarCus's avatar
MarCus committed
321

322
        arrival_map = np.zeros((self.rig_bins.size, self.npix))
MarCus's avatar
MarCus committed
323
        for i, rig in enumerate(self.rig_bins):
324
            lp = lens.get_lens_part(rig, cache=cache)
MarCus's avatar
MarCus committed
325
            eg_map_bin = self.cr_map[0] if self.cr_map.size == self.npix else self.cr_map[i]
MarCus's avatar
MarCus committed
326
            lensed_map = lp.dot(eg_map_bin)
327
328
            if not cache:
                del lp.data, lp.indptr, lp.indices
MarCus's avatar
MarCus committed
329
            arrival_map[i] = lensed_map / np.sum(lensed_map) if np.sum(lensed_map) > 0 else 1. / self.npix
330

331
        self.lensed = True
MarCus's avatar
MarCus committed
332
        self.cr_map = arrival_map
333

334
    def apply_exposure(self, a0=-35.25, zmax=60):
MarCus's avatar
MarCus committed
335
        """
336
337
338
        Apply the exposure (coverage) of any experiment (default: AUGER) to the observed maps.

        :param a0: equatorial declination [deg] of the experiment (default: AUGER, a0=-35.25 deg)
MarCus's avatar
MarCus committed
339
        :type a0: float, int
340
        :param zmax: maximum zenith angle [deg] for the events
341
        :return: no return
MarCus's avatar
MarCus committed
342
        """
Marcus Wirtz's avatar
Marcus Wirtz committed
343
344
345
346
347
        self.exposure.update({'map': hpt.exposure_pdf(self.nside, a0, zmax), 'a0': a0, 'zmax': zmax})
        if self.cr_map is None:
            self.cr_map = np.reshape(self.exposure['map'], (1, self.npix))
        else:
            self.cr_map = self.cr_map * self.exposure['map']
MarCus's avatar
MarCus committed
348
        self.cr_map /= np.sum(self.cr_map, axis=-1)[:, np.newaxis]
349

350
    def arrival_setup(self, fsig):
MarCus's avatar
MarCus committed
351
        """
352
353
354
        Creates the realizations of the arrival maps.

        :param fsig: signal fraction of cosmic rays per set (signal = originate from sources)
MarCus's avatar
MarCus committed
355
        :type fsig: float
356
        :return: no return
MarCus's avatar
MarCus committed
357
        """
358
359
        dtype = np.uint16 if self.nside <= 64 else np.uint32
        pixel = np.zeros(self.shape).astype(dtype)
360
361
362

        # Setup the signal part
        n_sig = int(fsig * self.ncrs)
363
        signal_label = np.in1d(range(self.ncrs), np.arange(0, n_sig, 1))
364
365
        if n_sig == 0:
            pass
366
367
368
369
        elif (self.cr_map is None) and (self.sources is None):
            pixel[:, signal_label] = np.random.choice(self.npix, (self.nsets, n_sig))
        elif (self.cr_map is None):
            self._fix_point_source()
370
371
372
373
374
375
376
377
378
            if self.cr_map.size == self.npix:
                pixel[:, signal_label] = np.random.choice(self.npix, (self.nsets, n_sig), p=np.hstack(self.cr_map))
            else:
                for i, rig in enumerate(self.rig_bins):
                    mask_rig = (rig == self.rigidities) * signal_label  # type: np.ndarray
                    n = np.sum(mask_rig)
                    if n == 0:
                        continue
                    pixel[mask_rig] = np.random.choice(self.npix, n, p=self.cr_map[i])
379
        elif np.sum(self.cr_map) > 0:
MarCus's avatar
MarCus committed
380
            if self.cr_map.size == self.npix:
381
                pixel[:, signal_label] = np.random.choice(self.npix, (self.nsets, n_sig), p=np.hstack(self.cr_map))
382
383
            else:
                for i, rig in enumerate(self.rig_bins):
384
                    mask_rig = (rig == self.rigidities) * signal_label  # type: np.ndarray
385
                    n = np.sum(mask_rig)
386
387
                    if n == 0:
                        continue
MarCus's avatar
MarCus committed
388
                    pixel[mask_rig] = np.random.choice(self.npix, n, p=self.cr_map[i])
389
390
        else:
            raise Exception("No signal probability to sample signal from!")
391

392
        # Setup the background part
MarCus's avatar
MarCus committed
393
        n_back = self.ncrs - n_sig
Marcus Wirtz's avatar
Marcus Wirtz committed
394
        bpdf = self.exposure['map'] if self.exposure['map'] is not None else np.ones(self.npix) / float(self.npix)
395
        pixel[:, np.invert(signal_label)] = np.random.choice(self.npix, (self.nsets, n_back), p=bpdf)
396

397
        self.signal_label = np.repeat(signal_label[np.newaxis], self.nsets, axis=0)
398
        self.crs['pixel'] = pixel
399

400
401
402
    def apply_uncertainties(self, err_e=None, err_a=None, err_xmax=None):
        """ Apply measurement uncertainties (see BaseSimulation.apply_uncertainties()). """
        BaseSimulation.apply_uncertainties(self, err_e, err_a, err_xmax)
403

404
    def get_data(self, convert_all=False, shuffle=False):
MarCus's avatar
MarCus committed
405
        """
406
        Returns the data in the form of the cosmic_rays.CosmicRaysSets() container.
407
408
409

        :param convert_all: if True, also vectors and lon/lat of the cosmic ray sets are saved (more memory expensive)
        :type convert_all: bool
410
        :return: Instance of cosmic_rays.CosmicRaysSets() classes
411

412
                 Example:
413
                 sim = ObservedBound()
414
                 ...
415
                 crs = sim.get_data(convert_all=True)
416
417
418
419
420
                 pixel = crs['pixel']
                 lon = crs['lon']
                 lat = crs['lat']
                 log10e = crs['log10e']
                 charge = crs['charge']
MarCus's avatar
MarCus committed
421
        """
422
        if convert_all is True:
423
            if not hasattr(self.crs, 'lon') or not hasattr(self.crs, 'lat'):
424
                self.convert_pixel(convert_all=True)
425
426
        if shuffle:
            self.shuffle_events()
427
        self.crs['signal_label'] = self.signal_label
428
        return self.crs
429

430
    def convert_pixel(self, keyword='vecs', convert_all=False):
431
432
        """
        Converts pixelized information stored under key 'pixel' to vectors (keyword='vecs')
433
        or angles (keyword='angles'), accessible via 'lon', 'lat'. When convert_all is True, both are saved.
434
435
436
437
438
439
        This can be used at a later stage, if convert_all was set to False originally.
        """
        shape = (-1, self.shape[0], self.shape[1])
        if self.exposure['map'] is not None:
            a0 = self.exposure['a0']
            zmax = self.exposure['zmax']
440
            vecs = hpt.rand_exposure_vec_in_pix(self.nside, np.hstack(self.crs['pixel']), a0, zmax).reshape(shape)
441
442
        else:
            vecs = hpt.rand_vec_in_pix(self.nside, np.hstack(self.crs['pixel'])).reshape(shape)
443
444
445
446
447
448
449
450
451
452
453
454
455
        if keyword == 'vecs' or convert_all is True:
            if hasattr(self.crs, 'lon') and hasattr(self.crs, 'lat') and not all:
                raise Exception('Not useful to convert pixels to vecs, information already there!')
            self.crs['vecs'] = vecs
        if keyword == 'angles' or convert_all is True:
            if keyword == 'angles' and not convert_all:
                if hasattr(self.crs, 'vecs') and not convert_all:
                    raise Exception('Not useful to convert pixels to angles, information already there!')
            lon, lat = coord.vec2ang(vecs)
            self.crs['lon'] = lon.reshape(self.shape)
            self.crs['lat'] = lat.reshape(self.shape)
        else:
            raise Exception('keyword not understood! Use angles or vecs!')
456

457
458
459
460
461
462
463
464
465
    def _fix_point_source(self):
        """ Internal function to set non-smeared sources automatically """
        print("Warning: No extragalactic smearing of the sources performed before lensing (smear_sources). "
              "Sources are considered point-like.")
        eg_map = np.zeros((1, self.npix))
        weights = self.source_fluxes if self.source_fluxes is not None else 1.
        eg_map[:, hpt.vec2pix(self.nside, *self.sources)] = weights
        self.cr_map = eg_map

466

467
class SourceBound(BaseSimulation):
468
    """
469
    Class to simulate cosmic ray arrival scenario by sources located at the sky, including energies, charges, smearing.
470
471
472
473
    This is a source bound simulation, thus energies and composition is set at the sources and might differ at Earth.
    """

    def __init__(self, nsets, ncrs):
474
475
        """
        Initialization of SourceBound simulation.
476

477
478
479
480
481
482
483
484
        :param nsets: number of simulated cosmic ray sets
        :param ncrs: number of cosmic rays per set
        """
        BaseSimulation.__init__(self, nsets, ncrs)

        self.arrival_matrix, self.source_matrix = None, None
        self.dis_bins, self.log10e_bins = None, None
        self.energy_setting = None
485
        self.charge_weights = None
486
        self.universe = SourceGeometry(nsets)
487

488
    def set_energy(self, log10e_min, gamma=-2, log10_cut=20., rig_cut=True):
489
490
491
492
        """
        Define energy spectrum and cut off energy of sources.

        :param log10e_min: Minimum threshold energy for observation
493
494
        :param gamma: Spectral index of energy spectrum at sources
        :param log10_cut: Maximum cut-off energy or rigidity for sources
495
        :param rig_cut: if True, log10_cut refers to a rigidity cut
496
        """
497
        self.energy_setting = {'log10e_min': log10e_min, 'gamma': gamma, 'log10_cut': log10_cut, 'rig_cut': rig_cut}
498
499
500

    def set_charges(self, charges):
        """
501
        Define fraction of charge groups in form of dictionary (e.g. {'h':0.5, 'fe':0.5}) at source
502
        or as keyword 'first_minimum'/'second_minimum' from Auger's best fit paper (arXiv:1612.07155)
503
504
505
506
507
508
509
510
511
512
513
514
        If string is given, gamma and Rcut are also set to the respective best fit values.

        :param charges: dictionary hosting the fractions of injected elements (H, He, N, Si, Fe possible)
                        or string ('first_minimum'/'second_minimum')
        """
        if isinstance(charges, dict):
            fraction = np.sum([charges[key] for key in charges])
            assert np.abs(fraction - 1) < 1e-4, "Fractions of charges dictionary must be normalized!"
            self.charge_weights = charges

        elif isinstance(charges, str):
            if charges == 'first_minimum':
515
516
                self.energy_setting['gamma'], self.energy_setting['log10_cut'] = -0.87, 18.62
                self.charge_weights = {'n': 0.88, 'si': 0.12}
517
            elif charges == 'second_minimum':
518
519
520
521
522
523
524
525
                self.energy_setting['gamma'], self.energy_setting['log10_cut'] = 1.03, 18.21
                self.charge_weights = {'h': 0.6794, 'he': 0.31, 'n': 0.01, 'si': 0.0006}
            elif charges == 'first_minimum_walz':
                self.energy_setting['gamma'], self.energy_setting['log10_cut'] = -0.62, 18.56
                self.charge_weights = {'h': 0.001, 'he': 0.001, 'n': 0.985, 'fe': 0.012}
            elif charges == 'second_minimum_walz':
                self.energy_setting['gamma'], self.energy_setting['log10_cut'] = -2.03, 19.88
                self.charge_weights = {'he': 0.003, 'n': 0.92, 'fe': 0.077}
526
527
528
529
530
            else:
                raise Exception('Charge keyword not understood (first_minimum/second_minimum)')
            self.energy_setting['rig_cut'] = True
        else:
            raise Exception('Charge type not understood (dictionary or string)')
531

532
    def set_sources(self, source_density, fluxes=None, n_src=100):
533
534
535
536
537
        """
        Define source density or directly positions and optional weights (cosmic ray luminosity).

        :param source_density: source density (in 1 / Mpc^3) or array of shape (3, n_sources)
                               that point towards the center of the sources or keyword 'sbg'
538
        :param fluxes: corresponding cosmic ray fluxes of the sources of shape (n_sources).
539
        :param n_src: Number of point sources to be considered.
540
541
        :return: no return
        """
542
543
        self.universe.set_sources(source_density, fluxes, n_src)
        self.crs['sources'] = self.universe.sources
544

545
    def attenuate(self, library_path=PATH+'/simulation/crpropa3__emin_18.5__emax_21.npz', inside_fraction=None):
546
547
548
549
550
551
        """
        Apply attenuation for far away sources based on CRPropa simulations

        :param library_path: Input library file to use.
        """
        # Prepare the arrival and source matrix by reweighting
552
553
        self._prepare_arrival_matrix(library_path)
        # Assign source allocation of cosmic rays
554
        self._allocate_sources(inside_fraction)
555
        # Assign the arrival directions of the cosmic rays
556
        self._set_arrival_directions()
557
558
559
560
561
        # Assign charges and energies of the cosmic rays
        self._set_charges_energies()

    def set_xmax(self, z2a='double', model='EPOS-LHC'):
        """
562
        Calculate Xmax by Gumbel distribution for the simulated energies and charges.
563
564
565
566
        For more information refer to BaseSimulation.set_xmax().
        """
        return BaseSimulation.set_xmax(self, z2a, model)

567
568
569
570
571
572
573
574
    def smear_sources(self, delta, dynamic=None):
        """
        Smears the source positions with a fisher distribution of width delta (optional dynamic smearing).

        :param delta: either the constant width of the fisher distribution or dynamic (delta / R[10 EV]), in radians
        :param dynamic: if True, function applies dynamic smearing (delta / R[EV]) with value delta at 10 EV rigidity
        :return: no return
        """
575
576
        assert self.universe.sources is not None, "Cannot smear sources without positions!"
        assert self.crs is not None, "Cannot smear sources without calling attenuate() before!"
577

578
579
580
581
        mask_close = self.crs['source_labels'] >= 0
        d = delta if dynamic is None else delta / \
            10 ** (self.crs['log10e'] - np.log10(self.crs['charge']) - 19.)[mask_close]
        self.crs['vecs'][:, mask_close] = coord.rand_fisher_vec(self.crs['vecs'][:, mask_close], kappa=1/d**2)
582

583
    def get_data(self, shuffle=False):
584
        """
585
586
587
588
        Returns the data in the form of the cosmic_rays.CosmicRaysSets() container.

        :return: Instance of cosmic_rays.CosmicRaysSets() classes
        """
589
590
        if shuffle:
            self.shuffle_events()
591
        self.crs['signal_label'] = self.signal_label
592
593
594
595
596
        return self.crs

    def _prepare_arrival_matrix(self, library_path):
        """
        Prepare the arrival and source matrix by reweighting
597

598
599
        :param library_path: Input library file to use.
        """
600
601
602
        if (self.energy_setting is None) or (self.charge_weights is None):
            raise Exception("You have to define energies and charges before (set_energy() and set_charges())!")

603
        data = np.load(library_path, allow_pickle=True)
604
605
606
        self.dis_bins, self.log10e_bins = data['distances'], data['log10e_bins']
        if np.median(np.min(self.universe.distances, axis=1)) < np.min(self.dis_bins):
            print("Warning: Distance binning of simulation starts too far outside (%s Mpc)! " % np.min(self.dis_bins))
607
            print("Required would be: %.2fMpc." % np.median(np.min(self.universe.distances, axis=1)))
608
        # Assign distance indices for all simulated soures, shape: (self.nsets, self.universe.n_src)
609
        dis_bin_idx = self.universe.distance_indices(self.dis_bins)
610

611
        charge = {'h': 1, 'he': 2, 'n': 7, 'si': 14, 'fe': 26}
612
613
        self.source_matrix = np.zeros((self.nsets, self.universe.n_src, 5))
        self.arrival_matrix = np.zeros((self.dis_bins.size, 5, len(self.log10e_bins)-1))
614
615
        for key in self.charge_weights:
            f = self.charge_weights[key]
616
617
            if f == 0:
                continue
618
            fractions = data['fractions'].item()[key]
619
            # reweight to spectral index (simulated gamma=-1) and apply energy / rigidity cut
teresa.bister's avatar
teresa.bister committed
620
            fractions = self._reweight_spectrum(fractions, charge[key])
621
            # as log-space binning the width of the distance bin is increasing with distance
622
623
            self.source_matrix += f * np.sum(fractions, axis=-1)[dis_bin_idx]
            self.arrival_matrix += f * fractions
624
625
        # Account for optional luminosity weights of sources
        self.source_matrix *= coord.atleast_kd(self.universe.source_fluxes, k=self.source_matrix.ndim)
626

teresa.bister's avatar
teresa.bister committed
627
628
    def _reweight_spectrum(self, fractions, c):
        """ Internal function to reweight to desired energy spectrum and rigidity cut """
629
630
631
632
633
634
635
        assert fractions.ndim == 4, "Element arrival matrix fraction must have 4 dimensions!"
        bin_center = (self.log10e_bins[:-1] + self.log10e_bins[1:]) / 2.
        # reweight spectrum (simulated is gamma=-1 as resulting from equally binning in log space)
        fractions *= 10**((self.energy_setting['gamma'] + 1)*(coord.atleast_kd(bin_center, fractions.ndim) - 19))
        # Apply the rigidity cut
        log10_cut = self.energy_setting['log10_cut']
        if self.energy_setting['rig_cut']:
636
637
638
639
640
641
642
643
644
645
646
            # convert rigidity cut (log10_cut) to corresponding energy cut (log10e_cut)
            log10e_cut = log10_cut + np.log10(c)
            fcut = np.ones(bin_center.size)
            # cut function from combined fit paper: https://arxiv.org/pdf/1612.07155.pdf
            fcut[bin_center > log10e_cut] = np.exp(1 - 10**(bin_center[bin_center > log10e_cut] - log10e_cut))
            fractions *= coord.atleast_kd(fcut, fractions.ndim)
        else:
            if log10_cut <= self.energy_setting['log10e_min']:
                print("Warning: element with charge Z=%i is not injected with rigidity cut of" % c)
                print("%s and log10e > %s" % (self.energy_setting['log10_cut'], self.energy_setting['log10e_min']))
            fractions[bin_center > log10_cut, :, :, :] = 0
647
648
        fractions[:, :, :, bin_center < self.energy_setting['log10e_min']] = 0.  # energy threshold (measured)
        fractions = np.sum(fractions, axis=0)   # sum over injected energies to reweight the spectrum
649
650
        return fractions

651
652
    def _get_inside_fraction(self):
        """ Internal function to determine fraction of cosmic rays which come from inside rmax (signal cosmic rays) """
653
654
        distance_fractions = self.arrival_matrix * coord.atleast_kd(self.dis_bins, 3)
        inside_fraction = np.sum(distance_fractions[self.dis_bins <= self.universe.rmax]) / np.sum(distance_fractions)
655
        self.crs['inside_fraction'] = inside_fraction
656
        return inside_fraction
657

658
659
660
661
    def _allocate_sources(self, inside_fraction=None):
        """ Internal function to assign the source allocation of the signal (inside rmax) cosmic rays """
        if inside_fraction is None:
            inside_fraction = self._get_inside_fraction()
662
663
664
665
666
        # Sample for each set the number of CRs coming from inside and outside rmax
        nsplit = np.random.multinomial(self.ncrs, [inside_fraction, 1-inside_fraction], size=self.nsets).T
        # Assign the CRs from inside rmax to their separate sources (by index label)
        source_labels = -np.ones(self.shape).astype(int)
        n_max = np.max(nsplit[0])
667
        source_labels[:, :n_max] = sample_from_m_distributions(self.source_matrix.sum(axis=-1), n_max)
668
669
        nrange = np.tile(np.arange(self.ncrs), self.nsets).reshape(self.shape)
        mask_close = nrange < nsplit[0][:, np.newaxis]  # Create mask for CRs inside rmax
670
        source_labels[~mask_close] = -1  # correct the ones resulting by max(nsplit[0])
671
672
673
674
675
676

        # Checks if for all sets there is at least one explicit simulated source WITHOUT cosmic-ray contribution
        range_src = np.arange(self.universe.n_src)
        if not np.apply_along_axis(lambda x: np.in1d(range_src, x, invert=True).any(), axis=1, arr=source_labels).all():
            print("Warning: too few sources simulated! Set keyword 'n_src' higher.")

677
        self.crs['source_labels'] = source_labels
678
679
        occ = np.apply_along_axis(lambda x: np.bincount(x+1)[x+1], axis=1, arr=source_labels)
        self.signal_label = (occ >= 2) & (source_labels >= 0)
680
681
682
683
684
685
        return source_labels

    def _set_arrival_directions(self):
        """ Internal function to sample the arrival directions """
        # First, set random direction of all events
        vecs = coord.rand_vec(self.shape)
686
        # Set source directions of simulated sources
687
688
        mask = self.crs['source_labels'] >= 0   # mask all cosmic rays which originate within rmax
        vecs[:, mask] = self.universe.sources[:, np.argwhere(mask)[:, 0], self.crs['source_labels'][mask]]
689
690
        self.crs['vecs'] = vecs
        distances = np.zeros(self.shape)
691
        distances[mask] = self.universe.distances[np.where(mask)[0], self.crs['source_labels'][mask]]
692
        self.crs['distances'] = distances
693
        return vecs
694

695
    def _set_charges_energies(self):
696
        """ Internal function to assign charges and energies of all cosmic rays """
697
698
        log10e = np.zeros(self.shape)
        charge = np.zeros(self.shape)
699
        c = [1, 2, 7, 14, 26]
700
        d_dis, d_log10e = np.diff(np.log10(self.dis_bins))[0], np.diff(self.log10e_bins)[0]
701

702
        # Assign charges and energies of background cosmic rays
703
        arrival_mat_far = self.arrival_matrix * coord.atleast_kd(self.dis_bins, 3)
704
705
706
        mask_out = self.dis_bins >= self.universe.rmax
        arrival_mat_far[~mask_out] = 0
        arrival_mat_far /= arrival_mat_far.sum()
707
        mask_close = self.crs['source_labels'] >= 0
708
709
        arrival_idx = np.random.choice(arrival_mat_far.size, size=np.sum(~mask_close), p=arrival_mat_far.flatten())
        idx = np.unravel_index(arrival_idx, arrival_mat_far.shape)
710
        _d = 10**(np.log10(self.dis_bins)[idx[0]] + (np.random.random(idx[0].size) - 0.5) * d_dis)
711
        _c, _lge = np.array(c)[idx[1]], self.log10e_bins[:-1][idx[2]]
712
713
714
715
        perm = np.arange(np.sum(~mask_close))
        np.random.shuffle(perm)
        log10e[~mask_close] = _lge[perm]
        charge[~mask_close] = _c[perm]
Marcus Wirtz's avatar
Marcus Wirtz committed
716
        self.crs['distances'][~mask_close] = _d[perm]
717
718

        # Assign charges and energies of close-by cosmic rays
719
        dis_bin_idx = self.universe.distance_indices(self.dis_bins)
720
        for i in range(1 + np.max(dis_bin_idx)):
Marcus Wirtz's avatar
Marcus Wirtz committed
721
            # assign distance indices to CRs
722
            cr_idx = dis_bin_idx[np.where(mask_close)[0], self.crs['source_labels'][mask_close]]
723
724
725
            mask = cr_idx == i
            if np.sum(mask) == 0:
                continue
726
            w_mat = self.arrival_matrix[i] / self.arrival_matrix[i].sum()
727
            arrival_idx = np.random.choice(w_mat.size, size=np.sum(mask), p=w_mat.flatten())
728
            idx = np.unravel_index(arrival_idx, w_mat.shape)
729
730
731
732
            perm = np.arange(np.sum(mask))
            np.random.shuffle(perm)
            _mask = np.copy(mask_close)
            _mask[_mask] = mask
733
            charge[_mask], log10e[_mask] = np.array(c)[idx[0]][perm], self.log10e_bins[:-1][idx[1]][perm]
734
735
736
        log10e += d_log10e * np.random.random(self.shape)
        self.crs['log10e'] = log10e
        self.crs['charge'] = charge
737
        return log10e, charge
738

739
740
    def _get_charge_id(self):
        """ Return charge id of universe """
741
742
743
        chargegroups = ['h', 'he', 'n', 'si', 'fe']
        return ''.join(['__%s_%s' % (key, self.charge_weights[key]) for key in chargegroups
                        if key in self.charge_weights])
Marcus Wirtz's avatar
Marcus Wirtz committed
744

745
    def _select_representative_set(self):  # pragma: no cover
Marcus Wirtz's avatar
Marcus Wirtz committed
746
747
748
749
750
751
752
        """ Select a representative set in terms of anisotropies """
        from scipy.stats import mode
        src_labels = np.copy(self.crs['source_labels']).astype(float)
        src_labels[src_labels < 0] = np.nan
        _, counts = mode(src_labels, axis=1, nan_policy='omit')
        return np.argsort(np.squeeze(counts))[int(self.nsets/2.)]

753
754
755
756
757
758
759
760
761
762
763
764
765
766
    def _get_strongest_sources(self, idx=None, min_number_src=5):
        """ Select strongest sources """
        ns = np.array([np.sum(self.crs['source_labels'][idx] == k) for k in range(self.universe.n_src)])
        n_t = self.ncrs
        n_thres = ns >= n_t
        while (np.sum(n_thres) < min_number_src) & (n_t >= 3):
            n_t = int(0.8 * n_t)
            n_thres = ns >= n_t
        try:
            src_idx = np.squeeze(np.argwhere(n_thres))[np.argsort(ns[n_thres])]
        except IndexError:
            src_idx = []
        return src_idx, ns

767
    def plot_spectrum(self, opath=None):  # pragma: no cover
Marcus Wirtz's avatar
Marcus Wirtz committed
768
        """ Plot energy spectrum """
769
        import matplotlib.pyplot as plt
770
771
772
773
        from scipy.interpolate import interp1d
        from astrotools import statistics as st
        log10e = np.array(self.crs['log10e'])

774
        plt.figure(figsize=(6, 4))
775
776
777
778
779
780
        dspectrum = auger.SPECTRA_DICT[17]
        log10e_center = dspectrum['logE']
        # log10e_center = log10e_center   # [log10e_center + 0.05 >= self.energy_setting['log10e_min']]
        flux = (10 ** log10e_center) ** 3 * dspectrum['mean']
        flux_high = (10 ** log10e_center) ** 3 * dspectrum['stathi']
        flux_low = (10 ** log10e_center) ** 3 * dspectrum['statlo']
781
        plt.errorbar(log10e_center, flux, yerr=[flux_low, flux_high], color='red',
782
783
784
785
                     fmt='.', linewidth=1, markersize=8, capsize=0, label='Auger 2017')

        log10e_bins = np.arange(np.round(np.min(log10e), 1), np.max(log10e) + 0.1, 0.1)
        n, bins = np.histogram(log10e.flatten(), bins=log10e_bins)
786
787
788
789
790
791
        nall = np.apply_along_axis(lambda x: np.histogram(x, bins=log10e_bins)[0], 1, log10e)
        flux_sim = (10 ** st.mid(bins)) ** 2 * n / self.nsets
        flux_unc = (10 ** st.mid(bins)) ** 2 * nall
        norm = flux[np.argmin(np.abs(log10e_center - st.mid(bins)[0]))] / flux_sim[0]
        plt.errorbar(st.mid(bins), norm * flux_sim, marker='.', ls='none', color='k', xerr=0.05,
                     yerr=np.std(norm*flux_unc, axis=0), zorder=-1)
792
        # plot arriving composition (spline approximation)
793
794
        colors = ['firebrick', 'darkorange', 'forestgreen', 'deepskyblue', 'darkblue']
        e, c = ['h', 'he', 'n', 'si', 'fe'], [1, 2, 7, 14, 26]
795
        for i, ei in enumerate(e):
796
797
798
            mask = self.crs['charge'] == c[i]
            if np.sum(mask) == 0:
                continue
799
            _n, _bins = np.histogram(log10e[mask].flatten(), bins=log10e_bins)
800
801
            _flux_sim = (10 ** st.mid(_bins)) ** 2 * _n / self.nsets
            smooth_spline = interp1d(st.mid(_bins), norm * _flux_sim, kind='cubic', bounds_error=False)
802
803
            x = np.linspace(log10e_bins[0], log10e_bins[-1], 100)
            plt.plot(x, smooth_spline(x), color=colors[i], label=ei)
804
        plt.yscale('log')
805
        plt.xlim([self.energy_setting['log10e_min'] - 0.01, np.max(log10e) + 0.07])
806
807
        plt.ylim([1e36, 1e38])
        plt.legend(loc='lower left', fontsize=12)
Marcus Wirtz's avatar