dkist-processing-cryonirsp 1.3.4__py3-none-any.whl

dkist_processing_cryonirsp/tests/test_sp_geometric.py

@@ -0,0 +1,319 @@
+import json
+
+import numpy as np
+import pytest
+from dkist_processing_common._util.scratch import WorkflowFileSystem
+from dkist_processing_common.codecs.fits import fits_hdulist_encoder
+from dkist_processing_common.tests.conftest import FakeGQLClient
+from dkist_processing_math import transform
+
+from dkist_processing_cryonirsp.models.exposure_conditions import AllowableOpticalDensityFilterNames
+from dkist_processing_cryonirsp.models.exposure_conditions import ExposureConditions
+from dkist_processing_cryonirsp.models.tags import CryonirspTag
+from dkist_processing_cryonirsp.tasks.sp_geometric import SPGeometricCalibration
+from dkist_processing_cryonirsp.tests.conftest import cryonirsp_testing_parameters_factory
+from dkist_processing_cryonirsp.tests.conftest import CryonirspConstantsDb
+from dkist_processing_cryonirsp.tests.conftest import generate_214_l0_fits_frame
+from dkist_processing_cryonirsp.tests.header_models import CryonirspHeadersValidSPSolarGainFrames
+
+
+@pytest.fixture(scope="function")
+def geometric_calibration_task_that_completes(
+    tmp_path, recipe_run_id, assign_input_dataset_doc_to_task, init_cryonirsp_constants_db
+):
+    # This fixture makes data that look enough like real data that all of the feature detection stuff at least runs
+    # through (mostly this is an issue for the angle calculation). It would be great to contrive data that
+    # produce a geometric calibration with real numbers that can be checked, but for now we'll rely on the grogu
+    # tests for that. In other words, this fixture just tests if the machinery of the task completes and some object
+    # (ANY object) is written correctly.
+    number_of_modstates = 1
+    number_of_beams = 2
+    exposure_time = 20.0  # From CryonirspHeadersValidSolarGainFrames fixture
+    exposure_conditions = ExposureConditions(
+        exposure_time, AllowableOpticalDensityFilterNames.OPEN.value
+    )
+    data_shape_int = 100, 98
+    data_shape_raw = 100, 196
+    dataset_shape = 1, *data_shape_raw
+    constants_db = CryonirspConstantsDb(
+        NUM_MODSTATES=number_of_modstates,
+        SOLAR_GAIN_EXPOSURE_CONDITIONS_LIST=(exposure_conditions,),
+        ARM_ID="SP",
+    )
+    init_cryonirsp_constants_db(recipe_run_id, constants_db)
+    with SPGeometricCalibration(
+        recipe_run_id=recipe_run_id,
+        workflow_name="sp_geometric_calibration",
+        workflow_version="VX.Y",
+    ) as task:
+        try:  # This try... block is here to make sure the dbs get cleaned up if there's a failure in the fixture
+            task.scratch = WorkflowFileSystem(
+                scratch_base_path=tmp_path, recipe_run_id=recipe_run_id
+            )
+            param_class = cryonirsp_testing_parameters_factory(param_path=tmp_path)
+            assign_input_dataset_doc_to_task(task, param_class())
+
+            task.angles = [0.0, 0.0]
+            task.offsets = np.zeros((number_of_beams, 2))
+            task.shifts = np.zeros(data_shape_int[0])
+
+            # Create the intermediate frames needed
+            # Create fake bad pixel map
+            task.intermediate_frame_write_arrays(
+                arrays=np.zeros(data_shape_raw),
+                task_tag=CryonirspTag.task_bad_pixel_map(),
+            )
+            for beam in range(1, number_of_beams + 1):
+                dark_cal = np.ones(data_shape_int) * 3.0
+                task.intermediate_frame_write_arrays(
+                    arrays=dark_cal,
+                    beam=beam,
+                    task_tag=CryonirspTag.task_dark(),
+                    exposure_conditions=exposure_conditions,
+                )
+
+                # Need a lamp for each beam
+                lamp_gain = np.ones(data_shape_int)
+                task.intermediate_frame_write_arrays(
+                    arrays=lamp_gain,
+                    beam=beam,
+                    task_tag=CryonirspTag.task_lamp_gain(),
+                    exposure_conditions=exposure_conditions,
+                )
+
+                # And a beam border intermediate array
+                border = data_shape_raw[1] // 2
+                task.intermediate_frame_write_arrays(
+                    arrays=np.array(
+                        [
+                            0,
+                            data_shape_int[0],
+                            ((beam - 1) * border),
+                            (border + (beam - 1) * border),
+                        ]
+                    ),
+                    task_tag=CryonirspTag.task_beam_boundaries(),
+                    beam=beam,
+                )
+
+            # Create the raw data, which is based on two beams per frame
+            beam1 = 1
+            beam2 = 2
+            ds = CryonirspHeadersValidSPSolarGainFrames(
+                dataset_shape=dataset_shape,
+                array_shape=dataset_shape,
+                time_delta=10,
+            )
+            header = ds.header()
+            # Create each beam of the solar gain image, rotate and translate them, and then combine them into one
+            true_solar_1 = np.ones(data_shape_int)
+            true_solar_1[data_shape_int[0] // 2, :] += 5
+            true_solar_1[:, data_shape_int[1] // 2] += 5
+            true_solar_2 = np.copy(true_solar_1)
+            # Now add the beam number to each beam in the array
+            true_solar_1 += beam1
+            true_solar_2 += beam2
+            translated_solar_1 = next(
+                transform.translate_arrays(arrays=true_solar_1, translation=task.offsets[0])
+            )
+            translated_solar_2 = next(
+                transform.translate_arrays(arrays=true_solar_2, translation=task.offsets[1])
+            )
+            distorted_solar_1 = next(
+                transform.rotate_arrays_about_point(arrays=translated_solar_1, angle=task.angles[0])
+            )
+            distorted_solar_2 = next(
+                transform.rotate_arrays_about_point(arrays=translated_solar_2, angle=task.angles[1])
+            )
+
+            raw_solar_1 = distorted_solar_1 + dark_cal
+            raw_solar_2 = distorted_solar_2 + dark_cal
+            raw_solar = np.concatenate((raw_solar_1, raw_solar_2), axis=1)
+            solar_hdul = generate_214_l0_fits_frame(data=raw_solar, s122_header=header)
+            task.write(
+                data=solar_hdul,
+                tags=[
+                    CryonirspTag.linearized(),
+                    CryonirspTag.task_solar_gain(),
+                    CryonirspTag.frame(),
+                    CryonirspTag.beam(beam),
+                    CryonirspTag.exposure_conditions(exposure_conditions),
+                ],
+                encoder=fits_hdulist_encoder,
+            )
+
+            yield task, data_shape_int
+        finally:
+            task._purge()
+
+
+@pytest.fixture(scope="function")
+def geometric_calibration_task_with_simple_raw_data(
+    tmp_path, recipe_run_id, assign_input_dataset_doc_to_task, init_cryonirsp_constants_db
+):
+    number_of_modstates = 1
+    number_of_beams = 2
+    exposure_time = 20.0  # From CryonirspHeadersValidSolarGainFrames fixture
+    exposure_conditions = ExposureConditions(
+        exposure_time, AllowableOpticalDensityFilterNames.OPEN.value
+    )
+    data_shape_int = 100, 98
+    data_shape_raw = 100, 196
+    constants_db = CryonirspConstantsDb(
+        NUM_MODSTATES=number_of_modstates,
+        SOLAR_GAIN_EXPOSURE_CONDITIONS_LIST=(exposure_conditions,),
+        ARM_ID="SP",
+    )
+    init_cryonirsp_constants_db(recipe_run_id, constants_db)
+    with SPGeometricCalibration(
+        recipe_run_id=recipe_run_id,
+        workflow_name="sp_geometric_calibration",
+        workflow_version="VX.Y",
+    ) as task:
+        try:  # This try... block is here to make sure the dbs get cleaned up if there's a failure in the fixture
+            task.scratch = WorkflowFileSystem(
+                scratch_base_path=tmp_path, recipe_run_id=recipe_run_id
+            )
+            param_class = cryonirsp_testing_parameters_factory(param_path=tmp_path)
+            assign_input_dataset_doc_to_task(task, param_class())
+
+            # Create the intermediate frames needed
+            # Create a fake bad pixel map
+            task.intermediate_frame_write_arrays(
+                arrays=np.zeros(data_shape_raw),
+                task_tag=CryonirspTag.task_bad_pixel_map(),
+            )
+            for beam in range(1, number_of_beams + 1):
+                dark_cal = np.ones(data_shape_int) * 3.0
+                task.intermediate_frame_write_arrays(
+                    arrays=dark_cal,
+                    beam=beam,
+                    task_tag=CryonirspTag.task_dark(),
+                    exposure_conditions=exposure_conditions,
+                )
+
+                # Need a lamp for each beam
+                lamp_gain = np.ones(data_shape_int)
+                task.intermediate_frame_write_arrays(
+                    arrays=lamp_gain,
+                    beam=beam,
+                    task_tag=CryonirspTag.task_lamp_gain(),
+                    exposure_conditions=exposure_conditions,
+                )
+
+                # And a beam border intermediate array
+                border = data_shape_raw[1] // 2
+                task.intermediate_frame_write_arrays(
+                    arrays=np.array(
+                        [
+                            0,
+                            data_shape_int[0],
+                            ((beam - 1) * border),
+                            (border + (beam - 1) * border),
+                        ]
+                    ),
+                    task_tag=CryonirspTag.task_beam_boundaries(),
+                    beam=beam,
+                )
+
+                # Let's write a dark with the wrong exposure time, just to make sure it doesn't get used
+                task.intermediate_frame_write_arrays(
+                    arrays=np.ones(data_shape_int) * 1e6,
+                    beam=beam,
+                    task_tag=CryonirspTag.task_dark(),
+                    exposure_conditions=ExposureConditions(
+                        exposure_time**2, AllowableOpticalDensityFilterNames.OPEN.value
+                    ),
+                )
+
+            # Create the raw data, which is based on two beams per frame
+            beam1 = 1
+            beam2 = 2
+            dark_cal_two_beams = np.concatenate((dark_cal, dark_cal), axis=1)
+            ds = CryonirspHeadersValidSPSolarGainFrames(
+                dataset_shape=(1,) + data_shape_raw,
+                array_shape=(1,) + data_shape_raw,
+                time_delta=10,
+            )
+            header = ds.header()
+            true_solar = np.ones(data_shape_raw)
+            # Now add the beam number to each beam in the array
+            true_solar[:, : data_shape_int[1]] += beam1
+            true_solar[:, data_shape_int[1] :] += beam2
+            raw_solar = true_solar + dark_cal_two_beams
+            solar_hdul = generate_214_l0_fits_frame(data=raw_solar, s122_header=header)
+            task.write(
+                data=solar_hdul,
+                tags=[
+                    CryonirspTag.linearized(),
+                    CryonirspTag.task_solar_gain(),
+                    CryonirspTag.frame(),
+                    CryonirspTag.exposure_conditions(exposure_conditions),
+                ],
+                encoder=fits_hdulist_encoder,
+            )
+
+            yield task, data_shape_int, true_solar
+        finally:
+            task._purge()
+
+
+def test_geometric_task(geometric_calibration_task_that_completes, mocker):
+    """
+    Given: A set of raw solar gain images and necessary intermediate calibrations
+    When: Running the geometric task
+    Then: The damn thing runs and makes outputs that at least are the right type
+    """
+    # See the note in the fixture above: this test does NOT test for accuracy of the calibration
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+    )
+    task, beam_shape = geometric_calibration_task_that_completes
+    task()
+    for beam in range(1, task.constants.num_beams + 1):
+        assert type(task.intermediate_frame_load_angle(beam=beam)) is float
+        assert task.intermediate_frame_load_spec_shift(beam=beam).shape[0] == beam_shape[0]
+        assert task.intermediate_frame_load_state_offset(beam=beam).shape == (2,)
+
+    quality_files = task.read(tags=[CryonirspTag.quality("TASK_TYPES")])
+    for file in quality_files:
+        with file.open() as f:
+            data = json.load(f)
+            assert isinstance(data, dict)
+            assert data["total_frames"] == task.scratch.count_all(
+                tags=[
+                    CryonirspTag.linearized(),
+                    CryonirspTag.frame(),
+                    CryonirspTag.task_solar_gain(),
+                ]
+            )
+
+
+def test_basic_corrections(geometric_calibration_task_with_simple_raw_data):
+    """
+    Given: A set of raw solar gain images and necessary intermediate calibrations
+    When: Doing basic dark and lamp gain corrections
+    Then: The corrections are applied correctly
+    """
+    task, beam_shape, true_solar = geometric_calibration_task_with_simple_raw_data
+    flipped_true_solar = np.flip(true_solar, axis=1)
+    task.do_basic_corrections()
+
+    # Positive test on the flipped true solar image
+    for beam in range(1, task.constants.num_beams + 1):
+        if beam == 1:
+            expected = flipped_true_solar[:, 0 : beam_shape[1]]
+        else:
+            expected = flipped_true_solar[:, beam_shape[1] :]
+        array = task.basic_dark_bp_corrected_data(beam=beam)
+        np.testing.assert_equal(expected, array)
+
+    # Negative test on the original true solar image
+    for beam in range(1, task.constants.num_beams + 1):
+        if beam == 1:
+            expected = true_solar[:, 0 : beam_shape[1]]
+        else:
+            expected = true_solar[:, beam_shape[1] :]
+        array = task.basic_dark_bp_corrected_data(beam=beam)
+        with pytest.raises(AssertionError):
+            np.testing.assert_equal(expected, array)

dkist_processing_cryonirsp/tests/test_sp_make_movie_frames.py

@@ -0,0 +1,121 @@
+from datetime import datetime
+
+import numpy as np
+import pytest
+from astropy.io import fits
+from dkist_header_validator import spec122_validator
+from dkist_processing_common._util.scratch import WorkflowFileSystem
+from dkist_processing_common.codecs.fits import fits_hdulist_encoder
+from dkist_processing_common.tests.conftest import FakeGQLClient
+
+from dkist_processing_cryonirsp.models.tags import CryonirspTag
+from dkist_processing_cryonirsp.tasks.make_movie_frames import SPMakeCryonirspMovieFrames
+from dkist_processing_cryonirsp.tests.conftest import CryonirspConstantsDb
+from dkist_processing_cryonirsp.tests.conftest import generate_fits_frame
+from dkist_processing_cryonirsp.tests.header_models import CryonirspHeadersValidObserveFrames
+
+
+@pytest.fixture(
+    scope="function",
+    params=[
+        pytest.param((True, 2), id="polarimetric"),
+        pytest.param((False, 2), id="intensity-only"),
+        pytest.param((False, 1), id="single_scan_step"),
+    ],
+)
+def sp_movie_frames_task(tmp_path, recipe_run_id, init_cryonirsp_constants_db, request):
+    is_polarimetric, scan_steps = request.param
+    map_scans = 3
+    array_shape = (3, 4)
+    if is_polarimetric:
+        num_mod = 8
+        spin_mode = "Continuous"
+    else:
+        num_mod = 1
+        spin_mode = "None"
+    constants_db = CryonirspConstantsDb(
+        NUM_MODSTATES=num_mod,
+        MODULATOR_SPIN_MODE=spin_mode,
+        NUM_SCAN_STEPS=scan_steps,
+        NUM_MAP_SCANS=map_scans,
+        TIME_OBS_LIST=(datetime.now().isoformat("T"),),
+    )
+    init_cryonirsp_constants_db(recipe_run_id, constants_db)
+    with SPMakeCryonirspMovieFrames(
+        recipe_run_id=recipe_run_id, workflow_name="make_movie_frames", workflow_version="VX.Y"
+    ) as task:
+        try:  # This try... block is here to make sure the dbs get cleaned up if there's a failure in the fixture
+            meas_num = 1  # Use only the first measurement if there are multiple measurements.
+            task.scratch = WorkflowFileSystem(
+                scratch_base_path=tmp_path, recipe_run_id=recipe_run_id
+            )
+            start_time = datetime.now()
+            for stokes_state in ["I", "Q", "U", "V"]:
+                for map_scan in range(1, map_scans + 1):
+                    for scan_step in range(0, scan_steps + 1):
+                        ds = CryonirspHeadersValidObserveFrames(
+                            dataset_shape=(2, *array_shape),
+                            array_shape=(1, *array_shape),
+                            time_delta=10,
+                            num_map_scans=map_scans,
+                            map_scan=map_scan,
+                            num_scan_steps=scan_steps,
+                            scan_step=scan_step,
+                            num_modstates=1,
+                            modstate=1,
+                            start_time=start_time,
+                            num_meas=1,
+                            meas_num=1,
+                            arm_id="SP",
+                        )
+                        header_generator = (
+                            spec122_validator.validate_and_translate_to_214_l0(
+                                d.header(), return_type=fits.HDUList
+                            )[0].header
+                            for d in ds
+                        )
+                        hdul = generate_fits_frame(
+                            header_generator=header_generator,
+                            shape=(1, *array_shape),
+                        )
+                        task.write(
+                            data=hdul,
+                            tags=[
+                                CryonirspTag.calibrated(),
+                                CryonirspTag.frame(),
+                                CryonirspTag.map_scan(map_scan),
+                                CryonirspTag.scan_step(scan_step),
+                                CryonirspTag.stokes(stokes_state),
+                                CryonirspTag.meas_num(meas_num),
+                            ],
+                            encoder=fits_hdulist_encoder,
+                        )
+            yield task, map_scans, scan_steps, array_shape, is_polarimetric
+        finally:
+            task._purge()
+
+
+def test_sp_make_movie_frames(sp_movie_frames_task, mocker):
+    """
+    Given: A SPMakeCryonirspMovieFrames task
+    When: Calling the task instance
+    Then: a fits file is made for each scan containing the movie frame for that scan
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+    )
+    task, map_scans, scan_steps, array_shape, is_polarimetric = sp_movie_frames_task
+    if scan_steps > 1:
+        expected_shape = (scan_steps, array_shape[0])
+    else:
+        expected_shape = array_shape
+    if is_polarimetric:
+        expected_shape = tuple(np.array(expected_shape) * 2)
+    task()
+    assert len(list(task.read(tags=[CryonirspTag.movie_frame()]))) == map_scans
+    for filepath in task.read(tags=[CryonirspTag.movie_frame()]):
+        assert filepath.exists()
+        hdul = fits.open(filepath)
+        assert hdul[0].header["INSTRUME"] == "CRYO-NIRSP"
+        # Multiply by 2 because a single map is (axis_length, steps) but there are 4 stokes in a 2x2 array
+        assert hdul[0].data.shape == expected_shape