test_cosmic_rays.py 40.7 KB
Newer Older
1
import os
Martin Urban's avatar
Martin Urban committed
2
3
4
import unittest

import numpy as np
5

6
from astrotools.cosmic_rays import CosmicRaysBase, CosmicRaysSets
7
from astrotools import coord, healpytools as hpt
Martin Urban's avatar
Martin Urban committed
8

9
user = os.getenv('USER')
Marcus Wirtz's avatar
Marcus Wirtz committed
10
np.random.seed(0)
Martin Urban's avatar
Martin Urban committed
11
12
13


class TestCosmicRays(unittest.TestCase):
14
15
16
17

    def setUp(self):
        self.ncrs = 10

Martin Urban's avatar
Martin Urban committed
18
    def test_01_n_cosmic_rays(self):
19
20
        crs = CosmicRaysBase(self.ncrs)
        self.assertEqual(crs.ncrs, self.ncrs)
Martin Urban's avatar
Martin Urban committed
21

22
23
        crs = CosmicRaysBase(self.ncrs)
        self.assertEqual(crs.ncrs, int(self.ncrs))
Marcus Wirtz's avatar
Marcus Wirtz committed
24
25
26
27
        ncrs = 10.2
        with self.assertRaises(TypeError):
            CosmicRaysBase(ncrs)

Martin Urban's avatar
Martin Urban committed
28
    def test_02_set_energy(self):
29
        crs = CosmicRaysBase(self.ncrs)
30
        crs["log10e"] = 1 + np.arange(self.ncrs)
Martin Urban's avatar
Martin Urban committed
31
        # noinspection PyTypeChecker,PyUnresolvedReferences
32
33
        self.assertTrue(len(crs.log10e()) == self.ncrs)
        self.assertTrue(len(crs["log10e"]) == self.ncrs)
34
        self.assertTrue(np.all(crs.log10e() > 0))
35
        self.assertTrue(np.all(crs["log10e"] > 0))
Martin Urban's avatar
Martin Urban committed
36
37

    def test_03_set_new_element(self):
38
39
        crs = CosmicRaysBase(self.ncrs)
        crs["karl"] = np.random.uniform(-10, -1, self.ncrs)
40
41
        # noinspection PyTypeChecker,PyUnresolvedReferences
        self.assertTrue(np.all(crs.karl() < 0))
Martin Urban's avatar
Martin Urban committed
42
43
        self.assertTrue(np.all(crs["karl"] < 0))

44
45
        crs = CosmicRaysBase(self.ncrs)
        crs.set("karl", np.random.uniform(-10, -1, self.ncrs))
46
47
48
        self.assertTrue(np.all(crs.karl() < 0))
        self.assertTrue(np.all(crs["karl"] < 0))

Martin Urban's avatar
Martin Urban committed
49
    def test_04_numpy_magic(self):
50
51
52
53
54
55
        crs = CosmicRaysBase(self.ncrs)
        crs["karl"] = np.random.uniform(-10, -1, self.ncrs)
        crs["klaus"] = np.linspace(-10, 10, self.ncrs)
        crs["log10e"] = np.zeros(self.ncrs)
        self.assertEqual(len(crs["log10e"][crs["karl"] <= 0]), self.ncrs)
        self.assertEqual(len(crs["log10e"][crs["klaus"] <= 0]), int(self.ncrs / 2))
Martin Urban's avatar
Martin Urban committed
56

57
    def test_05_copy_int(self):
58
        crs = CosmicRaysBase(self.ncrs)
59
        key = "an_int"
Martin Urban's avatar
Martin Urban committed
60
61
62
63
64
        crs[key] = 10
        crs2 = CosmicRaysBase(crs)
        crs[key] = -2
        # noinspection PyTypeChecker
        self.assertTrue(np.all(crs2[key] == 10))
65
        self.assertTrue(np.all(crs[key] == -2))
Martin Urban's avatar
Martin Urban committed
66

67
    def test_06_copy_array(self):
68
        crs = CosmicRaysBase(self.ncrs)
69
        key = "an_array"
70
        array = np.random.random(self.ncrs)
71
72
        crs[key] = array
        crs2 = CosmicRaysBase(crs)
73
        crs[key] = np.random.random(self.ncrs)
74
75
76
77
        # noinspection PyTypeChecker
        self.assertTrue(np.allclose(array, crs2[key]))

    def test_07_setting_an_element_as_list(self):
78
79
80
81
        crs = CosmicRaysBase(self.ncrs)
        length = np.random.randint(2, 6, self.ncrs)
        random_idx = np.random.randint(0, self.ncrs)
        crs["likelihoods"] = [np.random.uniform(1, 10, length[i]) for i in range(self.ncrs)]
Martin Urban's avatar
Martin Urban committed
82
83
        self.assertEqual(len(crs["likelihoods"][random_idx]), length[random_idx])

84
    def test_08a_saving_and_loading(self):
85
86
        crs = CosmicRaysBase(self.ncrs)
        length = np.random.randint(2, 6, self.ncrs)
Martin Urban's avatar
Martin Urban committed
87
        key = "karl"
88
        crs[key] = [np.random.uniform(1, 10, length[i]) for i in range(self.ncrs)]
89
        fname = "/tmp/test-base08a-%s.npy" % user
Martin Urban's avatar
Martin Urban committed
90
91
92
        crs.save(fname)
        crs3 = CosmicRaysBase(fname)
        # noinspection PyTypeChecker
93
        os.remove(fname)
94
        self.assertTrue(np.all([np.all(crs3[key][i] == crs[key][i]) for i in range(self.ncrs)]))
Martin Urban's avatar
Martin Urban committed
95

96
    def test_08b_saving_and_loading_pickle(self):
97
98
        crs = CosmicRaysBase(self.ncrs)
        length = np.random.randint(2, 6, self.ncrs)
