cosmic_rays.py 25.1 KB
Newer Older
1
# -*- coding: utf-8 -*-
2
"""
Marcus Wirtz's avatar
Marcus Wirtz committed
3
4
Contains the cosmic rays base class which allows to store arbitrary properties for the cosmic rays and
makes them accesseble via key or getter function.
5
6
The second class describes sets of cosmic rays as needed for larger studies.
"""
7
import matplotlib.pyplot as plt
Martin Urban's avatar
Martin Urban committed
8
import numpy as np
9

10
from astrotools import container, coord, healpytools as hpt, obs, skymap
Martin Urban's avatar
Martin Urban committed
11

12
DTYPE_TEMPLATE = []
Marcus Wirtz's avatar
Marcus Wirtz committed
13
PHYS_ENERGIES = ['e', 'log10e', 'energy', 'E']
14
PHYS_DIRECTIONS = ['vecs', 'lon', 'lat', 'pixel', 'pix']
15
SHAPE_FLEXIBLE = ['charge'] + PHYS_ENERGIES
16

Martin Urban's avatar
Martin Urban committed
17

18
def plot_eventmap(crs, opath=None, **kwargs):  # pragma: no cover
19
20
21
    """
    Function to plot a scatter skymap of the cosmic rays

Marcus Wirtz's avatar
Marcus Wirtz committed
22
    :param crs: CosmicRaysBase object
Marcus Wirtz's avatar
Marcus Wirtz committed
23
    :param opath: Output path for the image, default is None
Marcus Wirtz's avatar
Marcus Wirtz committed
24
    :type opath: str
25
26
    :param kwargs:

27
28
           - c: quantity that is supposed to occur in colorbar, e.g. energy of the cosmic rays
           - cblabel: label for the colorbar
29
           - nside: Healpy resolution of the 'pixel' array in the cosmic ray class
Marcus Wirtz's avatar
Marcus Wirtz committed
30
           - additional named keywords passed to matplotlib.scatter
31
           - fontsize: Scales the fontsize in the image
32
    :return: figure, axis of the scatter plot
33
    """
34
    vecs = crs['vecs']
Marcus Wirtz's avatar
Marcus Wirtz committed
35
    c = kwargs.pop('c', crs['log10e'] if len(set(PHYS_ENERGIES) & set(crs.keys())) > 0 else None)
36
    return skymap.scatter(vecs, c=c, opath=opath, **kwargs)
37
38


39
def plot_heatmap(crs, opath=None, **kwargs):  # pragma: no cover
40
41
42
    """
    Function to plot a scatter skymap of the cosmic rays

Marcus Wirtz's avatar
Marcus Wirtz committed
43
    :param crs: CosmicRaysBase object
44
    :param opath: Output path for the image, default is None
Marcus Wirtz's avatar
Marcus Wirtz committed
45
    :type opath: str
46
47
48
    :param kwargs:

           - nside: Healpy resolution of the 'pixel' array in the cosmic ray class
Marcus Wirtz's avatar
Marcus Wirtz committed
49
           - additional named keywords passed to matplotlib.pcolormesh
50
    :return: figure of the heatmap, colorbar
51
    """
52
    nside = crs['nside'] if 'nside' in crs.keys() else kwargs.pop('nside', 64)
53
    pixel = crs['pixel']
54
55

    count = np.bincount(pixel, minlength=hpt.nside2npix(nside))
56
    return skymap.heatmap(count, opath=opath, **kwargs)
57
58


59
def plot_energy_spectrum(crs, xlabel='log$_{10}$(Energy / eV)', ylabel='entries', fontsize=16, bw=0.05,
Marcus Wirtz's avatar
Marcus Wirtz committed
60
                         opath=None, **kwargs):  # pragma: no cover
