simulations.py 45.1 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
73
74
75
76
77
78
79
80

    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')):
            raise Exception("Use function set_energy() and set_charges() before using function set_xmax.")
        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']

81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
    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))

        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)

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

100
101
102
103
    def shuffle_events(self):
        """
        Independently shuffle the cosmic rays of each set.
        """
104
        # This function can be simplified in the future using np.take_along_axis()
105
        shuffle_ids = np.random.permutation(np.prod(self.shape)).reshape(self.shape)
106
        shuffle_ids = np.argsort(shuffle_ids, axis=1)
107
108
        sets_ids = np.repeat(np.arange(self.nsets), self.ncrs).reshape(self.shape)
        for _key in self.crs.shape_array.dtype.names:
109
            self.crs[_key] = self.crs[_key][sets_ids, shuffle_ids]
110
111
        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)]
112
        self.signal_label = self.signal_label[sets_ids, shuffle_ids]
113
        self.crs['signal_label'] = self.signal_label
114

115
116

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

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

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

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

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

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

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

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

179
180
        return self.crs['log10e']

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

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

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

207
208
        return self.crs['charge']

209
210
211
212
213
214
215
    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)

216
    def set_sources(self, sources, fluxes=None):
MarCus's avatar
MarCus committed
217
        """
218
        Define source position and optional weights (cosmic ray luminosity).
219

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

245
    def set_rigidity_bins(self, lens_or_bins, cover_rigidity=True):
MarCus's avatar
MarCus committed
246
        """
247
        Defines the bin centers of the rigidities.
248
249

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

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

MarCus's avatar
MarCus committed
275
        return self.rig_bins
276

277
    def smear_sources(self, delta, dynamic=None):
MarCus's avatar
MarCus committed
278
        """
279
        Smears the source positions with a fisher distribution of width delta (optional dynamic smearing).
280

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

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

MarCus's avatar
MarCus committed
300
    def lensing_map(self, lens, cache=None):
MarCus's avatar
MarCus committed
301
        """
302
        Apply a galactic magnetic field to the extragalactic map.
303

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

MarCus's avatar
MarCus committed
311
        if self.rig_bins is None:
312
313
            self.set_rigidity_bins(lens)

MarCus's avatar
MarCus committed
314
        if self.cr_map is None:
315
            self._fix_point_source()
MarCus's avatar
MarCus committed
316

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

326
        self.lensed = True
MarCus's avatar
MarCus committed
327
        self.cr_map = arrival_map
328

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

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

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

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

383
        self.signal_label = np.repeat(signal_label[np.newaxis], self.nsets, axis=0)
384
        self.crs['pixel'] = pixel
385

386
387
388
    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)
389

390
    def get_data(self, convert_all=False, shuffle=False):
MarCus's avatar
MarCus committed
391
        """
392
        Returns the data in the form of the cosmic_rays.CosmicRaysSets() container.
393
394
395

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

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

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

442
443
444
445
446
447
448
449
450
    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

451

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

462
463
464
465
466
467
468
469
        :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
470
        self.charge_weights = None
471
        self.universe = SourceGeometry(nsets)
472

473
    def set_energy(self, log10e_min, gamma=-2, log10_cut=20., rig_cut=True):
474
475
476
477
        """
        Define energy spectrum and cut off energy of sources.

        :param log10e_min: Minimum threshold energy for observation
478
479
480
        :param gamma: Spectral index of energy spectrum at sources
        :param log10_cut: Maximum cut-off energy or rigidity for sources
        :param rig_cut: if True, lgo10_cut refers to a rigidity cut
481
        """
482
        self.energy_setting = {'log10e_min': log10e_min, 'gamma': gamma, 'log10_cut': log10_cut, 'rig_cut': rig_cut}
483
484
485

    def set_charges(self, charges):
        """
486
        Define fraction of charge groups in form of dictionary (e.g. {'h':0.5, 'fe':0.5}) at source.
487

488
        :param charges: dictionary hosting the fractions of injected elements ('h', 'he', ...)
489
        """
490
        fraction = np.sum([charges[key] for key in charges])
491
        assert np.abs(fraction - 1) < 1e-4, "Fractions of charges dictionary must be normalized!"
492
        self.charge_weights = charges
493

494
    def set_sources(self, source_density, fluxes=None, n_src=100):
495
496
497
498
499
        """
        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'
500
        :param fluxes: corresponding cosmic ray fluxes of the sources of shape (n_sources).
501
        :param n_src: Number of point sources to be considered.
