dkist-processing-cryonirsp 1.3.4__py3-none-any.whl

dkist_processing_cryonirsp/tasks/ci_science.py

@@ -0,0 +1,169 @@
+"""Cryonirsp CI science calibration task."""
+from collections import defaultdict
+
+from dkist_service_configuration.logging import logger
+
+from dkist_processing_cryonirsp.models.exposure_conditions import ExposureConditions
+from dkist_processing_cryonirsp.models.tags import CryonirspTag
+from dkist_processing_cryonirsp.parsers.cryonirsp_l0_fits_access import CryonirspL0FitsAccess
+from dkist_processing_cryonirsp.tasks.science_base import CalibrationCollection
+from dkist_processing_cryonirsp.tasks.science_base import ScienceCalibrationBase
+
+__all__ = ["CIScienceCalibration"]
+
+
+class CIScienceCalibration(ScienceCalibrationBase):
+    """Task class for Cryonirsp CI science calibration of polarized and non-polarized data."""
+
+    CI_BEAM = 1
+
+    def calibrate_and_write_frames(self, calibrations: CalibrationCollection):
+        """
+        Completely calibrate all science frames.
+
+        - Apply dark and gain corrections
+        - Demodulate if needed
+        - Apply telescope correction, if needed
+        - Write calibrated arrays
+        """
+        for exposure_conditions in self.constants.observe_exposure_conditions_list:
+            for map_scan in range(1, self.constants.num_map_scans + 1):
+                for scan_step in range(1, self.constants.num_scan_steps + 1):
+                    for meas_num in range(1, self.constants.num_meas + 1):
+                        if self.constants.correct_for_polarization:
+                            calibrated_object = self.calibrate_polarimetric_frames(
+                                exposure_conditions=exposure_conditions,
+                                map_scan=map_scan,
+                                scan_step=scan_step,
+                                meas_num=meas_num,
+                                calibrations=calibrations,
+                            )
+                        else:
+                            calibrated_object = self.calibrate_intensity_only_frames(
+                                exposure_conditions=exposure_conditions,
+                                map_scan=map_scan,
+                                scan_step=scan_step,
+                                meas_num=meas_num,
+                                calibrations=calibrations,
+                            )
+
+                        logging_str = f"{exposure_conditions = }, {map_scan = }, {scan_step = } and {meas_num = }"
+                        logger.info(f"Writing calibrated arrays for {logging_str}")
+                        self.write_calibrated_object(
+                            calibrated_object,
+                            map_scan=map_scan,
+                            scan_step=scan_step,
+                            meas_num=meas_num,
+                        )
+
+    def calibrate_polarimetric_frames(
+        self,
+        *,
+        exposure_conditions: ExposureConditions,
+        map_scan: int,
+        scan_step: int,
+        meas_num: int,
+        calibrations: CalibrationCollection,
+    ) -> CryonirspL0FitsAccess:
+        """
+        Completely calibrate polarimetric science frames.
+
+        - Apply dark and gain corrections
+        - Demodulate
+        - Apply telescope correction
+        """
+        logging_str = f"{exposure_conditions = }, {map_scan = }, {scan_step = } and {meas_num = }"
+        logger.info(f"Processing polarimetric observe frames from {logging_str}")
+
+        intermediate_array, intermediate_header = self.correct_and_demodulate(
+            beam=self.CI_BEAM,
+            meas_num=meas_num,
+            scan_step=scan_step,
+            map_scan=map_scan,
+            exposure_conditions=exposure_conditions,
+            calibrations=calibrations,
+        )
+
+        intermediate_object = self.wrap_array_and_header_in_fits_access(
+            intermediate_array, intermediate_header
+        )
+
+        logger.info(f"Correcting telescope polarization for {logging_str}")
+        calibrated_object = self.telescope_polarization_correction(intermediate_object)
+
+        return calibrated_object
+
+    def calibrate_intensity_only_frames(
+        self,
+        *,
+        exposure_conditions: ExposureConditions,
+        map_scan: int,
+        scan_step: int,
+        meas_num: int,
+        calibrations: CalibrationCollection,
+    ) -> CryonirspL0FitsAccess:
+        """
+        Completely calibrate non-polarimetric science frames.
+
+        - Apply dark and gain corrections
+        """
+        logging_str = f"{exposure_conditions = }, {map_scan = }, {scan_step = } and {meas_num = }"
+        logger.info(f"Processing Stokes-I observe frames from {logging_str}")
+        intermediate_array, intermediate_header = self.apply_basic_corrections(
+            beam=self.CI_BEAM,
+            modstate=1,
+            meas_num=meas_num,
+            scan_step=scan_step,
+            map_scan=map_scan,
+            exposure_conditions=exposure_conditions,
+            calibrations=calibrations,
+        )
+        intermediate_header = self.compute_date_keys(intermediate_header)
+
+        intermediate_array = self.add_stokes_dimension_to_intensity_only_array(intermediate_array)
+
+        calibrated_object = self.wrap_array_and_header_in_fits_access(
+            intermediate_array, intermediate_header
+        )
+
+        return calibrated_object
+
+    def collect_calibration_objects(self) -> CalibrationCollection:
+        """
+        Collect *all* calibration for all modstates, and exposure times.
+
+        Doing this once here prevents lots of reads as we reduce the science data.
+        """
+        dark_dict = defaultdict(dict)
+        solar_dict = dict()
+        demod_dict = dict() if self.constants.correct_for_polarization else None
+
+        beam = self.CI_BEAM
+        # Load the dark arrays
+        for exposure_conditions in self.constants.observe_exposure_conditions_list:
+            dark_array = self.intermediate_frame_load_dark_array(
+                beam=beam, exposure_conditions=exposure_conditions
+            )
+            dark_dict[CryonirspTag.beam(beam)][
+                CryonirspTag.exposure_conditions(exposure_conditions)
+            ] = dark_array
+
+        # Load the gain arrays
+        solar_dict[CryonirspTag.beam(beam)] = self.intermediate_frame_load_solar_gain_array(
+            beam=beam,
+        )
+
+        # Load the demod matrices
+        if self.constants.correct_for_polarization:
+            demod_dict[CryonirspTag.beam(beam)] = self.intermediate_frame_load_demod_matrices(
+                beam=beam
+            )
+
+        return CalibrationCollection(
+            dark=dark_dict,
+            angle=None,
+            spec_shift=None,
+            state_offset=None,
+            solar_gain=solar_dict,
+            demod_matrices=demod_dict,
+        )