Martin Urban's avatar
Martin Urban committed
99
100
        key = "karl"
        key2 = "production_date"
101
        crs[key] = [np.random.uniform(1, 10, length[i]) for i in range(self.ncrs)]
Martin Urban's avatar
Martin Urban committed
102
        crs[key2] = "YYYY-MM-DD-HH-MM-SS"
103
        fname = "/tmp/test-base08b-%s.pkl" % user
Martin Urban's avatar
Martin Urban committed
104
105
        crs.save(fname)
        crs3 = CosmicRaysBase(fname)
106
        os.remove(fname)
Martin Urban's avatar
Martin Urban committed
107
        # noinspection PyTypeChecker
108
        self.assertTrue(np.all([np.all(crs3[key][i] == crs[key][i]) for i in range(self.ncrs)]))
Martin Urban's avatar
Martin Urban committed
109
        # noinspection PyTypeChecker,PyUnresolvedReferences
110
        self.assertTrue(np.all([np.all(crs3.karl()[i] == crs.karl()[i]) for i in range(self.ncrs)]))
Martin Urban's avatar
Martin Urban committed
111
112
        self.assertTrue(crs3[key2] == crs[key2])

113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
    def test_08c_saving_and_loading_npz(self):
        crs = CosmicRaysBase(self.ncrs)
        length = np.random.randint(2, 6, self.ncrs)
        key = "karl"
        key2 = "production_date"
        crs[key] = [np.random.uniform(1, 10, length[i]) for i in range(self.ncrs)]
        crs[key2] = "YYYY-MM-DD-HH-MM-SS"
        fname = "/tmp/test-base08c-%s.npz" % user
        crs.save(fname)
        crs3 = CosmicRaysBase(fname)
        os.remove(fname)
        # noinspection PyTypeChecker
        self.assertTrue(np.all([np.all(crs3[key][i] == crs[key][i]) for i in range(self.ncrs)]))
        # noinspection PyTypeChecker,PyUnresolvedReferences
        self.assertTrue(np.all([np.all(crs3.karl()[i] == crs.karl()[i]) for i in range(self.ncrs)]))
        self.assertTrue(crs3[key2] == crs[key2])

130
131
132
133
134
135
136
137
138
139
    def test_08d_saving_and_loading_no_shape(self):
        crs = CosmicRaysBase(self.ncrs)
        key = "karl"
        crs[key] = "no_shaped_array"
        fname = "/tmp/test-base08d-%s.npz" % user
        crs.save(fname)
        crs3 = CosmicRaysBase(fname)
        os.remove(fname)
        self.assertTrue(crs3[key] == crs[key])

140
141
142
143
144
145
146
147
148
149
150
151
    def test_09_loading_with_sets(self):
        sets_shape = (10, 5)
        fname = "/tmp/test-base09-%s" % user
        crs_sets = CosmicRaysSets(sets_shape)
        crs_sets['array'] = np.random.random(sets_shape)

        for ending in ['.npz', '.npy', '.pkl']:
            crs_sets.save(fname + ending)
            with self.assertRaises(AttributeError):
                CosmicRaysBase(fname + ending)
            os.remove(fname + ending)

152
    def test_10_start_from_dict(self):
Martin Urban's avatar
Martin Urban committed
153
154
155
156
        cosmic_rays_dtype = np.dtype([("log10e", float), ("xmax", float), ("time", str), ("other", object)])
        crs = CosmicRaysBase(cosmic_rays_dtype)
        self.assertEqual(crs.ncrs, 0)

157
    def test_11a_add_crs(self):
Martin Urban's avatar
Martin Urban committed
158
        cosmic_rays_dtype = np.dtype([("log10e", float), ("xmax", float), ("time", "|S8"), ("other", object)])
159
160
161
162
        new_crs = np.zeros(shape=self.ncrs, dtype=[("log10e", float), ("xmax", float), ("time", "|S2")])
        new_crs["log10e"] = np.random.exponential(1, self.ncrs)
        new_crs["xmax"] = np.random.uniform(800, 900, self.ncrs)
        new_crs["time"] = ["0"] * self.ncrs
163
164

        crs = CosmicRaysBase(cosmic_rays_dtype)
165
        crs.add_cosmic_rays(CosmicRaysBase(new_crs))
166
        self.assertEqual(crs.ncrs, self.ncrs)
Martin Urban's avatar
Martin Urban committed
167
        # noinspection PyTypeChecker
168
        self.assertTrue(np.all(crs["time"] == b"0"))
Martin Urban's avatar
Martin Urban committed
169
170
171
        self.assertEqual(crs["time"].dtype, "|S8")
        # noinspection PyTypeChecker
        self.assertTrue(np.all(crs["xmax"] > 0))
172
173
174
175
176
177
178

        crs2 = CosmicRaysBase(cosmic_rays_dtype)
        crs2.add_shape_array(new_crs)
        self.assertEqual(crs.ncrs, crs2.ncrs)
        self.assertTrue(np.all(crs2["time"] == b"0"))
        # noinspection PyTypeChecker
        self.assertTrue(np.allclose(new_crs["xmax"], crs2["xmax"]))
Martin Urban's avatar
Martin Urban committed
179

180
181
182
183
184
185
186
187
    def test_11a_add_crs_from_cosmic_rays_class(self):
        ncrs1 = 10
        ncrs2 = 40
        crs1 = CosmicRaysBase(ncrs1)
        crs2 = CosmicRaysBase(ncrs2)
        crs1.add_cosmic_rays(crs2)
        self.assertEqual(crs1.ncrs, ncrs1 + ncrs2)

188
    def test_12_len(self):
189
        crs = CosmicRaysBase(self.ncrs)
Martin Urban's avatar
Martin Urban committed
190
191
192
        # noinspection PyTypeChecker
        self.assertEqual(len(crs), crs.ncrs)

193
    def test_13_add_new_keys(self):
