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

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

9

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

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


36
37
38
39
40
41
42
43
44
45
46
47
48
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)


49
50
51
52
53
54
55
56
57
58
59
60
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
61
        self.signal_label = None
62
63
64
65
66
67
68
69
70
71
72
73

    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')):
74
75
            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""")
76
77
78
79
80
81
82
        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']

83
84
85
86
87
88
89
90
91
92
93
    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))

94
95
96
97
98
        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)
99
100
101
102

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

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

118
119

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

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

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

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

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

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

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

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

182
183
        return self.crs['log10e']

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

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

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

213
214
        return self.crs['charge']

215
216
217
218
219
220
221
    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)

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

402
403
404
    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)
405

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

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

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

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

458
459
460
461
462
463
464
465
466
    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

467

468
class SourceBound(BaseSimulation):
469
    """
470
    Class to simulate cosmic ray arrival scenario by sources located at the sky, including energies, charges, smearing.
471
472
473
474
    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):
475
476
        """
        Initialization of SourceBound simulation.
477

478
479
480
481
482
483
484
485
        :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
486
        self.charge_weights = None
487
        self.universe = SourceGeometry(nsets)
488

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

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

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

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

        elif isinstance(charges, str):
515
            # (CTG, CTD) according to naming in Auger's combined fit paper (arXiv:1612.07155)
516
            if (charges == 'first_minimum_CTG') or (charges == 'first_minimum'):
517
518
                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}
519
520
521
522
523
524
            elif charges == 'second_minimum_CTG':
                self.energy_setting['gamma'], self.energy_setting['log10_cut'] = -0.87, 18.62
                self.charge_weights = {'n': 0.88, 'si': 0.12}
            elif charges == 'first_minimum_CTD':
                self.energy_setting['gamma'], self.energy_setting['log10_cut'] = 1.47, 18.15
                self.charge_weights = {'h': 0.4494, 'he': 0.52, 'n': 0.03, 'si': 0.0006}
525
526
527
528
529
530
            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}
531
532
533
534
535
            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)')
536

537
    def set_sources(self, source_density=None, sources=None, fluxes=None, n_src=None, rmax=None):
538
539
540
541
542
        """
        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'
543
        :param fluxes: corresponding cosmic ray fluxes of the sources of shape (n_sources).
544
        :param n_src: Number of point sources to be considered.
545
546
        :return: no return
        """
547
548
549
550
551
552
        if (n_src is not None) and (n_src > self.ncrs):
            print("Number of sources should not be larger than number of cosmic rays. Automatically set n_src = ncrs.")
            n_src = self.ncrs
        elif n_src is None:
            n_src = max(30, min(500, self.ncrs // 10))
        self.universe.set_sources(source_density, sources, fluxes, n_src, rmax)
553
        self.crs['sources'] = self.universe.sources
554
        self.crs['source_density'] = source_density
555

556
    def attenuate(self, library_path=None, inside_fraction=None, evolution=0):
557
558
559
560
        """
        Apply attenuation for far away sources based on CRPropa simulations

        :param library_path: Input library file to use.
561
562
        :param evolution: m for source evolution of the homogeneous background
                          (right now: not for foreground sources)
563
564
        """
        # Prepare the arrival and source matrix by reweighting
565
        self._prepare_arrival_matrix(library_path, evolution=evolution)
566
        # Assign source allocation of cosmic rays
567
        self._allocate_sources(inside_fraction)
568
        # Assign the arrival directions of the cosmic rays
569
        self._set_arrival_directions()
570
        # Assign charges and energies of the cosmic rays
571
        self._set_charges_energies(evolution=evolution)
572
573
574

    def set_xmax(self, z2a='double', model='EPOS-LHC'):
        """
575
        Calculate Xmax by Gumbel distribution for the simulated energies and charges.
576
577
578
579
        For more information refer to BaseSimulation.set_xmax().
        """
        return BaseSimulation.set_xmax(self, z2a, model)

580
581
582
583
584
585
586
587
    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
        """
