Coverage for core / src / sensorkit / std / sensor.py: 89%

1# SPDX-License-Identifier: Apache-2.0 

2from __future__ import annotations 

3 

4import asyncio 

5import contextlib 

6from collections.abc import Coroutine 

7from typing import Any, Literal 

8 

9from loguru import logger 

10from pydantic import BaseModel, Field 

11 

12import sensorkit.api as sk 

13from sensorkit.astro.common import AltAzPointing, RADecPointing, ReferenceFrame, SitePosition 

14from sensorkit.astro.target import CatalogTarget, FrameTarget, ICRSTarget, TLETarget 

15from sensorkit.std.collect import Collect, StandardCollectTask 

16from sensorkit.std.enclosure import CloseEnclosure, OpenEnclosure 

17from sensorkit.std.instrument import Binning, CameraCapture, ConfigureCameraSensor 

18from sensorkit.std.mount import AxisRates, FollowTarget 

19from sensorkit.std.optics import CloseMirrorCover, OpenMirrorCover, SetFilter 

20from sensorkit.std.traits import Connect, Deinit, Init, Stop 

21from sensorkit.std.weather import BasicWeather 

22 

23 

24class Sensor: 

25 """High-level client to a set of devices comprising a logical sensor.""" 

26 

27 def __init__(self, impl: sk.ControllerImpl, devices: SensorDevices, policies: SensorPolicies): 

28 self.policies = policies 

29 self.mount = impl.use_device( 

30 devices.mount, 

31 subscribe=[AltAzPointing, RADecPointing, AxisRates], 

32 ) if devices.mount else None 

33 self.camera = impl.use_device(devices.camera) if devices.camera else None 

34 self.focuser = impl.use_device(devices.focuser) if devices.focuser else None 

35 self.rotator = impl.use_device(devices.rotator) if devices.rotator else None 

36 self.mirror_cover = impl.use_device(devices.mirror_cover) if devices.mirror_cover else None 

37 self.filter_wheel = impl.use_device(devices.filter_wheel) if devices.filter_wheel else None 

38 self.dome = impl.use_device(devices.dome) if devices.dome else None 

39 self.weather = impl.use_device(devices.weather, 

40 subscribe=[BasicWeather], 

41 ) if devices.weather else None 

42 

43 async def init_dome(self): 

44 """Initialize the dome and open it, if one is configured.""" 

45 if not self.dome: 

46 return 

47 

48 if self.policies.concurrent_dome_init_open: 

49 logger.info("Initializing and opening the dome") 

50 async with asyncio.TaskGroup() as tg: 

51 tg.create_task( 

52 asyncio.wait_for( 

53 self.dome.command(Init()), 

54 self.policies.dome_init_timeout, 

55 ) 

56 ) 

57 tg.create_task( 

58 asyncio.wait_for( 

59 self.dome.command(OpenEnclosure()), 

60 self.policies.dome_open_close_timeout, 

61 ) 

62 ) 

63 return 

64 

65 logger.info("Initializing the dome") 

66 async with asyncio.timeout(self.policies.dome_init_timeout): 

67 await self.dome.command(Init()) 

68 

69 logger.info("Opening the dome") 

70 async with asyncio.timeout(self.policies.dome_open_close_timeout): 

71 await self.dome.command(OpenEnclosure()) 

72 

73 async def deinit_dome(self): 

74 """Stop the dome, close it, and deinitialize it, if one is configured.""" 

75 if not self.dome: 

76 return 

77 

78 # Halt any leftover motion before closing, so nothing aborts the close mid-flight. 

79 with contextlib.suppress(Exception): 

80 await self.dome.command(Stop()) 

81 

82 if self.policies.concurrent_dome_deinit_close: 

83 logger.info("Closing and deinitializing the dome") 

84 async with asyncio.TaskGroup() as tg: 

85 tg.create_task( 

86 asyncio.wait_for( 

87 self.dome.command(CloseEnclosure()), 

88 self.policies.dome_open_close_timeout, 

89 ) 

90 ) 

91 tg.create_task( 

92 asyncio.wait_for( 

93 self.dome.command(Deinit()), 

94 self.policies.dome_deinit_timeout, 

95 ) 

96 ) 

97 return 

98 

99 logger.info("Closing the dome") 

100 async with asyncio.timeout(self.policies.dome_open_close_timeout): 

101 await self.dome.command(CloseEnclosure()) 

102 

103 logger.info("Deinitializing the dome") 

