test_simulations.py 19.8 KB
Newer Older
1
import unittest
2
import os
3
import numpy as np
4

5
from astrotools import coord, gamale, healpytools as hpt
6
from astrotools.simulations import PATH, ObservedBound, SourceBound
7
8
9

nside = 64
ncrs = 1000
MarCus's avatar
MarCus committed
10
nsets = 10
11
np.random.seed(0)
MarCus's avatar
MarCus committed
12

13

14
class TestObservedBound(unittest.TestCase):
15

16
17
18
19
20
21
    def setUp(self):
        self.nside = 64
        self.nsets = 10
        self.ncrs = 1000
        self.shape = (self.nsets, self.ncrs)

22
    def test_01_n_cosmic_rays(self):
23
24
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
        self.assertEqual(sim.ncrs, self.ncrs)
25

26
    def test_02_nsets(self):
27
28
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
        self.assertEqual(sim.nsets, self.nsets)
29
30

    def test_03_keyword_setup(self):
31
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
32
        sim.set_energy(log10e_min=19.)
33
34
        sim.set_charges(charge='mixed')
        sim.set_xmax('double')
35
        sim.set_sources(sources='sbg')
36
        sim.smear_sources(delta=0.1)
37
38
        sim.apply_exposure()
        sim.arrival_setup(fsig=0.4)
MarCus's avatar
MarCus committed
39
        crs = sim.get_data(convert_all=True)
40
        self.assertEqual(crs['pixel'].shape, crs['lon'].shape, crs['log10e'].shape)
41
42

    def test_04_set_energy_charge_arrays(self):
43
44
45
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
        log10e = np.random.random(size=self.shape)
        charge = np.random.randint(0, 10, size=self.shape)
46
        sim.set_energy(log10e_min=log10e)
47
        sim.set_charges(charge=charge)
48
49
        crs = sim.get_data()
        self.assertTrue(np.allclose(crs['log10e'], log10e) and np.allclose(crs['charge'], charge))
50

51
52
53
        sim2 = ObservedBound(self.nside, self.nsets, self.ncrs)
        log10e = np.random.random(self.ncrs)
        charge = np.random.random(self.ncrs)
54
55
56
57
        sim2.set_energy(log10e)
        sim2.set_charges(charge)
        self.assertTrue(np.allclose(sim2.crs['log10e'], log10e) and np.allclose(sim2.crs['charge'], charge))

58
59
60
        sim3 = ObservedBound(self.nside, self.nsets, self.ncrs)
        log10e = np.random.random(self.nsets)
        charge = np.random.random(self.nsets)
61
62
63
64
        with self.assertRaises(Exception):
            sim3.set_energy(log10e)
            sim3.set_charges(log10e)

65
66
        sim4 = ObservedBound(self.nside, self.nsets, self.ncrs)
        sim4.set_charges({'h': 0.5, 'he': 0.25, 'n': 0.24, 'si': 0.01})
67
68
69
70
        self.assertTrue(np.abs(np.sum(sim4.crs['charge'] == 1) / float(self.nsets * self.ncrs) - 0.5) < 0.02)
        self.assertTrue(np.abs(np.sum(sim4.crs['charge'] == 2) / float(self.nsets * self.ncrs) - 0.25) < 0.02)
        self.assertTrue(np.abs(np.sum(sim4.crs['charge'] == 7) / float(self.nsets * self.ncrs) - 0.25) < 0.02)
        self.assertTrue(np.abs(np.sum(sim4.crs['charge'] == 14) / float(self.nsets * self.ncrs) - 0.01) < 0.01)
71

72
73
    def test_05_set_n_random_sources(self):
        n = 5
Marcus Wirtz's avatar
Marcus Wirtz committed
74
        fluxes = np.random.random(n)
75
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
Marcus Wirtz's avatar
Marcus Wirtz committed
76
        sim.set_sources(n, fluxes=fluxes)
77
        self.assertTrue(sim.sources.shape[1] == n)
Marcus Wirtz's avatar
Marcus Wirtz committed
78
        self.assertTrue(np.allclose(fluxes, sim.source_fluxes))
79
80
81

    def test_06_set_n_sources(self):
        v_src = np.random.rand(30).reshape((3, 10))
Marcus Wirtz's avatar
Marcus Wirtz committed
82
        fluxes = np.random.random(10)