dkist_processing_cryonirsp/tasks/cryonirsp_base.py

@@ -0,0 +1,67 @@
+"""CryoNIRSP base class."""
+from abc import ABC
+
+from dkist_processing_common.tasks import WorkflowTaskBase
+from dkist_processing_common.tasks.mixin.input_dataset import InputDatasetMixin
+from dkist_processing_common.tasks.mixin.quality import QualityMixin
+
+from dkist_processing_cryonirsp.models.constants import CryonirspConstants
+from dkist_processing_cryonirsp.models.parameters import CryonirspParameters
+from dkist_processing_cryonirsp.tasks.mixin.beam_access import BeamAccessMixin
+from dkist_processing_cryonirsp.tasks.mixin.corrections import CorrectionsMixin
+from dkist_processing_cryonirsp.tasks.mixin.intermediate_frame import (
+    IntermediateFrameMixin,
+)
+from dkist_processing_cryonirsp.tasks.mixin.linearized_frame import (
+    LinearizedFrameMixin,
+)
+
+
+class CryonirspTaskBase(
+    WorkflowTaskBase,
+    InputDatasetMixin,
+    BeamAccessMixin,
+    LinearizedFrameMixin,
+    IntermediateFrameMixin,
+    CorrectionsMixin,
+    QualityMixin,
+    ABC,
+):
+    """
+    Task class for base CryoNIRSP tasks that occur after linearization.
+
+    Parameters
+    ----------
+    recipe_run_id : int
+        id of the recipe run used to identify the workflow run this task is part of
+    workflow_name : str
+        name of the workflow to which this instance of the task belongs
+    workflow_version : str
+        version of the workflow to which this instance of the task belongs
+    """
+
+    # So tab completion shows all the Cryonirsp constants
+    constants: CryonirspConstants
+
+    @property
+    def constants_model_class(self):
+        """Get CryoNIRSP pipeline constants."""
+        return CryonirspConstants
+
+    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 = CryonirspParameters(
+            self.input_dataset_parameters,
+            obs_ip_start_time=self.constants.obs_ip_start_time,
+            wavelength=self.constants.wavelength,
+            arm_id=self.constants.arm_id,
+        )

