diff --git a/python/lsst/ts/standardscripts/maintel/__init__.py b/python/lsst/ts/standardscripts/maintel/__init__.py
index 27c01a834..fc5ac9c2b 100644
--- a/python/lsst/ts/standardscripts/maintel/__init__.py
+++ b/python/lsst/ts/standardscripts/maintel/__init__.py
@@ -45,6 +45,7 @@
from .standby_mtcs import *
from .stop import *
from .stop_rotator import *
+from .take_aos_sequence_balanced_comcam import *
from .take_aos_sequence_comcam import *
from .take_image_anycam import *
from .take_image_comcam import *
diff --git a/python/lsst/ts/standardscripts/maintel/take_aos_sequence_balanced_comcam.py b/python/lsst/ts/standardscripts/maintel/take_aos_sequence_balanced_comcam.py
new file mode 100644
index 000000000..06cbc947a
--- /dev/null
+++ b/python/lsst/ts/standardscripts/maintel/take_aos_sequence_balanced_comcam.py
@@ -0,0 +1,382 @@
+# This file is part of ts_standardscripts
+#
+# Developed for the LSST Telescope and Site Systems.
+# This product includes software developed by the LSST Project
+# (https://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see .
+
+__all__ = ["TakeAOSSequenceBalancedComCam", "Mode"]
+
+import asyncio
+import enum
+import json
+import types
+
+import yaml
+from lsst.ts import salobj
+from lsst.ts.observatory.control.maintel.comcam import ComCam
+from lsst.ts.observatory.control.maintel.mtcs import MTCS
+
+from ..base_block_script import BaseBlockScript
+
+
+class Mode(enum.IntEnum):
+ TRIPLET = enum.auto()
+ INTRA = enum.auto()
+ EXTRA = enum.auto()
+ PAIR = enum.auto()
+
+
+class TakeAOSSequenceBalancedComCam(BaseBlockScript):
+ """Take aos sequence, either triplet (intra-focal, extra-focal
+ and in-focus images), intra doublets (intra and in-focus) or extra
+ doublets (extra and in-focus) sequences with ComCam.
+
+ This version splits the dz offset evenly between the camera and M2
+ hexapods.
+
+ Parameters
+ ----------
+ index : `int`
+ Index of Script SAL component.
+
+ Notes
+ -----
+ **Checkpoints**
+
+ * sequence {n} of {m}: before taking a sequence.
+
+ """
+
+ def __init__(self, index, descr="Take AOS sequence with ComCam.") -> None:
+ super().__init__(index=index, descr=descr)
+
+ self.config = None
+ self.mtcs = None
+ self.camera = None
+ self.ocps = None
+ self.current_z_position = 0
+ self.n_images = 9
+
+ @classmethod
+ def get_schema(cls) -> dict:
+ schema_yaml = f"""
+ $schema: http://json-schema.org/draft-07/schema#
+ $id: https://github.com/lsst-ts/ts_standardscripts/maintel/TakeAOSSequenceComCam.yaml
+ title: TakeAOSSequenceComCam v1
+ description: Configuration for TakeAOSSequenceComCam.
+ type: object
+ properties:
+ filter:
+ description: Filter name or ID; if omitted the filter is not changed.
+ anyOf:
+ - type: string
+ - type: integer
+ minimum: 1
+ - type: "null"
+ default: null
+ exposure_time:
+ description: The exposure time to use when taking images (sec).
+ type: number
+ default: 30.
+ dz:
+ description: De-focus to apply when acquiring the intra/extra focal images (microns).
+ type: number
+ default: 1500.
+ n_sequences:
+ description: Number of aos sequences.
+ type: integer
+ default: 1
+ mode:
+ description: >-
+ Mode of operation. Options are 'triplet' (default), 'intra' or 'extra'.
+ type: string
+ default: TRIPLET
+ enum: {[mode.name for mode in Mode]}
+ program:
+ description: >-
+ Optional name of the program this dataset belongs to.
+ type: string
+ default: AOSSEQUENCE
+ reason:
+ description: Optional reason for taking the data.
+ anyOf:
+ - type: string
+ - type: "null"
+ default: null
+ note:
+ description: A descriptive note about the image being taken.
+ anyOf:
+ - type: string
+ - type: "null"
+ default: null
+ ignore:
+ description: >-
+ CSCs from the group to ignore in status check. Name must
+ match those in self.group.components, e.g.; hexapod_1.
+ type: array
+ items:
+ type: string
+ additionalProperties: false
+ """
+ return yaml.safe_load(schema_yaml)
+
+ async def configure(self, config: types.SimpleNamespace) -> None:
+ """Configure script.
+
+ Parameters
+ ----------
+ config : `types.SimpleNamespace`
+ Script configuration, as defined by `schema`.
+ """
+ # Configure tcs and camera
+ await self.configure_tcs()
+ await self.configure_camera()
+
+ if hasattr(config, "ignore"):
+ for comp in config.ignore:
+ if comp in self.mtcs.components_attr:
+ self.log.debug(f"Ignoring MTCS component {comp}.")
+ setattr(self.mtcs.check, comp, False)
+ elif comp in self.camera.components_attr:
+ self.log.debug(f"Ignoring Camera component {comp}.")
+ setattr(self.camera.check, comp, False)
+ else:
+ self.log.warning(
+ f"Component {comp} not in CSC Groups. "
+ f"Must be one of {self.mtcs.components_attr} or "
+ f"{self.camera.components_attr}. Ignoring."
+ )
+
+ # Set filter
+ self.filter = config.filter
+
+ # Set exposure time
+ self.exposure_time = config.exposure_time
+
+ # Set intra/extra focal offsets
+ self.dz = config.dz
+
+ # Set maximum number of iterations
+ self.n_sequences = config.n_sequences
+
+ self.mode = getattr(Mode, config.mode)
+
+ # Set program, reason and note
+ self.program = config.program
+ self.reason = config.reason
+ self.note = config.note
+
+ def set_metadata(self, metadata: salobj.type_hints.BaseMsgType) -> None:
+ """Sets script metadata.
+
+ Parameters
+ ----------
+ metadata : `salobj.type_hints.BaseMsgType`
+ Script metadata topic. The information is set on the topic
+ directly.
+ """
+ # Estimated duration is maximum number of iterations multiplied by
+ # 3 or 2 multiplied by the time it takes to take an image
+ # plus estimation on reading out the images (10s)
+ number_of_images = 3 if self.mode == Mode.TRIPLET else 2
+
+ metadata.duration = (
+ self.n_sequences
+ * number_of_images
+ * (
+ self.exposure_time
+ + self.camera.read_out_time
+ + self.camera.shutter_time
+ )
+ )
+ metadata.filter = f"{self.filter}"
+
+ async def assert_feasibility(self) -> None:
+ """Verify that the telescope and camera are in a feasible state to
+ execute the script.
+ """
+ await asyncio.gather(
+ self.mtcs.assert_all_enabled(), self.camera.assert_all_enabled()
+ )
+
+ async def configure_camera(self) -> None:
+ """Handle creating ComCam object and waiting for remote to start."""
+ if self.camera is None:
+ self.log.debug("Creating Camera.")
+
+ self.camera = ComCam(
+ self.domain,
+ log=self.log,
+ tcs_ready_to_take_data=self.mtcs.ready_to_take_data,
+ )
+ await self.camera.start_task
+ else:
+ self.log.debug("Camera already defined, skipping.")
+
+ if self.ocps is None:
+ self.log.debug("Create OCPS remote.")
+
+ self.ocps = salobj.Remote(self.domain, "OCPS", 101)
+
+ await self.ocps.start_task
+
+ async def configure_tcs(self) -> None:
+ """Handle creating MTCS object and waiting for remote to start."""
+ if self.mtcs is None:
+ self.log.debug("Creating MTCS.")
+ self.mtcs = MTCS(
+ domain=self.domain,
+ log=self.log,
+ )
+ await self.mtcs.start_task
+ else:
+ self.log.debug("MTCS already defined, skipping.")
+
+ async def take_aos_sequence(self) -> None:
+ """Take out-of-focus sequence images."""
+ supplemented_group_id = self.next_supplemented_group_id()
+
+ if (
+ self.mode == Mode.TRIPLET
+ or self.mode == Mode.INTRA
+ or self.mode == Mode.PAIR
+ ):
+ self.log.debug("Moving to intra-focal position")
+
+ # Move the camera and M2 hexapods to the target z position
+ # Offset split in half and shared between each hexapod
+ z_offset = (-self.dz - self.current_z_position) / 2
+ await self.mtcs.offset_camera_hexapod(x=0, y=0, z=z_offset, u=0, v=0)
+ await self.mtcs.offset_m2_hexapod(x=0, y=0, z=z_offset, u=0, v=0)
+ self.current_z_position = -self.dz
+
+ self.log.info("Taking intra-focal image")
+ self.camera.rem.ccoods.evt_imageInOODS.flush()
+ intra_visit_id = await self.camera.take_cwfs(
+ exptime=self.exposure_time,
+ n=1,
+ group_id=supplemented_group_id,
+ filter=self.filter,
+ reason="INTRA" + ("" if self.reason is None else f"_{self.reason}"),
+ program=self.program,
+ note=self.note,
+ )
+
+ if (
+ self.mode == Mode.TRIPLET
+ or self.mode == Mode.EXTRA
+ or self.mode == Mode.PAIR
+ ):
+ self.log.debug("Moving to extra-focal position")
+
+ # Move the camera and M2 hexapods to the target z position
+ # Offset split in half and shared between each hexapod
+ z_offset = (self.dz - self.current_z_position) / 2
+ await self.mtcs.offset_camera_hexapod(x=0, y=0, z=z_offset, u=0, v=0)
+ await self.mtcs.offset_m2_hexapod(x=0, y=0, z=z_offset, u=0, v=0)
+ self.current_z_position = self.dz
+
+ self.log.info("Taking extra-focal image")
+
+ self.camera.rem.ccoods.evt_imageInOODS.flush()
+ extra_visit_id = await self.camera.take_cwfs(
+ exptime=self.exposure_time,
+ n=1,
+ group_id=supplemented_group_id,
+ filter=self.filter,
+ reason="EXTRA" + ("" if self.reason is None else f"_{self.reason}"),
+ program=self.program,
+ note=self.note,
+ )
+
+ if self.mode == Mode.TRIPLET or self.mode == Mode.PAIR:
+ self.log.debug("Waiting for images to be ingested in OODS.")
+ extra_image_ingested = False
+ while not extra_image_ingested:
+ try:
+ image_in_oods = await self.camera.rem.ccoods.evt_imageInOODS.next(
+ flush=False, timeout=self.exposure_time
+ )
+ try:
+ image_name_split = image_in_oods.obsid.split("_")
+ image_index = int(
+ f"{image_name_split[-2]}{image_name_split[-1][1:]}"
+ )
+ extra_image_ingested = image_index == extra_visit_id[0]
+ except Exception:
+ self.log.exception(
+ "Failed to parse image name into index for {image_in_oods.obsid}."
+ )
+
+ self.log.info(
+ f"Image {image_in_oods.obsid} {image_in_oods.raft} {image_in_oods.sensor} ingested."
+ )
+
+ except asyncio.TimeoutError:
+ self.log.warning(
+ "Timeout waiting for images to ingest. Continuing."
+ )
+ break
+ self.log.info("Send processing request to RA OCPS.")
+ config = {
+ "LSSTComCam-FROM-OCS_DONUTPAIR": f"{intra_visit_id[0]},{extra_visit_id[0]}"
+ }
+ ocps_execute_task = asyncio.create_task(
+ self.ocps.cmd_execute.set_start(
+ config=json.dumps(config),
+ timeout=self.camera.fast_timeout,
+ )
+ )
+
+ self.log.debug("Moving to in-focus position")
+
+ # Move the camera and M2 hexapods to the target z position
+ # Offset split in half and shared between each hexapod
+ z_offset = (-self.current_z_position) / 2
+ await self.mtcs.offset_camera_hexapod(x=0, y=0, z=z_offset, u=0, v=0)
+ await self.mtcs.offset_m2_hexapod(x=0, y=0, z=z_offset, u=0, v=0)
+ self.current_z_position = 0
+
+ if self.mode != Mode.PAIR:
+ self.log.info("Taking in-focus image")
+ self.camera.rem.ccoods.evt_imageInOODS.flush()
+ await self.camera.take_acq(
+ exptime=self.exposure_time,
+ n=1,
+ group_id=self.group_id,
+ filter=self.filter,
+ reason="INFOCUS" + ("" if self.reason is None else f"_{self.reason}"),
+ program=self.program,
+ note=self.note,
+ )
+
+ if self.mode == Mode.TRIPLET:
+ try:
+ await ocps_execute_task
+ except Exception:
+ self.log.exception("Executing OCPS task failed. Ignoring.")
+
+ async def run_block(self) -> None:
+ """Execute script operations."""
+ await self.assert_feasibility()
+
+ for i in range(self.n_sequences):
+ self.log.info(f"Starting aos sequence {i+1} of {self.n_sequences}")
+ await self.checkpoint(f"out-of-focus sequence {i+1} of {self.n_sequences}")
+
+ await self.take_aos_sequence()
diff --git a/tests/test_maintel_take_aos_sequence_balanced_comcam.py b/tests/test_maintel_take_aos_sequence_balanced_comcam.py
new file mode 100644
index 000000000..601b18898
--- /dev/null
+++ b/tests/test_maintel_take_aos_sequence_balanced_comcam.py
@@ -0,0 +1,268 @@
+# This file is part of ts_standardscripts
+#
+# Developed for the LSST Telescope and Site Systems.
+# This product includes software developed by the LSST Project
+# (https://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program. If not, see .
+
+import types
+import unittest
+from unittest.mock import patch
+
+from lsst.ts import standardscripts
+from lsst.ts.idl.enums.Script import ScriptState
+from lsst.ts.observatory.control.maintel.comcam import ComCam, ComCamUsages
+from lsst.ts.observatory.control.maintel.mtcs import MTCS, MTCSUsages
+from lsst.ts.standardscripts.maintel import Mode, TakeAOSSequenceBalancedComCam
+from lsst.ts.utils import index_generator
+
+index_gen = index_generator()
+
+
+class TestTakeAOSSequenceBalancedComCam(
+ standardscripts.BaseScriptTestCase, unittest.IsolatedAsyncioTestCase
+):
+ async def basic_make_script(self, index):
+ self.script = TakeAOSSequenceBalancedComCam(index=index)
+
+ self.script.mtcs = MTCS(
+ domain=self.script.domain,
+ intended_usage=MTCSUsages.DryTest,
+ log=self.script.log,
+ )
+
+ self.script.camera = ComCam(
+ domain=self.script.domain,
+ intended_usage=ComCamUsages.DryTest,
+ log=self.script.log,
+ )
+
+ self.script.ocps = unittest.mock.AsyncMock()
+
+ self.script.mtcs.offset_camera_hexapod = unittest.mock.AsyncMock()
+ self.script.camera.expose = unittest.mock.AsyncMock(
+ side_effect=self._get_visit_id
+ )
+ self.script.camera.setup_instrument = unittest.mock.AsyncMock()
+ self.script.camera.rem.ccoods = unittest.mock.AsyncMock()
+ self.script.camera.rem.ccoods.configure_mock(
+ **{
+ "evt_imageInOODS.next.side_effect": self._get_next_image_in_oods,
+ }
+ )
+
+ self._dayobs = 2024111900000
+ self._visit_index = next(index_gen)
+
+ return (self.script,)
+
+ async def _get_visit_id(self, *args, **kwargs):
+ self._visit_index = next(index_gen)
+ return [self._dayobs + self._visit_index]
+
+ async def _get_next_image_in_oods(self, *args, **kwargs):
+ return types.SimpleNamespace(
+ obsid=f"CC_O_{int(self._dayobs/100000)}_{self._visit_index:06d}",
+ raft=0,
+ sensor=0,
+ )
+
+ async def test_configure(self):
+ async with self.make_script():
+ exposure_time = 15.0
+ filter = "g"
+ dz = 2000.0
+ n_sequences = 15
+ mode = "INTRA"
+
+ await self.configure_script(
+ filter=filter,
+ exposure_time=exposure_time,
+ dz=dz,
+ n_sequences=n_sequences,
+ mode=mode,
+ )
+ assert self.script.exposure_time == exposure_time
+ assert self.script.filter == filter
+ assert self.script.dz == 2000.0
+ assert self.script.n_sequences == n_sequences
+ assert self.script.mode == Mode.INTRA
+
+ async def test_configure_ignore(self):
+ async with self.make_script():
+ self.script.mtcs.check.mtmount = True
+ self.script.mtcs.check.mtrotator = True
+ self.script.mtcs.check.mtm2 = True
+ self.script.camera.check.ccoods = True
+
+ exposure_time = 15.0
+ filter = "g"
+ dz = 2000.0
+ n_sequences = 15
+ mode = "INTRA"
+ ignore = ["mtrotator", "mtm2", "ccoods"]
+
+ await self.configure_script(
+ filter=filter,
+ exposure_time=exposure_time,
+ dz=dz,
+ n_sequences=n_sequences,
+ ignore=ignore,
+ mode=mode,
+ )
+ assert self.script.exposure_time == exposure_time
+ assert self.script.filter == filter
+ assert self.script.dz == 2000.0
+ assert self.script.n_sequences == n_sequences
+ assert self.script.mode == Mode.INTRA
+ assert self.script.mtcs.check.mtmount
+ assert not self.script.mtcs.check.mtrotator
+ assert not self.script.mtcs.check.mtm2
+ assert not self.script.camera.check.ccoods
+
+ async def run_take_triplets_test(
+ self, mock_ready_to_take_data=None, expect_exception=None
+ ):
+ async with self.make_script():
+ self.script.camera.ready_to_take_data = mock_ready_to_take_data
+
+ exposure_time = 15.0
+ filter = "g"
+ dz = 2000.0
+ n_sequences = 3
+ mode = "TRIPLET"
+
+ await self.configure_script(
+ filter=filter,
+ exposure_time=exposure_time,
+ dz=dz,
+ n_sequences=n_sequences,
+ mode=mode,
+ )
+
+ # Wrap `take_cwfs` and `take_acq` to count calls
+ with patch.object(
+ self.script.camera, "take_cwfs", wraps=self.script.camera.take_cwfs
+ ) as mock_take_cwfs, patch.object(
+ self.script.camera, "take_acq", wraps=self.script.camera.take_acq
+ ) as mock_take_acq:
+
+ if expect_exception is not None:
+ await self.run_script(expected_final_state=ScriptState.FAILED)
+ self.assertEqual(self.script.state.state, ScriptState.FAILED)
+ self.assertIn(
+ str(mock_ready_to_take_data.side_effect),
+ self.script.state.reason,
+ )
+ # the first image taken is type cwfs and in this case
+ # it should throw and exception for TCS not being ready
+ expected_take_cwfs_calls = 1
+ expected_take_acq_calls = 0
+ else:
+ await self.run_script()
+ self.assertEqual(self.script.state.state, ScriptState.DONE)
+ expected_take_cwfs_calls = n_sequences * 2
+ expected_take_acq_calls = n_sequences
+
+ expected_tcs_ready_calls = (
+ expected_take_cwfs_calls + expected_take_acq_calls
+ )
+ if expected_take_acq_calls > 0:
+ # number of calls to the expose method
+ # in BaseCamera.take_imgtype
+ expected_expose_calls = expected_tcs_ready_calls
+ else:
+ expected_expose_calls = 0
+
+ if mock_ready_to_take_data is not None:
+ self.assertEqual(
+ mock_ready_to_take_data.await_count,
+ expected_tcs_ready_calls,
+ f"ready_to_take_data was called {mock_ready_to_take_data.await_count} times, "
+ f"expected {expected_tcs_ready_calls}",
+ )
+ else:
+ with self.assertRaises(AttributeError):
+ self.script.camera.ready_to_take_data.assert_not_called()
+
+ self.assertEqual(
+ self.script.camera.expose.await_count,
+ expected_expose_calls,
+ f"expose was called {self.script.camera.expose.await_count} times, "
+ f"expected {expected_expose_calls}",
+ )
+ self.assertEqual(
+ mock_take_cwfs.await_count,
+ expected_take_cwfs_calls,
+ f"take_cwfs was called {mock_take_cwfs.await_count} times, "
+ f"expected {expected_take_cwfs_calls}",
+ )
+ self.assertEqual(
+ mock_take_acq.await_count,
+ expected_take_acq_calls,
+ f"take_acq was called {mock_take_acq.await_count} times, "
+ f"expected {expected_take_acq_calls}",
+ )
+
+ async def test_take_triplets(self):
+ await self.run_take_triplets_test()
+
+ async def test_take_triplets_tcs_ready(self):
+ mock_ready = unittest.mock.AsyncMock(return_value=None)
+ await self.run_take_triplets_test(
+ mock_ready_to_take_data=mock_ready,
+ )
+
+ async def test_take_triplets_tcs_not_ready(self):
+ mock_ready = unittest.mock.AsyncMock(side_effect=RuntimeError("TCS not ready"))
+ await self.run_take_triplets_test(
+ mock_ready_to_take_data=mock_ready, expect_exception=RuntimeError
+ )
+
+ async def test_take_doublet(self):
+ async with self.make_script():
+ self.script.camera.take_cwfs = unittest.mock.AsyncMock()
+ self.script.camera.take_acq = unittest.mock.AsyncMock()
+
+ exposure_time = 15.0
+ filter = "g"
+ dz = 2000.0
+ n_sequences = 3
+ mode = "INTRA"
+
+ await self.configure_script(
+ filter=filter,
+ exposure_time=exposure_time,
+ dz=dz,
+ n_sequences=n_sequences,
+ mode=mode,
+ )
+
+ await self.run_script()
+
+ assert n_sequences == self.script.camera.take_cwfs.await_count
+ assert n_sequences == self.script.camera.take_acq.await_count
+
+ async def test_executable_lsstcam(self) -> None:
+ """Test that the script is executable."""
+ scripts_dir = standardscripts.get_scripts_dir()
+ script_path = scripts_dir / "maintel" / "take_aos_sequence_balanced_comcam.py"
+ await self.check_executable(script_path)
+
+
+if __name__ == "__main__":
+ unittest.main()