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

dkist_processing_trend/tests/local_trial_workflows/local_trial_helpers.py

@@ -0,0 +1,488 @@
+import json
+import os
+import shutil
+from dataclasses import asdict
+from pathlib import Path
+from random import randint
+
+import asdf
+from astropy.io import fits
+from dkist_header_validator import spec122_validator
+from dkist_processing_common.codecs.basemodel import basemodel_encoder
+from dkist_processing_common.models.input_dataset import InputDatasetPartDocumentList
+from dkist_processing_common.models.task_name import TaskName
+from dkist_processing_common.tasks.base import WorkflowTaskBase
+from dkist_service_configuration.logging import logger
+
+from dkist_processing_trend.models.constants import TrendBudName
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.models.task_name import TrendTaskName
+from dkist_processing_trend.tests.conftest import TrendInputDatasetParameterValues
+
+
+def translate_task(instrument: str, suffix: str):
+    class Translate122To214L0(WorkflowTaskBase):
+        def run(self) -> None:
+            raw_dir = (
+                Path(self.scratch.scratch_base_path)
+                / f"{instrument.upper()}{self.recipe_run_id:03n}"
+            )
+            if not self.scratch.workflow_base_path.exists():
+                self.scratch.workflow_base_path.mkdir(parents=True, exist_ok=False)
+
+            if not raw_dir.exists():
+                raise FileNotFoundError(
+                    f"Expected to find a raw {instrument.upper()}{self.recipe_run_id:03n} folder in {self.scratch.scratch_base_path}"
+                )
+
+            for file in raw_dir.glob(f"*.{suffix}"):
+                translated_file_name = self.scratch.workflow_base_path / os.path.basename(file)
+                logger.info(f"Translating {file} -> {translated_file_name}")
+                with fits.open(file, memmap=False) as hdl:
+                    # Handle both compressed and uncompressed files...
+                    hdu = hdl[1] if len(hdl) > 1 else hdl[0]
+                    header = spec122_validator.validate_and_translate_to_214_l0(
+                        hdu.header, return_type=fits.HDUList
+                    )[0].header
+
+                    comp_hdu = fits.CompImageHDU(header=header, data=hdu.data)
+                    comp_hdl = fits.HDUList([fits.PrimaryHDU(), comp_hdu])
+                    comp_hdl.writeto(translated_file_name, overwrite=True)
+
+                comp_hdl.close()
+                del comp_hdl
+
+    return Translate122To214L0
+
+
+def tag_inputs_task(suffix: str):
+    class TagInputs(WorkflowTaskBase):
+        def run(self) -> None:
+            logger.info(f"Looking in {self.scratch.workflow_base_path.absolute()}")
+            input_file_list = list(self.scratch.workflow_base_path.glob(f"*.{suffix}"))
+            if len(input_file_list) == 0:
+                raise FileNotFoundError(
+                    f"Did not find any files matching '*.{suffix}' in {self.scratch.workflow_base_path}"
+                )
+            for file in input_file_list:
+                logger.info(f"Found {file}")
+                self.tag(path=file, tags=[TrendTag.input(), TrendTag.frame()])
+
+    return TagInputs
+
+
+def input_dataset_parameter_task(parameter_data: TrendInputDatasetParameterValues):
+
+    class CreateInputDatasetParameterDocument(WorkflowTaskBase):
+        def run(self) -> None:
+            relative_path = "input_dataset_parameters.json"
+            self.write(
+                data=InputDatasetPartDocumentList(
+                    doc_list=self.input_dataset_document_simple_parameters_part
+                ),
+                relative_path=relative_path,
+                tags=TrendTag.input_dataset_parameters(),
+                encoder=basemodel_encoder,
+                overwrite=True,
+            )
+            logger.info(f"Wrote input dataset parameter doc to {relative_path}")
+
+        @property
+        def input_dataset_document_simple_parameters_part(self):
+            parameters_list = []
+            value_id = randint(1000, 2000)
+            for pn, pv in asdict(parameter_data).items():
+                values = [
+                    {
+                        "parameterValueId": value_id,
+                        "parameterValue": json.dumps(pv),
+                        "parameterValueStartDate": "1946-11-20",
+                    }
+                ]
+                parameter = {"parameterName": pn, "parameterValues": values}
+                parameters_list.append(parameter)
+
+            return parameters_list
+
+    return CreateInputDatasetParameterDocument
+
+
+def save_parsing_task(
+    tag_list: list[str], save_file: str, save_file_tags: bool = True, save_constants: bool = True
+):
+    class SaveParsing(WorkflowTaskBase):
+        """Save the result of parsing (constants and tags) to an asdf file."""
+
+        @property
+        def relative_save_file(self) -> str:
+            return save_file
+
+        def run(self):
+            if save_file_tags:
+                file_tag_dict = self.get_input_tags()
+            else:
+                logger.info("Skipping saving of file tags")
+                file_tag_dict = dict()
+            if save_constants:
+                constant_dict = self.get_constants()
+            else:
+                logger.info("Skipping saving of constants")
+                constant_dict = dict()
+
+            full_save_file = self.scratch.workflow_base_path / self.relative_save_file
+            tree = {"file_tag_dict": file_tag_dict, "constants_dict": constant_dict}
+            af = asdf.AsdfFile(tree)
+            af.write_to(full_save_file)
+            logger.info(f"Saved input tags to {full_save_file}")
+
+        def get_input_tags(self) -> dict[str, list[str]]:
+            file_tag_dict = dict()
+            path_list = self.read(tags=tag_list)
+            for p in path_list:
+                tags = self.tags(p)
+                file_tag_dict[str(p)] = tags
+
+            return file_tag_dict
+
+        def get_constants(self) -> dict[str, str | float | list]:
+            constants_dict = dict()
+            for c in self.constants._db_dict.keys():
+                constants_dict[c] = self.constants._db_dict[c]
+
+            return constants_dict
+
+    return SaveParsing
+
+
+def load_parsing_task(save_file: str):
+    class LoadParsing(WorkflowTaskBase):
+        """Load tags and constants into the database."""
+
+        @property
+        def relative_save_file(self) -> str:
+            return save_file
+
+        def run(self):
+            full_save_file = self.scratch.workflow_base_path / self.relative_save_file
+            with asdf.open(full_save_file) as af:
+                file_tag_dict = af.tree["file_tag_dict"]
+                self.tag_input_files(file_tag_dict)
+
+                constants_dict = af.tree["constants_dict"]
+                self.populate_constants(constants_dict)
+
+            logger.info(f"Loaded tags and constants from {full_save_file}")
+
+        def tag_input_files(self, file_tag_dict: dict[str, list[str]]):
+            for f, t in file_tag_dict.items():
+                if not os.path.exists(f):
+                    pass
+                    # raise FileNotFoundError(f"Expected to find {f}, but it doesn't exist.")
+                else:
+                    self.tag(path=f, tags=t)
+
+        def populate_constants(self, constants_dict: dict[str, str | int | float]) -> None:
+            # First we purge all constants because a previous load might have polluted the DB
+            self.constants._purge()
+            for c, v in constants_dict.items():
+                logger.info(f"Setting value of {c} to {v}")
+                self.constants._update({c: v})
+
+    return LoadParsing
+
+
+def get_arm_id_list(workflow_base_path: Path, parsing_file_name: str | Path) -> list[str | int]:
+    with asdf.open(workflow_base_path / parsing_file_name) as af:
+        constants = af.tree["constants_dict"]
+        arm_id_list = constants[TrendBudName.arm_id_list]
+
+    return arm_id_list
+
+
+class SaveTaskTags(WorkflowTaskBase):
+    """Base task for saving all INTERMEDIATE files with the same 'TASK'"""
+
+    @property
+    def task_str(self) -> str:
+        return "TASK"
+
+    @property
+    def arm_id(self) -> str | int | None:
+        return None
+
+    @property
+    def relative_save_file(self) -> str:
+        return "default_sav.asdf"
+
+    @property
+    def tag_lists_to_save(self) -> list[list[str]] | list[str]:
+        intermediate_tags = [TrendTag.task(self.task_str), TrendTag.intermediate()]
+        if self.arm_id is not None:
+            intermediate_tags.append(TrendTag.arm_id(self.arm_id))
+        return [intermediate_tags]
+
+    def run(self):
+        file_tag_dict = dict()
+        tag_list_list = self.tag_lists_to_save
+        if isinstance(tag_list_list[0], str):
+            tag_list_list = [tag_list_list]
+
+        for tags_to_save in tag_list_list:
+            path_list = self.read(tags=tags_to_save)
+            save_dir = self.scratch.workflow_base_path / Path(self.relative_save_file).stem
+            save_dir.mkdir(exist_ok=True)
+            for p in path_list:
+                copied_path = shutil.copy(str(p), save_dir)
+                tags = self.tags(p)
+                file_tag_dict[copied_path] = tags
+
+        full_save_file = self.scratch.workflow_base_path / self.relative_save_file
+        tree = {"file_tag_dict": file_tag_dict}
+        af = asdf.AsdfFile(tree)
+        af.write_to(full_save_file)
+        logger.info(f"Saved {self.task_str} to {full_save_file}")
+
+
+class LoadTaskTags(WorkflowTaskBase):
+    """Base task for loading file/tag associations from a previously saved set."""
+
+    @property
+    def relative_save_file(self) -> str:
+        return "default_sav.asdf"
+
+    def run(self):
+        full_save_file = self.scratch.workflow_base_path / self.relative_save_file
+        with asdf.open(full_save_file) as af:
+            for f, t in af.tree["file_tag_dict"].items():
+                self.tag(path=f, tags=t)
+        logger.info(f"Loaded database entries from {full_save_file}")
+
+
+def save_dark_cal_task(arm_id: str | int):
+    class SaveDarkCal(SaveTaskTags):
+        @property
+        def task_str(self) -> str:
+            return TaskName.dark.value
+
+        @property
+        def arm_id(self):
+            return arm_id
+
+        @property
+        def relative_save_file(self) -> str:
+            return f"dark_cal_{arm_id}.asdf"
+
+    return SaveDarkCal
+
+
+def load_dark_cal_task(arm_id: str | int):
+    class LoadDarkCal(LoadTaskTags):
+        @property
+        def relative_save_file(self) -> str:
+            return f"dark_cal_{arm_id}.asdf"
+
+    return LoadDarkCal
+
+
+def save_gain_cal_task(arm_id: str | int):
+    class SaveGainCal(SaveTaskTags):
+        @property
+        def task_str(self) -> str:
+            return TaskName.gain.value
+
+        @property
+        def arm_id(self):
+            return arm_id
+
+        @property
+        def relative_save_file(self) -> str:
+            return f"gain_cal_{arm_id}.asdf"
+
+    return SaveGainCal
+
+
+def load_gain_cal_task(arm_id: str | int):
+    class LoadGainCal(LoadTaskTags):
+        @property
+        def relative_save_file(self) -> str:
+            return f"gain_cal_{arm_id}.asdf"
+
+    return LoadGainCal
+
+
+def save_prepped_cal_task(arm_id: str | int):
+    class SavePreppedData(SaveTaskTags):
+        @property
+        def tag_lists_to_save(self):
+            return [
+                [
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_local_pac_input(),
+                ],
+                [
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_global_pac_input(),
+                ],
+            ]
+
+        @property
+        def arm_id(self):
+            return arm_id
+
+        @property
+        def relative_save_file(self) -> str:
+            return f"prepped_cal_{arm_id}.asdf"
+
+    return SavePreppedData
+
+
+def load_prepped_cal_task(arm_id: str | int):
+    class LoadPreppedData(LoadTaskTags):
+        @property
+        def relative_save_file(self) -> str:
+            return f"prepped_cal_{arm_id}.asdf"
+
+    return LoadPreppedData
+
+
+def save_pac_fit_task(arm_id: str | int):
+    class SavePacFit(SaveTaskTags):
+        @property
+        def tag_lists_to_save(self):
+            return [
+                [
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_calibration_sequence(),
+                ],
+                [
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_best_fit_parameters(),
+                ],
+                [
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_best_fit_demodulation_matrices(),
+                ],
+                [
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_best_fit_flux(),
+                ],
+                [
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_fit_residuals(),
+                ],
+            ]
+
+        @property
+        def arm_id(self):
+            return arm_id
+
+        @property
+        def relative_save_file(self) -> str:
+            return f"pac_fit_{arm_id}.asdf"
+
+    return SavePacFit
+
+
+def load_pac_fit_task(arm_id: str | int):
+    class LoadPacFit(LoadTaskTags):
+        @property
+        def relative_save_file(self) -> str:
+            return f"pac_fit_{arm_id}.asdf"
+
+    return LoadPacFit
+
+
+def save_processed_demod_cal_task(arm_id: str | int):
+    class SaveProcessedDemodCal(SaveTaskTags):
+        @property
+        def task_str(self) -> str:
+            return TrendTaskName.processed_demodulation_matrices.value
+
+        @property
+        def arm_id(self):
+            return arm_id
+
+        @property
+        def relative_save_file(self) -> str:
+            return f"processed_demod_cal_{arm_id}.asdf"
+
+    return SaveProcessedDemodCal
+
+
+def load_processed_demod_cal_task(arm_id: str | int):
+    class LoadGainCal(LoadTaskTags):
+        @property
+        def relative_save_file(self) -> str:
+            return f"processed_demod_cal_{arm_id}.asdf"
+
+    return LoadGainCal
+
+
+def save_dmpd_task(arm_id: str | int):
+    class SaveDmpd(SaveTaskTags):
+        @property
+        def tag_lists_to_save(self):
+            return [
+                [
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_single_beam_demodulated_polcal_data(),
+                ],
+                [
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_demodulated_polcal_data(),
+                ],
+            ]
+
+        @property
+        def arm_id(self):
+            return arm_id
+
+        @property
+        def relative_save_file(self) -> str:
+            return f"dmpd_{arm_id}.asdf"
+
+    return SaveDmpd
+
+
+def load_dmpd_task(arm_id: str | int):
+    class LoadDmpd(LoadTaskTags):
+        @property
+        def relative_save_file(self) -> str:
+            return f"dmpd_{arm_id}.asdf"
+
+    return LoadDmpd
+
+
+def save_visp_geo_cal_task(arm_id: str | int):
+    class SaveVispGeoCal(SaveTaskTags):
+        @property
+        def task_str(self) -> str:
+            return TrendTaskName.visp_geometric_calibration.value
+
+        @property
+        def arm_id(self):
+            return arm_id
+
+        @property
+        def relative_save_file(self) -> str:
+            return f"visp_geo_cal_{arm_id}.asdf"
+
+    return SaveVispGeoCal
+
+
+def load_visp_geo_cal_task(arm_id: str | int):
+    class LoadVispGeoCal(LoadTaskTags):
+        @property
+        def relative_save_file(self) -> str:
+            return f"visp_geo_cal_{arm_id}.asdf"
+
+    return LoadVispGeoCal

