dkist-processing-trend 0.1.0__py3-none-any.whl

dkist_processing_trend/tests/test_parse.py

@@ -0,0 +1,276 @@
+import json
+
+import astropy.units as u
+import pytest
+from astropy.time.core import Time
+from astropy.time.core import TimeDelta
+from dkist_processing_common._util.scratch import WorkflowFileSystem
+from dkist_processing_common.models.constants import BudName
+
+from dkist_processing_trend.models.constants import TrendBudName
+from dkist_processing_trend.models.instrument import InstrumentName
+from dkist_processing_trend.models.parameters import TrendParsingParameters
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.tasks.parse import ParseL0TrendInputData
+from dkist_processing_trend.tests.conftest import ARM_ANGLE_FOR_ARM
+from dkist_processing_trend.tests.conftest import CRYO_OD_FILTER_FOR_ARM
+from dkist_processing_trend.tests.conftest import CRYO_SLIT_WIDTH_FOR_ARM
+from dkist_processing_trend.tests.conftest import GRATING_ANGLE_FOR_ARM
+from dkist_processing_trend.tests.conftest import GRATING_CONSTANT_FOR_ARM
+from dkist_processing_trend.tests.conftest import WAVELENGTH_FOR_ARM
+from dkist_processing_trend.tests.conftest import TrendInputDatasetParameterValues
+from dkist_processing_trend.tests.conftest import write_calibration_sequence_frames_to_task
+
+
+@pytest.fixture
+def input_parse_task(tmp_path, recipe_run_id, assign_input_dataset_doc_to_task):
+    with ParseL0TrendInputData(
+        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)
+        assign_input_dataset_doc_to_task(
+            task=task,
+            parameter_values=TrendInputDatasetParameterValues(),
+        )
+        task.parameters = TrendParsingParameters(scratch=task.scratch)
+
+        yield task
+        task._purge()
+
+
+@pytest.mark.parametrize(
+    "instrument, arm_id_list, op_id_list, expected_op_id_constant, experiment_id_list, expected_experiment_id_constant, proposal_id_list, expected_proposal_id_constant, hls_list, expected_hls_version_constant",
+    [
+        pytest.param(
+            InstrumentName.visp,
+            [1, 2, 3],
+            ["OP1"],
+            ["OP1"],
+            ["EX1"],
+            ["EX1"],
+            ["PROP1"],
+            ["PROP1"],
+            ["hlsv1"],
+            ["hlsv1"],
+            id="visp",
+        ),
+        pytest.param(
+            InstrumentName.cryonirsp,
+            ["CI", "SP"],
+            ["OP1"],
+            ["OP1"],
+            ["EX1"],
+            ["EX1"],
+            ["PROP1"],
+            ["PROP1"],
+            ["hlsv1"],
+            ["hlsv1"],
+            id="cryonirsp",
+        ),
+        pytest.param(
+            InstrumentName.dlnirsp,
+            ["VIS", "JBand", "HBand"],
+            ["OP1"],
+            ["OP1"],
+            ["EX1"],
+            ["EX1"],
+            ["PROP1"],
+            ["PROP1"],
+            ["hlsv1"],
+            ["hlsv1"],
+            id="dlnirsp",
+        ),
+        pytest.param(
+            InstrumentName.visp,
+            [1, 2],
+            ["OP2", "OP1", "OP2"],
+            ["OP2", "OP1"],
+            ["EX6", "EX6", "EX7"],
+            ["EX6", "EX7"],
+            ["PROP1", "PROP1", "PROP1"],
+            ["PROP1"],
+            ["hlsv1", "hlsv2", "hlsv2"],
+            ["hlsv2", "hlsv1"],
+            id="multi_OP",
+        ),
+    ],
+)
+def test_parse_input_data(
+    input_parse_task,
+    instrument,
+    arm_id_list,
+    op_id_list,
+    expected_op_id_constant,
+    experiment_id_list,
+    expected_experiment_id_constant,
+    proposal_id_list,
+    expected_proposal_id_constant,
+    hls_list,
+    expected_hls_version_constant,
+    visp_grating_angle,
+    visp_arm_position,
+    visp_grating_constant,
+    cryonirsp_od_filter_name,
+    cryonirsp_grating_angle,
+    cryonirsp_grating_littrow_angle,
+    cryonirsp_grating_constant,
+    cryonirsp_slit_width,
+    dlnirsp_grating_angle,
+    dlnirsp_arm_position,
+    dlnirsp_grating_constant,
+    small_calibration_sequence,
+    small_num_gain_steps,
+    small_num_dark_steps,
+    small_num_cs_steps,
+):
+    """
+    Given: A set of INPUT frames and a Parse task
+    When: Parsing the frames
+    Then: The frames are tagged correctly and constants are populated correctly
+    """
+
+    task = input_parse_task
+
+    pol_status, pol_theta, ret_status, ret_theta, dark_status = small_calibration_sequence
+    num_modstates = 3
+    num_ops = len(op_id_list)
+    start_time = Time("2026-02-01T15:00:12")
+
+    for i, (op_id, ex_id, prop_id, hls_v) in enumerate(
+        zip(op_id_list, experiment_id_list, proposal_id_list, hls_list)
+    ):
+        polcal_op_start_time = (start_time + i * TimeDelta(1 * u.day)).fits
+        for arm_id in arm_id_list:
+            write_calibration_sequence_frames_to_task(
+                task=task,
+                instrument=instrument,
+                arm_id=arm_id,
+                wavelength=WAVELENGTH_FOR_ARM[arm_id],
+                num_modstates=num_modstates,
+                pol_status=pol_status,
+                pol_theta=pol_theta,
+                ret_status=ret_status,
+                ret_theta=ret_theta,
+                dark_status=dark_status,
+                array_shape=(2, 2),
+                input_tags=[TrendTag.input()],
+                op_start_time=polcal_op_start_time,
+                op_id=op_id,
+                experiment_id=ex_id,
+                proposal_id=prop_id,
+                hls_version=hls_v,
+                visp_grating_angle=GRATING_ANGLE_FOR_ARM[arm_id],
+                visp_arm_position=ARM_ANGLE_FOR_ARM[arm_id],
+                visp_grating_constant=GRATING_CONSTANT_FOR_ARM[arm_id],
+                cryonirsp_od_filter_name=CRYO_OD_FILTER_FOR_ARM[arm_id],
+                cryonirsp_grating_angle=GRATING_ANGLE_FOR_ARM[arm_id],
+                cryonirsp_grating_littrow_angle=ARM_ANGLE_FOR_ARM[arm_id],
+                cryonirsp_grating_constant=GRATING_CONSTANT_FOR_ARM[arm_id],
+                cryonirsp_slit_width=CRYO_SLIT_WIDTH_FOR_ARM[arm_id],
+                dlnirsp_grating_angle=GRATING_ANGLE_FOR_ARM[arm_id],
+                dlnirsp_arm_position=ARM_ANGLE_FOR_ARM[arm_id],
+                dlnirsp_grating_constant=GRATING_CONSTANT_FOR_ARM[arm_id],
+            )
+
+    task()
+
+    # Tags applied correctly
+    for arm_id in arm_id_list:
+        for step in range(small_num_cs_steps * num_ops):
+            for modstate in range(1, num_modstates + 1):
+                assert (
+                    task.count(
+                        tags=[
+                            TrendTag.task_polcal(),
+                            TrendTag.arm_id(arm_id),
+                            TrendTag.cs_step(step),
+                            TrendTag.modstate(modstate),
+                        ]
+                    )
+                    == 1
+                )
+
+    for arm_id in arm_id_list:
+        for modstate in range(1, num_modstates + 1):
+            assert (
+                task.count(
+                    tags=[
+                        TrendTag.arm_id(arm_id),
+                        TrendTag.modstate(modstate),
+                        TrendTag.task_polcal_gain(),
+                    ]
+                )
+                == small_num_gain_steps * num_ops
+            )
+            assert (
+                task.count(
+                    tags=[
+                        TrendTag.arm_id(arm_id),
+                        TrendTag.modstate(modstate),
+                        TrendTag.task_polcal_dark(),
+                    ]
+                )
+                == small_num_dark_steps * num_ops
+            )
+
+    # Constants parsed correctly
+    # Wrap both side in `Time(...).fits` to ensure the same formatting of decimal places in string
+    assert (
+        Time(task.constants._db_dict[TrendBudName.earliest_ip_start_time]).fits == start_time.fits
+    )
+    assert task.constants._db_dict[TrendBudName.wavelength_for_arm] == {
+        json.dumps(arm): WAVELENGTH_FOR_ARM[arm] for arm in arm_id_list
+    }
+    assert task.constants._db_dict[BudName.instrument] == instrument
+    assert task.constants._db_dict[BudName.num_cs_steps] == small_num_cs_steps * num_ops
+    assert task.constants._db_dict[BudName.num_modstates] == num_modstates
+    assert sorted(task.constants._db_dict[TrendBudName.arm_id_list]) == sorted(arm_id_list)
+    assert (
+        task.constants._db_dict[TrendBudName.observing_program_id_list] == expected_op_id_constant
+    )
+    assert (
+        task.constants._db_dict[TrendBudName.experiment_id_list] == expected_experiment_id_constant
+    )
+    assert task.constants._db_dict[TrendBudName.proposal_id_list] == expected_proposal_id_constant
+    assert task.constants._db_dict[TrendBudName.hls_version_list] == expected_hls_version_constant
+
+    if instrument == InstrumentName.visp:
+        assert task.constants._db_dict[TrendBudName.visp_grating_angle_for_arm] == {
+            json.dumps(arm): GRATING_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+        }
+        assert task.constants._db_dict[TrendBudName.visp_arm_position_for_arm] == {
+            json.dumps(arm): ARM_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+        }
+        assert task.constants._db_dict[TrendBudName.visp_grating_constant_for_arm] == {
+            json.dumps(arm): GRATING_CONSTANT_FOR_ARM[arm] for arm in arm_id_list
+        }
+
+    if instrument == InstrumentName.cryonirsp:
+        assert task.constants._db_dict[TrendBudName.cryonirsp_od_filter_name_for_arm] == {
+            json.dumps(arm): CRYO_OD_FILTER_FOR_ARM[arm] for arm in arm_id_list
+        }
+        assert task.constants._db_dict[TrendBudName.cryonirsp_grating_angle_for_arm] == {
+            json.dumps(arm): GRATING_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+        }
+        assert task.constants._db_dict[TrendBudName.cryonirsp_grating_littrow_angle_for_arm] == {
+            json.dumps(arm): ARM_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+        }
+        assert task.constants._db_dict[TrendBudName.cryonirsp_grating_constant_for_arm] == {
+            json.dumps(arm): GRATING_CONSTANT_FOR_ARM[arm] for arm in arm_id_list
+        }
+        assert task.constants._db_dict[TrendBudName.cryonirsp_slit_width_for_arm] == {
+            json.dumps(arm): CRYO_SLIT_WIDTH_FOR_ARM[arm] for arm in arm_id_list
+        }
+
+    if instrument == InstrumentName.dlnirsp:
+        assert task.constants._db_dict[TrendBudName.dlnirsp_grating_angle_for_arm] == {
+            json.dumps(arm): GRATING_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+        }
+        assert task.constants._db_dict[TrendBudName.dlnirsp_arm_position_for_arm] == {
+            json.dumps(arm): ARM_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+        }
+        assert task.constants._db_dict[TrendBudName.dlnirsp_grating_constant_for_arm] == {
+            json.dumps(arm): GRATING_CONSTANT_FOR_ARM[arm] for arm in arm_id_list
+        }

