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

dkist_processing_trend/tasks/parse.py

@@ -0,0 +1,169 @@
+"""Tasks for parsing both IR and VIS data."""
+
+from dkist_processing_common.models.constants import BudName
+from dkist_processing_common.models.fits_access import MetadataKey
+from dkist_processing_common.models.flower_pot import Stem
+from dkist_processing_common.models.tags import StemName
+from dkist_processing_common.parsers.cs_step import CSStepFlower
+from dkist_processing_common.parsers.cs_step import NumCSStepBud
+from dkist_processing_common.parsers.id_bud import ContributingIdsBud
+from dkist_processing_common.parsers.single_value_single_key_flower import (
+    SingleValueSingleKeyFlower,
+)
+from dkist_processing_common.parsers.task import PolcalTaskFlower
+from dkist_processing_common.parsers.unique_bud import UniqueBud
+from dkist_processing_common.tasks import ParseL0InputDataBase
+
+from dkist_processing_trend.models.constants import TrendBudName
+from dkist_processing_trend.models.fits_access import TrendMetadataKey
+from dkist_processing_trend.models.instrument import InstrumentName
+from dkist_processing_trend.models.parameters import TrendParsingParameters
+from dkist_processing_trend.models.tags import TrendStemName
+from dkist_processing_trend.parsers.arm_id import ArmIdLookupBud
+from dkist_processing_trend.parsers.arm_id import InstrumentArmIdLookupBud
+from dkist_processing_trend.parsers.arm_id import TrendArmIdListBud
+from dkist_processing_trend.parsers.time import EarliestIpStartTime
+from dkist_processing_trend.parsers.trend_l0_fits_access import TrendL0FitsAccess
+
+__all__ = ["ParseL0TrendInputData"]
+
+
+class ParseL0TrendInputData(ParseL0InputDataBase):
+    """
+    Parse either LINEARIZED data (IR instruments) or INPUT data (VIS instruments).
+
+    Parameters
+    ----------
+    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
+    """
+
+    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 = TrendParsingParameters(scratch=self.scratch)
+
+    @property
+    def fits_parsing_class(self):
+        """FITS access class to use in this task."""
+        return TrendL0FitsAccess
+
+    @property
+    def constant_buds(self) -> list[Stem]:
+        """Return the list of Buds to use when parsing constants."""
+        return [
+            EarliestIpStartTime(),
+            UniqueBud(constant_name=BudName.instrument, metadata_key=MetadataKey.instrument),
+            NumCSStepBud(max_cs_step_time_sec=self.parameters.max_cs_step_time_sec),
+            UniqueBud(
+                constant_name=BudName.num_modstates,
+                metadata_key=TrendMetadataKey.number_of_modulator_states,
+            ),
+            ArmIdLookupBud(
+                constant_name=TrendBudName.wavelength_for_arm, metadata_key=MetadataKey.wavelength
+            ),
+            TrendArmIdListBud(),
+            ContributingIdsBud(
+                constant_name=TrendBudName.observing_program_id_list,
+                metadata_key=MetadataKey.observing_program_execution_id,
+            ),
+            ContributingIdsBud(
+                constant_name=TrendBudName.experiment_id_list,
+                metadata_key=MetadataKey.experiment_id,
+            ),
+            ContributingIdsBud(
+                constant_name=TrendBudName.proposal_id_list, metadata_key=MetadataKey.proposal_id
+            ),
+            ContributingIdsBud(
+                constant_name=TrendBudName.hls_version_list, metadata_key=MetadataKey.hls_version
+            ),
+            # ViSP
+            InstrumentArmIdLookupBud(
+                constant_name=TrendBudName.visp_grating_angle_for_arm,
+                metadata_key=TrendMetadataKey.visp_grating_angle,
+                instruments=InstrumentName.visp,
+            ),
+            InstrumentArmIdLookupBud(
+                constant_name=TrendBudName.visp_arm_position_for_arm,
+                metadata_key=TrendMetadataKey.visp_arm_position,
+                instruments=InstrumentName.visp,
+            ),
+            InstrumentArmIdLookupBud(
+                constant_name=TrendBudName.visp_grating_constant_for_arm,
+                metadata_key=TrendMetadataKey.visp_grating_constant,
+                instruments=InstrumentName.visp,
+            ),
+            # CryoNIRSP
+            InstrumentArmIdLookupBud(
+                constant_name=TrendBudName.cryonirsp_od_filter_name_for_arm,
+                metadata_key=TrendMetadataKey.cryonirsp_od_filter_name,
+                instruments=InstrumentName.cryonirsp,
+            ),
+            InstrumentArmIdLookupBud(
+                constant_name=TrendBudName.cryonirsp_grating_angle_for_arm,
+                metadata_key=TrendMetadataKey.cryonirsp_grating_angle,
+                instruments=InstrumentName.cryonirsp,
+            ),
+            InstrumentArmIdLookupBud(
+                constant_name=TrendBudName.cryonirsp_grating_littrow_angle_for_arm,
+                metadata_key=TrendMetadataKey.cryonirsp_grating_littrow_angle,
+                instruments=InstrumentName.cryonirsp,
+            ),
+            InstrumentArmIdLookupBud(
+                constant_name=TrendBudName.cryonirsp_grating_constant_for_arm,
+                metadata_key=TrendMetadataKey.cryonirsp_grating_constant,
+                instruments=InstrumentName.cryonirsp,
+            ),
+            InstrumentArmIdLookupBud(
+                constant_name=TrendBudName.cryonirsp_slit_width_for_arm,
+                metadata_key=TrendMetadataKey.cryonirsp_slit_width,
+                instruments=InstrumentName.cryonirsp,
+            ),
+            # DLNIRSP
+            InstrumentArmIdLookupBud(
+                constant_name=TrendBudName.dlnirsp_grating_angle_for_arm,
+                metadata_key=TrendMetadataKey.dlnirsp_grating_angle,
+                instruments=InstrumentName.dlnirsp,
+            ),
+            InstrumentArmIdLookupBud(
+                constant_name=TrendBudName.dlnirsp_arm_position_for_arm,
+                metadata_key=TrendMetadataKey.dlnirsp_arm_position,
+                instruments=InstrumentName.dlnirsp,
+            ),
+            InstrumentArmIdLookupBud(
+                constant_name=TrendBudName.dlnirsp_grating_constant_for_arm,
+                metadata_key=TrendMetadataKey.dlnirsp_grating_constant,
+                instruments=InstrumentName.dlnirsp,
+            ),
+        ]
+
+    @property
+    def tag_flowers(self) -> list[Stem]:
+        """Return the list of Flowers to use when parsing tags."""
+        return [
+            SingleValueSingleKeyFlower(
+                tag_stem_name=StemName.task, metadata_key=MetadataKey.ip_task_type
+            ),
+            CSStepFlower(max_cs_step_time_sec=self.parameters.max_cs_step_time_sec),
+            SingleValueSingleKeyFlower(
+                tag_stem_name=StemName.modstate, metadata_key=TrendMetadataKey.modulator_state
+            ),
+            SingleValueSingleKeyFlower(
+                tag_stem_name=TrendStemName.arm_id, metadata_key=TrendMetadataKey.arm_id
+            ),
+            PolcalTaskFlower(),
+        ]

