dkist-processing-cryonirsp 1.4.15__py3-none-any.whl → 1.4.16rc1__py3-none-any.whl

dkist_processing_cryonirsp/tasks/mixin/intermediate_frame.py

@@ -1,193 +0,0 @@
-"""Mixin class providing support for loading and writing intermediate arrays."""
-from collections.abc import Generator
-from collections.abc import Iterable
-
-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_service_configuration.logging import logger
-
-from dkist_processing_cryonirsp.models.exposure_conditions import ExposureConditions
-from dkist_processing_cryonirsp.models.tags import CryonirspTag
-
-
-class IntermediateFrameMixin:
-    """Mixin for methods that support easy loading and writing of intermediate frames."""
-
-    def intermediate_frame_load_intermediate_arrays(
-        self,
-        tags: [str],
-    ) -> Generator[np.ndarray, None, None]:
-        """Yield a generator that produces ndarrays for the requested tags."""
-        tags = list(set([CryonirspTag.intermediate(), CryonirspTag.frame()] + tags))
-
-        if self.scratch.count_all(tags=tags) == 0:
-            raise RuntimeError(f"No files found matching {tags =}")
-        for hdu in self.read(
-            tags=tags,
-            decoder=fits_hdu_decoder,
-        ):
-            yield hdu.data
-
-    def intermediate_frame_load_beam_boundaries(self, beam: int) -> np.ndarray:
-        """Return array containing beam boundaries for a given beam."""
-        tags = [CryonirspTag.task_beam_boundaries(), CryonirspTag.beam(beam)]
-        return next(self.intermediate_frame_load_intermediate_arrays(tags=tags))
-
-    def intermediate_frame_load_full_bad_pixel_map(self) -> np.ndarray:
-        """Return the full array version of the bad-pixel map."""
-        tags = [CryonirspTag.task_bad_pixel_map()]
-        bad_pixel_map = next(self.intermediate_frame_load_intermediate_arrays(tags=tags))
-        return bad_pixel_map
-
-    def intermediate_frame_load_bad_pixel_map(self, beam: int) -> np.ndarray:
-        """Return bad-pixel map for given beam."""
-        bad_pixel_map = self.intermediate_frame_load_full_bad_pixel_map()
-        return self.beam_access_get_beam(bad_pixel_map, beam)
-
-    def intermediate_frame_load_dark_array(
-        self,
-        exposure_conditions: ExposureConditions,
-        beam: int,
-    ) -> np.ndarray:
-        """Load an existing dark array."""
-        tags = [
-            CryonirspTag.task_dark(),
-            CryonirspTag.exposure_conditions(exposure_conditions),
-            CryonirspTag.beam(beam),
-        ]
-        return next(self.intermediate_frame_load_intermediate_arrays(tags=tags))
-
-    def intermediate_frame_load_polcal_dark_array(
-        self,
-        beam: int,
-        exposure_conditions: ExposureConditions,
-    ) -> np.ndarray:
-        """Load an existing polcal dark array."""
-        tags = [
-            CryonirspTag.task_polcal_dark(),
-            CryonirspTag.exposure_conditions(exposure_conditions),
-            CryonirspTag.beam(beam),
-        ]
-        return next(self.intermediate_frame_load_intermediate_arrays(tags=tags))
-
-    def intermediate_frame_load_solar_gain_array(self, beam: int) -> np.ndarray:
-        """Load an existing solar gain array."""
-        tags = [CryonirspTag.task_solar_gain(), CryonirspTag.beam(beam)]
-        return next(self.intermediate_frame_load_intermediate_arrays(tags=tags))
-
-    def intermediate_frame_load_lamp_gain_array(self, beam: int) -> np.ndarray:
-        """Load an existing lamp gain array."""
-        tags = [CryonirspTag.task_lamp_gain(), CryonirspTag.beam(beam)]
-        return next(self.intermediate_frame_load_intermediate_arrays(tags=tags))
-
-    def intermediate_frame_load_polcal_gain_array(
-        self,
-        beam: int,
-        exposure_conditions: ExposureConditions,
-    ) -> np.ndarray:
-        """Load an existing polcal gain array."""
-        tags = [
-            CryonirspTag.task_polcal_gain(),
-            CryonirspTag.exposure_conditions(exposure_conditions),
-            CryonirspTag.beam(beam),
-        ]
-        return next(self.intermediate_frame_load_intermediate_arrays(tags=tags))
-
-    def intermediate_frame_load_demod_matrices(
-        self,
-        beam: int,
-    ) -> np.ndarray:
-        """Load existing demodulation matrices."""
-        tags = [CryonirspTag.task_demodulation_matrices(), CryonirspTag.beam(beam)]
-        return next(self.intermediate_frame_load_intermediate_arrays(tags=tags))
-
-    def intermediate_frame_load_angle(self, beam: int) -> float:
-        """Return the geometric correction angle for a given beam."""
-        tags = [CryonirspTag.task_geometric_angle(), CryonirspTag.beam(beam)]
-        angle_array = next(self.intermediate_frame_load_intermediate_arrays(tags=tags))
-        return float(angle_array[0])
-
-    def intermediate_frame_load_state_offset(
-        self,
-        beam: int,
-    ) -> np.ndarray:
-        """Return state offset shifts for a given beam."""
-        tags = [CryonirspTag.task_geometric_offset(), CryonirspTag.beam(beam)]
-        offset = next(self.intermediate_frame_load_intermediate_arrays(tags=tags))
-        return offset
-
-    def intermediate_frame_load_spec_shift(self, beam: int) -> np.ndarray:
-        """Return spectral shifts for a given beam."""
-        tags = [CryonirspTag.task_geometric_sepectral_shifts(), CryonirspTag.beam(beam)]
-        shifts = next(self.intermediate_frame_load_intermediate_arrays(tags=tags))
-        return shifts
-
-    def intermediate_frame_write_arrays(
-        self,
-        arrays: Iterable[np.ndarray] | np.ndarray,
-        task_tag: str | None = None,
-        task: str | None = None,
-        headers: Iterable[fits.Header] | fits.Header | None = None,
-        beam: int | None = None,
-        modstate: int | None = None,
-        map_scan: int | None = None,
-        scan_step: int | None = None,
-        exposure_time: float | None = None,
-        exposure_conditions: ExposureConditions | None = None,
-        meas_num: int | None = None,
-        cs_step: int | None = None,
-    ) -> None:
-        """Write an intermediate fits files given a list of input arrays and headers."""
-        if task_tag is not None and task is not None:
-            raise ValueError("Cannot specify both the raw 'task' and a formatted 'task_tag'.")
-        if task_tag is None and task is None:
-            raise ValueError("Must specify exactly one of raw 'task' or formatted 'task_tag'.")
-
-        if task is not None:
-            task_tag = CryonirspTag.task(task)
-        tags = [CryonirspTag.intermediate(), CryonirspTag.frame(), task_tag]
-        for arg, tag_func in zip(
-            [
-                beam,
-                modstate,
-                map_scan,
-                scan_step,
-                exposure_time,
-                meas_num,
-                cs_step,
-                exposure_conditions,
-            ],
-            [
-                CryonirspTag.beam,
-                CryonirspTag.modstate,
-                CryonirspTag.map_scan,
-                CryonirspTag.scan_step,
-                CryonirspTag.exposure_time,
-                CryonirspTag.meas_num,
-                CryonirspTag.cs_step,
-                CryonirspTag.exposure_conditions,
-            ],
-        ):
-            if arg is not None:
-                tags.append(tag_func(arg))
-
-        arrays = [arrays] if isinstance(arrays, np.ndarray) else arrays
-        if headers is not None:
-            headers = [headers] if isinstance(headers, fits.Header) else headers
-        else:
-            headers = [None] * len(list(arrays))
-
-        filenames = []
-        for array, header in zip(arrays, headers):
-            hdul = fits.HDUList([fits.PrimaryHDU(data=array, header=header)])
-            path = str(
-                self.write(
-                    data=hdul,
-                    tags=tags,
-                    encoder=fits_hdulist_encoder,
-                )
-            )
-            filenames.append(path)
-        logger.info(f"Wrote intermediate file(s) for {tags = } to {filenames}")