dkist_processing_trend/tests/test_arm_class_factory.py

@@ -0,0 +1,39 @@
+import pytest
+from dkist_processing_core import TaskBase
+
+from dkist_processing_trend.tasks.trend_base import TrendArmTaskBase
+from dkist_processing_trend.tests.conftest import TrendTestingConstantsDatabase
+from dkist_processing_trend.workflows.arm_class_factory import arm_class_factory
+
+
+class DummyArmClass(TrendArmTaskBase):
+    """A real cool class"""
+
+    def run(self):
+        return 5
+
+
+@pytest.mark.parametrize("arm_id", [pytest.param("foo", id="str"), pytest.param(3, id="int")])
+def test_arm_class_factory(arm_id, recipe_run_id, link_constants_db):
+    """
+    Given: A Task class that subclasses `TrendArmTaskBase`
+    When: Creating a new version of that task with a pre-defined arm_id
+    Then: The new class is a `TaskBase`, has the correct arm_id, and retains the docstring of its parent
+    """
+    link_constants_db(
+        recipe_run_id,
+        TrendTestingConstantsDatabase(),
+    )
+    arm_class = arm_class_factory(DummyArmClass, arm_id=arm_id)
+    assert issubclass(arm_class, DummyArmClass)
+    assert issubclass(arm_class, TaskBase)
+    assert arm_class.__doc__ == "A real cool class"
+    assert arm_class.__name__ == f"DummyArmClass{str(arm_id).capitalize()}"
+    assert arm_class.__qualname__ == f"DummyArmClass{str(arm_id).capitalize()}"
+    assert arm_class.__module__ == DummyArmClass.__module__
+
+    with arm_class(
+        recipe_run_id=recipe_run_id, workflow_version="version", workflow_name="name"
+    ) as task:
+        assert task.arm_id == arm_id
+        assert task.run() == 5

