/
test_energy.py
173 lines (142 loc) · 4.76 KB
/
test_energy.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
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
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
# Licensed under a 3-clause BSD style license - see LICENSE.rst
from __future__ import absolute_import, division, print_function, unicode_literals
import astropy.units as u
from numpy.testing import assert_equal
from ...spectrum import Energy, EnergyBounds
from astropy.tests.helper import remote_data
from ...datasets import get_path
from astropy.io import fits
import numpy as np
def test_Energy():
# Explicit constructor call
energy = Energy([1, 3, 6, 8, 12], 'TeV')
actual = type(energy).__module__
desired = 'gammapy.spectrum.energy'
assert_equal(actual, desired)
val = u.Quantity([1, 3, 6, 8, 12], 'TeV')
actual = Energy(val, 'GeV')
desired = Energy((1, 3, 6, 8, 12), 'TeV')
assert_equal(actual, desired)
# View casting
energy = val.view(Energy)
actual = type(energy).__module__
desired = 'gammapy.spectrum.energy'
assert_equal(actual, desired)
# New from template
energy = Energy([0, 1, 2, 3, 4, 5], 'eV')
energy2 = energy[1:3]
actual = energy2
desired = Energy([1, 2], 'eV')
assert_equal(actual, desired)
actual = energy2.nbins
desired = 2
assert_equal(actual, desired)
actual = energy2.unit
desired = u.eV
assert_equal(actual, desired)
# Equal log spacing
energy = Energy.equal_log_spacing(1 * u.GeV, 10 * u.TeV, 6)
actual = energy[0]
desired = Energy(1 * u.GeV, 'TeV')
assert_equal(actual, desired)
energy = Energy.equal_log_spacing(2, 6, 3, 'GeV')
actual = energy.nbins
desired = 3
assert_equal(actual, desired)
# range + nbins
erange = energy.range.value
bins = energy.nbins
actual = np.logspace(np.log10(erange[0]), np.log10(erange[1]), bins)
desired = energy.value
assert_equal(actual, desired)
# Input string
e_string = '10 TeV'
actual = Energy(e_string)
desired = Energy(10, 'TeV')
assert_equal(actual, desired)
e_string = u'10 TeV'
actual = Energy(e_string)
desired = Energy(10, 'TeV')
assert_equal(actual, desired)
@remote_data
def test_EnergyBounds():
val = u.Quantity([1, 2, 3, 4, 5], 'TeV')
actual = EnergyBounds(val, 'GeV')
desired = EnergyBounds((1, 2, 3, 4, 5), 'TeV')
assert_equal(actual, desired)
# View casting
energy = val.view(EnergyBounds)
actual = type(energy).__module__
desired = 'gammapy.spectrum.energy'
assert_equal(actual, desired)
# New from template
energy = EnergyBounds([0, 1, 2, 3, 4, 5], 'keV')
energy2 = energy[1:4]
actual = energy2
desired = EnergyBounds([1, 2, 3], 'keV')
assert_equal(actual, desired)
actual = energy2.nbins
desired = 2
assert_equal(actual, desired)
actual = energy2.unit
desired = u.keV
assert_equal(actual, desired)
# Equal log spacing
energy = EnergyBounds.equal_log_spacing(1 * u.TeV, 10 * u.TeV, 10)
actual = energy.nbins
desired = 10
assert_equal(actual, desired)
# Log centers
center = energy.log_centers
actual = type(center).__module__
desired = 'gammapy.spectrum.energy'
assert_equal(actual, desired)
# Upper/lower bounds
actual = energy.upper_bounds
desired = energy[1:]
assert_equal(actual, desired)
actual = energy.lower_bounds
desired = energy[:-1]
assert_equal(actual, desired)
lower = [1, 3, 4, 5]
upper = [3, 4, 5, 8]
actual = EnergyBounds.from_lower_and_upper_bounds(lower, upper, 'TeV')
desired = EnergyBounds([1, 3, 4, 5, 8], 'TeV')
assert_equal(actual, desired)
# Range
erange = energy.range
actual = erange[0]
desired = energy[0]
assert_equal(actual, desired)
actual = erange[1]
desired = energy[-1]
assert_equal(actual, desired)
# Bands
bands = energy.bands
actual = bands[0]
desired = energy[1] - energy[0]
assert_equal(actual, desired)
# read EBOUNDS extension
filename = get_path("../test_datasets/irf/hess/ogip/run_rmf60741.fits",
location='remote')
hdulist = fits.open(filename)
ebounds = EnergyBounds.from_ebounds(hdulist['EBOUNDS'])
desired = hdulist['EBOUNDS'].data['E_MAX'][-1]
actual = ebounds[-1].value
assert_equal(actual, desired)
# read MATRIX extension
ebounds = EnergyBounds.from_rmf_matrix(hdulist['MATRIX'])
desired = hdulist['MATRIX'].data['ENERG_LO'][3]
actual = ebounds[3].value
assert_equal(actual, desired)
def test_EnergyBounds_write(tmpdir):
ebounds = EnergyBounds.equal_log_spacing(1 * u.TeV, 10 * u.TeV, 10)
writename = str(tmpdir.join('ebounds_test.fits'))
hdu = ebounds.to_ebounds()
prim_hdu = fits.PrimaryHDU()
hdulist = fits.HDUList([prim_hdu, hdu])
hdulist.writeto(writename)
ebounds2 = EnergyBounds.from_ebounds(hdulist[1])
actual = ebounds2
desired = ebounds
assert_equal(actual, desired)