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

dkist_processing_trend/tasks/trend_base.py

@@ -0,0 +1,97 @@
+"""Base classes for all Trend science tasks."""
+
+from abc import ABC
+
+from dkist_processing_common.tasks.base import WorkflowTaskBase
+from dkist_service_configuration.logging import logger
+
+from dkist_processing_trend.models.constants import TrendConstants
+from dkist_processing_trend.models.parameters import TrendParameters
+
+
+class TrendTaskBase(
+    WorkflowTaskBase,
+    ABC,
+):
+    """
+    Task class for base Trend tasks.
+
+    Parameters
+    ----------
+    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
+    """
+
+    # So tab completion shows all the constants
+    constants: TrendConstants
+
+    @property
+    def constants_model_class(self):
+        """Get Trend pipeline constants."""
+        return TrendConstants
+
+    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 = TrendParameters(
+            scratch=self.scratch,
+            obs_ip_start_time=self.constants.earliest_ip_start_time,
+            instrument=self.constants.instrument,
+        )
+
+
+class TrendArmTaskBase(TrendTaskBase, ABC):
+    """
+    Task class base for Trend tasks that operate on a single instrument arm.
+
+    Similar to `TrendTaskBase` except it exposes the `self.arm_id` parameter that returns the arm ID this class was
+    instantiated with.
+
+    Another important feature of this class is that the `run` method will NOT be called if `self.arm_id <arm_id>` is not found
+    in `self.constants.arm_id_list <dkist_processing_trend.models.constants.TrendConstants.arm_id_list>` (because data for this arm doesn't exist in the input dataset).
+
+    Parameters
+    ----------
+    arm_id
+        id of the arm that this task will 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 __init__(
+        self, arm_id: str | int, 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.arm_id = arm_id
+
+    def pre_run(self) -> None:
+        """Check if `self.arm_id <arm_id>` is in the list of arms present in the dataset and run `run` if so."""
+        super().pre_run()
+        if self.arm_id not in self.constants.arm_id_list:
+            logger.info(f"This dataset has no data for arm {self.arm_id}. Nothing to do.")
+            self.run = lambda: None

dkist_processing_trend/tasks/trend_output_data.py

@@ -0,0 +1,167 @@
+"""Task for transferring trend data from scratch to datacenter stores."""
+
+from datetime import datetime
+from pathlib import Path
+from typing import Iterable
+
+from dkist_processing_common.models.message import CatalogFrameMessage
+from dkist_processing_common.models.message import CatalogFrameMessageBody
+from dkist_processing_common.tasks.mixin.globus import GlobusMixin
+from dkist_processing_common.tasks.mixin.interservice_bus import InterserviceBusMixin
+from dkist_processing_common.tasks.output_data_base import OutputDataBase
+from dkist_processing_common.tasks.output_data_base import TransferDataBase
+from dkist_service_configuration.logging import logger
+
+from dkist_processing_trend.models.constants import TrendConstants
+
+__all__ = ["TransferTrendData", "PublishTrendCatalogMessages"]
+
+
+class TrendOutputDataBase(OutputDataBase):
+    """
+    Base class that defines the destination folders for trend data.
+
+    Parameters
+    ----------
+    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
+    """
+
+    # So tab completion shows all the constants
+    constants: TrendConstants
+
+    @property
+    def constants_model_class(self):
+        """Define the constants class used to access the constants db."""
+        return TrendConstants
+
+    @property
+    def destination_root_folder(self) -> Path:
+        """
+        Define the root folder.
+
+        "trend/polcal"
+        """
+        return Path("trend") / "polcal"
+
+    @property
+    def destination_folder(self) -> Path:
+        """
+        Define the destination folder for this trend run.
+
+        "trend/polcal/{EARLIEST_IP_START_TIME}/{INSTRUMENT}"
+        """
+        # E.g., 1999-01-02T12:34:56.12352 -> 19990102T123456
+        formatted_date = datetime.fromisoformat(self.constants.earliest_ip_start_time).strftime(
+            "%Y%m%dT%H%M%S"
+        )
+        return self.destination_root_folder / formatted_date / self.constants.instrument
+
+
+class TransferTrendData(TrendOutputDataBase, TransferDataBase, GlobusMixin):
+    """
+    Task class for transferring processed trend data to the object store.
+
+    Parameters
+    ----------
+    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 transfer_objects(self):
+        """Transfer output frames."""
+        with self.telemetry_span("Upload output frames"):
+            self.transfer_output_frames()
+
+    def transfer_output_frames(self):
+        """Create a Globus transfer for all output data, as well as any available dataset extras."""
+        output_transfer_items = self.build_output_frame_transfer_list()
+
+        logger.info(
+            f"Preparing globus transfer {len(output_transfer_items)} items. "
+            f"recipe_run_id={self.recipe_run_id}. "
+            f"transfer_items={output_transfer_items[:3]}..."
+        )
+
+        self.globus_transfer_scratch_to_object_store(
+            transfer_items=output_transfer_items,
+            label=f"Transfer trend output frames for recipe_run_id {self.recipe_run_id}",
+        )
+
+
+class PublishTrendCatalogMessages(TrendOutputDataBase, InterserviceBusMixin):
+    """
+    Task class for publishing catalog messages related to the frames transferred to the object store.
+
+    Parameters
+    ----------
+    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 run(self) -> None:
+        """Run method for this task."""
+        with self.telemetry_span("Gather output data"):
+            frames = self.read(tags=self.output_frame_tags)
+
+        with self.telemetry_span("Create message objects"):
+            messages = self.frame_messages(paths=frames)
+            frame_message_count = len(messages)
+
+        with self.telemetry_span(f"Publish messages: {frame_message_count = }"):
+            self.interservice_bus_publish(messages=messages)
+
+    def frame_messages(self, paths: Iterable[Path]) -> list[CatalogFrameMessage]:
+        """
+        Create the frame messages.
+
+        Parameters
+        ----------
+        paths
+            The input paths for which to publish frame messages
+        folder_modifier
+            A subdirectory to use if the files in paths are not in the base directory
+
+        Returns
+        -------
+        A list of frame messages
+        """
+        message_bodies = [
+            CatalogFrameMessageBody(
+                objectName=self.format_object_key(path=p),
+                conversationId=str(self.recipe_run_id),
+                bucket=self.destination_bucket,
+            )
+            for p in paths
+        ]
+        messages = [CatalogFrameMessage(body=body) for body in message_bodies]
+        return messages
+
+    def rollback(self):
+        """
+        Warn that the metadata-store and the interservice bus retain the effect of this tasks execution.
+
+        Rolling back this task may not be achievable without other action.
+        """
+        super().rollback()
+        logger.warning(
+            f"Modifications to the metadata store and the interservice bus were not rolled back."
+        )

dkist_processing_trend/tasks/visp/__init__.py

@@ -0,0 +1,6 @@
+"""Package for tasks specific to ViSP data."""
+
+from dkist_processing_trend.tasks.visp.visp_dmpd import *
+from dkist_processing_trend.tasks.visp.visp_geometric import *
+from dkist_processing_trend.tasks.visp.visp_prep_fit_data import *
+from dkist_processing_trend.tasks.visp.visp_process_demod import *

dkist_processing_trend/tasks/visp/visp_dmpd.py

@@ -0,0 +1,405 @@
+"""Task for producing ViSP Demodulated PolcalData (DMPD)."""
+
+from dataclasses import dataclass
+from functools import cached_property
+
+import numpy as np
+from dkist_processing_common.codecs.asdf import asdf_decoder
+from dkist_processing_common.codecs.fits import fits_array_decoder
+from dkist_processing_common.codecs.fits import fits_array_encoder
+from dkist_processing_math.arithmetic import divide_arrays_by_array
+from dkist_processing_math.arithmetic import subtract_array_from_arrays
+from dkist_processing_math.linear_algebra import nd_left_matrix_multiply
+from dkist_processing_math.statistics import average_numpy_arrays
+from dkist_service_configuration.logging import logger
+from scipy.ndimage import affine_transform
+
+from dkist_processing_trend.models.fit_options import FitOptions
+from dkist_processing_trend.models.instrument_options import VispInstrumentOptions
+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__ = ["VispDemodulatedPolcalData"]
+
+
+@dataclass
+class VispCalibrationCollection:
+    """Dataclass to hold all calibration objects needed for ViSP "science" calibration."""
+
+    dark: np.ndarray
+    gain: np.ndarray
+    angle: dict
+    shift: dict
+    demod_matrices: dict
+
+    @cached_property
+    def beams_overlap_slice(self) -> tuple[slice, slice]:
+        """
+        Compute array slices that will extract the largest region with overlap from both beams.
+
+        This is done by considering the shifts 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 both beams
+        flat_shifts = list(self.shift.values())
+        # Split out into an x list and a y list
+        all_x_shifts, all_y_shifts = zip(*flat_shifts)
+        all_x_shifts = np.array(all_x_shifts)
+        all_y_shifts = np.array(all_y_shifts)
+
+        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 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.
+        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 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 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.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
+
+        # 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)
+
+        logger.info(f"Trimming non-overlapping beam edges by ({x_slice}, {y_slice})")
+        return x_slice, y_slice
+
+
+class VispDemodulatedPolcalData(TrendArmTaskBase):
+    """
+    Task class for producing demodulated polcal data (DMPD) for ViSP.
+
+    Demodulated polcal data is essentially a "science" reduction of the input POLCAL 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
+    """
+
+    record_provenance = True
+
+    def run(self):
+        """
+        Loop over all PAC and instrument options and produce a set of demodulated polcal data for each.
+
+        The "set" includes data for each beam prior to combination and then a final array for the combined data.
+
+        Generally, the algorithm for each set of PAC and instrument options is:
+
+        #. Collect all calibration objects needed (dark, gain, geometric, demodulation matrices)
+        #. Apply dark and gain corrections
+        #. Split by beam and demodulated
+        #. Apply geometric corrections
+        #. Combined beams
+        #. Save each beam's output and the combined output
+        """
+        for fit_options in self.parameters.fit_options_list:
+            for inst_options in self.parameters.instrument_processing_options:
+                log_str = f"fit option {fit_options.name} and instrument option {inst_options.name}"
+
+                logger.info(f"Collecting calibration objects for {log_str}")
+                calibrations = self.collect_calibration_objects(
+                    fit_options=fit_options, instrument_options=inst_options
+                )
+                beam_array = extract_visp_beam(
+                    calibrations.gain, beam=1, beam_border=self.parameters.visp_beam_border
+                )
+                beam_shape = beam_array.shape
+                x_slice, y_slice = calibrations.beams_overlap_slice
+                cut_shape = beam_array[x_slice, y_slice].shape
+
+                with self.telemetry_span(f"Processing {self.constants.num_cs_steps} CS steps"):
+                    beam1_full_dmpd = np.full(
+                        (1, self.constants.num_cs_steps, 4, *cut_shape), np.nan
+                    )
+                    beam2_full_dmpd = np.full_like(beam1_full_dmpd, np.nan)
+                    combined_full_dmpd = np.full_like(beam1_full_dmpd, np.nan)
+                    for cs_step in range(self.constants.num_cs_steps):
+                        logger.info(
+                            f"Processing demodulated polcal data for {cs_step = } with {log_str}"
+                        )
+                        beam1_dmpd, beam2_dmpd, combined_dmpd = self.process_single_cs_step(
+                            cs_step=cs_step, calibrations=calibrations, beam_shape=beam_shape
+                        )
+
+                        # np.moveaxis here because the stokes axis is the last axis in `self.process_single_cs_step`
+                        # (to make math easier), but we want the DMPD outputs to have stokes before the array axes.
+                        beam1_full_dmpd[0, cs_step, :, :, :] = np.moveaxis(beam1_dmpd, -1, 0)
+                        beam2_full_dmpd[0, cs_step, :, :, :] = np.moveaxis(beam2_dmpd, -1, 0)
+                        combined_full_dmpd[0, cs_step, :, :, :] = np.moveaxis(combined_dmpd, -1, 0)
+
+                self.write_dmpd(
+                    beam1_full_dmpd,
+                    beam=1,
+                    fit_options=fit_options,
+                    instrument_options=inst_options,
+                )
+                self.write_dmpd(
+                    beam2_full_dmpd,
+                    beam=2,
+                    fit_options=fit_options,
+                    instrument_options=inst_options,
+                )
+                self.write_dmpd(
+                    beam1_full_dmpd,
+                    fit_options=fit_options,
+                    instrument_options=inst_options,
+                    beam=None,
+                )
+
+    def collect_calibration_objects(
+        self, fit_options: FitOptions, instrument_options: VispInstrumentOptions
+    ) -> VispCalibrationCollection:
+        """Collect the INTERMEDIATE calibration objects needed to correct demodulated polcal data."""
+        angle_dict = dict()
+        shift_dict = dict()
+        demod_dict = dict()
+
+        base_tags = [TrendTag.intermediate(), TrendTag.arm_id(self.arm_id)]
+
+        # Dark
+        ######
+        full_dark_array = next(
+            self.read(
+                tags=base_tags + [TrendTag.frame(), TrendTag.task_dark()],
+                decoder=fits_array_decoder,
+            )
+        )
+
+        # Gain
+        #######
+        full_gain_array = next(
+            self.read(
+                tags=base_tags + [TrendTag.frame(), TrendTag.task_gain()],
+                decoder=fits_array_decoder,
+            )
+        )
+
+        for beam in [1, 2]:
+            beam_tags = base_tags + [TrendTag.beam(beam)]
+
+            # Geo
+            #######
+            geo_dict = next(
+                self.read(
+                    tags=beam_tags + [TrendTag.task_visp_geometric_calibration()],
+                    decoder=asdf_decoder,
+                )
+            )
+            angle_dict[TrendTag.beam(beam)] = geo_dict["angle_rad"]
+            shift_dict[TrendTag.beam(beam)] = geo_dict["shift"]
+
+            # Demod
+            #######
+            demod_matrices = next(
+                self.read(
+                    tags=beam_tags
+                    + [
+                        TrendTag.pac_fit_options(fit_options.name),
+                        TrendTag.instrument_processing_options(instrument_options.name),
+                        TrendTag.task_processed_demodulation_matrices(),
+                    ],
+                    decoder=fits_array_decoder,
+                )
+            )
+            demod_dict[TrendTag.beam(beam)] = demod_matrices
+
+        return VispCalibrationCollection(
+            dark=full_dark_array,
+            gain=full_gain_array,
+            angle=angle_dict,
+            shift=shift_dict,
+            demod_matrices=demod_dict,
+        )
+
+    def process_single_cs_step(
+        self,
+        cs_step: int,
+        calibrations: VispCalibrationCollection,
+        beam_shape: tuple[int, int],
+    ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+        """
+        Fully process the data for a single CS step.
+
+        To simplify the code, the dark and gain corrections are done on the full-frame before beam extraction.
+        Demodulation and geometric corrections are applied on a per-beam basis.
+
+        All outputs are trimmed to only include the regions that exist in both beams.
+
+        Returns
+        -------
+        The corrected beam 1 array
+        The corrected beam 2 array
+        The corrected combined array
+        """
+        modstate_full_frame_data_list = []
+        for modstate in range(1, self.constants.num_modstates + 1):
+            tags = [
+                TrendTag.input(),
+                TrendTag.frame(),
+                TrendTag.arm_id(self.arm_id),
+                TrendTag.task_polcal(),
+                TrendTag.cs_step(cs_step),
+                TrendTag.modstate(modstate),
+            ]
+            input_arrays = self.read(tags=tags, decoder=fits_array_decoder)
+            avg_input = average_numpy_arrays(input_arrays)
+
+            dark_corrected = subtract_array_from_arrays(
+                arrays=avg_input, array_to_subtract=calibrations.dark
+            )
+            gain_corrected = next(
+                divide_arrays_by_array(arrays=dark_corrected, array_to_divide_by=calibrations.gain)
+            )
+
+            modstate_full_frame_data_list.append(gain_corrected)
+
+        final_beam_data_list = []
+        for beam in [1, 2]:
+            demod_matrices = calibrations.demod_matrices[TrendTag.beam(beam)]
+            beam_stack = np.full(beam_shape + (self.constants.num_modstates,), np.nan)
+            for modstate in range(1, self.constants.num_modstates + 1):
+                full_frame_data = modstate_full_frame_data_list[modstate - 1]
+                beam_data = extract_visp_beam(
+                    full_frame_data, beam=beam, beam_border=self.parameters.visp_beam_border
+                )
+                beam_stack[:, :, modstate - 1] = beam_data
+
+            demodulated_data = nd_left_matrix_multiply(
+                vector_stack=beam_stack,
+                matrix_stack=demod_matrices,
+            )
+
+            geo_corrected = np.full_like(demodulated_data, np.nan)
+            for s in range(demodulated_data.shape[-1]):
+                geo_corrected[:, :, s] = self.correct_geometry(
+                    demodulated_data[:, :, s],
+                    angle=calibrations.angle[TrendTag.beam(beam)],
+                    shift=calibrations.shift[TrendTag.beam(beam)],
+                )
+
+            x_slice, y_slice = calibrations.beams_overlap_slice
+            cut_array = geo_corrected[x_slice, y_slice, :]
+
+            final_beam_data_list.append(cut_array)
+
+        combined_data = self.combine_beams(
+            beam1_data=final_beam_data_list[0], beam2_data=final_beam_data_list[1]
+        )
+
+        return final_beam_data_list[0], final_beam_data_list[1], combined_data
+
+    def combine_beams(self, beam1_data: np.ndarray, beam2_data: np.ndarray) -> np.ndarray:
+        r"""
+        Combine polarimetric beams so that polarization states are normalized by the intensity state (Stokes I).
+
+        In other words:
+
+        .. math::
+
+                I_{avg} & = (I_1 + I_2) / 2
+
+                Q_{avg} & = I_{avg} (Q_1 / I_1 + Q_2 / I_2) / 2
+
+        ...and the same for U and V
+        """
+        avg_data = np.zeros_like(beam1_data)
+
+        # Rely on the fact that the Stokes states are in order after demodulation
+        avg_I = (beam1_data[:, :, 0] + beam2_data[:, :, 0]) / 2.0
+        avg_data[:, :, 0] = avg_I
+
+        for stokes in range(1, 4):
+            beam1_norm = beam1_data[:, :, stokes] / beam1_data[:, :, 0]
+            beam2_norm = beam2_data[:, :, stokes] / beam2_data[:, :, 0]
+            avg_data[:, :, stokes] = avg_I * (beam1_norm + beam2_norm) / 2.0
+
+        return avg_data
+
+    def correct_geometry(
+        self, array: np.ndarray, angle: float, shift: tuple[float, float]
+    ) -> np.ndarray:
+        """
+        Rotate and shift an array.
+
+        Note that the input ``angle`` and ``shift`` arguments are the *amount to rotate/shift by*. They are passed straight
+        through to the affine transform.
+        """
+        array_shape = array.shape
+        cosine_angle = np.cos(angle)
+        sine_angle = np.sin(angle)
+        inverse_rotation_matrix = np.array(
+            [[cosine_angle, -sine_angle], [sine_angle, cosine_angle]],
+            dtype=np.float64,
+        )
+
+        rotation_center = 0.5 * (np.asarray(array_shape, dtype=np.float64) - 1.0)
+        transform_offset = rotation_center + shift - inverse_rotation_matrix @ rotation_center
+
+        corrected_array = affine_transform(
+            array,
+            matrix=inverse_rotation_matrix,
+            offset=transform_offset,
+            order=5,
+            mode="nearest",
+            cval=np.nan,
+            prefilter=True,
+        )
+        return corrected_array
+
+    def write_dmpd(
+        self,
+        array: np.ndarray,
+        beam: int | None,
+        fit_options: FitOptions,
+        instrument_options: VispInstrumentOptions,
+    ) -> None:
+        """
+        Write a demodulated polcal data output to scratch.
+
+        If ``beam`` is not `None` then the output will be tagged as a `~dkist_processing_trend.models.tags.TrendTag.task_single_beam_demodulated_polcal_data`,
+        otherwise the output is tagged with `~dkist_processing_trend.models.tags.TrendTag.task_demodulated_polcal_data`.
+        """
+        tags = [
+            TrendTag.intermediate(),
+            TrendTag.arm_id(self.arm_id),
+            TrendTag.pac_fit_options(fit_options.name),
+            TrendTag.instrument_processing_options(instrument_options.name),
+        ]
+        if beam is not None:
+            tags += [TrendTag.beam(beam), TrendTag.task_single_beam_demodulated_polcal_data()]
+        else:
+            tags.append(TrendTag.task_demodulated_polcal_data())
+        self.write(data=array, tags=tags, encoder=fits_array_encoder)

dkist_processing_trend/tasks/visp/visp_extract_beam.py

@@ -0,0 +1,14 @@
+"""Function to split apart the two ViSP beams."""
+
+from typing import Literal
+
+import numpy as np
+
+
+def extract_visp_beam(array: np.ndarray, beam: Literal[1, 2], beam_border: int) -> np.ndarray:
+    """Extract a single beam from a raw, dual-beam array."""
+    if beam == 1:
+        return np.copy(array[:beam_border, ...])
+    if beam == 2:
+        return np.copy(array[beam_border:, ...][::-1, :])
+    raise ValueError(f"ViSP only has beams 1 and 2. Can't extract data for beam '{beam}'.")