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

dkist_processing_trend/parsers/arm_id.py

@@ -0,0 +1,103 @@
+"""Bud for parsing the full set of input ARM IDs."""
+
+import json
+from collections import defaultdict
+from enum import StrEnum
+from typing import DefaultDict
+
+from dkist_processing_common.models.flower_pot import SetStem
+from dkist_processing_common.models.flower_pot import SpilledDirt
+
+from dkist_processing_trend.models.constants import TrendBudName
+from dkist_processing_trend.models.fits_access import TrendMetadataKey
+from dkist_processing_trend.parsers.trend_l0_fits_access import TrendL0FitsAccess
+
+
+class TrendArmIdListBud(SetStem):
+    """
+    Bud for parsing the set of ARM IDs present in the input data.
+
+    The actual value stored in the db is a `list` only because `set` isn't serializable with JSON.
+    """
+
+    def __init__(self):
+        super().__init__(stem_name=TrendBudName.arm_id_list)
+        self.metadata_key = TrendMetadataKey.arm_id.name
+
+    def setter(self, fits_obj: TrendL0FitsAccess) -> int | str:
+        """Return the arm_id property of the given FitsAccess object."""
+        return getattr(fits_obj, self.metadata_key)
+
+    def getter(self) -> list[int | str]:
+        """Return the set of values as a list for database storage."""
+        return list(self.value_set)
+
+
+class ArmIdLookupBud(SetStem):
+    """
+    Bud that stores a header key as a function of arm_id.
+
+    Each arm is allowed to only have a single value of the given key.
+    """
+
+    def __init__(self, constant_name: str, metadata_key: StrEnum):
+        super().__init__(stem_name=constant_name)
+        self.value_metadata_key = metadata_key.name
+        self.mapping: DefaultDict[float, set[float]] = defaultdict(set)
+
+    def setter(self, fits_obj: TrendL0FitsAccess) -> None:
+        """
+        Update the mapping dictionary.
+
+        Parameters
+        ----------
+        fits_obj
+            The input fits object
+
+        Returns
+        -------
+        Updates the dictionary and returns None
+        """
+        arm_id = fits_obj.arm_id
+        value = getattr(fits_obj, self.value_metadata_key)
+        self.mapping[arm_id].add(value)
+        return None
+
+    def getter(self) -> dict[str, float | str]:
+        """
+        Check that each arm only contains a single value and convert each arm's `set` to a single value.
+
+        The output dict has its keys wrapped in `json.loads` to simplify the extraction of both `int` and `str` arm IDs
+        in `~dkist_processing_trend.models.constants.TrendConstants`.
+        """
+        output_mapping = dict()
+        error_list = []
+        for arm_id, value_set in self.mapping.items():
+            if len(value_set) > 1:
+                error_list.append(f"{arm_id = }: {value_set}")
+
+            output_mapping[json.dumps(arm_id)] = value_set.pop()
+
+        if len(error_list) > 0:
+            error_str = f"Some arms have multiple values for {self.value_metadata_key}: "
+            for arm_err in error_list:
+                error_str += f"{arm_err} "
+            raise ValueError(error_str)
+
+        return output_mapping
+
+
+class InstrumentArmIdLookupBud(ArmIdLookupBud):
+    """Subclass of `ArmIdLookupBud` that only ingests data for a given instrument."""
+
+    def __init__(self, constant_name: str, metadata_key: StrEnum, instruments: str | list[str]):
+        super().__init__(constant_name=constant_name, metadata_key=metadata_key)
+        if isinstance(instruments, str):
+            instruments = [instruments]
+        self.instruments = [i.casefold() for i in instruments]
+
+    def setter(self, fits_obj: TrendL0FitsAccess) -> None:
+        """Update the mapping dict only if the input object's instrument attribute matches the instrument associated with this instance."""
+        if fits_obj.instrument.casefold() not in self.instruments:
+            return SpilledDirt
+        return super().setter(fits_obj)

dkist_processing_trend/parsers/instrument_unique_bud.py

