dkist-processing-test 1.21.3__py3-none-any.whl

dkist_processing_test/__init__.py

@@ -0,0 +1,11 @@
+from importlib.metadata import PackageNotFoundError
+from importlib.metadata import version
+
+from dkist_service_configuration.logging import logger  # first import to call logging.BasicConfig
+
+
+try:
+    __version__ = version(distribution_name=__name__)
+except PackageNotFoundError:
+    # package is not installed
+    __version__ = "unknown"

dkist_processing_test/config.py

@@ -0,0 +1,10 @@
+"""Configurations for the dkist-processing-test package."""
+from dkist_processing_common.config import DKISTProcessingCommonConfiguration
+
+
+class DKISTProcessingTestConfigurations(DKISTProcessingCommonConfiguration):
+    pass  # nothing custom yet
+
+
+dkist_processing_test_configurations = DKISTProcessingTestConfigurations()
+dkist_processing_test_configurations.log_configurations()

dkist_processing_test/models/constants.py

@@ -0,0 +1,18 @@
+from dkist_processing_common.models.constants import BudName
+from dkist_processing_common.models.constants import ConstantsBase
+
+
+class TestConstants(ConstantsBase):
+    """
+    Constants for the test instrument.
+
+    This class is only used on the `TestQualityL0Metrics` task to allow us to check for `num_modstates`.
+    """
+
+    @property
+    def num_modstates(self) -> int:
+        """Return the number of modstates."""
+        # Use .get with default because integration tests use VBI, which doesn't have a modstate key and thus the db
+        # entry won't be there.
+        # In other words, we get the actual db value in unit tests and 1 in integration tests
+        return self._db_dict.get(BudName.num_modstates, 1)

dkist_processing_test/models/parameters.py

@@ -0,0 +1,35 @@
+from datetime import datetime
+
+import numpy as np
+from dkist_processing_common.models.parameters import ParameterBase
+from dkist_processing_common.models.parameters import ParameterWavelengthMixin
+
+
+class TestParameters(ParameterBase, ParameterWavelengthMixin):
+    """Class to test loading parameters from a file."""
+
+    @property
+    def randomness(self) -> (float, float):
+        """A dummy parameter that requires loading a file."""
+        param_dict = self._find_most_recent_past_value(
+            "test_random_data", start_date=datetime.now()
+        )
+        data = self._load_param_value_from_fits(param_dict)
+        mean = np.nanmean(data)
+        std = np.nanstd(data)
+
+        return mean, std
+
+    @property
+    def constant(self) -> float:
+        """A dummy parameter that depends on the same file as a different parameter."""
+        param_dict = self._find_most_recent_past_value("test_random_data")
+        data = self._load_param_value_from_fits(param_dict, hdu=1)
+        constant = np.median(data)
+
+        return float(constant)
+
+    @property
+    def wavelength_category(self) -> str:
+        """A dummy parameter that depends on wavelength."""
+        return self._find_parameter_closest_wavelength("test_wavelength_category")

dkist_processing_test/tasks/__init__.py

@@ -0,0 +1,10 @@
+from dkist_processing_test.tasks.exercise_numba import *
+from dkist_processing_test.tasks.fail import *
+from dkist_processing_test.tasks.fake_science import *
+from dkist_processing_test.tasks.high_memory import *
+from dkist_processing_test.tasks.movie import *
+from dkist_processing_test.tasks.noop import *
+from dkist_processing_test.tasks.parse import *
+from dkist_processing_test.tasks.quality import *
+from dkist_processing_test.tasks.trial_output_data import *
+from dkist_processing_test.tasks.write_l1 import *

dkist_processing_test/tasks/exercise_numba.py