588
589
        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!"
590

591
592
593
594
        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)
595

596
    def get_data(self, shuffle=False):
597
        """
598
599
600
601
        Returns the data in the form of the cosmic_rays.CosmicRaysSets() container.

        :return: Instance of cosmic_rays.CosmicRaysSets() classes
        """
602
603
        if shuffle:
            self.shuffle_events()
604
        self.crs['signal_label'] = self.signal_label
605
606
        return self.crs

607
    def _prepare_arrival_matrix(self, library_path, evolution):
608
609
        """
        Prepare the arrival and source matrix by reweighting
610

611
612
        :param library_path: Input library file to use.
        """
613
614
615
        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())!")

616
617
        if library_path is None:
            library_path = PATH + '/simulation/crpropa3__emin_18.5__emax_21.0__IRB_Gilmore12.npz'
618
        data = np.load(library_path, allow_pickle=True)
619
        self.dis_bins, self.log10e_bins = data['distances'], data['log10e_bins']
620
621
622
623
624
625
        try:
            self.redshift_bins = data['redshift']
        except KeyError:  # only needed if simulation with not flat evolution is wanted
            if evolution != 0:
                import crpropa as crp
                self.redshift_bins = np.array([crp.comovingDistance2Redshift(dis * crp.Mpc) for dis in self.dis_bins])
626
627
628
629
630
631
632
633
        if self.universe.has_sources():
            # Assign distance indices for all simulated soures, shape: (self.nsets, self.universe.n_src)
            dis_bin_idx = self.universe.distance_indices(self.dis_bins)
            min_dis = np.min(self.dis_bins)
            if np.median(np.min(self.universe.distances, axis=1)) < min_dis:
                print("Warning: Distance binning of simulation starts too far outside (%s Mpc)! " % min_dis)
                print("\tRequired would be: %.2fMpc." % np.median(np.min(self.universe.distances, axis=1)))
                print("\tThis is only relevant if there is substantial attenuation below %s Mpc." % min_dis)
634

635
        charge = {'h': 1, 'he': 2, 'n': 7, 'si': 14, 'fe': 26}
636
637
        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))
638
639
        for key in self.charge_weights:
            f = self.charge_weights[key]
640
641
            if f == 0:
                continue
642
            fractions = data['fractions'].item()[key]   # dimensions: (energy_in, distances, charges_out, energy_out)
643
            # reweight to spectral index (simulated gamma=-1) and apply energy / rigidity cut
644
            fractions = self._reweight_spectrum(fractions, charge[key])  # dim: (distances, charges_out, energy_out)
645
            self.arrival_matrix += f * fractions
646
647
648
            if self.universe.has_sources():
                self.source_matrix += f * np.sum(fractions, axis=-1)[dis_bin_idx]

649
        # Account for optional luminosity weights of sources
650
651
        if self.universe.has_sources():
            self.source_matrix *= coord.atleast_kd(self.universe.source_fluxes, k=self.source_matrix.ndim)
652

teresa.bister's avatar
teresa.bister committed
653
654
    def _reweight_spectrum(self, fractions, c):
        """ Internal function to reweight to desired energy spectrum and rigidity cut """
655
656
657
        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)
658
        fractions *= 10**((self.energy_setting['gamma'] + 1)*(coord.atleast_kd(bin_center, fractions.ndim) - 19))
659
660
661
        # Apply the rigidity cut
        log10_cut = self.energy_setting['log10_cut']
        if self.energy_setting['rig_cut']:
662
663
664
665
666
667
668
669
670
671
672
            # 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
673
674
        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
675
676
        return fractions

677
678
    def _get_inside_fraction(self):
        """ Internal function to determine fraction of cosmic rays which come from inside rmax (signal cosmic rays) """
679
        # as log-space binning the width of the distance bin is increasing with distance (d log(x) = 1/x * dx)
