dkist-processing-visp 4.0.0__py3-none-any.whl → 5.0.0__py3-none-any.whl

dkist_processing_visp/tests/test_geometric.py

@@ -358,3 +358,43 @@ def test_basic_corrections(geometric_calibration_task, assign_input_dataset_doc_
             lamp_array = task.basic_corrected_lamp_data(beam=beam, modstate=modstate)
             np.testing.assert_equal(expected, solar_array)
             np.testing.assert_equal(expected, lamp_array)
+
+
+def test_line_zones(geometric_calibration_task):
+    """
+    Given: A spectrum with some absorption lines
+    When: Computing zones around the lines
+    Then: Correct results are returned
+    """
+
+    # NOTE that it does not test for removal of overlapping regions
+    def gaussian(x, amp, mu, sig):
+        return amp * np.exp(-np.power(x - mu, 2.0) / (2 * np.power(sig, 2.0)))
+
+    spec = np.ones(1000) * 100
+    x = np.arange(1000.0)
+    expected = []
+    for m, s in zip([100.0, 300.0, 700], [10.0, 20.0, 5.0]):
+        spec -= gaussian(x, 40, m, s)
+        hwhm = s * 2.355 / 2
+        expected.append((np.floor(m - hwhm).astype(int), np.ceil(m + hwhm).astype(int)))
+
+    task = geometric_calibration_task[0]
+
+    zones = task.compute_line_zones(spec[:, None], bg_order=0, rel_height=0.5)
+    assert zones == expected
+
+
+def test_identify_overlapping_zones(geometric_calibration_task):
+    """
+    Given: A list of zone borders that contain overlapping zones
+    When: Identifying zones that overlap
+    Then: The smaller of the overlapping zones are identified for removal
+    """
+    rips = np.array([100, 110, 220, 200])
+    lips = np.array([150, 120, 230, 250])
+
+    task = geometric_calibration_task[0]
+
+    idx_to_remove = task.identify_overlapping_zones(rips, lips)
+    assert idx_to_remove == [1, 2]

dkist_processing_visp/tests/test_instrument_polarization.py

@@ -255,7 +255,8 @@ def test_instrument_polarization_calibration_task(
         task=task, num_modstates=num_modstates, data_shape=intermediate_shape
     )
     write_dummy_intermediate_solar_cals_to_task(
-        task=task, data_shape=intermediate_shape, num_modstates=num_modstates
+        task=task,
+        data_shape=intermediate_shape,
     )
     write_input_polcals_to_task(
         task=task,

dkist_processing_visp/tests/test_lamp.py

@@ -23,9 +23,10 @@ def make_lamp_array_data(
     frame: VispHeadersInputLampGainFrames, dark_signal: float, beam_border: int
 ):
     num_raw_frames_per_fpa = frame.header()["CAM__014"]
+    modstate = frame.current_modstate("")  # Weird signature due to @key_function
     data = np.zeros(frame.array_shape)
-    data[0, :beam_border, :] = (1.1 + dark_signal) * num_raw_frames_per_fpa
-    data[0, beam_border:, :] = (1.2 + dark_signal) * num_raw_frames_per_fpa
+    data[0, :beam_border, :] = (1.0 + 0.1 * modstate + dark_signal) * num_raw_frames_per_fpa
+    data[0, beam_border:, :] = (2.0 + 0.1 * modstate + dark_signal) * num_raw_frames_per_fpa
 
     return data
 
@@ -132,27 +133,21 @@ def test_lamp_calibration_task(
         VispTag.task_lamp_gain(),
         VispTag.intermediate(),
     ]
-    assert len(list(task.read(tags=tags))) == num_modstates * 2  # 2 from beams
-
-    for i in range(num_modstates):
-        for beam in [1, 2]:
-            tags = [
-                VispTag.task_lamp_gain(),
-                VispTag.intermediate(),
-                VispTag.modstate(i + 1),
-                VispTag.beam(beam),
-            ]
-            files = list(task.read(tags=tags))
-            assert len(files) == 1
-            hdu = fits.open(files[0])[0]
-            np.testing.assert_allclose(hdu.data, np.ones((10, 10)) * (1 + (0.1 * beam)))
+    assert len(list(task.read(tags=tags))) == 2  # One per beam
 
-    tags = [
-        VispTag.task_lamp_gain(),
-        VispTag.intermediate(),
-    ]
-    for filepath in task.read(tags=tags):
-        assert filepath.exists()
+    for beam in [1, 2]:
+        tags = [
+            VispTag.task_lamp_gain(),
+            VispTag.intermediate(),
+            VispTag.beam(beam),
+        ]
+        files = list(task.read(tags=tags))
+        assert len(files) == 1
+
+        expected_signal = beam + np.mean(np.arange(1, num_modstates + 1)) * 0.1
+
+        hdu = fits.open(files[0])[0]
+        np.testing.assert_allclose(hdu.data, np.ones(intermediate_shape) * expected_signal)
 
     quality_files = task.read(tags=[Tag.quality("TASK_TYPES")])
     for file in quality_files:

dkist_processing_visp/tests/test_parameters.py

@@ -1,6 +1,7 @@
 from dataclasses import asdict
 from dataclasses import dataclass
 
+import astropy.units as u
 import numpy as np
 import pytest
 from hypothesis import HealthCheck
@@ -8,6 +9,7 @@ from hypothesis import example
 from hypothesis import given
 from hypothesis import settings
 from hypothesis import strategies as st
+from pydantic import BaseModel
 
 from dkist_processing_visp.models.parameters import VispParameters
 from dkist_processing_visp.models.parameters import VispParsingParameters
@@ -15,10 +17,27 @@ from dkist_processing_visp.tasks.visp_base import VispTaskBase
 from dkist_processing_visp.tests.conftest import VispConstantsDb
 from dkist_processing_visp.tests.conftest import VispInputDatasetParameterValues
 
-# The property names of all parameters on `VispParsingParameters`
-PARSE_PARAMETER_NAMES = [
-    k for k, v in vars(VispParsingParameters).items() if isinstance(v, property)
-]
+
+@pytest.fixture(scope="session")
+def parse_parameter_names() -> list[str]:
+    # The property names of all parameters on `VispParsingParameters`
+    return [k for k, v in vars(VispParsingParameters).items() if isinstance(v, property)]
+
+
+@pytest.fixture(scope="session")
+def arm_parameter_names() -> list[str]:
+    return [
+        "wavecal_camera_lens_parameters",
+    ]
+
+
+@pytest.fixture(scope="session")
+def unit_parameter_names_and_units() -> dict[str, u.Unit | list[u.Unit]]:
+    return {
+        "solar_vignette_crval_bounds_px": u.pix,
+        "wavecal_camera_lens_parameters": [u.m, u.m / u.nm, u.m / u.nm**2],
+        "wavecal_pixel_pitch_micron_per_pix": u.um / u.pix,
+    }
 
 
 @pytest.fixture(scope="function")
@@ -27,6 +46,7 @@ def basic_science_task_with_parameter_mixin(
     assign_input_dataset_doc_to_task,
     init_visp_constants_db,
     testing_obs_ip_start_time,
+    arm_id,
 ):
     def make_task(
         parameters_part: dataclass,
@@ -48,8 +68,9 @@ def basic_science_task_with_parameter_mixin(
                 parameters_part,
                 parameter_class=parameter_class,
                 obs_ip_start_time=obs_ip_start_time,
+                arm_id=arm_id,
             )
-            yield task, parameters_part
+            yield task, asdict(parameters_part)
         except:
             raise
         finally:
@@ -58,7 +79,14 @@ def basic_science_task_with_parameter_mixin(
     return make_task
 
 
-def test_non_wave_parameters(basic_science_task_with_parameter_mixin):
+@pytest.mark.parametrize("arm_id", [pytest.param("1"), pytest.param("2"), pytest.param("3")])
+def test_non_wave_parameters(
+    basic_science_task_with_parameter_mixin,
+    parse_parameter_names,
+    arm_parameter_names,
+    unit_parameter_names_and_units,
+    arm_id,
+):
     """
     Given: A Science task with the parameter mixin
     When: Accessing properties for parameters that do not depend on wavelength
@@ -68,15 +96,41 @@ def test_non_wave_parameters(basic_science_task_with_parameter_mixin):
         basic_science_task_with_parameter_mixin(VispInputDatasetParameterValues())
     )
     task_param_attr = task.parameters
-    for pn, pv in asdict(expected).items():
-        property_name = pn.removeprefix("visp_")
+    parameter_properties = [k for k, v in vars(VispParameters).items() if isinstance(v, property)]
+    for parameter_name in parameter_properties:
+        pn = f"visp_{parameter_name}"
+        if parameter_name in arm_parameter_names:
+            pn = f"{pn}_{arm_id}"
+        pv = expected[pn]
+        is_wavelength_param = isinstance(pv, dict) and "wavelength" in pv
         if (
-            type(pv) is not dict and property_name not in PARSE_PARAMETER_NAMES
-        ):  # Don't test wavelength dependent parameters
-            assert getattr(task_param_attr, property_name) == pv
+            parameter_name not in parse_parameter_names
+            and not is_wavelength_param
+            and parameter_name != "solar_vignette_wavecal_fit_kwargs"
+        ):
+            param_obj_value = getattr(task_param_attr, parameter_name)
+            if isinstance(pv, tuple):
+                pv = list(pv)
+
+            if parameter_name in unit_parameter_names_and_units:
+                expected_units = unit_parameter_names_and_units[parameter_name]
+                if not isinstance(param_obj_value, list):
+                    param_obj_value = [param_obj_value]
+                    pv = [pv]
+                    expected_units = [expected_units]
 
+                assert all([param_obj_value[i].value == pv[i] for i in range(len(pv))])
+                assert all(
+                    [param_obj_value[i].unit == expected_units[i] for i in range(len(pv))]
+                ), f"Units {[v.unit for v in param_obj_value]} does not match expected {expected_units}"
+            elif isinstance(param_obj_value, BaseModel):
+                assert param_obj_value.model_dump() == pv
+            else:
+                assert getattr(task_param_attr, parameter_name) == pv
 
-def test_parse_parameters(basic_science_task_with_parameter_mixin):
+
+@pytest.mark.parametrize("arm_id", ["1"])
+def test_parse_parameters(basic_science_task_with_parameter_mixin, parse_parameter_names):
     """
     Given: A Science task with Parsing parameters
     When: Accessing properties for Parse parameters
@@ -90,16 +144,19 @@ def test_parse_parameters(basic_science_task_with_parameter_mixin):
         )
     )
     task_param_attr = task.parameters
-    for pn, pv in asdict(expected).items():
+    for pn, pv in expected.items():
         property_name = pn.removeprefix("visp_")
-        if property_name in PARSE_PARAMETER_NAMES and type(pv) is not dict:
+        if property_name in parse_parameter_names and type(pv) is not dict:
             assert getattr(task_param_attr, property_name) == pv
 
 
 @given(wave=st.floats(min_value=500.0, max_value=2000.0))
 @settings(suppress_health_check=[HealthCheck.function_scoped_fixture])
 @example(wave=492.5)
-def test_wave_parameters(basic_science_task_with_parameter_mixin, wave):
+@pytest.mark.parametrize("arm_id", ["1"])
+def test_wave_parameters(
+    basic_science_task_with_parameter_mixin, parse_parameter_names, arm_parameter_names, wave
+):
     """
     Given: A Science task with the paramter mixin
     When: Accessing properties for parameters that depend on wavelength
@@ -110,10 +167,55 @@ def test_wave_parameters(basic_science_task_with_parameter_mixin, wave):
     )
     task_param_attr = task.parameters
     task_param_attr._wavelength = wave
-    pwaves = np.array(expected.visp_solar_zone_normalization_percentile.wavelength)
+    pwaves = np.array(expected["visp_geo_zone_normalization_percentile"]["wavelength"])
     midpoints = 0.5 * (pwaves[1:] + pwaves[:-1])
     idx = np.sum(midpoints < wave)
-    for pn, pv in asdict(expected).items():
+    for pn, pv in expected.items():
         property_name = pn.removeprefix("visp_")
-        if type(pv) is dict and property_name not in PARSE_PARAMETER_NAMES:
+        is_wavelength_param = isinstance(pv, dict) and "wavelength" in pv
+        if is_wavelength_param and property_name not in parse_parameter_names + arm_parameter_names:
             assert getattr(task_param_attr, property_name) == pv["values"][idx]
+
+
+class AnyInt:
+    pass
+
+
+@pytest.mark.parametrize("arm_id", [pytest.param("1")])
+@pytest.mark.parametrize(
+    "db_value, expected",
+    [
+        pytest.param({"method": "nelder"}, {"method": "nelder"}, id="non_rng_method"),
+        pytest.param(
+            {"method": "basinhopping"}, {"method": "basinhopping", "rng": AnyInt}, id="random_rng"
+        ),
+        pytest.param(
+            {"method": "differential_evolution", "rng": 6.28},
+            {"method": "differential_evolution", "rng": 6.28},
+            id="override_rng",
+        ),
+        pytest.param(dict(), dict(), id="no_kwargs"),
+    ],
+)
+def test_fit_kwarg_parameters(
+    basic_science_task_with_parameter_mixin,
+    db_value,
+    expected,
+):
+    """
+    Given: A Science task with the parameter mixin
+    When: Accessing properties for parameters that do not depend on wavelength
+    Then: The correct value is returned
+    """
+    task, _ = next(
+        basic_science_task_with_parameter_mixin(
+            VispInputDatasetParameterValues(visp_solar_vignette_wavecal_fit_kwargs=db_value)
+        )
+    )
+    kwarg_dict = task.parameters.solar_vignette_wavecal_fit_kwargs
+    assert kwarg_dict.keys() == expected.keys()
+    for k in expected.keys():
+        if expected[k] is AnyInt:
+            assert type(kwarg_dict[k]) is int
+        else:
+            assert expected[k] == kwarg_dict[k]

dkist_processing_visp/tests/test_parse.py

@@ -86,6 +86,7 @@ def write_input_dark_frames_to_task(
     time_delta: float = 10.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputDarkFrames(
@@ -94,6 +95,7 @@ def write_input_dark_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -109,6 +111,7 @@ def write_input_lamp_frames_to_task(
     time_delta: float = 10.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputLampGainFrames(
@@ -117,6 +120,7 @@ def write_input_lamp_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -132,6 +136,7 @@ def write_input_solar_frames_to_task(
     time_delta: float = 10.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputSolarGainFrames(
@@ -140,6 +145,7 @@ def write_input_solar_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -155,6 +161,7 @@ def write_input_polcal_frames_to_task(
     time_delta: float = 30.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputPolcalFrames(
@@ -163,6 +170,7 @@ def write_input_polcal_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -178,6 +186,7 @@ def write_input_polcal_dark_frames_to_task(
     time_delta: float = 30.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputPolcalDarkFrames(
@@ -186,6 +195,7 @@ def write_input_polcal_dark_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -201,6 +211,7 @@ def write_input_polcal_gain_frames_to_task(
     time_delta: float = 30.0,
     num_modstates: int = 2,
     data_shape: tuple[int, int] = (2, 2),
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = VispHeadersInputPolcalGainFrames(
@@ -209,6 +220,7 @@ def write_input_polcal_gain_frames_to_task(
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
         num_modstates=num_modstates,
+        **kwargs,
     )
 
     num_written_frames = write_frames_to_task(
@@ -227,6 +239,7 @@ def write_input_observe_frames_to_task(
     time_delta: float = 10.0,
     data_shape: tuple[int, int] = (2, 2),
     obs_dataset_class=VispHeadersValidObserveFrames,
+    **kwargs,
 ):
     array_shape = (1, *data_shape)
     dataset = obs_dataset_class(
@@ -237,6 +250,7 @@ def write_input_observe_frames_to_task(
         num_modstates=num_modstates,
         exp_time=exp_time,
         readout_exp_time=readout_exp_time,
+        **kwargs,
     )
     num_written_frames = write_frames_to_task(
         task=task, frame_generator=dataset, extra_tags=[VispTag.input()]
@@ -272,6 +286,11 @@ def write_input_cal_frames_to_task(
     solar_exp_time,
     polcal_exp_time,
     num_modstates,
+    testing_arm_id,
+    testing_solar_ip_start_time,
+    testing_grating_constant,
+    testing_grating_angle,
+    testing_arm_position,
 ):
     def write_frames_to_task(task):
         for readout_exp_time in required_dark_readout_exp_times:
@@ -280,6 +299,7 @@ def write_input_cal_frames_to_task(
                 readout_exp_time=readout_exp_time,
                 exp_time=dark_exp_time,
                 num_modstates=num_modstates,
+                arm_id=testing_arm_id,
             )
 
         write_input_lamp_frames_to_task(
@@ -287,30 +307,39 @@ def write_input_cal_frames_to_task(
             readout_exp_time=lamp_readout_exp_time,
             exp_time=lamp_exp_time,
             num_modstates=num_modstates,
+            arm_id=testing_arm_id,
         )
         write_input_solar_frames_to_task(
             task=task,
             readout_exp_time=solar_readout_exp_time,
             exp_time=solar_exp_time,
             num_modstates=num_modstates,
+            arm_id=testing_arm_id,
+            ip_start_time=testing_solar_ip_start_time,
+            grating_constant=testing_grating_constant,
+            grating_angle=testing_grating_angle,
+            arm_position=testing_arm_position,
         )
         write_input_polcal_frames_to_task(
             task=task,
             readout_exp_time=polcal_readout_exp_time,
             exp_time=polcal_exp_time,
             num_modstates=num_modstates,
+            arm_id=testing_arm_id,
         )
         write_input_polcal_dark_frames_to_task(
             task=task,
             readout_exp_time=polcal_readout_exp_time,
             exp_time=polcal_exp_time,
             num_modstates=num_modstates,
+            arm_id=testing_arm_id,
         )
         write_input_polcal_gain_frames_to_task(
             task,
             readout_exp_time=polcal_readout_exp_time,
             exp_time=polcal_exp_time,
             num_modstates=num_modstates,
+            arm_id=testing_arm_id,
         )
 
     return write_frames_to_task
@@ -350,6 +379,10 @@ def test_parse_visp_input_data(
     num_modstates,
     mocker,
     fake_gql_client,
+    testing_arm_id,
+    testing_grating_constant,
+    testing_grating_angle,
+    testing_arm_position,
 ):
     """
     Given: A ParseVispInputData task
@@ -369,6 +402,10 @@ def test_parse_visp_input_data(
             num_steps=3,
             readout_exp_time=obs_readout_exp_time,
             exp_time=obs_exp_time,
+            arm_id=testing_arm_id,
+            grating_constant=testing_grating_constant,
+            grating_angle=testing_grating_angle,
+            arm_position=testing_arm_position,
         )
 
     # When
@@ -404,6 +441,12 @@ def test_parse_visp_input_data_constants(
     polcal_exp_time,
     observe_exp_times,
     num_modstates,
+    testing_arm_id,
+    testing_obs_ip_start_time,
+    testing_solar_ip_start_time,
+    testing_grating_constant,
+    testing_grating_angle,
+    testing_arm_position,
 ):
     """
     Given: A ParseVispInputData task
@@ -426,16 +469,29 @@ def test_parse_visp_input_data_constants(
             num_steps=num_steps,
             readout_exp_time=obs_readout_exp_time,
             exp_time=obs_exp_time,
+            arm_id=testing_arm_id,
+            ip_start_time=testing_obs_ip_start_time,
+            grating_constant=testing_grating_constant,
+            grating_angle=testing_grating_angle,
+            arm_position=testing_arm_position,
         )
 
     # When
     task()
     # Then
+    assert task.constants._db_dict["ARM_ID"] == testing_arm_id
     expected_dark_readout_exp_times = [
         lamp_readout_exp_time,
         solar_readout_exp_time,
     ] + observe_readout_exp_times
-    assert task.constants._db_dict["OBS_IP_START_TIME"] == "2022-11-28T13:55:00"
+    assert task.constants._db_dict["OBS_IP_START_TIME"] == testing_obs_ip_start_time
+    assert task.constants._db_dict["INCIDENT_LIGHT_ANGLE_DEG"] == -1 * testing_grating_angle
+    assert (
+        task.constants._db_dict["REFLECTED_LIGHT_ANGLE_DEG"]
+        == -1 * testing_grating_angle + testing_arm_position
+    )
+    assert task.constants._db_dict["GRATING_CONSTANT_INVERSE_MM"] == testing_grating_constant
+    assert task.constants._db_dict["SOLAR_GAIN_IP_START_TIME"] == testing_solar_ip_start_time
     assert task.constants._db_dict["NUM_MODSTATES"] == num_modstates
     assert task.constants._db_dict["NUM_MAP_SCANS"] == num_maps_per_readout_exp_time * len(
         observe_readout_exp_times
@@ -468,6 +524,10 @@ def test_parse_visp_values(
     num_modstates,
     mocker,
     fake_gql_client,
+    testing_arm_id,
+    testing_grating_constant,
+    testing_grating_angle,
+    testing_arm_position,
 ):
     """
     :Given: A valid parse input task
@@ -486,6 +546,10 @@ def test_parse_visp_values(
         num_maps=1,
         num_steps=1,
         num_modstates=num_modstates,
+        arm_id=testing_arm_id,
+        grating_constant=testing_grating_constant,
+        grating_angle=testing_grating_angle,
+        arm_position=testing_arm_position,
     )
 
     task()
@@ -585,6 +649,10 @@ def test_intensity_observes_and_polarimetric_cals(
     observe_exp_times,
     mocker,
     fake_gql_client,
+    testing_arm_id,
+    testing_grating_constant,
+    testing_grating_angle,
+    testing_arm_position,
 ):
     """
     :Given: Data where the observe frames are in intensity mode and the calibration frames are in polarimetric mode
@@ -604,6 +672,10 @@ def test_intensity_observes_and_polarimetric_cals(
         readout_exp_time=observe_readout_exp_times[0],
         exp_time=observe_exp_times[0],
         obs_dataset_class=VispHeadersIntensityObserveFrames,
+        arm_id=testing_arm_id,
+        grating_constant=testing_grating_constant,
+        grating_angle=testing_grating_angle,
+        arm_position=testing_arm_position,
     )
     task()
     assert task.constants._db_dict["NUM_MODSTATES"] == 1

dkist_processing_visp/tests/test_science.py

@@ -144,7 +144,7 @@ def dummy_calibration_collection():
 
     dark_dict = {VispTag.beam(beam): {VispTag.readout_exp_time(0.04): np.zeros(intermediate_shape)}}
     background_dict = {VispTag.beam(beam): np.zeros(intermediate_shape)}
-    solar_dict = {VispTag.beam(beam): {VispTag.modstate(modstate): np.ones(intermediate_shape)}}
+    solar_dict = {VispTag.beam(beam): np.ones(intermediate_shape)}
     angle_dict = {VispTag.beam(beam): 0.0}
     spec_dict = {VispTag.beam(beam): np.zeros(intermediate_shape[1])}
     offset_dict = {VispTag.beam(beam): {VispTag.modstate(modstate): np.zeros(2)}}
@@ -252,7 +252,7 @@ def calibration_collection_with_full_overlap_slice() -> CalibrationCollection:
 
 @pytest.mark.parametrize(
     "background_on",
-    [pytest.param(True, id="Background on"), pytest.param(False, id="Background off")],
+    [pytest.param(True, id="background_on"), pytest.param(False, id="background_off")],
 )
 def test_science_calibration_task(
     science_calibration_task,
@@ -308,7 +308,8 @@ def test_science_calibration_task(
         task=task, num_modstates=num_modstates, data_shape=intermediate_shape, offsets=offsets
     )
     write_dummy_intermediate_solar_cals_to_task(
-        task=task, data_shape=intermediate_shape, num_modstates=num_modstates
+        task=task,
+        data_shape=intermediate_shape,
     )
     write_demod_matrices_to_task(task=task, num_modstates=num_modstates)
     write_input_observe_frames_to_task(
@@ -658,9 +659,7 @@ def test_generate_nan_mask(science_calibration_task, dummy_calibration_collectio
     calibration_collection, _, _ = dummy_calibration_collection
     beam = 1
     modstate = 1
-    solar_gain_array = calibration_collection.solar_gain[VispTag.beam(beam)][
-        VispTag.modstate(modstate)
-    ]
+    solar_gain_array = calibration_collection.solar_gain[VispTag.beam(beam)]
     angle = calibration_collection.angle[VispTag.beam(beam)]
     spec_shift = calibration_collection.spec_shift[VispTag.beam(beam)]
     state_offset = calibration_collection.state_offset[VispTag.beam(beam)][