104 async with asyncio.timeout(self.policies.dome_deinit_timeout): 

105 await self.dome.command(Deinit()) 

106 

107 async def init_mount(self): 

108 """Initialize the mount.""" 

109 if self.mount: 

110 logger.info("Initializing the mount") 

111 async with asyncio.timeout(self.policies.mount_init_timeout): 

112 await self.mount.command(Init()) 

113 

114 async def deinit_mount(self): 

115 """Deinitialize the mount.""" 

116 if self.mount: 

117 logger.info("Deinitializing the mount") 

118 await self.mount.command(Deinit()) 

119 

120 async def init_mirror_cover(self): 

121 """Open the mirror cover if one is configured.""" 

122 if self.mirror_cover: 

123 # Start opening the mirror cover. 

124 logger.info("Opening mirror cover") 

125 async with asyncio.timeout(self.policies.mirror_cover_open_close_timeout): 

126 await self.mirror_cover.command(OpenMirrorCover()) 

127 

128 async def deinit_mirror_cover(self): 

129 """Close the mirror cover if one is configured.""" 

130 if self.mirror_cover: 

131 # Start closing the mirror cover. 

132 logger.info("Closing mirror cover") 

133 async with asyncio.timeout(self.policies.mirror_cover_open_close_timeout): 

134 await self.mirror_cover.command(CloseMirrorCover()) 

135 

136 async def init_all(self, tg: asyncio.TaskGroup): 

137 """Init the mount and open the dome and mirror cover, if configured.""" 

138 tasks = [] 

139 

140 # Open the dome, if any. 

141 tasks.append(tg.create_task(self.init_dome())) 

142 

143 if not self.policies.concurrent_dome_and_mount_init: 

144 await asyncio.wait(tasks) 

145 

146 # Initialize the mount. 

147 tasks.append(tg.create_task(self.init_mount())) 

148 

149 if not self.policies.concurrent_mount_and_mirror_cover_init: 

150 await asyncio.wait(tasks) 

151 

152 # Open the mirror cover, if any. 

153 tasks.append(tg.create_task(self.init_mirror_cover())) 

154 

155 # Wait for all operations to complete. 

156 await asyncio.wait(tasks) 

157 

158 async def stop_all(self): 

159 """Issue Stop commands to the mount, dome, and mirror cover, suppressing any errors.""" 

160 if self.mount: 

161 with contextlib.suppress(Exception): 

162 await self.mount.command(Stop()) 

163 

164 if self.dome: 

165 with contextlib.suppress(Exception): 

166 await self.dome.command(Stop()) 

167 

168 if self.mirror_cover: 

169 with contextlib.suppress(Exception): 

170 await self.mirror_cover.command(Stop()) 

171 

172 

173@sk.declare_controller 

174class SensorControl: 

175 """A Controller that controls a mount and camera.""" 

176 

177 def __init__(self, config: SensorConfig): 

178 self.config = config 

179 

180 @sk.on_attach 

181 async def controller_init(self): 

182 controller = sk.controller() 

183 

184 # TODO: Phase out when UI code is updated to use ControllerInfo and SensorConfig for this 

185 # info. 

186 await controller.kv_put_model( 

187 Capabilities( 

188 tasks=[h.__name__ for h in controller._task_handlers.keys()], 

189 devices=self.config.devices, 

190 ) 

191 ) 

192 

193 self.sensor = Sensor( 

194 controller, 

195 self.config.devices, 

196 self.config.policies, 

197 ) 

198 

199 await controller.kv_put_model(self.config.site_position) 

200 

201 @sk.task_handler 

202 async def sensor_init(self, task: sk.InitTask): 

203 """Attempt to start the sensor.""" 

204 try: 

205 async with asyncio.TaskGroup() as tg: 

206 await self.sensor.init_all(tg) 

207 except* BaseException: 

208 with contextlib.suppress(BaseException): 

209 await self.sensor.stop_all() 

210 

211 raise 

212 

213 logger.info(f"Sensor '{sk.controller().entity}' is ready to operate") 

214 

215 @sk.task_handler 

216 async def sensor_standby(self, task: sk.StandbyTask): 

217 """Put the sensor in standby mode.""" 

218 # FIXME: Presently this is a synonym for init. Semantics should be dictated by config. 

219 try: 

220 async with asyncio.TaskGroup() as tg: 

221 await self.sensor.init_all(tg) 

222 except* BaseException: 

223 with contextlib.suppress(BaseException): 

