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

dkist_processing_cryonirsp/tasks/make_movie_frames.py

@@ -0,0 +1,302 @@
+"""Cryonirsp make movie frames task."""
+from abc import ABC
+from abc import abstractmethod
+
+import numpy as np
+from astropy.io import fits
+from astropy.visualization import ZScaleInterval
+from dkist_processing_common.codecs.fits import fits_access_decoder
+from dkist_processing_common.codecs.fits import fits_hdulist_encoder
+from dkist_service_configuration.logging import logger
+
+from dkist_processing_cryonirsp.models.tags import CryonirspTag
+from dkist_processing_cryonirsp.parsers.cryonirsp_l1_fits_access import CryonirspL1FitsAccess
+from dkist_processing_cryonirsp.tasks.cryonirsp_base import CryonirspTaskBase
+
+__all__ = ["MakeCryonirspMovieFrames", "SPMakeCryonirspMovieFrames"]
+
+
+class MakeCryonirspMovieFramesBase(CryonirspTaskBase, ABC):
+    """Create CryoNIRSP movie frames common functionality for the Context Imager (CI) and Spectropolarimeter (SP)."""
+
+    def run(self):
+        """Create movie frames using all stokes states if they exist, otherwise only use intensity."""
+        if self.constants.correct_for_polarization:
+            with self.apm_task_step("Make full stokes movie"):
+                self.make_full_stokes_movie_frames()
+        else:
+            with self.apm_task_step("Make intensity only movie"):
+                self.make_intensity_movie_frames()
+
+    @abstractmethod
+    def make_full_stokes_movie_frames(self):
+        """Make a movie that combines each of the stokes frames into a single movie frame."""
+
+    @abstractmethod
+    def make_intensity_movie_frames(self):
+        """Make a movie out of stokes I frames."""
+
+    @staticmethod
+    def scale_for_rendering(data: np.ndarray):
+        """Scale the calibrated frame data using a normalization function to facilitate display as a movie frame."""
+        zscale = ZScaleInterval()
+        return zscale(data)
+
+    @staticmethod
+    def grid_movie_frame(
+        top_left: np.ndarray,
+        top_right: np.ndarray,
+        bottom_left: np.ndarray,
+        bottom_right: np.ndarray,
+    ) -> np.ndarray:
+        """Combine multiple arrays into a 2x2 grid."""
+        result = np.concatenate(
+            (
+                np.concatenate((top_left, top_right), axis=1),
+                np.concatenate((bottom_left, bottom_right), axis=1),
+            ),
+            axis=0,
+        )
+        return result
+
+    def get_movie_header(
+        self, map_scan: int, scan_step: int, meas_num: int, stokes_state: str
+    ) -> fits.Header:
+        """Create a header to use on a movie frame based on a calibrated frame."""
+        calibrated_frame = next(
+            self.read(
+                tags=[
+                    CryonirspTag.frame(),
+                    CryonirspTag.calibrated(),
+                    CryonirspTag.stokes(stokes_state),
+                    CryonirspTag.meas_num(meas_num),
+                    CryonirspTag.map_scan(map_scan),
+                    CryonirspTag.scan_step(scan_step),
+                ],
+                decoder=fits_access_decoder,
+                fits_access_class=CryonirspL1FitsAccess,
+            )
+        )
+        header = calibrated_frame.header
+        return header
+
+
+# See note below on `SPMakeCryonirspMovieFrames`
+class MakeCryonirspMovieFrames(MakeCryonirspMovieFramesBase):
+    """Make CryoNIRSP movie frames for the Context Imager and tag with CryonirspTag.movie_frame()."""
+
+    def make_full_stokes_movie_frames(self):
+        """Make a movie that combines each of the stokes frames into a single movie frame."""
+        self.make_intensity_movie_frames()
+
+    def make_intensity_movie_frames(self):
+        """Make a movie out of stokes I frames."""
+        for map_scan in range(1, self.constants.num_map_scans + 1):
+            for scan_step in range(1, self.constants.num_scan_steps + 1):
+                logger.info(f"Generate Stokes-I movie frame for {map_scan=} and {scan_step=}")
+                movie_frame_hdu = self.make_intensity_frame(map_scan=map_scan, scan_step=scan_step)
+                logger.info(f"Writing Stokes-I movie frame for {map_scan=} and {scan_step=}")
+                self.write(
+                    data=fits.HDUList([movie_frame_hdu]),
+                    tags=[
+                        CryonirspTag.map_scan(map_scan),
+                        CryonirspTag.scan_step(scan_step),
+                        CryonirspTag.movie_frame(),
+                    ],
+                    encoder=fits_hdulist_encoder,
+                )
+
+    def make_full_stokes_frame(self, map_scan: int, scan_step: int) -> fits.PrimaryHDU:
+        """Create a movie frame (data + header) with the stokes frames (IQUV) combined into top left, top right, bottom left, bottom right quadrants respectively."""
+        meas_num = 1  # Use only the first measurement if there are multiple measurements.
+        stokes_i = self.get_movie_frame(
+            map_scan=map_scan, scan_step=scan_step, meas_num=meas_num, stokes_state="I"
+        )
+        stokes_q = self.get_movie_frame(
+            map_scan=map_scan, scan_step=scan_step, meas_num=meas_num, stokes_state="Q"
+        )
+        stokes_u = self.get_movie_frame(
+            map_scan=map_scan, scan_step=scan_step, meas_num=meas_num, stokes_state="U"
+        )
+        stokes_v = self.get_movie_frame(
+            map_scan=map_scan, scan_step=scan_step, meas_num=meas_num, stokes_state="V"
+        )
+        movie_frame = self.grid_movie_frame(
+            top_left=stokes_i,
+            top_right=stokes_q,
+            bottom_left=stokes_u,
+            bottom_right=stokes_v,
+        )
+        movie_header = self.get_movie_header(
+            map_scan=map_scan, scan_step=scan_step, meas_num=meas_num, stokes_state="I"
+        )
+        return fits.PrimaryHDU(header=movie_header, data=movie_frame)
+
+    def make_intensity_frame(self, map_scan: int, scan_step: int) -> fits.PrimaryHDU:
+        """Create a movie frame (data + header) with just the stokes I frames."""
+        meas_num = 1  # Use only the first measurement if there are multiple measurements.
+        stokes_i = self.get_movie_frame(
+            map_scan=map_scan, scan_step=scan_step, meas_num=meas_num, stokes_state="I"
+        )
+        movie_header = self.get_movie_header(
+            map_scan=map_scan, scan_step=scan_step, meas_num=meas_num, stokes_state="I"
+        )
+        return fits.PrimaryHDU(header=movie_header, data=stokes_i)
+
+    def get_movie_frame(
+        self, map_scan: int, scan_step: int, meas_num: int, stokes_state: str
+    ) -> np.ndarray:
+        """Retrieve the calibrated frame data for the first frame which matches the input parameters and transform it into a movie frame (i.e. normalize the values)."""
+        calibrated_frame = next(
+            self.read(
+                tags=[
+                    CryonirspTag.frame(),
+                    CryonirspTag.calibrated(),
+                    CryonirspTag.stokes(stokes_state),
+                    CryonirspTag.meas_num(meas_num),
+                    CryonirspTag.map_scan(map_scan),
+                    CryonirspTag.scan_step(scan_step),
+                ],
+                decoder=fits_access_decoder,
+                fits_access_class=CryonirspL1FitsAccess,
+            )
+        )
+        movie_frame = self.scale_for_rendering(calibrated_frame.data)
+        return movie_frame
+
+
+# NOTE:
+# This task isn't used right now because the SP movies need some better handling of the wavelength dimension.
+# For the time being both arms use the `MakeCryonirspMovieFrames` task above.
+# See PR #91 for more information.
+class SPMakeCryonirspMovieFrames(MakeCryonirspMovieFramesBase):
+    """Make CryoNIRSP movie frames for the Spectropolarimeter and tag with CryonirspTag.movie_frame()."""
+
+    def make_full_stokes_movie_frames(self):
+        """Make a movie that combines each of the stokes frames into a single movie frame."""
+        for map_scan in range(1, self.constants.num_map_scans + 1):
+            logger.info(f"Generate full stokes movie frame for {map_scan=}")
+            movie_frame_hdu = self.make_full_stokes_frame(map_scan=map_scan)
+
+            logger.info(f"Writing full stokes movie frame for {map_scan=}")
+            self.write(
+                data=fits.HDUList([movie_frame_hdu]),
+                tags=[
+                    CryonirspTag.map_scan(map_scan),
+                    CryonirspTag.movie_frame(),
+                ],
+                encoder=fits_hdulist_encoder,
+            )
+
+    def make_intensity_movie_frames(self):
+        """Make a movie out of stokes I frames."""
+        for map_scan in range(1, self.constants.num_map_scans + 1):
+            logger.info(f"Generate intensity movie frame for {map_scan=}")
+            movie_frame_hdu = self.make_intensity_frame(map_scan=map_scan)
+
+            logger.info(f"Writing intensity movie frame for {map_scan=}")
+            self.write(
+                data=fits.HDUList([movie_frame_hdu]),
+                tags=[
+                    CryonirspTag.map_scan(map_scan),
+                    CryonirspTag.movie_frame(),
+                ],
+                encoder=fits_hdulist_encoder,
+            )
+
+    def make_full_stokes_frame(self, map_scan: int) -> fits.PrimaryHDU:
+        """Create a movie frame (data + header) with the stokes frames (IQUV) combined into top left, top right, bottom left, bottom right quadrants respectively."""
+        meas_num = 1  # Use only the first measurement if there are multiple measurements.
+        stokes_i = self.get_movie_frame(map_scan=map_scan, meas_num=meas_num, stokes_state="I")
+        stokes_q = self.get_movie_frame(map_scan=map_scan, meas_num=meas_num, stokes_state="Q")
+        stokes_u = self.get_movie_frame(map_scan=map_scan, meas_num=meas_num, stokes_state="U")
+        stokes_v = self.get_movie_frame(map_scan=map_scan, meas_num=meas_num, stokes_state="V")
+        movie_frame = self.grid_movie_frame(
+            top_left=stokes_i,
+            top_right=stokes_q,
+            bottom_left=stokes_u,
+            bottom_right=stokes_v,
+        )
+        # Here we make the full map have the same header as the first scan step.
+        movie_header = self.get_movie_header(
+            map_scan=map_scan, scan_step=1, meas_num=meas_num, stokes_state="I"
+        )
+        return fits.PrimaryHDU(header=movie_header, data=movie_frame)
+
+    def make_intensity_frame(self, map_scan: int) -> fits.PrimaryHDU:
+        """Create a movie frame (data + header) with just the stokes I frames."""
+        meas_num = 1  # Use only the first measurement if there are multiple measurements.
+        stokes_i = self.get_movie_frame(map_scan=map_scan, meas_num=meas_num, stokes_state="I")
+        movie_header = self.get_movie_header(
+            map_scan=map_scan, scan_step=1, meas_num=meas_num, stokes_state="I"
+        )
+        return fits.PrimaryHDU(header=movie_header, data=stokes_i)
+
+    def get_movie_frame(self, map_scan: int, meas_num: int, stokes_state: str) -> np.ndarray:
+        """Retrieve the calibrated frame data for the first frame which matches the input parameters and transform it into a movie frame (i.e. normalize the values)."""
+        if self.constants.num_scan_steps == 1:
+            calibrated_frame = self.get_spectral_calibrated_frame(
+                map_scan=map_scan, meas_num=meas_num, stokes_state=stokes_state
+            )
+        else:
+            calibrated_frame = self.get_integrated_calibrated_frame(
+                map_scan=map_scan, meas_num=meas_num, stokes_state=stokes_state
+            )
+        movie_frame = self.scale_for_rendering(calibrated_frame.data)
+        return movie_frame
+
+    def get_spectral_calibrated_frame(
+        self, map_scan: int, meas_num: int, stokes_state: str
+    ) -> np.ndarray:
+        """Retrieve a calibrated frame for a single scan step (no integration)."""
+        scan_step = 1  # There is only a single scan step in a spectral movie
+        calibrated_frame = next(
+            self.read(
+                tags=[
+                    CryonirspTag.frame(),
+                    CryonirspTag.calibrated(),
+                    CryonirspTag.stokes(stokes_state),
+                    CryonirspTag.meas_num(meas_num),
+                    CryonirspTag.map_scan(map_scan),
+                    CryonirspTag.scan_step(scan_step),
+                ],
+                decoder=fits_access_decoder,
+                fits_access_class=CryonirspL1FitsAccess,
+            )
+        )
+        return calibrated_frame.data
+
+    def get_integrated_calibrated_frame(
+        self, map_scan: int, meas_num: int, stokes_state: str
+    ) -> np.ndarray:
+        """Retrieve a frame that has been integrated across scan steps."""
+        data = self.get_spectral_calibrated_frame(
+            map_scan=map_scan, meas_num=meas_num, stokes_state=stokes_state
+        )
+        integrated_data = self.integrate_spectral_frame(data)
+        integrated_arrays = [integrated_data]
+        for scan_step in range(2, self.constants.num_scan_steps + 1):
+            calibrated_frame = next(
+                self.read(
+                    tags=[
+                        CryonirspTag.frame(),
+                        CryonirspTag.calibrated(),
+                        CryonirspTag.stokes(stokes_state),
+                        CryonirspTag.meas_num(meas_num),
+                        CryonirspTag.map_scan(map_scan),
+                        CryonirspTag.scan_step(scan_step),
+                    ],
+                    decoder=fits_access_decoder,
+                    fits_access_class=CryonirspL1FitsAccess,
+                )
+            )
+            integrated_data = self.integrate_spectral_frame(calibrated_frame.data)
+            integrated_arrays.append(integrated_data)
+        full_frame = np.vstack(integrated_arrays)
+        return full_frame
+
+    @staticmethod
+    def integrate_spectral_frame(data: np.ndarray) -> np.ndarray:
+        """Integrate spectral frame."""
+        wavelength_integrated_data = np.sum(np.abs(data), axis=1)
+        return wavelength_integrated_data

