simulations.py 55.4 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
        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)]
114
        self.signal_label = self.signal_label[sets_ids, shuffle_ids]
115

116
117

class ObservedBound(BaseSimulation):
MarCus's avatar
MarCus committed
118
    """
MarCus's avatar
MarCus committed
119
120
121
    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
122
    """
Martin Urban's avatar
Martin Urban committed
123

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

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

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

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

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

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

161
162
163
164
        if isinstance(log10e_min, np.ndarray):
            if log10e_min.shape == self.shape:
                self.crs['log10e'] = log10e_min
            elif log10e_min.size == self.ncrs:
165
166
                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)
167
            else:
168
                raise Exception("Shape of input energies not in format (nsets, ncrs).")
169
        elif isinstance(log10e_min, (float, np.float, int, np.int)):
170
171
172
173
174
175
176
            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
177
178
179
        else:
            raise Exception("Input of emin could not be understood.")

180
181
        return self.crs['log10e']

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

186
        :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
187
        :type: charge: Union[np.ndarray, str, float]
188
        :return: charges
MarCus's avatar
MarCus committed
189
        """
190
191
        assert 'charge' not in self.crs.keys(), "Try to re-assign charges!"

192
193
        if isinstance(charge, np.ndarray):
            if charge.shape == self.shape:
194
                self.crs['charge'] = charge
195
            elif charge.size == self.ncrs:
196
                self.crs['charge'] = np.tile(charge, self.nsets).reshape(self.shape)
197
            else:
198
                raise Exception("Shape of input charges not in format (nsets, ncrs).")
199
        elif isinstance(charge, (float, np.float, int, np.int)):
marcus's avatar
marcus committed
200
            self.crs['charge'] = charge
201
        elif isinstance(charge, str):
MarCus's avatar
MarCus committed
202
203
            if not hasattr(self.crs, 'log10e'):
                raise Exception("Use function set_energy() before accessing a composition model.")
204
            self.crs['charge'] = getattr(CompositionModel(self.shape, self.crs['log10e']), charge.lower())(**kwargs)
205
206
207
        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])
208
209
210
        else:
            raise Exception("Input of charge could not be understood.")

211
212
        return self.crs['charge']

213
214
215
216
217
218
219
    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)

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

224
225
        :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']
226
        :param fluxes: corresponding cosmic ray fluxes of the sources of shape (n_sources).
227
        :return: no return
MarCus's avatar
MarCus committed
228
        """
229
        if isinstance(sources, np.ndarray):
230
231
232
233
            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
234
            self.sources = sources
Marcus Wirtz's avatar
Marcus Wirtz committed
235
236
237
            if fluxes is not None:
                assert fluxes.size == len(sources.T)
                self.source_fluxes = fluxes
238
        elif isinstance(sources, (int, np.int)):
MarCus's avatar
MarCus committed
239
            src_pix = np.random.randint(0, self.npix, sources)
240
            self.sources = np.array(hpt.pix2vec(self.nside, src_pix))
Marcus Wirtz's avatar
Marcus Wirtz committed
241
242
243
            if fluxes is not None:
                assert fluxes.size == sources
                self.source_fluxes = fluxes
244
        elif isinstance(sources, str):
Marcus Wirtz's avatar
Marcus Wirtz committed
245
            self.sources, self.source_fluxes = getattr(SourceScenario(), sources.lower())()[:2]
246
247
248
        else:
            raise Exception("Source scenario not understood.")

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

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

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

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

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

285
286
        :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
287
        :return: no return
MarCus's avatar
MarCus committed
288
        """
289
290
291
        if self.sources is None:
            raise Exception("Cannot smear sources without positions.")

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

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

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

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

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

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

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

333
    def apply_exposure(self, a0=-35.25, zmax=60):
MarCus's avatar
MarCus committed
334
        """
335
336
337
        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
338
        :type a0: float, int
339
        :param zmax: maximum zenith angle [deg] for the events
340
        :return: no return
MarCus's avatar
MarCus committed
341
        """
Marcus Wirtz's avatar
Marcus Wirtz committed
342
343
344
345
346
        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
347
        self.cr_map /= np.sum(self.cr_map, axis=-1)[:, np.newaxis]
348

349
    def arrival_setup(self, fsig):
MarCus's avatar
MarCus committed
350
        """
351
352
353
        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
354
        :type fsig: float
355
        :return: no return
