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

dkist_processing_trend/tasks/run_pac_fitter.py

@@ -0,0 +1,392 @@
+"""Task for running processed POLCAL data through the `dkist_processing_pac` fitter."""
+
+import numpy as np
+from dkist_processing_common.codecs.asdf import asdf_decoder
+from dkist_processing_common.codecs.fits import fits_array_encoder
+from dkist_processing_pac.fitter.fitter_parameters import CU_PARAMS
+from dkist_processing_pac.fitter.fitting_core import compare_I
+from dkist_processing_pac.fitter.polcal_fitter import PolcalFitter
+from dkist_processing_pac.input_data.drawer import Drawer
+from dkist_processing_pac.input_data.dresser import Dresser
+from dkist_service_configuration.logging import logger
+
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.parsers.trend_l0_fits_access import TrendL0FitsAccess
+from dkist_processing_trend.tasks.trend_base import TrendArmTaskBase
+
+__all__ = ["RunPacFitter"]
+
+# The are the names of parameters as used by `dkist-processing-pac`, which are slightly different than their "common" names
+ORDERED_PARAMETER_NAMES = [
+    "x12",
+    "t12",
+    "x34",
+    "t34",
+    "x56",
+    "t56",
+    "Q_in",
+    "U_in",
+    "V_in",
+    "t_pol",
+    "t_ret",
+    "ret0r",
+    "ret045",
+    "ret0h",
+    "py",
+]
+
+
+class RunPacFitter(TrendArmTaskBase):
+    """
+    Run prepared POLCAL data through `dkist_processing_pac` to fit the CU parameters and produce demodulation matrices.
+
+    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):
+        """
+        Run all prepared POLCAL datasets through the `~dkist_processing_pac.fitter.polcal_fitter.PolcalFitter` and save the results.
+
+        Here is where we loop over all PA&C fitting options described in the
+        `~dkist_processing_trend.models.parameters.TrendParameters.fit_options_list`. Each run gets its own unique set of
+        outputs.
+        """
+        num_beams = 1 if self.arm_id == "CI" else 2
+        cs_config_saved = False
+
+        for inst_options in self.parameters.instrument_processing_options:
+            inst_option_name = inst_options.name
+            for beam in range(1, num_beams + 1):
+                logger.info(f"Loading PAC input data for {beam = }")
+                global_input_dict = next(
+                    self.read(
+                        tags=[
+                            TrendTag.intermediate(),
+                            TrendTag.arm_id(self.arm_id),
+                            TrendTag.instrument_processing_options(inst_option_name),
+                            TrendTag.task_global_pac_input(),
+                            TrendTag.beam(beam),
+                        ],
+                        decoder=asdf_decoder,
+                    )
+                )
+                local_input_dict = next(
+                    self.read(
+                        tags=[
+                            TrendTag.intermediate(),
+                            TrendTag.arm_id(self.arm_id),
+                            TrendTag.instrument_processing_options(inst_option_name),
+                            TrendTag.task_local_pac_input(),
+                            TrendTag.beam(beam),
+                        ],
+                        decoder=asdf_decoder,
+                    )
+                )
+
+                if not cs_config_saved:
+                    logger.info("Writing CS GOS configuration")
+                    self.write_calibration_sequence(global_input_dict)
+                    cs_config_saved = True
+
+                for fit_options in self.parameters.fit_options_list:
+                    with self.telemetry_span(
+                        f"Run PAC fit for inst processing option = {inst_option_name}, {beam = }, and fit options = {fit_options.name}"
+                    ):
+                        logger.info(
+                            f"Running fits for inst processing option = {inst_option_name}, {beam = }, and fit options = {fit_options.name}"
+                        )
+                        global_dresser = self.populate_dresser(
+                            global_input_dict, remove_I_trend=fit_options.remove_I_trend
+                        )
+                        local_dresser = self.populate_dresser(
+                            local_input_dict, remove_I_trend=fit_options.remove_I_trend
+                        )
+
+                        pac_fitter = PolcalFitter(
+                            local_dresser=local_dresser,
+                            global_dresser=global_dresser,
+                            fit_mode=fit_options.fit_mode_name,
+                            init_set=fit_options.init_set_name,
+                            inherit_global_vary_in_local_fit=True,
+                            suppress_local_starting_values=True,
+                            fit_TM=False,
+                        )
+
+                    with self.telemetry_span(
+                        f"Saving fit results for {beam = } and fit options = {fit_options.name}"
+                    ):
+                        common_kwargs = {
+                            "beam": beam,
+                            "inst_options_name": inst_option_name,
+                            "fit_options_name": fit_options.name,
+                        }
+                        logger.info("Writing best-fit parameters")
+                        self.write_best_fit_parameters(pac_fitter, **common_kwargs)
+
+                        logger.info("Writing best-fit demodulation matrices")
+                        self.write_best_fit_demodulation_matrices(pac_fitter, **common_kwargs)
+
+                        logger.info("Writing best-fit flux and residuals")
+                        self.write_best_fit_flux_and_residuals(pac_fitter, **common_kwargs)
+
+    def populate_dresser(
+        self,
+        pac_input_dict: dict[int, list[dict[str, dict | np.ndarray]]],
+        remove_I_trend: bool,
+        skip_darks: bool = True,
+    ) -> Dresser:
+        """
+        Convert a raw set of dicts that live on scratch as an ASDF file into a `~dkist_processing_pac.input_data.dresser.Dresser`.
+
+        `Drawers <dkist_processing_pac.input_data.drawer.Drawer>` expect to be instantiated with a dict containing lists
+        of `~dkist_processing_trend.parsers.trend_l0_fits_access.TrendL0FitsAccess` objects but we store these
+        "pac input data" as lists of dicts containing "header" and "data" keys. So the first step is to convert these
+        raw dictionaries into `~dkist_processing_trend.parsers.trend_l0_fits_access.TrendL0FitsAccess`  objects, then we
+        can use them to instantiate `Drawers <dkist_processing_pac.input_data.drawer.Drawer>` and finally a
+        `~dkist_processing_pac.input_data.dresser.Dresser`.
+        """
+        dresser = Dresser()
+        fits_access_dict = dict()
+
+        for cs_step in range(self.constants.num_cs_steps):
+            cs_step_dict_list = pac_input_dict[cs_step]
+            cs_step_fits_access_list = []
+
+            for cs_step_dict in cs_step_dict_list:
+                fits_access_obj = TrendL0FitsAccess(
+                    header=cs_step_dict["header"], data=cs_step_dict["data"]
+                )
+                cs_step_fits_access_list.append(fits_access_obj)
+
+            fits_access_dict[cs_step] = cs_step_fits_access_list
+
+        dresser.add_drawer(
+            Drawer(fits_access_dict, remove_I_trend=remove_I_trend, skip_darks=skip_darks)
+        )
+        return dresser
+
+    def write_calibration_sequence(
+        self, pac_input_dict: dict[int, list[dict[str, dict | np.ndarray]]]
+    ) -> None:
+        """
+        Write the Calibration Sequence configuration to scratch.
+
+        The result has shape ``(K, 8, N)``, where ``K`` is the number of polcal OPs/calibration sequences, and ``N`` is
+        the number of steps per OP/sequence. The length-8 dimension corresponds to the 8 optical parameters that describe the
+        calibration sequence.
+
+        If the different polcal OPs have a different number of steps then ``N`` is the length of the longest OP and the data
+        for shorter OPs is padded with NaNs.
+        """
+        dresser = self.populate_dresser(pac_input_dict, remove_I_trend=False, skip_darks=False)
+
+        longest_cs_num_steps = max(dresser.drawer_step_list)
+        full_data = np.empty((dresser.numdrawers, 8, longest_cs_num_steps))
+
+        for cs_num, drawer in enumerate(dresser.drawers):
+            drawer_data_stack = np.vstack(
+                [
+                    drawer.pol_in,
+                    drawer.theta_pol_steps,
+                    drawer.ret_in,
+                    drawer.theta_ret_steps,
+                    drawer.dark_in,
+                    drawer.azimuth,
+                    drawer.elevation,
+                    drawer.table_angle,
+                ]
+            )
+
+            if step_deficit := drawer.numsteps - longest_cs_num_steps > 0:
+                drawer_data_stack = np.pad(
+                    drawer_data_stack, ((0, 0), (0, step_deficit)), constant_values=np.nan
+                )
+
+            full_data[cs_num, :, :] = drawer_data_stack
+
+        self.write(
+            data=full_data,
+            tags=[
+                TrendTag.intermediate(),
+                TrendTag.arm_id(self.arm_id),
+                TrendTag.task_calibration_sequence(),
+            ],
+            encoder=fits_array_encoder,
+        )
+
+    def write_best_fit_parameters(
+        self, fitter: PolcalFitter, beam: int, inst_options_name: str, fit_options_name: str
+    ):
+        """
+        Collect ALL best-fit parameter values (including those fixed in the fit) and write to a single array.
+
+        The full output shape is ``(*FOV_shape, K, 15 + N, 3)``, where FOV_shape is the shape of the input POLCAL data,
+        K is the number of polcal OPs/calibration sequences, and N is the number of steps in each sequence. The last
+        dimension contains the best-fit value in its 0th index, the initial value at index 1, and whether that parameter
+        varied in the fit at index 2. Note that many parameters are held to be the same across polcal OPs; these
+        parameters will have the same values for all slices of the "K" index.
+
+        If the different polcal OPs have a different number of steps then N is the length of the longest OP and the data
+        for shorter OPs is padded with NaNs.
+        """
+        fov_shape = fitter.local_objects.dresser.shape
+        num_cs = fitter.local_objects.dresser.numdrawers
+        steps_per_cs = fitter.local_objects.dresser.drawer_step_list
+        longest_cs_num_steps = max(steps_per_cs)
+        init_parameters = fitter.local_objects.init_parameters
+        fit_parameters = fitter.fit_parameters
+
+        full_output = np.empty(fov_shape + (num_cs, 15 + longest_cs_num_steps, 3))
+
+        num_points = np.prod(fov_shape)
+        for point in range(num_points):
+            idx = np.unravel_index(point, fov_shape)
+            point_parameters = fit_parameters[idx]
+            init_point_parameters = init_parameters[idx]
+
+            for cs in range(num_cs):
+                cs_data = np.empty((15 + steps_per_cs[cs], 3))
+
+                # These are all parameters except for the I_sys params, which we'll add in a moment
+                for p, param_name in enumerate(ORDERED_PARAMETER_NAMES):
+
+                    if param_name in CU_PARAMS:
+                        param_name = f"{param_name}_CS{cs:02n}"
+
+                    if param_name == "py":
+                        value = fitter.global_objects.calibration_unit.py
+                        vary = False
+                        init_value = value
+                    else:
+                        parameter = point_parameters[param_name]
+                        value = parameter.value
+                        vary = parameter.vary
+                        init_value = init_point_parameters[param_name].value
+
+                    cs_data[p] = np.array([value, init_value, vary])
+
+                # Add I_sys parameters, the number of which depends on the number of *total* CS steps (across all
+                # calibration sequences).
+                for p, s in enumerate(range(steps_per_cs[cs]), start=15):
+                    param_name = f"I_sys_CS{cs:02n}_step{s:02n}"
+                    parameter = point_parameters[param_name]
+                    value = parameter.value
+                    vary = parameter.vary
+                    init_value = init_point_parameters[param_name].value
+                    cs_data[p] = np.array([value, init_value, vary])
+
+                if step_deficit := steps_per_cs[cs] - longest_cs_num_steps > 0:
+                    cs_data = np.pad(cs_data, ((0, step_deficit), (0, 0)), constant_values=np.nan)
+
+                full_output[*idx, cs] = cs_data
+
+        self.write(
+            data=full_output,
+            tags=[
+                TrendTag.intermediate(),
+                TrendTag.arm_id(self.arm_id),
+                TrendTag.beam(beam),
+                TrendTag.instrument_processing_options(inst_options_name),
+                TrendTag.pac_fit_options(fit_options_name),
+                TrendTag.task_best_fit_parameters(),
+            ],
+            encoder=fits_array_encoder,
+        )
+
+    def write_best_fit_demodulation_matrices(
+        self, fitter: PolcalFitter, beam: int, inst_options_name: str, fit_options_name: str
+    ):
+        """
+        Write the best-fit demodulation matrices to scratch.
+
+        The shape will be ``(*FOV_shape, 4, M)``, where FOV_shape is the shape of the input data and M is the number of
+        modulation states.
+        """
+        demod_matrices = fitter.demodulation_matrices
+
+        self.write(
+            data=demod_matrices,
+            tags=[
+                TrendTag.intermediate(),
+                TrendTag.arm_id(self.arm_id),
+                TrendTag.beam(beam),
+                TrendTag.instrument_processing_options(inst_options_name),
+                TrendTag.pac_fit_options(fit_options_name),
+                TrendTag.task_best_fit_demodulation_matrices(),
+            ],
+            encoder=fits_array_encoder,
+        )
+
+    def write_best_fit_flux_and_residuals(
+        self, fitter: PolcalFitter, beam: int, inst_options_name: str, fit_options_name: str
+    ):
+        """Compute the best-fit flux and fit residuals and write them to scratch."""
+        fit_container = fitter.local_objects
+        TM = fit_container.telescope
+        CM = fit_container.calibration_unit
+        fov_shape = fit_container.dresser.shape
+        num_mod = fit_container.dresser.nummod
+        num_steps = fit_container.dresser.numsteps
+        num_points = np.prod(fov_shape)
+
+        flux_array = np.zeros(fov_shape + (num_mod, num_steps))
+        residual_array = np.zeros_like(flux_array)
+
+        for i in range(num_points):
+            idx = np.unravel_index(i, fov_shape)
+            I_cal, I_unc = fit_container.dresser[idx]
+            fit_params = fit_container.fit_parameters[idx]
+            modmat = np.zeros((I_cal.shape[0], 4), dtype=np.float64)
+            flat_residual = compare_I(
+                params=fit_params,
+                I_cal=I_cal,
+                I_unc=I_unc,
+                TM=TM,
+                CM=CM,
+                modmat=modmat,
+                use_M12=True,
+            )
+            diff = np.reshape(flat_residual, (num_mod, num_steps))
+            residual_array[*idx, :, :] = diff
+
+            flux = diff * I_unc + I_cal
+            flux_array[*idx, :, :] = flux
+
+        self.write(
+            data=flux_array,
+            tags=[
+                TrendTag.intermediate(),
+                TrendTag.arm_id(self.arm_id),
+                TrendTag.beam(beam),
+                TrendTag.instrument_processing_options(inst_options_name),
+                TrendTag.pac_fit_options(fit_options_name),
+                TrendTag.task_best_fit_flux(),
+            ],
+            encoder=fits_array_encoder,
+        )
+        self.write(
+            data=residual_array,
+            tags=[
+                TrendTag.intermediate(),
+                TrendTag.arm_id(self.arm_id),
+                TrendTag.beam(beam),
+                TrendTag.instrument_processing_options(inst_options_name),
+                TrendTag.pac_fit_options(fit_options_name),
+                TrendTag.task_fit_residuals(),
+            ],
+            encoder=fits_array_encoder,
+        )

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 *