Coverage for python/lsst/pipe/base/prerequisite_helpers.py: 33%

202 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-30 02:55 -0700

1# This file is part of pipe_base. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This software is dual licensed under the GNU General Public License and also 

10# under a 3-clause BSD license. Recipients may choose which of these licenses 

11# to use; please see the files gpl-3.0.txt and/or bsd_license.txt, 

12# respectively. If you choose the GPL option then the following text applies 

13# (but note that there is still no warranty even if you opt for BSD instead): 

14# 

15# This program is free software: you can redistribute it and/or modify 

16# it under the terms of the GNU General Public License as published by 

17# the Free Software Foundation, either version 3 of the License, or 

18# (at your option) any later version. 

19# 

20# This program is distributed in the hope that it will be useful, 

21# but WITHOUT ANY WARRANTY; without even the implied warranty of 

22# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

23# GNU General Public License for more details. 

24# 

25# You should have received a copy of the GNU General Public License 

26# along with this program. If not, see <http://www.gnu.org/licenses/>. 

27 

28"""Helper classes for finding prerequisite input datasets during 

29QuantumGraph generation. 

30""" 

31 

32from __future__ import annotations 

33 

34__all__ = ( 

35 "SkyPixBoundsBuilder", 

36 "PrerequisiteFinder", 

37 "PrerequisiteBounds", 

38 "TimespanBuilder", 

39 "PrerequisiteInfo", 

40) 

41 

42import dataclasses 

43from abc import ABC, abstractmethod 

44from collections.abc import Callable, Iterable, Mapping, Sequence 

45from typing import cast 

46 

47from lsst.daf.butler import ( 

48 Butler, 

49 DataCoordinate, 

50 DatasetRef, 

51 DatasetType, 

52 DimensionElement, 

53 Registry, 

54 SkyPixDimension, 

55 Timespan, 

56) 

57from lsst.daf.butler.registry import MissingDatasetTypeError 

58from lsst.sphgeom import RangeSet, Region 

59 

60from .pipeline_graph import DatasetTypeNode, PipelineGraph, ReadEdge, TaskNode 

61 

62 

63@dataclasses.dataclass 

64class PrerequisiteInfo: 

65 """A QuantumGraph-generation helper class that manages the searches for all 

66 prerequisite input connections for a task. 

67 

68 Parameters 

69 ---------- 

70 task_node : `TaskNode` 

71 The relevant node. 

72 pipeline_graph : `PipelineGraph` 

73 The pipeline graph. 

74 """ 

75 

76 bounds: PrerequisiteBounds 

77 """Another helper object that manages the spatial/temporal bounds of the 

78 task's quanta. 

79 """ 

80 

81 finders: dict[str, PrerequisiteFinder] 

82 """Mapping of helper objects responsible for a single prerequisite input 

83 connection. 

84 

85 Keys are connection names. Elements of this dictionary should be removed 

86 by implementations of `QuantumGraphBuilder.process_subgraph` to take 

87 responsibility for finding them away from the the `QuantumGraphBuilder` 

88 base class. 

89 """ 

90 

91 def __init__(self, task_node: TaskNode, pipeline_graph: PipelineGraph): 

92 self.bounds = PrerequisiteBounds(task_node) 

93 self.finders = { 

94 edge.connection_name: PrerequisiteFinder(edge, self.bounds, pipeline_graph) 

95 for edge in task_node.prerequisite_inputs.values() 

96 } 

97 

98 def update_bounds(self) -> None: 

99 """Inspect the current state of `finders` and update `bounds` to 

100 reflect the needs of only the finders that remain. 

101 """ 

102 self.bounds.all_dataset_skypix.clear() 

103 self.bounds.any_dataset_has_timespan = False 

104 for finder in self.finders.values(): 

105 self.bounds.all_dataset_skypix.update(finder.dataset_skypix) 

106 self.bounds.any_dataset_has_timespan = ( 

107 self.bounds.any_dataset_has_timespan or finder.dataset_has_timespan 

108 ) 

109 

110 

111class PrerequisiteFinder: 