@@ -0,0 +1,40 @@
+"""Version of UniqueBud that only parses frame from a specific instrument."""
+
+from enum import StrEnum
+
+from dkist_processing_common.models.flower_pot import SpilledDirt
+from dkist_processing_common.parsers.unique_bud import UniqueBud
+
+from dkist_processing_trend.parsers.trend_l0_fits_access import TrendL0FitsAccess
+
+
+class InstrumentUniqueBud(UniqueBud):
+    """
+    Subclass of `UniqueBud` that only considers objects that have specific task types.
+
+    Parameters
+    ----------
+    constant_name
+        The name for the constant to be defined
+
+    metadata_key
+        The metadata key associated with the constant
+
+    instruments
+        Only consider frames whose header INSTRUME value is in this list
+    """
+
+    def __init__(
+        self, constant_name: str, metadata_key: str | StrEnum, instruments: str | list[str]
+    ):
+        super().__init__(constant_name=constant_name, metadata_key=metadata_key)
+        if isinstance(instruments, str):
+            instruments = [instruments]
+        self.instruments = [i.casefold() for i in instruments]
+
+    def setter(self, fits_obj: TrendL0FitsAccess) -> int | float | bool | str:
+        """Ingest an object's metadata key only if its instrument property is in the allowed list."""
+        if fits_obj.instrument.casefold() in self.instruments:
+            return super().setter(fits_obj)
+
+        return SpilledDirt

dkist_processing_trend/parsers/time.py

@@ -0,0 +1,27 @@
+"""Stems that parse time-related constants and tags."""
+
+from dkist_processing_common.models.flower_pot import SetStem
+
+from dkist_processing_trend.models.constants import TrendBudName
+from dkist_processing_trend.parsers.trend_l0_fits_access import TrendL0FitsAccess
+
+
+class EarliestIpStartTime(SetStem):
+    """
+    Bud for parsing the earliest IP start time in the input data.
+
+    This is akin to the concept of "OBSERVE IP start time" in the instrument pipelines. The wrinkle is that a POLCAL OP
+    consists of multiple IPs, so with this Bud we consider the "time of observation" to be the earliest IP we have. This
+    time is used for choosing parameter values, so high precision is not needed.
+    """
+
+    def __init__(self):
+        super().__init__(stem_name=TrendBudName.earliest_ip_start_time)
+
+    def setter(self, fits_obj: TrendL0FitsAccess) -> str:
+        """Return the IP start time of the given FitsAccess object."""
+        return fits_obj.ip_start_time
+
+    def getter(self) -> str:
+        """Return the earliest IP start time in the set of all start times."""
+        return min(self.value_set)

dkist_processing_trend/parsers/trend_l0_fits_access.py