83
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
Marcus Wirtz's avatar
Marcus Wirtz committed
84
        sim.set_sources(v_src, fluxes=fluxes)
85
        self.assertTrue(np.allclose(v_src, sim.sources))
Marcus Wirtz's avatar
Marcus Wirtz committed
86
        self.assertTrue(np.allclose(fluxes, sim.source_fluxes))
87
88

    def test_07_smear_sources_dynamically(self):
89
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
90
        sim.set_energy(log10e_min=19.)
91
        sim.set_charges('AUGER')
marcus's avatar
marcus committed
92
        sim.set_sources(1)
93
        sim.set_rigidity_bins(np.arange(17., 20.5, 0.02))
94
        sim.smear_sources(delta=0.1, dynamic=True)
95
        sim.arrival_setup(1.)
marcus's avatar
marcus committed
96
97
98
99
100
101
102
        crs = sim.get_data(convert_all=True)
        rigs = sim.rigidities
        rig_med = np.median(rigs)
        vecs1 = coord.ang2vec(crs['lon'][rigs >= rig_med], crs['lat'][rigs >= rig_med])
        vecs2 = coord.ang2vec(crs['lon'][rigs < rig_med], crs['lat'][rigs < rig_med])
        # Higher rigidities experience higher deflections
        self.assertTrue(np.mean(coord.angle(vecs1, sim.sources)) < np.mean(coord.angle(vecs2, sim.sources)))
103
104

    def test_08_isotropy(self):
105
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
marcus's avatar
marcus committed
106
107
        sim.arrival_setup(0.)
        crs = sim.get_data(convert_all=True)
108
109
110
        x = np.abs(np.mean(crs['vecs'][0]))
        y = np.abs(np.mean(crs['vecs'][1]))
        z = np.abs(np.mean(crs['vecs'][2]))
marcus's avatar
marcus committed
111
112
113
        self.assertTrue((x < 0.03) & (y < 0.03) & (z < 0.03))

    def test_09_exposure(self):
114
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
115
116
        sim.apply_exposure()
        sim.arrival_setup(0.)
marcus's avatar
marcus committed
117
118
119
120
121
122
123
        crs = sim.get_data(convert_all=True)
        vecs_eq = coord.gal2eq(coord.ang2vec(np.hstack(crs['lon']), np.hstack(crs['lat'])))
        lon_eq, lat_eq = coord.vec2ang(vecs_eq)
        self.assertTrue(np.abs(np.mean(lon_eq)) < 0.05)
        self.assertTrue((np.mean(lat_eq) < -0.5) & (np.mean(lat_eq) > - 0.55))

    def test_10_charge(self):
124
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
marcus's avatar
marcus committed
125
126
        charge = 2
        sim.set_charges(charge)
127
        self.assertTrue(np.all(sim.crs['charge'] == charge))
marcus's avatar
marcus committed
128
129

    def test_11_xmax_setup(self):
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
        sim1 = ObservedBound(self.nside, self.nsets, self.ncrs)
        sim1.set_energy(19.)
        sim1.set_charges(2)
        sim1.set_xmax('stable')
        check1 = (sim1.crs['xmax'] > 500) & (sim1.crs['xmax'] < 1200)

        sim2 = ObservedBound(self.nside, self.nsets, self.ncrs)
        sim2.set_energy(19.)
        sim2.set_charges(2)
        sim2.set_xmax('stable')
        check2 = (sim2.crs['xmax'] > 500) & (sim2.crs['xmax'] < 1200)

        sim3 = ObservedBound(self.nside, self.nsets, self.ncrs)
        sim3.set_energy(19.)
        sim3.set_charges(2)
        sim3.set_xmax('empiric')
        check3 = (sim3.crs['xmax'] > 500) & (sim3.crs['xmax'] < 1200)

        sim4 = ObservedBound(self.nside, self.nsets, self.ncrs)
        sim4.set_energy(19.)
        sim4.set_charges(2)
        sim4.set_xmax('double')
        check4 = (sim4.crs['xmax'] > 500) & (sim4.crs['xmax'] < 1200)

        self.assertTrue(check1.all() and check2.all() and check3.all() and check4.all())
marcus's avatar
marcus committed
155
156

    def test_12_xmax_mass(self):