194
195
196
197
        crs = CosmicRaysBase(self.ncrs)
        crs["log10e"] = np.zeros(self.ncrs)
        crs["C_best_fit"] = np.ones(self.ncrs, dtype=float)
        crs["C_best_fit_object"] = np.ones(self.ncrs, dtype=[("C_best_fit_object", object)])
198
199
200
201
202
        crs["rigidities_fit"] = crs["log10e"]
        # noinspection PyTypeChecker
        self.assertTrue(np.all(crs["C_best_fit"] == 1))
        # noinspection PyTypeChecker
        self.assertTrue(np.all(crs["rigidities_fit"] == crs["log10e"]))
203

204
    def test_13a_change_via_function(self):
205
206
        crs = CosmicRaysBase(self.ncrs)
        crs["log10e"] = np.zeros(self.ncrs)
207
208
209
        crs.log10e(1)
        self.assertTrue(np.all(crs["log10e"] == 1))
        self.assertTrue(np.all(crs.log10e() == 1))
210

211
    def test_14_access_by_id(self):
212
        idx = 8
213
        crs = CosmicRaysBase(self.ncrs)
214
        # crs["C_best_fit"] = np.ones(ncrs, dtype=[("C_best_fit", np.float64)])
215
        crs["C_best_fit"] = np.ones(self.ncrs, dtype=float)
216
217
        self.assertEqual(crs[idx]["C_best_fit"], 1)

218
    def test_15_iteration(self):
219
        crs = CosmicRaysBase(self.ncrs)
220
221
        crs["idx"] = np.arange(self.ncrs)
        crs["random"] = np.random.random(self.ncrs)
Martin Urban's avatar
Martin Urban committed
222
        for i, cr in enumerate(crs):
223
224
225
            self.assertTrue(cr.dtype.names == ('idx', 'random'))
            self.assertTrue(cr["idx"] == crs["idx"][i])
            self.assertTrue(cr["random"] == crs["random"][i])
Martin Urban's avatar
Martin Urban committed
226

227
    @unittest.skipIf(os.path.isfile("/.dockerenv"), "Plotting in Docker environment is not possible!")
228
    def test_16_plotting(self):
229
230
231
232
233
        ncrs = 1000
        crs = CosmicRaysBase(ncrs)
        crs['pixel'] = np.random.randint(0, 49152, ncrs)
        crs['log10e'] = 17. + 2.5 * np.random.random(ncrs)

234
        fname = "/tmp/energy_spectrum-%s.png" % user
235
        crs.plot_energy_spectrum(opath=fname)
236
        crs.plot_eventmap()
237
        crs.plot_healpy_map()
238
239
        self.assertTrue(os.path.isfile(fname))
        os.remove(fname)
Martin Urban's avatar
Martin Urban committed
240

241
    def test_17_initialize_with_array(self):
242
243
244
245
        energies = np.array(np.random.uniform(18, 20, 100), dtype=[("Energy", float)])
        crs = CosmicRaysBase(energies)
        # noinspection PyTypeChecker
        self.assertTrue(np.all(crs.get("Energy") >= 18))
246
        self.assertTrue("Energy" in crs.shape_array.dtype.names)
247
248
        # noinspection PyTypeChecker,PyUnresolvedReferences
        self.assertTrue(np.all(crs.Energy() >= 18))
249

250
    def test_17a_initialize_with_2d_array(self):
251
252
253
        energies = np.random.uniform(18, 20, self.ncrs)
        ids = np.arange(self.ncrs)
        crs_arr = np.empty((self.ncrs,), dtype=np.dtype([("energy", "f"), ("cr_id", "i")]))
254
255
256
257
258
        crs_arr["energy"] = energies
        crs_arr["cr_id"] = ids
        crs = CosmicRaysBase(crs_arr)
        # noinspection PyTypeChecker
        self.assertTrue(np.all(crs.get("energy") >= 18))
259
        self.assertTrue(np.all(crs.get("cr_id") == np.arange(self.ncrs)))
260
261
        self.assertTrue("energy" in crs.shape_array.dtype.names)
        self.assertTrue("cr_id" in crs.shape_array.dtype.names)
262
263
        # noinspection PyTypeChecker,PyUnresolvedReferences
        self.assertTrue(np.all(crs.energy() >= 18))
264

265
    def test_18_combine_keys(self):
266
        crs = CosmicRaysBase(self.ncrs)
267
        self.assertTrue(len(crs.keys()) == 0)
268
        crs['array'] = np.random.randint(0, 49152, self.ncrs)
269
270
        crs['ndarray'] = np.random.random((5, 2))
        crs['float'] = 5
271
272
273
        self.assertTrue('ndarray' in crs.keys())
        self.assertTrue('array' in crs.keys())
        self.assertTrue('float' in crs.keys())
274
275

    def test_19_keys_available(self):
276
277
        crs = CosmicRaysBase(self.ncrs)
        crs['array'] = np.random.randint(0, 49152, self.ncrs)
MarCus's avatar
MarCus committed
278
        crs['ndarray'] = np.random.random((5, 2))
279
280
        self.assertTrue('array' in crs.keys())
        self.assertTrue('ndarray' in crs.keys())
Marcus Wirtz's avatar
Marcus Wirtz committed
281

282
    def test_20_enumerate(self):
283
        crs = CosmicRaysBase(initializer=self.ncrs)
284
285
286
        imax1, imax2 = 0, 0
        for i, _ in enumerate(crs):
            imax1 = i + 1
287
        crs["energy"] = np.random.uniform(0, 1, self.ncrs)
288
289
        for i, _ in enumerate(crs):
            imax2 = i + 1
290
291
        self.assertTrue(imax1 == self.ncrs)
        self.assertTrue(imax2 == self.ncrs)
Marcus Wirtz's avatar
Marcus Wirtz committed
292