dkist_processing_cryonirsp/tasks/mixin/__init__.py

@@ -0,0 +1 @@
+"""Init."""

dkist_processing_cryonirsp/tasks/mixin/beam_access.py

@@ -0,0 +1,52 @@
+"""Mixin to support extracting the desired beam from an input image on-the-fly."""
+from functools import cached_property
+
+import numpy as np
+
+
+class BeamAccessMixin:
+    """Mixin that supports extracting the desired beam from an input image on-the-fly."""
+
+    def beam_access_get_beam(self, array: np.ndarray, beam: int) -> np.ndarray:
+        """
+        Extract a single beam array from a dual-beam array.
+
+        Parameters
+        ----------
+        array
+            The input dual-beam array
+        beam
+            The desired beam to extract
+
+        Returns
+        -------
+        An ndarray containing the extracted beam
+        """
+        boundaries = self.beam_boundaries[beam]
+        spatial_min, spatial_max, spectral_min, spectral_max = boundaries
+
+        if (
+            spatial_min < 0
+            or spatial_max > array.shape[0]
+            or spectral_min < 0
+            or spectral_max > array.shape[1]
+        ):
+            raise IndexError(
+                f"beam_access_get_boundaries exceed array bounds: {boundaries = }, {array.shape = }."
+            )
+
+        return np.copy(array[spatial_min:spatial_max, spectral_min:spectral_max])
+
+    @cached_property
+    def beam_boundaries(self) -> dict[int, np.ndarray]:
+        """
+        Load the beam boundaries from their respective files and return as a boundary dict.
+
+        Returns
+        -------
+        beam_boundary dict
+        """
+        boundaries = dict()
+        for beam in range(1, self.constants.num_beams + 1):
+            boundaries[beam] = self.intermediate_frame_load_beam_boundaries(beam=beam)
+        return boundaries