dkist_processing_trend/tasks/prepare_fit_data_base.py

@@ -0,0 +1,257 @@
+"""Base task for preparing input data for the PAC fitter."""
+
+from abc import ABC
+from abc import abstractmethod
+from collections import defaultdict
+
+import numpy as np
+from dkist_processing_common.codecs.asdf import asdf_encoder
+from dkist_processing_common.codecs.fits import fits_array_decoder
+from dkist_processing_common.codecs.fits import fits_header_decoder
+from dkist_processing_math.arithmetic import divide_arrays_by_array
+from dkist_processing_math.arithmetic import subtract_array_from_arrays
+from dkist_processing_math.statistics import average_numpy_arrays
+from dkist_service_configuration.logging import logger
+
+from dkist_processing_trend.models.instrument_options import InstrumentOptionsBase
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.parsers.trend_l0_fits_access import TrendL0FitsAccess
+from dkist_processing_trend.tasks.trend_base import TrendArmTaskBase
+
+
+class PrepareFitDataBase(TrendArmTaskBase, ABC):
+    """
+    Base class for processing INPUT data to prepare it for fitting by `~dkist_processing_trend.tasks.run_pac_fitter.RunPacFitter`.
+
+    This class provides a populated `run` method that does steps common to all instruments (e.g., dark and gain correction).
+
+    Instrument-specific subclasses need to define the following methods:
+
+    * `extract_beam` - For getting a certain beam from a full-FOV array
+    * `apply_global_instrument_options` - Use any `*InstrumentOptions` to produce the "global" PAC input data for a single cs step, beam, and modstate
+    * `apply_local_instrument_options` - Use any `*InstrumentOptions` to produce the "local" PAC input data for a single cs step, beam, and modstate
+
+    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):
+        """
+        Prepare input data for fitting with `dkist_processing_pac`.
+
+        At a high level this involves three steps:
+
+        #. Perform dark and gain corrections to all input frames
+        #. Apply any instrument-specific processing (e.g., binning) to the resulting frames
+        #. Organize the results into the structure expected by `~dkist_processing_trend.tasks.run_pac_fitter.RunPacFitter`
+
+        This method should not be overloaded by instrument-specific subclasses. Rather they should define `extract_beam`,
+        `apply_global_instrument_options`, and `apply_local_instrument_options`, which are called from this method.
+        """
+        logger.info("Loading INTERMEDIATE dark frame")
+        avg_dark = next(
+            self.read(
+                tags=[
+                    TrendTag.intermediate(),
+                    TrendTag.frame(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_dark(),
+                ],
+                decoder=fits_array_decoder,
+            )
+        )
+
+        logger.info("Loading INTERMEDIATE gain frame")
+        avg_gain = next(
+            self.read(
+                tags=[
+                    TrendTag.intermediate(),
+                    TrendTag.frame(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_gain(),
+                ],
+                decoder=fits_array_decoder,
+            )
+        )
+
+        beam_dark_dict = {
+            beam: self.extract_beam(array=avg_dark, beam=beam)
+            for beam in range(1, self.num_beams + 1)
+        }
+        beam_gain_dict = {
+            beam: self.extract_beam(array=avg_gain, beam=beam)
+            for beam in range(1, self.num_beams + 1)
+        }
+
+        with self.telemetry_span("Prepare polcal fit inputs"):
+            for inst_options in self.parameters.instrument_processing_options:
+                for beam in range(1, self.num_beams + 1):
+                    global_obj_dict = defaultdict(list)
+                    local_obj_dict = defaultdict(list)
+                    beam_dark = beam_dark_dict[beam]
+                    beam_gain = beam_gain_dict[beam]
+                    for cs_step in range(self.constants.num_cs_steps):
+                        for modstate in range(1, self.constants.num_modstates + 1):
+                            logger.info(
+                                f"Processing instrument options = {inst_options.name}, {beam = }, {cs_step = }, and {modstate = }"
+                            )
+                            global_obj, local_obj = self.process_single_modstate(
+                                instrument_options=inst_options,
+                                cs_step=cs_step,
+                                modstate=modstate,
+                                beam=beam,
+                                avg_dark=beam_dark,
+                                avg_gain=beam_gain,
+                            )
+                            global_obj_dict[cs_step].append(global_obj)
+                            local_obj_dict[cs_step].append(local_obj)
+
+                    logger.info(
+                        f"Writing PAC inputs for instrument options = {inst_options.name} and {beam = }"
+                    )
+                    self.write_pac_inputs(
+                        global_obj_dict=global_obj_dict,
+                        local_obj_dict=local_obj_dict,
+                        instrument_options=inst_options,
+                        beam=beam,
+                    )
+
+    @property
+    def num_beams(self) -> int:
+        """Return the number of beams associated with this instance's arm_id."""
+        if self.arm_id == "CI":
+            return 1
+        return 2
+
+    def process_single_modstate(
+        self,
+        instrument_options: InstrumentOptionsBase,
+        cs_step: int,
+        modstate: int,
+        beam: int,
+        avg_dark: np.ndarray,
+        avg_gain: np.ndarray,
+    ) -> tuple[TrendL0FitsAccess, TrendL0FitsAccess]:
+        """
+        Process data for a single cs step, beam, and modstate.
+
+        This involves doing dark and gain correction before passing the result to the instrument-defined methods used to
+        extract global and local PAC input data.
+
+        If more than one file is found for the given ``(cs_step, beam, modstate)`` tuple then they are averaged together.
+        In this case, the header in the returned `~dkist_processing_trend.parsers.trend_l0_fits_access.TrendL0FitsAccess`
+        objects is randomly chosen from the input frames.
+        """
+        tags = [
+            TrendTag.task_polcal(),
+            TrendTag.arm_id(self.arm_id),
+            TrendTag.cs_step(cs_step),
+            TrendTag.modstate(modstate),
+        ]
+        logger.info(f"Found {self.count(tags)} frame(s)")
+
+        arrays = self.read(tags=tags, decoder=fits_array_decoder)
+        avg_array = average_numpy_arrays(arrays)
+
+        beam_array = self.extract_beam(array=avg_array, beam=beam)
+
+        dark_corrected_array = subtract_array_from_arrays(
+            arrays=beam_array, array_to_subtract=avg_dark
+        )
+        gain_corrected_array = next(
+            divide_arrays_by_array(arrays=dark_corrected_array, array_to_divide_by=avg_gain)
+        )
+
+        global_array = self.apply_global_instrument_options(
+            array=gain_corrected_array, instrument_options=instrument_options
+        )
+        local_array = self.apply_local_instrument_options(
+            array=gain_corrected_array, instrument_options=instrument_options
+        )
+
+        # Any header from the set of input data will do
+        header = next(self.read(tags=tags, decoder=fits_header_decoder))
+
+        return TrendL0FitsAccess(header=header, data=global_array), TrendL0FitsAccess(
+            header=header, data=local_array
+        )
+
+    def write_pac_inputs(
+        self,
+        global_obj_dict: dict[int, list[TrendL0FitsAccess]],
+        local_obj_dict: dict[int, list[TrendL0FitsAccess]],
+        instrument_options: InstrumentOptionsBase,
+        beam: int,
+    ) -> None:
+        """
+        Write the dicts containing the global and local data structures to disk.
+
+        Each of these data structures is essentially a `~dkist_processing_pac.input_data.drawer.Drawer`, but saved as
+        dicts in an ASDF file for easier serialization. This method converts each
+        `~dkist_processing_trend.parsers.trend_l0_fits_access.TrendL0FitsAccess` object into a dict of "header" and
+        "data" prior to writing.
+        """
+        global_output_dict = {
+            cs_step: [
+                {"header": fits_obj.header_dict, "data": fits_obj.data}
+                for fits_obj in modstate_list
+            ]
+            for cs_step, modstate_list in global_obj_dict.items()
+        }
+        local_output_dict = {
+            cs_step: [
+                {"header": fits_obj.header_dict, "data": fits_obj.data}
+                for fits_obj in modstate_list
+            ]
+            for cs_step, modstate_list in local_obj_dict.items()
+        }
+
+        base_tags = [
+            TrendTag.intermediate(),
+            TrendTag.arm_id(self.arm_id),
+            TrendTag.instrument_processing_options(instrument_options.name),
+            TrendTag.beam(beam),
+        ]
+        self.write(
+            data=global_output_dict,
+            tags=base_tags + [TrendTag.task_global_pac_input()],
+            encoder=asdf_encoder,
+        )
+        self.write(
+            data=local_output_dict,
+            tags=base_tags + [TrendTag.task_local_pac_input()],
+            encoder=asdf_encoder,
+        )
+
+    @abstractmethod
+    def extract_beam(self, array: np.ndarray, beam: int) -> np.ndarray:
+        """Extract the given beam from the input array."""
+
+    @abstractmethod
+    def apply_global_instrument_options(
+        self, array: np.ndarray, instrument_options: InstrumentOptionsBase
+    ) -> np.ndarray:
+        """
+        Do whatever instrument-specific processing is needed to produce a global PAC data point.
+
+        The shape of the final output should be (1, 1).
+        """
+
+    @abstractmethod
+    def apply_local_instrument_options(
+        self, array: np.ndarray, instrument_options: InstrumentOptionsBase
+    ) -> np.ndarray:
+        """Do whatever instrument-specific processing is needed to produce a local PAC data point."""