112 """A QuantumGraph-generation helper class that manages the searches for a 

113 prerequisite input connection. 

114 

115 Parameters 

116 ---------- 

117 edge : `pipeline_graph.ReadEdge` 

118 A `~pipeline_graph.PipelineGraph` edge that represents a single 

119 prerequisite input connection. 

120 bounds : `PrerequisiteBounds` 

121 Another helper object that manages the spatial/temporal bounds of the 

122 task's quanta, shared by all prerequisite inputs for that task. 

123 pipeline_graph : `pipeline_graph.PipelineGraph` 

124 Graph representation of the pipeline. 

125 

126 Notes 

127 ----- 

128 `PrerequisiteFinder` instances are usually constructed by a 

129 `PrerequisiteInfo` instance, which is in turn constructed by and attached 

130 to the base `QuantumGraphBuilder` when a new builder is constructed. During 

131 the `QuantumGraphBuilder.process_subgraph` hook implemented by a builder 

132 subclass, prerequisite inputs may be found in other ways (e.g. via bulk 

133 queries), as long as the results are consistent with the finder's 

134 attributes, and this is indicated to the base `QuantumGraphBuilder` by 

135 removing those finder instances after those prerequisites have been found 

136 and added to a `QuantumGraphSkeleton`. Finder instances that remain in the 

137 builder are used by calling `PrerequisiteFinder.find` on each quantum 

138 later in `QuantumGraphBuilder.build`. 

139 """ 

140 

141 def __init__( 

142 self, 

143 edge: ReadEdge, 

144 bounds: PrerequisiteBounds, 

145 pipeline_graph: PipelineGraph, 

146 ): 

147 self.edge = edge 

148 self._bounds = bounds 

149 self.dataset_type_node = pipeline_graph.dataset_types[edge.parent_dataset_type_name] 

150 self.lookup_function = self.task_node.get_lookup_function(edge.connection_name) 

151 self.dataset_skypix = {} 

152 self.dataset_other_spatial = {} 

153 self.dataset_has_timespan = False 

154 self.constraint_dimensions = self.task_node.dimensions 

155 if self.lookup_function is None: 

156 for family in self.dataset_type_node.dimensions.spatial - self.task_node.dimensions.spatial: 

157 best_spatial_element = family.choose( 

158 self.dataset_type_node.dimensions.elements.names, 

159 self.dataset_type_node.dimensions.universe, 

160 ) 

161 if isinstance(best_spatial_element, SkyPixDimension): 

162 self.dataset_skypix[best_spatial_element.name] = best_spatial_element 

163 else: 

164 self.dataset_other_spatial[best_spatial_element.name] = cast( 

165 DimensionElement, best_spatial_element 

166 ) 

167 self.dataset_has_timespan = bool( 

168 # If the task dimensions has a temporal family that isn't in 

169 # the dataset type (i.e. "observation_timespans", like visit 

170 # or exposure)... 

171 self.task_node.dimensions.temporal 

172 - self.dataset_type_node.dimensions.temporal 

173 ) and ( 

174 # ...and the dataset type has a temporal family that isn't in 

175 # the task dimensions, or is a calibration, the prerequisite 

176 # search needs a temporal join. Note that the default 

177 # dimension universe only has one temporal dimension family, so 

178 # in practice this just means "calibration lookups when visit 

179 # or exposure is in the task dimensions". 

180 self.dataset_type_node.is_calibration 

181 or bool(self.dataset_type_node.dimensions.temporal - self.task_node.dimensions.temporal) 

182 ) 

183 new_constraint_dimensions = set() 

184 universe = self.task_node.dimensions.universe 

185 for dimension_name in self.task_node.dimensions.names: 

186 if dimension_name in self.dataset_type_node.dimensions.names: 

187 new_constraint_dimensions.add(dimension_name) 

188 else: 

189 dimension = universe[dimension_name] 

190 if not (dimension.spatial or dimension.temporal): 

191 new_constraint_dimensions.add(dimension_name) 

192 self.constraint_dimensions = universe.conform(new_constraint_dimensions) 

193 

194 edge: ReadEdge 