dkist_processing_trend/tests/test_prep_fit_data_base.py

@@ -0,0 +1,233 @@
+import numpy as np
+import pytest
+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_array_encoder
+
+from dkist_processing_trend.models.fits_access import CryonirspMetadataKey
+from dkist_processing_trend.models.instrument import InstrumentName
+from dkist_processing_trend.models.instrument_options import CryonirspInstrumentOptions
+from dkist_processing_trend.models.parameters import TrendParameters
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.tasks.prepare_fit_data_base import PrepareFitDataBase
+from dkist_processing_trend.tests.conftest import ModulatedCSStepHeaders
+from dkist_processing_trend.tests.conftest import TrendInputDatasetParameterValues
+from dkist_processing_trend.tests.conftest import TrendTestingConstantsDatabase
+from dkist_processing_trend.tests.conftest import write_calibration_sequence_frames_to_task
+
+
+class InstrumentPrepData(PrepareFitDataBase):
+
+    def extract_beam(self, array: np.ndarray, beam: int) -> np.ndarray:
+        return array[beam - 1 :: 2, :]
+
+    def apply_global_instrument_options(
+        self, array: np.ndarray, instrument_options: CryonirspInstrumentOptions
+    ) -> np.ndarray:
+        return np.mean(array)[None, None]
+
+    def apply_local_instrument_options(
+        self, array: np.ndarray, instrument_options: CryonirspInstrumentOptions
+    ) -> np.ndarray:
+        return np.mean(array, axis=instrument_options.num_spatial_px)
+
+
+@pytest.fixture
+def instrument_options_list() -> dict[str, str | int]:
+    # pydantic forces the structure of these dicts. See `apply_local_instrument_options` above for how these parameters
+    # are actually used.
+    return [
+        {"name": "axis0_option", "num_spatial_px": 0, "num_spectral_px": 1},
+        {"name": "axis1_option", "num_spatial_px": 1, "num_spectral_px": 1},
+    ]
+
+
+@pytest.fixture
+def num_modstates() -> int:
+    return 4
+
+
+@pytest.fixture
+def dark_signal() -> float:
+    return 6.28
+
+
+@pytest.fixture
+def gain_signal() -> float:
+    return 3.14
+
+
+@pytest.fixture
+def make_data_function(dark_signal, gain_signal):
+    def make_data(frame: ModulatedCSStepHeaders) -> np.ndarray:
+        shape = frame.array_shape[1:]
+        cs_step = frame.cs_step_num
+        modstate = frame.current_modstate
+
+        true_signal = cs_step * 1000 + modstate * 100
+        raw_signal = true_signal * gain_signal + dark_signal
+        return np.full(shape, raw_signal)
+
+    return make_data
+
+
+@pytest.fixture
+def instrument_fit_data_task(
+    recipe_run_id,
+    tmp_path,
+    link_constants_db,
+    assign_input_dataset_doc_to_task,
+    arm_id,
+    instrument_options_list,
+    small_num_cs_steps,
+    num_modstates,
+):
+    link_constants_db(
+        recipe_run_id,
+        TrendTestingConstantsDatabase(
+            INSTRUMENT=InstrumentName.cryonirsp,
+            ARM_ID_LIST=["CI", "SP"],
+            NUM_CS_STEPS=small_num_cs_steps,
+            NUM_MODSTATES=num_modstates,
+        ),
+    )
+    with InstrumentPrepData(
+        arm_id=arm_id, recipe_run_id=recipe_run_id, workflow_name="name", workflow_version="version"
+    ) as task:
+        task.scratch = WorkflowFileSystem(recipe_run_id=recipe_run_id, scratch_base_path=tmp_path)
+        assign_input_dataset_doc_to_task(
+            task,
+            TrendInputDatasetParameterValues(
+                trend_instrument_processing_options_cryonirsp=instrument_options_list
+            ),
+        )
+        task.parameters = TrendParameters(
+            scratch=task.scratch,
+            obs_ip_start_time=task.constants.earliest_ip_start_time,
+            instrument=task.constants.instrument,
+        )
+
+        yield task
+        task._purge()
+
+
+@pytest.mark.parametrize("arm_id", [pytest.param("SP"), pytest.param("CI")])
+def test_prepare_data_base(
+    instrument_fit_data_task,
+    instrument_options_list,
+    small_calibration_sequence,
+    small_num_cs_steps,
+    num_modstates,
+    arm_id,
+    make_data_function,
+    dark_signal,
+    gain_signal,
+    mocker,
+    fake_gql_client,
+):
+    """
+    Given: A dummy child class of `PrepareFitDataBase` and a correct set of INPUT and INTERMEDIATE data.
+    When: Running the task
+    Then: The expected outputs are produced and they have the expected structure and data values
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
+    )
+    num_beams = 1 if arm_id == "CI" else 2
+    shape = (8, 9)
+    task = instrument_fit_data_task
+
+    pol_status, pol_theta, ret_status, ret_theta, dark_status = small_calibration_sequence
+
+    write_calibration_sequence_frames_to_task(
+        task=task,
+        instrument=InstrumentName.cryonirsp,
+        arm_id=arm_id,
+        wavelength=1074.7,
+        num_modstates=num_modstates,
+        pol_status=pol_status,
+        pol_theta=pol_theta,
+        ret_status=ret_status,
+        ret_theta=ret_theta,
+        dark_status=dark_status,
+        array_shape=shape,
+        tags=[TrendTag.arm_id(arm_id)],
+        data_func=make_data_function,
+    )
+
+    # INTERMEIDATE dark and gain frames
+    task.write(
+        data=np.full(shape, dark_signal),
+        tags=[
+            TrendTag.intermediate(),
+            TrendTag.frame(),
+            TrendTag.arm_id(arm_id),
+            TrendTag.task_dark(),
+        ],
+        encoder=fits_array_encoder,
+    )
+    task.write(
+        data=np.full(shape, gain_signal),
+        tags=[
+            TrendTag.intermediate(),
+            TrendTag.frame(),
+            TrendTag.arm_id(arm_id),
+            TrendTag.task_gain(),
+        ],
+        encoder=fits_array_encoder,
+    )
+
+    task()
+
+    for inst_opt in instrument_options_list:
+
+        # See `InstrumentPrepData.apply_local_instrument_options` above for why these values are what they are
+        if inst_opt["num_spatial_px"] == 0:
+            expected_local_shape = (9,)
+        elif inst_opt["num_spatial_px"] == 1:
+            expected_local_shape = (4,)
+
+        for beam in range(1, num_beams + 1):
+            global_tags = [
+                TrendTag.intermediate(),
+                TrendTag.arm_id(arm_id),
+                TrendTag.beam(beam),
+                TrendTag.instrument_processing_options(inst_opt["name"]),
+                TrendTag.task_global_pac_input(),
+            ]
+            local_tags = [
+                TrendTag.intermediate(),
+                TrendTag.arm_id(arm_id),
+                TrendTag.beam(beam),
+                TrendTag.instrument_processing_options(inst_opt["name"]),
+                TrendTag.task_local_pac_input(),
+            ]
+            assert task.count(global_tags) == 1
+            assert task.count(local_tags) == 1
+
+            global_dict = next(task.read(tags=global_tags, decoder=asdf_decoder))
+            local_dict = next(task.read(tags=local_tags, decoder=asdf_decoder))
+            for cs_step in range(small_num_cs_steps):
+                global_modstate_list = global_dict[cs_step]
+                local_modstate_list = local_dict[cs_step]
+                assert isinstance(global_modstate_list, list)
+                assert isinstance(local_modstate_list, list)
+                assert len(global_modstate_list) == len(local_modstate_list) == num_modstates
+                for m in range(num_modstates):
+                    global_mod_dict = global_modstate_list[m]
+                    global_header = global_mod_dict["header"]
+                    assert global_header[CryonirspMetadataKey.modulator_state] == m + 1
+                    global_data = global_mod_dict["data"]
+                    assert global_data.shape == (1, 1)
+                    np.testing.assert_array_equal(global_data, 1000 * cs_step + 100 * (m + 1))
+
+                    local_mod_dict = local_modstate_list[m]
+                    local_header = local_mod_dict["header"]
+                    assert local_header[CryonirspMetadataKey.modulator_state] == m + 1
+                    local_data = local_mod_dict["data"]
+                    assert local_data.shape == expected_local_shape
+                    np.testing.assert_array_equal(local_data, 1000 * cs_step + 100 * (m + 1))
+
+            # These asserts are checking that there aren't any extra CS steps beyond what we expect
+            assert small_num_cs_steps not in global_dict.keys()
+            assert small_num_cs_steps not in local_dict.keys()

dkist_processing_trend/tests/test_publish_catalog_messages.py

@@ -0,0 +1,45 @@
+from pathlib import Path
+
+import pytest
+
+from dkist_processing_trend.tasks.trend_output_data import PublishTrendCatalogMessages
+from dkist_processing_trend.tests.conftest import TrendTestingConstantsDatabase
+
+
+@pytest.fixture
+def publish_catalog_and_quality_messages_task(
+    recipe_run_id, link_constants_db, mocker, fake_gql_client
+):
+    instrument = "DOPE"
+    full_date = "2024-01-02T12:34:56.2341234234"
+    dir_date = "20240102T123456"
+    link_constants_db(
+        recipe_run_id,
+        TrendTestingConstantsDatabase(INSTRUMENT=instrument, EARLIEST_IP_START_TIME=full_date),
+    )
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
+    )
+    with PublishTrendCatalogMessages(
+        recipe_run_id=recipe_run_id,
+        workflow_name="workflow_name",
+        workflow_version="workflow_version",
+    ) as task:
+        yield task, instrument, dir_date
+        task.constants._purge()
+
+
+def test_frame_messages(publish_catalog_and_quality_messages_task):
+    """
+    Given: a PublishCatalogAndQualityMessages task
+    When: creating frame messages
+    Then: the attributes are correctly populated
+    """
+    task, instrument, dir_date = publish_catalog_and_quality_messages_task
+    filenames = [f"test_frame_{i}.ext" for i in range(10)]
+    filepaths = [Path(f"a/b/c/{filename}") for filename in filenames]
+    frame_messages = task.frame_messages(paths=filepaths)
+    for message in frame_messages:
+        assert message.body.bucket == "data"
+        assert str(Path(message.body.objectName).parent) == f"trend/polcal/{dir_date}/{instrument}"
+        assert message.body.conversationId == str(task.recipe_run_id)