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

dkist_processing_trend/tests/test_workflows.py

@@ -0,0 +1,10 @@
+"""Test integrity of workflows."""
+
+from dkist_processing_core.build_utils import validate_workflows
+
+from dkist_processing_trend import workflows
+
+
+def test_workflow_integrity():
+    """Validate workflow to ensure acyclic-ness and export compilation"""
+    validate_workflows(workflows)

dkist_processing_trend/tests/test_write_trend.py

@@ -0,0 +1,460 @@
+import json
+
+import numpy as np
+import pytest
+from astropy.io import fits
+from dkist_processing_common._util.scratch import WorkflowFileSystem
+from dkist_processing_common.codecs.fits import fits_array_encoder
+
+from dkist_processing_trend.models.instrument import InstrumentName
+from dkist_processing_trend.models.parameters import TrendParameters
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.tasks.write_trend import WriteTrendOutputs
+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 TrendTestingConstantsDatabase
+
+
+@pytest.fixture(scope="session")
+def num_polcal_ops() -> int:
+    return 1
+
+
+@pytest.fixture(scope="session")
+def num_cs_steps() -> int:
+    return 3
+
+
+@pytest.fixture(scope="session")
+def num_modstates() -> int:
+    return 4
+
+
+@pytest.fixture(scope="session")
+def fov_shape() -> tuple[int, int]:
+    return (12, 11)
+
+
+def num_beams_from_arm_id(arm_id) -> int:
+    if arm_id == "CI":
+        return 1
+    return 2
+
+
+@pytest.fixture(scope="session")
+def fit_options_list() -> list[dict[str, str | bool]]:
+    return [
+        {
+            "name": "option1",
+            "remove_I_trend": True,
+            "fit_mode_name": "use_M12",
+            "init_set_name": "OCCal_VIS",
+        },
+        {
+            "name": "option2",
+            "remove_I_trend": False,
+            "fit_mode_name": "use_M12_I_sys_per_step",
+            "init_set_name": "OCCal_VIS",
+        },
+    ]
+
+
+@pytest.fixture(scope="session")
+def visp_instrument_options_list() -> list[dict]:
+    return [
+        {
+            "name": "full",
+            "num_spatial_px": 2560,
+            "num_spectral_px": 1000,
+            "spatial_smoothing_fit_order": 22,
+        },
+        {
+            "name": "single",
+            "num_spatial_px": 1,
+            "num_spectral_px": 1,
+            "spatial_smoothing_fit_order": 1,
+        },
+    ]
+
+
+@pytest.fixture(scope="session")
+def cryonirsp_instrument_options_list() -> list[dict]:
+    return [
+        {"name": "full", "num_spatial_px": 1000, "num_spectral_px": 1000},
+        {"name": "single", "num_spatial_px": 1, "num_spectral_px": 1},
+    ]
+
+
+@pytest.fixture(scope="session")
+def dlnirsp_instrument_options_list() -> list[dict]:
+    return [
+        {
+            "name": "full",
+            "num_spatial_px": 1280,
+            "num_spectral_px": 1,
+            "spatial_smoothing_fit_order": 13,
+        },
+        {
+            "name": "single",
+            "num_spatial_px": 1,
+            "num_spectral_px": 1,
+            "spatial_smoothing_fit_order": 1,
+        },
+    ]
+
+
+@pytest.fixture
+def write_trend_task(
+    tmp_path,
+    recipe_run_id,
+    link_constants_db,
+    assign_input_dataset_doc_to_task,
+    num_cs_steps,
+    num_modstates,
+    instrument,
+    fit_options_list,
+    visp_instrument_options_list,
+    cryonirsp_instrument_options_list,
+    dlnirsp_instrument_options_list,
+    arm_id_list,
+):
+    link_constants_db(
+        recipe_run_id,
+        TrendTestingConstantsDatabase(
+            INSTRUMENT=instrument,
+            NUM_CS_STEPS=num_cs_steps,
+            NUM_MODSTATES=num_modstates,
+            ARM_ID_LIST=arm_id_list,
+            WAVELENGTH_FOR_ARM={json.dumps(arm): WAVELENGTH_FOR_ARM[arm] for arm in arm_id_list},
+            VISP_GRATING_ANGLE_FOR_ARM={
+                json.dumps(arm): GRATING_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+            },
+            VISP_ARM_POSITION_FOR_ARM={
+                json.dumps(arm): ARM_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+            },
+            VISP_GRATING_CONSTANT_FOR_ARM={
+                json.dumps(arm): GRATING_CONSTANT_FOR_ARM[arm] for arm in arm_id_list
+            },
+            CRYONIRSP_GRATING_ANGLE_FOR_ARM={
+                json.dumps(arm): GRATING_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+            },
+            CRYONIRSP_GRATING_LITTROW_ANGLE_FOR_ARM={
+                json.dumps(arm): ARM_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+            },
+            CRYONIRSP_GRATING_CONSTANT_FOR_ARM={
+                json.dumps(arm): GRATING_CONSTANT_FOR_ARM[arm] for arm in arm_id_list
+            },
+            CRYONIRSP_OD_FILTER_NAME_FOR_ARM={
+                json.dumps(arm): CRYO_OD_FILTER_FOR_ARM[arm] for arm in arm_id_list
+            },
+            CRYONIRSP_SLIT_WIDTH_FOR_ARM={
+                json.dumps(arm): CRYO_SLIT_WIDTH_FOR_ARM[arm] for arm in arm_id_list
+            },
+            DLNIRSP_GRATING_ANGLE_FOR_ARM={
+                json.dumps(arm): GRATING_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+            },
+            DLNIRSP_ARM_POSITION_FOR_ARM={
+                json.dumps(arm): ARM_ANGLE_FOR_ARM[arm] for arm in arm_id_list
+            },
+            DLNIRSP_GRATING_CONSTANT_FOR_ARM={
+                json.dumps(arm): GRATING_CONSTANT_FOR_ARM[arm] for arm in arm_id_list
+            },
+        ),
+    )
+    with WriteTrendOutputs(
+        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_fit_options_list=fit_options_list,
+                trend_instrument_processing_options_visp=visp_instrument_options_list,
+                trend_instrument_processing_options_cryonirsp=cryonirsp_instrument_options_list,
+                trend_instrument_processing_options_dlnirsp=dlnirsp_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()
+
+
+def write_pre_output_files_to_task(
+    task: WriteTrendOutputs,
+    raw_calibration_sequence_data_shape: tuple[int, int, int],
+    raw_fit_parameter_shape: tuple[int, ...],
+    raw_demod_shape: tuple[int, ...],
+    raw_flux_shape: tuple[int, ...],
+    raw_dmpd_shape: tuple[int, ...],
+):
+
+    inst_options_list = task.parameters.instrument_processing_options
+    fit_options_list = task.parameters.fit_options_list
+
+    rng = np.random.default_rng()
+
+    for arm_id in task.constants.arm_id_list:
+        # Calibration sequence
+        data = rng.random(raw_calibration_sequence_data_shape, dtype=np.float64)
+        task.write(
+            data=data,
+            tags=[
+                TrendTag.intermediate(),
+                TrendTag.arm_id(arm_id),
+                TrendTag.task_calibration_sequence(),
+            ],
+            encoder=fits_array_encoder,
+        )
+
+        for inst_opt in inst_options_list:
+            inst_opt_name = inst_opt.name
+            for fit_opt in fit_options_list:
+                fit_opt_name = fit_opt.name
+
+                # DMPD doesn't depend on beam
+                data = rng.random(raw_dmpd_shape, dtype=np.float64)
+                task.write(
+                    data=data,
+                    tags=[
+                        TrendTag.intermediate(),
+                        TrendTag.arm_id(arm_id),
+                        TrendTag.instrument_processing_options(inst_opt_name),
+                        TrendTag.pac_fit_options(fit_opt_name),
+                        TrendTag.task_demodulated_polcal_data(),
+                    ],
+                    encoder=fits_array_encoder,
+                )
+
+                for beam in range(1, num_beams_from_arm_id(arm_id) + 1):
+                    # Fit parameters
+                    data = rng.random(raw_fit_parameter_shape, dtype=np.float64)
+                    task.write(
+                        data=data,
+                        tags=[
+                            TrendTag.intermediate(),
+                            TrendTag.arm_id(arm_id),
+                            TrendTag.beam(beam),
+                            TrendTag.instrument_processing_options(inst_opt_name),
+                            TrendTag.pac_fit_options(fit_opt_name),
+                            TrendTag.task_best_fit_parameters(),
+                        ],
+                        encoder=fits_array_encoder,
+                    )
+
+                    # Demodulation matrices
+                    data = rng.random(raw_demod_shape, dtype=np.float64)
+                    task.write(
+                        data=data,
+                        tags=[
+                            TrendTag.intermediate(),
+                            TrendTag.arm_id(arm_id),
+                            TrendTag.beam(beam),
+                            TrendTag.instrument_processing_options(inst_opt_name),
+                            TrendTag.pac_fit_options(fit_opt_name),
+                            TrendTag.task_best_fit_demodulation_matrices(),
+                        ],
+                        encoder=fits_array_encoder,
+                    )
+
+                    # Flux
+                    data = rng.random(raw_flux_shape, dtype=np.float64)
+                    task.write(
+                        data=data,
+                        tags=[
+                            TrendTag.intermediate(),
+                            TrendTag.arm_id(arm_id),
+                            TrendTag.beam(beam),
+                            TrendTag.instrument_processing_options(inst_opt_name),
+                            TrendTag.pac_fit_options(fit_opt_name),
+                            TrendTag.task_best_fit_flux(),
+                        ],
+                        encoder=fits_array_encoder,
+                    )
+
+                    # Residuals
+                    data = rng.random(raw_flux_shape, dtype=np.float64)
+                    task.write(
+                        data=data,
+                        tags=[
+                            TrendTag.intermediate(),
+                            TrendTag.arm_id(arm_id),
+                            TrendTag.beam(beam),
+                            TrendTag.instrument_processing_options(inst_opt_name),
+                            TrendTag.pac_fit_options(fit_opt_name),
+                            TrendTag.task_fit_residuals(),
+                        ],
+                        encoder=fits_array_encoder,
+                    )
+
+                    if arm_id != "CI":
+                        # Single-beam DMPD
+                        data = rng.random(raw_dmpd_shape, dtype=np.float64)
+                        task.write(
+                            data=data,
+                            tags=[
+                                TrendTag.intermediate(),
+                                TrendTag.arm_id(arm_id),
+                                TrendTag.beam(beam),
+                                TrendTag.instrument_processing_options(inst_opt_name),
+                                TrendTag.pac_fit_options(fit_opt_name),
+                                TrendTag.task_single_beam_demodulated_polcal_data(),
+                            ],
+                            encoder=fits_array_encoder,
+                        )
+
+
+@pytest.mark.parametrize(
+    "instrument, arm_id_list",
+    [
+        pytest.param(InstrumentName.visp, [1], id="visp"),
+        pytest.param(InstrumentName.cryonirsp, ["CI"], id="cryo-ci"),
+        pytest.param(InstrumentName.cryonirsp, ["SP"], id="cryo-sp"),
+        pytest.param(InstrumentName.dlnirsp, ["HBand"], id="dlnirsp"),
+        pytest.param(InstrumentName.visp, [1, 2, 3], id="visp-multi-arm"),
+        pytest.param(InstrumentName.cryonirsp, ["CI", "SP"], id="cryo-multi-arm"),
+        pytest.param(InstrumentName.dlnirsp, ["VIS", "JBand", "HBand"], id="dlnirsp-multi-arm"),
+    ],
+)
+def test_write_trend_task(
+    write_trend_task,
+    instrument,
+    arm_id_list,
+    num_polcal_ops,
+    num_cs_steps,
+    num_modstates,
+    fov_shape,
+    fit_options_list,
+    visp_instrument_options_list,
+    cryonirsp_instrument_options_list,
+    dlnirsp_instrument_options_list,
+    mocker,
+    fake_gql_client,
+):
+    """
+    Given: A `WriteTrendOutputs` task with some pre-output data on disk
+    When: Running the task
+    Then: The correct number of outputs is written, and the structure of each is correct
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
+    )
+
+    raw_cs_shape = (num_polcal_ops, num_modstates, num_cs_steps)
+    raw_fit_param_shape = (*fov_shape, num_polcal_ops, 15 + num_cs_steps, 3)
+    raw_demod_shape = (*fov_shape, 4, num_modstates)
+    raw_flux_shape = (*fov_shape, num_modstates, num_cs_steps)
+    raw_dmpd_shape = (num_polcal_ops, num_cs_steps, *fov_shape)
+
+    task = write_trend_task
+
+    write_pre_output_files_to_task(
+        task=task,
+        raw_calibration_sequence_data_shape=raw_cs_shape,
+        raw_fit_parameter_shape=raw_fit_param_shape,
+        raw_demod_shape=raw_demod_shape,
+        raw_flux_shape=raw_flux_shape,
+        raw_dmpd_shape=raw_dmpd_shape,
+    )
+
+    task()
+
+    match instrument:
+        case InstrumentName.visp:
+            inst_opt_list = visp_instrument_options_list
+        case InstrumentName.cryonirsp:
+            inst_opt_list = cryonirsp_instrument_options_list
+        case InstrumentName.dlnirsp:
+            inst_opt_list = dlnirsp_instrument_options_list
+        case _:
+            raise ValueError(
+                f"Do not recognize {instrument = }. Your fixtures/parametrization are jacked."
+            )
+
+    for arm_id in arm_id_list:
+        num_beams = num_beams_from_arm_id(arm_id)
+        expected_cs_shape = (num_polcal_ops, num_modstates, num_cs_steps)
+        expected_fit_param_shape = (num_beams, *fov_shape, num_polcal_ops, 15 + num_cs_steps, 3)
+        expected_modmat_shape = (num_beams, *fov_shape, num_modstates, 4)
+        expected_flux_shape = (num_beams, *fov_shape, num_modstates, num_cs_steps)
+        expected_dmpd_shape = (num_polcal_ops, num_cs_steps, *fov_shape)
+        expected_num_hdus = 9
+        expected_hdu_name_order = [
+            "Inventory",
+            "Calibration sequence",
+            "Fit info",
+            "Fit parameters",
+            "Modulation matrices",
+            "Best fit flux",
+            "Best fit residuals",
+            "Demodulated polcal data",
+        ]
+        expected_data_shape_list = [
+            # None for HDUs that have no data
+            None,
+            expected_cs_shape,
+            None,
+            expected_fit_param_shape,
+            expected_modmat_shape,
+            expected_flux_shape,
+            expected_flux_shape,
+            expected_dmpd_shape,
+        ]
+
+        if arm_id != "CI":
+            expected_num_hdus += 2
+            expected_hdu_name_order.insert(-1, "Beam 1 demodulated polcal data")
+            expected_hdu_name_order.insert(-1, "Beam 2 demodulated polcal data")
+            expected_data_shape_list += [expected_dmpd_shape, expected_dmpd_shape]
+
+        for fit_opt in fit_options_list:
+            for inst_opt in inst_opt_list:
+                tags = [
+                    TrendTag.output(),
+                    TrendTag.frame(),
+                    TrendTag.arm_id(arm_id),
+                    TrendTag.pac_fit_options(fit_opt["name"]),
+                    TrendTag.instrument_processing_options(inst_opt["name"]),
+                ]
+                assert task.count(tags) == 1
+                full_path = next(task.read(tags))
+                filename = full_path.name
+                assert filename.startswith("trend_polcal")
+                assert f"_{instrument}_" in filename
+                assert f"_{arm_id}_" in filename
+                assert f"_{fit_opt["name"]}-{inst_opt["name"]}_" in filename
+                assert filename.endswith(".fits")
+
+                with fits.open(full_path) as hdul:
+                    assert len(hdul) == expected_num_hdus
+                    assert isinstance(hdul[0], fits.PrimaryHDU)
+                    assert [hdu.name for hdu in hdul[1:]] == expected_hdu_name_order
+
+                    for hdu, expected_shape in zip(
+                        hdul[1:],
+                        expected_data_shape_list,
+                    ):
+                        if hdu.name == "Inventory":
+                            assert "FRAMEVOL" in hdu.header
+                        else:
+                            assert "FRAMEVOL" not in hdu.header
+                        if hdu.data is not None:
+                            assert isinstance(hdu, fits.CompImageHDU)
+                            assert hdu.data.dtype is np.dtype(np.float32)
+                            assert hdu.header["BITPIX"] == -32
+                            assert (
+                                hdu.data.shape == expected_shape
+                            ), f"Data shape {hdu.data.shape} does not match expected {expected_shape} for HDU {hdu.name}"
+                        else:
+                            assert isinstance(hdu, fits.ImageHDU)
+
+                    # TODO: Test validation

dkist_processing_trend/workflows/__init__.py

@@ -0,0 +1,3 @@
+"""Workflow package."""
+
+from dkist_processing_trend.config import dkist_processing_trend_configurations

dkist_processing_trend/workflows/visp.py

@@ -0,0 +1,58 @@
+"""Trend pipelines for ViSP."""
+
+from dkist_processing_common.tasks import Teardown
+from dkist_processing_common.tasks import TransferL0Data
+from dkist_processing_core import Workflow
+
+from dkist_processing_trend.models.instrument import VispArms
+from dkist_processing_trend.tasks import DarkCalibration
+from dkist_processing_trend.tasks import GainCalibration
+from dkist_processing_trend.tasks import ParseL0TrendInputData
+from dkist_processing_trend.tasks import PublishTrendCatalogMessages
+from dkist_processing_trend.tasks import RunPacFitter
+from dkist_processing_trend.tasks import TransferTrendData
+from dkist_processing_trend.tasks import WriteTrendOutputs
+from dkist_processing_trend.tasks.arm_task_factory import retreive_arm_task
+from dkist_processing_trend.tasks.visp import VispDemodulatedPolcalData
+from dkist_processing_trend.tasks.visp import VispGeometricCalibration
+from dkist_processing_trend.tasks.visp import VispPrepareFitData
+from dkist_processing_trend.tasks.visp import VispProcessDemodulationMatrices
+
+visp_polcal_pipeline = Workflow(
+    category="visp",
+    input_data="l0",
+    output_data="trend",
+    detail="polcal",
+    workflow_package=__package__,
+)
+visp_polcal_pipeline.add_node(task=TransferL0Data, upstreams=None)
+visp_polcal_pipeline.add_node(task=ParseL0TrendInputData, upstreams=TransferL0Data)
+
+output_upstreams = []
+for arm_id in VispArms:
+    dark_task = retreive_arm_task(DarkCalibration, arm_id=arm_id)
+    gain_task = retreive_arm_task(GainCalibration, arm_id=arm_id)
+    geo_task = retreive_arm_task(VispGeometricCalibration, arm_id=arm_id)
+    prep_data_task = retreive_arm_task(VispPrepareFitData, arm_id=arm_id)
+    run_fit_task = retreive_arm_task(RunPacFitter, arm_id=arm_id)
+    process_demod_task = retreive_arm_task(VispProcessDemodulationMatrices, arm_id=arm_id)
+    dmpd_task = retreive_arm_task(VispDemodulatedPolcalData, arm_id=arm_id)
+
+    # Generate demod flow
+    visp_polcal_pipeline.add_node(task=dark_task, upstreams=ParseL0TrendInputData)
+    visp_polcal_pipeline.add_node(task=gain_task, upstreams=dark_task)
+    visp_polcal_pipeline.add_node(task=prep_data_task, upstreams=gain_task)
+    visp_polcal_pipeline.add_node(task=run_fit_task, upstreams=prep_data_task)
+
+    # DMPD flow
+    visp_polcal_pipeline.add_node(task=geo_task, upstreams=ParseL0TrendInputData)
+    visp_polcal_pipeline.add_node(task=process_demod_task, upstreams=run_fit_task)
+    visp_polcal_pipeline.add_node(task=dmpd_task, upstreams=[geo_task, process_demod_task])
+
+    output_upstreams += [dmpd_task]
+
+# Output
+visp_polcal_pipeline.add_node(task=WriteTrendOutputs, upstreams=output_upstreams)
+visp_polcal_pipeline.add_node(task=TransferTrendData, upstreams=WriteTrendOutputs)
+visp_polcal_pipeline.add_node(task=PublishTrendCatalogMessages, upstreams=TransferTrendData)
+visp_polcal_pipeline.add_node(task=Teardown, upstreams=PublishTrendCatalogMessages)