680
681
682
        dist_frac = self.arrival_matrix * coord.atleast_kd(self.dis_bins, 3)
        rmax = np.atleast_1d(self.universe.rmax)
        inside_fraction = np.array([np.sum(dist_frac[self.dis_bins <= r]) for r in rmax]) / np.sum(dist_frac)
683
        return inside_fraction
684

685
686
687
688
    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()
689
        self.crs['inside_fraction'] = inside_fraction
690
691
692
693
694
        # Sample for each set the number of CRs coming from inside the horizon rmax
        std = np.sqrt(self.ncrs * inside_fraction * (1 - inside_fraction))      # approximate fluctutation from
        n_fluc = np.random.normal(scale=std, size=np.size(inside_fraction))     # multinomial distribution
        n_inside = np.rint(np.clip(self.ncrs * inside_fraction + n_fluc, 0, self.ncrs)).astype(int)
        # Assign the CRs from inside rmax to their separate sources by index label (-1 for outside rmax / no source)
695
        source_labels = -np.ones(self.shape).astype(int)
696
        n_max = np.max(n_inside)  # max events from inside over all sets
697
        source_labels[:, :n_max] = sample_from_m_distributions(self.source_matrix.sum(axis=-1), n_max)
698
        nrange = np.tile(np.arange(self.ncrs), self.nsets).reshape(self.shape)
699
700
        mask_close = nrange < n_inside[:, np.newaxis]  # Create mask for CRs inside rmax
        source_labels[~mask_close] = -1  # correct the ones resulting by max(n_inside)
701

702
        # Checks if for ALL sets there is at least one explicitly simulated source WITHOUT cosmic-ray contribution
703
        range_src = np.arange(self.universe.n_src)
704
705
706
        check = np.apply_along_axis(lambda x: np.in1d(range_src, x, invert=True).any(), axis=1, arr=source_labels).all()
        if (not check):
            print("Warning: not enough sources. Set keyword 'n_src' (currently %s) higher?" % self.universe.n_src)
707

708
        self.crs['source_labels'] = source_labels
709
710
        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)
711
712
713
714
715
716
        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)
717
        distances = np.zeros(self.shape)
718
719
720
721
722
723
        if self.universe.has_sources():
            # Set source directions of simulated sources
            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]]
            distances[mask] = self.universe.distances[np.where(mask)[0], self.crs['source_labels'][mask]]
        self.crs['vecs'] = vecs
724
        self.crs['distances'] = distances
725
        return vecs
726

727
    def _set_charges_energies(self, evolution=0):
728
        """ Internal function to assign charges and energies of all cosmic rays """
729
730
        log10e = np.zeros(self.shape)
        charge = np.zeros(self.shape)
731
        c = [1, 2, 7, 14, 26]
732
        d_dis, d_log10e = np.diff(np.log10(self.dis_bins))[0], np.diff(self.log10e_bins)[0]
733

734
        # Assign distances, charges and energies of background cosmic rays
735
        mask_close = self.crs['source_labels'] >= 0
736
        if np.sum(~mask_close) > 0:
737
            # as log-space binning the width of the distance bin is increasing with distance (d log(x) = 1/x * dx)
738
739
740
741
742
            if evolution != 0:
                dis_bins_e = self.dis_bins*(1 + self.redshift_bins)**evolution
                arrival_mat_far = self.arrival_matrix * coord.atleast_kd(dis_bins_e, 3)
            else:
                arrival_mat_far = self.arrival_matrix * coord.atleast_kd(self.dis_bins, 3)
743
            # determine for each set the distance bin of the horizon rmax
744
            dis_idx = np.argmin(np.abs(np.atleast_1d(self.universe.rmax)[None] - self.dis_bins[:, None]), axis=0)
745
746
            # loop over all occuring distance bins and fill events of the respective sets
            for _dis in np.sort(np.unique(dis_idx)):
