added test for non square case
@@ -27,15 +27,25 @@ class TestEnergyDispersion:
)
def test_from_diagonal_matrix(self):
- e_true = [1, 3, 7, 10] * u.TeV
- edisp = EnergyDispersion.from_diagonal_matrix(e_true)
+ e_true = [0.5, 1, 2, 4, 6] * u.TeV
+ e_reco = [2, 4, 6] * u.TeV
- assert edisp.pdf_matrix.shape == (3, 3)
- assert_equal(edisp.pdf_matrix[0][0], 1)
- assert_equal(edisp.pdf_matrix[2][0], 0)
+ edisp = EnergyDispersion.from_diagonal_matrix(e_true, e_reco)
- assert edisp.e_reco.bins.unit == 'TeV'
+ assert edisp.pdf_matrix.shape == (4, 2)
+ expected = np.array( [[0, 0],
+ [0, 0],
+ [1, 0],
+ [0, 1]]
+ )
+ assert_equal(edisp.pdf_matrix, expected)
+
+ # Test square matrix
+ edisp = EnergyDispersion.from_diagonal_matrix(e_true)
assert_allclose(edisp.e_reco.bins.value, e_true.value)
+ assert edisp.e_reco.bins.unit == 'TeV'
+ assert_equal(edisp.pdf_matrix[0][0], 1)
+ assert_equal(edisp.pdf_matrix[2][0], 0)
def test_str(self):
assert 'EnergyDispersion' in str(self.edisp)