61
62
    """
    Function to plot the energy spectrum of the cosmic ray set
63

Marcus Wirtz's avatar
Marcus Wirtz committed
64
    :param crs: CosmicRaysBase object
65
    :param xlabel: label for the x-axis
Marcus Wirtz's avatar
Marcus Wirtz committed
66
    :type xlabel: str
67
    :param ylabel: label for the y-axis
Marcus Wirtz's avatar
Marcus Wirtz committed
68
    :type ylabel: str
69
    :param fontsize: Scales the fontsize in the image.
Marcus Wirtz's avatar
Marcus Wirtz committed
70
    :type fontsize: int
71
    :param bw: bin width for the histogram
Marcus Wirtz's avatar
Marcus Wirtz committed
72
    :type bw: float
Marcus Wirtz's avatar
Marcus Wirtz committed
73
    :param opath: Output path for the image, default is None
Marcus Wirtz's avatar
Marcus Wirtz committed
74
75
    :type opath: str
    :param kwargs: additional named keywords passed to matplotlib.pyplot.hist
76
    :return: bins of the histogram
77
78
    """
    log10e = crs['log10e']
79
80
81
82
83
84
85
86
87
88
    if 'bins' not in kwargs:
        bins = np.arange(17., 20.6, bw)
        bins = bins[(bins >= np.min(log10e) - 0.1) & (bins <= np.max(log10e) + 0.1)]
        kwargs['bins'] = bins

    kwargs.setdefault('color', 'k')
    kwargs.setdefault('fill', None)
    kwargs.setdefault('histtype', 'step')
    plt.hist(log10e, **kwargs)

89
90
91
92
93
94
95
    plt.xticks(fontsize=fontsize - 4)
    plt.yticks(fontsize=fontsize - 4)
    plt.xlabel(xlabel, fontsize=fontsize)
    plt.ylabel(ylabel, fontsize=fontsize)
    if opath is not None:
        plt.savefig(opath, bbox_inches='tight')
        plt.clf()
96
    return bins
97
98


99
# TODO: Do not allow names with leading underscore (if before self.__dict__.update)
100
class CosmicRaysBase(container.DataContainer):
101
    """ Cosmic rays base class meant for inheritance """
Marcus Wirtz's avatar
Marcus Wirtz committed
102

103
    def __init__(self, initializer=None, coord_system='gal'):
104
        # Inherits all functionalities from container.DataContainer object
105
        super(CosmicRaysBase, self).__init__(initializer)
106
        self.type = "CosmicRays"
107
        self.coord_system = coord_system
Martin Urban's avatar
Martin Urban committed
108
109

    def __getitem__(self, key):
110
111
112
113
114
115
116
117

        if isinstance(key, (int, np.integer, np.ndarray, slice)):
            crs = CosmicRaysBase(self.shape_array[key])
            for k in self.general_object_store.keys():
                to_copy = self.get(k)
                if isinstance(to_copy, (np.ndarray, list)):
                    if len(to_copy) == self.ncrs:
                        to_copy = to_copy[key]
118
119
                if (k == 'vecs'):
                    to_copy = to_copy[:, key]
120
121
122
                crs.__setitem__(k, to_copy)
            return crs
        if key in self.general_object_store.keys():
123
            if key in SHAPE_FLEXIBLE:
124
                return self.general_object_store[key] * np.ones(self.ncrs, dtype='int')
125
126
127
128
129
130
131
132
            return self.general_object_store[key]
        if key in self.shape_array.dtype.names:
            return self.shape_array[key]

        if len(self._similar_key(key)) > 0:
            return self._get_values_similar_key(self._similar_key(key).pop(), key)

        raise ValueError("Key '%s' does not exist, no info stored under similar keys was found" % key)
Martin Urban's avatar
Martin Urban committed
133
134

    def __setitem__(self, key, value):
135
136
        if key in self.shape_array.dtype.names:
            self.shape_array[key] = value
137
            if len(self._similar_key(key)) > 1:
138
139
                print("Warning: Your cosmic rays object contains data stored under a key similar to %s. "
                      "Changing one without the other may lead to problems." % key)
Martin Urban's avatar
Martin Urban committed
140
            return
