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

dkist_processing_trend/tests/conftest.py

@@ -0,0 +1,718 @@
+import json
+from dataclasses import asdict
+from dataclasses import dataclass
+from dataclasses import field
+from dataclasses import is_dataclass
+from datetime import datetime
+from functools import partial
+from itertools import product
+from random import randint
+from typing import Type
+
+import astropy.units as u
+import numpy as np
+import pytest
+from astropy import coordinates
+from astropy.io import fits
+from astropy.time import Time
+from dkist_data_simulator.dataset import key_function
+from dkist_data_simulator.spec122 import Spec122Dataset
+from dkist_processing_common.codecs.basemodel import basemodel_encoder
+from dkist_processing_common.codecs.fits import fits_array_encoder
+from dkist_processing_common.codecs.fits import fits_hdulist_encoder
+from dkist_processing_common.models.input_dataset import InputDatasetPartDocumentList
+from dkist_processing_common.tasks.base import WorkflowTaskBase
+
+# Despite what an IDE might think, the following import is used by most of the tests. Don't delete it!
+from dkist_processing_common.tests.mock_metadata_store import fake_gql_client
+
+from dkist_processing_trend.models.constants import TrendConstants
+from dkist_processing_trend.models.instrument import InstrumentName
+from dkist_processing_trend.models.tags import TrendTag
+
+
+@pytest.fixture()
+def recipe_run_id():
+    return randint(0, 99999)
+
+
+@dataclass
+class TrendInputDatasetParameterValues:
+    trend_max_cs_step_time_sec: float = 180.0
+    trend_fit_options_list: tuple[dict] = (
+        {
+            "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",
+        },
+    )
+    trend_instrument_processing_options_visp: tuple[dict] = (
+        {
+            "name": "visp_option1",
+            "num_spatial_px": 2560,
+            "num_spectral_px": 1,
+            "spatial_smoothing_fit_order": 3,
+        },
+        {
+            "name": "visp_option2",
+            "num_spatial_px": 1,
+            "num_spectral_px": 1,
+            "spatial_smoothing_fit_order": 1,
+        },
+    )
+    trend_instrument_processing_options_cryonirsp: tuple[dict] = (
+        {"name": "cryonirsp_option1", "num_spatial_px": 1, "num_spectral_px": 1},
+        {"name": "cryonirsp_option2", "num_spatial_px": 2, "num_spectral_px": 1},
+    )
+    trend_instrument_processing_options_dlnirsp: tuple[dict] = (
+        {
+            "name": "dlnirsp_option1",
+            "num_spatial_px": 1000,
+            "num_spectral_px": 1000,
+            "spatial_smoothing_fit_order": 3,
+        },
+        {
+            "name": "dlnirsp_option2",
+            "num_spatial_px": 1,
+            "num_spectral_px": 1,
+            "spatial_smoothing_fit_order": 1,
+        },
+    )
+    trend_visp_beam_border: int = 1000
+    trend_visp_geo_upsample_factor: int = 100
+    trend_visp_hairline_median_spatial_smoothing_width_px: int = 30
+    trend_visp_hairline_fraction: float = 0.11
+    trend_visp_hairline_mask_spatial_smoothing_width_px: float = 1.0
+    trend_visp_hairline_mask_gaussian_peak_cutoff_fraction: float = 0.02
+    trend_visp_polcal_spatial_median_filter_width_px: int = 10
+    trend_visp_polcal_demod_spatial_smooth_min_samples: float = 0.9
+    trend_visp_polcal_demod_upsample_order: int = 3
+
+
+WAVELENGTH_FOR_ARM = {
+    # Visp
+    1: 854.2,
+    2: 589.6,
+    3: 630.2,
+    # Cryo
+    "CI": 1083.0,
+    "SP": 1074.7,
+    # DL
+    "VIS": 854.2,
+    "JBand": 1083.0,
+    "HBand": 1565.0,
+}
+
+
+GRATING_ANGLE_FOR_ARM = {  # Visp
+    1: 9.871,
+    2: 9.872,
+    3: 9.873,
+    # Cryo
+    "CI": 123.4,
+    "SP": 123.45,
+    # DL
+    "VIS": 2.22,
+    "JBand": 2.222,
+    "HBand": 2.2222,
+}
+
+ARM_ANGLE_FOR_ARM = {
+    # Visp
+    1: 8.761,
+    2: 8.762,
+    3: 8.763,
+    # Cryo
+    "CI": 12.34,
+    "SP": 12.345,
+    # DL
+    "VIS": 3.33,
+    "JBand": 3.333,
+    "HBand": 3.3333,
+}
+GRATING_CONSTANT_FOR_ARM = {
+    # Visp
+    1: 9871,
+    2: 9872,
+    3: 9873,
+    # Cryo
+    "CI": 1234,
+    "SP": 12345,
+    # DL
+    "VIS": 4444,
+    "JBand": 4444.4,
+    "HBand": 4444.44,
+}
+
+CRYO_OD_FILTER_FOR_ARM = {
+    "CI": "CIFILT",
+    "SP": "SPFILT",
+    # Need dummy values for other instrument's arm so tests can slice into this dict with any arm
+    1: None,
+    2: None,
+    3: None,
+    "VIS": None,
+    "JBand": None,
+    "HBand": None,
+}
+CRYO_SLIT_WIDTH_FOR_ARM = {
+    "CI": 0.1234,
+    "SP": 0.12345,
+    # Need dummy values for other instrument's arm so tests can slice into this dict with any arm
+    1: None,
+    2: None,
+    3: None,
+    "VIS": None,
+    "JBand": None,
+    "HBand": None,
+}
+
+
+@dataclass
+class TrendTestingConstantsDatabase:
+    INSTRUMENT: str = "VISP"
+    EARLIEST_IP_START_TIME: str = "2026-05-01"
+    ARM_ID_LIST: tuple[str | int] = (1, 2, 3)
+    WAVELENGTH_FOR_ARM: dict[str, float] = field(default_factory=lambda: {"1": 123.4, "2": 567.8})
+    NUM_CS_STEPS: int = 7  # Number of steps in `small_calibration_sequence`
+    NUM_MODSTATES: int = 4
+    OBSERVING_PROGRAM_ID_LIST: tuple[str] = ("TOP1", "TOP2")
+    EXPERIMENT_ID_LIST: tuple[str] = "EX1"
+    PROPOSAL_ID_LIST: tuple[str] = "PROP5"
+    HLS_VERSION_LIST: tuple[str] = ("Tau 6.2.8", "Pi 3.14")
+    VISP_GRATING_ANGLE_FOR_ARM: dict[str, float] = field(
+        default_factory=lambda: {"1": 9.87, "2": 9.872}
+    )
+    VISP_ARM_POSITION_FOR_ARM: dict[str, float] = field(
+        default_factory=lambda: {"1": 8.76, "2": 8.762}
+    )
+    VISP_GRATING_CONSTANT_FOR_ARM: dict[str, float] = field(
+        default_factory=lambda: {"1": 9870, "2": 9872}
+    )
+    CRYONIRSP_OD_FILTER_NAME_FOR_ARM: dict[str, str] = field(
+        default_factory=lambda: {'"CI"': "CRYOFILT"}
+    )
+    CRYONIRSP_GRATING_ANGLE_FOR_ARM: dict[str, float] = field(
+        default_factory=lambda: {'"CI"': 123.4, '"SP"': 123.5}
+    )
+    CRYONIRSP_GRATING_LITTROW_ANGLE_FOR_ARM: dict[str, float] = field(
+        default_factory=lambda: {'"CI"': 12.34, '"SP"': 12.35}
+    )
+    CRYONIRSP_GRATING_CONSTANT_FOR_ARM: dict[str, float] = field(
+        default_factory=lambda: {'"SP"': 1234}
+    )
+    CRYONIRSP_SLIT_WIDTH_FOR_ARM: dict[str, float] = field(default_factory=lambda: {'"CI"': 0.1234})
+    DLNIRSP_GRATING_ANGLE_FOR_ARM: dict[str, float] = field(
+        default_factory=lambda: {'"VIS"': 2.22, '"JBand"': 2.222, '"HBand"': 2.2222}
+    )
+    DLNIRSP_ARM_POSITION_FOR_ARM: dict[str, float] = field(
+        default_factory=lambda: {'"JBand"': 3.33}
+    )
+    DLNIRSP_GRATING_CONSTANT_FOR_ARM: dict[str, float] = field(
+        default_factory=lambda: {'"HBand"': 4444}
+    )
+
+
+@pytest.fixture(scope="session")
+def visp_grating_angle() -> float:
+    return 1.23
+
+
+@pytest.fixture(scope="session")
+def visp_arm_position() -> float:
+    return 3.45
+
+
+@pytest.fixture(scope="session")
+def visp_grating_constant() -> float:
+    return 6789
+
+
+@pytest.fixture(scope="session")
+def cryonirsp_od_filter_name() -> str:
+    return "FILTER"
+
+
+@pytest.fixture(scope="session")
+def cryonirsp_grating_angle() -> float:
+    return 4.56
+
+
+@pytest.fixture(scope="session")
+def cryonirsp_grating_littrow_angle() -> float:
+    return 5.67
+
+
+@pytest.fixture(scope="session")
+def cryonirsp_grating_constant() -> float:
+    return 7.89
+
+
+@pytest.fixture(scope="session")
+def cryonirsp_slit_width() -> float:
+    return 0.3
+
+
+@pytest.fixture(scope="session")
+def dlnirsp_grating_angle() -> float:
+    return 3.21
+
+
+@pytest.fixture(scope="session")
+def dlnirsp_arm_position() -> float:
+    return 5.43
+
+
+@pytest.fixture(scope="session")
+def dlnirsp_grating_constant() -> float:
+    return 9876
+
+
+@pytest.fixture(scope="session")
+def input_dataset_document_simple_parameters_part():
+    def get_input_dataset_parameters_part(parameters: TrendInputDatasetParameterValues):
+        parameters_list = []
+        value_id = randint(1000, 2000)
+        for pn, pv in asdict(parameters).items():
+            values = [
+                {
+                    "parameterValueId": value_id,
+                    "parameterValue": json.dumps(pv),
+                    "parameterValueStartDate": "1946-11-20",  # Remember Duane Allman
+                }
+            ]
+            parameter = {"parameterName": pn, "parameterValues": values}
+            parameters_list.append(parameter)
+        return parameters_list
+
+    return get_input_dataset_parameters_part
+
+
+@pytest.fixture(scope="session")
+def assign_input_dataset_doc_to_task(
+    input_dataset_document_simple_parameters_part,
+):
+    # NOTE: Unlike the other instruments, this fixture does NOT set the `task.parameters` attribute; it only writes the
+    # doc to scratch and tags it.
+    # This means you'll need to add `task.parameters = ParameterClass(...)` after using this fixture
+    def update_task(
+        task: WorkflowTaskBase,
+        parameter_values,
+    ):
+        task.write(
+            data=InputDatasetPartDocumentList(
+                doc_list=input_dataset_document_simple_parameters_part(parameter_values)
+            ),
+            tags=TrendTag.input_dataset_parameters(),
+            encoder=basemodel_encoder,
+        )
+
+    return update_task
+
+
+@pytest.fixture
+def link_constants_db():
+    def constants_linker(recipe_run_id: int, constants_obj):
+        """Take a dataclass (or dict) containing a constants DB and link it to a specific recipe run id."""
+        if is_dataclass(constants_obj):
+            constants_obj = asdict(constants_obj)
+        constants = TrendConstants(recipe_run_id=recipe_run_id, task_name="test")
+        constants._purge()
+        constants._update(constants_obj)
+
+    return constants_linker
+
+
+class ModulatedCSStepHeaders(Spec122Dataset):
+    def __init__(
+        self,
+        instrument: InstrumentName,
+        arm_id: int | str,
+        wavelength: float,
+        num_modstates: int,
+        pol_status: str,
+        pol_theta: float,
+        ret_status: str,
+        ret_theta: float,
+        dark_status: str,
+        cs_step_num: int,
+        array_shape: tuple[int, int],
+        exp_time_ms: float,
+        start_date: str = "2023-01-01T01:23:45",
+        modstate_length_sec: float = 0.5,
+        num_frames_per_cs_step: int = 1,
+        op_id: str = "Cool_OP",
+        experiment_id: str = "exp1",
+        proposal_id: str = "prop_1",
+        hls_version: str = "hls_v1",
+        visp_grating_angle: float = 6.28,
+        visp_arm_position: float = 6.2831,
+        visp_grating_constant: float = 62831,
+        cryonirsp_od_filter_name: str = "cryonirsp_filter",
+        cryonirsp_grating_angle: float = 8.67,
+        cryonirsp_grating_littrow_angle: float = 6.78,
+        cryonirsp_grating_constant: float = 7890,
+        cryonirsp_slit_width: float = 0.123,
+        dlnirsp_grating_angle: float = 11.22,
+        dlnirsp_arm_position: float = 22.11,
+        dlnirsp_grating_constant: float = 1231,
+    ):
+        indices = list(product(range(1, num_modstates + 1), range(num_frames_per_cs_step)))
+
+        self.index_to_modstate = [ms for (ms, _) in indices]
+        self.index_to_cs_step_frame = [csf for (_, csf) in indices]
+
+        # This doesn't become a header key, but helps with making fake data
+        self.cs_step_num = cs_step_num
+        self.pol_status = pol_status
+        self.pol_theta = pol_theta
+        self.ret_status = ret_status
+        self.ret_theta = ret_theta
+        self.dark_status = dark_status
+        self.start_date = Time(start_date)
+        self.modstate_length_sec = modstate_length_sec
+        self.num_frames_per_cs_step = num_frames_per_cs_step
+
+        if len(array_shape) == 3:
+            if array_shape[0] != 1:
+                raise ValueError(f"{array_shape = } is weird")
+        elif len(array_shape) == 2:
+            array_shape = (1, *array_shape)
+
+        dataset_shape = (num_modstates * self.num_frames_per_cs_step, *array_shape)
+        inst_schema_table_name = InstrumentName(instrument).name
+        super().__init__(
+            instrument=inst_schema_table_name,
+            dataset_shape=dataset_shape,
+            array_shape=array_shape,
+            start_time=datetime.fromisoformat(start_date),
+            time_delta=modstate_length_sec,
+            file_schema="level0_spec214",
+        )
+
+        self.add_constant_key("WAVELNTH", wavelength)
+        self.add_constant_key("DKIST004", "polcal")
+        self.add_constant_key("CAM__004", exp_time_ms)
+        self.add_constant_key("ID___008", op_id)
+        self.add_constant_key("ID___012", experiment_id)
+        self.add_constant_key("ID___013", proposal_id)
+        self.add_constant_key("ID___014", hls_version)
+
+        # These should all be 122 keys
+        if instrument == InstrumentName.visp:
+            self.add_constant_key("VISP_001", arm_id)
+            self.add_constant_key("VISP_010", num_modstates)
+            self.add_generator_function(key="VISP_011", function=modstate_key_function)
+            self.add_constant_key("VISP_002", visp_arm_position)
+            self.add_constant_key("VISP_013", visp_grating_constant)
+            self.add_constant_key("VISP_015", visp_grating_angle)
+
+        if instrument == InstrumentName.cryonirsp:
+            self.add_constant_key("CRSP_001", arm_id)
+            self.add_constant_key("CRSP_041", num_modstates)
+            self.add_generator_function(key="CRSP_042", function=modstate_key_function)
+            self.add_constant_key("CRSP_048", cryonirsp_od_filter_name)
+            self.add_constant_key("CRSP_074", cryonirsp_grating_angle)
+            self.add_constant_key("CRSP_077", cryonirsp_grating_constant)
+            self.add_constant_key("CRSP_079", cryonirsp_grating_littrow_angle)
+            self.add_constant_key("CRSP_082", cryonirsp_slit_width)
+
+        if instrument == InstrumentName.dlnirsp:
+            self.add_constant_key("DLN__001", arm_id)
+            self.add_constant_key("DLN__014", num_modstates)
+            self.add_generator_function(key="DLN__015", function=modstate_key_function)
+            self.add_constant_key("DLN__002", dlnirsp_arm_position)
+            self.add_constant_key("DLN__017", dlnirsp_grating_constant)
+            self.add_constant_key("DLN__019", dlnirsp_grating_angle)
+
+    @key_function("PAC__004")
+    def polarizer_status(self, key: str) -> str:
+        return self.pol_status
+
+    @key_function("PAC__005")
+    def polarizer_angle(self, key: str) -> float | str:
+        return "none" if self.pol_status == "clear" else self.pol_theta
+
+    @key_function("PAC__006")
+    def retarder_status(self, key: str) -> str:
+        return self.ret_status
+
+    @key_function("PAC__007")
+    def retarder_angle(self, key: str) -> float | str:
+        return "none" if self.ret_status == "clear" else self.ret_theta
+
+    @key_function("PAC__008")
+    def gos_level3_status(self, key: str) -> str:
+        return self.dark_status
+
+    @key_function("TAZIMUTH", "TELEVATN", "TTBLANGL")
+    def telescope_geometry(self, key: str):
+        return compute_telgeom(Time(self.date_obs(key), format="fits"))[key]
+
+    @property
+    def current_cs_step_frame(self) -> int:
+        # There is no header value for CS step frame; used to help data generator fixtures with tagging
+        return self.index_to_cs_step_frame[self.index]
+
+    @property
+    def current_modstate(self) -> int:
+        # This is a property because the actual key generator function is added in __init__ because the associated
+        # key is not static but based on the instrument.
+        # This property is here for convenience in tests and fixtures
+        return self.index_to_modstate[self.index]
+
+
+def modstate_key_function(dataset: ModulatedCSStepHeaders, key: str) -> int:
+    return dataset.current_modstate
+
+
+def compute_telgeom(time_hst: Time):
+    dkist_lon = (156 + 15 / 60.0 + 21.7 / 3600.0) * (-1)
+    dkist_lat = 20 + 42 / 60.0 + 27.0 / 3600.0
+    hel = 3040.4
+    hloc = coordinates.EarthLocation.from_geodetic(dkist_lon, dkist_lat, hel)
+    sun_body = coordinates.get_body("sun", time_hst, hloc)  # get the solar ephemeris
+    azel_frame = coordinates.AltAz(obstime=time_hst, location=hloc)  # Horizon coords
+    sun_altaz = sun_body.transform_to(azel_frame)  # Sun in horizon coords
+    alt = sun_altaz.alt.value  # Extract altitude
+    azi = sun_altaz.az.value  # Extract azimuth
+
+    tableang = alt - azi
+
+    return {"TELEVATN": alt, "TAZIMUTH": azi, "TTBLANGL": tableang}
+
+
+def make_random_data(frame: Spec122Dataset) -> np.ndarray:
+    shape = frame.array_shape[1:]
+    data = np.random.random(shape)
+
+    return data
+
+
+def make_cs_data(
+    frame: ModulatedCSStepHeaders, dark_signal: float, clear_signal: float
+) -> np.ndarray:
+
+    shape = frame.array_shape[1:]
+    clear_signal += frame.current_modstate
+    if frame.pol_status == "clear" and frame.ret_status == "clear":
+        if frame.dark_status == "DarkShutter":
+            value = dark_signal
+        else:
+            value = clear_signal + dark_signal
+    else:
+        value = (
+            frame.cs_step_num * 10000.0 + frame.current_modstate * 100.0
+        ) * clear_signal + dark_signal
+
+    data = np.full(shape, value)
+
+    return data
+
+
+def tag_on_modstate(frame: ModulatedCSStepHeaders) -> list[str]:
+    modstate = frame.current_modstate
+    return [TrendTag.modstate(modstate)]
+
+
+def write_frames_to_task(
+    task: WorkflowTaskBase,
+    frame_generator: Spec122Dataset,
+    data_func: callable = make_random_data,
+    extra_tags: list[str] | None = None,
+    tag_func: callable = lambda x: [],
+    compress_files: bool = False,
+):
+    if not extra_tags:
+        extra_tags = []
+    tags = [TrendTag.frame()] + extra_tags
+
+    num_frames = 0
+    for frame in frame_generator:
+        header = frame.header()
+        data = data_func(frame)
+        frame_tags = tags + tag_func(frame)
+        if compress_files:
+            hdul = fits.HDUList([fits.PrimaryHDU(), fits.CompImageHDU(data=data, header=header)])
+            task.write(data=hdul, tags=frame_tags, encoder=fits_hdulist_encoder)
+        else:
+            task.write(data=data, header=header, tags=frame_tags, encoder=fits_array_encoder)
+        num_frames += 1
+
+    return num_frames
+
+
+def write_calibration_sequence_frames_to_task(
+    task: Type[WorkflowTaskBase],
+    instrument: InstrumentName,
+    arm_id: int | str,
+    wavelength: float,
+    pol_status: list[str],
+    pol_theta: list[float],
+    ret_status: list[str],
+    ret_theta: list[float],
+    dark_status: list[str],
+    num_modstates: int = 4,
+    exp_time_ms: float = 7.0,
+    dark_signal: float = 5.0,
+    clear_signal: float = 10.0,
+    array_shape: tuple[int, ...] = (10, 10),
+    num_frames_per_cs_step: int = 1,
+    input_tags: list[str] | None = None,
+    data_func: callable = None,
+    tags: list[str] | None = None,
+    op_start_time: str = "2023-01-01T01:23:45",
+    op_id: str = "Cool_OPID",
+    experiment_id: str = "cool exp",
+    proposal_id: str = "very cool proposal",
+    hls_version: str = "HLS_VERSION_x.x",
+    visp_grating_angle: float = 3.14,
+    visp_arm_position: float = 3.14159,
+    visp_grating_constant: float = 31415,
+    cryonirsp_od_filter_name: str = "cryonirsp_filter",
+    cryonirsp_grating_angle: float = 2.34,
+    cryonirsp_grating_littrow_angle: float = 3.45,
+    cryonirsp_grating_constant: float = 4.56,
+    cryonirsp_slit_width: float = 0.45,
+    dlnirsp_grating_angle: float = 4.55,
+    dlnirsp_arm_position: float = 4.5678,
+    dlnirsp_grating_constant: float = 5678,
+) -> int:
+    """
+    Write a full Calibration Sequence to a task's scratch.
+
+    If `input_tags` is specified then the files will be written with ONLY those tags. Otherwise they will always
+    have CS_STEP, MODSTATE, TASK_POLCAL, and (maybe) TASK_POLCAL_DARK or TASK_POLCAL_GAIN tags along with whatever is in
+    `extra_tags`. The FRAME tag is *always* included.
+    """
+    if data_func is None:
+        data_func = partial(make_cs_data, dark_signal=dark_signal, clear_signal=clear_signal)
+
+    num_frames = 0
+    base_time = Time(op_start_time)
+
+    for step, (pol_s, pol_t, ret_s, ret_t, dark_s) in enumerate(
+        zip(pol_status, pol_theta, ret_status, ret_theta, dark_status)
+    ):
+        step_start = (base_time + (step * 60) * u.s).fits
+
+        dataset = ModulatedCSStepHeaders(
+            instrument=instrument,
+            arm_id=arm_id,
+            wavelength=wavelength,
+            num_modstates=num_modstates,
+            pol_status=pol_s,
+            pol_theta=pol_t,
+            ret_status=ret_s,
+            ret_theta=ret_t,
+            dark_status=dark_s,
+            cs_step_num=step,
+            array_shape=array_shape,
+            exp_time_ms=exp_time_ms,
+            num_frames_per_cs_step=num_frames_per_cs_step,
+            start_date=step_start,
+            op_id=op_id,
+            experiment_id=experiment_id,
+            proposal_id=proposal_id,
+            hls_version=hls_version,
+            visp_grating_angle=visp_grating_angle,
+            visp_arm_position=visp_arm_position,
+            visp_grating_constant=visp_grating_constant,
+            cryonirsp_od_filter_name=cryonirsp_od_filter_name,
+            cryonirsp_grating_angle=cryonirsp_grating_angle,
+            cryonirsp_grating_littrow_angle=cryonirsp_grating_littrow_angle,
+            cryonirsp_grating_constant=cryonirsp_grating_constant,
+            cryonirsp_slit_width=cryonirsp_slit_width,
+            dlnirsp_grating_angle=dlnirsp_grating_angle,
+            dlnirsp_arm_position=dlnirsp_arm_position,
+            dlnirsp_grating_constant=dlnirsp_grating_constant,
+        )
+
+        if tags is None:
+            tags = []
+        if input_tags:
+            tag_func = lambda x: []
+            extra_tags = input_tags
+        else:
+            tag_func = tag_on_modstate
+            extra_tags = list(
+                set(
+                    tags
+                    + [
+                        TrendTag.cs_step(step),
+                        TrendTag.task_polcal(),
+                    ]
+                )
+            )
+
+            if pol_s == "clear" and ret_s == "clear":
+                if dark_s == "DarkShutter":
+                    extra_tags += [TrendTag.task_polcal_dark()]
+                else:
+                    extra_tags += [TrendTag.task_polcal_gain()]
+
+        num_frames += write_frames_to_task(
+            task=task,
+            frame_generator=dataset,
+            data_func=data_func,
+            extra_tags=extra_tags,
+            tag_func=tag_func,
+        )
+
+    return num_frames
+
+
+@pytest.fixture(scope="session")
+def small_calibration_sequence() -> tuple[list, list, list, list, list]:
+    # Make up a Calibration sequence. Mostly random except for two clears and two darks at start and end
+    pol_status = [
+        "clear",
+        "clear",
+        "Sapphire Polarizer",
+        "Sapphire Polarizer",
+        "Sapphire Polarizer",
+        "clear",
+        "clear",
+    ]
+    pol_theta = [0.0, 0.0, 60.0, 0.0, 120.0, 0.0, 0.0]
+    ret_status = ["clear", "clear", "clear", "SiO2 SAR", "clear", "clear", "clear"]
+    ret_theta = [0.0, 0.0, 0.0, 45.0, 0.0, 0.0, 0.0]
+    dark_status = [
+        "DarkShutter",
+        "FieldStop (5arcmin)",
+        "FieldStop (5arcmin)",
+        "FieldStop (5arcmin)",
+        "FieldStop (5arcmin)",
+        "FieldStop (5arcmin)",
+        "DarkShutter",
+    ]
+
+    return pol_status, pol_theta, ret_status, ret_theta, dark_status
+
+
+@pytest.fixture(scope="session")
+def small_num_cs_steps(small_calibration_sequence) -> int:
+    return len(small_calibration_sequence[0])
+
+
+@pytest.fixture(scope="session")
+def small_num_gain_steps(small_calibration_sequence) -> int:
+    num_gains = 0
+    for ps, pt, rs, rt, ds in zip(*small_calibration_sequence):
+        if ps == "clear" and rs == "clear" and ds != "DarkShutter":
+            num_gains += 1
+
+    return num_gains
+
+
+@pytest.fixture(scope="session")
+def small_num_dark_steps(small_calibration_sequence) -> int:
+    dark_status = small_calibration_sequence[-1]
+    return dark_status.count("DarkShutter")