dkist_processing_cryonirsp/tasks/mixin/linearized_frame.py

@@ -1,309 +0,0 @@
-"""Helper to manage input data."""
-from collections.abc import Generator
-
-import numpy as np
-from astropy.io import fits
-from dkist_processing_common.codecs.fits import fits_access_decoder
-from dkist_processing_common.models.fits_access import FitsAccessBase
-from dkist_processing_common.models.task_name import TaskName
-
-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
-
-
-class LinearizedFrameMixin:
-    """
-    Mixin for methods that support easy loading of linearity corrected frames.
-
-    Throughout this class wse assume that if beam is None, then it means to return the entire array.
-    """
-
-    def linearized_frame_fits_access_generator(
-        self,
-        tags: [str],
-        beam: int,
-    ) -> Generator[CryonirspL0FitsAccess, None, None]:
-        """Load linearized fits frames based on given tags and extract the desired beam."""
-        tags = list(set([CryonirspTag.linearized(), CryonirspTag.frame()] + tags))
-
-        full_frame_generator = self.read(
-            tags=tags,
-            decoder=fits_access_decoder,
-            fits_access_class=CryonirspL0FitsAccess,
-        )
-        for item in full_frame_generator:
-            header = item.header
-            flipped_data = self.linearized_frame_flip_dispersion_axis(item)
-            data = self.beam_access_get_beam(array=flipped_data, beam=beam)
-            hdu = fits.PrimaryHDU(header=header, data=data)
-            yield CryonirspL0FitsAccess(hdu=hdu)
-
-    def linearized_frame_full_array_generator(
-        self,
-        tags: [str],
-    ) -> Generator[np.ndarray, None, None]:
-        """Load linearized fits frames based on given tags."""
-        tags = list(set([CryonirspTag.linearized(), CryonirspTag.frame()] + tags))
-        frame_generator = self.read(
-            tags=tags,
-            decoder=fits_access_decoder,
-            fits_access_class=CryonirspL0FitsAccess,
-        )
-        for frame in frame_generator:
-            yield self.linearized_frame_flip_dispersion_axis(frame)
-
-    def linearized_frame_dark_array_generator(
-        self, exposure_conditions: ExposureConditions, beam: int
-    ) -> Generator[np.ndarray, None, None]:
-        """
-        Return a fits access generator with dark frames for the given parameters.
-
-        Parameters
-        ----------
-        exposure_conditions
-            The current exposure conditions
-        modstate
-            The current modulator state
-        beam
-            The current beam number
-
-        Returns
-        -------
-        A fits access generator based on the inputs provided
-        """
-        tags = [CryonirspTag.task_dark(), CryonirspTag.exposure_conditions(exposure_conditions)]
-        dark_array_fits_access = self.linearized_frame_fits_access_generator(tags=tags, beam=beam)
-        for array in dark_array_fits_access:
-            yield array.data
-
-    def linearized_frame_polcal_dark_array_generator(
-        self,
-        exposure_conditions: ExposureConditions,
-        beam: int,
-    ) -> Generator[np.ndarray, None, None]:
-        """
-        Return a fits access generator with polcal dark frames based on the given parameters.
-
-        Parameters
-        ----------
-        beam
-            The beam from which to return data
-        exposure_conditions
-            The current exposure conditions
-
-        Returns
-        -------
-        A fits access generator based on the inputs provided
-        """
-        tags = [
-            CryonirspTag.task_polcal_dark(),
-            CryonirspTag.exposure_conditions(exposure_conditions),
-        ]
-        polcal_dark_array_fits_access = self.linearized_frame_fits_access_generator(
-            tags=tags,
-            beam=beam,
-        )
-        for array in polcal_dark_array_fits_access:
-            yield array.data
-
-    def linearized_frame_gain_array_generator(
-        self,
-        gain_type: str,
-        beam: int,
-        exposure_conditions: ExposureConditions | None = None,
-    ) -> Generator[np.ndarray, None, None]:
-        """
-        Return a fits access generator with gain frames based on the parameters provided.
-
-        Parameters
-        ----------
-        gain_type
-            The gain type
-        exposure_conditions
-            The current exposure conditions
-        beam
-            The current beam number
-
-        Returns
-        -------
-        A fits access generator based on the inputs provided
-        """
-        tags = [CryonirspTag.task(gain_type)]
-        if exposure_conditions is not None:
-            tags.append(CryonirspTag.exposure_conditions(exposure_conditions))
-        gain_array_fits_access = self.linearized_frame_fits_access_generator(
-            tags=tags,
-            beam=beam,
-        )
-        for array in gain_array_fits_access:
-            yield array.data
-
-    def linearized_frame_lamp_gain_array_generator(
-        self,
-        beam: int,
-        exposure_conditions: float | None = None,
-    ) -> Generator[np.ndarray, None, None]:
-        """
-        Return a fits access generator with lamp gain frames based on the parameters provided.
-
-        Parameters
-        ----------
-        exposure_conditions
-            The current exposure conditions
-        beam
-            The current beam number
-
-        Returns
-        -------
-        A fits access generator based on the inputs provided
-        """
-        return self.linearized_frame_gain_array_generator(
-            gain_type=TaskName.lamp_gain.value, beam=beam, exposure_conditions=exposure_conditions
-        )
-
-    def linearized_frame_solar_gain_array_generator(
-        self,
-        beam: int,
-        exposure_conditions: float | None = None,
-    ) -> Generator[np.ndarray, None, None]:
-        """
-        Return a fits access generator with solar gain frames based on the parameters provided.
-
-        Parameters
-        ----------
-        exposure_conditions
-            The current exposure conditions
-        beam
-            The current beam number
-
-        Returns
-        -------
-        A fits access generator based on the inputs provided
-        """
-        return self.linearized_frame_gain_array_generator(
-            gain_type=TaskName.solar_gain.value, beam=beam, exposure_conditions=exposure_conditions
-        )
-
-    def linearized_frame_polcal_gain_array_generator(
-        self,
-        exposure_conditions: ExposureConditions,
-        beam: int,
-    ) -> Generator[np.ndarray, None, None]:
-        """
-        Return a fits access generator with polcal gain frames based on the given parameters.
-
-        Parameters
-        ----------
-        beam
-            The beam from which to return data
-        exposure_conditions
-            The current exposure conditions
-
-        Returns
-        -------
-        A fits access generator based on the inputs provided
-        """
-        tags = [
-            CryonirspTag.task_polcal_gain(),
-            CryonirspTag.exposure_conditions(exposure_conditions),
-        ]
-        polcal_gain_array_fits_access = self.linearized_frame_fits_access_generator(
-            tags=tags,
-            beam=beam,
-        )
-        for array in polcal_gain_array_fits_access:
-            yield array.data
-
-    def linearized_frame_observe_fits_access_generator(
-        self,
-        modstate: int,
-        meas_num: int,
-        scan_step: int,
-        map_scan: int,
-        exposure_conditions: ExposureConditions,
-        beam: int,
-    ) -> Generator[FitsAccessBase, None, None]:
-        """
-        Return a fits access generator of observe frames based on the parameters provided.
-
-        Parameters
-        ----------
-        beam
-            The beam from which to return data
-        modstate
-            The current modulator state
-        meas_num
-            The current measurement number
-        scan_step
-            The current scan step
-        map_scan
-            The current map_scan number
-        exposure_conditions
-            The current exposure conditions
-
-        Returns
-        -------
-        A fits access generator based on the inputs provided
-        """
-        tags = [
-            CryonirspTag.task_observe(),
-            CryonirspTag.modstate(modstate),
-            CryonirspTag.meas_num(meas_num),
-            CryonirspTag.scan_step(scan_step),
-            CryonirspTag.map_scan(map_scan),
-            CryonirspTag.exposure_conditions(exposure_conditions),
-        ]
-        return self.linearized_frame_fits_access_generator(
-            tags=tags,
-            beam=beam,
-        )
-
-    def linearized_frame_polcal_fits_access_generator(
-        self,
-        modstate: int,
-        cs_step: int,
-        exposure_conditions: ExposureConditions,
-        beam: int,
-    ) -> Generator[FitsAccessBase, None, None]:
-        """
-        Return a fits access generator of polcal frames based on the parameters provided.
-
-        Parameters
-        ----------
-        modstate
-            The current modulator state
-        cs_step
-            The current cs_step number
-        exposure_conditions
-            The current exposure conditions
-        beam
-            The current beam number
-
-        Returns
-        -------
-        A fits access generator based on the inputs provided
-        """
-        tags = [
-            CryonirspTag.task_polcal(),
-            CryonirspTag.modstate(modstate),
-            CryonirspTag.cs_step(cs_step),
-            CryonirspTag.exposure_conditions(exposure_conditions),
-        ]
-        return self.linearized_frame_fits_access_generator(
-            tags=tags,
-            beam=beam,
-        )
-
-    @staticmethod
-    def linearized_frame_flip_dispersion_axis(frame: CryonirspL0FitsAccess) -> np.ndarray:
-        """
-        Flip dispersion axis. Cryo's wavelength decreases from left to right, so we flip it here to match the other instruments.
-
-        :param frame: frame to be flipped
-        :return: flipped array (SP), array (CI)
-        """
-        if frame.arm_id == "SP":
-            return np.flip(frame.data, 1)
-        elif frame.arm_id == "CI":
-            return frame.data