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

dkist_processing_cryonirsp/tests/test_bad_pixel_maps.py

@@ -0,0 +1,115 @@
+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_service_configuration.logging import logger
+
+from dkist_processing_cryonirsp.models.tags import CryonirspTag
+from dkist_processing_cryonirsp.tasks.bad_pixel_map import BadPixelMapCalibration
+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_fits_frame
+from dkist_processing_cryonirsp.tests.header_models import CryonirspHeadersValidCISolarGainFrames
+
+
+@pytest.fixture(scope="function", params=["CI", "SP"])
+def compute_bad_pixel_map_task(
+    tmp_path,
+    recipe_run_id,
+    assign_input_dataset_doc_to_task,
+    init_cryonirsp_constants_db,
+    request,
+):
+    arm_id = request.param
+    if arm_id == "SP":
+        dataset_shape = (1, 100, 200)
+        array_shape = (1, 100, 200)
+    else:
+        dataset_shape = (1, 100, 100)
+        array_shape = (1, 100, 100)
+    constants_db = CryonirspConstantsDb(ARM_ID=arm_id)
+    init_cryonirsp_constants_db(recipe_run_id, constants_db)
+    with BadPixelMapCalibration(
+        recipe_run_id=recipe_run_id,
+        workflow_name="sp_compute_bad_pixel_map",
+        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())
+            start_time = datetime.now()
+            ds = CryonirspHeadersValidCISolarGainFrames(
+                dataset_shape=dataset_shape,
+                array_shape=array_shape,
+                time_delta=10,
+                start_time=start_time,
+            )
+            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=array_shape)
+            # Create an array with a random number of zero and hot values
+            rng = np.random.default_rng()
+            # Generate rabdomly the number of bad pixels to be used in trange(50, 100)
+            num_bad_pixels = rng.integers(50, 100)
+            # Let 2/3 of the bad pixels be hot
+            num_hot_pixels = num_bad_pixels * 2 // 3
+            # Let the remaining 1/3 be zero
+            num_zero_pixels = num_bad_pixels - num_hot_pixels
+            logger.debug(f"{num_bad_pixels = }, {num_hot_pixels = }, {num_zero_pixels = }")
+            nelem = np.prod(array_shape)
+            array = 1000.0 * np.ones(nelem)
+            # Need choice here with replace=False to avoid generating duplicates
+            bad_pixel_locs = rng.choice(nelem, size=num_bad_pixels, replace=False)
+            hot_pixel_locs = sorted(bad_pixel_locs[:num_hot_pixels])
+            zero_pixel_locs = sorted(bad_pixel_locs[num_hot_pixels:])
+            array[zero_pixel_locs] = 0.0
+            array[hot_pixel_locs] = 2000.0
+            array = array.reshape(array_shape[1:])
+            hdul[0].data = array
+            task.write(
+                data=hdul,
+                tags=[
+                    CryonirspTag.linearized(),
+                    CryonirspTag.task_solar_gain(),
+                    CryonirspTag.frame(),
+                ],
+                encoder=fits_hdulist_encoder,
+            )
+            yield task, num_zero_pixels, num_hot_pixels
+        finally:
+            task._purge()
+
+
+def test_compute_bad_pixel_map_task(compute_bad_pixel_map_task, mocker):
+    """
+    Given: An BadPixelMapCalibration task
+    When: Calling the task instance with known input data
+    Then: The correct beam boundary values are created and saved as intermediate files
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+    )
+    # Given
+    task, num_zeros, num_hot = compute_bad_pixel_map_task
+    # When
+    task()
+    # Then
+    tags = [CryonirspTag.task_bad_pixel_map()]
+    bad_pixel_map_paths = list(task.read(tags))
+    assert len(bad_pixel_map_paths) == 1
+    bad_pixel_map_hdul = fits.open(bad_pixel_map_paths[0])
+    bad_pixel_map = bad_pixel_map_hdul[0].data
+    num_bad_pixels = bad_pixel_map.sum()
+    assert num_bad_pixels == num_zeros + num_hot

dkist_processing_cryonirsp/tests/test_ci_beam_boundaries.py

@@ -0,0 +1,106 @@
+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.ci_beam_boundaries import CIBeamBoundariesCalibration
+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_fits_frame
+from dkist_processing_cryonirsp.tests.header_models import CryonirspHeadersValidCISolarGainFrames
+
+
+@pytest.fixture(scope="function")
+def compute_beam_boundaries_task(
+    tmp_path,
+    recipe_run_id,
+    assign_input_dataset_doc_to_task,
+    init_cryonirsp_constants_db,
+):
+    arm_id = "CI"
+    dataset_shape = (1, 100, 100)
+    array_shape = (1, 100, 100)
+    constants_db = CryonirspConstantsDb(ARM_ID=arm_id)
+    init_cryonirsp_constants_db(recipe_run_id, constants_db)
+    with CIBeamBoundariesCalibration(
+        recipe_run_id=recipe_run_id,
+        workflow_name="ci_compute_beam_boundaries",
+        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 fake bad pixel map
+            task.intermediate_frame_write_arrays(
+                arrays=np.zeros(array_shape[1:]),
+                task_tag=CryonirspTag.task_bad_pixel_map(),
+            )
+            start_time = datetime.now()
+            ds = CryonirspHeadersValidCISolarGainFrames(
+                dataset_shape=dataset_shape,
+                array_shape=array_shape,
+                time_delta=10,
+                start_time=start_time,
+            )
+            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=array_shape)
+            # Tweak data to form a beam illumination pattern
+            # Data from generate_fits_frame are value 150
+            array = hdul[0].data
+            # Initial illumination borders that are made up. Precise border depends on the algorithm.
+            # [0:0, y_min:y_max, x_min:x_max]
+            array[:, 7:-5, 3:-8] = 1000.0
+            # Put some large vertical streaks in the image to help the shift measurement converge
+            minus_streak_pos = array_shape[2] // 4
+            plus_streak_pos = 3 * array_shape[2] // 4
+            array[:, :, minus_streak_pos - 5 : minus_streak_pos + 5] += 100
+            array[:, :, plus_streak_pos - 5 : plus_streak_pos + 5] += 100
+            hdul[0].data = array
+            task.write(
+                data=hdul,
+                tags=[
+                    CryonirspTag.linearized(),
+                    CryonirspTag.task_solar_gain(),
+                    CryonirspTag.frame(),
+                ],
+                encoder=fits_hdulist_encoder,
+            )
+            yield task, arm_id
+        finally:
+            task._purge()
+
+
+def test_compute_beam_boundaries_task(compute_beam_boundaries_task, mocker):
+    """
+    Given: A CIBeamBoundariesCalibration task
+    When: Calling the task instance with known input data
+    Then: The correct beam boundary values are created and saved as intermediate files
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+    )
+    # Given
+    task, arm_id = compute_beam_boundaries_task
+    # When
+    task()
+    # Then
+    beam_1_tags = [CryonirspTag.task_beam_boundaries(), CryonirspTag.beam(1)]
+    beam_1_boundary = np.array([8, 94, 4, 91])
+    files_found = list(task.read(tags=beam_1_tags))
+    assert len(files_found) == 1
+    array = fits.open(files_found[0])[0].data
+    assert np.array_equal(array, beam_1_boundary)

