Merge branch 'main' into rtbdb_implementation

Author: BradleySappington

src/wfi_reference_pipeline/reference_types/tests/test_readnoise.py

@@ -1,3 +1,4 @@
+import asdf
 import numpy as np
 import pytest
 
@@ -9,9 +10,42 @@
 )
 from wfi_reference_pipeline.reference_types.readnoise.readnoise import ReadNoise
 from wfi_reference_pipeline.resources.make_test_meta import MakeTestMeta
+from wfi_reference_pipeline.utilities.simulate_reads import simulate_dark_reads
 
+# Set smaller test data size
+TEST_DETECTOR_PIXEL_COUNT = 32 
 
-@pytest.fixture
+
+@pytest.fixture(scope="module")
+def simulated_reads_filelist(tmp_path_factory, ref_type_data_factory):
+
+    data_path = tmp_path_factory.mktemp("data")
+
+    file_list = []
+
+    for i in range(1, 4):
+        cube_data = ref_type_data_factory(i)
+        
+        curr_path = data_path / f"data_num_{i}.asdf"
+        curr_path.parent.mkdir(parents=True, exist_ok=True)
+
+        tree = {
+            "roman" : {
+                "data" : cube_data
+            }
+        }
+
+        af = asdf.AsdfFile(tree)
+
+        af.write_to(curr_path)
+
+        file_list.append(str(curr_path))
+
+    # Return the file list
+    return file_list
+
+
+@pytest.fixture(scope="module")
 def valid_meta_data():
     """Fixture for generating valid meta_data for ReadNoise class."""
     test_meta = MakeTestMeta(ref_type=REF_TYPE_READNOISE)
@@ -24,27 +58,39 @@ def valid_ref_type_data_array():
     return np.random.random((DETECTOR_PIXEL_X_COUNT, DETECTOR_PIXEL_Y_COUNT))  # Simulate a valid read noise image
 
 
-@pytest.fixture
-def valid_ref_type_data_cube():
-    """Fixture for generating valid ref_type_data cube (read noise cube)."""
-    return np.random.random((3, DETECTOR_PIXEL_X_COUNT, DETECTOR_PIXEL_Y_COUNT))  # Simulate a valid read noise image
+# NOTE: This factory makes a new read each call
+#    We can introduce a _cache dictionary to store in the future if generating becomes too expensive
+#    I'm currently not concerned as long as we use the TEST_DETECTOR_PIXEL_COUNT to keep data small
+@pytest.fixture(scope="session")
+def ref_type_data_factory():
+    """Factory fixture for generating valid ref_type_data cubes with N reads."""
+    def _make_cube(num_reads):
+        cube_data, _ = simulate_dark_reads(num_reads, ni=TEST_DETECTOR_PIXEL_COUNT)
+        return cube_data
+    
+    return _make_cube
 
 
 @pytest.fixture
 def readnoise_object_with_data_array(valid_meta_data, valid_ref_type_data_array):
     """Fixture for initializing a ReadNoise object with a valid data array."""
     readnoise_object_with_data_array = ReadNoise(meta_data=valid_meta_data,
                                                  ref_type_data=valid_ref_type_data_array)
-    yield readnoise_object_with_data_array
+    return readnoise_object_with_data_array
 
 
 @pytest.fixture
-def readnoise_object_with_data_cube(valid_meta_data, valid_ref_type_data_cube):
+def readnoise_object_with_data_cube(valid_meta_data, ref_type_data_factory):
     """Fixture for initializing a ReadNoise object with a valid data cube."""
+    ref_type_data = ref_type_data_factory(3)
     readnoise_object_with_data_cube = ReadNoise(meta_data=valid_meta_data,
-                                                ref_type_data=valid_ref_type_data_cube)
-    yield readnoise_object_with_data_cube
+                                                ref_type_data=ref_type_data)
+    return readnoise_object_with_data_cube
 
