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

dkist_processing_trend/tests/test_constants.py

@@ -0,0 +1,120 @@
+import json
+from dataclasses import asdict
+from dataclasses import dataclass
+from dataclasses import field
+
+import pytest
+from dkist_processing_common.tasks.base import WorkflowTaskBase
+
+from dkist_processing_trend.models.constants import TrendBudName
+from dkist_processing_trend.models.constants import TrendConstants
+
+
+@dataclass
+class TrendOnlyConstantDB:
+    earliest_ip_start_time: str = "1900-01-01"
+    arm_id_list: tuple[str | int] = (1, 2, 3)
+    wavelength_for_arm: dict[str, float] = field(default_factory=lambda: {"1": 123.4, '"CI"': 567})
+    observing_program_id_list: tuple[str] = ("TOP1", "TOP2")
+    experiment_id_list: tuple[str] = "EX1"
+    proposal_id_list: tuple[str] = "PROP5"
+    hls_version_list: tuple[str] = ("Tau 6.2.8", "Pi 3.14")
+    visp_grating_angle_for_arm: dict[str, float] = field(
+        default_factory=lambda: {"1": 9.87, "2": 9.872}
+    )
+    visp_arm_position_for_arm: dict[str, float] = field(
+        default_factory=lambda: {"1": 8.76, "2": 8.762}
+    )
+    visp_grating_constant_for_arm: dict[str, float] = field(
+        default_factory=lambda: {"1": 9870, "2": 9872}
+    )
+    cryonirsp_od_filter_name_for_arm: dict[str, str] = field(
+        default_factory=lambda: {'"CI"': "CRYOFILT"}
+    )
+    cryonirsp_grating_angle_for_arm: dict[str, float] = field(
+        default_factory=lambda: {'"CI"': 123.4, '"SP"': 123.5}
+    )
+    cryonirsp_grating_littrow_angle_for_arm: dict[str, float] = field(
+        default_factory=lambda: {'"CI"': 12.34, '"SP"': 12.35}
+    )
+    cryonirsp_grating_constant_for_arm: dict[str, float] = field(
+        default_factory=lambda: {'"SP"': 1234}
+    )
+    cryonirsp_slit_width_for_arm: dict[str, float] = field(default_factory=lambda: {'"CI"': 0.1234})
+    dlnirsp_grating_angle_for_arm: dict[str, float] = field(
+        default_factory=lambda: {'"VIS"': 2.22, '"JBand"': 2.222, '"HBand"': 2.2222}
+    )
+    dlnirsp_arm_position_for_arm: dict[str, float] = field(
+        default_factory=lambda: {'"JBand"': 3.33}
+    )
+    dlnirsp_grating_constant_for_arm: dict[str, float] = field(
+        default_factory=lambda: {'"HBand"': 4444}
+    )
+    # And one "common" constant just to confirm subclassing
+    instrument: str = "an instrument"
+
+
+@pytest.fixture(scope="session")
+def arm_id_lookup_constant_names():
+    return {
+        "wavelength_for_arm",
+        "visp_grating_angle_for_arm",
+        "visp_arm_position_for_arm",
+        "visp_grating_constant_for_arm",
+        "cryonirsp_od_filter_name_for_arm",
+        "cryonirsp_grating_angle_for_arm",
+        "cryonirsp_grating_littrow_angle_for_arm",
+        "cryonirsp_grating_constant_for_arm",
+        "cryonirsp_slit_width_for_arm",
+        "dlnirsp_grating_angle_for_arm",
+        "dlnirsp_arm_position_for_arm",
+        "dlnirsp_grating_constant_for_arm",
+    }
+
+
+@pytest.fixture
+def trend_constants_db():
+    values = TrendOnlyConstantDB()
+    return {k.upper(): v for k, v in asdict(values).items()}
+
+
+@pytest.fixture
+def task_with_trend_constants(recipe_run_id, link_constants_db, trend_constants_db):
+    class Task(WorkflowTaskBase):
+        @property
+        def constants_model_class(self):
+            return TrendConstants
+
+        def run(self): ...
+
+    link_constants_db(recipe_run_id, trend_constants_db)
+    with Task(
+        recipe_run_id=recipe_run_id,
+        workflow_name="workflow_name",
+        workflow_version="workflow_version",
+    ) as task:
+        yield task
+        task._purge()
+
+
+def test_trend_constants(
+    task_with_trend_constants, trend_constants_db, arm_id_lookup_constant_names
+):
+    """
+    Given: A task whose `constants_model_class` property is `TrendConstants`
+    When: Accessing the constants implied by `TrendBudName`
+    Then: The property on `.constants` matches the correct value in the constants db
+    """
+    task = task_with_trend_constants
+    task_constants = task.constants
+    assert task_constants.instrument == trend_constants_db["INSTRUMENT"]
+    for constant in TrendBudName:
+        expected_value = trend_constants_db[constant.value]
+        if isinstance(expected_value, tuple):
+            # B/c JSON serializes/deserializes lists as tuples
+            expected_value = list(expected_value)
+        if constant.name in arm_id_lookup_constant_names:
+            expected_value = {json.loads(k): v for k, v in expected_value.items()}
+        assert (
+            getattr(task_constants, constant.name) == expected_value
+        ), f"{constant.name} does not match"

