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

dkist_processing_trend/tests/local_trial_workflows/l0_to_trend_visp_polcal.py

@@ -0,0 +1,294 @@
+import argparse
+import sys
+from pathlib import Path
+
+from dkist_processing_common.manual import ManualProcessing
+from dkist_service_configuration.logging import logger
+
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.tasks import DarkCalibration
+from dkist_processing_trend.tasks import GainCalibration
+from dkist_processing_trend.tasks import ParseL0TrendInputData
+from dkist_processing_trend.tasks import RunPacFitter
+from dkist_processing_trend.tasks.visp import VispDemodulatedPolcalData
+from dkist_processing_trend.tasks.visp import VispGeometricCalibration
+from dkist_processing_trend.tasks.visp import VispPrepareFitData
+from dkist_processing_trend.tasks.visp import VispProcessDemodulationMatrices
+from dkist_processing_trend.tasks.write_trend import WriteTrendOutputs
+from dkist_processing_trend.tests.conftest import TrendInputDatasetParameterValues
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import get_arm_id_list
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import (
+    input_dataset_parameter_task,
+)
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import (
+    load_dark_cal_task,
+)
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import load_dmpd_task
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import (
+    load_gain_cal_task,
+)
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import load_pac_fit_task
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import load_parsing_task
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import (
+    load_prepped_cal_task,
+)
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import (
+    load_processed_demod_cal_task,
+)
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import (
+    load_visp_geo_cal_task,
+)
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import (
+    save_dark_cal_task,
+)
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import save_dmpd_task
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import (
+    save_gain_cal_task,
+)
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import save_pac_fit_task
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import save_parsing_task
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import (
+    save_prepped_cal_task,
+)
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import (
+    save_processed_demod_cal_task,
+)
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import (
+    save_visp_geo_cal_task,
+)
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import tag_inputs_task
+from dkist_processing_trend.tests.local_trial_workflows.local_trial_helpers import translate_task
+from dkist_processing_trend.workflows.arm_class_factory import arm_class_factory
+
+
+def main(
+    scratch_path: str,
+    suffix: str = "FITS",
+    recipe_run_id: int = 2,
+    skip_translation: bool = False,
+    only_translate: bool = False,
+    load_input_parsing: bool = False,
+    load_dark: bool = False,
+    load_gain: bool = False,
+    load_geometric: bool = False,
+    load_prepped: bool = False,
+    load_pac_fit: bool = False,
+    load_processed_demod: bool = False,
+    load_dmpd: bool = False,
+):
+    with ManualProcessing(
+        workflow_path=Path(scratch_path),
+        recipe_run_id=recipe_run_id,
+        testing=True,
+        workflow_name="visp-l0-to-trend-polcal",
+        workflow_version="GROGU",
+    ) as manual_processing_run:
+        if not skip_translation:
+            manual_processing_run.run_task(task=translate_task(instrument="visp", suffix=suffix))
+        if only_translate:
+            return
+
+        manual_processing_run.run_task(task=tag_inputs_task(suffix))
+
+        parameter_values = TrendInputDatasetParameterValues(
+            trend_fit_options_list=[
+                {
+                    "name": "standard",
+                    "remove_I_trend": True,
+                    "fit_mode_name": "use_M12_I_sys_per_step",
+                    "init_set_name": "OCCal_VIS",
+                }
+            ],
+            trend_instrument_processing_options_visp=[
+                {
+                    "name": "full_spatial",
+                    "num_spatial_px": 256,
+                    "num_spectral_px": 1,
+                    "spatial_smoothing_fit_order": 17,
+                },
+                {
+                    "name": "single_point",
+                    "num_spatial_px": 1,
+                    "num_spectral_px": 1,
+                    "spatial_smoothing_fit_order": 1,
+                },
+            ],
+        )
+        manual_processing_run.run_task(task=input_dataset_parameter_task(parameter_values))
+
+        if load_input_parsing:
+            manual_processing_run.run_task(task=load_parsing_task(save_file="input_parsing.asdf"))
+        else:
+            manual_processing_run.run_task(task=ParseL0TrendInputData)
+            manual_processing_run.run_task(
+                task=save_parsing_task(
+                    tag_list=[TrendTag.input()],
+                    save_file="input_parsing.asdf",
+                )
+            )
+
+        arm_id_list = get_arm_id_list(
+            workflow_base_path=Path(scratch_path) / str(recipe_run_id),
+            parsing_file_name="input_parsing.asdf",
+        )
+
+        logger.info(f"Found {arm_id_list = }")
+
+        for arm_id in arm_id_list:
+            logger.info(f"Working on {arm_id = }")
+            dark_task = arm_class_factory(DarkCalibration, arm_id=arm_id)
+            gain_task = arm_class_factory(GainCalibration, arm_id=arm_id)
+            geo_task = arm_class_factory(VispGeometricCalibration, arm_id=arm_id)
+            prep_data_task = arm_class_factory(VispPrepareFitData, arm_id=arm_id)
+            run_fit_task = arm_class_factory(RunPacFitter, arm_id=arm_id)
+            process_demod_task = arm_class_factory(VispProcessDemodulationMatrices, arm_id=arm_id)
+            dmpd_task = arm_class_factory(VispDemodulatedPolcalData, arm_id=arm_id)
+
+            if load_dark:
+                manual_processing_run.run_task(load_dark_cal_task(arm_id=arm_id))
+            else:
+                manual_processing_run.run_task(dark_task)
+                manual_processing_run.run_task(save_dark_cal_task(arm_id=arm_id))
+
+            if load_gain:
+                manual_processing_run.run_task(load_gain_cal_task(arm_id=arm_id))
+            else:
+                manual_processing_run.run_task(gain_task)
+                manual_processing_run.run_task(save_gain_cal_task(arm_id=arm_id))
+
+            if load_geometric:
+                manual_processing_run.run_task(load_visp_geo_cal_task(arm_id=arm_id))
+            else:
+                manual_processing_run.run_task(geo_task)
+                manual_processing_run.run_task(save_visp_geo_cal_task(arm_id=arm_id))
+
+            if load_prepped:
+                manual_processing_run.run_task(load_prepped_cal_task(arm_id=arm_id))
+            else:
+                manual_processing_run.run_task(prep_data_task)
+                manual_processing_run.run_task(save_prepped_cal_task(arm_id=arm_id))
+
+            if load_pac_fit:
+                manual_processing_run.run_task(load_pac_fit_task(arm_id=arm_id))
+            else:
+                manual_processing_run.run_task(run_fit_task)
+                manual_processing_run.run_task(save_pac_fit_task(arm_id=arm_id))
+
+            if load_processed_demod:
+                manual_processing_run.run_task(load_processed_demod_cal_task(arm_id=arm_id))
+            else:
+                manual_processing_run.run_task(process_demod_task)
+                manual_processing_run.run_task(save_processed_demod_cal_task(arm_id=arm_id))
+
+            if load_dmpd:
+                manual_processing_run.run_task(load_dmpd_task(arm_id=arm_id))
+            else:
+                manual_processing_run.run_task(dmpd_task)
+                manual_processing_run.run_task(save_dmpd_task(arm_id=arm_id))
+
+        manual_processing_run.run_task(WriteTrendOutputs)
+        if any(
+            [
+                load_dark,
+                load_gain,
+                load_geometric,
+                load_prepped,
+                load_pac_fit,
+                load_processed_demod,
+                load_dmpd,
+            ]
+        ):
+            logger.info("NOT counting provenance records because some tasks were skipped")
+        else:
+            manual_processing_run.count_provenance()
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(
+        description="Run an end-to-end test of the DLNirsp DC Science pipeline"
+    )
+    parser.add_argument("scratch_path", help="Location to use as the DC 'scratch' disk")
+    parser.add_argument(
+        "-i",
+        "--run-id",
+        help="Which subdir to use. This will become the recipe run id",
+        type=int,
+        default=4,
+    )
+    parser.add_argument("--suffix", help="File suffix to treat as INPUT frames", default="FITS")
+    parser.add_argument(
+        "-T",
+        "--skip-translation",
+        help="Skip the translation of raw 122 l0 frames to 214 l0",
+        action="store_true",
+    )
+    parser.add_argument(
+        "-t", "--only-translate", help="Do ONLY the translation step", action="store_true"
+    )
+    parser.add_argument(
+        "-I",
+        "--load-input-parsing",
+        help="Load tags and constants on input files",
+        action="store_true",
+    )
+    parser.add_argument(
+        "-D",
+        "--load-dark",
+        help="Load dark calibration from previously saved run",
+        action="store_true",
+    )
+    parser.add_argument(
+        "-G",
+        "--load-gain",
+        help="Load gain calibration from previously saved run",
+        action="store_true",
+    )
+    parser.add_argument(
+        "-C",
+        "--load-geometric",
+        help="Load geometric calibration from previously saved run",
+        action="store_true",
+    )
+    parser.add_argument(
+        "-P",
+        "--load-prepped",
+        help="Load prepped fit data from previously saved run",
+        action="store_true",
+    )
+    parser.add_argument(
+        "-F",
+        "--load-pac-fit",
+        help="Load pac fit data from previously saved run",
+        action="store_true",
+    )
+    parser.add_argument(
+        "-E",
+        "--load-processed-demod",
+        help="Load processed demod data from previously saved run",
+        action="store_true",
+    )
+    parser.add_argument(
+        "-M",
+        "--load-dmpd",
+        help="Load DMPD data from previously saved run",
+        action="store_true",
+    )
+
+    args = parser.parse_args()
+    sys.exit(
+        main(
+            scratch_path=args.scratch_path,
+            suffix=args.suffix,
+            recipe_run_id=args.run_id,
+            skip_translation=args.skip_translation,
+            only_translate=args.only_translate,
+            load_input_parsing=args.load_input_parsing,
+            load_dark=args.load_dark,
+            load_gain=args.load_gain,
+            load_geometric=args.load_geometric,
+            load_prepped=args.load_prepped,
+            load_pac_fit=args.load_pac_fit,
+            load_processed_demod=args.load_processed_demod,
+            load_dmpd=args.load_dmpd,
+        )
+    )