224 await self.sensor.stop_all() 

225 

226 raise 

227 

228 logger.info(f"Sensor '{sk.controller().entity}' is standing by") 

229 

230 @sk.task_handler 

231 async def sensor_collect(self, task: StandardCollectTask): 

232 """Execute a StandardCollectTask: slew, configure camera, capture frames.""" 

233 if not self.sensor.mount or not self.sensor.camera: 

234 raise RuntimeError("Standard collect requires a mount and a camera") 

235 

236 # Perform the device commands to do the collect! 

237 logger.info("Moving to target") 

238 await self.sensor.mount.command(FollowTarget(target=task.target)) 

239 

240 logger.info("Reached target") 

241 

242 if task.camera_params.filter_name is not None and self.sensor.filter_wheel is not None: 

243 await self.sensor.filter_wheel.command( 

244 SetFilter(filter=task.camera_params.filter_name) 

245 ) 

246 

247 # Configure camera capture parameters. 

248 if task.camera_params.readout_mode is not None: 

249 await self.sensor.camera.command( 

250 ConfigureCameraSensor(readout_mode=task.camera_params.readout_mode) 

251 ) 

252 

253 if task.camera_params.gain is not None: 

254 await self.sensor.camera.command( 

255 ConfigureCameraSensor(gain=task.camera_params.gain) 

256 ) 

257 

258 if None not in (task.camera_params.binning_x, task.camera_params.binning_y): 

259 await self.sensor.camera.command( 

260 ConfigureCameraSensor( 

261 binning=Binning( 

262 x=task.camera_params.binning_x, 

263 y=task.camera_params.binning_y, 

264 ), 

265 ) 

266 ) 

267 

268 # Set base context for all frames. 

269 collect = Collect( 

270 target=task.target, 

271 params=task.camera_params, 

272 target_id=( 

273 task.target_id 

274 if task.target_id 

275 else task.target.tle.norad_id 

276 if isinstance(task.target, TLETarget) 

277 else task.target.object 

278 if isinstance(task.target, CatalogTarget) 

279 else None 

280 ), 

281 ) 

282 await sk.controller().update_context(self.config.site_position, collect) 

283 

284 # For inherently sidereal targets (stars), the initial FollowTarget already 

285 # establishes sidereal tracking — mark as sidereal from the start so the frame 

286 # loop never issues a redundant tracking switch. 

287 target_is_sidereal = isinstance(task.target, (ICRSTarget, CatalogTarget)) 

288 currently_sidereal = target_is_sidereal 

289 

290 # Capture the requested frames. 

291 for frame_num in range(0, task.camera_params.frame_count): 

292 collect.frame_number = frame_num 

293 want_sidereal = target_is_sidereal or frame_num in task.sidereal_frames 

294 

295 if want_sidereal and not currently_sidereal: 

296 # Hold the current RA/Dec under sidereal tracking. 

297 logger.info(f"Frame #{frame_num+1} of {task.camera_params.frame_count}: switching to sidereal track") 

298 collect.target = FrameTarget(frame=ReferenceFrame.ICRF) 

299 await self.sensor.mount.command(FollowTarget(target=collect.target)) 

300 currently_sidereal = True 

301 elif not want_sidereal and currently_sidereal: 

302 # Resume following the original target. 

303 logger.info(f"Frame #{frame_num+1} of {task.camera_params.frame_count}: resuming target track") 

304 collect.target = task.target 

305 await self.sensor.mount.command(FollowTarget(target=collect.target)) 

306 currently_sidereal = False 

307 

308 logger.info(f"Acquiring frame #{frame_num+1} of {task.camera_params.frame_count}") 

309 context = await sk.controller().update_context(collect) 

310 _add_compat_context(context) 

311 

312 await self.sensor.camera.command( 

313 CameraCapture( 

314 integration_time=task.camera_params.integration_time_seconds, 

315 frame_type=task.camera_params.frame_type, 

316 context=context, 

317 ) 

318 ) 

319 

320 # Stop motion when collection completes 

321 await self.sensor.mount.command(Stop()) 

322 

323 @sk.task_handler 

324 async def sensor_recover(self, task: sk.RecoverTask): 

325 """Reconnect to all devices and stop any in-progress motion.""" 

326 

327 def assert_success_or_unsupported(results: tuple[Any]): 

328 logger.debug(f"{results=}") 

329 

330 for result in results: 

331 if isinstance(result, BaseException) and ( 

332 not isinstance(result, sk.CallError) or "Request rejected" not in str(result) 

333 ): 