195 """The `~pipeline_graph.PipelineGraph` edge that represents the 

196 prerequisite input connection. 

197 """ 

198 

199 dataset_type_node: DatasetTypeNode 

200 """The `~pipeline_graph.PipelineGraph` node that represents the dataset 

201 type of this connection. 

202 

203 This always uses the registry storage class and is never a component 

204 dataset type. 

205 """ 

206 

207 lookup_function: ( 

208 Callable[[DatasetType, Registry, DataCoordinate, Sequence[str]], Iterable[DatasetRef]] | None 

209 ) 

210 """A task-provided callback for finding these datasets. 

211 

212 If this is not `None`, it must be used to ensure correct behavior. 

213 """ 

214 

215 dataset_skypix: dict[str, SkyPixDimension] 

216 """Dimensions representing a pixelization of the sky used by the dataset 

217 type for this connection that are also not part of the task's dimensions. 

218 

219 Keys are dimension names. It is at least extremely rare for this 

220 dictionary to have more than one element. 

221 """ 

222 

223 dataset_other_spatial: dict[str, DimensionElement] 

224 """Spatial dimensions other than sky pixelizations used by the dataset type 

225 for this connection that are also not part of the task's dimensions. 

226 """ 

227 

228 dataset_has_timespan: bool 

229 """Whether the dataset has a timespan that should be used in the lookup, 

230 either because it is a calibration dataset or because it has temporal 

231 dimensions that are not part of the tasks's dimensions. 

232 """ 

233 

234 @property 

235 def task_node(self) -> TaskNode: 

236 """The `~pipeline_graph.PipelineGraph` node that represents the task 

237 for this connection. 

238 """ 

239 return self._bounds.task_node 

240 

241 def find( 

242 self, 

243 butler: Butler, 

244 input_collections: Sequence[str], 

245 data_id: DataCoordinate, 

246 skypix_bounds: Mapping[str, RangeSet], 

247 timespan: Timespan | None, 

248 ) -> list[DatasetRef]: 

249 """Find prerequisite input datasets for a single quantum. 

250 

251 Parameters 

252 ---------- 

253 butler : `lsst.daf.butler.Butler` 

254 Butler client to use for queries. 

255 input_collections : `~collections.abc.Sequence` [ `str` ] 

256 Sequence of collections to search, in order. 

257 data_id : `lsst.daf.butler.DataCoordinate` 

258 Data ID for the quantum. 

259 skypix_bounds : `Mapping` [ `str`, `lsst.sphgeom.RangeSet` ] 

260 The spatial bounds of this quantum in various skypix dimensions. 

261 Keys are skypix dimension names (a superset of those in 

262 `dataset_skypix`) and values are sets of integer pixel ID ranges. 

263 timespan : `lsst.daf.butler.Timespan` or `None` 

264 The temporal bounds of this quantum. Guaranteed to not be `None` 

265 if `dataset_has_timespan` is `True`. 

266 

267 Returns 

268 ------- 

269 refs : `list` [ `lsst.daf.butler.DatasetRef` ] 

270 Dataset references. These use 

271 ``self.dataset_type_node.dataset_type``, which may differ from the 

272 connection's dataset type in storage class or [lack of] component. 

273 

274 Raises 

275 ------ 

276 NotImplementedError 

277 Raised for certain relationships between task and dataset type 

278 dimensions that are possible to define but not believed to be 

279 useful in practice. These errors occur late rather than early in 

280 order to allow a `QuantumGraphBuilder` subclass to handle them 

281 first, in case an unusual task's needs must be met by a custom 

282 builder class anyway. 

283 """ 

284 if self.lookup_function: 

285 # If there is a lookup function, just use it; nothing else matters. 

286 return [ 

287 self.dataset_type_node.generalize_ref(ref) 

288 for ref in self.lookup_function( 

289 self.edge.adapt_dataset_type(self.dataset_type_node.dataset_type), 

290 butler.registry, 

291 data_id, 

292 input_collections, 

293 ) 

294 if ref is not None 

295 ] 

296 if self.dataset_type_node.dimensions <= self.constraint_dimensions: 

297 # If this is a calibration dataset and the dataset doesn't have 

