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

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

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

14

15
class TestObservedBound(unittest.TestCase):
16

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

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

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

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

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

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

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

66
67
        sim4 = ObservedBound(self.nside, self.nsets, self.ncrs)
        sim4.set_charges({'h': 0.5, 'he': 0.25, 'n': 0.24, 'si': 0.01})
68
69
70
71
        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)
72

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

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

    def test_07_smear_sources_dynamically(self):
90
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
91
        sim.set_energy(log10e_min=19.)
92
        sim.set_charges('AUGER')
marcus's avatar
marcus committed
93
        sim.set_sources(1)
94
        sim.set_rigidity_bins(np.arange(17., 20.5, 0.02))
95
        sim.smear_sources(delta=0.1, dynamic=True)
96
        sim.arrival_setup(1.)
marcus's avatar
marcus committed
97
98
99
100
101
102
103
        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)))
104
105

    def test_08_isotropy(self):
106
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
marcus's avatar
marcus committed
107
108
        sim.arrival_setup(0.)
        crs = sim.get_data(convert_all=True)
109
110
111
        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
112
113
114
        self.assertTrue((x < 0.03) & (y < 0.03) & (z < 0.03))

    def test_09_exposure(self):
115
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
116
117
        sim.apply_exposure()
        sim.arrival_setup(0.)
marcus's avatar
marcus committed
118
119
120
121
122
123
124
        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):
125
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
marcus's avatar
marcus committed
126
127
        charge = 2
        sim.set_charges(charge)
128
        self.assertTrue(np.all(sim.crs['charge'] == charge))
marcus's avatar
marcus committed
129
130

    def test_11_xmax_setup(self):
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
        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
156
157

    def test_12_xmax_mass(self):
158
159
        sim1 = ObservedBound(self.nside, self.nsets, self.ncrs)
        sim2 = ObservedBound(self.nside, self.nsets, self.ncrs)
marcus's avatar
marcus committed
160
161
        sim1.set_energy(19.)
        sim2.set_energy(19.)
162
        sim1.set_charges('equal')
marcus's avatar
marcus committed
163
164
165
166
167
168
169
        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):
170
171
172
173
        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
174
175
176
177
178
179
        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']))
180

181
    def test_14_lensing_map(self):
182
183
        lens_path = os.path.dirname(os.path.realpath(__file__)) + '/toy-lens/jf12-regular.cfg'
        toy_lens = gamale.Lens(lens_path)
184
        nside = toy_lens.nside
185
186
        sim = ObservedBound(nside, self.nsets, self.ncrs)
        sim.set_energy(19.*np.ones(self.shape))
187
188
189
190
191
192
193
194
195
196
        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))

197
198
    def test_15_exposure(self):
        nsets = 100
199
        sim = ObservedBound(self.nside, nsets, self.ncrs)
200
201
202
203
204
205
206
207
        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())

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

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

225
226
227
    def test_18_energy_spectra(self):
        nsets = 100
        sim = ObservedBound(self.nside, nsets, self.ncrs)
228
        log10e_power_3 = sim.set_energy(log10e_min=19., log10e_max=21., energy_spectrum='power_law', gamma=-3)
229
230
231
232
233
234
        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)
235
        log10e_power_4 = sim2.set_energy(log10e_min=19., log10e_max=21., energy_spectrum='power_law', gamma=-4)
236
237
238
        # higher energies for flatter spectrum
        self.assertTrue(np.mean(log10e_power_3) > np.mean(log10e_power_4))
        sim3 = ObservedBound(self.nside, nsets, self.ncrs)
239
        log10e_auger_fit = sim3.set_energy(log10e_min=19., log10e_max=21., energy_spectrum='auger_fit')
240
241
242
        self.assertTrue(np.mean(log10e_power_3) > np.mean(log10e_auger_fit))
        self.assertTrue(np.mean(log10e_power_4) < np.mean(log10e_auger_fit))

243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
    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())

263
    def test_20_exposure_issue(self):
264
        sim = ObservedBound(nside=4, nsets=nsets, ncrs=ncrs)
265
266
        sim.apply_exposure(a0=-35.25, zmax=60)
        sim.arrival_setup(0.)
267
        crs = sim.get_data(convert_all=True)
268
269
270
        _, 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)