141
        super(CosmicRaysBase, self).__setitem__(key, value)
Martin Urban's avatar
Martin Urban committed
142

Niklas Uwe Langner's avatar
Niklas Uwe Langner committed
143
    def __add__(self, other):
144
        return CosmicRaysBase([self, other])
Niklas Uwe Langner's avatar
Niklas Uwe Langner committed
145

146
147
    def _similar_key(self, key):
        """
148
        Helper function to check for keys describing the same physical data eg. 'vecs' and 'pixel'.
149
150
        """
        key_list = self.keys()
151
        common_keys = []
Marcus Wirtz's avatar
Marcus Wirtz committed
152
153
        if key in PHYS_DIRECTIONS:
            common_keys = set(PHYS_DIRECTIONS) & set(key_list)
154
            if key not in ['lon', 'lat']:
Marcus Wirtz's avatar
Marcus Wirtz committed
155
                if ('lon' in common_keys) and ('lat' not in common_keys):
156
                    common_keys.discard('lon')
Marcus Wirtz's avatar
Marcus Wirtz committed
157
                if ('lat' in common_keys) and ('lon' not in common_keys):
158
                    common_keys.discard('lat')
159
                common_keys = sorted(common_keys, key=PHYS_DIRECTIONS.index, reverse=True)
Marcus Wirtz's avatar
Marcus Wirtz committed
160
161
        elif key in PHYS_ENERGIES:
            common_keys = set(PHYS_ENERGIES) & set(key_list)
162
        return common_keys
163
164
165
166
167
168

    def _get_values_similar_key(self, similar_key, orig_key):
        """
        Helper function to get values stored under a different physical key in the correctly
        transformed way, together with _similar_key()
        """
169
        store = self.shape_array if similar_key in list(self.shape_array.dtype.names) else self.general_object_store
170
171
172
        if orig_key in ['e', 'energy', 'E']:
            if similar_key in ['e', 'energy', 'E']:
                return store[similar_key]
173
            return 10**np.array(store[similar_key])
174
        if orig_key == 'log10e':
175
            return np.log10(store[similar_key])
176
177
178
179
180
181
182
183
        return self._direction_transformation(similar_key, orig_key)

    def _direction_transformation(self, similar_key, orig_key):
        """
        Helper function to get values stored under a different physical key in the correctly
        transformed way specifically only for directions
        """
        nside = self.general_object_store['nside'] if 'nside' in self.keys() else 64
184
        store = self.shape_array if similar_key in list(self.shape_array.dtype.names) else self.general_object_store
185
        if orig_key == 'vecs':
186
            if ('lon' in similar_key) or ('lat' in similar_key):
187
188
                return hpt.ang2vec(store['lon'], store['lat'])
            return hpt.pix2vec(nside, store[similar_key])
Marcus Wirtz's avatar
Marcus Wirtz committed
189
        if ('pix' in orig_key):
190
191
            if 'pix' in similar_key:
                return store[similar_key]
192
193
            if similar_key == 'vecs':
                return hpt.vec2pix(nside, store['vecs'])
Marcus Wirtz's avatar
Marcus Wirtz committed
194
            return hpt.ang2pix(nside, store['lon'], store['lat'])
195
196
197
198
199
200
        if similar_key == 'vecs':
            lon, lat = hpt.vec2ang(store['vecs'])
        else:
            lon, lat = hpt.pix2ang(nside, store[similar_key])
        return lon if orig_key == 'lon' else lat

Martin Urban's avatar
Martin Urban committed
201
202
203
204
    def add_cosmic_rays(self, crs):
        """
        Function to add cosmic rays to the already existing set of cosmic rays

Marcus Wirtz's avatar
Marcus Wirtz committed
205
206
207
        :param crs: numpy array with cosmic rays. The cosmic rays must not contain
                    all original keys. Missing keys are set to zero. If additional
                    keys are provided, they are ignored.
Martin Urban's avatar
Martin Urban committed
208
        """