293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
    def test_21a_assign_with_list(self):
        a = np.random.random((3, self.ncrs))
        vecs = np.random.random((3, 3, self.ncrs))
        info = "info"
        crs1 = CosmicRaysBase(self.ncrs)
        crs1["a"], crs1["vecs"], crs1["info"] = a[0], vecs[:, 0], info
        crs2 = CosmicRaysBase(self.ncrs)
        crs2["a"], crs2["vecs"], crs2["info"] = a[1], vecs[:, 1], info
        crs3 = CosmicRaysBase(self.ncrs)
        crs3["a"], crs3["vecs"], crs3["info"] = a[2], vecs[:, 2], info

        crs_merged = CosmicRaysBase([crs1, crs2, crs3])
        self.assertTrue(np.allclose(np.hstack(a), crs_merged["a"]))
        self.assertTrue(np.allclose(np.reshape(vecs, (3, -1)), crs_merged["vecs"]))
        self.assertTrue(info == crs_merged["info"])
        self.assertTrue(crs_merged.ncrs == 3 * self.ncrs)
309
        self.assertTrue(crs_merged.shape == (crs_merged.ncrs, ))
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326

    def test_21b_assign_with_list_unequal_cr_number(self):
        a1 = np.random.random(10)
        a2 = np.random.random(5)
        vecs1 = np.random.random((3, 10))
        vecs2 = np.random.random((3, 5))
        info = "info"
        crs1 = CosmicRaysBase(10)
        crs1["a"], crs1["vecs"], crs1["info"] = a1, vecs1, info
        crs2 = CosmicRaysBase(5)
        crs2["a"], crs2["vecs"], crs2["info"] = a2, vecs2, info

        crs_merged = CosmicRaysBase([crs1, crs2])
        self.assertTrue(np.allclose(np.concatenate([a1, a2]), crs_merged["a"]))
        self.assertTrue(np.allclose(np.concatenate([vecs1, vecs2], axis=1), crs_merged["vecs"]))
        self.assertTrue(info == crs_merged["info"])
        self.assertTrue(crs_merged.ncrs == 15)
327
        self.assertTrue(crs_merged.shape == (crs_merged.ncrs, ))
328
329
330
331
332

    def test_22_numpy_integer(self):
        n = np.int16(64)
        crs = CosmicRaysBase(n)
        self.assertTrue(crs.ncrs == n)
333

334
    def test_23_access_non_existing_element(self):
335
336
        crs = CosmicRaysBase(self.ncrs)
        crs['array'] = np.random.randint(0, 49152, self.ncrs)
337
        with self.assertRaises(ValueError):
338
            crs["non_existing"]
339

340
    def test_24_set_unallowed_items(self):
341
        crs = CosmicRaysBase(self.ncrs)
342
343
344
345
        with self.assertRaises(KeyError):
            # case where the user does use the value as key and vice versa
            crs[[1, 2, 3]] = "key"

346
    def test_25a_slicing_base(self):
347
348
        crs = CosmicRaysBase(self.ncrs)
        energy = np.random.random(self.ncrs)
349
        crs['energy'] = energy
350
        crs['vecs'] = coord.rand_vec(self.ncrs)
351
        crs_sub = crs[energy < 0.3]
352
        self.assertTrue(hasattr(crs_sub, 'keys'))
353
        self.assertTrue(len(crs_sub) < self.ncrs)
354
        self.assertTrue(len(crs_sub) == crs_sub.ncrs)
355
        self.assertTrue(len(crs_sub['energy']) == len(crs_sub))
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
        self.assertTrue(np.shape(crs_sub['vecs']) == (3, len(crs_sub)))
        self.assertTrue(crs.type == 'CosmicRays')
        self.assertTrue(crs_sub.type == 'CosmicRays')

    def test_25b_slicing_base(self):
        crs = CosmicRaysBase(self.ncrs)
        energy = np.random.random(self.ncrs)
        crs['energy'] = energy
        crs['vecs'] = coord.rand_vec(self.ncrs)
        n_sub = 3
        crs_sub = crs[np.random.randint(low=0, high=self.ncrs, size=n_sub)]
        self.assertTrue(hasattr(crs_sub, 'keys'))
        self.assertTrue(len(crs_sub) == n_sub)
        self.assertTrue(len(crs_sub['energy']) == n_sub)
        self.assertTrue(np.shape(crs_sub['vecs']) == (3, n_sub))
371
372
        self.assertTrue(crs.type == 'CosmicRays')
        self.assertTrue(crs_sub.type == 'CosmicRays')
373

374
375
376
377
378
379
380
381
382
383
384
    def test_26_set_unfortunate_length_of_string(self):
        _str = 'hallo'
        ncrs = len(_str)
        crs = CosmicRaysBase(ncrs)
        crs['feature1'] = 2 * _str
        crs['feature2'] = _str
        crs['feature1'] = _str
        self.assertTrue(('feature1') in crs.keys())
        self.assertTrue((crs['feature1'] == _str) and (crs['feature2'] == _str))
        self.assertTrue(('feature2') in crs.keys())

385
386
387
388
    def test_27_save_readable(self):
        crs = CosmicRaysBase(self.ncrs)
        keys = ['a_str', 'a_float', 'an_int', 'ndarray', 'array', 'custom_array']
        entries = ['karl', 42.42, 4, np.random.random(size=(2, self.ncrs)),
389
                   1 + np.random.random(self.ncrs), np.random.random(17)]
390
391
392
        for _key, _entry in zip(keys, entries):
            crs[_key] = _entry

393
394
        outpath = '/tmp/cosmicraysbase-readable-%s.npz' % user
        crs.save_readable(outpath)
395

396
        array = np.genfromtxt(outpath)
397
        self.assertTrue(np.allclose(array, entries[keys.index('array')]))
398
399
        with open(outpath, 'r') as rl:
            lines = "".join(rl.readlines())
400
        for _key, _entry in zip(keys, entries):