271
272
273
274
275
276
277
278
279
280
281
282
283
    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)
284
285
286
287
        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))
288
289
290
291
292
293
294
295

        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)

296
297
    def test_22_shuffle(self):
        sim = ObservedBound(self.nside, self.nsets, self.ncrs)
298
        src_vecs = np.array([1, 1, 1])
299
300
        src_pix = hpt.vec2pix(sim.nside, src_vecs)
        sim.set_sources(src_vecs[:, np.newaxis])
301

302
        sim.arrival_setup(fsig=0.1)
303
        crs = sim.get_data(convert_all=False, shuffle=True)
304
        self.assertTrue(np.all(src_pix == crs['pixel'][sim.signal_label]))
305
        self.assertTrue(np.all(src_pix == crs['pixel'][crs['signal_label'].astype(bool)]))
306

307

308
class TestSourceBound(unittest.TestCase):
309

310
    def setUp(self):
311
        self.nsets = 120
312
313
314
315
316
317
318
319
320
321
322
323
        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
324
325
        self.assertTrue((sim.universe.rmax > 20) & (sim.universe.rmax < 40))
        dmin = np.min(sim.universe.distances, axis=-1)
326
        self.assertTrue((np.median(dmin) > 3) & (np.median(dmin) < 10))
327
328
329

        sim = SourceBound(self.nsets, self.ncrs)
        sim.set_sources(source_density=1)
Marcus Wirtz's avatar
Marcus Wirtz committed
330
        dmin = np.min(sim.universe.distances, axis=-1)
331
332
333
334
        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
335
        dmin = np.min(sim.universe.distances, axis=-1)
336
337
        self.assertTrue((np.median(dmin) > 20))

338
339
340
341
342
343
344
345
346
347
348
        densities = np.array([1e-6, 1e-3, 1])
        n_each = 100
        sim = SourceBound(int(n_each * len(densities)), self.ncrs)
        sim.set_sources(source_density=np.repeat(densities, n_each))
        dmin = np.min(sim.universe.distances[:100], axis=-1)
        self.assertTrue((np.median(dmin) > 20))
        dmin = np.min(sim.universe.distances[100:200], axis=-1)
        self.assertTrue((np.median(dmin) > 3) & (np.median(dmin) < 10))
        dmin = np.min(sim.universe.distances[200:], axis=-1)
        self.assertTrue((np.median(dmin) < 1))

349
    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
350
    def test_03a_fluxes(self):
351
        sim = SourceBound(self.nsets, self.ncrs)
352
        sim.set_energy(gamma=-2, log10e_min=19.6, log10_cut=20.5, rig_cut=False)
353
        sim.set_charges(charges={'h': 1.})
354
355
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
356
        sim.smear_sources(np.deg2rad(3))
357
358
359
360
361
362
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()

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

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

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

        sim = SourceBound(self.nsets, self.ncrs)
393
394
395
396
397
398
399
400
401
402
403
404
405
406
        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)
407
        sim.set_charges(charges={'fe': 1.})
408
409
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
410
        sim.smear_sources(np.deg2rad(3))
411
412
413
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()
414
        mask_inside_10 = sim.crs['distances'] <= 30
415
        fraction_inside = np.sum(mask_inside_10) / float(sim.ncrs * sim.nsets)
416
        self.assertTrue(fraction_inside < 0.2)
417

418
    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
419
    def test_03b_fluxes_multiple_densities(self):
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
        densities = [1e-5, 1e-3, 1e-1]
        bound_source = [[120, 180], [30, 60], [3, 15]]

        for i, d in enumerate(densities):
            # make a separate simulation for each source density
            sim = SourceBound(self.nsets, 1000)
            sim.set_energy(log10e_min=19.6)
            sim.set_charges('first_minimum_CTG')
            sim.set_sources(source_density=d)
            sim.attenuate()
            # check if number of events from the strongest source of the representative set is within the expected range
            idx_select = sim._select_representative_set()
            ns = np.array([np.sum(sim.crs['source_labels'][idx_select] == k) for k in range(sim.universe.n_src)]).max()
            self.assertTrue((ns >= bound_source[i][0]) & (ns <= bound_source[i][1]))

        # same check as above but with all source densities simulated simultaneously
        n_each = 100
        sim = SourceBound(int(n_each * len(densities)), 1000)
        sim.set_energy(log10e_min=19.6)
        sim.set_charges('first_minimum_CTG')
        sim.set_sources(source_density=np.repeat(densities, n_each))
        sim.attenuate()
        for i in range(len(densities)):
            mask = (np.arange(sim.nsets) < (i+1)*n_each) & (np.arange(sim.nsets) >= i*n_each)
            idx_select = sim._select_representative_set(mask=mask)
            ns = np.array([np.sum(sim.crs['source_labels'][idx_select] == k) for k in range(sim.universe.n_src)]).max()
            self.assertTrue((ns >= bound_source[i][0]) & (ns <= bound_source[i][1]))