209
210
211
        if not isinstance(crs, CosmicRaysBase):
            raise Exception("You need to add a CosmicRaysBase object!")
        self.add_shape_array(crs.get_array())
212

213
214
215
216
217
218
    def sensitivity_2pt(self, niso=1000, bins=180, **kwargs):
        """
        Function to calculate the sensitivity by the 2pt-auto-correlation over a scrambling
        of the right ascension coordinates.

        :param niso: Number of isotropic sets to calculate.
219
        :param bins: Number of angular bins, 180 correspond to 1 degree binning (np.linspace(0, np.pi, bins+1).
220
        :param kwargs: additional named arguments passed to obs.two_pt_auto()
221
        :return: pvalues in the shape (bins)
222
223
224
225
226
227
228
229
230
231
232
233
234
        """
        kwargs.setdefault('cumulative', True)
        vec_crs = self.get('vecs')
        _, dec = coord.vec2ang(coord.gal2eq(vec_crs))

        # calculate auto correlation for isotropic scrambled data
        _ac_iso = np.zeros((niso, bins))
        for i in range(niso):
            _vecs = coord.ang2vec(coord.rand_phi(self.ncrs), dec)
            _ac_iso[i] = obs.two_pt_auto(_vecs, bins, **kwargs)

        # calculate p-value by comparing the true sets with the isotropic ones
        _ac_crs = obs.two_pt_auto(vec_crs, bins, **kwargs)
235
        pvals = np.sum(_ac_iso >= _ac_crs[np.newaxis], axis=0) / float(niso)
236
237
        return pvals

238
    def plot_eventmap(self, **kwargs):  # pragma: no cover
239
240
        """
        Function to plot a scatter skymap of the cosmic rays
241

Marcus Wirtz's avatar
Marcus Wirtz committed
242
        :param kwargs: additional named arguments passed to plot_eventmap().
243
        """
244
        return plot_eventmap(self, **kwargs)
245

246
    def plot_heatmap(self, **kwargs):  # pragma: no cover
247
248
249
        """
        Function to plot a healpy skymap of the cosmic rays

Marcus Wirtz's avatar
Marcus Wirtz committed
250
        :param kwargs: additional named arguments passed to plot_healpy_map().
251
        """
252
253
254
255
256
        return plot_heatmap(self, **kwargs)

    def plot_healpy_map(self, **kwargs):  # pragma: no cover
        """ Forwards to function plot_heatmap() """
        return self.plot_heatmap(**kwargs)
257

258
    def plot_energy_spectrum(self, **kwargs):  # pragma: no cover
259
        """
Marcus Wirtz's avatar
Design    
Marcus Wirtz committed
260
        Function to plot the energy spectrum of the cosmic rays
261

Marcus Wirtz's avatar
Marcus Wirtz committed
262
        :param kwargs: additional named arguments passed to plot_energy_spectum().
263
        """
264
        return plot_energy_spectrum(self, **kwargs)
265
266
267
268


class CosmicRaysSets(CosmicRaysBase):
    """Set of cosmic rays """
269

270
    def __init__(self, nsets=None, ncrs=None):
271
        self.type = "CosmicRaysSet"
272
        if nsets is None:
273
            CosmicRaysBase.__init__(self, initializer=None)
Teresa Bister's avatar
Teresa Bister committed
274
275
            self.type = "CosmicRaysSet"

276
        # noinspection PyUnresolvedReferences
Teresa Bister's avatar
Teresa Bister committed
277
278
        elif isinstance(nsets, str):
            self.load(nsets)
279
        elif isinstance(nsets, (tuple, float, int, np.integer)):
280
281
282
            self.nsets = nsets[0] if isinstance(nsets, tuple) else nsets
            ncrs = nsets[1] if isinstance(nsets, tuple) else ncrs

283
            # Set the shape first as this is required for __setitem__ used by copy from CosmicRaysBase
284
            CosmicRaysBase.__init__(self, initializer=ncrs * self.nsets)