dkist_processing_cryonirsp/tests/test_ci_science.py

@@ -0,0 +1,355 @@
+import json
+import random
+from datetime import datetime
+
+import numpy as np
+import pytest
+from astropy.io import fits
+from astropy.time import Time
+from astropy.time import TimeDelta
+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.exposure_conditions import AllowableOpticalDensityFilterNames
+from dkist_processing_cryonirsp.models.exposure_conditions import ExposureConditions
+from dkist_processing_cryonirsp.models.tags import CryonirspStemName
+from dkist_processing_cryonirsp.models.tags import CryonirspTag
+from dkist_processing_cryonirsp.tasks.ci_science import CIScienceCalibration
+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_fits_frame
+from dkist_processing_cryonirsp.tests.header_models import CryonirspHeadersValidObserveFrames
+
+# from dkist_processing_common.models.tags import Tag
+
+
+@pytest.fixture(scope="function", params=["Full Stokes", "Stokes-I"])
+def ci_science_calibration_task(
+    tmp_path,
+    recipe_run_id,
+    assign_input_dataset_doc_to_task,
+    init_cryonirsp_constants_db,
+    request,
+):
+    num_map_scans = 2
+    num_scan_steps = 2
+    num_meas = 1
+    exposure_time = 0.02  # From CryonirspHeadersValidObserveFrames fixture
+    exposure_conditions = ExposureConditions(
+        exposure_time, AllowableOpticalDensityFilterNames.OPEN.value
+    )
+    if request.param == "Full Stokes":
+        num_modstates = 2
+    else:
+        num_modstates = 1
+    array_shape = (1, 20, 20)
+    intermediate_shape = array_shape[1:]
+    dataset_shape = (num_map_scans * num_scan_steps * num_modstates,) + array_shape[1:]
+
+    constants_db = CryonirspConstantsDb(
+        ARM_ID="CI",
+        NUM_MODSTATES=num_modstates,
+        NUM_MAP_SCANS=num_map_scans,
+        NUM_SCAN_STEPS=num_scan_steps,
+        NUM_BEAMS=1,
+        OBSERVE_EXPOSURE_CONDITIONS_LIST=(exposure_conditions,),
+        MODULATOR_SPIN_MODE="Continuous" if request.param == "Full Stokes" else "Off",
+    )
+    init_cryonirsp_constants_db(recipe_run_id, constants_db)
+    with CIScienceCalibration(
+        recipe_run_id=recipe_run_id, workflow_name="ci_science_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
+            all_zeros = np.zeros(intermediate_shape)
+            all_ones = np.ones(intermediate_shape)
+
+            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())
+
+            # Need a beam boundary file
+            task.intermediate_frame_write_arrays(
+                arrays=np.array([0, intermediate_shape[0], 0, intermediate_shape[1]]),
+                task_tag=CryonirspTag.task_beam_boundaries(),
+                beam=1,
+            )
+            # Create fake bad pixel map
+            task.intermediate_frame_write_arrays(
+                arrays=np.zeros(array_shape[1:]), task_tag=CryonirspTag.task_bad_pixel_map()
+            )
+
+            # Create fake demodulation matrices
+            demod_matrices = np.zeros((1, 1, 4, num_modstates))
+            for modstate in range(num_modstates):
+                demod_matrices[0, 0, :, modstate] = [1, 2, 3, 4]
+            demod_hdul = fits.HDUList([fits.PrimaryHDU(data=demod_matrices)])
+            task.write(
+                data=demod_hdul,
+                tags=[
+                    CryonirspTag.intermediate(),
+                    CryonirspTag.frame(),
+                    CryonirspTag.task_demodulation_matrices(),
+                    CryonirspTag.beam(1),
+                ],
+                encoder=fits_hdulist_encoder,
+            )
+
+            # Create fake dark intermediate arrays
+            task.intermediate_frame_write_arrays(
+                all_zeros,
+                beam=1,
+                task_tag=CryonirspTag.task_dark(),
+                exposure_conditions=exposure_conditions,
+            )
+
+            # Create fake lamp and solar gain intermediate arrays
+            for modstate in range(1, num_modstates + 1):
+                gain_hdul = fits.HDUList([fits.PrimaryHDU(data=all_ones)])
+                task.write(
+                    data=gain_hdul,
+                    tags=[
+                        CryonirspTag.intermediate(),
+                        CryonirspTag.frame(),
+                        CryonirspTag.task_lamp_gain(),
+                        CryonirspTag.beam(1),
+                        CryonirspTag.modstate(modstate),
+                    ],
+                    encoder=fits_hdulist_encoder,
+                )
+                task.write(
+                    data=gain_hdul,
+                    tags=[
+                        CryonirspTag.intermediate(),
+                        CryonirspTag.frame(),
+                        CryonirspTag.task_solar_gain(),
+                        CryonirspTag.beam(1),
+                        CryonirspTag.modstate(modstate),
+                    ],
+                    encoder=fits_hdulist_encoder,
+                )
+
+            # Create fake observe arrays
+            start_time = datetime.now()
+            for map_scan in range(1, num_map_scans + 1):
+                for scan_step in range(1, num_scan_steps + 1):
+                    for modstate in range(1, num_modstates + 1):
+                        for meas_num in range(1, num_meas + 1):
+                            ds = CryonirspHeadersValidObserveFrames(
+                                dataset_shape=dataset_shape,
+                                array_shape=array_shape,
+                                time_delta=10,
+                                scan_step=scan_step,
+                                num_scan_steps=num_scan_steps,
+                                num_map_scans=num_map_scans,
+                                map_scan=map_scan,
+                                num_modstates=num_modstates,
+                                modstate=modstate,
+                                start_time=start_time,
+                                num_meas=num_meas,
+                                meas_num=meas_num,
+                                arm_id="CI",
+                            )
+                            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=array_shape
+                            )
+                            header = hdul[0].header
+                            task.write(
+                                data=hdul,
+                                tags=[
+                                    CryonirspTag.task_observe(),
+                                    CryonirspTag.scan_step(scan_step),
+                                    CryonirspTag.map_scan(map_scan),
+                                    CryonirspTag.modstate(modstate),
+                                    CryonirspTag.linearized(),
+                                    CryonirspTag.frame(),
+                                    CryonirspTag.exposure_conditions(exposure_conditions),
+                                    CryonirspTag.meas_num(meas_num),
+                                ],
+                                encoder=fits_hdulist_encoder,
+                            )
+            yield task, request.param, header, intermediate_shape
+        finally:
+            task._purge()
+
+
+@pytest.fixture(scope="session")
+def ci_headers_with_dates() -> tuple[list[fits.Header], str, int, int]:
+    num_headers = 5
+    start_time = "1969-12-06T18:00:00"
+    exp_time = 12
+    time_delta = 10
+    ds = CryonirspHeadersValidObserveFrames(
+        dataset_shape=(num_headers, 4, 4),
+        array_shape=(1, 4, 4),
+        time_delta=time_delta,
+        num_map_scans=1,
+        map_scan=1,
+        num_scan_steps=1,
+        scan_step=1,
+        num_meas=1,
+        meas_num=1,
+        num_modstates=num_headers,
+        modstate=1,
+        start_time=datetime.fromisoformat(start_time),
+        arm_id="CI",
+    )
+    headers = [
+        spec122_validator.validate_and_translate_to_214_l0(d.header(), return_type=fits.HDUList)[
+            0
+        ].header
+        for d in ds
+    ]
+    random.shuffle(headers)  # Shuffle to make sure they're not already in time order
+    for h in headers:
+        h["XPOSURE"] = exp_time  # Exposure time, in ms
+
+    return headers, start_time, exp_time, time_delta
+
+
+@pytest.fixture(scope="session")
+def ci_compressed_headers_with_dates(
+    ci_headers_with_dates,
+) -> tuple[list[fits.Header], str, int, int]:
+    headers, start_time, exp_time, time_delta = ci_headers_with_dates
+    comp_headers = [fits.hdu.compressed.CompImageHeader(h, h) for h in headers]
+    return comp_headers, start_time, exp_time, time_delta
+
+
+def test_ci_science_calibration_task(ci_science_calibration_task, mocker):
+    """
+    Given: A CIScienceCalibration task
+    When: Calling the task instance
+    Then: There are the expected number of science frames with the correct tags applied and the headers have been correctly updated
+    """
+
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+    )
+
+    # When
+    task, polarization_mode, og_header, og_single_beam_shape = ci_science_calibration_task
+    task()
+
+    # 1 from re-dummification
+    expected_final_shape = (1, og_single_beam_shape[0], og_single_beam_shape[1])
+
+    # Then
+    tags = [
+        CryonirspTag.calibrated(),
+        CryonirspTag.frame(),
+    ]
+    files = list(task.read(tags=tags))
+    if polarization_mode == "Full Stokes":
+        # 2 scan steps * 2 map scans * 4 stokes params = 16 frames
+        assert len(files) == 16
+    elif polarization_mode == "Stokes-I":
+        # 2 scan steps * 2 map scans * 1 stokes param = 4 frames
+        assert len(files) == 4
+    for file in files:
+        hdul = fits.open(file)
+        assert len(hdul) == 1
+        hdu = hdul[0]
+        assert type(hdul[0]) is fits.PrimaryHDU
+        assert hdu.data.shape == expected_final_shape
+        assert "DATE-BEG" in hdu.header.keys()
+        assert "DATE-END" in hdu.header.keys()
+        if polarization_mode == "Full Stokes":
+            assert "POL_NOIS" in hdu.header.keys()
+            assert "POL_SENS" in hdu.header.keys()
+
+        # Check that scan step keys were updated
+        scan_step = [
+            int(t.split("_")[-1]) for t in task.tags(file) if CryonirspStemName.scan_step.value in t
+        ][0]
+
+        assert hdu.header["CNNUMSCN"] == 2
+        assert hdu.header["CNCURSCN"] == scan_step
+
+    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_observe()]
+            )
+
+
+def test_compute_ci_date_keys(ci_headers_with_dates, recipe_run_id, init_cryonirsp_constants_db):
+    """
+    Given: A set of CI headers with different DATE-OBS values
+    When: Computing the time over which the headers were taken
+    Then: A header with correct DATE-BEG, DATE-END, and DATE-AVG keys is returned
+    """
+    headers, start_time, exp_time, time_delta = ci_headers_with_dates
+    constants_db = CryonirspConstantsDb()
+    init_cryonirsp_constants_db(recipe_run_id, constants_db)
+    with CIScienceCalibration(
+        recipe_run_id=recipe_run_id, workflow_name="science_calibration", workflow_version="VX.Y"
+    ) as task:
+        final_header = task.compute_date_keys(headers)
+        final_header_from_single = task.compute_date_keys(headers[0])
+
+    date_end = (
+        Time(start_time)
+        + (len(headers) - 1) * TimeDelta(time_delta, format="sec")
+        + TimeDelta(exp_time / 1000.0, format="sec")
+    ).isot
+
+    assert final_header["DATE-BEG"] == start_time
+    assert final_header["DATE-END"] == date_end
+
+    date_end_from_single = (
+        Time(headers[0]["DATE-BEG"])
+        # + TimeDelta(time_delta, format="sec")
+        + TimeDelta(exp_time / 1000.0, format="sec")
+    ).isot
+
+    assert final_header_from_single["DATE-BEG"] == headers[0]["DATE-BEG"]
+    assert final_header_from_single["DATE-END"] == date_end_from_single
+
+
+def test_compute_ci_date_keys_compressed_headers(
+    ci_compressed_headers_with_dates, recipe_run_id, init_cryonirsp_constants_db
+):
+    """
+    Given: A set of CI compressed headers with different DATE-OBS values
+    When: Computing the time over which the headers were taken
+    Then: A header with correct DATE-BEG, DATE-END, and DATE-AVG keys is returned
+    """
+    headers, start_time, exp_time, time_delta = ci_compressed_headers_with_dates
+    constants_db = CryonirspConstantsDb()
+    init_cryonirsp_constants_db(recipe_run_id, constants_db)
+    with CIScienceCalibration(
+        recipe_run_id=recipe_run_id, workflow_name="science_calibration", workflow_version="VX.Y"
+    ) as task:
+        final_header = task.compute_date_keys(headers)
+        final_header_from_single = task.compute_date_keys(headers[0])
+
+    date_end = (
+        Time(start_time)
+        + (len(headers) - 1) * TimeDelta(time_delta, format="sec")
+        + TimeDelta(exp_time / 1000.0, format="sec")
+    ).isot
+
+    assert final_header["DATE-BEG"] == start_time
+    assert final_header["DATE-END"] == date_end
+
+    date_end_from_single = (
+        Time(headers[0]["DATE-BEG"]) + TimeDelta(exp_time / 1000.0, format="sec")
+    ).isot
+
+    assert final_header_from_single["DATE-BEG"] == headers[0]["DATE-BEG"]
+    assert final_header_from_single["DATE-END"] == date_end_from_single