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

dkist_processing_trend/tasks/visp/visp_prep_fit_data.py

@@ -0,0 +1,162 @@
+"""Task for preparing ViSP data for fitting by `dkist_processing_pac`."""
+
+import numpy as np
+import scipy.ndimage as spnd
+
+from dkist_processing_trend.models.instrument_options import VispInstrumentOptions
+from dkist_processing_trend.tasks.prepare_fit_data_base import PrepareFitDataBase
+from dkist_processing_trend.tasks.visp.visp_extract_beam import extract_visp_beam
+
+__all__ = ["VispPrepareFitData"]
+
+
+class VispPrepareFitData(PrepareFitDataBase):
+    """
+    Subclass of `~dkist_processing_trend.tasks.prepare_fit_data_base.PrepareFitDataBase` that adds functionality for working with ViSP data.
+
+    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
+    """
+
+    def extract_beam(self, array: np.ndarray, beam: int) -> np.ndarray:
+        """Extract a given beam array from a full ViSP frame."""
+        beam_border = self.parameters.visp_beam_border
+        return extract_visp_beam(array=array, beam=beam, beam_border=beam_border)
+
+    def apply_global_instrument_options(
+        self, array: np.ndarray, instrument_options: VispInstrumentOptions
+    ) -> np.ndarray:
+        """
+        Prepare global ViSP data for PAC fitting.
+
+        We first mask the slit hairlines and then simply compute a global median of the array.
+        """
+        filtered_array = self.mask_hairlines(array)
+
+        return np.nanmedian(filtered_array)[None, None]
+
+    def apply_local_instrument_options(
+        self, array: np.ndarray, instrument_options: VispInstrumentOptions
+    ) -> np.ndarray:
+        """
+        Prepare local ViSP data for PAC fitting.
+
+        Algorithm:
+
+        #. Mask the slit hairlines
+        #. Collapse the array by computing the median along the spectral dimension
+        #. Smooth the result in the spatial dimension
+        #. Bin the array in the spatial dimension using the local median to compute the value of each bin
+        """
+        filtered_array = self.mask_hairlines(array)
+
+        # Add back in a dummy spectral dimension so things stay 2D, which helps thinking about it later on.
+        spectral_binned_array = np.nanmedian(filtered_array, axis=0)[None, :]
+
+        spatially_smoothed_array = spnd.median_filter(
+            spectral_binned_array,
+            # The 1 below means we don't smooth in the spectral dimension
+            size=(1, self.parameters.visp_polcal_spatial_median_filter_width_px),
+        )
+
+        local_array = self.downsample_spatial_dimension_local_median(
+            spatially_smoothed_array, instrument_options.num_spatial_px
+        )
+
+        return local_array
+
+    def mask_hairlines(self, array: np.ndarray) -> np.ndarray:
+        # This method is copied directly from the `CorrectionsMixin` in `dkist-processing-visp`
+        """
+        Mask hairlines from an array.
+
+        The hairlines will be replaced with data from a median-filtered version of the input array.
+
+        Hairlines are identified by first subtracting a spatially smoothed copy of the array and then looking for pixels
+        that have large differences. This works because the hairlines are the only features that are sharp in the
+        spatial dimension. The identified hairlines are then slightly smoothed spatially to ensure that their
+        higher-flux wings are correctly masked.
+        """
+        filtered_array = self.median_filter_array_for_hairline_identification(array)
+        hairline_locations = self.find_hairline_pixels(
+            input_array=array, filtered_array=filtered_array
+        )
+
+        # Replace hairline pixels with data from the spatially-filtered array
+        array[hairline_locations] = filtered_array[hairline_locations]
+
+        return array
+
+    def median_filter_array_for_hairline_identification(self, array: np.ndarray) -> np.ndarray:
+        # This method is copied directly from the `CorrectionsMixin` in `dkist-processing-visp`
+        """
+        Small helper to separate out the median filter step of hairline identification.
+
+        This step has been factored out so that functions that need the filtered array for further processing can avoid
+        repeating this expensive computation.
+        """
+        # The size=(1, X) means we only smooth in the spatial dimension (1st axis)
+        filtered_array = spnd.median_filter(
+            array, size=(1, self.parameters.visp_hairline_median_spatial_smoothing_width_px)
+        )
+        return filtered_array
+
+    def find_hairline_pixels(
+        self, input_array: np.ndarray, filtered_array: np.ndarray
+    ) -> np.ndarray:
+        # This method is copied directly from the `CorrectionsMixin` in `dkist-processing-visp`
+        """
+        Find pixels that likely correspond to hairlines.
+
+        This also slightly smooths the identified pixels so that high-flux wings of the hairlines are included.
+        """
+        diff = (input_array - filtered_array) / filtered_array
+        hairline_locations = np.abs(diff) > self.parameters.visp_hairline_fraction
+
+        # Now smooth the hairline mask in the spatial dimension to capture the higher-flux wings
+        mask_array = np.zeros_like(input_array)
+        mask_array[hairline_locations] = 1.0
+        mask_array = spnd.gaussian_filter1d(
+            mask_array, self.parameters.visp_hairline_mask_spatial_smoothing_width_px, axis=1
+        )
+
+        hairline_locations = np.where(
+            mask_array
+            > mask_array.max() * self.parameters.visp_hairline_mask_gaussian_peak_cutoff_fraction
+        )
+
+        return hairline_locations
+
+    @staticmethod
+    def downsample_spatial_dimension_local_median(
+        data: np.ndarray, num_spatial_bins: int
+    ) -> np.ndarray:
+        # Copied directly from the `DownsampleMixin` of `dkist-processing-visp`
+        """Resample a stack of spectra along the spatial dimension.
+
+        This is a separate function because calling `skimage.measure.block_reduce` on the entire (large) array all at once
+        creates a huge memory strain that is not needed since we're only reducing over one of the dimensions. Instead,
+        this function does some very cool tricks with reshaping and base numpy to keep the memory footprint lower.
+        """
+        # Taken from the amazing answer in
+        #  https://stackoverflow.com/questions/44527579/whats-the-best-way-to-downsample-a-numpy-array
+        num_wave, num_spat_pos = data.shape
+
+        if (num_spat_pos / num_spatial_bins) % 1 != 0:
+            raise ValueError(
+                f"The number of spatial bins must evenly divide the spatial dimension. {num_spatial_bins} bins do not evenly divide {num_spat_pos}."
+            )
+
+        reshaped = data.reshape((num_wave, num_spatial_bins, num_spat_pos // num_spatial_bins))
+        return np.median(reshaped, axis=2)

dkist_processing_trend/tasks/visp/visp_process_demod.py

@@ -0,0 +1,236 @@
+"""Task for producing full-frame ViSP demodulation matrices from the output of PAC fitting."""
+
+from typing import Literal
+
+import numpy as np
+from dkist_processing_common.codecs.fits import fits_array_decoder
+from dkist_processing_common.codecs.fits import fits_array_encoder
+from dkist_processing_math.transform.binning import resize_arrays
+from dkist_service_configuration.logging import logger
+from sklearn.linear_model import RANSACRegressor
+from sklearn.pipeline import Pipeline
+from sklearn.pipeline import make_pipeline
+from sklearn.preprocessing import PolynomialFeatures
+from sklearn.preprocessing import RobustScaler
+
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.tasks.trend_base import TrendArmTaskBase
+from dkist_processing_trend.tasks.visp.visp_extract_beam import extract_visp_beam
+
+__all__ = ["VispProcessDemodulationMatrices"]
+
+
+class VispProcessDemodulationMatrices(TrendArmTaskBase):
+    """
+    Task class for producing full-frame demodulation ViSP matrices from the output of PAC fits.
+
+    This task mimics the functionality (and code) of the `InstrumentPolarizationCalibration <https://docs.dkist.nso.edu/projects/visp/en/stable/autoapi/dkist_processing_visp/tasks/instrument_polarization/index.html>`_
+    in the L1 ViSP pipeline; it first smooths the fit modulation matrices in the spatial dimension before upsampling to
+    the full-frame shape.
+
+    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):
+        """
+        Smooth and upsample raw best-fit demodulation matrices to match the full ViSP beam size.
+
+        First, the MODulation matrices are smoothed element-by-element along the spatial dimension by fitting with a polynomial.
+        This also has the effect of upsampling the spatial dimension to the full-beam size.
+        Finally, the wavelength dimension is upsampled to the full-beam size via interpolation.
+        """
+        for beam in [1, 2]:
+            full_beam_shape = self.get_full_beam_shape(beam)
+            logger.info(f"{full_beam_shape = }")
+            for fit_options in self.parameters.fit_options_list:
+                for instrument_options in self.parameters.instrument_processing_options:
+                    log_str = f"fit options '{fit_options.name}', inst options '{instrument_options.name}', and {beam = }"
+                    base_tags = [
+                        TrendTag.intermediate(),
+                        TrendTag.pac_fit_options(fit_options.name),
+                        TrendTag.instrument_processing_options(instrument_options.name),
+                        TrendTag.beam(beam),
+                        TrendTag.arm_id(self.arm_id),
+                    ]
+                    with self.telemetry_span(f"Resampling demodulation matrices for {log_str}"):
+                        raw_demod_matrices = next(
+                            self.read(
+                                tags=base_tags + [TrendTag.task_best_fit_demodulation_matrices()],
+                                decoder=fits_array_decoder,
+                            )
+                        )
+
+                        logger.info(f"Smoothing demodulation matrices for {log_str}")
+                        smoothed_demod = self.smooth_demod_matrices(
+                            demod_matrices=raw_demod_matrices,
+                            fit_order=instrument_options.spatial_smoothing_fit_order,
+                            num_full_slit_pos=full_beam_shape[1],
+                        )
+
+                        # Reshaping the demodulation matrix to get rid of unit length dimensions
+                        logger.info(f"Resampling demodulation matrices for {log_str}")
+                        final_demod_matrices = self.reshape_demod_matrices(
+                            smoothed_demod, full_beam_shape
+                        )
+                        logger.info(
+                            f"Shape of resampled demodulation matrices: {final_demod_matrices.shape}"
+                        )
+
+                    with self.telemetry_span(f"Writing demodulation matrices for {log_str}"):
+                        self.write(
+                            data=final_demod_matrices,
+                            tags=base_tags + [TrendTag.task_processed_demodulation_matrices()],
+                            encoder=fits_array_encoder,
+                        )
+
+    def get_full_beam_shape(self, beam: Literal[1, 2]) -> tuple[int, int]:
+        """Read an INTERMEDIATE gain frame to measure the shape of a single beam array."""
+        gain_array = next(
+            self.read(
+                tags=[
+                    TrendTag.intermediate(),
+                    TrendTag.frame(),
+                    TrendTag.arm_id(self.arm_id),
+                    TrendTag.task_gain(),
+                ],
+                decoder=fits_array_decoder,
+            )
+        )
+        beam_array = extract_visp_beam(
+            gain_array, beam=beam, beam_border=self.parameters.visp_beam_border
+        )
+        return beam_array.shape
+
+    def smooth_demod_matrices(
+        self, demod_matrices: np.ndarray, fit_order: int, num_full_slit_pos: int
+    ) -> np.ndarray:
+        """
+        Smooth demodulation matrices in the spatial dimension.
+
+        The output will fully sample the spatial dimension so, as a side effect, this function also up-samples any data
+        that were binned spatially.
+
+        Smoothing is done using a RANSAC regression estimator to perform a polynomial fit with high resilience to outliers.
+
+        This method is taken directly from the ViSP L1 pipeline.
+        """
+        # We need to smooth the *modulation* (not DEmodulation) matrices to preserve the normalization of the Stokes-I
+        # values. linalg.pinv makes this easy
+        modulation_matrices = np.linalg.pinv(demod_matrices)
+
+        num_wave, num_binned_slit_pos, num_mod, num_stokes = modulation_matrices.shape
+
+        smoothed_mod = np.zeros((num_wave, num_full_slit_pos, num_mod, num_stokes))
+
+        # The binned abscissa is the "x" locations of the bins along the full spatial range
+        # The full abscissa is the "x" locations of all spatial pixels
+        # Add dummy dimensions because sklearn requires it
+        binned_abscissa = np.linspace(
+            start=0, stop=num_full_slit_pos, num=num_binned_slit_pos, endpoint=False
+        )[:, None]
+        full_abscissa = np.arange(num_full_slit_pos)[:, None]
+
+        model = self.build_RANSAC_model(fit_order=fit_order)
+        for w in range(num_wave):
+            for m in range(num_mod):
+                for s in range(num_stokes):
+                    curve = modulation_matrices[w, :, m, s]
+
+                    # Clean weirdo pixels
+                    fill_value = np.nanmedian(curve)
+                    curve[~np.isfinite(curve)] = fill_value
+
+                    model.fit(binned_abscissa, curve)
+                    fit_curve = model.predict(full_abscissa)
+                    smoothed_mod[w, :, m, s] = fit_curve
+
+        # Now compute the inverse again to return the DEmodulation matrices
+        smoothed_demod = np.linalg.pinv(smoothed_mod)
+
+        return smoothed_demod
+
+    def build_RANSAC_model(self, fit_order: int) -> Pipeline:
+        """
+        Build a scikit-learn pipeline from a set of estimators.
+
+        This method is taken directly from the ViSP L1 pipeline.
+        """
+        # PolynomialFeatures casts the pipeline as a polynomial fit
+        # see https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.PolynomialFeatures.html#sklearn.preprocessing.PolynomialFeatures
+        poly_feature = PolynomialFeatures(degree=fit_order)
+
+        # RobustScaler is a scale factor that is robust to outliers
+        # see https://scikit-learn.org/stable/modules/generated/sklearn.preprocessing.RobustScaler.html#sklearn.preprocessing.RobustScaler
+        scaler = RobustScaler()
+
+        # The RANSAC regressor iteratively sub-samples the input data and constructs a model with this sub-sample.
+        # The method used allows it to be robust to outliers.
+        # see https://scikit-learn.org/stable/modules/linear_model.html#ransac-regression
+        RANSAC = RANSACRegressor(
+            min_samples=self.parameters.visp_polcal_demod_spatial_smooth_min_samples
+        )
+
+        return make_pipeline(poly_feature, scaler, RANSAC)
+
+    def reshape_demod_matrices(
+        self, demod_matrices: np.ndarray, full_beam_shape: tuple[int, int]
+    ) -> np.ndarray:
+        """Upsample demodulation matrices to match the full beam size.
+
+        Given an input set of demodulation matrices with shape ``(X', Y', 4, M)`` resample the output to shape
+        ``(X, Y, 4, M)``, where ``X'`` and ``Y'`` are the binned size of the beam FOV, ``X`` and ``Y`` are the full beam shape, and ``M`` is the
+        number of modulator states.
+
+        If only a single demodulation matrix was made then it is returned as a single array with shape ``(4, M)``.
+
+
+        This method is taken directly from the ViSP L1 pipeline.
+
+        Parameters
+        ----------
+        demod_matrices
+            A set of demodulation matrices with shape ``(X', Y', 4, M)``
+
+        Returns
+        -------
+        If ``X'`` or ``Y'`` > 1 then upsampled matrices that are the full beam size ``(X, Y, 4, M)``. If ``X' == Y' == 1`` then a single matric for the whole FOV with shape ``(4, M)``
+
+        """
+        if len(demod_matrices.shape) != 4:
+            raise ValueError(
+                f"Expected demodulation matrices to have 4 dimensions. Got shape {demod_matrices.shape}"
+            )
+
+        # The non-demodulation matrix part of the larger array
+        data_shape = demod_matrices.shape[:2]
+        # The shape of a single demodulation matrix
+        demod_shape = demod_matrices.shape[-2:]
+        logger.info(f"Demodulation FOV sampling shape: {data_shape}")
+        logger.info(f"Demodulation matrix shape: {demod_shape}")
+        if data_shape == (1, 1):
+            # A single modulation matrix can be used directly, so just return it after removing extraneous dimensions
+            logger.info(f"Single demodulation matrix detected")
+            return demod_matrices[0, 0, :, :]
+
+        target_shape = full_beam_shape + demod_shape
+        logger.info(f"Target full-frame demodulation shape: {target_shape}")
+        full_frame_matrices = next(
+            resize_arrays(
+                demod_matrices, target_shape, order=self.parameters.visp_polcal_demod_upsample_order
+            )
+        )
+        return full_frame_matrices