298 # any dimensions that aren't constrained by the quantum data 

299 # ID, we know there'll only be one result, and that means we 

300 # can call Butler.find_dataset, which takes a timespan. Note 

301 # that the AllDimensionsQuantumGraphBuilder subclass will 

302 # intercept this case in order to optimize it when: 

303 # 

304 # - PipelineTaskConnections.getTemporalBoundsConnections is 

305 # empty; 

306 # 

307 # - the quantum data IDs have temporal dimensions; 

308 # 

309 # and when that happens PrerequisiteFinder.find never gets 

310 # called. 

311 try: 

312 ref = butler.find_dataset( 

313 self.dataset_type_node.dataset_type, 

314 data_id.subset(self.constraint_dimensions), 

315 collections=input_collections, 

316 timespan=timespan, 

317 ) 

318 except MissingDatasetTypeError: 

319 ref = None 

320 return [ref] if ref is not None else [] 

321 elif self.dataset_has_timespan: 

322 extra_dimensions = self.dataset_type_node.dimensions.names - self.constraint_dimensions.names 

323 raise NotImplementedError( 

324 f"No support for calibration lookup {self.task_node.label}.{self.edge.connection_name} " 

325 f"with dimension(s) {extra_dimensions} not fully constrained by the task. " 

326 "Please create a feature-request ticket and use a lookup function in the meantime." 

327 ) 

328 if self.dataset_skypix: 

329 if not self.dataset_has_timespan and not self.dataset_other_spatial: 

330 # If the dataset has skypix dimensions but is not otherwise 

331 # spatial or temporal (this describes reference catalogs and 

332 # things like them), we can stuff the skypix IDs we want into 

333 # the query via bind parameters and call queryDatasets. Once 

334 # again AllDimensionsQuantumGraphBuilder will often intercept 

335 # this case in order to optimize it, when: 

336 # 

337 # - PipelineTaskConnections.getSpatialBoundsConnections is 

338 # empty; 

339 # 

340 # - the quantum data IDs have spatial dimensions; 

341 # 

342 # and when that happens PrerequisiteFinder.find never gets 

343 # called. 

344 where_terms: list[str] = [] 

345 bind: dict[str, list[int]] = {} 

346 for name in self.dataset_skypix: 

347 where_terms.append(f"{name} IN ({name}_pixels)") 

348 pixels: list[int] = [] 

349 for begin, end in skypix_bounds[name]: 

350 pixels.extend(range(begin, end)) 

351 bind[f"{name}_pixels"] = pixels 

352 try: 

353 return list( 

354 butler.registry.queryDatasets( 

355 self.dataset_type_node.dataset_type, 

356 collections=input_collections, 

357 dataId=data_id.subset(self.constraint_dimensions), 

358 where=" AND ".join(where_terms), 

359 bind=bind, 

360 findFirst=True, 

361 ).expanded() 

362 ) 

363 except MissingDatasetTypeError: 

364 return [] 

365 else: 

366 raise NotImplementedError( 

367 f"No support for skypix lookup {self.task_node.label}.{self.edge.connection_name} " 

368 "that requires additional spatial and/or temporal constraints. " 

369 "Please create a feature-request ticket and use a lookup function in the meantime." 

370 ) 

371 if self._bounds.spatial_connections or self._bounds.temporal_connections: 

372 raise NotImplementedError( 

373 f"No support for prerequisite lookup {self.task_node.label}.{self.edge.connection_name} " 

374 "that requires other connections to determine spatial or temporal bounds but does not " 

375 "fit into one of our standard cases. " 

376 "Please create a feature-request ticket and use a lookup function in the meantime." 

377 ) 

378 # If the spatial/temporal bounds are not customized, and the dataset 

379 # doesn't have any skypix dimensions, a vanilla queryDatasets call 

380 # should work. This case should always be optimized by 

381 # AllDimensionsQuantumGraphBuilder as well. Note that we use the 

382 # original quantum data ID here, not those with constraint_dimensions 

383 # that strips out the spatial/temporal stuff, because here we want the 

384 # butler query system to handle the spatial/temporal stuff like it 