MarCus's avatar
MarCus committed
356
        """
357
358
        dtype = np.uint16 if self.nside <= 64 else np.uint32
        pixel = np.zeros(self.shape).astype(dtype)
359
360
361

        # Setup the signal part
        n_sig = int(fsig * self.ncrs)
362
        signal_label = np.in1d(range(self.ncrs), np.arange(0, n_sig, 1))
363
364
        if n_sig == 0:
            pass
365
366
367
368
        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()
369
        elif np.sum(self.cr_map) > 0:
MarCus's avatar
MarCus committed
370
            if self.cr_map.size == self.npix:
371
                pixel[:, signal_label] = np.random.choice(self.npix, (self.nsets, n_sig), p=np.hstack(self.cr_map))
372
373
            else:
                for i, rig in enumerate(self.rig_bins):
374
                    mask_rig = (rig == self.rigidities) * signal_label  # type: np.ndarray
375
                    n = np.sum(mask_rig)
376
377
                    if n == 0:
                        continue
MarCus's avatar
MarCus committed
378
                    pixel[mask_rig] = np.random.choice(self.npix, n, p=self.cr_map[i])
379
380
        else:
            raise Exception("No signal probability to sample signal from!")
381

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

387
        self.signal_label = np.repeat(signal_label[np.newaxis], self.nsets, axis=0)
388
        self.crs['pixel'] = pixel
389

390
391
392
    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)
393

394
    def get_data(self, convert_all=False, shuffle=False):
MarCus's avatar
MarCus committed
395
        """
396
        Returns the data in the form of the cosmic_rays.CosmicRaysSets() container.
397
398
399

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

402
                 Example:
403
                 sim = ObservedBound()
404
                 ...
405
                 crs = sim.get_data(convert_all=True)
406
407
408
409
410
                 pixel = crs['pixel']
                 lon = crs['lon']
                 lat = crs['lat']
                 log10e = crs['log10e']
                 charge = crs['charge']
MarCus's avatar
MarCus committed
411
        """
412
        if convert_all is True:
413
            if not hasattr(self.crs, 'lon') or not hasattr(self.crs, 'lat'):
414
                self.convert_pixel(convert_all=True)
415
416
        if shuffle:
            self.shuffle_events()
417
        self.crs['signal_label'] = self.signal_label
418
        return self.crs
419

420
    def convert_pixel(self, keyword='vecs', convert_all=False):
421
422
        """
        Converts pixelized information stored under key 'pixel' to vectors (keyword='vecs')
423
        or angles (keyword='angles'), accessible via 'lon', 'lat'. When convert_all is True, both are saved.
424
425
426
427
428
429
        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']
430
            vecs = hpt.rand_exposure_vec_in_pix(self.nside, np.hstack(self.crs['pixel']), a0, zmax).reshape(shape)
431
432
        else:
            vecs = hpt.rand_vec_in_pix(self.nside, np.hstack(self.crs['pixel'])).reshape(shape)
433
434
435
436
437
438
439
440
441
442
443
444
445
        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!')
446

447
448
449
450
451
452
453
454
455
    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

456

457
class SourceBound(BaseSimulation):
458
    """
459
    Class to simulate cosmic ray arrival scenario by sources located at the sky, including energies, charges, smearing.
460
461
462
463
    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):
464
465
        """
        Initialization of SourceBound simulation.
466

467
468
469
470
471
472
473
474
        :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
475
        self.charge_weights = None
476
        self.universe = SourceGeometry(nsets)
477

478
    def set_energy(self, log10e_min, gamma=-2, log10_cut=20., rig_cut=True):
479
480
481
482
        """
        Define energy spectrum and cut off energy of sources.

        :param log10e_min: Minimum threshold energy for observation
483
484
        :param gamma: Spectral index of energy spectrum at sources
        :param log10_cut: Maximum cut-off energy or rigidity for sources
485
        :param rig_cut: if True, log10_cut refers to a rigidity cut
486
        """
487
        self.energy_setting = {'log10e_min': log10e_min, 'gamma': gamma, 'log10_cut': log10_cut, 'rig_cut': rig_cut}
488
489
490

    def set_charges(self, charges):
        """
491
        Define fraction of charge groups in form of dictionary (e.g. {'h':0.5, 'fe':0.5}) at source
492
        or as keyword 'first_minimum'/'second_minimum' from Auger's best fit paper (arXiv:1612.07155)
493
494
495
496
497
498
499
500
501
502
503
504
        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':
505
506
                self.energy_setting['gamma'], self.energy_setting['log10_cut'] = -0.87, 18.62
                self.charge_weights = {'n': 0.88, 'si': 0.12}
507
            elif charges == 'second_minimum':
508
509
510
511
512
513
514
515
                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}
516
517
518
519
520
            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)')
521

522
    def set_sources(self, source_density, fluxes=None, n_src=100):
523
524
525
526
527
        """
        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'
528
        :param fluxes: corresponding cosmic ray fluxes of the sources of shape (n_sources).
529
        :param n_src: Number of point sources to be considered.