@@ -0,0 +1,42 @@
+"""
+Basic exercising of numba
+"""
+import timeit
+
+import numpy as np
+from dkist_processing_common.tasks import WorkflowTaskBase
+from dkist_service_configuration.logging import logger
+from numba import njit
+
+__all__ = ["ExerciseNumba"]
+
+
+class ExerciseNumba(WorkflowTaskBase):
+    def run(self):
+        bubblesort_numba = njit(self.bubblesort)
+        original = np.linspace(0.0, 10.0, 1001)
+        shuffled_1 = original.copy()
+        np.random.shuffle(shuffled_1)
+        shuffled_2 = shuffled_1.copy()
+        foo_1 = timeit.Timer(lambda: self.bubblesort(shuffled_1), globals=globals())
+        time_1 = foo_1.timeit(100)
+        foo_2 = timeit.Timer(lambda: bubblesort_numba(shuffled_2), globals=globals())
+        time_2 = foo_2.timeit(100)
+        speedup = time_1 / time_2
+        logger.info(f"Normal task execution time: {time_1} secs")
+        logger.info(f"Numba task execution time: {time_2} secs")
+        logger.info(f"ExerciseNumba: Achieved a speedup of {speedup} using numba.")
+        self.speedup = speedup
+        self.sorted_array = shuffled_2
+
+    @staticmethod
+    def bubblesort(x):
+        """Simple bubblesort algorithm copied from numba documentation"""
+        n = len(x)
+        for end in range(n, 1, -1):
+            for i in range(end - 1):
+                cur = x[i]
+                if cur > x[i + 1]:
+                    tmp = x[i]
+                    x[i] = x[i + 1]
+                    x[i + 1] = tmp

dkist_processing_test/tasks/fail.py

@@ -0,0 +1,11 @@
+"""
+Test task for infrastructure integration that will always fail
+"""
+from dkist_processing_core import TaskBase
+
+__all__ = ["FailTask"]
+
+
+class FailTask(TaskBase):
+    def run(self) -> None:
+        raise RuntimeError("Failure is guaranteed")

dkist_processing_test/tasks/fake_science.py

@@ -0,0 +1,101 @@
+"""
+Fake science task
+"""
+import numpy as np
+from astropy.io import fits
+from dkist_processing_common.codecs.fits import fits_array_encoder
+from dkist_processing_common.codecs.fits import fits_hdu_decoder
+from dkist_processing_common.codecs.fits import fits_hdulist_encoder
+from dkist_processing_common.codecs.json import json_encoder
+from dkist_processing_common.models.tags import Tag
+from dkist_processing_common.tasks import WorkflowTaskBase
+from dkist_processing_common.tasks.mixin.input_dataset import InputDatasetMixin
+
+from dkist_processing_test.models.parameters import TestParameters
+
+__all__ = ["GenerateCalibratedData"]
+
+
+class GenerateCalibratedData(WorkflowTaskBase, InputDatasetMixin):
+
+    record_provenance = True
+
+    def __init__(
+        self,
+        recipe_run_id: int,
+        workflow_name: str,
+        workflow_version: str,
+    ):
+        super().__init__(
+            recipe_run_id=recipe_run_id,
+            workflow_name=workflow_name,
+            workflow_version=workflow_version,
+        )
+        self.parameters = TestParameters(
+            self.input_dataset_parameters,
+            obs_ip_start_time=self.constants.obs_ip_start_time,
+            wavelength=2.0,
+        )
+
+    def run(self):
+        rng = np.random.default_rng()
+        with self.apm_task_step("Create debug frame"):
+            self.write(
+                data=np.arange(10), tags=[Tag.frame(), Tag.debug()], encoder=fits_array_encoder
+            )
+
+        with self.apm_task_step("Creating intermediate frame"):
+            self.write(
+                data=np.arange(5),
+                tags=[Tag.frame(), Tag.intermediate(), Tag.task("DUMMY")],
+                encoder=fits_array_encoder,
+            )
+
+        with self.apm_task_step("Creating unique frames"):
+            for _ in range(2):
+                self.write(data=np.arange(3), tags=["FOO", "BAR"], encoder=fits_array_encoder)
+
+            self.write(data={"test": "dictionary"}, tags=["BAZ"], encoder=json_encoder)
+
+        with self.apm_task_step("Creating frames that won't be used"):
+            self.write(data=b"123", tags=[Tag.intermediate(), Tag.task("NOT_USED"), Tag.frame()])
+            self.write(data=b"123", tags=["FOO"])
+
+        with self.apm_task_step("Loop over inputs"):
+            count = 1  # keep a running count to increment the dsps repeat number
+            for hdu in self.read(tags=Tag.input(), decoder=fits_hdu_decoder):
+                header = hdu.header
+                with self.apm_processing_step("Doing some calculations"):
+                    header["DSPSNUM"] = count
+                    data = hdu.data
+
+                    # Just do some weird crap. We don't use the loaded random array directly so that we
+                    # don't have to care that the shapes are the same as the "real" data.
+                    random_signal = rng.normal(*self.parameters.randomness, size=data.shape)
+                    data = (
+                        data + random_signal
+                    )  # Needs to be like this because data will start as int-type
+                    data += self.parameters.constant
+
+                    # Add needed VBI L1 keys that would be computed during real VBI science
+                    header["VBINMOSC"] = self.constants.num_dsps_repeats
+                    header["VBICMOSC"] = count
+
+                    output_hdu = fits.PrimaryHDU(data=data, header=header)
+
+                    wavelength_category = self.parameters.wavelength_category
+                    header["WAVECAT"] = wavelength_category
+
+                with self.apm_writing_step("Writing data"):
+                    output_hdul = fits.HDUList([output_hdu])
+                    self.write(
+                        data=output_hdul,
+                        tags=[
+                            Tag.calibrated(),
+                            Tag.frame(),
+                            Tag.stokes("I"),
+                            Tag.dsps_repeat(count),
+                        ],
+                        encoder=fits_hdulist_encoder,
+                    )
+                count += 1

