dkist-processing-visp 3.3.0__py3-none-any.whl → 5.1.1__py3-none-any.whl

dkist_processing_visp/tasks/science.py

@@ -1,4 +1,5 @@
 """ViSP science calibration task. See :doc:`this page </science_calibration>` for more information."""
+
 from collections import defaultdict
 from dataclasses import dataclass
 from functools import cached_property
@@ -11,6 +12,7 @@ from astropy.time import TimeDelta
 from dkist_processing_common.codecs.fits import fits_access_decoder
 from dkist_processing_common.codecs.fits import fits_array_decoder
 from dkist_processing_common.codecs.fits import fits_hdulist_encoder
+from dkist_processing_common.models.fits_access import MetadataKey
 from dkist_processing_common.models.task_name import TaskName
 from dkist_processing_common.tasks.mixin.quality import QualityMixin
 from dkist_processing_math.arithmetic import divide_arrays_by_array
@@ -20,6 +22,7 @@ from dkist_processing_math.statistics import average_numpy_arrays
 from dkist_processing_pac.optics.telescope import Telescope
 from dkist_service_configuration.logging import logger
 
+from dkist_processing_visp.models.fits_access import VispMetadataKey
 from dkist_processing_visp.models.tags import VispTag
 from dkist_processing_visp.parsers.visp_l0_fits_access import VispL0FitsAccess
 from dkist_processing_visp.tasks.mixin.beam_access import BeamAccessMixin
@@ -48,6 +51,9 @@ class CalibrationCollection:
 
         This is done by considering that state offset values computed by the GeometricCalibration task. Any sub-pixel
         overlaps are rounded to the next integer that still guarantees overlap.
+
+        When "start pixels" are mentioned, those are pixels being counted from zero on a given axis in the positive direction.
+        When "end pixels" are mentioned, those are pixels being counted from the end of a given axis in the negative direction.
         """
         logger.info("Computing beam overlap slices")
         # This will be a flat list of (x, y) pairs for all modstates and beams
@@ -60,25 +66,33 @@ class CalibrationCollection:
         logger.info(f"All x shifts: {all_x_shifts}")
         logger.info(f"All y shifts: {all_y_shifts}")
 
-        # The amount we need to "slice in" from the front of the array is simply the maximum positive shift
+        # The amount we need to "slice in" from the start of the array is equivalent to the absolute value of the most negative shift.
         # The call to `np.ceil` ensures that the integer rounding doesn't allow non-overlap regions to leak in.
-        max_x = int(np.ceil(np.max(all_x_shifts)))
-        max_y = int(np.ceil(np.max(all_y_shifts)))
+        start_pixels_to_slice_x = int(np.ceil(abs(np.min(all_x_shifts))))
+        start_pixels_to_slice_y = int(np.ceil(abs(np.min(all_y_shifts))))
 
-        # The amount we need to "chop off" the end of the array is the most negative negative shift.
+        # The amount we need to "chop off" the end of the array is the most positive shift.
         #
         # Here we rely on the fact that the fiducial array's shift is *always* (0, 0)
-        #  (see `geometric.compute_modstate_offset`). Thus, if there are no negative shifts then the following lines
-        #  will result in None. This is required for slicing because array[x:0] is no good. So if the min is 0 then we
-        #  end up with array[x:None] which goes all the way to the end of the array.
+        # (see `geometric.compute_modstate_offset`). Thus, if there are no negative shifts then the following lines
+        # will result in None. This is required for slicing because array[x:0] is no good. So if the max is 0 then we
+        # end up with array[x:None] which goes all the way to the end of the array.
         #
-        # The call to `np.floor` ensures that the integer rounding doesn't allow non-overlap regions to leak in.
-        #  (because more negative slices will cut out more data).
-        min_x = int(np.floor(np.min(all_x_shifts))) or None
-        min_y = int(np.floor(np.min(all_y_shifts))) or None
+        # The call to `np.ceil` ensures that the integer rounding doesn't allow non-overlap regions to leak in.
+        # (because more negative slices will cut out more data).
+        end_pixels_to_slice_x = int(np.ceil(np.max(all_x_shifts))) or None
+        end_pixels_to_slice_y = int(np.ceil(np.max(all_y_shifts))) or None
+
+        # As the pixels to remove from the end of axes is given as a positive number, we need to make it negative for slicing.
+        if end_pixels_to_slice_x is not None:
+            end_pixels_to_slice_x *= -1
+
+        if end_pixels_to_slice_y is not None:
+            end_pixels_to_slice_y *= -1
 
-        x_slice = slice(max_x, min_x)
-        y_slice = slice(max_y, min_y)
+        # Construct the slices
+        x_slice = slice(start_pixels_to_slice_x, end_pixels_to_slice_x)
+        y_slice = slice(start_pixels_to_slice_y, end_pixels_to_slice_y)
 
         return x_slice, y_slice
 
@@ -125,17 +139,17 @@ class ScienceCalibration(
         None
 
         """
