added tests

python/tests/test_calibration.py

@@ -1,6 +1,7 @@
 import pytest 
 import numpy as np
-from aare import apply_calibration
+
+import aare
 
 def test_apply_calibration_small_data():
     # The raw data consists of 10 4x5 images
@@ -27,7 +28,7 @@ def test_apply_calibration_small_data():
 
     
 
-    data = apply_calibration(raw, pd = pedestal, cal = calibration)
+    data = aare.apply_calibration(raw, pd = pedestal, cal = calibration)
 
 
     # The formula that is applied is:
@@ -41,3 +42,94 @@ def test_apply_calibration_small_data():
     assert data[2,2,2] == 0
     assert data[0,1,1] == 0
     assert data[1,3,0] == 0
+
+
+@pytest.fixture
+def raw_data_3x2x2():
+    raw = np.zeros((3, 2, 2), dtype=np.uint16)
+    raw[0, 0, 0] = 100
+    raw[1,0, 0] =  200
+    raw[2, 0, 0] = 300
+
+    raw[0, 0, 1] = (1<<14) + 100
+    raw[1, 0, 1] = (1<<14) + 200
+    raw[2, 0, 1] = (1<<14) + 300
+
+    raw[0, 1, 0] = (1<<14) + 37
+    raw[1, 1, 0] =  38
+    raw[2, 1, 0] = (3<<14) + 39
+
+    raw[0, 1, 1] = (3<<14) + 100
+    raw[1, 1, 1] = (3<<14) + 200
+    raw[2, 1, 1] = (3<<14) + 300
+    return raw
+
+def test_calculate_pedestal(raw_data_3x2x2):
+    # Calculate the pedestal
+    pd = aare.calculate_pedestal(raw_data_3x2x2)
+    assert pd.shape == (3, 2, 2)
+    assert pd.dtype == np.float64
+    assert pd[0, 0, 0] == 200
+    assert pd[1, 0, 0] == 0
+    assert pd[2, 0, 0] == 0
+
+    assert pd[0, 0, 1] == 0
+    assert pd[1, 0, 1] == 200
+    assert pd[2, 0, 1] == 0
+
+    assert pd[0, 1, 0] == 38
+    assert pd[1, 1, 0] == 37
+    assert pd[2, 1, 0] == 39
+
+    assert pd[0, 1, 1] == 0
+    assert pd[1, 1, 1] == 0
+    assert pd[2, 1, 1] == 200
+
+def test_calculate_pedestal_float(raw_data_3x2x2):
+    #results should be the same for float
+    pd2 = aare.calculate_pedestal_float(raw_data_3x2x2)
+    assert pd2.shape == (3, 2, 2)
+    assert pd2.dtype == np.float32
+    assert pd2[0, 0, 0] == 200
+    assert pd2[1, 0, 0] == 0
+    assert pd2[2, 0, 0] == 0
+
+    assert pd2[0, 0, 1] == 0
+    assert pd2[1, 0, 1] == 200
+    assert pd2[2, 0, 1] == 0
+
+    assert pd2[0, 1, 0] == 38
+    assert pd2[1, 1, 0] == 37
+    assert pd2[2, 1, 0] == 39
+
+    assert pd2[0, 1, 1] == 0
+    assert pd2[1, 1, 1] == 0
+    assert pd2[2, 1, 1] == 200
+
+def test_calculate_pedestal_g0(raw_data_3x2x2):
+    pd = aare.calculate_pedestal_g0(raw_data_3x2x2)
+    assert pd.shape == (2, 2)
+    assert pd.dtype == np.float64
+    assert pd[0, 0] == 200
+    assert pd[1, 0] == 38
+    assert pd[0, 1] == 0
+    assert pd[1, 1] == 0
+
+def test_calculate_pedestal_g0_float(raw_data_3x2x2):
+    pd = aare.calculate_pedestal_g0_float(raw_data_3x2x2)
+    assert pd.shape == (2, 2)
+    assert pd.dtype == np.float32
+    assert pd[0, 0] == 200
+    assert pd[1, 0] == 38
+    assert pd[0, 1] == 0
+    assert pd[1, 1] == 0
+
+def test_count_switching_pixels(raw_data_3x2x2):
+    # Count the number of pixels that switched gain
+    count = aare.count_switching_pixels(raw_data_3x2x2)
+    assert count.shape == (2, 2)
+    assert count.sum() == 8
+    assert count[0, 0] == 0 
+    assert count[1, 0] == 2 
+    assert count[0, 1] == 3 
+    assert count[1, 1] == 3