dkist_processing_cryonirsp/tasks/dark.py

@@ -0,0 +1,98 @@
+"""CryoNIRSP dark calibration task."""
+from dkist_processing_common.models.task_name import TaskName
+from dkist_processing_math.statistics import average_numpy_arrays
+from dkist_service_configuration.logging import logger
+
+from dkist_processing_cryonirsp.models.tags import CryonirspTag
+from dkist_processing_cryonirsp.tasks.cryonirsp_base import CryonirspTaskBase
+
+__all__ = ["DarkCalibration"]
+
+
+class DarkCalibration(CryonirspTaskBase):
+    """Task class for calculation of the averaged dark frame for a CryoNIRSP calibration run.
+
+    Parameters
+    ----------
+    recipe_run_id : int
+        id of the recipe run used to identify the workflow run this task is part of
+    workflow_name : str
+        name of the workflow to which this instance of the task belongs
+    workflow_version : str
+        version of the workflow to which this instance of the task belongs
+
+    """
+
+    record_provenance = True
+
+    def run(self):
+        """
+        Run method for the task.
+
+        - Gather input dark frames
+        - Calculate average dark
+        - Write average dark
+        - Record quality metrics
+
+        Returns
+        -------
+        None
+
+        """
+        target_exposure_conditions = (
+            self.constants.non_dark_and_non_polcal_task_exposure_conditions_list
+        )
+
+        logger.info(f"{target_exposure_conditions = }")
+        with self.apm_task_step(
+            f"Calculating dark frames for {self.constants.num_beams} beams and {len(target_exposure_conditions)} exp times"
+        ):
+            total_dark_frames_used = 0
+            for exposure_conditions in target_exposure_conditions:
+                for beam in range(1, self.constants.num_beams + 1):
+                    logger.info(
+                        f"Gathering input dark frames for {exposure_conditions = } and {beam = }"
+                    )
+                    dark_tags = [
+                        CryonirspTag.linearized(),
+                        CryonirspTag.frame(),
+                        CryonirspTag.task_dark(),
+                        CryonirspTag.exposure_conditions(exposure_conditions),
+                    ]
+                    current_exp_dark_count = self.scratch.count_all(tags=dark_tags)
+                    total_dark_frames_used += current_exp_dark_count
+                    linearized_dark_arrays = self.linearized_frame_dark_array_generator(
+                        exposure_conditions=exposure_conditions, beam=beam
+                    )
+
+                    with self.apm_processing_step(
+                        f"Calculating dark for {exposure_conditions = } and {beam = }"
+                    ):
+                        averaged_dark_array = average_numpy_arrays(linearized_dark_arrays)
+
+                    with self.apm_writing_step(
+                        f"Writing dark for {exposure_conditions = } {beam = }"
+                    ):
+                        self.intermediate_frame_write_arrays(
+                            averaged_dark_array,
+                            beam=beam,
+                            task_tag=CryonirspTag.task_dark(),
+                            exposure_conditions=exposure_conditions,
+                        )
+
+        with self.apm_processing_step("Computing and logging quality metrics"):
+            no_of_raw_dark_frames: int = self.scratch.count_all(
+                tags=[
+                    CryonirspTag.linearized(),
+                    CryonirspTag.frame(),
+                    CryonirspTag.task_dark(),
+                ],
+            )
+            unused_count = int(
+                no_of_raw_dark_frames - (total_dark_frames_used // self.constants.num_beams)
+            )
+            self.quality_store_task_type_counts(
+                task_type=TaskName.dark.value,
+                total_frames=no_of_raw_dark_frames,
+                frames_not_used=unused_count,
+            )

dkist_processing_cryonirsp/tasks/gain.py

@@ -0,0 +1,251 @@
+"""Cryo gain task."""
+from abc import abstractmethod
+from typing import Callable
+
+import numpy as np
+from dkist_processing_common.models.task_name import TaskName
+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_cryonirsp.models.exposure_conditions import ExposureConditions
+from dkist_processing_cryonirsp.models.tags import CryonirspTag
+from dkist_processing_cryonirsp.tasks.cryonirsp_base import CryonirspTaskBase
+
+__all__ = ["LampGainCalibration", "CISolarGainCalibration"]
+
+
+class GainCalibrationBase(CryonirspTaskBase):
+    """
+    Base task class for calculation of average lamp or solar gains for CI and average lamp gains for SP.
+
+    Parameters
+    ----------
+    recipe_run_id : int
+        id of the recipe run used to identify the workflow run this task is part of
+    workflow_name : str
+        name of the workflow to which this instance of the task belongs
+    workflow_version : str
+        version of the workflow to which this instance of the task belongs
+
+    """
+
+    @property
+    @abstractmethod
+    def gain_type(self) -> str:
+        """Return the gain type, SOLAR_GAIN or LAMP_GAIN."""
+        pass
+
+    @property
+    @abstractmethod
+    def exposure_conditions(self) -> [ExposureConditions]:
+        """Return the exposure conditions list."""
+        pass
+
+    @property
+    @abstractmethod
+    def gain_array_generator(self) -> Callable:
+        """Return the gain array generator to use based on the gain type."""
+        pass
+
+    @property
+    @abstractmethod
+    def normalize_gain_switch(self) -> bool:
+        """If True then the final gain image is normalized to have a mean of 1."""
+        pass
+
+    record_provenance = True
+
+    def run(self):
+        """
+        Execute the task.
+
+        For each exposure time and beam:
+            - Gather input gain and averaged dark arrays
+            - Calculate average array
+            - Normalize average array
+            - Write average gain array
+            - Record quality metrics
+
+        Returns
+        -------
+        None
+
+        """
+        target_exposure_conditions = self.exposure_conditions
+
+        logger.info(f"{target_exposure_conditions = }")
+        with self.apm_task_step(
+            f"Generate {self.gain_type} for {len(target_exposure_conditions)} exposure times"
+        ):
+            for exposure_conditions in target_exposure_conditions:
+                # NB: By using num_beams = 1 for CI, this method works for both CI and SP
+                for beam in range(1, self.constants.num_beams + 1):
+                    apm_str = f"{beam = } and {exposure_conditions = }"
+                    with self.apm_processing_step(f"Remove dark signal for {apm_str}"):
+                        dark_array = self.intermediate_frame_load_dark_array(
+                            beam=beam, exposure_conditions=exposure_conditions
+                        )
+
+                        avg_gain_array = self.compute_average_gain_array(
+                            beam=beam, exposure_conditions=exposure_conditions
+                        )
+
+                        dark_corrected_gain_array = next(
+                            subtract_array_from_arrays(avg_gain_array, dark_array)
+                        )
+
+                    with self.apm_processing_step(f"Correct bad pixels for {apm_str}"):
+                        bad_pixel_map = self.intermediate_frame_load_bad_pixel_map(beam=beam)
+                        bad_pixel_corrected_array = self.corrections_correct_bad_pixels(
+                            dark_corrected_gain_array, bad_pixel_map
+                        )
+
+                    if self.normalize_gain_switch:
+                        with self.apm_processing_step(f"Normalize final gain for {apm_str}"):
+                            normalized_gain_array = self.normalize_gain(bad_pixel_corrected_array)
+                    else:
+                        normalized_gain_array = bad_pixel_corrected_array
+
+                    with self.apm_writing_step(
+                        f"Writing gain array for {beam = } and {exposure_conditions = }"
+                    ):
+                        self.intermediate_frame_write_arrays(
+                            normalized_gain_array,
+                            beam=beam,
+                            task=self.gain_type,
+                        )
+
+        with self.apm_processing_step("Computing and logging quality metrics"):
+            no_of_raw_gain_frames: int = self.scratch.count_all(
+                tags=[
+                    CryonirspTag.linearized(),
+                    CryonirspTag.frame(),
+                    CryonirspTag.task(self.gain_type),
+                ],
+            )
+
+            self.quality_store_task_type_counts(
+                task_type=self.gain_type, total_frames=no_of_raw_gain_frames
+            )
+
+    def compute_average_gain_array(
+        self,
+        beam: int,
+        exposure_conditions: ExposureConditions,
+    ) -> np.ndarray:
+        """
+        Compute average gain array for a given exposure conditions and beam.
+
+        Parameters
+        ----------
+        beam : int
+            The number of the beam
+
+        exposure_conditions : float
+            Exposure time
+
+        Returns
+        -------
+        np.ndarray
+        """
+        linearized_gain_arrays = self.gain_array_generator(
+            beam=beam, exposure_conditions=exposure_conditions
+        )
+        averaged_gain_data = average_numpy_arrays(linearized_gain_arrays)
+        return averaged_gain_data
+
+    @staticmethod
+    def normalize_gain(gain_array: np.ndarray) -> np.ndarray:
+        """
+        Normalize gain to a mean of 1.
+
+        Find any residual pixels that are zero valued and set them to 1.
+
+        Parameters
+        ----------
+        gain_array : np.ndarray
+            Dark corrected gain array
+
+        Returns
+        -------
+        np.ndarray
+            Normalized dark-corrected gain array
+
+        """
+        avg = np.nanmean(gain_array)
+        gain_array /= avg
+
+        return gain_array
+
+
+class LampGainCalibration(GainCalibrationBase):
+    """
+    Task class for calculation of an average lamp gain frame for a CryoNIRSP CI or SP calibration run.
+
+    Parameters
+    ----------
+    recipe_run_id : int
+        id of the recipe run used to identify the workflow run this task is part of
+    workflow_name : str
+        name of the workflow to which this instance of the task belongs
+    workflow_version : str
+        version of the workflow to which this instance of the task belongs
+
+    """
+
+    @property
+    def gain_type(self) -> str:
+        """Return the gain type, SOLAR_GAIN or LAMP_GAIN."""
+        return TaskName.lamp_gain.value
+
+    @property
+    def exposure_conditions(self) -> [ExposureConditions]:
+        """Return the exposure conditions list."""
+        return self.constants.lamp_gain_exposure_conditions_list
+
+    @property
+    def gain_array_generator(self) -> Callable:
+        """Return the gain array generator to use based on the gain type."""
+        return self.linearized_frame_lamp_gain_array_generator
+
+    @property
+    def normalize_gain_switch(self) -> True:
+        """Lamp gains should be normalized."""
+        return True
+
+
+class CISolarGainCalibration(GainCalibrationBase):
+    """
+    Task class for calculation of an average solar gain frame for a CryoNIRSP CI calibration run.
+
+    Parameters
+    ----------
+    recipe_run_id : int
+        id of the recipe run used to identify the workflow run this task is part of
+    workflow_name : str
+        name of the workflow to which this instance of the task belongs
+    workflow_version : str
+        version of the workflow to which this instance of the task belongs
+
+    """
+
+    @property
+    def gain_type(self) -> str:
+        """Return the gain type, SOLAR_GAIN or LAMP_GAIN."""
+        return TaskName.solar_gain.value
+
+    @property
+    def exposure_conditions(self) -> [ExposureConditions]:
+        """Return the exposure conditions list."""
+        return self.constants.solar_gain_exposure_conditions_list
+
+    @property
+    def gain_array_generator(self) -> Callable:
+        """Return the gain array generator to use based on the gain type."""
+        return self.linearized_frame_solar_gain_array_generator
+
+    @property
+    def normalize_gain_switch(self) -> False:
+        """We don't want to normalize and Solar Gain images."""
+        return False