285
286
            # this number has to be set again as it is overwritten by the init function.
            # It is important to set it before adding the index
Martin Urban's avatar
Martin Urban committed
287
            self.type = "CosmicRaysSet"
288
            self.ncrs = ncrs
289
            self.shape = (int(self.nsets), int(self.ncrs))
290
            self.general_object_store["shape"] = self.shape
291
        elif isinstance(nsets, (list, np.ndarray)):
292
            self._from_list(nsets)
293
        else:
Martin Urban's avatar
Martin Urban committed
294
            # copy case of a cosmic rays set
295
296
            try:
                if nsets.type == self.type:
Martin Urban's avatar
Martin Urban committed
297
298
                    self.general_object_store = {}
                    self.shape = nsets.shape
299
                    self.copy(nsets)
300
                    self._create_access_functions()
301
302
                    # _create_access_functions and the __setitem__ function from the CosmicRaysBase overwrite self.shape
                    self.shape = nsets.shape
303
304
305
306
307
            except AttributeError as e:
                raise AttributeError(str(e))
                # raise NotImplementedError("Trying to instantiate the CosmicRaysSets class with a non "
                #                           "supported type of cosmic_rays")

Marcus Wirtz's avatar
Marcus Wirtz committed
308
    def load(self, filename, **kwargs):
309
310
311
312
        """ Loads cosmic rays from a filename

        :param filename: filename from where to load
        :type filename: str
Marcus Wirtz's avatar
Marcus Wirtz committed
313
        :param kwargs: additional keyword arguments passed to numpy / pickle load functions
314
        """
Marcus Wirtz's avatar
Marcus Wirtz committed
315
        CosmicRaysBase.load(self, filename, **kwargs)
316
        self._create_access_functions()
317
318
319
        if (len(self.shape) == 1) or len(self.general_object_store["shape"]) == 1:
            raise AttributeError("Loading a CosmicRaysBase() object with the CosmicRaysSets() class. Use function "
                                 "cosmic_rays.CosmicRaysBase() instead.")
320
321
        self.ncrs = self.shape[1]
        self.nsets = self.shape[0]
322
323

    def _create_access_functions(self):
324
        super(CosmicRaysSets, self)._create_access_functions()
325
326
        if "shape" in self.general_object_store.keys():
            self.shape = self.general_object_store["shape"]
327

328
    def _from_list(self, l):
329
330
331
332
333
334
335
336
337
338
        types = np.array([type(elem) for elem in l])
        if np.all(types == CosmicRaysBase):
            _nsets, _ncrs = len(l), len(l[0])
        elif np.all(types == CosmicRaysSets):
            _nsets, _ncrs = sum([len(elem) for elem in l]), l[0].ncrs
        else:
            raise TypeError("All elements must be either of type CosmicRays or of type CosmicRaysSets")

        ncrs_each = np.array([elem.ncrs for elem in l])
        if not np.all(ncrs_each == _ncrs):
339
            raise ValueError("The number of cosmic rays must be the same in each set")
340

341
        keys = [sorted(elem.shape_array.dtype.names) for elem in l]
342
343
344
        joint_keys = np.array(["-".join(elem) for elem in keys])
        gos_keys = [sorted(elem.general_object_store.keys()) for elem in l]
        joint_gos_keys = np.array(["-".join(elem) for elem in gos_keys])
345
        if not np.all(joint_keys == joint_keys[0]) or not np.all(joint_gos_keys == joint_gos_keys[0]):
346
            raise AttributeError("All cosmic rays must have the same properties array and general object store")
347
348

        self.__init__(_nsets, _ncrs)
349
        for key in keys[0]:
350
            value = np.array([cr[key] for cr in l]).reshape(self.shape)
351
352
            self.__setitem__(key, value)
        for key in gos_keys[0]:
353
354
            if key == 'shape':
                continue
355
            value = np.array([cr[key] for cr in l])