+@pytest.fixture
+def readnoise_object_with_file_list(valid_meta_data, simulated_reads_filelist):
+    readnoise_object_with_file_list_obj = ReadNoise(meta_data=valid_meta_data, file_list=simulated_reads_filelist)
+    return readnoise_object_with_file_list_obj
 
 class TestReadNoise:
 
@@ -137,3 +183,65 @@ def test_readnoise_outfile_default(self, readnoise_object_with_data_array):
         """
         assert readnoise_object_with_data_array.outfile == "roman_readnoise.asdf"
 
+
+def test_make_readnoise_image_sets_correct_shape_pass(readnoise_object_with_file_list):
+
+    readnoise_object_with_file_list.make_readnoise_image()
+
+    assert readnoise_object_with_file_list.readnoise_image is not None
+    assert readnoise_object_with_file_list.readnoise_image.shape == (TEST_DETECTOR_PIXEL_COUNT, TEST_DETECTOR_PIXEL_COUNT)
+    
+
+
+def test_select_data_cube_sets_correct_num_reads_pass(readnoise_object_with_file_list):
+
+    # Check datacube doesn't exist
+    with pytest.raises(AttributeError):
+        _ = readnoise_object_with_file_list.data_cube
+
+    readnoise_object_with_file_list._select_data_cube_from_file_list()
+
+    assert readnoise_object_with_file_list.data_cube is not None
+    assert readnoise_object_with_file_list.data_cube.num_reads == 3
+
+def test_make_rate_image_updates_dimensions_pass(readnoise_object_with_data_cube):
+
+    readnoise_object_with_data_cube.make_rate_image_from_data_cube()
+
+    assert readnoise_object_with_data_cube.data_cube.rate_image.shape == (TEST_DETECTOR_PIXEL_COUNT, TEST_DETECTOR_PIXEL_COUNT)
+    assert readnoise_object_with_data_cube.data_cube.intercept_image.shape == (TEST_DETECTOR_PIXEL_COUNT, TEST_DETECTOR_PIXEL_COUNT)
+
+def test_comp_ramp_res_var_output_matches_pixel_count_pass(readnoise_object_with_data_cube, ref_type_data_factory, mocker):
+
+    # Mock a datacube with the necessary parts: ramp_model
+    mock_readnoise_datacube = mocker.Mock()
+
+    mock_readnoise_datacube.num_i_pixels = TEST_DETECTOR_PIXEL_COUNT
+    mock_readnoise_datacube.num_j_pixels = TEST_DETECTOR_PIXEL_COUNT
+    mock_readnoise_datacube.ramp_model = ref_type_data_factory(3)
+    mock_readnoise_datacube.data = ref_type_data_factory(3)
+
+    # Add datacube to the readnoise
+    mocker.patch.object(readnoise_object_with_data_cube, 'data_cube', mock_readnoise_datacube)
+
+    result = readnoise_object_with_data_cube.comp_ramp_res_var()
+
+    assert result.shape == (TEST_DETECTOR_PIXEL_COUNT, TEST_DETECTOR_PIXEL_COUNT)
+
+def est_comp_cds_noise_output_matches_pixel_count_pass(readnoise_object_with_data_cube, ref_type_data_factory, mocker):
+
+    # Mock a datacube with the necessary parts: ramp_model
+    mock_readnoise_datacube = mocker.Mock()
+
+    mock_readnoise_datacube.num_i_pixels = TEST_DETECTOR_PIXEL_COUNT
+    mock_readnoise_datacube.num_j_pixels = TEST_DETECTOR_PIXEL_COUNT
+    mock_readnoise_datacube.ramp_model = ref_type_data_factory(3)
+    mock_readnoise_datacube.data = ref_type_data_factory(3)
+    mock_readnoise_datacube.num_reads = 3
+
+    # Add datacube to the readnoise
+    mocker.patch.object(readnoise_object_with_data_cube, 'data_cube', mock_readnoise_datacube)
+
+    result = readnoise_object_with_data_cube.comp_cds_noise()
+
+    assert result.shape == (TEST_DETECTOR_PIXEL_COUNT, TEST_DETECTOR_PIXEL_COUNT)