157
158
        sim1 = ObservedBound(self.nside, self.nsets, self.ncrs)
        sim2 = ObservedBound(self.nside, self.nsets, self.ncrs)
marcus's avatar
marcus committed
159
160
        sim1.set_energy(19.)
        sim2.set_energy(19.)
161
        sim1.set_charges('equal')
marcus's avatar
marcus committed
162
163
164
165
166
167
168
        sim2.set_charges(26)
        sim1.set_xmax('double')
        sim2.set_xmax('double')
        # Xmax of iron should be smaller (interact higher in atmosphere)
        self.assertTrue(np.mean(sim1.crs['xmax']) > np.mean(sim2.crs['xmax']))

    def test_13_xmax_energy(self):
169
170
171
172
        sim1 = ObservedBound(self.nside, self.nsets, self.ncrs)
        sim2 = ObservedBound(self.nside, self.nsets, self.ncrs)
        sim1.set_energy(20. * np.ones(self.shape))
        sim2.set_energy(19. * np.ones(self.shape))
marcus's avatar
marcus committed
173
174
175
176
177
178
        sim1.set_charges(1)
        sim2.set_charges(1)
        sim1.set_xmax('double')
        sim2.set_xmax('double')
        # Xmax for higher energy is bigger
        self.assertTrue(np.mean(sim1.crs['xmax']) > np.mean(sim2.crs['xmax']))
179

180
    def test_14_lensing_map(self):
181
182
        lens_path = os.path.dirname(os.path.realpath(__file__)) + '/toy-lens/jf12-regular.cfg'
        toy_lens = gamale.Lens(lens_path)
183
        nside = toy_lens.nside
184
185
        sim = ObservedBound(nside, self.nsets, self.ncrs)
        sim.set_energy(19.*np.ones(self.shape))
186
187
188
189
190
191
192
193
194
195
        sim.set_charges(1)
        sim.set_xmax('empiric')
        sim.set_sources('gamma_agn')
        sim.lensing_map(toy_lens)
        # Xmax for higher energy is bigger
        self.assertTrue(sim.lensed)
        self.assertTrue(np.shape(sim.cr_map) == (1, sim.npix))
        self.assertAlmostEqual(np.sum(sim.cr_map), 1.)
        self.assertTrue(np.min(sim.cr_map) < np.max(sim.cr_map))

196
197
    def test_15_exposure(self):
        nsets = 100
198
        sim = ObservedBound(self.nside, nsets, self.ncrs)
199
200
201
202
203
204
205
206
        sim.apply_exposure(a0=-35.25, zmax=60)
        sim.arrival_setup(0.2)
        crs = sim.get_data(convert_all=True)
        lon, lat = np.hstack(crs['lon']), np.hstack(crs['lat'])
        ra, dec = coord.vec2ang(coord.gal2eq(coord.ang2vec(lon, lat)))
        exp = coord.exposure_equatorial(dec, a0=-35.25, zmax=60)
        self.assertTrue((exp > 0).all())

207
208
    def test_16_high_nside(self):
        nside_high = 256
209
        sim = ObservedBound(nside_high, nsets=self.nsets, ncrs=self.ncrs)
210
        self.assertTrue(sim.crs['nside'] == nside_high)
211
        self.assertTrue(sim.crs.nside() == nside_high)
212
213
214
        sim.arrival_setup(1.)
        self.assertTrue(np.max(sim.crs['pixel']) >= 0.8 * sim.npix)

215
216
217
    def test_17_energy_rigidity_set(self):
        ncrs = 100
        e = 19.5
218
219
        sim = ObservedBound(self.nside, self.nsets, ncrs)
        sim.set_energy(e * np.ones((self.nsets, ncrs)))
220
221
222
223
        sim.set_charges(2)
        sim.set_rigidity_bins(np.linspace(17., 20.48, 175))
        self.assertTrue((sim.crs['log10e'] == e).all())

224
225
226
    def test_18_energy_spectra(self):
        nsets = 100
        sim = ObservedBound(self.nside, nsets, self.ncrs)
227
        log10e_power_3 = sim.set_energy(log10e_min=19., log10e_max=21., energy_spectrum='power_law', gamma=-3)