356
357
            if key == 'vecs':
                value = np.swapaxes(value, 0, 1).reshape(3, -1, _ncrs)
358
359
            if np.all(value == value[0]):
                value = value[0]
360
            self.general_object_store[key] = value
361
        self.general_object_store["shape"] = self.shape
362

363
364
    def __setitem__(self, key, value):
        # casting into int is required to get python3 compatibility
365
        v = value.reshape(int(self.nsets * self.ncrs)) if np.shape(value) == self.shape else value
366
367
        # to avoid the overwriting we use this hack
        self.ncrs = self.ncrs * self.nsets
368
        super(CosmicRaysSets, self).__setitem__(key, v)
369
        # this number has to be set again as it is overwritten by the init function
370
        self.ncrs = int(self.ncrs / int(self.nsets))
371
372

    def __getitem__(self, key):
373
        # noinspection PyUnresolvedReferences
374
        if isinstance(key, (int, np.integer)):
375
            crs = CosmicRaysBase(self.shape_array.dtype)
376
377
            idx_begin = int(key * self.ncrs)
            idx_end = int((key + 1) * self.ncrs)
378
            crs.shape_array = self.shape_array[idx_begin:idx_end]
379
            for k in self.general_object_store.keys():
380
381
                if (k == 'shape'):
                    continue
382
383
384
385
                to_copy = self.get(k)
                if isinstance(to_copy, (np.ndarray, list)):
                    if len(to_copy) == self.nsets:
                        to_copy = to_copy[key]
386
387
                if (k == 'vecs'):
                    to_copy = to_copy[:, key]
388
                crs.__setitem__(k, to_copy)
389
390
            # The order is important
            crs.ncrs = self.ncrs
391
            crs.shape = (self.ncrs, )
392
            return crs
393
        if isinstance(key, (np.ndarray, slice)):
394
            return self._masking(key)
395
        if key in self.general_object_store.keys():
396
            if key in SHAPE_FLEXIBLE:
397
                return self.general_object_store[key] * np.ones((self.nsets, self.ncrs), dtype='int')
398
            return self.general_object_store[key]
399
400
        try:
            # casting into int is required to get python3 compatibility
401
            return np.reshape(self.shape_array[key], self.shape)
402
        except ValueError as e:
403
            if len(self._similar_key(key)) > 0:
404
                value = self._get_values_similar_key(self._similar_key(key).pop(), key)
405
406
407
                if value.size in (np.prod(self.shape), 3*np.prod(self.shape)):
                    shape = self.shape if value.size == np.prod(self.shape) else (-1,)+self.shape
                    return np.reshape(value, shape)
408
409
                raise Exception("Weird error occured, please report this incident with a minimal example!")

410
            raise ValueError("The key %s does not exist and the error message was %s" % (key, str(e)))
411

412
413
414
    def __len__(self):
        return int(self.nsets)

Niklas Uwe Langner's avatar
Niklas Uwe Langner committed
415
    def __add__(self, other):
416
        return CosmicRaysSets([self, other])
Niklas Uwe Langner's avatar
Niklas Uwe Langner committed
417

418
419
420
421
422
423
424
425
426
427
428
429
430
431
    def __iter__(self):
        self._current_idx = 0
        return self

    def __next__(self):
        return self.next()

    def next(self):
        """returns next element when iterating over all elements"""
        self._current_idx += 1
        if self._current_idx > self.nsets:
            raise StopIteration
        return self.__getitem__(self._current_idx - 1)

432
    def _masking(self, sl):
433
        mask = np.zeros(self.shape, dtype=bool)
434
        if isinstance(sl, slice):
435
            mask[sl] = True
Teresa Bister's avatar
Teresa Bister committed
436
            sl = mask
437
        if sl.dtype == bool:
438
439
440
            if sl.shape == (self.nsets,):
                nsets = np.sum(sl)
                ncrs = self.ncrs
441
442
                mask[sl, :] = True
                sl = np.where(mask)