401
            self.assertTrue(_key in lines)
402
        os.remove(outpath)
403

404
    def test_28a_similar_keys(self):
405
406
407
408
        crs = CosmicRaysBase(self.ncrs)
        nside = 64

        crs['pix'] = [1, 2, 3]
409
        self.assertTrue(np.array_equal(crs['pix'], crs['pixel']))
410
411
412
413
        self.assertTrue(np.array_equal(crs['pix'], hpt.ang2pix(nside, crs['lon'], crs['lat'])))
        self.assertTrue(np.array_equal(crs['pix'], hpt.vec2pix(nside, crs['vecs'])))

        crs = CosmicRaysBase(self.ncrs)
414
415
416
        crs['lon'] = [0.1, 0.2, 0.5]
        crs['lat'] = [-0.2, 0., 0.2]
        vecs = crs['vecs']
417
        self.assertTrue(vecs.shape == (3, 3))
418
419
420
421
        _lon, _lat = coord.vec2ang(vecs)
        self.assertTrue(np.allclose(crs['lon'], _lon))
        self.assertTrue(np.allclose(crs['lat'], _lat))

422
423
        crs = CosmicRaysBase(self.ncrs)
        crs['lon'] = [1., 2., -0.3]
424
        with self.assertRaises(ValueError):
425
            crs['pix']
426

427
428
    def test_28b_similar_keys(self):
        crs = CosmicRaysBase(self.ncrs)
429
        crs['e'] = [5, 2, 6, 1, 5, 3, 8, 7, 9, 4]
430
431
432
433
434
        self.assertTrue(np.array_equal(crs['e'], crs['energy']))
        self.assertTrue(np.array_equal(crs['e'], crs['E']))
        self.assertTrue(np.allclose(crs['e'], 10**np.array(crs['log10e'])))

        crs = CosmicRaysBase(self.ncrs)
435
        crs['log10e'] = [3, 4, 7, 1, 2, 5, 8, 9, 6, 2]
436
437
438
        self.assertTrue(np.allclose(crs['log10e'], np.log10(crs['e'])))
        self.assertTrue(np.allclose(crs['log10e'], np.log10(crs['energy'])))
        self.assertTrue(np.allclose(crs['log10e'], np.log10(crs['E'])))
439

440

441
class TestCosmicRaysSets(unittest.TestCase):
442
443
444
445
446
447

    def setUp(self):
        self.ncrs = 10
        self.nsets = 15
        self.shape = (self.nsets, self.ncrs)

448
    def test_01_create(self):
449
450
451
452
        crsset = CosmicRaysSets(self.shape)
        self.assertEqual(crsset.shape, self.shape)
        self.assertEqual(crsset.ncrs, self.ncrs)
        self.assertEqual(crsset.nsets, self.nsets)
453

454
    def test_01a_create_with_None(self):
455
456
        crsset = CosmicRaysSets(self.shape)
        log10e = np.random.random(self.shape)
Teresa Bister's avatar
Teresa Bister committed
457
        crsset['log10e'] = log10e
458
        outpath = "/tmp/cosmicraysset-%s.npz" % user
Teresa Bister's avatar
Teresa Bister committed
459
460
461
462
463
464
465
        crsset.save(outpath)

        crsset2 = CosmicRaysSets(None)
        crsset2.load(outpath)
        self.assertTrue('log10e' in crsset2.keys())
        self.assertTrue((crsset2['log10e'] == log10e).all())
        os.remove(outpath)
466
467
468
469
470
471
472

    def test_01b_create_with_fake_object(self):
        class test:
            def __init__(self):
                self.type = "CosmicRaysSet"
        with self.assertRaises(AttributeError):
            t = test()
473
            CosmicRaysSets(t)
474
475

    def test_02_get_element_from_set(self):
476
        crsset = CosmicRaysSets(self.shape)
477
        # noinspection PyTypeChecker
478
        crsset["log10e"] = np.zeros(shape=crsset.shape)
479
        # noinspection PyTypeChecker
480
        self.assertTrue(np.all(crsset["log10e"] == 0.))
481
        self.assertEqual(crsset["log10e"].shape, self.shape)
482
483

    def test_03_set_element(self):
484
485
        crsset = CosmicRaysSets(self.shape)
        energies = np.random.uniform(18, 20, size=self.shape)
486
487
        crsset["log10e"] = energies
        # noinspection PyTypeChecker
488
489
        self.assertTrue('log10e' in crsset.keys())
        self.assertTrue(len(crsset.log10e()) == self.nsets)
490
491
492
493
        self.assertTrue(np.all(crsset["log10e"] >= 18))

    def test_04_get_set_by_number(self):
        set_number = 3
494
        crsset = CosmicRaysSets(self.shape)
495
        crsset["an_array"] = np.random.random(self.shape)
496
497
        crsset["creator"] = "Martin"
        subset = crsset[set_number]
498
        self.assertTrue(isinstance(subset, CosmicRaysBase))
499
        self.assertTrue(len(subset), self.ncrs)
500
        self.assertTrue("shape" not in subset.keys())
501
        self.assertTrue(subset["creator"], "Martin")
502
        self.assertTrue(np.allclose(subset["an_array"], crsset["an_array"][set_number]))
503
        self.assertTrue(len(subset.shape_array), self.ncrs)
504
        self.assertTrue(subset.shape == (self.ncrs, ))
505

506
    # noinspection PyTypeChecker
507
508
    def test_05_set_in_subset(self):
        set_number = 3
509
        crsset = CosmicRaysSets(self.shape)
510
        crsset["creator"] = "Martin"
511
        crsset["log10e"] = np.zeros(shape=self.shape)
512
        subset = crsset[set_number]
513
        subset["log10e"] = np.random.uniform(18, 20, self.ncrs)
514
515
        # noinspection PyTypeChecker
        self.assertTrue(np.all(subset["log10e"] >= 18))
