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

dkist_processing_cryonirsp/tests/test_quality.py

@@ -0,0 +1,203 @@
+import json
+
+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_array_encoder
+from dkist_processing_common.codecs.fits import fits_hdulist_encoder
+from dkist_processing_common.models.task_name import TaskName
+from dkist_processing_common.tests.conftest import FakeGQLClient
+
+from dkist_processing_cryonirsp.models.tags import CryonirspTag
+from dkist_processing_cryonirsp.tasks.quality_metrics import CryonirspL0QualityMetrics
+from dkist_processing_cryonirsp.tasks.quality_metrics import CryonirspL1QualityMetrics
+from dkist_processing_cryonirsp.tests.conftest import CryonirspConstantsDb
+from dkist_processing_cryonirsp.tests.conftest import generate_214_l1_fits_frame
+from dkist_processing_cryonirsp.tests.header_models import Cryonirsp122ObserveFrames
+from dkist_processing_cryonirsp.tests.header_models import CryonirspHeaders
+
+
+@pytest.fixture
+def cryonirsp_l0_quality_task(recipe_run_id, num_modstates, init_cryonirsp_constants_db, tmp_path):
+    constants = CryonirspConstantsDb(
+        NUM_MODSTATES=num_modstates,
+        MODULATOR_SPIN_MODE="Stepped" if num_modstates > 1 else "Off",
+    )
+    init_cryonirsp_constants_db(recipe_run_id, constants)
+
+    with CryonirspL0QualityMetrics(
+        recipe_run_id=recipe_run_id,
+        workflow_name="workflow_name",
+        workflow_version="workflow_version",
+    ) as task:
+        task.scratch = WorkflowFileSystem(scratch_base_path=tmp_path, recipe_run_id=recipe_run_id)
+
+        yield task
+        task._purge()
+
+
+@pytest.fixture()
+def l0_quality_task_types() -> list[str]:
+    # The tasks types we want to build l0 metrics for
+    return [TaskName.lamp_gain.value, TaskName.dark.value]
+
+
+@pytest.fixture()
+def dataset_task_types(l0_quality_task_types) -> list[str]:
+    # The task types that exist in the dataset. I.e., a larger set than we want to build metrics for.
+    return l0_quality_task_types + [TaskName.solar_gain.value, TaskName.observe.value]
+
+
+@pytest.fixture
+def write_l0_task_frames_to_task(num_modstates, dataset_task_types):
+    array_shape = (1, 4, 4)
+    data = np.ones(array_shape)
+
+    def writer(task):
+        for task_type in dataset_task_types:
+            ds = CryonirspHeaders(
+                dataset_shape=(num_modstates,) + array_shape[-2:],
+                array_shape=array_shape,
+                time_delta=1.0,
+                file_schema="level0_spec214",
+            )
+            for modstate, frame in enumerate(ds, start=1):
+                header = frame.header()
+                task.write(
+                    data=data,
+                    header=header,
+                    tags=[
+                        CryonirspTag.linearized(),
+                        CryonirspTag.frame(),
+                        CryonirspTag.task(task_type),
+                        CryonirspTag.modstate(modstate),
+                    ],
+                    encoder=fits_array_encoder,
+                )
+
+    return writer
+
+
+@pytest.fixture(scope="function")
+def cryo_l1_quality_task(tmp_path, recipe_run_id, init_cryonirsp_constants_db):
+    num_map_scans = 3
+    num_scan_steps = 1
+    constants_db = CryonirspConstantsDb(
+        NUM_MAP_SCANS=num_map_scans,
+        NUM_SCAN_STEPS=num_scan_steps,
+    )
+    init_cryonirsp_constants_db(recipe_run_id, constants_db)
+    with CryonirspL1QualityMetrics(
+        recipe_run_id=recipe_run_id, workflow_name="science_calibration", workflow_version="VX.Y"
+    ) as task:
+        task.scratch = WorkflowFileSystem(scratch_base_path=tmp_path)
+
+        # Create fake stokes frames
+        for map_scan in range(1, num_map_scans + 1):
+            for scan_step in range(1, num_scan_steps + 1):
+                for stokes_param, index in zip(("I", "Q", "U", "V"), (1, 2, 3, 4)):
+                    ds = Cryonirsp122ObserveFrames(
+                        array_shape=(1, 10, 10),
+                        num_steps=num_scan_steps,
+                        num_map_scans=num_map_scans,
+                    )
+                    header_generator = (
+                        spec122_validator.validate_and_translate_to_214_l0(
+                            d.header(), return_type=fits.HDUList
+                        )[0].header
+                        for d in ds
+                    )
+
+                    hdul = generate_214_l1_fits_frame(s122_header=next(header_generator))
+                    hdul[1].header["DINDEX5"] = index
+                    task.write(
+                        data=hdul,
+                        tags=[
+                            CryonirspTag.calibrated(),
+                            CryonirspTag.frame(),
+                            CryonirspTag.stokes(stokes_param),
+                            CryonirspTag.scan_step(scan_step),
+                            CryonirspTag.map_scan(map_scan),
+                            CryonirspTag.meas_num(1),
+                        ],
+                        encoder=fits_hdulist_encoder,
+                    )
+
+        yield task
+        task._purge()
+
+
+@pytest.mark.parametrize(
+    "num_modstates", [pytest.param(4, id="polarimetric"), pytest.param(1, id="intensity")]
+)
+def test_l0_quality_task(
+    cryonirsp_l0_quality_task, num_modstates, write_l0_task_frames_to_task, l0_quality_task_types
+):
+    """
+    Given: A `CryonirspL0QualityMetrics` task and some INPUT frames tagged with their task type and modstate
+    When: Running the task
+    Then: The expected L0 quality metric files exist
+    """
+    # NOTE: We rely on the unit tests in `*-common` to verify the correct format/data of the metric files
+    task = cryonirsp_l0_quality_task
+    write_l0_task_frames_to_task(task)
+
+    task()
+
+    task_metric_names = ["FRAME_RMS", "FRAME_AVERAGE"]
+    for metric_name in task_metric_names:
+        for modstate in range(1, num_modstates + 1):
+            for task_type in l0_quality_task_types:
+                tags = [CryonirspTag.quality(metric_name), CryonirspTag.quality_task(task_type)]
+                if num_modstates > 1:
+                    tags.append(CryonirspTag.modstate(modstate))
+                files = list(task.read(tags))
+                assert len(files) == 1
+
+    global_metric_names = ["DATASET_AVERAGE", "DATASET_RMS"]
+    for metric_name in global_metric_names:
+        files = list(task.read(tags=[CryonirspTag.quality(metric_name)]))
+        assert len(files) > 0
+
+
+def test_l1_quality_task(cryo_l1_quality_task, mocker):
+    """
+    Given: A CryonirspL1QualityMetrics task
+    When: Calling the task instance
+    Then: A single sensitivity measurement and datetime is recorded for each map scan for each Stokes Q, U, and V,
+            and a single noise measurement and datetime is recorded for L1 file for each Stokes Q, U, and V
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+    )
+    # When
+    task = cryo_l1_quality_task
+    task()
+    # Then
+    num_map_scans = task.constants.num_map_scans
+    num_steps = task.constants.num_scan_steps
+    sensitivity_files = list(task.read(tags=[CryonirspTag.quality("SENSITIVITY")]))
+    assert len(sensitivity_files) == 4
+    for file in sensitivity_files:
+        with file.open() as f:
+            data = json.load(f)
+            assert isinstance(data, dict)
+            for time in data["x_values"]:
+                assert isinstance(time, str)
+            for sensitivity in data["y_values"]:
+                assert isinstance(sensitivity, float)
+            assert len(data["x_values"]) == len(data["y_values"]) == num_map_scans
+
+    noise_files = list(task.read(tags=[CryonirspTag.quality("NOISE")]))
+    assert len(noise_files) == 4
+    for file in noise_files:
+        with file.open() as f:
+            data = json.load(f)
+            assert isinstance(data, dict)
+            for time in data["x_values"]:
+                assert isinstance(time, str)
+            for noise in data["y_values"]:
+                assert isinstance(noise, float)
+            assert len(data["x_values"]) == len(data["y_values"]) == num_map_scans * num_steps

dkist_processing_cryonirsp/tests/test_sp_beam_boundaries.py

@@ -0,0 +1,112 @@
+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.sp_beam_boundaries import SPBeamBoundariesCalibration
+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 CryonirspHeadersValidSPSolarGainFrames
+
+
+@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 = "SP"
+    dataset_shape = (1, 100, 200)
+    array_shape = (1, 100, 200)
+    constants_db = CryonirspConstantsDb(ARM_ID=arm_id)
+    init_cryonirsp_constants_db(recipe_run_id, constants_db)
+    with SPBeamBoundariesCalibration(
+        recipe_run_id=recipe_run_id,
+        workflow_name="sp_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 = CryonirspHeadersValidSPSolarGainFrames(
+                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, v_min:v_max, h_min:h_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 - 1
+            plus_streak_pos = 3 * array_shape[2] // 4 + 1
+            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 SPBeamBoundariesCalibration 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([7, 94, 10, 88])
+    beam_2_tags = [CryonirspTag.task_beam_boundaries(), CryonirspTag.beam(2)]
+    beam_2_boundary = np.array([8, 95, 112, 190])
+    files_found = list(task.read(tags=beam_1_tags))
+    assert len(files_found) == 1
+    array1 = fits.open(files_found[0])[0].data
+    assert np.array_equal(array1, beam_1_boundary)
+    files_found = list(task.read(tags=beam_2_tags))
+    assert len(files_found) == 1
+    array2 = fits.open(files_found[0])[0].data
+    assert np.array_equal(array2, beam_2_boundary)

dkist_processing_cryonirsp/tests/test_sp_dispersion_axis_correction.py

@@ -0,0 +1,155 @@
+from datetime import datetime
+
+import numpy as np
+import pytest
+from astropy import units as u
+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.asdf import asdf_decoder
+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.sp_dispersion_axis_correction import (
+    SPDispersionAxisCorrection,
+)
+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 CryonirspHeadersValidSPSolarGainFrames
+
+
+@pytest.fixture(scope="function")
+def sp_dispersion_axis_correction_task(
+    tmp_path,
+    recipe_run_id,
+    assign_input_dataset_doc_to_task,
+    init_cryonirsp_constants_db,
+):
+    num_map_scans = 2
+    num_beams = 2
+    num_scan_steps = 2
+    num_modstates = 2
+    array_shape = (1, 30, 60)
+    char_spec_shape = (916,)
+    grating_position_deg = 62.505829779431224
+    grating_littrow_angle = -5.5
+    grating_constant = 31.6
+    dataset_shape = (num_beams * num_map_scans * num_scan_steps * num_modstates,) + array_shape[1:]
+
+    constants_db = CryonirspConstantsDb(
+        NUM_BEAMS=num_beams,
+        GRATING_POSITION_DEG=grating_position_deg,
+        GRATING_LITTROW_ANGLE_DEG=grating_littrow_angle,
+        GRATING_CONSTANT=grating_constant,
+        SOLAR_GAIN_IP_START_TIME="2021-01-01T00:00:00",
+    )
+    init_cryonirsp_constants_db(recipe_run_id, constants_db)
+    with SPDispersionAxisCorrection(
+        recipe_run_id=recipe_run_id,
+        workflow_name="sp_dispersion_axis_correction",
+        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_ones = np.ones(char_spec_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())
+
+            # Create fake solar charcteristic spectra
+            for beam in range(1, num_beams + 1):
+                char_spec_hdul = fits.HDUList([fits.PrimaryHDU(data=all_ones)])
+                task.write(
+                    data=char_spec_hdul,
+                    tags=[
+                        CryonirspTag.intermediate(),
+                        CryonirspTag.frame(),
+                        CryonirspTag.task_characteristic_spectra(),
+                        CryonirspTag.beam(beam),
+                    ],
+                    encoder=fits_hdulist_encoder,
+                )
+
+                # And a beam border intermediate array
+                task.intermediate_frame_write_arrays(
+                    arrays=np.array([0, 30, ((beam - 1) * 30), (30 + (beam - 1) * 30)]),
+                    task_tag=CryonirspTag.task_beam_boundaries(),
+                    beam=beam,
+                )
+
+                # Create fake linearized solar gain array with headers
+                start_time = datetime.now()
+                ds = CryonirspHeadersValidSPSolarGainFrames(
+                    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)
+                header = hdul[0].header
+                task.write(
+                    data=hdul,
+                    tags=[
+                        CryonirspTag.linearized(),
+                        CryonirspTag.frame(),
+                        CryonirspTag.task_solar_gain(),
+                        CryonirspTag.beam(beam),
+                    ],
+                    encoder=fits_hdulist_encoder,
+                )
+
+            yield task, header
+        except:
+            raise
+        finally:
+            task._purge()
+
+
+def test_sp_dispersion_axis_correction(sp_dispersion_axis_correction_task, mocker):
+    """
+    Given: A SPDispersionAxisCorrection task
+    When: Calling the task instance
+    Then: There are the expected number of intermediate fit frames with the correct tags applied and the values have been correctly fit
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=FakeGQLClient
+    )
+
+    task, original_header = sp_dispersion_axis_correction_task
+    task()
+    tags = [
+        CryonirspTag.task_spectral_fit(),
+        CryonirspTag.intermediate(),
+    ]
+    fit_result = next(task.read(tags=tags, decoder=asdf_decoder))
+    del fit_result["asdf_library"]
+    del fit_result["history"]
+
+    expected_wavelength_vector = task.compute_expected_wavelength_vector(
+        np.random.rand(100), original_header
+    )
+
+    (
+        telluric_atlas_wave,
+        telluric_atlas_trans,
+        solar_atlas_wave_air,
+        solar_atlas_trans_flipped,
+    ) = task.load_and_resample_fts_atlas(expected_wavelength_vector)
+
+    assert len(telluric_atlas_wave) == len(telluric_atlas_trans)
+    assert len(solar_atlas_wave_air) == len(solar_atlas_trans_flipped)
+
+    # make sure that the values have changed
+    assert original_header["CRVAL1"] != fit_result["CRVAL1"]
+    assert original_header["CDELT1"] != fit_result["CDELT1"]
+    assert original_header["CRVAL1A"] != fit_result["CRVAL1A"]
+    assert original_header["CDELT1A"] != fit_result["CDELT1A"]