443
444
445
446
447
            elif sl.shape == self.shape:
                ncrs_in_nsets = np.sum(sl, axis=1)
                ncrs = np.amax(ncrs_in_nsets)
                assert self.nsets == np.sum(ncrs_in_nsets == ncrs) + np.sum(ncrs_in_nsets == 0)
                nsets = np.sum(ncrs_in_nsets > 0)
448
                sl = np.where(sl)
449
450
            else:
                raise AssertionError("Slicing dimension is neither (nsets) nor (nsets, ncrs)")
451
        elif sl.dtype == int:
452
            assert (np.amin(sl) >= 0) & (np.amax(sl) < self.nsets)
453
            nsets = len(sl)
454
            ncrs = self.ncrs
455
        else:
456
            raise ValueError("Dtype of slicing ndarray not understood: %s" % (sl.dtype))
457

458
        crs = CosmicRaysSets(nsets, ncrs)
459
460
461
        for key_copy in self.keys():
            if key_copy not in crs.keys():
                to_copy = self.get(key_copy)
462
                # check if array needs to be sliced
463
                if isinstance(to_copy, np.ndarray):
464
                    if to_copy.shape == self.shape:
465
                        to_copy = to_copy[sl]
466
467
468
                    elif to_copy.shape == (self.nsets, ):
                        _sl = np.unique(sl[0]) if isinstance(sl, tuple) else sl
                        to_copy = to_copy[_sl]
469
470
471
                crs.__setitem__(key_copy, to_copy)
        return crs

Martin Urban's avatar
Martin Urban committed
472
473
474
475
    def _update_attributes(self):
        self.ncrs = self.shape[1]
        self.nsets = self.shape[0]

476
477
    def save_readable(self, fname, use_keys=None, **kwargs):
        """
478
        Saves cosmic ray class as ASCII file with general object store written to header.
479
480
481
482
483
484
485

        :param fname: file name of the outfile
        :type fname: str
        :param use_keys: list or tuple of keywords that will be used for the saved file
        :param kwargs: additional named keyword arguments passed to numpy.savetxt()
        """
        dump, header, fmt = self._prepare_readable_output(use_keys)
486
        dump = container.join_struct_arrays([np.repeat(np.arange(self.nsets), self.ncrs), dump])
487
488
489
490
491
492
493
494
        header_parts = header.split("\n")
        header_parts[-1] = ('\n' if len(header_parts) > 1 else '') + 'setID\t' + header_parts[-1]
        fmt.insert(0, '%i')
        kwargs.setdefault('header', "\n".join(header_parts))
        kwargs.setdefault('fmt', fmt)
        kwargs.setdefault('delimiter', '\t')
        np.savetxt(fname, dump, **kwargs)

495
    def sensitivity_2pt(self, set_idx=None, niso=1000, bins=180, **kwargs):
496
497
498
499
        """
        Function to calculate the sensitivity by the 2pt-auto-correlation over a scrambling
        of the right ascension coordinates.

500
501
        :param set_idx: If set, only this set number will be evaluated
        :param niso: Number of isotropic sets to calculate
502
        :param bins: Number of angular bins, 180 correspond to 1 degree binning (np.linspace(0, np.pi, bins+1).
503
        :param kwargs: additional named arguments passed to obs.two_pt_auto()
504
        :return: pvalues in the shape (self.nsets, bins)
505
506
507
508
509
510
511
512
513
514
515
516
        """
        kwargs.setdefault('cumulative', True)
        vec_crs = self.get('vecs')
        _, dec = coord.vec2ang(coord.gal2eq(np.reshape(vec_crs, (3, -1))))

        # calculate auto correlation for isotropic scrambled data
        _ac_iso = np.zeros((niso, bins))
        for i in range(niso):
            _vecs = coord.ang2vec(coord.rand_phi(self.ncrs), np.random.choice(dec, size=self.ncrs))
            _ac_iso[i] = obs.two_pt_auto(_vecs, bins, **kwargs)

        # calculate p-value by comparing the true sets with the isotropic ones