-        with self.apm_task_step("Loading calibration objects"):
+        with self.telemetry_span("Loading calibration objects"):
             calibrations = self.collect_calibration_objects()
 
-        with self.apm_task_step(
+        with self.telemetry_span(
             f"Processing Science Frames for "
             f"{self.constants.num_map_scans} map scans and "
             f"{self.constants.num_raster_steps} raster steps"
         ):
             self.process_frames(calibrations=calibrations)
 
-        with self.apm_processing_step("Computing and logging quality metrics"):
+        with self.telemetry_span("Computing and logging quality metrics"):
             no_of_raw_science_frames: int = self.scratch.count_all(
                 tags=[
                     VispTag.input(),
@@ -156,7 +170,7 @@ class ScienceCalibration(
         """
         dark_dict = defaultdict(dict)
         background_dict = dict()
-        solar_dict = defaultdict(dict)
+        solar_dict = dict()
         angle_dict = dict()
         state_offset_dict = defaultdict(dict)
         spec_shift_dict = dict()
@@ -211,6 +225,18 @@ class ScienceCalibration(
                 )
             )
 
+            # Solar
+            #######
+            solar_dict[VispTag.beam(beam)] = next(
+                self.read(
+                    tags=[
+                        VispTag.intermediate_frame(beam=beam),
+                        VispTag.task_solar_gain(),
+                    ],
+                    decoder=fits_array_decoder,
+                )
+            )
+
             # Demod
             #######
             if self.constants.correct_for_polarization:
@@ -225,18 +251,6 @@ class ScienceCalibration(
                 )
 
             for modstate in range(1, self.constants.num_modstates + 1):
-                # Solar
-                #######
-                solar_dict[VispTag.beam(beam)][VispTag.modstate(modstate)] = next(
-                    self.read(
-                        tags=[
-                            VispTag.intermediate_frame(beam=beam, modstate=modstate),
-                            VispTag.task_solar_gain(),
-                        ],
-                        decoder=fits_array_decoder,
-                    )
-                )
-
                 # State Offset
                 ##############
                 state_offset_dict[VispTag.beam(beam)][VispTag.modstate(modstate)] = next(
@@ -274,9 +288,10 @@ class ScienceCalibration(
                 for raster_step in range(0, self.constants.num_raster_steps):
                     beam_storage = dict()
                     header_storage = dict()
+                    nan_storage = dict()
                     for beam in range(1, self.constants.num_beams + 1):
                         apm_str = f"{map_scan = }, {raster_step = }, and {beam = }"
-                        with self.apm_processing_step(f"Basic corrections for {apm_str}"):
+                        with self.telemetry_span(f"Basic corrections for {apm_str}"):
                             # Initialize array_stack and headers
                             if self.constants.correct_for_polarization:
                                 logger.info(
@@ -285,6 +300,7 @@ class ScienceCalibration(
                                 (
                                     intermediate_array,
                                     intermediate_header,
+                                    nan_mask,
                                 ) = self.process_polarimetric_modstates(
                                     beam=beam,
                                     raster_step=raster_step,
@@ -296,7 +312,11 @@ class ScienceCalibration(
                                 logger.info(
                                     f"Processing spectrographic observe frames from {apm_str}"
                                 )
-                                intermediate_array, intermediate_header = self.correct_single_frame(
+                                (
+                                    intermediate_array,
+                                    intermediate_header,
+                                    nan_mask,
+                                ) = self.correct_single_frame(
                                     beam=beam,
                                     modstate=1,
                                     raster_step=raster_step,
@@ -307,20 +327,38 @@ class ScienceCalibration(
                                 intermediate_header = self.compute_date_keys(intermediate_header)
                             beam_storage[VispTag.beam(beam)] = intermediate_array
                             header_storage[VispTag.beam(beam)] = intermediate_header
+                            nan_storage[VispTag.beam(beam)] = nan_mask
 
-                    with self.apm_processing_step("Combining beams"):
+                    with self.telemetry_span("Combining beams"):
                         calibrated = self.combine_beams(beam_storage, header_storage, calibrations)
 
                     if self.constants.correct_for_polarization:
-                        with self.apm_processing_step("Correcting telescope polarization"):
+                        with self.telemetry_span("Correcting telescope polarization"):
                             calibrated = self.telescope_polarization_correction(calibrated)
 
+                    with self.telemetry_span("Combining NaN masks from beams"):
+                        cut_combined_nan_mask = self.combine_and_cut_nan_masks(
+                            list(nan_storage.values()), calibrations
+                        )
+
                     # Save the final output files
-                    with self.apm_writing_step("Writing calibrated arrays"):
+                    with self.telemetry_span("Writing calibrated arrays"):
                         self.write_calibrated_array(
-                            calibrated, map_scan=map_scan, calibrations=calibrations
+                            calibrated,
+                            map_scan=map_scan,
+                            calibrations=calibrations,
+                            nan_mask=cut_combined_nan_mask,
                         )
 
+    @staticmethod
+    def combine_and_cut_nan_masks(
+        nan_masks: list[np.ndarray], calibrations: CalibrationCollection
+    ) -> np.ndarray:
+        """Combine two NaN masks into one, cropping the result based on pre-calculated shifts."""
+        combined_nan_mask = np.logical_or.reduce(nan_masks)
+        x_slice, y_slice = calibrations.beams_overlap_slice
+        return combined_nan_mask[x_slice, y_slice]
+
     def process_polarimetric_modstates(
         self,
         beam: int,
@@ -328,7 +366,7 @@ class ScienceCalibration(
         map_scan: int,
         readout_exp_time: float,
         calibrations: CalibrationCollection,
-    ) -> tuple[np.ndarray, fits.Header]:
+    ) -> tuple[np.ndarray, fits.Header, np.ndarray]:
         """
         Process a single polarimetric beam as much as is possible.
 
@@ -340,11 +378,12 @@ class ScienceCalibration(
         ].shape
         array_stack = np.zeros(array_shape + (self.constants.num_modstates,))
         header_stack = []
+        nan_mask_stack = np.zeros(array_shape + (self.constants.num_modstates,))
 
-        with self.apm_processing_step(f"Correcting {self.constants.num_modstates} modstates"):
+        with self.telemetry_span(f"Correcting {self.constants.num_modstates} modstates"):
             for modstate in range(1, self.constants.num_modstates + 1):
                 # Correct the arrays
-                corrected_array, corrected_header = self.correct_single_frame(
+                corrected_array, corrected_header, nan_mask = self.correct_single_frame(
                     beam=beam,
                     modstate=modstate,
                     raster_step=raster_step,
@@ -355,15 +394,17 @@ class ScienceCalibration(
                 # Add this result to the 3D stack
                 array_stack[:, :, modstate - 1] = corrected_array
                 header_stack.append(corrected_header)
+                nan_mask_stack[:, :, modstate - 1] = nan_mask
 
-        with self.apm_processing_step("Applying instrument polarization correction"):
+        with self.telemetry_span("Applying instrument polarization correction"):
             intermediate_array = nd_left_matrix_multiply(
                 vector_stack=array_stack,
                 matrix_stack=calibrations.demod_matrices[VispTag.beam(beam)],
             )
             intermediate_header = self.compute_date_keys(header_stack)
 
-        return intermediate_array, intermediate_header
+        # The modulator state NaN masks are stacked along axis=2 with axis=0 & 1 being the array axes of one modstate
+        return intermediate_array, intermediate_header, np.logical_or.reduce(nan_mask_stack, axis=2)
 
     def combine_beams(
         self,
@@ -437,6 +478,7 @@ class ScienceCalibration(
         calibrated_object: VispL0FitsAccess,
         map_scan: int,
         calibrations: CalibrationCollection,
+        nan_mask: np.ndarray,
     ) -> None:
         """
         Write out calibrated science frames.
@@ -455,6 +497,9 @@ class ScienceCalibration(
         calibrations
             Calibration collection
 
+        nan_mask
+            A mask containing the known areas where data does not exist for both beams
+
         Returns
         -------
         None
@@ -470,7 +515,8 @@ class ScienceCalibration(
             stokes_I_data = calibrated_object.data[:, :, 0]
             for i, stokes_param in enumerate(self.constants.stokes_params):
                 stokes_data = calibrated_object.data[:, :, i]
-                final_data = self.re_dummy_data(stokes_data)
+                nan_masked_data = np.where(nan_mask, np.nan, stokes_data)
+                final_data = self.re_dummy_data(nan_masked_data)
                 pol_header = self.add_L1_pol_headers(final_header, stokes_data, stokes_I_data)
                 self.write_cal_array(
                     data=final_data,
@@ -480,7 +526,8 @@ class ScienceCalibration(
                     map_scan=map_scan,
                 )
         else:  # Only write stokes I
-            final_data = self.re_dummy_data(calibrated_object.data)
+            nan_masked_data = np.where(nan_mask, np.nan, calibrated_object.data)
+            final_data = self.re_dummy_data(nan_masked_data)
             self.write_cal_array(
                 data=final_data,
                 header=final_header,
@@ -497,7 +544,7 @@ class ScienceCalibration(
         map_scan: int,
         readout_exp_time: float,
         calibrations: CalibrationCollection,
-    ) -> tuple[np.ndarray, fits.Header]:
+    ) -> tuple[np.ndarray, fits.Header, np.ndarray]:
         """
         Apply basic corrections to a single frame.
 
@@ -539,7 +586,7 @@ class ScienceCalibration(
             VispTag.readout_exp_time(readout_exp_time)
         ]
         background_array = calibrations.background[VispTag.beam(beam)]
-        solar_gain_array = calibrations.solar_gain[VispTag.beam(beam)][VispTag.modstate(modstate)]
+        solar_gain_array = calibrations.solar_gain[VispTag.beam(beam)]
         angle = calibrations.angle[VispTag.beam(beam)]
         spec_shift = calibrations.spec_shift[VispTag.beam(beam)]
         state_offset = calibrations.state_offset[VispTag.beam(beam)][VispTag.modstate(modstate)]
@@ -592,7 +639,40 @@ class ScienceCalibration(
             self.corrections_remove_spec_geometry(geo_corrected_array, spec_shift)
         )
 
-        return spectral_corrected_array, observe_object.header
+        nan_mask = self.generate_nan_mask(
+            solar_corrected_array=solar_corrected_array,
+            state_offset=state_offset,
+            angle=angle,
+            spec_shift=spec_shift,
+        )
+
+        return (
+            spectral_corrected_array,
+            observe_object.header,
+            nan_mask,
+        )
+
+    def generate_nan_mask(
+        self,
+        solar_corrected_array: np.ndarray,
+        state_offset: np.ndarray,
+        angle: float,
+        spec_shift: np.ndarray,
+    ) -> np.ndarray:
+        """Calculate the NaN mask through geometric correction to be applied to the final L1 arrays."""
+        # Using a bi-cubic polynomial (order = 3) best converges to the desired result in the underlying fits
+        geo_corrected_with_nan = next(
+            self.corrections_correct_geometry(
+                solar_corrected_array, state_offset, angle, mode="constant", order=3, cval=np.nan
+            )
+        )
+        # Interpolating with nearest neighbor (order = 0) prevents NaN values from "taking over" the whole array
+        spectral_corrected_with_nan = next(
+            self.corrections_remove_spec_geometry(
+                geo_corrected_with_nan, spec_shift, cval=np.nan, order=0
+            )
+        )
+        return np.isnan(spectral_corrected_with_nan)
 
     def telescope_polarization_correction(
         self,
@@ -652,7 +732,7 @@ class ScienceCalibration(
         date_end = (Time(sorted_obj_list[-1].time_obs) + exp_time).isot
 
         header = sorted_obj_list[0].header
-        header["DATE-BEG"] = date_beg
+        header[MetadataKey.time_obs] = date_beg
         header["DATE-END"] = date_end
 
         return header