448
449
450
451
452
453
454
455
456
457
458
    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
    def test_03c_fluxes_no_sources(self):
        sim = SourceBound(self.nsets, self.ncrs)
        sim.set_energy(log10e_min=19.6)
        sim.set_charges('first_minimum')
        sim.set_sources()
        sim.attenuate()
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()

459
    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
460
    def test_04a_shuffle(self):
461
462
463
464
465
466
467
468
469
        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))

470
    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
471
472
473
474
475
476
477
478
479
480
481
482
    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))

483
    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
484
485
    def test_05a_first_minimum(self):
        sim = SourceBound(self.nsets, 1000)
486
        sim.set_energy(log10e_min=19.6)
487
        sim.set_charges('first_minimum_CTG')
488
489
490
491
492
493
494
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
        sim.smear_sources(np.deg2rad(3))
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()

495
496
497
498
        sim = SourceBound(self.nsets, 1000)
        sim.set_energy(log10e_min=19.6)
        sim.set_charges('first_minimum_CTD')
        sim.set_sources(source_density=1e-3)
499
        sim.attenuate(library_path=PATH + '/simulation/crpropa3__emin_18.5__emax_21.0__IRB_Dominguez11.npz')
500
        sim.smear_sources(np.deg2rad(3))
501
502
503
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()
504
505

    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
506
507
    def test_05b_second_minimum(self):
        sim = SourceBound(self.nsets, 1000)
508
509
        # 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})
510
        sim.set_energy(log10e_min=19.6)
511
        sim.set_charges('second_minimum_CTG')
512
513
514
        sim.set_sources(source_density=1e-3)
        sim.attenuate()
        sim.smear_sources(np.deg2rad(3))
Marcus Wirtz's avatar
Marcus Wirtz committed
515
516
517
        # sim.plot_arrivals()
        # sim.plot_spectrum()
        # sim.plot_distance()
518

519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
    def test_6a_source_density_from_samples(self):
        n_src = 10000
        distances = 50 * np.random.random((1, n_src))**(1/3)
        sources = coord.rand_fisher_vec(np.array([1, 0, 0]), kappa=1/np.deg2rad(30)**2, n=n_src) * distances
        sim = SourceBound(self.nsets, 1000)
        sim.set_energy(log10e_min=19.6)
        sim.set_charges('first_minimum')
        sim.set_sources(source_density=1e-4, sources=sources, n_src=300)
        sim.attenuate()
        sim.smear_sources(np.deg2rad(3))
        sim.plot_arrivals(opath='source_hybrid_1e-4.png')
        # sim.plot_spectrum()
        # sim.plot_distance()
        vecs_sum = np.mean(sim.crs['vecs'], axis=(-2, -1))
        self.assertTrue(coord.angle(vecs_sum, np.array([1, 0, 0]))[0] < 0.5)
        amplitude = 3 * np.sqrt(np.sum(vecs_sum**2))
        self.assertTrue(amplitude > 1.)

    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
    def test_6b_source_density_from_samples(self):
        densities = [1e-4, 1e-3, 1e-2]
        bound_amplitude = [1., 0.5, 0.1]

        # same check as above but with all source densities simulated simultaneously
        n_each = 100
        ncrs = 1000
        sim = SourceBound(int(n_each * len(densities)), ncrs)
        sim.set_energy(log10e_min=19.6)
        sim.set_charges('first_minimum')
        n_src = 10000
        distances = 50 * np.random.random((1, n_src))**(1/3)
        sources = coord.rand_fisher_vec(np.array([1, 0, 0]), kappa=1/np.deg2rad(30)**2, n=n_src) * distances
        sim.set_sources(source_density=np.repeat(densities, n_each), sources=sources, n_src=300)
        sim.attenuate()
        sim.smear_sources(np.deg2rad(3))
        for i in range(len(densities)):
            mask = (np.arange(sim.nsets) < (i+1)*n_each) & (np.arange(sim.nsets) >= i*n_each)
            sim.plot_arrivals(idx=sim._select_representative_set(mask), opath='source_hybrid_sd_%s.png' % densities[i])
            vecs = sim.crs['vecs'][:, mask]
            self.assertTrue(vecs.shape == (3, n_each, ncrs))
            vecs_sum = np.mean(vecs, axis=(-2, -1))
            self.assertTrue(coord.angle(vecs_sum, np.array([1, 0, 0]))[0] < 0.5)
            amplitude = 3 * np.sqrt(np.sum(vecs_sum**2))
            self.assertTrue(amplitude > bound_amplitude[i])