dkist_processing_test/tasks/high_memory.py

@@ -0,0 +1,20 @@
+"""
+Test task for the 'high_memory' resource queue
+"""
+from time import sleep
+
+from dkist_processing_core import TaskBase
+from dkist_service_configuration.logging import logger
+
+__all__ = ["HighMemoryTask"]
+
+
+class HighMemoryTask(TaskBase):
+    def run(self) -> None:
+        logger.info("Starting High Memory Task")
+        one_gibibyte = int(9.7e9)
+        use_memory = bytearray(one_gibibyte)
+        logger.info(f"Using Memory {one_gibibyte = }")
+        sleep(5)
+        use_memory = None
+        logger.info(f"Memory De-Referenced.")

dkist_processing_test/tasks/manual.py

@@ -0,0 +1,26 @@
+"""
+Tasks simulating manual intervention.
+Manual tasks are expected to write provenance records regardless of they always record provenance.
+"""
+from dkist_processing_common.tasks import WorkflowTaskBase
+
+
+__all__ = ["ManualWithProvenance", "ManualWithoutProvenance"]
+
+
+class ManualBase(WorkflowTaskBase):
+    def run(self):
+        with self.apm_task_step("NoOp"):
+            pass
+
+
+class ManualWithProvenance(ManualBase):
+
+    record_provenance = True
+    is_task_manual = True
+
+
+class ManualWithoutProvenance(ManualBase):
+
+    record_provenance = False
+    is_task_manual = True

dkist_processing_test/tasks/movie.py