502
503
        :return: no return
        """
504
505
        self.universe.set_sources(source_density, fluxes, n_src)
        self.crs['sources'] = self.universe.sources
506

507
    def attenuate(self, library_path=PATH+'/simulation/sim__emin_19.0__emax_20.5.npz'):
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
        """
        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
        inside_fraction = self._prepare_arrival_matrix(library_path)
        # Assign the arrival diretions of the cosmic rays
        self._set_arrival_directions(inside_fraction)
        # Assign charges and energies of the cosmic rays
        self._set_charges_energies()

    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)

527
528
529
530
531
532
533
534
    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
        """
535
536
        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!"
537

538
539
540
541
        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)
542

543
    def get_data(self, shuffle=False):
544
        """
545
546
547
548
        Returns the data in the form of the cosmic_rays.CosmicRaysSets() container.

        :return: Instance of cosmic_rays.CosmicRaysSets() classes
        """
549
550
        if shuffle:
            self.shuffle_events()
551
552
553
554
555
        return self.crs

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

557
558
        :param library_path: Input library file to use.
        """
559
560
561
        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())!")

562
        data = np.load(library_path, allow_pickle=True)
563
564
565
        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))
566
            print("Required would be: %.2fMpc." % np.median(np.min(self.universe.distances, axis=1)))
567
        # Assign distance indices for all simulated soures, shape: (self.nsets, self.universe.n_src)
568
        dis_bin_idx = self.universe.distance_indices(self.dis_bins)
569

570
        charge = {'h': 1, 'he': 2, 'n': 7, 'si': 14, 'fe': 26}
571
572
        self.source_matrix = np.zeros((self.nsets, self.universe.n_src, 4))
        self.arrival_matrix = np.zeros((self.dis_bins.size, 4, len(self.log10e_bins)-1))
573
574
        for key in self.charge_weights:
            f = self.charge_weights[key]
575
576
            if f == 0:
                continue
577
            fractions = data['fractions'].item()[key]
578
579
            # reweight to spectral index (simulated gamma=-1) and apply enrgy / rigidity cut
            fractions = self._reweight_spetrum(fractions, charge[key])
580
            # as log-space binning the width of the distance bin is increasing with distance
581
582
            self.source_matrix += f * np.sum(fractions, axis=-1)[dis_bin_idx]
            self.arrival_matrix += f * fractions
583

584
585
        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)
586
        return inside_fraction
587

588
589
590
591
592
593
594
595
596
    def _reweight_spetrum(self, fractions, c):
        """ Internal function to reweight to desired energy spetrum and rigidity cut """
        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']:
597
598
599
600
601
602
603
604
605
606
607
608
            # 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
        fractions = np.sum(fractions, axis=0)   # forget about injected energies
609
610
611
        fractions[:, :, bin_center < self.energy_setting['log10e_min']] = 0
        return fractions

612
    def _set_arrival_directions(self, inside_fraction):
613
614
615
616
        """ Internal function to sample the arrival directions """
        vecs = coord.rand_vec(self.shape)
        # 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
617
        self.source_matrix *= coord.atleast_kd(self.universe.source_fluxes, k=self.source_matrix.ndim)
618
619
620
        # 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])
621
        source_labels[:, :n_max] = sample_from_m_distributions(self.source_matrix.sum(axis=-1), n_max)
622
623
624
        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
        source_labels[~mask_close] = -1  # corret the ones resulting by max(nsplit[0])
625
        self.crs['source_labels'] = source_labels
626
627
        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)
628
        # Set source diretions of simulated sources
629
        vecs[:, mask_close] = self.universe.sources[:, np.argwhere(mask_close)[:, 0], source_labels[mask_close]]
630
631
        self.crs['vecs'] = vecs
        distances = np.zeros(self.shape)
632
        distances[mask_close] = self.universe.distances[np.where(mask_close)[0], source_labels[mask_close]]
633
634
        self.crs['distances'] = distances
        return mask_close
635

636
    def _set_charges_energies(self):
637
        """ Internal function to assign charges and energies of all cosmic rays """
638
639
        log10e = np.zeros(self.shape)
        charge = np.zeros(self.shape)
640
        c = [1, 2, 7, 26]
641
        d_dis, d_log10e = np.diff(np.log10(self.dis_bins))[0], np.diff(self.log10e_bins)[0]
642

643
        # Assign charges and energies of background cosmic rays
644
        arrival_mat_far = self.arrival_matrix * coord.atleast_kd(self.dis_bins, 3)