530
531
        :return: no return
        """
532
533
        self.universe.set_sources(source_density, fluxes, n_src)
        self.crs['sources'] = self.universe.sources
534

535
    def attenuate(self, library_path=PATH+'/simulation/crpropa3__emin_18.5__emax_21.npz', inside_fraction=None):
536
537
538
539
540
541
        """
        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
542
543
        self._prepare_arrival_matrix(library_path)
        # Assign source allocation of cosmic rays
544
        self._allocate_sources(inside_fraction)
545
        # Assign the arrival directions of the cosmic rays
546
        self._set_arrival_directions()
547
548
549
550
551
        # Assign charges and energies of the cosmic rays
        self._set_charges_energies()

    def set_xmax(self, z2a='double', model='EPOS-LHC'):
        """
552
        Calculate Xmax by Gumbel distribution for the simulated energies and charges.
553
554
555
556
        For more information refer to BaseSimulation.set_xmax().
        """
        return BaseSimulation.set_xmax(self, z2a, model)

557
558
559
560
561
562
563
564
    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
        """
565
566
        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!"
567

568
569
570
571
        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)
572

573
    def get_data(self, shuffle=False):
574
        """
575
576
577
578
        Returns the data in the form of the cosmic_rays.CosmicRaysSets() container.

        :return: Instance of cosmic_rays.CosmicRaysSets() classes
        """
579
580
        if shuffle:
            self.shuffle_events()
581
        self.crs['signal_label'] = self.signal_label
582
583
584
585
586
        return self.crs

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

588
589
        :param library_path: Input library file to use.
        """
590
591
592
        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())!")

593
        data = np.load(library_path, allow_pickle=True)
594
595
596
        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))
597
            print("Required would be: %.2fMpc." % np.median(np.min(self.universe.distances, axis=1)))
598
        # Assign distance indices for all simulated soures, shape: (self.nsets, self.universe.n_src)
599
        dis_bin_idx = self.universe.distance_indices(self.dis_bins)
600

601
        charge = {'h': 1, 'he': 2, 'n': 7, 'si': 14, 'fe': 26}
602
603
        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))
604
605
        for key in self.charge_weights:
            f = self.charge_weights[key]
606
607
            if f == 0:
                continue
608
            fractions = data['fractions'].item()[key]
609
            # reweight to spectral index (simulated gamma=-1) and apply energy / rigidity cut
teresa.bister's avatar
teresa.bister committed
610
            fractions = self._reweight_spectrum(fractions, charge[key])
611
            # as log-space binning the width of the distance bin is increasing with distance
612
613
            self.source_matrix += f * np.sum(fractions, axis=-1)[dis_bin_idx]
            self.arrival_matrix += f * fractions
614
615
        # Account for optional luminosity weights of sources
        self.source_matrix *= coord.atleast_kd(self.universe.source_fluxes, k=self.source_matrix.ndim)
616

teresa.bister's avatar
teresa.bister committed
617
618
    def _reweight_spectrum(self, fractions, c):
        """ Internal function to reweight to desired energy spectrum and rigidity cut """
619
620
621
622
623
624
625
        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']:
626
627
628
629
630
631
632
633
634
635
636
            # 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
637
638
        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
639
640
        return fractions

641
642
    def _get_inside_fraction(self):
        """ Internal function to determine fraction of cosmic rays which come from inside rmax (signal cosmic rays) """
643
644
        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)
645
        self.crs['inside_fraction'] = inside_fraction
646
        return inside_fraction
647

648
649
650
651
    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()
652
653
654
655
656
        # 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])
657
        source_labels[:, :n_max] = sample_from_m_distributions(self.source_matrix.sum(axis=-1), n_max)
658
659
        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
660
        source_labels[~mask_close] = -1  # correct the ones resulting by max(nsplit[0])
661
        self.crs['source_labels'] = source_labels
662
663
        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)
664
665
666
667
668
669
        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)
670
        # Set source directions of simulated sources
671
672
        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]]
673
674
        self.crs['vecs'] = vecs
        distances = np.zeros(self.shape)
675
        distances[mask] = self.universe.distances[np.where(mask)[0], self.crs['source_labels'][mask]]
676
        self.crs['distances'] = distances
677
        return vecs
678

679
    def _set_charges_energies(self):
680
        """ Internal function to assign charges and energies of all cosmic rays """
681
682
        log10e = np.zeros(self.shape)
        charge = np.zeros(self.shape)
683
        c = [1, 2, 7, 14, 26]
684
        d_dis, d_log10e = np.diff(np.log10(self.dis_bins))[0], np.diff(self.log10e_bins)[0]
685

686
        # Assign charges and energies of background cosmic rays
687
        arrival_mat_far = self.arrival_matrix * coord.atleast_kd(self.dis_bins, 3)