334 raise result 

335 

336 # Try to reconnect to all devices that support it. 

337 devices = sk.controller().all_devices() 

338 

339 logger.info("Reconnecting to devices...") 

340 assert_success_or_unsupported( 

341 await asyncio.gather( 

342 *(dev.client.command(Connect()) for dev in devices), 

343 return_exceptions=True, 

344 ) 

345 ) 

346 

347 # Try to stop all devices that support it. 

348 logger.info("Stopping device activity...") 

349 assert_success_or_unsupported( 

350 await asyncio.gather( 

351 *(dev.client.command(Stop()) for dev in devices), 

352 return_exceptions=True, 

353 ) 

354 ) 

355 

356 @sk.task_handler 

357 async def sensor_shutdown(self, task: sk.ShutdownTask): 

358 """Shut down the sensor. 

359 

360 Under the always_deinit_dome policy every step is attempted even if an earlier one 

361 fails, so the dome comes down regardless; the failures are raised once it is closed. 

362 Otherwise the first failure ends the shutdown. 

363 """ 

364 policies = self.config.policies 

365 failures: list[Exception] = [] 

366 

367 async def attempt(step: Coroutine[Any, Any, None]): 

368 """Run a shutdown step, recording rather than propagating its failure.""" 

369 try: 

370 await step 

371 except Exception as exc: 

372 logger.exception("Sensor shutdown step failed") 

373 failures.append(exc) 

374 

375 def raise_failures(): 

376 match failures: 

377 case []: 

378 return 

379 case [failure]: 

380 raise failure 

381 case _: 

382 raise ExceptionGroup("Sensor shutdown failed", failures) 

383 

384 await attempt(self.sensor.deinit_mirror_cover()) 

385 

386 if not policies.always_deinit_dome: 

387 raise_failures() 

388 

389 if policies.concurrent_dome_and_mount_deinit: 

390 # attempt() absorbs failures, so both steps always run to completion. 

391 await asyncio.gather( 

392 attempt(self.sensor.deinit_dome()), 

393 attempt(self.sensor.deinit_mount()), 

394 ) 

395 else: 

396 await attempt(self.sensor.deinit_mount()) 

397 

398 if not policies.always_deinit_dome: 

399 raise_failures() 

400 

401 await attempt(self.sensor.deinit_dome()) 

402 

403 raise_failures() 

404 

405 

406class SensorDevices(BaseModel): 

407 """Device entity references for sensor control.""" 

408 

409 mount: str | None = None 

410 camera: str | None = None 

411 focuser: str | None = None 

412 rotator: str | None = None 

413 filter_wheel: str | None = None 

414 mirror_cover: str | None = None 

415 dome: str | None = None 

416 weather: str | None = None 

417 

418 

419class SensorPolicies(BaseModel): 

420 """Policies for sensor control.""" 

421 

422 concurrent_dome_and_mount_init: bool = False 

423 """Whether the dome and mount can be initialized concurrently.""" 

424 

425 concurrent_dome_and_mount_deinit: bool = False 

426 """Whether the dome and mount can be deinitialized concurrently.""" 

427 

428 concurrent_dome_init_open: bool = False 

429 """Whether the dome can be initialized and opened concurrently.""" 

430 

431 concurrent_dome_deinit_close: bool = False 

432 """Whether the dome can be closed and deinitialized concurrently.""" 

433 

434 always_deinit_dome: bool = False 

435 """Whether the dome should be deinitialized even if other parts of deinitialization fail.""" 

436 

437 dome_open_close_timeout: float = 120.0 

438 """Timeout for fully opening or closing the dome.""" 

439 

440 dome_init_timeout: float = 300.0 

441 """Timeout for initializing the dome, including homing if required.""" 

442 

443 dome_deinit_timeout: float = 300.0 

444 """Timeout for deinitializing the dome, including parking if required.""" 

445 

446 concurrent_mount_and_mirror_cover_init: bool = False 

447 """Whether the mount and mirror cover can be initialized concurrently.""" 

448 

449 mirror_cover_open_close_timeout: float = 60.0 

450 """Timeout for fully opening or closing the mirror cover.""" 

451 

452 mount_init_timeout: float = 30.0 

453 """Timeout for powering the mount and enabling axis control.""" 

454 

455 mount_home_timeout: float = 300.0 

456 """Timeout for homing the mount.""" 

457 

