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

dkist_processing_trend/tests/test_visp.py

@@ -0,0 +1,874 @@
+from functools import partial
+
+import numpy as np
+import pytest
+import scipy.ndimage as spnd
+from astropy.io import fits
+from astropy.modeling.functional_models import Voigt1D
+from dkist_processing_common._util.scratch import WorkflowFileSystem
+from dkist_processing_common.codecs.asdf import asdf_decoder
+from dkist_processing_common.codecs.asdf import asdf_encoder
+from dkist_processing_common.codecs.fits import fits_array_decoder
+from dkist_processing_common.codecs.fits import fits_array_encoder
+
+from dkist_processing_trend.models.instrument import InstrumentName
+from dkist_processing_trend.models.instrument_options import VispInstrumentOptions
+from dkist_processing_trend.models.parameters import TrendParameters
+from dkist_processing_trend.models.tags import TrendTag
+from dkist_processing_trend.tasks.visp.visp_dmpd import VispCalibrationCollection
+from dkist_processing_trend.tasks.visp.visp_dmpd import VispDemodulatedPolcalData
+from dkist_processing_trend.tasks.visp.visp_extract_beam import extract_visp_beam
+from dkist_processing_trend.tasks.visp.visp_geometric import VispGeometricCalibration
+from dkist_processing_trend.tasks.visp.visp_prep_fit_data import VispPrepareFitData
+from dkist_processing_trend.tasks.visp.visp_process_demod import VispProcessDemodulationMatrices
+from dkist_processing_trend.tests.conftest import ModulatedCSStepHeaders
+from dkist_processing_trend.tests.conftest import TrendInputDatasetParameterValues
+from dkist_processing_trend.tests.conftest import TrendTestingConstantsDatabase
+from dkist_processing_trend.tests.conftest import write_calibration_sequence_frames_to_task
+
+
+def test_extract_beam():
+    """
+    Given: A task based on `VispTrendTaskBase` and an array
+    When: Extracting the ViSP beams from the array
+    Then: The expected array is returned
+    """
+    small_beam_border = 10
+    array = np.empty((small_beam_border * 3, 4))
+    beam2_signal = np.arange(array.shape[0] - small_beam_border)
+    array[:small_beam_border, :] = 1.1111
+    array[small_beam_border:, :] = beam2_signal[::-1][:, None]
+
+    beam1 = extract_visp_beam(array=array, beam=1, beam_border=small_beam_border)
+    np.testing.assert_array_equal(beam1, 1.1111)
+
+    beam2 = extract_visp_beam(array=array, beam=2, beam_border=small_beam_border)
+    np.testing.assert_array_equal(np.mean(beam2, axis=1), beam2_signal)
+    np.testing.assert_array_equal(np.std(beam2, axis=1), 0.0)
+
+    with pytest.raises(
+        ValueError, match="ViSP only has beams 1 and 2. Can't extract data for beam '3'."
+    ):
+        extract_visp_beam(array=array, beam=3, beam_border=small_beam_border)
+
+
+@pytest.fixture
+def geometric_task_beam_shape() -> tuple[int, int]:
+    return (700, 900)
+
+
+@pytest.fixture
+def tiny_calibration_sequence() -> tuple[list, list, list, list, list]:
+    # We want some clear steps and at least one non-clear step
+    pol_status = [
+        "clear",
+        "Sapphire Polarizer",
+        "clear",
+    ]
+    pol_theta = [0.0, 120.0, 0.0]
+    ret_status = ["clear", "SiO2 SAR", "clear"]
+    ret_theta = [0.0, 0.0, 0.0]
+    dark_status = [
+        "FieldStop (5arcmin)",
+        "FieldStop (5arcmin)",
+        "FieldStop (5arcmin)",
+    ]
+
+    return pol_status, pol_theta, ret_status, ret_theta, dark_status
+
+
+@pytest.fixture
+def visp_geometric_task(
+    tmp_path,
+    recipe_run_id,
+    geometric_task_beam_shape,
+    link_constants_db,
+    assign_input_dataset_doc_to_task,
+):
+    link_constants_db(recipe_run_id, TrendTestingConstantsDatabase(INSTRUMENT=InstrumentName.visp))
+
+    with VispGeometricCalibration(
+        arm_id=1, recipe_run_id=recipe_run_id, workflow_name="name", workflow_version="version"
+    ) as task:
+        task.scratch = WorkflowFileSystem(recipe_run_id=recipe_run_id, scratch_base_path=tmp_path)
+        assign_input_dataset_doc_to_task(
+            task=task,
+            parameter_values=TrendInputDatasetParameterValues(
+                trend_visp_beam_border=geometric_task_beam_shape[0]
+            ),
+        )
+        task.parameters = TrendParameters(
+            scratch=task.scratch,
+            obs_ip_start_time=task.constants.earliest_ip_start_time,
+            instrument=InstrumentName.visp,
+        )
+
+        yield task
+        task._purge()
+
+
+def make_polcal_data_for_geo_task(
+    frame: ModulatedCSStepHeaders,
+    offset: tuple[float, float],
+    angles_deg: tuple[float, float],
+) -> np.ndarray:
+    full_shape = frame.array_shape[1:]
+    if (
+        frame.pol_status != "clear"
+        or frame.ret_status != "clear"
+        or frame.dark_status == "DarkShutter"
+    ):
+        # Put trash in the non-clear frames so we'll crash hard if they're used
+        return np.full(full_shape, np.nan)
+
+    beam_shape = (full_shape[0] // 2, full_shape[1])
+
+    # Spectral line centers
+    line_centers = np.array([0.2, 0.4, 0.5, 0.8]) * beam_shape[0]
+    offset_line_centers = line_centers - offset[0]
+
+    # Hairline centers
+    hairlines = [50, beam_shape[1] - 120]
+
+    beam1_data = np.full(beam_shape, 100.0)
+    # Approximate absorption lines with a set of Voigt profiles
+    beam1_voigt_model = sum(
+        [
+            Voigt1D(x_0=c, amplitude_L=2 - 0.8 * abs(i - 2), fwhm_G=4, fwhm_L=2)
+            for i, c in enumerate(line_centers[1:], start=1)
+        ],
+        Voigt1D(x_0=line_centers[0], amplitude_L=0.4),
+    )
+    beam1_voight = beam1_voigt_model(np.arange(beam_shape[0]))
+    # Normalize so the strongest line imparts a 50% reduction in flux
+    beam1_voight = beam1_voight / beam1_voight.max() * 2 + 1
+    beam1_data /= beam1_voight[:, None]
+    for h in hairlines:
+        # Add hairlines at 10% the continuum
+        beam1_data[:, h : h + 5] *= 0.1
+
+    beam1_data = spnd.rotate(beam1_data, angle=angles_deg[0], reshape=False, mode="nearest")
+
+    # Same for beam2, but with offsets for line centers and hairlines
+    beam2_data = np.full(beam_shape, 200.0)
+    beam2_voigt_model = sum(
+        [
+            Voigt1D(x_0=c, amplitude_L=2 - 0.8 * abs(i - 2), fwhm_G=4, fwhm_L=2)
+            for i, c in enumerate(offset_line_centers[1:], start=1)
+        ],
+        Voigt1D(x_0=offset_line_centers[0], amplitude_L=0.4),
+    )
+    beam2_voight = beam2_voigt_model(np.arange(beam_shape[0]))
+    beam2_voight = beam2_voight / beam2_voight.max() * 2 + 1
+    beam2_data /= beam2_voight[:, None]
+    for h in hairlines:
+        b2_h = h - int(offset[1])
+        beam2_data[:, b2_h : b2_h + 5] *= 0.1
+
+    beam2_data = spnd.rotate(beam2_data, angle=angles_deg[1], reshape=False, mode="nearest")
+
+    full_frame_array = np.vstack([beam1_data, beam2_data[::-1, :]])
+
+    return full_frame_array
+
+
+def test_visp_geometric_task(
+    visp_geometric_task,
+    geometric_task_beam_shape,
+    tiny_calibration_sequence,
+    mocker,
+    fake_gql_client,
+):
+    """
+    Given: A VispGeometricCalibration task with INPUT polcal data that contains clear measurements with a known offset between the beams
+    When: Running the task
+    Then: The correct output is written and the computed shift matches what is expected
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
+    )
+
+    task = visp_geometric_task
+    pol_status, pol_theta, ret_status, ret_theta, dark_status = tiny_calibration_sequence
+    full_shape = (geometric_task_beam_shape[0] * 2, geometric_task_beam_shape[1])
+    offset = (23.2, 6.28)
+    angles_deg = (0.005, -0.01)
+    data_func = partial(make_polcal_data_for_geo_task, offset=offset, angles_deg=angles_deg)
+
+    write_calibration_sequence_frames_to_task(
+        task=task,
+        instrument=InstrumentName.visp,
+        arm_id=1,
+        wavelength=633.6,
+        pol_status=pol_status,
+        pol_theta=pol_theta,
+        ret_status=ret_status,
+        ret_theta=ret_theta,
+        dark_status=dark_status,
+        array_shape=full_shape,
+        data_func=data_func,
+        tags=[TrendTag.input(), TrendTag.arm_id(1)],
+    )
+
+    task()
+
+    for b, beam in enumerate([1, 2]):
+        tags = [
+            TrendTag.intermediate(),
+            TrendTag.arm_id(1),
+            TrendTag.beam(beam),
+            TrendTag.task_visp_geometric_calibration(),
+        ]
+        assert task.count(tags) == 1
+        geo_dict = next(task.read(tags=tags, decoder=asdf_decoder))
+
+        # Check that the shift is correct to within 1/2 pixel, which is acceptable for Trend stuff
+        computed_shift = geo_dict["shift"]
+        expected_offset = (0, 0) if beam == 1 else -1 * np.array(offset)
+        np.testing.assert_array_less(np.abs(computed_shift - expected_offset), 0.5)
+
+        # Check that the angles are correct to within 1%, which I mostly just made up based on real data
+        computed_angle = geo_dict["angle_rad"]
+        np.testing.assert_array_less(
+            np.abs(np.rad2deg(computed_angle) - angles_deg[b]) / angles_deg[b], 0.01
+        )
+
+
+@pytest.fixture
+def single_visp_instrument_options_list() -> list[dict]:
+    return [
+        {
+            "name": "visp_options",
+            "num_spatial_px": 10,
+            "num_spectral_px": 1,
+            "spatial_smoothing_fit_order": 3,
+        }
+    ]
+
+
+@pytest.fixture
+def small_beam_border() -> int:
+    return 100
+
+
+@pytest.fixture
+def visp_prep_fit_data_task(
+    recipe_run_id,
+    tmp_path,
+    link_constants_db,
+    assign_input_dataset_doc_to_task,
+    single_visp_instrument_options_list,
+    small_num_cs_steps,
+    small_beam_border,
+):
+    link_constants_db(
+        recipe_run_id,
+        TrendTestingConstantsDatabase(
+            INSTRUMENT=InstrumentName.visp,
+            ARM_ID_LIST=[1, 3],
+            NUM_CS_STEPS=small_num_cs_steps,
+        ),
+    )
+    with VispPrepareFitData(
+        arm_id=3, recipe_run_id=recipe_run_id, workflow_name="name", workflow_version="version"
+    ) as task:
+        task.scratch = WorkflowFileSystem(recipe_run_id=recipe_run_id, scratch_base_path=tmp_path)
+        assign_input_dataset_doc_to_task(
+            task,
+            TrendInputDatasetParameterValues(
+                trend_instrument_processing_options_visp=single_visp_instrument_options_list,
+                trend_visp_beam_border=small_beam_border,
+            ),
+        )
+        task.parameters = TrendParameters(
+            scratch=task.scratch,
+            obs_ip_start_time=task.constants.earliest_ip_start_time,
+            instrument=task.constants.instrument,
+        )
+
+        yield task
+        task._purge()
+
+
+def test_visp_prepare_fit_data(
+    visp_prep_fit_data_task,
+    single_visp_instrument_options_list,
+    small_calibration_sequence,
+    small_beam_border,
+    mocker,
+    fake_gql_client,
+):
+    """
+    Given: A `VispPrepareFitData` task and INPUT and INTERMEDIATE frames needed for the task to run
+    When: Running the task
+    Then: The dang thing runs and produces the outputs we expect
+
+    NOTE: This is just a test of "glue" of the subclass. See other tests for checks that the visp-specific methods function
+    correctly.
+    """
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
+    )
+    shape = (small_beam_border * 2, small_beam_border)
+    inst_opt = single_visp_instrument_options_list[0]
+    task = visp_prep_fit_data_task
+
+    pol_status, pol_theta, ret_status, ret_theta, dark_status = small_calibration_sequence
+
+    write_calibration_sequence_frames_to_task(
+        task=task,
+        instrument=InstrumentName.visp,
+        arm_id=3,
+        wavelength=360.2,
+        pol_status=pol_status,
+        pol_theta=pol_theta,
+        ret_status=ret_status,
+        ret_theta=ret_theta,
+        dark_status=dark_status,
+        array_shape=shape,
+        tags=[TrendTag.arm_id(3)],
+    )
+
+    # INTERMEIDATE dark and gain frames
+    task.write(
+        data=np.full(shape, 0.0),
+        tags=[
+            TrendTag.intermediate(),
+            TrendTag.frame(),
+            TrendTag.arm_id(3),
+            TrendTag.task_dark(),
+        ],
+        encoder=fits_array_encoder,
+    )
+    task.write(
+        data=np.full(shape, 1.0),
+        tags=[
+            TrendTag.intermediate(),
+            TrendTag.frame(),
+            TrendTag.arm_id(3),
+            TrendTag.task_gain(),
+        ],
+        encoder=fits_array_encoder,
+    )
+
+    task()
+
+    for beam in range(1, 2):
+        global_tags = [
+            TrendTag.intermediate(),
+            TrendTag.arm_id(3),
+            TrendTag.beam(beam),
+            TrendTag.instrument_processing_options(inst_opt["name"]),
+            TrendTag.task_global_pac_input(),
+        ]
+        local_tags = [
+            TrendTag.intermediate(),
+            TrendTag.arm_id(3),
+            TrendTag.beam(beam),
+            TrendTag.instrument_processing_options(inst_opt["name"]),
+            TrendTag.task_local_pac_input(),
+        ]
+        assert task.count(global_tags) == 1
+        assert task.count(local_tags) == 1
+
+
+def test_apply_visp_global_instrument_options(
+    visp_prep_fit_data_task, single_visp_instrument_options_list
+):
+    """
+    Given: A `VispPrepareFitData` task and an array
+    When: Converting the array into a "global" data point
+    Then: The correct output is returned
+    """
+    task = visp_prep_fit_data_task
+    inst_opt_object = VispInstrumentOptions.model_validate(single_visp_instrument_options_list[0])
+    array = np.full((100, 200), 100.0)
+    expected_median = np.median(array)
+
+    # Add some hairlines as sharp features in the spectral (0th) dimension
+    array[30, :] = 9999.9
+    array[70, :] = 121212.1
+
+    # And NaN's just b/c
+    array[34, 102] = np.nan
+
+    global_data = task.apply_global_instrument_options(
+        array=array, instrument_options=inst_opt_object
+    )
+    assert global_data.shape == (1, 1)
+    np.testing.assert_array_equal(global_data, expected_median)
+
+
+def test_apply_visp_local_instrument_options(
+    visp_prep_fit_data_task, single_visp_instrument_options_list
+):
+    """
+    Given: A `VispPrepareFitData` task and an array
+    When: Converting the array into a "local" data point
+    Then: The correct output is returned
+    """
+    task = visp_prep_fit_data_task
+    inst_opt_object = VispInstrumentOptions.model_validate(single_visp_instrument_options_list[0])
+    array = np.full((100, 200), 100.0)
+    expected_median = np.median(array)
+
+    # Add high-frequency signal in the spectral dimension that will disappear with a median
+    # Signal of [-1, 0, 1, ...] will have a median of 0 => no change to overall median
+    array += (1 - np.arange(100) % 3)[:, None]
+
+    # Add some spikes in the spatial dimension that will get smoothed out
+    array[:, 130] *= 1.2
+    array[:, 180] *= 1.5
+
+    # Add some hairlines as sharp features in the spectral dimension
+    array[30, :] = 9999.9
+    array[70, :] = 121212.1
+
+    # And NaN's just b/c
+    array[34, 102] = np.nan
+
+    local_data = task.apply_local_instrument_options(
+        array=array, instrument_options=inst_opt_object
+    )
+    assert local_data.shape == (1, inst_opt_object.num_spatial_px)
+    np.testing.assert_array_equal(local_data, expected_median)
+
+
+@pytest.fixture
+def multiple_visp_instrument_options_list() -> list[dict]:
+    return [
+        {
+            "name": "visp_option1",
+            "num_spatial_px": 10,
+            "num_spectral_px": 1,
+            "spatial_smoothing_fit_order": 3,
+        },
+        {
+            "name": "visp_option2",
+            "num_spatial_px": 5,
+            "num_spectral_px": 1,
+            "spatial_smoothing_fit_order": 2,
+        },
+    ]
+
+
+@pytest.fixture
+def multiple_pac_fit_options_list() -> list[dict]:
+    return [
+        {
+            "name": "option1",
+            "remove_I_trend": True,
+            "fit_mode_name": "fit_mode1",
+            "init_set_name": "init_set1",
+        },
+        {
+            "name": "option2",
+            "remove_I_trend": False,
+            "fit_mode_name": "fit_mode1",
+            "init_set_name": "init_set2",
+        },
+    ]
+
+
+@pytest.fixture
+def visp_process_demod_task(
+    tmp_path,
+    recipe_run_id,
+    link_constants_db,
+    assign_input_dataset_doc_to_task,
+    small_beam_border,
+    multiple_visp_instrument_options_list,
+    multiple_pac_fit_options_list,
+):
+
+    link_constants_db(
+        recipe_run_id,
+        TrendTestingConstantsDatabase(
+            INSTRUMENT=InstrumentName.visp,
+            ARM_ID_LIST=[2],
+        ),
+    )
+    with VispProcessDemodulationMatrices(
+        arm_id=2, recipe_run_id=recipe_run_id, workflow_name="name", workflow_version="version"
+    ) as task:
+        task.scratch = WorkflowFileSystem(recipe_run_id=recipe_run_id, scratch_base_path=tmp_path)
+        assign_input_dataset_doc_to_task(
+            task,
+            TrendInputDatasetParameterValues(
+                trend_fit_options_list=multiple_pac_fit_options_list,
+                trend_instrument_processing_options_visp=multiple_visp_instrument_options_list,
+                trend_visp_beam_border=small_beam_border,
+            ),
+        )
+        task.parameters = TrendParameters(
+            scratch=task.scratch,
+            obs_ip_start_time=task.constants.earliest_ip_start_time,
+            instrument=task.constants.instrument,
+        )
+
+        yield task
+        task._purge()
+
+
+def test_visp_process_demod_task(
+    visp_process_demod_task,
+    multiple_visp_instrument_options_list,
+    multiple_pac_fit_options_list,
+    small_beam_border,
+    mocker,
+    fake_gql_client,
+):
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
+    )
+    """
+    Given: A `VispProcessDemodulationMatrices` task and some best-fit demodulation matrices to mimic the output of running the PAC fitter
+    When: Running the task
+    Then: The task runs, the correct number of outputs is produced, and all processed demodulation matrices have the correct form
+    """
+
+    task = visp_process_demod_task
+
+    full_beam_shape = (small_beam_border, 102)
+    full_frame_shape = (full_beam_shape[0] * 2, full_beam_shape[1])
+    num_mod = 3
+    task.write(
+        data=np.empty(full_frame_shape),
+        tags=[
+            TrendTag.intermediate(),
+            TrendTag.frame(),
+            TrendTag.arm_id(task.arm_id),
+            TrendTag.task_gain(),
+        ],
+        encoder=fits_array_encoder,
+    )
+
+    for fit_options in multiple_pac_fit_options_list:
+        for instrument_options in multiple_visp_instrument_options_list:
+            num_spat = instrument_options["num_spatial_px"]
+            num_wave = instrument_options["num_spectral_px"]
+            for beam in [1, 2]:
+                # Add a linear signal so there's *something* for the fit to grab onto
+                # We divide by the linear signal instead of multiply because this matrix will be inverted for the actual fit.
+                raw_demod = (
+                    np.ones((num_wave, num_spat, 4, num_mod))
+                    / (np.arange(1, num_spat + 1) * 1000)[None, :, None, None]
+                )
+
+                task.write(
+                    data=raw_demod,
+                    tags=[
+                        TrendTag.intermediate(),
+                        TrendTag.arm_id(task.arm_id),
+                        TrendTag.beam(beam),
+                        TrendTag.pac_fit_options(fit_options["name"]),
+                        TrendTag.instrument_processing_options(instrument_options["name"]),
+                        TrendTag.task_best_fit_demodulation_matrices(),
+                    ],
+                    encoder=fits_array_encoder,
+                )
+
+    task()
+
+    for fit_options in multiple_pac_fit_options_list:
+        for instrument_options in multiple_visp_instrument_options_list:
+            for beam in [1, 2]:
+                tags = [
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(task.arm_id),
+                    TrendTag.beam(beam),
+                    TrendTag.pac_fit_options(fit_options["name"]),
+                    TrendTag.instrument_processing_options(instrument_options["name"]),
+                    TrendTag.task_processed_demodulation_matrices(),
+                ]
+                assert task.count(tags) == 1
+                processed_demod = next(task.read(tags=tags, decoder=fits_array_decoder))
+                assert processed_demod.shape == full_beam_shape + (4, num_mod)
+                # Test that the upsampling in the spectral dimension didn't do anything weird
+                np.testing.assert_almost_equal(np.std(processed_demod, axis=0), 0.0)
+
+
+def test_smooth_demod_matrices(visp_process_demod_task):
+    """
+    Given: A demodulation matrix with known form in the spatial direction
+    When: Smoothing that demodulation matri
+    Then: The result has the right shape and the correct polynomial form
+    """
+    task = visp_process_demod_task
+    order = 5
+    coeffs = np.linspace(1, 37, order + 1)
+    full_num_spat = 200
+    binned_num_spat = 50
+    full_signal = np.poly1d(coeffs)(np.arange(full_num_spat))
+    idx = np.linspace(0, full_num_spat, binned_num_spat, endpoint=False, dtype=int)
+    binned_signal = full_signal[idx]
+
+    # Divide by the signal because we fit the polynomial in the inverse
+    binned_demod = 1.0 / binned_signal[None, :, None, None]
+
+    smoothed_demod = task.smooth_demod_matrices(binned_demod, order, full_num_spat)
+
+    assert smoothed_demod.shape == (1, full_num_spat, 1, 1)
+
+    # Like above, fit the inverse of the signal
+    fit_coeffs = np.polyfit(np.arange(full_num_spat), 1.0 / smoothed_demod[0, :, 0, 0], order)
+    np.testing.assert_allclose(coeffs, fit_coeffs, rtol=1e-5)
+
+
+@pytest.fixture
+def visp_dmpd_task(
+    tmp_path,
+    recipe_run_id,
+    link_constants_db,
+    assign_input_dataset_doc_to_task,
+    small_beam_border,
+    multiple_visp_instrument_options_list,
+    multiple_pac_fit_options_list,
+):
+
+    link_constants_db(
+        recipe_run_id,
+        TrendTestingConstantsDatabase(
+            INSTRUMENT=InstrumentName.visp,
+            ARM_ID_LIST=[2],
+            NUM_MODSTATES=3,
+        ),
+    )
+    with VispDemodulatedPolcalData(
+        arm_id=2, recipe_run_id=recipe_run_id, workflow_name="name", workflow_version="version"
+    ) as task:
+        task.scratch = WorkflowFileSystem(recipe_run_id=recipe_run_id, scratch_base_path=tmp_path)
+        assign_input_dataset_doc_to_task(
+            task,
+            TrendInputDatasetParameterValues(
+                trend_fit_options_list=multiple_pac_fit_options_list,
+                trend_instrument_processing_options_visp=multiple_visp_instrument_options_list,
+                trend_visp_beam_border=small_beam_border,
+            ),
+        )
+        task.parameters = TrendParameters(
+            scratch=task.scratch,
+            obs_ip_start_time=task.constants.earliest_ip_start_time,
+            instrument=task.constants.instrument,
+        )
+
+        yield task
+        task._purge()
+
+
+def test_visp_dmpd_task(
+    visp_dmpd_task,
+    small_calibration_sequence,
+    small_num_cs_steps,
+    small_beam_border,
+    multiple_pac_fit_options_list,
+    multiple_visp_instrument_options_list,
+    mocker,
+    fake_gql_client,
+):
+    mocker.patch(
+        "dkist_processing_common.tasks.mixin.metadata_store.GraphQLClient", new=fake_gql_client
+    )
+    shape = (small_beam_border * 2, small_beam_border)
+    task = visp_dmpd_task
+
+    pol_status, pol_theta, ret_status, ret_theta, dark_status = small_calibration_sequence
+    num_mod = task.constants.num_modstates
+    write_calibration_sequence_frames_to_task(
+        task=task,
+        instrument=InstrumentName.visp,
+        num_modstates=num_mod,
+        arm_id=2,
+        wavelength=360.2,
+        pol_status=pol_status,
+        pol_theta=pol_theta,
+        ret_status=ret_status,
+        ret_theta=ret_theta,
+        dark_status=dark_status,
+        array_shape=shape,
+        tags=[TrendTag.input(), TrendTag.arm_id(2)],
+    )
+
+    # INTERMEIDATE dark, gain, geo, and processed demod frames
+    task.write(
+        data=np.full(shape, 0.0),
+        tags=[
+            TrendTag.intermediate(),
+            TrendTag.frame(),
+            TrendTag.arm_id(2),
+            TrendTag.task_dark(),
+        ],
+        encoder=fits_array_encoder,
+    )
+    task.write(
+        data=np.full(shape, 1.0),
+        tags=[
+            TrendTag.intermediate(),
+            TrendTag.frame(),
+            TrendTag.arm_id(2),
+            TrendTag.task_gain(),
+        ],
+        encoder=fits_array_encoder,
+    )
+    task.write(
+        data={"angle_rad": 0.004, "shift": np.array([0.0, 0.0])},
+        tags=[
+            TrendTag.intermediate(),
+            TrendTag.arm_id(2),
+            TrendTag.beam(1),
+            TrendTag.task_visp_geometric_calibration(),
+        ],
+        encoder=asdf_encoder,
+    )
+    task.write(
+        data={"angle_rad": -0.004, "shift": np.array([5.0, 12.3])},
+        tags=[
+            TrendTag.intermediate(),
+            TrendTag.arm_id(2),
+            TrendTag.beam(2),
+            TrendTag.task_visp_geometric_calibration(),
+        ],
+        encoder=asdf_encoder,
+    )
+    demod = np.random.random((shape[0] // 2, shape[1], 4, num_mod))
+    for fit_options in multiple_pac_fit_options_list:
+        for inst_options in multiple_visp_instrument_options_list:
+            task.write(
+                data=demod,
+                tags=[
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(2),
+                    TrendTag.pac_fit_options(fit_options["name"]),
+                    TrendTag.instrument_processing_options(inst_options["name"]),
+                    TrendTag.beam(1),
+                    TrendTag.task_processed_demodulation_matrices(),
+                ],
+                encoder=fits_array_encoder,
+            )
+            task.write(
+                data=demod,
+                tags=[
+                    TrendTag.intermediate(),
+                    TrendTag.arm_id(2),
+                    TrendTag.pac_fit_options(fit_options["name"]),
+                    TrendTag.instrument_processing_options(inst_options["name"]),
+                    TrendTag.beam(2),
+                    TrendTag.task_processed_demodulation_matrices(),
+                ],
+                encoder=fits_array_encoder,
+            )
+
+    task()
+
+    # -5 and -13 come from the shift values written in the geo intermediate above
+    expected_beam_shape = (shape[0] // 2 - 5, shape[1] - 13)
+    expected_full_shape = (1, small_num_cs_steps, 4, *expected_beam_shape)
+    for fit_options in multiple_pac_fit_options_list:
+        for inst_options in multiple_visp_instrument_options_list:
+            tags = [
+                TrendTag.intermediate(),
+                TrendTag.arm_id(2),
+                TrendTag.pac_fit_options(fit_options["name"]),
+                TrendTag.instrument_processing_options(inst_options["name"]),
+            ]
+            beam1_tags = tags + [
+                TrendTag.beam(1),
+                TrendTag.task_single_beam_demodulated_polcal_data(),
+            ]
+            beam2_tags = tags + [
+                TrendTag.beam(2),
+                TrendTag.task_single_beam_demodulated_polcal_data(),
+            ]
+            combined_tags = tags + [TrendTag.task_demodulated_polcal_data()]
+
+            assert task.count(tags=beam1_tags) == 1
+            assert task.count(tags=beam2_tags) == 1
+            assert task.count(tags=combined_tags) == 1
+
+            beam1_dmdp = next(task.read(tags=beam1_tags, decoder=fits_array_decoder))
+            beam2_dmdp = next(task.read(tags=beam2_tags, decoder=fits_array_decoder))
+            combined_dmdp = next(task.read(tags=combined_tags, decoder=fits_array_decoder))
+
+            assert beam1_dmdp.shape == expected_full_shape
+            # Testing >f4 instead of np.float32 b/c PrimaryHDUs change endianess
+            assert beam1_dmdp.dtype == np.dtype(">f4")
+            assert beam2_dmdp.shape == expected_full_shape
+            assert beam2_dmdp.dtype == np.dtype(">f4")
+            assert combined_dmdp.shape == expected_full_shape
+            assert combined_dmdp.dtype == np.dtype(">f4")
+
+
+@pytest.mark.parametrize(
+    "shifts, expected",
+    # Shifts have shape (num_beams, 2)
+    # So the inner-most lists below (e.g., [5.0, 6.0]) correspond to [x_shift, y_shit]
+    [
+        (
+            np.array(
+                [
+                    [0.0, 0.0],  # Beam 1
+                    [11.0, 10.0],  # Beam 2
+                ]
+            ),
+            [slice(0, -11, None), slice(0, -10, None)],
+        ),
+        (
+            np.array(
+                [
+                    [0.0, 0.0],  # Beam 1
+                    [-11.0, -10.0],  # Beam 2
+                ]
+            ),
+            [slice(11, None, None), slice(10, None, None)],
+        ),
+        (
+            np.array(
+                [  # This case would never happen, but is useful as a check of the logic
+                    [10.0, -2.0],  # Beam 1
+                    [-11.0, 10.0],  # Beam 2
+                ]
+            ),
+            [slice(11, -10, None), slice(2, -10, None)],
+        ),
+    ],
+    ids=["All positive", "All negative", "Positive and negative"],
+)
+def test_beam_overlap_slice(shifts, expected):
+    """
+    Given: A `VispCalibrationCollection` object with populated shifts
+    When: Computing the overlapping beam slices
+    Then: The correct values are returned
+    """
+    shift_dict = {i: v for i, v in enumerate(shifts)}
+    calibrations = VispCalibrationCollection(
+        dark=np.empty(2), gain=np.empty(2), angle=dict(), shift=shift_dict, demod_matrices=dict()
+    )
+    x_slice, y_slice = calibrations.beams_overlap_slice
+
+    assert x_slice == expected[0]
+    assert y_slice == expected[1]
+
+
+def test_combine_beams(visp_dmpd_task):
+    """
+    Given: A `VispDemodulatedPolcalData` task and set of calibrated array data
+    When: Combining the two beams
+    Then: The correct result is returned
+    """
+    task = visp_dmpd_task
+    shape = (10, 10, 4)
+    beam1 = np.ones(shape) + np.arange(4)[None, None, :]
+    beam2 = np.ones(shape) + np.arange(4)[::-1][None, None, :]
+    result = task.combine_beams(
+        beam1_data=beam1,
+        beam2_data=beam2,
+    )
+
+    x = np.arange(1, 5)
+    expected_I = np.ones((10, 10)) * 2.5
+    expected_Q = np.ones((10, 10)) * (x[1] / x[0] + x[-2] / x[-1]) / 2.0 * 2.5
+    expected_U = np.ones((10, 10)) * (x[2] / x[0] + x[-3] / x[-1]) / 2.0 * 2.5
+    expected_V = np.ones((10, 10)) * (x[3] / x[0] + x[-4] / x[-1]) / 2.0 * 2.5
+    expected = np.dstack([expected_I, expected_Q, expected_U, expected_V])
+
+    np.testing.assert_array_equal(result, expected)