517
518
519
520
521
        set_idx = np.arange(self.nsets) if set_idx is None else [set_idx]
        pvals = np.zeros((len(set_idx), bins))
        for i, idx in enumerate(set_idx):
            _ac_crs = obs.two_pt_auto(vec_crs[:, idx], bins, **kwargs)
            pvals[i] = np.sum(_ac_iso >= _ac_crs[np.newaxis], axis=0) / float(niso)
522
523
        return pvals

524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
    def add_cosmic_rays(self, crs):
        """
        Function to add cosmic rays to the already existing sets of cosmic rays.
        Number of sets must be equal.

        :param crs: CosmicRaysSet instance. The cosmic rays must not contain
                    all original keys. Missing keys are set to zero. If additional
                    keys are provided, they are ignored.
        """
        if not isinstance(crs, CosmicRaysSets):
            raise Exception("You need to add a CosmicRaysSet object!")
        if not self.nsets == crs.nsets:
            raise Exception("Adding CRs to existing CosmicRaysSet instance is only \
                             possible if they have same number of sets!")
        self.add_shape_array(crs.get_array())
539
        self.ncrs = self.ncrs + crs.ncrs
540
        self.shape = (self.nsets, self.ncrs)
541

542
    def plot_eventmap(self, setid=0, **kwargs):  # pragma: no cover
543
544
545
        """
        Function to plot a scatter skymap of the cosmic rays

Marcus Wirtz's avatar
Marcus Wirtz committed
546
547
548
        :param setid: id of the set which is plotted (default: 0)
        :type setid: int
        :param kwargs: additional named arguments passed to plot_eventmap().
549
550
        """
        # noinspection PyTypeChecker
551
        return plot_eventmap(self.get(setid), **kwargs)
552

553
    def plot_heatmap(self, setid=0, **kwargs):  # pragma: no cover
554
555
556
        """
        Function to plot a healpy map of the cosmic ray set

Marcus Wirtz's avatar
Marcus Wirtz committed
557
558
559
        :param setid: id of the set which is plotted (default: 0)
        :type setid: int
        :param kwargs: additional named arguments pased to plot_healpy_map().
560
561
        """
        # noinspection PyTypeChecker
562
        return plot_heatmap(self.get(setid), **kwargs)
563
564
565
566

    def plot_healpy_map(self, setid=0, **kwargs):  # pragma: no cover
        """ Forwards to function plot_heatmap() """
        self.plot_heatmap(setid, **kwargs)
567

568
    def plot_energy_spectrum(self, setid=0, **kwargs):  # pragma: no cover
569
570
571
        """
        Function to plot the energy spectrum of the cosmic ray set

Marcus Wirtz's avatar
Marcus Wirtz committed
572
573
574
        :param setid: id of the set which is plotted (default: 0)
        :type setid: int
        :param kwargs: additional named arguments pased to plot_energy_spectrum().
575
576
577
        """
        # noinspection PyTypeChecker
        crs = self.get(setid)
578
        return plot_energy_spectrum(crs, **kwargs)
579
580
581
582
583
584
585
586
587
588
589
590
591

    def shuffle_events(self):
        """
        Independently shuffle the cosmic rays of each set.
        """
        # This function can be simplified in the future using np.take_along_axis()
        shuffle_ids = np.random.permutation(np.prod(self.shape)).reshape(self.shape)
        shuffle_ids = np.argsort(shuffle_ids, axis=1)
        sets_ids = np.repeat(np.arange(self.nsets), self.ncrs).reshape(self.shape)
        for _key in self.shape_array.dtype.names:
            self.__setitem__(_key, self.__getitem__(_key)[sets_ids, shuffle_ids])
        sets_ids_3d = np.repeat(np.arange(3), np.prod(self.shape)).reshape((3,) + self.shape)
        self.__setitem__('vecs', self.__getitem__('vecs')[sets_ids_3d, sets_ids, np.stack([shuffle_ids] * 3)])