458 minimum_target_altitude_degrees: float | None = None 

459 """Minimum altitude for a target to be tracked during collection.""" 

460 

461 sun_separation_degrees: float | None = None 

462 """Distance a target must be away from the sun during collection.""" 

463 

464 moon_separation_degrees: float | None = None 

465 """Distance a target must be away from the moon during collection""" 

466 

467 

468class SensorConfig(BaseModel): 

469 """Configuration for standard sensor control.""" 

470 

471 controller_name: str 

472 devices: SensorDevices 

473 site_position: SitePosition 

474 policies: SensorPolicies = Field(default_factory=SensorPolicies) 

475 

476 

477sk.declare_config_section( 

478 "sensors", 

479 list[SensorConfig], 

480 id_source="by_subkey", 

481 service_path=__name__, 

482) 

483 

484 

485# TODO: Phase out when UI code is updated to use ControllerInfo and SensorConfig for this info. 

486class Capabilities(BaseModel): 

487 """Deprecated controller capability descriptor for the sensor service.""" 

488 

489 type: Literal["controller"] = "controller" 

490 tasks: list[str] 

491 devices: SensorDevices 

492 

493 

494@sk.service_entrypoint(version=sk.VERSION) 

495async def sensor_control_service(service: sk.Service): 

496 await service.register() 

497 

498 try: 

499 # Read configuration. 

500 config = await service.context.kv_get_model(SensorConfig) 

501 except sk.KVError as e: 

502 logger.error(f"Service {service.name} could not get SensorConfig ({e})") 

503 return 

504 

505 sc = SensorControl(config=config) 

506 service.include( 

507 sc, 

508 name=config.controller_name, 

509 ) 

510 

511 await service.run() 

512 

513 

514def _add_compat_context(context: sk.Context): 

515 # FIXME: Temporary to avoid breaking existing deployed config. 

516 compat = dict( 

517 ra="RADecPointing.right_ascension_hours * 15", 

518 ra_hms="RADecPointing.ra_hms", 

519 ra_rate="AxisRates.right_ascension.velocity * 3600", 

520 dec="RADecPointing.declination_degrees", 

521 dec_dms="RADecPointing.dec_dms", 

522 dec_rate="AxisRates.declination.velocity * 3600", 

523 alt="AltAzPointing.altitude_degrees", 

524 alt_rate="AxisRates.altitude.velocity * 3600", 

525 az="AltAzPointing.azimuth_degrees", 

526 az_rate="AxisRates.azimuth.velocity * 3600", 

527 track_mode="Collect.track_mode", 

528 target_id="Collect.target_id", 

529 target_name="Collect.target.tle.line0 if Collect.target.target_type == 'tle' else target_id", 

530 tle_line0="Collect.target.tle.line0 if Collect.target.target_type == 'tle' else None", 

531 tle_line1="Collect.target.tle.line1 if Collect.target.target_type == 'tle' else None", 

532 tle_line2="Collect.target.tle.line2 if Collect.target.target_type == 'tle' else None", 

533 frame_num="Collect.frame_number", 

534 frame_count="Collect.params.frame_count", 

535 integration_time_seconds="Collect.params.integration_time_seconds", 

536 binning_x="Collect.params.binning_x", 

537 binning_y="Collect.params.binning_y", 

538 filter_name="Collect.params.filter_name", 

539 elevation="SitePosition.altitude_km", 

540 latitude="SitePosition.latitude_degrees", 

541 longitude="SitePosition.longitude_degrees", 

542 task_id="TaskInfo.task_id", 

543 ) 

544 

545 for key, expr in compat.items(): 

546 with contextlib.suppress(Exception): 

547 context.set_value(key, context.eval(expr)) 

548 

549 # Fold legacy flat file_name/file_path keys into FileNameTemplate / FileInfo keywords. 

550 # A file_name is an input naming template; a file_path is an explicit output location. 

551 import pathlib 

552 

553 from sensorkit.data.filesys import FileInfo, FileNameTemplate 

554 

555 file_name = context.pop("file_name", None) 

556 file_path = context.pop("file_path", None) 

557 

558 if file_name is not None: 

559 logger.warning("The 'file_name' context key is deprecated. Use FileNameTemplate instead.") 

560 context.set(FileNameTemplate(template=file_name)) 

561 

562 if file_path is not None: 

563 logger.warning("The 'file_path' context key is deprecated. Use FileInfo instead.") 

564 context.set(FileInfo(path=pathlib.Path(file_path)))