@@ -0,0 +1,53 @@
+"""
+Fake MakeMovieFrames and AssembleTestMovie
+"""
+import numpy as np
+from astropy.io import fits
+from dkist_processing_common.codecs.fits import fits_hdu_decoder
+from dkist_processing_common.codecs.fits import fits_hdulist_encoder
+from dkist_processing_common.models.tags import Tag
+from dkist_processing_common.parsers.l1_fits_access import L1FitsAccess
+from dkist_processing_common.tasks import AssembleMovie
+from dkist_processing_common.tasks import WorkflowTaskBase
+from PIL import ImageDraw
+
+
+__all__ = ["MakeTestMovieFrames", "AssembleTestMovie"]
+
+
+class MakeTestMovieFrames(WorkflowTaskBase):
+    """
+    Take each output frame, copy the header and data and write out
+    as a movie frame
+    """
+
+    def run(self):
+        for d in range(1, self.constants.num_dsps_repeats + 1):
+            with self.apm_task_step(f"Workign on dsps repeat {d}"):
+                for hdu in self.read(
+                    tags=[Tag.calibrated(), Tag.dsps_repeat(d)], decoder=fits_hdu_decoder
+                ):
+                    header = hdu.header
+                    data = np.squeeze(hdu.data)
+                    output_hdu = fits.PrimaryHDU(data=data, header=header)
+                    output_hdul = fits.HDUList([output_hdu])
+
+                    with self.apm_writing_step("Writing data"):
+                        self.write(
+                            data=output_hdul,
+                            tags=[Tag.movie_frame(), Tag.dsps_repeat(d)],
+                            encoder=fits_hdulist_encoder,
+                        )
+
+
+class AssembleTestMovie(AssembleMovie):
+    """
+    A shell to extend the AssembleMovie class for the end-to-end test.
+    """
+
+    @property
+    def fits_parsing_class(self):
+        return L1FitsAccess
+
+    def write_overlay(self, draw: ImageDraw, fits_obj: L1FitsAccess) -> None:
+        pass

dkist_processing_test/tasks/noop.py

@@ -0,0 +1,15 @@
+"""
+Test task for infrastructure integration
+"""
+from dkist_processing_core import TaskBase
+
+__all__ = ["NoOpTask", "NoOpTask2"]
+
+
+class NoOpTask(TaskBase):
+    def run(self) -> None:
+        pass
+
+
+class NoOpTask2(NoOpTask):
+    pass

dkist_processing_test/tasks/parse.py

@@ -0,0 +1,88 @@
+"""Parse Task definition."""
+from typing import TypeVar
+
+from dkist_processing_common.models.constants import BudName
+from dkist_processing_common.models.flower_pot import SpilledDirt
+from dkist_processing_common.models.flower_pot import Stem
+from dkist_processing_common.models.flower_pot import Thorn
+from dkist_processing_common.models.tags import StemName
+from dkist_processing_common.parsers.l0_fits_access import L0FitsAccess
+from dkist_processing_common.parsers.single_value_single_key_flower import (
+    SingleValueSingleKeyFlower,
+)
+from dkist_processing_common.parsers.unique_bud import UniqueBud
+from dkist_processing_common.tasks import ParseL0InputDataBase
+
+__all__ = ["ParseL0TestInputData"]
+
+S = TypeVar("S", bound=Stem)
+
+
+class TotalDspsRepeatsBud(UniqueBud):
+    def __init__(self):
+        super().__init__(
+            constant_name=BudName.num_dsps_repeats.value, metadata_key="num_dsps_repeats"
+        )
+
+    def setter(self, fits_obj: L0FitsAccess):
+        if fits_obj.ip_task_type != "observe":
+            return SpilledDirt
+        return super().setter(fits_obj)
+
+
+class ObsIpStartTimeBud(UniqueBud):
+    def __init__(self):
+        super().__init__(
+            constant_name=BudName.obs_ip_start_time.value, metadata_key="ip_start_time"
+        )
+
+    def setter(self, fits_obj: L0FitsAccess):
+        if fits_obj.ip_task_type != "observe":
+            return SpilledDirt
+        return super().setter(fits_obj)
+
+
+class DspsRepeatNumberFlower(SingleValueSingleKeyFlower):
+    def __init__(self):
+        super().__init__(
+            tag_stem_name=StemName.dsps_repeat.value, metadata_key="current_dsps_repeat"
+        )
+
+    def setter(self, fits_obj: L0FitsAccess):
+        if fits_obj.ip_task_type != "observe":
+            return SpilledDirt
+        return super().setter(fits_obj)
+
+
+class PickyDummyBud(Stem):
+    """Exists to do literally nothing"""
+
+    def setter(self, fits_obj: L0FitsAccess):
+        if fits_obj.ip_task_type == "bad value":
+            raise ValueError("This task type is bad!")
+
+    def getter(self, key):
+        return Thorn
+
+
+class ParseL0TestInputData(ParseL0InputDataBase):
+    @property
+    def fits_parsing_class(self):
+        return L0FitsAccess
+
+    @property
+    def tag_flowers(self) -> list[S]:
+        return super().tag_flowers + [
+            SingleValueSingleKeyFlower(
+                tag_stem_name=StemName.task.value, metadata_key="ip_task_type"
+            ),
+            DspsRepeatNumberFlower(),
+        ]
+
+    @property
+    def constant_buds(self) -> list[S]:
+        return super().constant_buds + [
+            TotalDspsRepeatsBud(),
+            PickyDummyBud(stem_name="PICKY_BUD"),
+            ObsIpStartTimeBud(),
+        ]

