diff --git a/tests/test_jungfraulib.py b/tests/test_jungfraulib.py index c7fb5258ba2c0035cbad2668e4bdc35131866855..90e4842f26fa0f0f936b5da73f9b5191d5636f3f 100644 --- a/tests/test_jungfraulib.py +++ b/tests/test_jungfraulib.py @@ -1,7 +1,14 @@ +import numpy as np import pytest from extra_data import RunDirectory -from cal_tools.jungfraulib import JungfrauCtrl, sort_runs_by_gain +from cal_tools.enums import BadPixels +from cal_tools.jungfraulib import ( + JungfrauCtrl, + replace_pixels_gain_values, + replace_pixels_wrong_gain_badpixel, + sort_runs_by_gain, +) # TODO: replace with mocked RAW data as in tests/test_agipdlib.py JF = JungfrauCtrl( @@ -64,3 +71,64 @@ def test_sort_runs_by_gain(original_runs, sorted_runs): runs=original_runs, ctrl_src="FXE_XAD_JF1M/DET/CONTROL") assert validated_runs == sorted_runs + + +GAIN = np.zeros((1, 512, 1024), dtype=np.uint8) +NEW_GAIN = GAIN.copy() # all pixels in high gain. +GAIN[:, 0:255, 896:1024] = 3 # subsection of a chip is stuck in gain bit = 3 +BP = np.zeros((512, 1024, 3), dtype=np.uint32) +for i in range(3): + BP[0:255, 896:1024, i] = BadPixels.WRONG_GAIN_VALUE.value + +MEM_CELLS = 1 +KARABO_DA = "JNGFR04" +REPLACE_WRONG_GAIN_VALUE = 0 +WRONG_GAIN_PIXELS = [4, 0, 255, 896, 1024] + + +def test_replace_pixels_wrong_gain_badpixel(): + assert np.allclose( + NEW_GAIN, + replace_pixels_wrong_gain_badpixel( + gain_arr=GAIN, + new_gain_value=REPLACE_WRONG_GAIN_VALUE, + badpixel_map=BP, + memory_cells=MEM_CELLS, + ) + ) + + +def test_replace_pixels_gain_values(): + assert np.allclose( + NEW_GAIN, + replace_pixels_gain_values( + gain_arr=GAIN, + new_gain_value=REPLACE_WRONG_GAIN_VALUE, + wrong_gain_pixels=WRONG_GAIN_PIXELS, + karabo_da=KARABO_DA, + badpixel_map=BP, + memory_cells=MEM_CELLS, + ) + ) + assert np.allclose( + NEW_GAIN, + replace_pixels_gain_values( + gain_arr=GAIN, + new_gain_value=REPLACE_WRONG_GAIN_VALUE, + wrong_gain_pixels=WRONG_GAIN_PIXELS, + karabo_da=KARABO_DA, + badpixel_map=BP, + memory_cells=MEM_CELLS, + ) + ) + assert np.allclose( + GAIN, # No replacement for non-selected module. + replace_pixels_gain_values( + gain_arr=GAIN, + new_gain_value=REPLACE_WRONG_GAIN_VALUE, + wrong_gain_pixels=WRONG_GAIN_PIXELS, + karabo_da="JNGFR01", + badpixel_map=BP, + memory_cells=MEM_CELLS, + ) + )