516
517
        idx_begin = int(self.ncrs * set_number)
        idx_end = int(self.ncrs * (set_number + 1))
518
519
520
        self.assertTrue(np.all(crsset.shape_array[idx_begin:idx_end]["log10e"] >= 18))
        self.assertTrue(np.all(crsset.shape_array[0:idx_begin]["log10e"] == 0))
        self.assertTrue(np.all(crsset.shape_array[idx_end:]["log10e"] == 0))
521

522
    def test_06_copy(self):
523
        crs = CosmicRaysSets(self.shape)
524
        key = "an_int"
525
526
527
528
529
530
        crs[key] = 10
        crs2 = CosmicRaysSets(crs)
        crs[key] = -2
        # noinspection PyTypeChecker
        self.assertTrue(np.all(crs2[key] == 10))

531
    def test_07_copy_array(self):
532
        crs = CosmicRaysSets(self.shape)
533
        key = "an_array"
534
        array = np.random.random(self.shape)
535
536
        crs[key] = array
        crs2 = CosmicRaysSets(crs)
537
        crs[key] = np.random.random(self.shape)
538
        # noinspection PyTypeChecker
539
        self.assertTrue((key in crs.keys()) and (key in crs2.keys()))
540
541
542
543
544
545
        self.assertTrue(key not in crs.general_object_store.keys())
        self.assertTrue(key not in crs2.general_object_store.keys())
        self.assertTrue(np.allclose(array, crs2[key]))

    def test_08_copy_gos_array(self):
        # gos = general object store
546
        crs = CosmicRaysSets(self.shape)
547
        key = "an_array"
548
        array = np.random.random(self.ncrs)
549
550
        crs[key] = array
        crs2 = CosmicRaysSets(crs)
551
        crs[key] = np.random.random(self.ncrs)
552
        # noinspection PyTypeChecker
553
        self.assertTrue((key in crs.keys()) and (key in crs2.keys()))
554
555
556
557
558
        self.assertTrue(key in crs.general_object_store.keys())
        self.assertTrue(key in crs2.general_object_store.keys())
        self.assertTrue(np.allclose(array, crs2[key]))

    def test_09_save(self):
559
        outpath = "/tmp/test-sets09-%s.pkl" % user
560
        crsset = CosmicRaysSets(self.shape)
561
        crsset["creator"] = "Marcus"
562
563
564
        crsset["log10e"] = np.zeros(shape=crsset.shape)
        crsset.save(outpath)
        self.assertTrue(os.path.exists(outpath))
565
        os.remove(outpath)
566

567
    def test_09a_save_load(self):
568
        outpath = "/tmp/test-sets09a-%s.npz" % user
569
        crsset = CosmicRaysSets(self.shape)
570
571
572
573
        crsset["creator"] = "Marcus"
        crsset["log10e"] = np.zeros(shape=crsset.shape)
        crsset.save(outpath)

574
        crsset2 = CosmicRaysSets(self.shape)
575
        crsset2.load(outpath)
576
        os.remove(outpath)
577
578
        self.assertTrue("creator" in crsset2.keys())
        self.assertTrue("log10e" in crsset2.keys())
579
580
        self.assertTrue(np.allclose(crsset["log10e"], crsset2["log10e"]))

581
    def test_09b_save_load(self):
582
        outpath = "/tmp/test-sets09b-%s.npz" % user
583
        crsset = CosmicRaysSets(self.shape)
584
        crsset["log10e"] = np.zeros(shape=crsset.shape)
585
586
587
588
        crsset["creator"] = "Marcus"
        crsset.save(outpath)

        crsset2 = CosmicRaysSets(outpath)
589
        os.remove(outpath)
590
        self.assertTrue("creator" in crsset2.keys())
591

592
593
594
595
596
597
598
599
600
601
602
    def test_09c_loading_with_base(self):
        fname = "/tmp/test-base09-%s" % user
        crs = CosmicRaysBase(10)
        crs['array'] = np.random.random(10)

        for ending in ['.npz', '.npy', '.pkl']:
            crs.save(fname + ending)
            with self.assertRaises(AttributeError):
                CosmicRaysSets(fname + ending)
            os.remove(fname + ending)

603
    def test_10_create_from_filename(self):
604
        # Create first the set and save it to file
605
        outpath = "/tmp/cosmicraysset-%s.npz" % user
606
        crsset = CosmicRaysSets(self.shape)
607
608
609
610
611
        crsset["creator"] = "Martin"
        crsset["log10e"] = np.ones(shape=crsset.shape)
        crsset.save(outpath)
        # reload the set as a new cosmic rays set
        crsset2 = CosmicRaysSets(outpath)
612
        os.remove(outpath)
613
        self.assertTrue(crsset2["creator"] == "Martin")
614
        self.assertTrue(np.shape(crsset2["log10e"]) == self.shape)
615
616
617
        # noinspection PyTypeChecker
        self.assertTrue(np.all(crsset2["log10e"] == 1))

618
    @unittest.skipIf(os.path.isfile("/.dockerenv"), "Plotting in Docker environment is not possible!")
619
    def test_11_plot(self):
620
621
622
623
        ncrs = 100
        crs = CosmicRaysSets((self.nsets, ncrs))
        crs['pixel'] = np.random.randint(0, 49152, (self.nsets, ncrs))
        crs['log10e'] = 18. + 2.5 * np.random.random((self.nsets, ncrs))
624
625
626
627
628
629

        crs.plot_eventmap()
        crs.plot_energy_spectrum()
        crs.plot_healpy_map()
        self.assertTrue(True)

630
    @unittest.skipIf(os.path.isfile("/.dockerenv"), "Plotting in Docker environment is not possible!")
631
    def test_12_plot_from_loaded_cosmic_rays_set(self):