228
229
230
231
232
233
        ebin = 0.1
        for e in np.arange(19.1, 20.1, ebin):
            sum_low = np.sum((log10e_power_3 >= e-ebin) & (log10e_power_3 < e))
            sum_high = np.sum((log10e_power_3 >= e) & (log10e_power_3 < e+ebin))
            self.assertTrue(sum_low > sum_high)
        sim2 = ObservedBound(self.nside, nsets, self.ncrs)
234
        log10e_power_4 = sim2.set_energy(log10e_min=19., log10e_max=21., energy_spectrum='power_law', gamma=-4)
235
236
237
        # higher energies for flatter spectrum
        self.assertTrue(np.mean(log10e_power_3) > np.mean(log10e_power_4))
        sim3 = ObservedBound(self.nside, nsets, self.ncrs)
238
        log10e_auger_fit = sim3.set_energy(log10e_min=19., log10e_max=21., energy_spectrum='auger_fit')
239
240
241
        self.assertTrue(np.mean(log10e_power_3) > np.mean(log10e_auger_fit))
        self.assertTrue(np.mean(log10e_power_4) < np.mean(log10e_auger_fit))

242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
    def test_19_apply_uncertainties(self):
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
        log10e = sim.set_energy(log10e_min=19., log10e_max=21., energy_spectrum='power_law', gamma=-3)
        sim.set_charges('mixed')
        xmax = sim.set_xmax()
        sim.set_sources(10)
        sim.set_rigidity_bins(np.arange(17., 20.5, 0.02))
        sim.smear_sources(delta=0.1, dynamic=True)
        sim.arrival_setup(1.)
        vecs = hpt.pix2vec(sim.nside, np.hstack(sim.crs['pixel']))
        sim.apply_uncertainties(err_e=0.1, err_a=1, err_xmax=10)
        # check that array are not equal but deviations are smaller than 5 sigma
        self.assertTrue(not (log10e == sim.crs['log10e']).all())
        self.assertTrue((np.abs(10**(log10e - 18) - 10**(sim.crs['log10e'] - 18)) < 5*0.1*10**(log10e - 18)).all())
        self.assertTrue(not (xmax == sim.crs['xmax']).all())
        self.assertTrue((np.abs(xmax - sim.crs['xmax']) < 50).all())
        vec_unc = coord.ang2vec(np.hstack(sim.crs['lon']), np.hstack(sim.crs['lat']))
        self.assertTrue(not (coord.angle(vecs, vec_unc) == 0).all())
        self.assertTrue((coord.angle(vecs, vec_unc) < np.deg2rad(10)).all())

262
    def test_20_exposure_issue(self):
263
        sim = ObservedBound(nside=4, nsets=nsets, ncrs=ncrs)
264
265
        sim.apply_exposure(a0=-35.25, zmax=60)
        sim.arrival_setup(0.)
266
        crs = sim.get_data(convert_all=True)
267
268
269
        _, dec = coord.vec2ang(coord.gal2eq(crs['vecs'].reshape(3, -1)))
        self.assertTrue(np.sum(coord.exposure_equatorial(dec, a0=-35.25, zmax=60) <= 0) == 0)

270
271
272
273
274
275
276
277
278
279
280
281
282
    def test_21_convert_all(self):
        sim = ObservedBound(nside=4, nsets=nsets, ncrs=ncrs)
        sim.arrival_setup(0.)
        crs = sim.get_data(convert_all=False)
        keys = crs.keys()
        self.assertTrue('vecs' not in keys and 'lon' not in keys and 'lat' not in keys)

        vecs = crs['vecs']  # automatic from pixel center
        sim.convert_pixel('angles')  # converts pixel to lon / lat
        crs = sim.get_data(convert_all=False)
        keys = crs.keys()
        self.assertTrue('vecs' not in keys and 'lon' in keys and 'lat' in keys)
        _lon, _lat = coord.vec2ang(vecs)
283
284
285
286
        self.assertTrue(np.mean(abs(crs['lon'] - _lon) < 0.5))
        self.assertTrue(np.mean(abs(crs['lat'] - _lat) < 0.5))
        self.assertTrue(np.mean(abs(crs['lon'] - _lon) > 0))
        self.assertTrue(np.mean(abs(crs['lat'] - _lat) > 0))