dkist_processing_test/tasks/quality.py

@@ -0,0 +1,30 @@
+"""Quality task definition."""
+from typing import Iterable
+from typing import Type
+
+from dkist_processing_common.models.constants import ConstantsBase
+from dkist_processing_common.tasks import AssembleQualityData
+from dkist_processing_common.tasks import QualityL0Metrics
+
+__all__ = ["TestQualityL0Metrics", "TestAssembleQualityData"]
+
+from dkist_processing_test.models.constants import TestConstants
+
+
+class TestQualityL0Metrics(QualityL0Metrics):
+    @property
+    def constants_model_class(self) -> Type[ConstantsBase]:
+        return TestConstants
+
+    @property
+    def modstate_list(self) -> Iterable[int] | None:
+        if self.constants.num_modstates > 1:
+            return range(1, self.constants.num_modstates + 1)
+
+        return None
+
+
+class TestAssembleQualityData(AssembleQualityData):
+    @property
+    def polcal_label_list(self) -> list[str] | None:
+        return ["Beam 1"]

dkist_processing_test/tasks/trial_output_data.py

@@ -0,0 +1,46 @@
+"""Trial tasks"""
+from dkist_processing_common.tasks.mixin.globus import GlobusTransferItem
+from dkist_processing_common.tasks.trial_output_data import TransferTrialDataBase
+from dkist_service_configuration.logging import logger
+
+__all__ = ["TransferTestTrialData"]
+
+
+class TransferTestTrialData(TransferTrialDataBase):
+    @property
+    def intermediate_task_names(self) -> list[str]:
+        """Return a list of intermediate tasks we want to transfer.
+
+        Just a dummy task for testing.
+        """
+        return ["DUMMY"]
+
+    def build_transfer_list(self) -> list[GlobusTransferItem]:
+        """
+        Build a list containing all files we want to transfer to the trial environment.
+
+        For the purposes of testing we try to exercise all of the provided helper methods.
+        """
+        transfer_list = []
+
+        transfer_list += self.build_debug_frame_transfer_list()
+
+        transfer_list += self.build_intermediate_frame_transfer_list()
+
+        transfer_list += self.build_output_frame_transfer_list()
+
+        transfer_list += self.build_output_dataset_inventory_transfer_list()
+
+        transfer_list += self.build_output_asdf_transfer_list()
+
+        transfer_list += self.build_output_quality_data_transfer_list()
+
+        transfer_list += self.build_output_quality_report_transfer_list()
+
+        transfer_list += self.build_output_movie_transfer_list()
+
+        transfer_list += self.build_transfer_list_from_tag_lists([["FOO", "BAR"], ["BAZ"]])
+
+        logger.info(f"{transfer_list = }")
+
+        return transfer_list

dkist_processing_test/tasks/write_l1.py