dkist_processing_trend/tests/test_base_tasks.py

@@ -0,0 +1,86 @@
+import random
+
+import pytest
+from dkist_processing_common.codecs.str import str_decoder
+from dkist_processing_common.codecs.str import str_encoder
+
+from dkist_processing_trend.models.constants import TrendConstants
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.tasks.trend_base import TrendArmTaskBase
+from dkist_processing_trend.tasks.trend_base import TrendTaskBase
+from dkist_processing_trend.tests.conftest import TrendTestingConstantsDatabase
+
+
+@pytest.fixture
+def trend_science_task(recipe_run_id, link_constants_db):
+    class ScienceTask(TrendTaskBase):
+        def run(self): ...
+
+    link_constants_db(recipe_run_id, TrendTestingConstantsDatabase())
+    with ScienceTask(
+        recipe_run_id=recipe_run_id, workflow_name="name", workflow_version="version"
+    ) as task:
+        yield task
+        task._purge
+
+
+class ArmScienceTask(TrendArmTaskBase):
+    def run(self):
+        self.write(data="Ran .run!", tags=[TrendTag.task("ALL_DONE")], encoder=str_encoder)
+
+
+@pytest.fixture
+def make_trend_arm_science_task(recipe_run_id, link_constants_db):
+
+    def make_task(arm_id):
+        link_constants_db(recipe_run_id, TrendTestingConstantsDatabase())
+        with ArmScienceTask(
+            arm_id=arm_id,
+            recipe_run_id=recipe_run_id,
+            workflow_name="name",
+            workflow_version="version",
+        ) as task:
+            yield task
+            task._purge
+
+    return make_task
+
+
+def test_trend_base_task(trend_science_task):
+    """
+    Given: A task class that is a subclass of `TrendTaskBase`
+    When: Instantiating the class
+    Then: The class' `.constants` property is the correct type
+    """
+    task = trend_science_task
+    assert isinstance(task.constants, TrendConstants)
+    assert task.constants.arm_id_list == [1, 2, 3]
+
+
+@pytest.mark.parametrize(
+    "dataset_has_arm_data",
+    [pytest.param(False, id="arm_w_no_data"), pytest.param(True, id="arm_w_data")],
+)
+def test_trend_arm_base_task(make_trend_arm_science_task, dataset_has_arm_data):
+    """
+    Given: A task class that is a subclass of `TrendArmTaskBase`
+    When: Instantiating that class
+    Then: The `arm_id` argument is correctly set as the `.arm_id` property
+    """
+    constants_arm_id_list = TrendTestingConstantsDatabase().ARM_ID_LIST
+    if dataset_has_arm_data:
+        arm_id = random.choice(constants_arm_id_list)
+    else:
+        arm_id = "LUHILHDLLVUHFSHDUIAHDL"
+        assert arm_id not in constants_arm_id_list
+
+    task = next(make_trend_arm_science_task(arm_id=arm_id))
+    assert task.arm_id == arm_id
+
+    task()
+
+    if dataset_has_arm_data:
+        assert task.count(tags=[TrendTag.task("ALL_DONE")]) == 1
+        assert next(task.read(tags=[TrendTag.task("ALL_DONE")], decoder=str_decoder)) == "Ran .run!"
+    else:
+        assert task.count(tags=[TrendTag.task("ALL_DONE")]) == 0