632
633
634
635
        ncrs = 100
        crs = CosmicRaysSets((self.nsets, ncrs))
        crs['pixel'] = np.random.randint(0, 49152, (self.nsets, ncrs))
        crs['log10e'] = 18. + 2.5 * np.random.random((self.nsets, ncrs))
636
637
        outpath = "/tmp/test_08-%s.npy" % user
        crs.save(outpath)
638

639
640
641
642
643
        crs3 = CosmicRaysSets(outpath)
        crs3.plot_eventmap(opath=outpath.replace('.npy', '.png'))
        self.assertTrue(os.path.exists(outpath))
        self.assertTrue(os.path.exists(outpath.replace('.npy', '.png')))
        os.remove(outpath)
644

645
646
647
648
649
650
651
652
    def test_13a_change_via_function(self):
        crs = CosmicRaysSets(self.nsets, self.ncrs)
        crs["log10e"] = np.zeros((self.nsets, self.ncrs))
        crs.log10e(1)
        self.assertTrue(np.all(crs["log10e"] == 1))
        self.assertTrue(np.all(crs.log10e() == 1))

    def test_13b_keys_available(self):
653
654
655
656
        ncrs = 100
        crs = CosmicRaysSets((self.nsets, ncrs))
        crs['ndarray'] = np.random.randint(0, 49152, (self.nsets, ncrs))
        crs['array'] = np.random.random(self.nsets)
marcus's avatar
marcus committed
657
        crs['float'] = 5.
658
659
660
        self.assertTrue('ndarray' in crs.keys())
        self.assertTrue('array' in crs.keys())
        self.assertTrue('float' in crs.keys())
marcus's avatar
marcus committed
661

662
663
664
665
666
    def test_14_mask_subset(self):
        nsets = 100
        ndarray = np.random.randint(0, 49152, (nsets, self.ncrs))
        array = np.random.random(nsets)
        crs = CosmicRaysSets((nsets, self.ncrs))
667
668
        crs['ndarray'] = ndarray
        crs['array'] = array
669
670
        crs['string'] = 'blubb'
        crs['integer'] = 5
671
672
        mask = np.linspace(0, 1, nsets) < 0.3
        nsets_subset = 30
673
        crs_subset = crs[mask]
674
        self.assertTrue(nsets == crs.nsets)
675
        self.assertTrue((crs_subset.nsets == nsets_subset) and (crs_subset.ncrs == self.ncrs))
676
677
        self.assertTrue('ndarray' in crs_subset.keys())
        self.assertTrue('array' in crs_subset.keys())
678
        self.assertTrue(np.allclose(crs_subset['array'], array[mask]))
679
680
681
        self.assertEqual(crs_subset['string'], crs['string'])
        self.assertEqual(crs_subset['string'], 'blubb')
        self.assertEqual(crs_subset['integer'], crs['integer'])
682

Marcus Wirtz's avatar
Design    
Marcus Wirtz committed
683
    def test_15_indexing_subset(self):
684
685
686
        ndarray = np.random.randint(0, 49152, self.shape)
        array = np.random.random(self.nsets)
        crs = CosmicRaysSets(self.shape)
687
688
        crs['ndarray'] = ndarray
        crs['array'] = array
689
690
        crs['string'] = 'blubb'
        crs['integer'] = 5
691
        indexing = np.random.choice(np.arange(self.nsets), 10, replace=False)
692
        crs_subset = crs[indexing]
693
        self.assertTrue(self.nsets == crs.nsets)
694
695
        self.assertTrue('ndarray' in crs_subset.keys())
        self.assertTrue('array' in crs_subset.keys())
696
        self.assertTrue(crs_subset.nsets == 10)
697
        self.assertTrue(np.allclose(crs_subset['array'], array[indexing]))
698
699
700
        self.assertEqual(crs_subset['string'], crs['string'])
        self.assertEqual(crs_subset['string'], 'blubb')
        self.assertEqual(crs_subset['integer'], crs['integer'])
701

Marcus Wirtz's avatar
Design    
Marcus Wirtz committed
702
    def test_16_slicing_subset(self):
703
704
705
        ndarray = np.random.randint(0, 49152, self.shape)
        array = np.random.random(self.nsets)
        crs = CosmicRaysSets(self.shape)
706
707
        crs['ndarray'] = ndarray
        crs['array'] = array
708
709
        crs['string'] = 'blubb'
        crs['integer'] = 5
710
711
        low, up = 2, 10
        crs_subset = crs[low:up]
712
        self.assertTrue(self.nsets == crs.nsets)
713
714
        self.assertTrue('ndarray' in crs_subset.keys())
        self.assertTrue('array' in crs_subset.keys())
715
716
        self.assertTrue(crs_subset.nsets == int(up - low))
        self.assertTrue(np.allclose(crs_subset['array'], array[low:up]))
717
718
719
        self.assertEqual(crs_subset['string'], crs['string'])
        self.assertEqual(crs_subset['string'], 'blubb')
        self.assertEqual(crs_subset['integer'], crs['integer'])
720
721
        with self.assertRaises(ValueError):
            a = np.zeros(10, dtype=[("a", float)])
722
            crs[a]
723

Marcus Wirtz's avatar
Design    
Marcus Wirtz committed
724
    def test_17_access_functions(self):
725
726
727
728
729
730
731
        crs = CosmicRaysSets(self.shape)
        self.assertTrue(crs.shape, self.shape)
        crs["log10e"] = np.random.random(self.shape)
        self.assertTrue(crs["log10e"].shape, self.shape)
        self.assertTrue(crs.log10e().shape, self.shape)
        self.assertTrue(crs.shape, self.shape)
        creator = "Peter"
732
733
734
        crs["creator"] = creator
        self.assertTrue(crs["creator"], creator)
        self.assertTrue(crs.creator(), creator)
735
        self.assertTrue(crs.shape, self.shape)
736
        with self.assertRaises(TypeError):