385 # normally would. 

386 try: 

387 return list( 

388 butler.registry.queryDatasets( 

389 self.dataset_type_node.dataset_type, 

390 collections=input_collections, 

391 dataId=data_id, 

392 findFirst=True, 

393 ).expanded() 

394 ) 

395 except MissingDatasetTypeError: 

396 return [] 

397 

398 

399@dataclasses.dataclass 

400class PrerequisiteBounds: 

401 """A QuantumGraph-generation helper class that manages the spatial and 

402 temporal bounds of a tasks' quanta, for the purpose of finding 

403 prerequisite inputs. 

404 """ 

405 

406 task_node: TaskNode 

407 """The `~pipeline_graph.PipelineGraph` node that represents the task.""" 

408 

409 spatial_connections: frozenset[str] = dataclasses.field(init=False) 

410 """Regular input or output connections whose (assumed spatial) data IDs 

411 should be used to define the spatial bounds of this task's quanta. 

412 

413 See Also 

414 -------- 

415 PipelineTaskConnections.getSpatialBoundsConnections 

416 """ 

417 

418 temporal_connections: frozenset[str] = dataclasses.field(init=False) 

419 """Regular input or output connections whose (assumed temporal) data IDs 

420 should be used to define the temporal bounds of this task's quanta. 

421 

422 See Also 

423 -------- 

424 PipelineTaskConnections.getTemporalBoundsConnections 

425 """ 

426 

427 all_dataset_skypix: dict[str, SkyPixDimension] = dataclasses.field(default_factory=dict) 

428 """The union of all `PrerequisiteFinder.dataset_skypix` attributes for all 

429 (remaining) prerequisite finders for this task. 

430 """ 

431 

432 any_dataset_has_timespan: bool = dataclasses.field(default=False) 

433 """Whether any `PrerequisiteFinder.dataset_has_timespan` attribute is true 

434 for any (remaining) prerequisite finder for this task. 

435 """ 

436 

437 def __post_init__(self) -> None: 

438 self.spatial_connections = frozenset(self.task_node.get_spatial_bounds_connections()) 

439 self.temporal_connections = frozenset(self.task_node.get_temporal_bounds_connections()) 

440 

441 def make_skypix_bounds_builder(self, quantum_data_id: DataCoordinate) -> SkyPixBoundsBuilder: 

442 """Return an object that accumulates the appropriate spatial bounds for 

443 a quantum. 

444 

445 Parameters 

446 ---------- 

447 quantum_data_id : `lsst.daf.butler.DataCoordinate` 

448 Data ID for this quantum. 

449 

450 Returns 

451 ------- 

452 builder : `SkyPixBoundsBuilder` 

453 Object that accumulates the appropriate spatial bounds for a 

454 quantum. If the spatial bounds are not needed, this object will do 

455 nothing. 

456 """ 

457 if not self.all_dataset_skypix: 

458 return _TrivialSkyPixBoundsBuilder() 

459 if self.spatial_connections: 

460 return _ConnectionSkyPixBoundsBuilder( 

461 self.task_node, self.spatial_connections, self.all_dataset_skypix.values(), quantum_data_id 

462 ) 

463 if self.task_node.dimensions.spatial: 

464 return _QuantumOnlySkyPixBoundsBuilder(self.all_dataset_skypix.values(), quantum_data_id) 

465 else: 

466 return _UnboundedSkyPixBoundsBuilder(self.all_dataset_skypix.values()) 

467 

468 def make_timespan_builder(self, quantum_data_id: DataCoordinate) -> TimespanBuilder: 

469 """Return an object that accumulates the appropriate timespan for 

470 a quantum. 

471 

472 Parameters 

473 ---------- 

474 quantum_data_id : `lsst.daf.butler.DataCoordinate` 

475 Data ID for this quantum. 

476 

477 Returns 

478 ------- 

479 builder : `TimespanBuilder` 

480 Object that accumulates the appropriate timespan bounds for a 

481 quantum. If a timespan is not needed, this object will do nothing. 

482 """ 

483 if not self.any_dataset_has_timespan: 