dkist_processing_cryonirsp/tasks/mixin/corrections.py

@@ -0,0 +1,177 @@
+"""Helper for CRYO-Nirsp array corrections."""
+from collections.abc import Generator
+from collections.abc import Iterable
+
+import numpy as np
+import scipy.ndimage as spnd
+from dkist_processing_math.transform import affine_transform_arrays
+from dkist_processing_math.transform import rotate_arrays_about_point
+
+
+class CorrectionsMixin:
+    """Mixin to provide support for various array corrections used by the workflow tasks."""
+
+    @staticmethod
+    def corrections_correct_geometry(
+        arrays: Iterable[np.ndarray] | np.ndarray,
+        shift: np.ndarray = np.zeros(2),
+        angle: float = 0.0,
+    ) -> Generator[np.ndarray, None, None]:
+        """
+        Shift and then rotate data.
+
+        It applies the inverse of the given shift and angle.
+
+        Parameters
+        ----------
+        arrays
+            2D array(s) containing the data for the un-shifted beam
+
+        shift
+            The measured shift offset in the spectral dimension
+            between beams or between modulator states in a single beam.
+
+        angle
+            The angle between the slit and pixel axes.
+
+        Returns
+        -------
+        Generator
+            2D array(s) containing the data of the rotated and shifted beam
+        """
+        arrays = [arrays] if isinstance(arrays, np.ndarray) else arrays
+        for array in arrays:
+            # Need to recast type to fix endianness issue caused by `fits`
+            array = array.astype(np.float64)
+            array[np.where(array == np.inf)] = np.max(array[np.isfinite(array)])
+            array[np.where(array == -np.inf)] = np.min(array[np.isfinite(array)])
+            array[np.isnan(array)] = np.nanmedian(array)
+            translated = affine_transform_arrays(array, translation=-shift, mode="edge", order=1)
+            # rotate_arrays_about_point rotates the wrong way, so no negative sign here
+            yield next(rotate_arrays_about_point(translated, angle=angle, mode="edge", order=1))
+
+    @staticmethod
+    def corrections_distort_geometry(
+        arrays: Iterable[np.ndarray] | np.ndarray,
+        shift: np.ndarray = np.zeros(2),
+        angle: float = 0.0,
+    ) -> Generator[np.ndarray, None, None]:
+        """
+        Rotate and then shift data.
+
+        It applies the inverse of the given shift and angle.
+
+        Parameters
+        ----------
+        arrays
+            2D array(s) containing the data for the un-shifted beam
+
+        shift
+            The measured shift offset in the spectral dimension
+            between beams or between modulator states in a single beam.
+
+        angle
+            The angle between the slit and pixel axes.
+
+        Returns
+        -------
+        Generator
+            2D array(s) containing the data of the rotated and shifted beam
+        """
+        arrays = [arrays] if isinstance(arrays, np.ndarray) else arrays
+        for array in arrays:
+            # Need to recast type to fix endianness issue caused by `fits`
+            array = array.astype(np.float64)
+            array[np.where(array == np.inf)] = np.max(array[np.isfinite(array)])
+            array[np.where(array == -np.inf)] = np.min(array[np.isfinite(array)])
+            array[np.isnan(array)] = np.nanmedian(array)
+            # rotate_arrays_about_point rotates the wrong way, so no negative sign here
+            rotated = rotate_arrays_about_point(array, angle=angle, mode="edge")
+            yield next(affine_transform_arrays(rotated, translation=-shift, mode="edge"))
+
+    @staticmethod
+    def corrections_remove_spec_shifts(
+        arrays: Iterable[np.ndarray] | np.ndarray, spec_shift: np.ndarray
+    ) -> Generator[np.ndarray, None, None]:
+        """
+        Remove spectral curvature.
+
+        This is a pretty simple function that simply undoes the computed spectral shifts.
+
+        Parameters
+        ----------
+        arrays
+            2D array(s) containing the data for the un-distorted beam
+
+        spec_shift : np.ndarray
+            Array with shape (X), where X is the number of pixels in the spatial dimension.
+            This dimension gives the spectral shift.
+
+        Returns
+        -------
+        Generator
+            2D array(s) containing the data of the corrected beam
+
+        """
+        arrays = [arrays] if isinstance(arrays, np.ndarray) else arrays
+        for array in arrays:
+            numx = array.shape[0]
+            array_output = np.zeros(array.shape)
+            for j in range(numx):
+                array_output[j, :] = spnd.shift(array[j, :], -spec_shift[j], mode="nearest")
+            yield array_output
+
+    def corrections_correct_bad_pixels(
+        self,
+        array_to_fix: np.ndarray,
+        bad_pixel_map: np.ndarray,
+    ) -> np.ndarray:
+        """
+        Correct bad pixels in an array using a median filter.
+
+        Corrects only the bad pixels identified by the bad pixel map.
+
+        If enough pixels are found to suggest readout errors (> a given fraction of the array),
+        a global replacement value is used instead.
+
+        Parameters
+        ----------
+        array_to_fix
+            The array to be corrected
+
+        bad_pixel_map
+            The bad_pixel_map to use for the correction
+            An array of zeros with bad pixel locations set to 1
+
+        Returns
+        -------
+        ndarray
+            The corrected array
+        """
+        bad_pixel_fraction = np.sum(bad_pixel_map) / bad_pixel_map.size
+        if bad_pixel_fraction > self.parameters.corrections_bad_pixel_fraction_threshold:
+            global_median = np.nanmedian(array_to_fix)
+            array_to_fix[bad_pixel_map == 1] = global_median
+            return array_to_fix
+
+        kernel_size = self.parameters.corrections_bad_pixel_median_filter_size
+        half_kernel_size = kernel_size // 2
+
+        masked_array = np.ma.array(array_to_fix, mask=bad_pixel_map)
+
+        y_locs, x_locs = np.nonzero(bad_pixel_map)
+        num_y = bad_pixel_map.shape[0]
+        num_x = bad_pixel_map.shape[1]
+        for x, y in zip(x_locs, y_locs):
+            y_slice = slice(max(y - half_kernel_size, 0), min(y + half_kernel_size + 1, num_y))
+            x_slice = slice(max(x - half_kernel_size, 0), min(x + half_kernel_size + 1, num_x))
+
+            if self.constants.arm_id == "SP":
+                # Only compute median in spatial dimension so we don't blur things spectrally
+                med = np.ma.median(masked_array[y_slice, x])
+            else:
+                med = np.ma.median(masked_array[y_slice, x_slice])
+
+            masked_array[y, x] = med
+
+        return masked_array.data