287
288
289
290
291
292
293
294

        sim = ObservedBound(nside=4, nsets=nsets, ncrs=ncrs)
        sim.apply_exposure(a0=-35.25, zmax=60)
        sim.arrival_setup(0.)
        crs = sim.get_data(convert_all=True)
        keys = crs.keys()
        self.assertTrue('vecs' in keys and 'lon' in keys and 'lat' in keys)

295
296
297
298
299
300
301
302
303
304
    def test_22_shuffle(self):
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
        src_vecs = np.array([0, 0, 1])
        src_pix = hpt.vec2pix(sim.nside, src_vecs)
        sim.set_sources(src_vecs[:, np.newaxis])
        sim.arrival_setup(fsig=0.1)
        crs = sim.get_data(convert_all=True, shuffle=True)
        self.assertTrue(np.all(src_pix == crs['pixel'][sim.signal_label]))
        self.assertTrue(np.all(src_pix != crs['pixel'][~sim.signal_label]))

305

306
class TestSourceBound(unittest.TestCase):
307

308
    def setUp(self):
309
        self.nsets = 120
310
311
312
313
314
315
316
317
318
319
320
321
        self.ncrs = 1000
        self.shape = (self.nsets, self.ncrs)

    def test_01_init(self):
        sim = SourceBound(self.nsets, self.ncrs)
        self.assertEqual(sim.ncrs, self.ncrs)
        self.assertEqual(sim.nsets, self.nsets)
        self.assertEqual(sim.shape, (self.nsets, self.ncrs))

    def test_02_source_distance(self):
        sim = SourceBound(self.nsets, self.ncrs)
        sim.set_sources(source_density=1e-3)
Marcus Wirtz's avatar
Marcus Wirtz committed
322
323
        self.assertTrue((sim.universe.rmax > 20) & (sim.universe.rmax < 40))
        dmin = np.min(sim.universe.distances, axis=-1)
324
        self.assertTrue((np.median(dmin) > 3) & (np.median(dmin) < 10))
325
326
327

        sim = SourceBound(self.nsets, self.ncrs)
        sim.set_sources(source_density=1)
Marcus Wirtz's avatar
Marcus Wirtz committed
328
        dmin = np.min(sim.universe.distances, axis=-1)
329
330
331
332
        self.assertTrue((np.median(dmin) < 1))

        sim = SourceBound(self.nsets, self.ncrs)
        sim.set_sources(source_density=1e-6)
Marcus Wirtz's avatar
Marcus Wirtz committed
333
        dmin = np.min(sim.universe.distances, axis=-1)
334
335
336
337
        self.assertTrue((np.median(dmin) > 20))

    def test_03_fluxes(self):
        sim = SourceBound(self.nsets, self.ncrs)
338
        sim.set_energy(gamma=-2, log10e_min=19.6, log10_cut=20.5, rig_cut=False)
339
        sim.set_charges(charges={'h': 1.})
340
341
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
342
        sim.smear_sources(np.deg2rad(3))
343
344
345
346
347
348
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()

        # protons propagate very far: barely protons within rmax
        mask_inside_10 = sim.crs['distances'] <= 30
349
        fraction_inside = np.sum(mask_inside_10) / float(sim.ncrs * sim.nsets)
350
351
        self.assertTrue(fraction_inside < 0.15)

352
353
        sim = SourceBound(self.nsets, self.ncrs)
        sim.set_energy(gamma=-2, log10e_min=19.6, log10_cut=20.5, rig_cut=False)
354
        sim.set_charges(charges={'he': 1.})
355
356
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
357
        sim.smear_sources(np.deg2rad(3))
358
359
360
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()
361
        mask_inside_10 = sim.crs['distances'] <= 30
362
        fraction_inside = np.sum(mask_inside_10) / float(sim.ncrs * sim.nsets)
363
        self.assertTrue(fraction_inside < 0.25)
364
365

        sim = SourceBound(self.nsets, self.ncrs)
366
        sim.set_energy(gamma=-2, log10e_min=19.6, log10_cut=20.5, rig_cut=False)
367
        sim.set_charges(charges={'n': 1.})
368
369
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
370
        sim.smear_sources(np.deg2rad(3))
371
372
373
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()
374
        mask_inside_10 = sim.crs['distances'] <= 30
375
376
        fraction_inside = np.sum(mask_inside_10) / float(sim.ncrs * sim.nsets)
        self.assertTrue(fraction_inside > 0.7)