688
689
690
        mask_out = self.dis_bins >= self.universe.rmax
        arrival_mat_far[~mask_out] = 0
        arrival_mat_far /= arrival_mat_far.sum()
691
        mask_close = self.crs['source_labels'] >= 0
692
693
        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)
694
        _d = 10**(np.log10(self.dis_bins)[idx[0]] + (np.random.random(idx[0].size) - 0.5) * d_dis)
695
        _c, _lge = np.array(c)[idx[1]], self.log10e_bins[:-1][idx[2]]
696
697
698
699
        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
700
        self.crs['distances'][~mask_close] = _d[perm]
701
702

        # Assign charges and energies of close-by cosmic rays
703
        dis_bin_idx = self.universe.distance_indices(self.dis_bins)
704
        for i in range(1 + np.max(dis_bin_idx)):
Marcus Wirtz's avatar
Marcus Wirtz committed
705
            # assign distance indices to CRs
706
            cr_idx = dis_bin_idx[np.where(mask_close)[0], self.crs['source_labels'][mask_close]]
707
708
709
            mask = cr_idx == i
            if np.sum(mask) == 0:
                continue
710
            w_mat = self.arrival_matrix[i] / self.arrival_matrix[i].sum()
711
            arrival_idx = np.random.choice(w_mat.size, size=np.sum(mask), p=w_mat.flatten())
712
            idx = np.unravel_index(arrival_idx, w_mat.shape)
713
714
715
716
            perm = np.arange(np.sum(mask))
            np.random.shuffle(perm)
            _mask = np.copy(mask_close)
            _mask[_mask] = mask
717
            charge[_mask], log10e[_mask] = np.array(c)[idx[0]][perm], self.log10e_bins[:-1][idx[1]][perm]
718
719
720
        log10e += d_log10e * np.random.random(self.shape)
        self.crs['log10e'] = log10e
        self.crs['charge'] = charge
721
        return log10e, charge
722

723
724
    def _get_charge_id(self):
        """ Return charge id of universe """
725
726
727
        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
728

729
    def _select_representative_set(self):  # pragma: no cover
Marcus Wirtz's avatar
Marcus Wirtz committed
730
731
732
733
734
735
736
        """ 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.)]

737
738
739
740
741
742
743
744
745
746
747
748
749
750
    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

751
    def plot_spectrum(self, opath=None):  # pragma: no cover
Marcus Wirtz's avatar
Marcus Wirtz committed
752
        """ Plot energy spectrum """
753
        import matplotlib.pyplot as plt
754
755
756
757
        from scipy.interpolate import interp1d
        from astrotools import statistics as st
        log10e = np.array(self.crs['log10e'])

758
        plt.figure(figsize=(6, 4))
759
760
761
762
763
764
        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']
765
        plt.errorbar(log10e_center, flux, yerr=[flux_low, flux_high], color='red',
766
767
768
769
                     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)
770
771
772
773
774
775
        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)
776
        # plot arriving composition (spline approximation)
777
778
        colors = ['firebrick', 'darkorange', 'forestgreen', 'deepskyblue', 'darkblue']
        e, c = ['h', 'he', 'n', 'si', 'fe'], [1, 2, 7, 14, 26]
779
        for i, ei in enumerate(e):
780
781
782
            mask = self.crs['charge'] == c[i]
            if np.sum(mask) == 0:
                continue
783
            _n, _bins = np.histogram(log10e[mask].flatten(), bins=log10e_bins)
784
785
            _flux_sim = (10 ** st.mid(_bins)) ** 2 * _n / self.nsets
            smooth_spline = interp1d(st.mid(_bins), norm * _flux_sim, kind='cubic', bounds_error=False)
786
787
            x = np.linspace(log10e_bins[0], log10e_bins[-1], 100)
            plt.plot(x, smooth_spline(x), color=colors[i], label=ei)
788
        plt.yscale('log')
789
        plt.xlim([self.energy_setting['log10e_min'] - 0.01, np.max(log10e) + 0.07])
790
791
        plt.ylim([1e36, 1e38])
        plt.legend(loc='lower left', fontsize=12)
792
793
794
        yl = r'$E^{3} \, J(E)$ [km$^{-2}$ yr$^{-1}$ sr$^{-1}$ eV$^{2}$]'
        plt.ylabel(yl, fontsize=16)
        plt.xlabel(r'$\log_{10}$($E$/eV)', fontsize=16)
795
796
797
798
        if opath is None:
            opath = '/tmp/spectrum%s__emin_%s__ecut_%s.pdf' % (self._get_charge_id(), self.energy_setting['log10e_min'],
                                                               self.energy_setting['log10_cut'])
        plt.savefig(opath, bbox_inches='tight')
799
800
        plt.close()

801
    def plot_arrivals(</