dkist_processing_trend/tests/test_dark.py

@@ -0,0 +1,97 @@
+import numpy as np
+import pytest
+from dkist_processing_common._util.scratch import WorkflowFileSystem
+from dkist_processing_common.codecs.fits import fits_array_decoder
+
+from dkist_processing_trend.models.instrument import InstrumentName
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.tasks.dark import DarkCalibration
+from dkist_processing_trend.tests.conftest import ModulatedCSStepHeaders
+from dkist_processing_trend.tests.conftest import TrendTestingConstantsDatabase
+from dkist_processing_trend.tests.conftest import write_calibration_sequence_frames_to_task
+
+
+@pytest.fixture
+def arm_id() -> str:
+    return "arm2"
+
+
+@pytest.fixture
+def dark_task(tmp_path, recipe_run_id, link_constants_db, arm_id):
+
+    link_constants_db(recipe_run_id, TrendTestingConstantsDatabase(ARM_ID_LIST=(1, "arm3", "arm2")))
+    with DarkCalibration(
+        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)
+        yield task
+        task._purge
+
+
+def correct_arm_dark_signal_from_modstate(frame: ModulatedCSStepHeaders) -> np.ndarray:
+    shape = frame.array_shape
+    signal = frame.current_modstate * 100
+    return np.full(shape, signal)
+
+
+def wrong_arm_dark_signal_from_modstate(frame: ModulatedCSStepHeaders) -> np.ndarray:
+    return correct_arm_dark_signal_from_modstate(frame) * 1000
+
+
+def test_dark_task(dark_task, arm_id, mocker, fake_gql_client):
+    """
+    Given: A `DarkCalibration` task and some frames tagged with POLCAL_DARK in *two different arms*
+    When: Running the task instantiated with a single arm
+    Then: The correct dark frame is made and did not include data from arms other than the task was instantiated with
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
+    )
+    task = dark_task
+    shape = (3, 3)
+
+    # Write POLCAL_DARK data for the arm we want
+    write_calibration_sequence_frames_to_task(
+        task=task,
+        instrument=InstrumentName.visp,
+        arm_id=arm_id,
+        wavelength=636.0,
+        pol_status=["clear"],
+        pol_theta=[99.0],
+        ret_status=["clear"],
+        ret_theta=[99.0],
+        dark_status=["DarkShutter"],
+        num_modstates=4,
+        array_shape=shape,
+        data_func=correct_arm_dark_signal_from_modstate,
+        tags=[TrendTag.task_polcal_dark(), TrendTag.arm_id(arm_id)],
+    )
+
+    # Now the same, but for a different arm with totally different signal
+    write_calibration_sequence_frames_to_task(
+        task=task,
+        instrument=InstrumentName.visp,
+        arm_id="wrong_arm",
+        wavelength=636.0,
+        pol_status=["clear"],
+        pol_theta=[99.0],
+        ret_status=["clear"],
+        ret_theta=[99.0],
+        dark_status=["DarkShutter"],
+        num_modstates=4,
+        array_shape=(2, 2),  # Different!
+        data_func=wrong_arm_dark_signal_from_modstate,
+        tags=[TrendTag.task_polcal_dark(), TrendTag.arm_id("wrong_arm")],
+    )
+
+    task()
+
+    tags = [
+        TrendTag.arm_id(arm_id),
+        TrendTag.intermediate(),
+        TrendTag.frame(),
+        TrendTag.task_dark(),
+    ]
+    assert task.count(tags=tags) == 1
+    dark_array = next(task.read(tags=tags, decoder=fits_array_decoder))
+    np.testing.assert_array_equal(dark_array, 250.0)  # avg of [1, 2, 3, 4] * 100

dkist_processing_trend/tests/test_gain.py

@@ -0,0 +1,135 @@
+from functools import partial
+
+import numpy as np
+import pytest
+from dkist_processing_common._util.scratch import WorkflowFileSystem
+from dkist_processing_common.codecs.fits import fits_array_decoder
+from dkist_processing_common.codecs.fits import fits_array_encoder
+
+from dkist_processing_trend.models.instrument import InstrumentName
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.tasks.gain import GainCalibration
+from dkist_processing_trend.tests.conftest import ModulatedCSStepHeaders
+from dkist_processing_trend.tests.conftest import TrendTestingConstantsDatabase
+from dkist_processing_trend.tests.conftest import write_calibration_sequence_frames_to_task
+
+
+@pytest.fixture
+def arm_id() -> str:
+    return "arm2"
+
+
+@pytest.fixture
+def gain_task(tmp_path, recipe_run_id, link_constants_db, arm_id):
+
+    link_constants_db(recipe_run_id, TrendTestingConstantsDatabase(ARM_ID_LIST=(1, "arm3", "arm2")))
+    with GainCalibration(
+        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)
+        yield task
+        task._purge
+
+
+@pytest.fixture
+def dark_signal() -> float:
+    return 314.0
+
+
+def correct_arm_gain_signal_from_modstate(
+    frame: ModulatedCSStepHeaders, dark_signal: float
+) -> np.ndarray:
+    shape = frame.array_shape
+    signal = frame.current_modstate * 1000
+    return np.full(shape, signal) + dark_signal
+
+
+def wrong_arm_gain_signal_from_modstate(
+    frame: ModulatedCSStepHeaders, dark_signal: float
+) -> np.ndarray:
+    return correct_arm_gain_signal_from_modstate(frame, dark_signal=0) * 1000 + dark_signal
+
+
+def test_gain_task(gain_task, arm_id, dark_signal, mocker, fake_gql_client):
+    """
+    Given: A `GainCalibration` task and some frames tagged with POLCAL_GAIN in *two different arms*
+    When: Running the task instantiated with a single arm
+    Then: The correct gain frame is made and did not include data from arms other than the task was instantiated with
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
+    )
+    task = gain_task
+    shape = (3, 3)
+
+    # Write POLCAL_GAIN data for the arm we want
+    write_calibration_sequence_frames_to_task(
+        task=task,
+        instrument=InstrumentName.visp,
+        arm_id=arm_id,
+        wavelength=636.0,
+        pol_status=["clear"],
+        pol_theta=[99.0],
+        ret_status=["clear"],
+        ret_theta=[99.0],
+        dark_status=["clear"],
+        num_modstates=4,
+        array_shape=shape,
+        data_func=partial(correct_arm_gain_signal_from_modstate, dark_signal=dark_signal),
+        tags=[TrendTag.task_polcal_gain(), TrendTag.arm_id(arm_id)],
+    )
+
+    # Now the same, but for a different arm with totally different signal
+    write_calibration_sequence_frames_to_task(
+        task=task,
+        instrument=InstrumentName.visp,
+        arm_id="wrong_arm",
+        wavelength=636.0,
+        pol_status=["clear"],
+        pol_theta=[99.0],
+        ret_status=["clear"],
+        ret_theta=[99.0],
+        dark_status=["DarkShutter"],
+        num_modstates=4,
+        array_shape=(2, 2),  # Different!
+        data_func=partial(wrong_arm_gain_signal_from_modstate, dark_signal=dark_signal),
+        tags=[TrendTag.task_polcal_gain(), TrendTag.arm_id("wrong_arm")],
+    )
+
+    # Write two INTERMEDIATE DARK frames: one for the correct arm and one for a different arm
+    dark_array = np.full(shape, dark_signal)
+    wrong_arm_dark_array = dark_array * 100000
+    task.write(
+        data=dark_array,
+        tags=[
+            TrendTag.intermediate(),
+            TrendTag.frame(),
+            TrendTag.task_dark(),
+            TrendTag.arm_id(arm_id),
+        ],
+        encoder=fits_array_encoder,
+    )
+    task.write(
+        data=wrong_arm_dark_array,
+        tags=[
+            TrendTag.intermediate(),
+            TrendTag.frame(),
+            TrendTag.task_dark(),
+            TrendTag.arm_id("wrong_arm"),
+        ],
+        encoder=fits_array_encoder,
+    )
+
+    task()
+
+    tags = [
+        TrendTag.arm_id(arm_id),
+        TrendTag.intermediate(),
+        TrendTag.frame(),
+        TrendTag.task_gain(),
+    ]
+
+    expected_gain_value = 2500.0  # Mean of [1,2,3,4] * 1000 (dark signal is removed)
+    assert task.count(tags=tags) == 1
+    dark_array = next(task.read(tags=tags, decoder=fits_array_decoder))
+    np.testing.assert_array_equal(dark_array, expected_gain_value)