377
378

        sim = SourceBound(self.nsets, self.ncrs)
379
380
381
382
383
384
385
386
387
388
389
390
391
392
        sim.set_energy(gamma=-2, log10e_min=19.6, log10_cut=20.5, rig_cut=False)
        sim.set_charges(charges={'si': 1.})
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
        sim.smear_sources(np.deg2rad(3))
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()
        mask_inside_10 = sim.crs['distances'] <= 30
        fraction_inside = np.sum(mask_inside_10) / float(sim.ncrs * sim.nsets)
        self.assertTrue(fraction_inside < 0.4)

        sim = SourceBound(self.nsets, self.ncrs)
        sim.set_energy(gamma=-2, log10e_min=19.6, log10_cut=20.5, rig_cut=False)
393
        sim.set_charges(charges={'fe': 1.})
394
395
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
396
        sim.smear_sources(np.deg2rad(3))
397
398
399
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()
400
        mask_inside_10 = sim.crs['distances'] <= 30
401
        fraction_inside = np.sum(mask_inside_10) / float(sim.ncrs * sim.nsets)
402
        self.assertTrue(fraction_inside < 0.2)
403

404
    def test_04a_shuffle(self):
405
406
407
408
409
410
411
412
413
        sim = SourceBound(self.nsets, self.ncrs)
        sim.set_energy(gamma=-2, log10e_min=19.6)
        sim.set_charges(charges={'h': 1.})
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
        crs = sim.get_data(shuffle=True)
        src_labels = crs['source_labels']
        self.assertTrue(np.all(src_labels[sim.signal_label] != -1))

414
415
416
417
418
419
420
421
422
423
424
425
    def test_04b_shuffle(self):
        sim = SourceBound(self.nsets, self.ncrs)
        sim.set_energy(gamma=-2, log10e_min=19.6)
        sim.set_charges(charges={'h': 1.})
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
        crs = sim.get_data(shuffle=False)
        test1 = np.unique(crs['vecs'] * crs['log10e'])
        crs = sim.get_data(shuffle=True)
        test2 = np.unique(crs['vecs'] * crs['log10e'])
        self.assertTrue(np.all(test1 == test2))

426
427
    def test_05a_first_minimum(self):
        sim = SourceBound(self.nsets, 1000)
428
429
        # sim.set_energy(gamma=-0.96, log10e_min=19.6, log10_cut=18.68, rig_cut=True)
        # sim.set_charges(charges={'he': 0.673, 'n': 0.281, 'si': 0.046})
430
431
        sim.set_energy(log10e_min=19.6)
        sim.set_charges('first_minimum')
432
433
434
435
436
437
438
439
440
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
        sim.smear_sources(np.deg2rad(3))
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()

    def test_05b_second_minimum(self):
        sim = SourceBound(self.nsets, 1000)
441
442
        # sim.set_energy(gamma=-2.04, log10e_min=19.6, log10_cut=19.88, rig_cut=True)
        # sim.set_charges(charges={'n': 0.798, 'si': 0.202})
443
        sim.set_energy(log10e_min=19.6)
444
        sim.set_charges('second_minimum')
445
446
447
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
        sim.smear_sources(np.deg2rad(3))
Marcus Wirtz's avatar
Marcus Wirtz committed
448
449
450
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()
451

marcus's avatar
marcus committed
452

453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
class TestReweight(unittest.TestCase):

    def setUp(self):
        self.charge = {'h': 1, 'he': 2, 'n': 7, 'si': 14, 'fe': 26}
        data = np.load(PATH+'/simulation/crpropa3__emin_18.5__emax_21.npz', allow_pickle=True)
        self.fractions = data['fractions'].item()
        self.log10e_bins = data['log10e_bins']
        self.distances = data['distances']

    def test_01_dimensions(self):
        ne = self.log10e_bins.size - 1
        nd = self.distances.size
        for key in self.charge:
            self.assertTrue(self.fractions[key].shape == (ne, nd, 5, ne))

    def test_02_no_energy_gain(self):
        for key in self.charge:
            for i, lge in enumerate(self.log10e_bins[:-1]):
                frac = self.fractions[key]
                self.assertTrue(np.sum(frac[i][:, :, self.log10e_bins[:-1] > lge]) == 0.)

474
475
if __name__ == '__main__':
    unittest.main()