747
                arrival_mat_far[0:_dis] = 0   # set probabilities inside rmax to zero
748
749
750
751
752
753
754
755
756
757
758
759
760
761
                mask_pick = ~mask_close & (dis_idx == _dis)[:, np.newaxis]  # select background events of respective set
                n_pick = np.sum(mask_pick)
                arrival_mat_far /= arrival_mat_far.sum()    # normalize arrival matrix to then draw random samples
                arrival_idx = np.random.choice(arrival_mat_far.size, size=n_pick, p=arrival_mat_far.flatten())
                # get the indices in terms of the original shape (distances, charges, energies)
                idx = np.unravel_index(arrival_idx, arrival_mat_far.shape)
                _d = 10**(np.log10(self.dis_bins)[idx[0]] + (np.random.random(idx[0].size) - 0.5) * d_dis)
                _c, _lge = np.array(c)[idx[1]], self.log10e_bins[:-1][idx[2]]
                # fill in the distances, charges, and energies with a random permutation
                perm = np.arange(n_pick)
                np.random.shuffle(perm)
                log10e[mask_pick] = _lge[perm]
                charge[mask_pick] = _c[perm]
                self.crs['distances'][mask_pick] = _d[perm]
762
763

        # Assign charges and energies of close-by cosmic rays
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
        if np.sum(mask_close) > 0:
            dis_bin_idx = self.universe.distance_indices(self.dis_bins)
            for i in range(1 + np.max(dis_bin_idx)):
                # assign distance indices to CRs
                cr_idx = dis_bin_idx[np.where(mask_close)[0], self.crs['source_labels'][mask_close]]
                mask = cr_idx == i
                if np.sum(mask) == 0:
                    continue
                # procedure in analogy to code block above
                w_mat = self.arrival_matrix[i] / self.arrival_matrix[i].sum()
                arrival_idx = np.random.choice(w_mat.size, size=np.sum(mask), p=w_mat.flatten())
                idx = np.unravel_index(arrival_idx, w_mat.shape)
                perm = np.arange(np.sum(mask))
                np.random.shuffle(perm)
                _mask = np.copy(mask_close)
                _mask[_mask] = mask
                charge[_mask], log10e[_mask] = np.array(c)[idx[0]][perm], self.log10e_bins[:-1][idx[1]][perm]
            log10e += d_log10e * np.random.random(self.shape)

783
784
        self.crs['log10e'] = log10e
        self.crs['charge'] = charge
785
        return log10e, charge
786

787
788
    def _get_charge_id(self):
        """ Return charge id of universe """
789
790
791
        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
792

793
794
795
796
    def _select_representative_set(self, mask=None):  # pragma: no cover
        """ Select a representative set in terms of anisotropies. Returns index """
        labels = self.crs['source_labels'] if mask is None else self.crs['source_labels'][mask]
        src_labels = np.copy(labels).astype(float)
Marcus Wirtz's avatar
Marcus Wirtz committed
797
798
        src_labels[src_labels < 0] = np.nan
        _, counts = mode(src_labels, axis=1, nan_policy='omit')
799
800
        idx_select = np.argsort(np.squeeze(counts))[int(len(labels)/2.)]
        return idx_select if mask is None else np.arange(self.nsets)[mask][idx_select]
Marcus Wirtz's avatar
Marcus Wirtz committed
801

802
803
804
805
806
807
808
809
810
811
812
813
814
815
    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

816
    def plot_spectrum(self, opath=None):  # pragma: no cover
Marcus Wirtz's avatar
Marcus Wirtz committed
817
        """ Plot energy spectrum """
818
        import matplotlib.pyplot as plt
819
820
821
822
        from scipy.interpolate import interp1d
        from astrotools import statistics as st
        log10e = np.array(self.crs['log10e'])

823
        plt.figure(figsize=(6, 4))
824
825
826
827
828
829
        dspectrum = auger.SPECTRA_DICT[17]
        log10e_center = dspectrum['logE'