484 return _TrivialTimespanBuilder() 

485 if self.temporal_connections: 

486 return _ConnectionTimespanBuilder(self.task_node, self.temporal_connections, quantum_data_id) 

487 if self.task_node.dimensions.temporal: 

488 return _QuantumOnlyTimespanBuilder(quantum_data_id) 

489 else: 

490 return _UnboundedTimespanBuilder() 

491 

492 

493class SkyPixBoundsBuilder(ABC): 

494 """A base class for objects that accumulate the appropriate spatial bounds 

495 for a quantum. 

496 """ 

497 

498 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None: 

499 """Handle the skeleton graph node for a regular input/output connection 

500 for this quantum, including its data ID in the bounds if appropriate. 

501 

502 Parameters 

503 ---------- 

504 parent_dataset_type_name : `str` 

505 Name of the dataset type. Never a component dataset type name. 

506 data_id : `lsst.daf.butler.DataCoordinate` 

507 Data ID for the dataset. 

508 """ 

509 pass 

510 

511 @abstractmethod 

512 def finish(self) -> dict[str, RangeSet]: 

513 """Finish building the spatial bounds and return them. 

514 

515 Returns 

516 ------- 

517 bounds : `dict` [ `str`, `lsst.sphgeom.RangeSet` ] 

518 The spatial bounds of this quantum in various skypix dimensions. 

519 Keys are skypix dimension names and values are sets of integer 

520 pixel ID ranges. 

521 """ 

522 raise NotImplementedError() 

523 

524 

525class TimespanBuilder(ABC): 

526 """A base class for objects that accumulate the appropriate timespan 

527 for a quantum. 

528 """ 

529 

530 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None: 

531 """Handle the skeleton graph node for a regular input/output connection 

532 for this quantum, including its data ID in the bounds if appropriate. 

533 

534 Parameters 

535 ---------- 

536 parent_dataset_type_name : `str` 

537 Name of the dataset type. Never a component dataset type name. 

538 data_id : `lsst.daf.butler.DataCoordinate` 

539 Data ID for the dataset. 

540 """ 

541 pass 

542 

543 @abstractmethod 

544 def finish(self) -> Timespan | None: 

545 """Finish building the timespan and return it. 

546 

547 Returns 

548 ------- 

549 timespan : `lsst.daf.butler.Timespan` or `None` 

550 The timespan of this quantum, or `None` if it is known to not be 

551 needed. 

552 """ 

553 raise NotImplementedError() 

554 

555 

556class _TrivialSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

557 """Implementation of `SkyPixBoundsBuilder` for when no skypix bounds are 

558 needed. 

559 """ 

560 

561 def finish(self) -> dict[str, RangeSet]: 

562 return {} 

563 

564 

565class _TrivialTimespanBuilder(TimespanBuilder): 

566 """Implementation of `TimespanBuilder` for when no timespan is needed.""" 

567 

568 def finish(self) -> None: 

569 return None 

570 

571 

572class _QuantumOnlySkyPixBoundsBuilder(SkyPixBoundsBuilder): 

573 """Implementation of `SkyPixBoundsBuilder` for when the quantum data IDs 

574 provide the only relevant spatial regions. 

575 """ 

576 

577 def __init__(self, dimensions: Iterable[SkyPixDimension], quantum_data_id: DataCoordinate) -> None: 

578 self._region = quantum_data_id.region 

579 self._dimensions = dimensions 

580 

581 def finish(self) -> dict[str, RangeSet]: 

582 return { 

583 dimension.name: dimension.pixelization.envelope(self._region) for dimension in self._dimensions 

584 } 

585 

586 

587class _QuantumOnlyTimespanBuilder(TimespanBuilder): 

588 """Implementation of `TimespanBuilder` for when the quantum data IDs 

589 provide the only relevant timespans. 

590 """ 

591 

592 def __init__(self, quantum_data_id: DataCoordinate) -> None: 

593 self._timespan = cast(Timespan, quantum_data_id.timespan) 

594 

595 def finish(self) -> Timespan: 

596 return self._timespan 

597 

598 

599class _UnboundedSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