645
646
647
        mask_out = self.dis_bins >= self.universe.rmax
        arrival_mat_far[~mask_out] = 0
        arrival_mat_far /= arrival_mat_far.sum()
648
        mask_close = self.crs['source_labels'] >= 0
649
650
        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)
651
        _d = 10**(np.log10(self.dis_bins)[idx[0]] + (np.random.random(idx[0].size) - 0.5) * d_dis)
652
        _c, _lge = np.array(c)[idx[1]], self.log10e_bins[:-1][idx[2]]
653
654
655
656
        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
657
        self.crs['distances'][~mask_close] = _d[perm]
658
659

        # Assign charges and energies of close-by cosmic rays
660
        dis_bin_idx = self.universe.distance_indices(self.dis_bins)
661
        for i in range(1 + np.max(dis_bin_idx)):
Marcus Wirtz's avatar
Marcus Wirtz committed
662
            # assign distance indices to CRs
663
            cr_idx = dis_bin_idx[np.where(mask_close)[0], self.crs['source_labels'][mask_close]]
664
665
666
            mask = cr_idx == i
            if np.sum(mask) == 0:
                continue
667
            w_mat = self.arrival_matrix[i] / self.arrival_matrix[i].sum()
668
            arrival_idx = np.random.choice(w_mat.size, size=np.sum(mask), p=w_mat.flatten())
669
            idx = np.unravel_index(arrival_idx, w_mat.shape)
670
671
672
673
            perm = np.arange(np.sum(mask))
            np.random.shuffle(perm)
            _mask = np.copy(mask_close)
            _mask[_mask] = mask
674
            charge[_mask], log10e[_mask] = np.array(c)[idx[0]][perm], self.log10e_bins[:-1][idx[1]][perm]
675
676
677
678
        log10e += d_log10e * np.random.random(self.shape)
        self.crs['log10e'] = log10e
        self.crs['charge'] = charge

679
680
681
    def _get_charge_id(self):
        """ Return charge id of universe """
        return ''.join(['__%s_%s' % (key, self.charge_weights[key]) for key in self.charge_weights])
Marcus Wirtz's avatar
Marcus Wirtz committed
682