class TestSourceGeometry(unittest.TestCase):

    def setUp(self):
        self.nsets = 50
        self.n_src = 100

    def test_01_no_soures(self):
        universe = SourceGeometry(self.nsets)
        universe.set_sources()
        self.assertTrue((universe.n_src == 0) and (universe.rmax == 0))

        universe2 = SourceGeometry(self.nsets)
        universe2.set_sources(source_density=1e-3, n_src=0)
        self.assertTrue((universe2.n_src == 0) and (universe2.rmax == 0))

    def test_02_source_density(self):
        universe = SourceGeometry(self.nsets)
        universe.set_sources(source_density=1e-3, n_src=self.n_src)
        self.assertTrue(np.abs(universe.rmax - 28.8) < 0.1)

        universe2 = SourceGeometry(self.nsets)
        universe2.set_sources(source_density=1e-3, rmax=universe.rmax)
        self.assertTrue((universe2.n_src == self.n_src) and (universe.rmax == universe2.rmax))

        sources = coord.rand_vec(1000) * 30 * np.random.random((1, 1000))**(1/3)
        distances = np.sqrt(np.sum(sources**2, axis=0))
        universe = SourceGeometry(self.nsets)
        universe.set_sources(source_density=1e-3, sources=sources, n_src=self.n_src)
        self.assertTrue(np.abs(universe.rmax - 28.8) < 0.1)
        self.assertTrue(np.all(universe.distances <= universe.rmax))
        self.assertTrue(np.shape(universe.distances) == (self.nsets, self.n_src))
        self.assertTrue(not np.allclose(universe.distances[0], universe.distances[1]))
        for dis in universe.distances:
            self.assertTrue(np.isin(dis, distances).all())

        universe2 = SourceGeometry(self.nsets)
        universe2.set_sources(source_density=1e-3, rmax=universe.rmax)
        self.assertTrue((universe2.n_src == self.n_src) and (universe.rmax == universe2.rmax))

        universe = SourceGeometry(self.nsets)
        with self.assertRaises(AssertionError):
            universe.set_sources(source_density=1, sources=sources)
        with self.assertRaises(AssertionError):
            universe.set_sources(source_density=1, sources=sources, n_src=100)

    def test_03_sources(self):
        sources = coord.rand_vec(1000) * 30 * np.random.random((1, 1000))**(1/3)
        universe = SourceGeometry(self.nsets)
        universe.set_sources(sources=sources)
        self.assertTrue(np.allclose(universe.sources, sources))

marcus's avatar
marcus committed
617

618
619
class TestReweight(unittest.TestCase):

620
    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
621
622
    def setUp(self):
        self.charge = {'h': 1, 'he': 2, 'n': 7, 'si': 14, 'fe': 26}
623
        data = np.load(PATH+'/simulation/crpropa3__emin_18.5__emax_21.0__IRB_Gilmore12.npz', allow_pickle=True)
624
625
626
627
        self.fractions = data['fractions'].item()
        self.log10e_bins = data['log10e_bins']
        self.distances = data['distances']

628
    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
629
630
631
632
633
634
    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))

635
    @unittest.skipIf(sys.version_info < (3, 0), "Simulation libraries work only for python 3.")
636
637
638
639
640
641
    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.)

642

643
644
if __name__ == '__main__':
    unittest.main()