dkist_processing_trend/tests/test_parameters.py

@@ -0,0 +1,149 @@
+from dataclasses import asdict
+
+import pytest
+from dkist_processing_common._util.scratch import WorkflowFileSystem
+from dkist_processing_common.tasks.base import WorkflowTaskBase
+
+from dkist_processing_trend.models.fit_options import FitOptions
+from dkist_processing_trend.models.instrument import InstrumentName
+from dkist_processing_trend.models.instrument_options import CryonirspInstrumentOptions
+from dkist_processing_trend.models.instrument_options import DlnirspInstrumentOptions
+from dkist_processing_trend.models.instrument_options import VispInstrumentOptions
+from dkist_processing_trend.models.parameters import TrendParameters
+from dkist_processing_trend.models.parameters import TrendParsingParameters
+from dkist_processing_trend.tasks.trend_base import TrendTaskBase
+from dkist_processing_trend.tests.conftest import TrendInputDatasetParameterValues
+from dkist_processing_trend.tests.conftest import TrendTestingConstantsDatabase
+
+
+@pytest.fixture(scope="session")
+def parse_parameter_names() -> set[str]:
+    # The property names of all parameter properties on `TrendParsingParameters`
+    return set(k for k, v in vars(TrendParsingParameters).items() if isinstance(v, property))
+
+
+@pytest.fixture(scope="session")
+def parameter_names() -> set[str]:
+    return set(k for k, v in vars(TrendParameters).items() if isinstance(v, property))
+
+
+@pytest.fixture(scope="session")
+def instrument_parameter_raw_names() -> set[str]:
+    return {"trend_instrument_processing_options_visp"}
+
+
+@pytest.fixture(scope="session")
+def instrument_parameter_property_names(instrument_parameter_raw_names) -> set[str]:
+    return {
+        p.removeprefix("trend_")
+        .replace("_visp", "")
+        .replace("_cryonirsp", "")
+        .replace("_dlnirsp", "")
+        for p in instrument_parameter_raw_names
+    }
+
+
+class Task(WorkflowTaskBase):
+    def run(self):
+        pass
+
+
+class ScienceTask(TrendTaskBase):
+    def run(self):
+        pass
+
+
+@pytest.fixture
+def task_with_parsing_parameters(tmp_path, recipe_run_id, assign_input_dataset_doc_to_task):
+    with Task(
+        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)
+        param_values = TrendInputDatasetParameterValues()
+        assign_input_dataset_doc_to_task(
+            task=task,
+            parameter_values=param_values,
+        )
+        task.parameters = TrendParsingParameters(scratch=task.scratch)
+
+        yield task, param_values
+        task._purge()
+
+
+@pytest.fixture
+def task_with_parameters(
+    tmp_path, recipe_run_id, assign_input_dataset_doc_to_task, link_constants_db, instrument
+):
+    link_constants_db(recipe_run_id, TrendTestingConstantsDatabase(INSTRUMENT=instrument))
+    with ScienceTask(
+        recipe_run_id=recipe_run_id, workflow_name="name", workflow_version="version"
+    ) as task:
+        task.scratch = WorkflowFileSystem(scratch_base_path=tmp_path, recipe_run_id=recipe_run_id)
+        param_values = TrendInputDatasetParameterValues()
+        assign_input_dataset_doc_to_task(task, parameter_values=param_values)
+        task.parameters = TrendParameters(
+            scratch=task.scratch,
+            obs_ip_start_time=task.constants.earliest_ip_start_time,
+            instrument=task.constants.instrument,
+        )
+
+        yield task, param_values
+        task._purge()
+
+
+def test_parse_parameters(task_with_parsing_parameters, parse_parameter_names):
+    """
+    Given: A Task whose `parameter` property is an instance of `TrendParsingParameters`
+    When: Accessing properties for Parse parameters
+    Then: The correct value is returned
+    """
+    task, expected_param_values = task_with_parsing_parameters
+    tested_parse_parameters = set()
+    task_param_attr = task.parameters
+    for pn, pv in asdict(expected_param_values).items():
+        property_name = pn.removeprefix("trend_")
+        if property_name in parse_parameter_names and type(pv) is not dict:
+            assert getattr(task_param_attr, property_name) == pv
+            tested_parse_parameters.add(property_name)
+
+    assert (
+        tested_parse_parameters == parse_parameter_names
+    ), f"Did not find test values for parameters {parse_parameter_names - tested_parse_parameters}"
+
+
+@pytest.mark.parametrize(
+    "instrument", [InstrumentName.visp, InstrumentName.cryonirsp, InstrumentName.dlnirsp]
+)
+def test_parameters(
+    task_with_parameters, instrument, parameter_names, instrument_parameter_property_names
+):
+    """
+    Given: A task with `TrendParameters` as its parameters object
+    When: Using the `.parameters` property to access parameters
+    Then: The correct values are returned
+    """
+    task, expected_param_values = task_with_parameters
+    task_param_attr = task.parameters
+    expected_dict = asdict(expected_param_values)
+
+    for parameter_name in parameter_names:
+        db_name = f"trend_{parameter_name}"
+        if parameter_name in instrument_parameter_property_names:
+            db_name = f"{db_name}_{task_param_attr._instrument}"
+        expected = expected_dict[db_name]
+        if parameter_name == "fit_options_list":
+            expected = [FitOptions.model_validate(i) for i in expected]
+        if parameter_name == "instrument_processing_options":
+            match instrument:
+                case InstrumentName.visp:
+                    expected = [VispInstrumentOptions.model_validate(i) for i in expected]
+                case InstrumentName.cryonirsp:
+                    expected = [CryonirspInstrumentOptions.model_validate(i) for i in expected]
+                case InstrumentName.dlnirsp:
+                    expected = [DlnirspInstrumentOptions.model_validate(i) for i in expected]
+
+        if isinstance(expected, tuple):
+            expected = list(expected)
+        assert getattr(task_param_attr, parameter_name) == expected