@@ -0,0 +1,123 @@
+"""Trend FitsAccess classes for accessing fits header values as properties."""
+
+import numpy as np
+from astropy.io import fits
+from dkist_processing_common.parsers.l0_fits_access import L0FitsAccess
+
+from dkist_processing_trend.models.fits_access import CryonirspMetadataKey
+from dkist_processing_trend.models.fits_access import DlnirspMetadataKey
+from dkist_processing_trend.models.fits_access import TrendMetadataKey
+from dkist_processing_trend.models.fits_access import VispMetadataKey
+from dkist_processing_trend.models.instrument import InstrumentName
+
+
+class IncorrectInstrumentError(Exception):
+    """Error to indicate we're trying to access a property that only exists for a different instrument."""
+
+
+class InstrumentSpecificProperty:
+    """
+    A class for defining properties that will produce a custom (and useful) error when accessed.
+
+    The error is an indication that the user is trying to access an instrument-specific property for a `TrendL0FitsAccess`
+    object that doesn't correspond to that particular instrument.
+    """
+
+    def __init__(self, instrument: str):
+        self.instrument = instrument
+
+    # Inspired by https://stackoverflow.com/questions/56997102/python-notimplementederror-for-instance-attributes
+    def __get__(self, *args, **kwargs):
+        raise IncorrectInstrumentError(
+            f"This property is only defined for data from {self.instrument}, which this frame is not."
+        )
+
+
+class TrendL0FitsAccess(L0FitsAccess):
+    """
+    Class to provide easy access to L0 headers for input/linearized trend data.
+
+    This class contains logic that allows the same property to look at different header keys depending on instrument.
+
+    Parameters
+    ----------
+    header
+        The header from which to parse higher-level properties
+
+    data
+        A data array to associate with the current instance
+
+    name
+        The name of the file that was loaded into this FitsAccess object
+
+    auto_squeeze
+        When set to True, dimensions of length 1 will be removed from the array
+    """
+
+    visp_grating_angle: float = InstrumentSpecificProperty(InstrumentName.visp)
+    visp_arm_position: float = InstrumentSpecificProperty(InstrumentName.visp)
+    visp_grating_constant: float = InstrumentSpecificProperty(InstrumentName.visp)
+    cryonirsp_od_filter_name: str = InstrumentSpecificProperty(InstrumentName.cryonirsp)
+    cryonirsp_grating_angle: float = InstrumentSpecificProperty(InstrumentName.cryonirsp)
+    cryonirsp_grating_littrow_angle: float = InstrumentSpecificProperty(InstrumentName.cryonirsp)
+    cryonirsp_grating_constant: float = InstrumentSpecificProperty(InstrumentName.cryonirsp)
+    cryonirsp_slit_width: float = InstrumentSpecificProperty(InstrumentName.cryonirsp)
+    dlnirsp_grating_angle: float = InstrumentSpecificProperty(InstrumentName.dlnirsp)
+    dlnirsp_arm_position: float = InstrumentSpecificProperty(InstrumentName.dlnirsp)
+    dlnirsp_grating_constant: float = InstrumentSpecificProperty(InstrumentName.dlnirsp)
+
+    def __init__(
+        self,
+        *,
+        header: fits.Header,
+        data: np.ndarray | None = None,
+        name: str | None = None,
+        auto_squeeze: bool = True,
+    ):
+        super().__init__(header=header, data=data, name=name, auto_squeeze=auto_squeeze)
+
+        match self.instrument:
+            case InstrumentName.visp:
+                self.number_of_modulator_states: int = header[
+                    VispMetadataKey.number_of_modulator_states
+                ]
+                self.modulator_state: int = header[VispMetadataKey.modulator_state]
+                self.arm_id: int = header[VispMetadataKey.arm_id]
+                self.visp_grating_angle: float = header[TrendMetadataKey.visp_grating_angle]
+                self.visp_arm_position: float = header[TrendMetadataKey.visp_arm_position]
+                self.visp_grating_constant: float = header[TrendMetadataKey.visp_grating_constant]
+
+            case InstrumentName.cryonirsp:
+                self.number_of_modulator_states: int = header[
+                    CryonirspMetadataKey.number_of_modulator_states
+                ]
+                self.modulator_state: int = header[CryonirspMetadataKey.modulator_state]
+                self.arm_id: str = header[CryonirspMetadataKey.arm_id]
+                self.cryonirsp_od_filter_name: str = header[
+                    TrendMetadataKey.cryonirsp_od_filter_name
+                ]
+                self.cryonirsp_grating_angle: float = header[
+                    TrendMetadataKey.cryonirsp_grating_angle
+                ]
+                self.cryonirsp_grating_littrow_angle: float = header[
+                    TrendMetadataKey.cryonirsp_grating_littrow_angle
+                ]
+                self.cryonirsp_grating_constant: float = header[
+                    TrendMetadataKey.cryonirsp_grating_constant
+                ]
+                self.cryonirsp_slit_width: float = header[TrendMetadataKey.cryonirsp_slit_width]
+
+            case InstrumentName.dlnirsp:
+                self.number_of_modulator_states: int = header[
+                    DlnirspMetadataKey.number_of_modulator_states
+                ]
+                self.modulator_state: int = header[DlnirspMetadataKey.modulator_state]
+                self.arm_id: int = header[DlnirspMetadataKey.arm_id]
+                self.dlnirsp_grating_angle: float = header[TrendMetadataKey.dlnirsp_grating_angle]
+                self.dlnirsp_arm_position: float = header[TrendMetadataKey.dlnirsp_arm_position]
+                self.dlnirsp_grating_constant: float = header[
+                    TrendMetadataKey.dlnirsp_grating_constant
+                ]
+
+            case _:
+                raise ValueError(f"Did not expect instrument = {self.instrument}")

dkist_processing_trend/tasks/__init__.py

@@ -0,0 +1,27 @@
+"""Task modules and classes that are used to build up a workflow."""
+
+from dkist_processing_trend.tasks.dark import *
+from dkist_processing_trend.tasks.gain import *
+from dkist_processing_trend.tasks.parse import *
+from dkist_processing_trend.tasks.run_pac_fitter import *
+from dkist_processing_trend.tasks.trend_output_data import *
+from dkist_processing_trend.tasks.write_trend import *
+
+# These imports need to be after all the task imports to avoid circularity
+from dkist_processing_trend.tasks.initialize_arm_tasks import (  # isort: skip
+    initialize_cryonirsp_arm_tasks,
+)
+from dkist_processing_trend.tasks.initialize_arm_tasks import (  # isort: skip
+    initialize_dlnirsp_arm_tasks,
+)
+from dkist_processing_trend.tasks.initialize_arm_tasks import (  # isort: skip
+    initialize_visp_arm_tasks,
+)
+
+# Not used directly, but ensures that the common_configurations are logged early in a task run
+from dkist_processing_trend.config import dkist_processing_trend_configurations  # isort: skip
+
+# Initialize the arm task classes for each instrument
+initialize_visp_arm_tasks()
+initialize_cryonirsp_arm_tasks()
+initialize_dlnirsp_arm_tasks()