Marcus Wirtz's avatar
Marcus Wirtz committed
737
            crs.shape()
738

Marcus Wirtz's avatar
Design    
Marcus Wirtz committed
739
    def test_18_failing_get(self):
740
741
        creator = "Peter"
        crs = CosmicRaysSets(self.shape)
742
        crs["creator"] = creator
743
        with self.assertRaises((ValueError, TypeError)):
744
            crs[["test"]]
745

746
747
748
749
750
751
752
753
754
755
756
757
758
    def test_19_iteration(self):
        crs_set = CosmicRaysSets(self.shape)
        crs_set["creator"] = "Marcus"
        crs_set["idx"] = np.random.random(self.shape)
        crs_set["log10e"] = np.random.random(self.shape)
        for i, _crs in enumerate(crs_set):
            self.assertTrue(isinstance(_crs, CosmicRaysBase))
            self.assertTrue(_crs.ncrs == self.ncrs)
            self.assertTrue(_crs["creator"] == crs_set["creator"])
            self.assertTrue(np.allclose(_crs["idx"], crs_set["idx"][i]))
            self.assertTrue(np.allclose(_crs["log10e"], crs_set["log10e"][i]))

    def test_20a_create_from_crs_list_base(self):
759
760
        crs1 = CosmicRaysBase(self.ncrs)
        crs1["log10e"] = np.random.rand(self.ncrs)
761
        crs1["name"] = "set1"
762
763
        crs2 = CosmicRaysBase(self.ncrs)
        crs2["log10e"] = np.random.rand(self.ncrs)
764
765
766
        crs2["name"] = "set2"
        crs_s = [crs1, crs2]
        crs_set = CosmicRaysSets(crs_s)
767
768
        self.assertTrue(crs_set.shape == (2, self.ncrs))
        self.assertTrue(np.all(crs_set[0]["log10e"] == crs1["log10e"]))
769
770
771
772
        self.assertTrue(np.all(crs_set[1]["log10e"] == crs2["log10e"]))
        self.assertTrue(len(crs_set["name"]) == 2)
        self.assertTrue(np.all([crs_set["name"][i] == "set%i" % (i+1) for i in range(2)]))

773
774
    def test_20b_create_from_crs_list_sets(self):
        crs1 = CosmicRaysSets(self.shape)
775
        crs1["nside"] = 64
776
        crs1["log10e"] = np.random.random(self.shape)
777
        crs1["vecs"] = coord.rand_vec(self.shape)
778
        crs2 = CosmicRaysSets(self.shape)
779
        crs2["nside"] = 64
780
        crs2["log10e"] = np.random.random(self.shape)
781
        crs2["vecs"] = coord.rand_vec(self.shape)
782
783
        crs_b_all = []
        for crs_s in [crs1, crs2]:
784
785
            for cr_b in crs_s:
                crs_b_all.append(cr_b)
786
787
        crs_merged_base = CosmicRaysSets(crs_b_all)
        self.assertTrue(crs_merged_base.shape == (2*self.nsets, self.ncrs))
788
789
        self.assertTrue(crs_merged_base["shape"] == (2*self.nsets, self.ncrs))
        self.assertTrue(crs_merged_base["vecs"].shape == (3, 2*self.nsets, self.ncrs))
790
791
792
        self.assertTrue(sorted(crs_merged_base.keys()) == sorted(crs1.keys()))
        self.assertTrue(np.all(crs_merged_base["log10e"][0] == crs1["log10e"][0]))
        self.assertTrue(np.all(crs_merged_base["vecs"][:, 0] == crs1["vecs"][:, 0]))
793
794
795

        crs_merged_set = CosmicRaysSets([crs1, crs2])
        self.assertTrue(crs_merged_set.shape == (2*self.nsets, self.ncrs))
796
797
        self.assertTrue(crs_merged_set["shape"] == (2*self.nsets, self.ncrs))
        self.assertTrue(crs_merged_set["vecs"].shape == (3, 2*self.nsets, self.ncrs))
798
799
800
        self.assertTrue(sorted(crs_merged_set.keys()) == sorted(crs1.keys()))
        self.assertTrue(np.allclose(crs_merged_set["log10e"], crs_merged_base["log10e"]))
        self.assertTrue(np.allclose(crs_merged_set["vecs"], crs_merged_base["vecs"]))
801
802

    def test_21a_create_from_crs_non_cosmic_rays_list(self):
803
804
        log10e1 = np.random.rand(self.ncrs)
        log10e2 = np.random.rand(self.ncrs)
805
806
        crs_s = [log10e1, log10e2]
        with self.assertRaises(TypeError):
807
            CosmicRaysSets(crs_s)
808

809
    def test_21b_create_from_crs_list_unequal_nr_crs(self):
810
811
812
813
814
815
816
817
818
        ncrs1, ncrs2 = 10, 11
        crs1 = CosmicRaysBase(ncrs1)
        crs1["log10e"] = np.random.rand(ncrs1)
        crs1["name"] = "set1"
        crs2 = CosmicRaysBase(ncrs2)
        crs2["log10e"] = np.random.rand(ncrs2)
        crs2["name"] = "set2"
        crs_s = [crs1, crs2]
        with self.assertRaises(ValueError):
819
820
            CosmicRaysSets(crs_s)

821
    def test_21c_create_from_crs_list_unequal_attributes(self):
822
823
        crs1 = CosmicRaysBase(self.ncrs)
        crs1["log10e"] = np.random.rand(self.ncrs)
824
        crs1["name"] = "set1"
825
826
        crs2 = CosmicRaysBase(self.ncrs)
        crs2["log10e"] = np.random.rand(self.ncrs)
827
828
829
830
        crs2["name"] = "set2"
        crs1["not_in_1"] = "test"
        crs_s = [crs1, crs2]
        with self.assertRaises(AttributeError):
831
            CosmicRaysSets(crs_s)