@@ -0,0 +1,64 @@
+from typing import Literal
+
+import astropy.units as u
+from astropy.io import fits
+from astropy.time import Time
+from astropy.time import TimeDelta
+from dkist_processing_common.models.wavelength import WavelengthRange
+from dkist_processing_common.tasks import WriteL1Frame
+
+__all__ = ["WriteL1Data"]
+
+
+class WriteL1Data(WriteL1Frame):
+    def add_dataset_headers(
+        self, header: fits.Header, stokes: Literal["I", "Q", "U", "V"]
+    ) -> fits.Header:
+        header["DAAXES"] = 2
+        header["DEAXES"] = 1
+        header["DNAXIS"] = 3
+        header["LEVEL"] = 1
+        header["WAVEREF"] = "Air"
+        header["WAVEUNIT"] = -9
+        header["DINDEX3"] = 3
+        header["DNAXIS1"] = header["NAXIS1"]
+        header["DNAXIS2"] = header["NAXIS2"]
+        header["DNAXIS3"] = 10
+        header["DPNAME1"] = "spatial x"
+        header["DPNAME2"] = "spatial y"
+        header["DPNAME3"] = "frame number"
+        header["DTYPE1"] = "SPATIAL"
+        header["DTYPE2"] = "SPATIAL"
+        header["DTYPE3"] = "TEMPORAL"
+        header["DUNIT1"] = "arcsec"
+        header["DUNIT2"] = "arcsec"
+        header["DUNIT3"] = "s"
+        header["DWNAME1"] = "helioprojective longitude"
+        header["DWNAME2"] = "helioprojective latitude"
+        header["DWNAME3"] = "time"
+        header["NBIN"] = 1
+        for i in range(1, header["NAXIS"] + 1):
+            header[f"NBIN{i}"] = 1
+
+        return header
+
+    def calculate_date_end(self, header: fits.Header) -> str:
+        """
+        Calculate the VBI specific version of the "DATE-END" keyword.
+
+        Parameters
+        ----------
+        header
+            The input fits header
+
+        Returns
+        -------
+        The isot formatted string of the DATE-END keyword value
+        """
+        return (
+            Time(header["DATE-BEG"], format="isot", precision=6)
+            + TimeDelta(float(header["TEXPOSUR"]) / 1000, format="sec")
+        ).to_value("isot")
+
+    def get_wavelength_range(self, header: fits.Header) -> WavelengthRange:
+        return WavelengthRange(min=1075.0 * u.nm, max=1085.0 * u.nm)

dkist_processing_test/tests/conftest.py

@@ -0,0 +1,64 @@
+"""Global test fixture configuration"""
+from random import randint
+
+import numpy as np
+import pytest
+from astropy.io import fits
+from dkist_data_simulator.spec122 import Spec122Dataset
+from dkist_header_validator.translator import translate_spec122_to_spec214_l0
+
+
+@pytest.fixture()
+def recipe_run_id():
+    return randint(0, 99999)
+
+
+class S122Headers(Spec122Dataset):
+    def __init__(
+        self,
+        array_shape: tuple[int, ...],
+        num_steps: int = 4,
+        num_exp_per_step: int = 1,
+        num_dsps_repeats: int = 5,
+        time_delta: float = 10.0,
+        instrument: str = "vbi",
+    ):
+        dataset_shape = (num_exp_per_step * num_steps * num_dsps_repeats,) + array_shape[-2:]
+        super().__init__(
+            dataset_shape=dataset_shape,
+            array_shape=array_shape,
+            time_delta=time_delta,
+            instrument=instrument,
+        )
+        self.num_steps = num_steps
+        self.num_exp_per_step = num_exp_per_step
+        self.num_dsps_repeats = num_dsps_repeats
+
+
+def generate_214_l0_fits_frame(
+    s122_header: fits.Header, data: np.ndarray | None = None
+) -> fits.HDUList:
+    """Convert S122 header into 214 L0"""
+    if data is None:
+        data = np.ones((1, 10, 10))
+    translated_header = translate_spec122_to_spec214_l0(s122_header)
+    del translated_header["COMMENT"]
+    hdu = fits.PrimaryHDU(data=data, header=fits.Header(translated_header))
+    return fits.HDUList([hdu])
+
+
+@pytest.fixture(scope="session")
+def parameter_file_object_key() -> str:
+    return "random.fits"
+
+
+@pytest.fixture(scope="session")
+def random_parameter_hdulist() -> (fits.HDUList, float, float, float):
+    rng = np.random.default_rng()
+    mu, std = 10.0, 2.0
+    const = 5.0
+    rand_data = rng.normal(mu, std, size=(100, 100))
+    const_data = np.ones((10, 10)) * const
+    hdul = fits.HDUList([fits.PrimaryHDU(rand_data), fits.ImageHDU(const_data)])
+
+    return hdul, mu, std, const