dkist_processing_trend/tasks/arm_task_factory.py

@@ -0,0 +1,57 @@
+"""
+Functions for defining and retreiving 'partial' classes with their `__init__` methods pre-hydrated with an ``arm_id``.
+
+This is used to ease the definition of `~dkist_processing_trend.tasks.trend_base.TrendArmTaskBase` with a list of ``arm_ids``.
+"""
+
+# Inspired by https://stackoverflow.com/questions/38911146/python-equivalent-of-functools-partial-for-a-class-constructor
+
+from functools import partialmethod
+from importlib import import_module
+
+
+def _arm_task_name(cls_name: str, arm_id: str | int) -> str:
+    """Define common logic for smashing a class name and arm_id together."""
+    return f"{cls_name}{str(arm_id).capitalize()}"
+
+
+def initialize_arm_task(cls, arm_id: str | int) -> None:
+    """
+    Define a new class that subclasses ``cls`` with its `__init__` pre-populated with the `arm_id` argument.
+
+    This function also updates the attrs of the module of ``cls`` so that the new class appears to be a member of that
+    same module.
+    """
+
+    class ArmTaskClass(cls):
+        __init__ = partialmethod(cls.__init__, arm_id=arm_id)
+        __doc__ = cls.__doc__
+
+    # Give the class a new name based on the arm_id. This is required for the DAG to keep the same task with different
+    # arm_ids separate.
+    ArmTaskClass.__name__ = _arm_task_name(cls.__name__, arm_id)
+    ArmTaskClass.__qualname__ = _arm_task_name(cls.__qualname__, arm_id)
+
+    # Modify the import path so the new class can be imported from its parent's module.
+    ArmTaskClass.__module__ = cls.__module__
+    mod = import_module(cls.__module__)
+    setattr(mod, ArmTaskClass.__name__, ArmTaskClass)
+
+    return
+
+
+def retreive_arm_task(cls, arm_id: str | int):
+    """
+    Return the "arm task" version of a given ``cls`` and ``arm_id``.
+
+    In other words, return the arm task class that was created by `initialize_arm_task`.
+    """
+    parent_module = import_module(cls.__module__)
+    arm_class_name = _arm_task_name(cls.__name__, arm_id)
+    try:
+        arm_class = getattr(parent_module, arm_class_name)
+    except AttributeError as e:
+        raise AttributeError(
+            f"Can't find {arm_class_name} in {parent_module}. Was it initialized with `initialize_arm_task`?"
+        ) from e
+    return arm_class

dkist_processing_trend/tasks/dark.py