600 """Implementation of `SkyPixBoundsBuilder` for when the bounds cover the 

601 full sky. 

602 """ 

603 

604 def __init__(self, dimensions: Iterable[SkyPixDimension]): 

605 self._dimensions = dimensions 

606 

607 def finish(self) -> dict[str, RangeSet]: 

608 return {dimension.name: dimension.pixelization.universe() for dimension in self._dimensions} 

609 

610 

611class _UnboundedTimespanBuilder(TimespanBuilder): 

612 """Implementation of `TimespanBuilder` for when the timespan covers all 

613 time. 

614 """ 

615 

616 def finish(self) -> Timespan: 

617 return Timespan(None, None) 

618 

619 

620class _ConnectionSkyPixBoundsBuilder(SkyPixBoundsBuilder): 

621 """Implementation of `SkyPixBoundsBuilder` for when other input or output 

622 connections contribute to the spatial bounds. 

623 """ 

624 

625 def __init__( 

626 self, 

627 task_node: TaskNode, 

628 bounds_connections: frozenset[str], 

629 dimensions: Iterable[SkyPixDimension], 

630 quantum_data_id: DataCoordinate, 

631 ) -> None: 

632 self._dimensions = dimensions 

633 self._regions: list[Region] = [] 

634 if task_node.dimensions.spatial: 

635 self._regions.append(quantum_data_id.region) 

636 self._dataset_type_names: set[str] = set() 

637 for connection_name in bounds_connections: 

638 if edge := task_node.inputs.get(connection_name): 

639 self._dataset_type_names.add(edge.parent_dataset_type_name) 

640 else: 

641 self._dataset_type_names.add(task_node.outputs[connection_name].parent_dataset_type_name) 

642 # Note that we end up raising if the input is a prerequisite (and 

643 # hence not in task_node.inputs or task_node.outputs); this 

644 # justifies the cast in `handle_dataset`. 

645 

646 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None: 

647 if parent_dataset_type_name in self._dataset_type_names: 

648 self._regions.append(data_id.region) 

649 

650 def finish(self) -> dict[str, RangeSet]: 

651 result = {} 

652 for dimension in self._dimensions: 

653 bounds = RangeSet() 

654 for region in self._regions: 

655 bounds |= dimension.pixelization.envelope(region) 

656 result[dimension.name] = bounds 

657 return result 

658 

659 

660class _ConnectionTimespanBuilder(TimespanBuilder): 

661 """Implementation of `TimespanBuilder` for when other input or output 

662 connections contribute to the timespan. 

663 """ 

664 

665 def __init__( 

666 self, 

667 task_node: TaskNode, 

668 bounds_connections: frozenset[str], 

669 quantum_data_id: DataCoordinate, 

670 ) -> None: 

671 timespan = ( 

672 cast(Timespan, quantum_data_id.timespan) 

673 if task_node.dimensions.temporal 

674 else Timespan.makeEmpty() 

675 ) 

676 self._begin_nsec = timespan.nsec[0] 

677 self._end_nsec = timespan.nsec[1] 

678 self._dataset_type_names = set() 

679 for connection_name in bounds_connections: 

680 if edge := task_node.inputs.get(connection_name): 

681 self._dataset_type_names.add(edge.parent_dataset_type_name) 

682 else: 

683 self._dataset_type_names.add(task_node.outputs[connection_name].parent_dataset_type_name) 

684 # Note that we end up raising if the input is a prerequisite (and 

685 # hence not in task_node.inputs or task_node.outputs); this 

686 # justifies the cast in `handle_dataset`. 

687 

688 def handle_dataset(self, parent_dataset_type_name: str, data_id: DataCoordinate) -> None: 

689 if parent_dataset_type_name in self._dataset_type_names: 

690 nsec = cast(Timespan, data_id.timespan).nsec 

691 self._begin_nsec = min(self._begin_nsec, nsec[0]) 

692 self._end_nsec = max(self._end_nsec, nsec[1]) 

693 

694 def finish(self) -> Timespan: 

695 return Timespan(None, None, _nsec=(self._begin_nsec, self._end_nsec))