Marcus Wirtz's avatar
Marcus Wirtz committed
683
684
685
686
687
688
689
690
    def _select_representative_set(self):
        """ 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.)]

691
    def plot_spectrum(self):
Marcus Wirtz's avatar
Marcus Wirtz committed
692
        """ Plot energy spectrum """
693
        import matplotlib.pyplot as plt
694
695
696
697
        from scipy.interpolate import interp1d
        from astrotools import statistics as st
        log10e = np.array(self.crs['log10e'])

698
        plt.figure(figsize=(6, 4))
699
700
701
702
703
704
        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']
705
        plt.errorbar(log10e_center, flux, yerr=[flux_low, flux_high], color='red',
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
                     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)
        flux_sim = (10 ** st.mid(bins)) ** 2 * n
        norm = np.sum(np.array(10 ** log10e_center * dspectrum['mean'])[dspectrum['logE'] > np.min(bins)])
        plt.errorbar(st.mid(bins), flux_sim / np.sum(n) * norm, xerr=0.05, marker='.', ls='none', color='k')
        # plot arriving composition (spline approximation)
        e, c = ['h', 'he', 'n', 'si-fe'], [1, 2, 7, 26]
        for i, ei in enumerate(e):
            _log10e = log10e[self.crs['charge'] == c[i]]
            _n, _bins = np.histogram(_log10e.flatten(), bins=log10e_bins)
            flux_sim = (10 ** st.mid(bins)) ** 2 * _n
            _bins = np.linspace(np.min(log10e), np.max(log10e), 300)
            smooth_spline = interp1d(st.mid(bins), flux_sim / np.sum(n) * norm, kind='cubic', bounds_error=False)
            plt.plot(_bins, smooth_spline(_bins), color='C%i' % i, label=ei)
722
        plt.yscale('log')
723
724
725
726
727
728
        plt.xlim([self.energy_setting['log10e_min'] - 0.01, np.max(log10e) + 0.07])
        plt.ylim([5e35, 1e38])
        plt.legend(loc='lower left', fontsize=14)
        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)
729
730
        plt.savefig('/tmp/spectrum%s__emin_%s__ecut_%s.pdf' % (self._get_charge_id(), self.energy_setting['log10e_min'],
                                                               self.energy_setting['log10_cut']), bbox_inches='tight')
731
732
733
        plt.close()

    def plot_arrivals(self, idx=None):
Marcus Wirtz's avatar
Marcus Wirtz committed
734
        """ Plot arrival map """
735
736
        import matplotlib.pyplot as plt
        from astrotools import skymap
Marcus Wirtz's avatar
Marcus Wirtz committed
737
        idx = self._select_representative_set() if idx is None else idx
738
        ns = np.array([np.sum(self.crs['source_labels'][idx] == k) for k in range(self.universe.n_src)])
739
740
741
        n_more_3 = ns >= 3
        src_idx = np.squeeze(np.argwhere(n_more_3))[np.argsort(ns[n_more_3])]
        labels_p = np.copy(self.crs['source_labels'][idx])
742
        labels_p[~np.in1d(labels_p, src_idx) & (labels_p >= 0)] = 10*self.universe.n_src
743
744
        for j, idxj in enumerate(src_idx):
            labels_p[labels_p == idxj] = j
745
        cmap = plt.get_cmap('jet', len(src_idx))
746
747
        cmap.set_over('k')
        cmap.set_under('gray')
748
749
750
751
752
        if len(src_idx) > 0:
            skymap.eventmap(self.crs['vecs'][:, idx], c=labels_p, cmap=cmap, cblabel='Source ID',
                            cticks=np.arange(-1, len(src_idx), 1), vmin=-0.5, vmax=len(src_idx)-0.5)
        else:
            skymap.eventmap(self.crs['vecs'][:, idx], c='gray')
753
        lon_src, lat_src = coord.vec2ang(self.universe.sources[:, idx])
754
755
        plt.scatter(-lon_src, lat_src, c='k', marker='*', s=2*ns)
        ns = np.sort(ns)[::-1]
756
        plt.title('Strongest sources: (%i, %i, %i)' % (ns[0], ns[1], ns[2]), fontsize=15)
757
758
        plt.savefig('/tmp/arrival%s__emin_%s__ecut_%s.pdf' % (self._get_charge_id(), self.energy_setting['log10e_min'],
                                                              self.energy_setting['log10_cut']), bbox_inches='tight')
759
760
        plt.close()

Marcus Wirtz's avatar
Marcus Wirtz committed
761
762
    def plot_distance(self):
        """ Plot distance histogram """
763
        import matplotlib.pyplot as plt
764
        e, c = ['h', 'he', 'n', 'si-fe'], [1, 2, 7, 26]
Marcus Wirtz's avatar
Marcus Wirtz committed
765
        bins = np.linspace(0, np.max(self.crs['distances']), 50)
766
        distances = np.array(self.crs['distances'])
767
        plt.figure(figsize=(6, 4))
768
        plt.hist(distances.flatten(), bins=bins, histtype='step', color='gray')
Marcus Wirtz's avatar
Marcus Wirtz committed
769
        for i, ei in enumerate(e):
770
771
            _distances = distances[self.crs['charge'] == c[i]]
            plt.hist(_distances, bins=bins, histtype='step', color='C%i' % i, label=ei)
Marcus Wirtz's avatar
Marcus Wirtz committed
772
773
774
        plt.legend(loc=0)
        plt.xlabel('d / Mpc', fontsize=14)
        plt.ylabel('counts', fontsize=14)
775
776
        plt.savefig('/tmp/distance%s__emin_%s__ecut_%s.pdf' % (self._get_charge_id(), self.energy_setting['log10e_min'],
                                                               self.energy_setting['log10_cut']), bbox_inches='tight')
777
        plt.close()
778
779


780
class SourceGeometry:
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
    """
    Class to set up a 3d Universe out of isotropically distributed sources.
    """

    def __init__(self, nsets):
        self.nsets = nsets
        self.rmax = None
        self.n_src = None
        self.sources = None
        self.source_fluxes = None
        self.distances = None

    def set_sources(self, source_density, fluxes, n_src):
        """
        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'
799
        :param fluxes: corresponding cosmic ray fluxes of the sources of shape (n_sources).
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
        :param n_src: Number of point sources to be considered.
        :return: no return
        """
        self.n_src = n_src
        if isinstance(source_density, (int, float, np.int, np.float)):
            # maximum radius for one source per cosmic ray (isotropy condition)
            self.rmax = (3*n_src/4/np.pi/source_density)**(1/3.)
            self.sources = coord.rand_vec((self.nsets, n_src))
            # random radius in volume
            self.distances = self.rmax * (np.random.random((self.nsets, n_src)))**(1/3.)
            self.source_fluxes = 1 / self.distances**2
        elif isinstance(source_density, np.ndarray):
            source_density = np.reshape(source_density, (3, -1))