@@ -0,0 +1,56 @@
+"""Task for computing the average of the DARK CS steps."""
+
+from dkist_processing_common.codecs.fits import fits_array_decoder
+from dkist_processing_common.codecs.fits import fits_array_encoder
+from dkist_processing_math.statistics import average_numpy_arrays
+
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.tasks.trend_base import TrendArmTaskBase
+
+__all__ = ["DarkCalibration"]
+
+
+class DarkCalibration(TrendArmTaskBase):
+    """
+    Generate an INTERMEDIATE dark frame by averaging all data from the "dark" CS steps.
+
+    Parameters
+    ----------
+    arm_id
+        id of the instrument arm to operate on
+
+    recipe_run_id
+        id of the recipe run used to identify the workflow run this task is part of
+
+    workflow_name
+        name of the workflow to which this instance of the task belongs
+
+    workflow_version
+        version of the workflow to which this instance of the task belongs
+    """
+
+    record_provenance = True
+
+    def run(self):
+        """
+        Produce a single, average dark frame from all of the dark CS steps.
+
+        In other words, collect all frames tagged as POLCAL_DARK, average them together, and write the result to disk.
+        """
+        with self.telemetry_span(f"Producing average dark for arm {self.arm_id}"):
+            input_dark_arrays = self.read(
+                tags=[TrendTag.task_polcal_dark(), TrendTag.arm_id(self.arm_id)],
+                decoder=fits_array_decoder,
+            )
+            averaged_dark_array = average_numpy_arrays(input_dark_arrays)
+
+            self.write(
+                data=averaged_dark_array,
+                tags=[
+                    TrendTag.intermediate(),
+                    TrendTag.frame(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_dark(),
+                ],
+                encoder=fits_array_encoder,
+            )

dkist_processing_trend/tasks/gain.py

@@ -0,0 +1,74 @@
+"""Task for computing the average of the CLEAR CS steps."""
+
+from dkist_processing_common.codecs.fits import fits_array_decoder
+from dkist_processing_common.codecs.fits import fits_array_encoder
+from dkist_processing_math.arithmetic import subtract_array_from_arrays
+from dkist_processing_math.statistics import average_numpy_arrays
+
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.tasks.trend_base import TrendArmTaskBase
+
+__all__ = ["GainCalibration"]
+
+
+class GainCalibration(TrendArmTaskBase):
+    """
+    Generate an INTERMEDIATE gain frame by averaging all data from the "clear" CS steps.
+
+    Parameters
+    ----------
+    arm_id
+        id of the instrument arm to operate on
+
+    recipe_run_id
+        id of the recipe run used to identify the workflow run this task is part of
+
+    workflow_name
+        name of the workflow to which this instance of the task belongs
+
+    workflow_version
+        version of the workflow to which this instance of the task belongs
+    """
+
+    record_provenance = True
+
+    def run(self):
+        """
+        Produce a single, average gain frame from all of the clear CS steps.
+
+        In other words, collect all frames tagged as POLCAL_GAIN, average them together, subtract the INTERMEDIATE DARK
+        signal, and write the result to disk.
+        """
+        with self.telemetry_span(f"Producing average gain for arm {self.arm_id}"):
+            input_gain_arrays = self.read(
+                tags=[TrendTag.task_polcal_gain(), TrendTag.arm_id(self.arm_id)],
+                decoder=fits_array_decoder,
+            )
+            averaged_gain_array = average_numpy_arrays(input_gain_arrays)
+
+            dark_array = next(
+                self.read(
+                    tags=[
+                        TrendTag.intermediate(),
+                        TrendTag.frame(),
+                        TrendTag.task_dark(),
+                        TrendTag.arm_id(self.arm_id),
+                    ],
+                    decoder=fits_array_decoder,
+                )
+            )
+
+            final_gain = next(
+                subtract_array_from_arrays(arrays=averaged_gain_array, array_to_subtract=dark_array)
+            )
+
+            self.write(
+                data=final_gain,
+                tags=[
+                    TrendTag.intermediate(),
+                    TrendTag.frame(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_gain(),
+                ],
+                encoder=fits_array_encoder,
+            )

dkist_processing_trend/tasks/initialize_arm_tasks.py

@@ -0,0 +1,50 @@
+"""
+Initialize "arm class" versions of all the arm-specific tasks needed by each instrument.
+
+The functions in this module need to be called prior to importing the respective instrument's workflow.
+"""
+
+from dkist_processing_trend.models.instrument import CryonirspArms
+from dkist_processing_trend.models.instrument import DlnirspArms
+from dkist_processing_trend.models.instrument import VispArms
+from dkist_processing_trend.tasks.arm_task_factory import initialize_arm_task
+from dkist_processing_trend.tasks.dark import DarkCalibration
+from dkist_processing_trend.tasks.gain import GainCalibration
+from dkist_processing_trend.tasks.run_pac_fitter import RunPacFitter
+from dkist_processing_trend.tasks.visp.visp_dmpd import VispDemodulatedPolcalData
+from dkist_processing_trend.tasks.visp.visp_geometric import VispGeometricCalibration
+from dkist_processing_trend.tasks.visp.visp_prep_fit_data import VispPrepareFitData
+from dkist_processing_trend.tasks.visp.visp_process_demod import VispProcessDemodulationMatrices
+
+VISP_ARM_TASKS = [
+    DarkCalibration,
+    GainCalibration,
+    VispGeometricCalibration,
+    VispPrepareFitData,
+    RunPacFitter,
+    VispProcessDemodulationMatrices,
+    VispDemodulatedPolcalData,
+]
+CRYONIRSP_ARM_TASKS = [DarkCalibration, GainCalibration, RunPacFitter]
+DLNIRSP_ARM_TASKS = [DarkCalibration, GainCalibration, RunPacFitter]
+
+
+def initialize_visp_arm_tasks():
+    """Create arm versions of all arm-specific ViSP tasks."""
+    for arm_id in VispArms:
+        for task_class in VISP_ARM_TASKS:
+            initialize_arm_task(task_class, arm_id=arm_id)
+
+
+def initialize_cryonirsp_arm_tasks():
+    """Create arm versions of all arm-specific CryoNIRSP tasks."""
+    for arm_id in CryonirspArms:
+        for task_class in CRYONIRSP_ARM_TASKS:
+            initialize_arm_task(task_class, arm_id=arm_id)
+
+
+def initialize_dlnirsp_arm_tasks():
+    """Create arm versions of all arm-specific DLNIRSP tasks."""
+    for arm_id in DlnirspArms:
+        for task_class in DLNIRSP_ARM_TASKS:
+            initialize_arm_task(task_class, arm_id=arm_id)