Coverage for python/lsst/analysis/ap/ppdb.py: 95%

1# This file is part of analysis_ap. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (https://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 program is free software: you can redistribute it and/or modify 

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

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

12# (at your option) any later version. 

13# 

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

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

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

17# GNU General Public License for more details. 

18# 

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

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

21 

22"""Load PPDB catalogs (DiaObjects, DiaSources, DiaForcedSources) through the 

23Rubin Science Platform TAP service. 

24 

25These tools deliberately use the scientist-facing TAP interface (ADQL over 

26``pyvo``) rather than a direct database connection, so that exercising them 

27also exercises the interface scientists will use once the PPDB goes live. 

28 

29The prototype PPDB currently lives at the ``data-int`` RSP environment 

30(``https://data-int.lsst.cloud/api/ppdbtap``). Tooling running elsewhere 

31(e.g. the USDF RSP) reaches it cross-environment with a ``data-int`` RSP 

32token carrying the ``read:tap`` scope, supplied via the ``RSP_TOKEN`` 

33environment variable. 

34 

35The ``DiaObject`` table is versioned: a single object accumulates multiple 

36rows over time and only the row with ``validityEndMjdTai IS NULL`` is the 

37current version. Every DiaObject query here applies that filter by default 

38so callers never accidentally retrieve (or double-count) stale versions. 

39``DiaSource`` and ``DiaForcedSource`` are append-only and are not versioned. 

40 

41The production PPDB will allow cone searches only on ``DiaObject``. Sources 

42and forced sources for a region must therefore be loaded object-first: 

43cone-search DiaObjects, then load the sources of those objects. The 

44``load_sources_for_region`` / ``load_forced_sources_for_region`` methods do 

45this for you and are the encouraged entry points. ``load_sources`` / 

46``load_forced_sources`` require a ``diaObjectId`` (scalar or iterable) and 

47take no spatial argument. Direct, object-unrestricted access to the source 

48tables is available only through the explicit, prototype-only 

49``load_sources_by_cone`` / ``load_forced_sources_by_cone`` (spatial) and 

50``load_sources_by_time_window`` / ``load_forced_sources_by_time_window`` 

51(temporal) methods, which warn when used. 

52 

53All loaders return `astropy.table.Table` objects. Their columns are returned 

54in SDM-schema order, since the TAP service returns ``SELECT *`` columns 

55alphabetically; passing an explicit ``columns`` list overrides this with the 

56order you request. 

57""" 

58 

59__all__ = [ 

60 "PpdbTap", 

61 "DiaObjectLightCurve", 

62 "region_from_exposure", 

63 "DEFAULT_PPDB_TAP_URL", 

64 "DEFAULT_PADDING_ARCSEC", 

65 "DEFAULT_ROW_LIMIT", 

66] 

67 

68import dataclasses 

69import functools 

70import logging 

71import numbers 

72import os 

73import time 

74from importlib.resources import files 

75 

76import astropy.table 

77import pyvo 

78import requests 

79 

80import lsst.geom 

81from lsst.pipe.tasks.schemaUtils import readSdmSchemaFile 

82 

83_log = logging.getLogger(__name__) 

84 

85# TAP endpoint of the prototype PPDB (the ``data-int`` RSP environment). 

86DEFAULT_PPDB_TAP_URL = "https://data-int.lsst.cloud/api/ppdbtap" 

87 

88# Default padding (arcseconds) loaded around an exposure footprint. 

89DEFAULT_PADDING_ARCSEC = 5.0 

90 

91# Default cap on rows returned by a single query, to guard against 

92# accidentally pulling the whole (multi-million row) PPDB. Pass ``limit=None`` 

93# to disable. 

94DEFAULT_ROW_LIMIT = 100000 

95 

96# Valid LSST band names; used to validate band filters before they are 

97# interpolated into an ADQL string. 

98_VALID_BANDS = frozenset("ugrizy") 

99 

100# The PPDB tables share the APDB SDM schema (``sdm_schemas/apdb.yaml``); its 

101# column order is the canonical, scientifically-grouped order. Loaders use it 

102# to reorder results, because the TAP service returns ``SELECT *`` columns 

103# alphabetically (its ``tap_schema.columns`` has no ``column_index``). Both 

104# helpers are memoized, so the schema file is parsed at most once. 

105 

106 

107@functools.cache 

108def _sdm_schema(): 

109 """Parse and return the APDB SDM schema (shared by the PPDB tables).""" 

110 path = os.fspath(files("lsst.sdm.schemas").joinpath("apdb.yaml")) 

111 return readSdmSchemaFile(path) 

112 

113 

114@functools.cache 

115def _sdm_column_order(table_name): 

116 """Return the SDM-schema column order for a PPDB table, or ``()``. 

117 

118 Parameters 

119 ---------- 

120 table_name : `str` 

121 Unqualified table name, e.g. ``"DiaSource"``. 

122 

123 Returns 

124 ------- 

125 columns : `tuple` [`str`] 

126 Column names in schema order, or an empty tuple if the table is not 

127 in the SDM schema. 

128 """ 

129 table_def = _sdm_schema().get(table_name) 

130 if table_def is None: 130 ↛ 131line 130 didn't jump to line 131 because the condition on line 130 was never true

131 return () 

132 return tuple(column.name for column in table_def.columns) 

133 

134 

135def region_from_exposure(exposure, padding=DEFAULT_PADDING_ARCSEC): 

136 """Compute a cone (ra, dec, radius) covering an exposure's sky footprint. 

137 

138 The cone is the smallest circle centered on the exposure's bounding-box 

139 center that contains all four corners, expanded by ``padding`` 

140 arcseconds. A circumscribing circle (rather than the exact polygon) keeps 

141 the resulting ADQL robust: ``CIRCLE`` is universally supported by TAP 

142 services, whereas ``POLYGON`` support and winding conventions vary. 

143 

144 Parameters 

145 ---------- 

146 exposure : `lsst.afw.image.Exposure` 

147 Exposure whose WCS and bounding box define the region. 

148 padding : `float`, optional 

149 Margin to add around the footprint, in arcseconds. 

150 

151 Returns 

152 ------- 

153 ra : `float` 

154 Cone center right ascension, in degrees (ICRS). 

155 dec : `float` 

156 Cone center declination, in degrees (ICRS). 

157 radius : `float` 

158 Cone radius, in degrees. 

159 

160 Raises 

161 ------ 

162 ValueError 

163 If the exposure has no WCS, an empty bounding box, or negative 

164 padding. 

165 """ 

166 if padding < 0: 

167 raise ValueError(f"padding must be non-negative, got {padding}") 

168 wcs = exposure.getWcs() 

169 if wcs is None: 

170 raise ValueError("Exposure has no WCS; cannot compute a sky region.") 

171 bbox = exposure.getBBox() 

172 if bbox.isEmpty(): 172 ↛ 173line 172 didn't jump to line 173 because the condition on line 172 was never true

173 raise ValueError("Exposure has an empty bounding box.") 

174 

175 center = wcs.pixelToSky(bbox.getCenter()) 

176 corners = [wcs.pixelToSky(lsst.geom.Point2D(corner)) 

177 for corner in bbox.getCorners()] 

178 radius = max(center.separation(corner) for corner in corners) 

179 radius = radius + padding * lsst.geom.arcseconds 

180 return (center.getRa().asDegrees(), 

181 center.getDec().asDegrees(), 

182 radius.asDegrees()) 

183 

184 

185@dataclasses.dataclass 

186class DiaObjectLightCurve: 

187 """The full PPDB record for a single diaObject. 

188 

189 Attributes 

190 ---------- 

191 diaObjectId : `int` 

192 Identifier of the diaObject. 

193 dia_object : `astropy.table.Row` 

194 The current (latest-version) DiaObject row. 

195 dia_sources : `astropy.table.Table` 

196 All DiaSources associated with the object, time-ordered. 

197 dia_forced_sources : `astropy.table.Table` 

198 All DiaForcedSources for the object, time-ordered. 

199 """ 

200 

201 diaObjectId: int 

202 dia_object: "astropy.table.Row" 

203 dia_sources: astropy.table.Table 

204 dia_forced_sources: astropy.table.Table 

205 

206 @property 

207 def n_sources(self): 

208 """Number of associated DiaSources (`int`).""" 

209 return len(self.dia_sources) 

210 

211 @property 

212 def n_forced_sources(self): 

213 """Number of associated DiaForcedSources (`int`).""" 

214 return len(self.dia_forced_sources) 

215 

216 def __repr__(self): 

217 return (f"DiaObjectLightCurve(diaObjectId={self.diaObjectId}, " 

218 f"n_sources={self.n_sources}, " 

219 f"n_forced_sources={self.n_forced_sources})") 

220 

221 

222class PpdbTap: 

223 """Load PPDB catalogs through the RSP TAP service as astropy Tables. 

224 

225 Parameters 

226 ---------- 

227 url : `str`, optional 

228 TAP endpoint to query. Defaults to the prototype PPDB at ``data-int``. 

229 token : `str`, optional 

230 RSP token with the ``read:tap`` scope. If not given (and ``service`` 

231 is also not given), the ``RSP_TOKEN`` environment variable is used. 

232 service : `pyvo.dal.TAPService`, optional 

233 Pre-built TAP service. When supplied, ``url`` and ``token`` are 

234 ignored. Intended for use with unit tests. 

235 log : `logging.Logger`, optional 

236 Logger to use; defaults to the module logger. 

237 

238 Notes 

239 ----- 

240 The token is read only from the environment or the constructor argument 

241 and is never logged. 

242 """ 

243 

244 def __init__(self, url=DEFAULT_PPDB_TAP_URL, *, token=None, 

245 service=None, log=None): 

246 self.url = url 

247 self.log = log if log is not None else _log 

248 if service is not None: 

249 self._service = service 

250 return 

251 if token is None: 

252 token = os.environ.get("RSP_TOKEN") 

253 if not token: 

254 raise RuntimeError( 

255 "No RSP token available for the PPDB TAP service. Generate " 

256 "a token with the 'read:tap' scope on data-int " 

257 "(https://data-int.lsst.cloud) and set it in the RSP_TOKEN " 

258 "environment variable, or pass token=...") 

259 session = requests.Session() 

260 session.headers["Authorization"] = f"Bearer {token}" 

261 self._service = pyvo.dal.TAPService(url, session=session) 

262 

263 # ------------------------------------------------------------------ 

264 # Query execution 

265 # ------------------------------------------------------------------ 

266 def run_query(self, adql, *, maxrec=None, order_table=None): 

267 """Run an ADQL query and return the result as an astropy Table. 

268 

269 Parameters 

270 ---------- 

271 adql : `str` 

272 The ADQL query to execute. 

273 maxrec : `int`, optional 

274 Server-side cap on the number of records returned. 

275 order_table : `str`, optional 

276 If given (e.g. ``"DiaSource"``), reorder the result columns into 

277 that table's SDM-schema order. The loaders pass this so that 

278 ``SELECT *`` results, which the TAP service returns alphabetically, 

279 come back in the schema's logical column order. Columns not in the 

280 schema are appended in their original order. 

281 

282 Returns 

283 ------- 

284 table : `astropy.table.Table` 

285 The query result. 

286 """ 

287 self.log.debug("Executing ADQL query: %s", adql) 

288 start = time.perf_counter() 

289 results = self._service.run_async(adql, maxrec=maxrec) 

290 table = results.to_table() 

291 if order_table is not None: 

292 table = self._order_columns(table, order_table) 

293 elapsed = time.perf_counter() - start 

294 self.log.info("Query returned %d rows in %.3f s", len(table), elapsed) 

295 return table 

296 

297 @staticmethod 

298 def _order_columns(table, table_name): 

299 """Reorder a result table's columns into SDM-schema order. 

300 

301 Every result column is expected to be defined in the SDM schema for 

302 ``table_name`` (the PPDB does not add columns outside it), so columns 

303 are returned in schema order. Returns the table unchanged if the 

304 schema is unknown or the order already matches. 

305 

306 Raises 

307 ------ 

308 RuntimeError 

309 If the result contains a column not present in the SDM schema, 

310 which signals an unexpected schema mismatch. 

311 """ 

312 schema_order = _sdm_column_order(table_name) 

313 if not schema_order: 313 ↛ 314line 313 didn't jump to line 314 because the condition on line 313 was never true

314 return table 

315 schema_set = set(schema_order) 

316 unknown = [c for c in table.colnames if c not in schema_set] 

317 if unknown: 

318 raise RuntimeError( 

319 f"ppdb.{table_name} returned columns not in the SDM schema: " 

320 f"{unknown}") 

321 new_order = [c for c in schema_order if c in table.colnames] 

322 if new_order == table.colnames: 

323 return table 

324 return table[new_order] 

325 

326 # ------------------------------------------------------------------ 

327 # ADQL construction helpers 

328 # ------------------------------------------------------------------ 

329 @staticmethod 

330 def _select_clause(columns, limit): 

331 """Build the ``SELECT [TOP n] cols`` prefix of a query.""" 

332 cols = "*" if not columns else ", ".join(columns) 

333 top = f"TOP {int(limit)} " if limit is not None else "" 

334 return f"SELECT {top}{cols}" 

335 

336 @staticmethod 

337 def _cone_clause(ra, dec, radius): 

338 """Build an ADQL cone-search predicate (all angles in degrees).""" 

339 return (f"CONTAINS(POINT('ICRS', ra, dec), " 

340 f"CIRCLE('ICRS', {float(ra)}, {float(dec)}, " 

341 f"{float(radius)})) = 1") 

342 

343 def _resolve_cone(self, ra, dec, radius, exposure, padding): 

344 """Resolve spatial arguments into a (ra, dec, radius) cone or None. 

345 

346 Either an ``exposure`` or an explicit ``ra``/``dec``/``radius`` triple 

347 may be given, but not both. Returns None when no spatial constraint 

348 was requested. Explicit ``radius`` is in degrees; ``padding`` (used 

349 only with ``exposure``) is in arcseconds. 

350 """ 

351 if exposure is not None: 

352 if any(v is not None for v in (ra, dec, radius)): 

353 raise ValueError("Specify either exposure= or ra/dec/radius, " 

354 "not both.") 

355 return region_from_exposure(exposure, padding=padding) 

356 if ra is None and dec is None and radius is None: 

357 return None 

358 if None in (ra, dec, radius): 

359 raise ValueError("ra, dec, and radius must all be provided for a " 

360 "cone search.") 

361 return float(ra), float(dec), float(radius) 

362 

363 @staticmethod 

364 def _band_clause(bands): 

365 """Build an ADQL ``band IN (...)`` predicate, validating each band.""" 

366 quoted = [] 

367 for band in bands: 

368 if band not in _VALID_BANDS: 

369 raise ValueError(f"Unknown band {band!r}; expected one of " 

370 f"{sorted(_VALID_BANDS)}.") 

371 quoted.append(f"'{band}'") 

372 return f"band IN ({', '.join(quoted)})" 

373 

374 def _finish(self, table, limit, table_name): 

375 """Warn on probable truncation and return the table unchanged.""" 

376 if limit is not None and len(table) == limit: 

377 self.log.warning( 

378 "%s query returned exactly the row limit (%d); results are " 

379 "likely truncated. Narrow the query or raise limit.", 

380 table_name, limit) 

381 return table 

382 

383 # ------------------------------------------------------------------ 

384 # DiaObject loaders 

385 # ------------------------------------------------------------------ 

386 def load_objects(self, *, ra=None, dec=None, radius=None, exposure=None, 

387 padding=DEFAULT_PADDING_ARCSEC, columns=None, 

388 limit=DEFAULT_ROW_LIMIT, latest=True): 

389 """Load DiaObjects, optionally within a spatial region. 

390 

391 Parameters 

392 ---------- 

393 ra, dec, radius : `float`, optional 

394 Cone-search center and radius, in degrees. All three must be 

395 given together; mutually exclusive with ``exposure``. 

396 exposure : `lsst.afw.image.Exposure`, optional 

397 Load objects covering this exposure's footprint (see 

398 `region_from_exposure`). Mutually exclusive with 

399 ``ra``/``dec``/``radius``. 

400 padding : `float`, optional 

401 Margin around the exposure footprint, in arcseconds (only used 

402 with ``exposure``). 

403 columns : `list` [`str`], optional 

404 Columns to select; defaults to all columns. 

405 limit : `int`, optional 

406 Maximum number of rows to return; None means no limit. 

407 latest : `bool`, optional 

408 If True (default), return only the current version of each object 

409 (``validityEndMjdTai IS NULL``). Setting this False returns every 

410 historical version and is rarely what you want. 

411 

412 Returns 

413 ------- 

414 objects : `astropy.table.Table` 

415 The matching DiaObjects. 

416 """ 

417 cone = self._resolve_cone(ra, dec, radius, exposure, padding) 

418 clauses = [] 

419 if latest: 

420 clauses.append("validityEndMjdTai IS NULL") 

421 if cone is not None: 

422 clauses.append(self._cone_clause(*cone)) 

423 where = (" WHERE " + " AND ".join(clauses)) if clauses else "" 

424 adql = (f"{self._select_clause(columns, limit)} FROM ppdb.DiaObject" 

425 f"{where} ORDER BY diaObjectId") 

426 order_table = "DiaObject" if columns is None else None 

427 table = self._finish(self.run_query(adql, order_table=order_table), 

428 limit, "DiaObject") 

429 if latest and len(table) and "diaObjectId" in table.colnames: 

430 n_unique = len(set(table["diaObjectId"].tolist())) 

431 if n_unique != len(table): 

432 self.log.warning( 

433 "DiaObject query returned %d rows but only %d unique " 

434 "diaObjectIds despite validityEndMjdTai IS NULL; the PPDB " 

435 "may contain duplicate current versions.", 

436 len(table), n_unique) 

437 return table 

438 

439 def load_object(self, diaObjectId, *, columns=None): 

440 """Load the current version of a single DiaObject. 

441 

442 Parameters 

443 ---------- 

444 diaObjectId : `int` 

445 Identifier of the object to load. 

446 columns : `list` [`str`], optional 

447 Columns to select; defaults to all columns. 

448 

449 Returns 

450 ------- 

451 object : `astropy.table.Row` 

452 The current DiaObject row. 

453 

454 Raises 

455 ------ 

456 RuntimeError 

457 If no current version of the object exists. 

458 """ 

459 adql = (f"{self._select_clause(columns, None)} FROM ppdb.DiaObject " 

460 f"WHERE validityEndMjdTai IS NULL AND " 

461 f"diaObjectId = {int(diaObjectId)}") 

462 order_table = "DiaObject" if columns is None else None 

463 table = self.run_query(adql, order_table=order_table) 

464 if len(table) == 0: 

465 raise RuntimeError( 

466 f"diaObjectId={diaObjectId} not found in ppdb.DiaObject") 

467 if len(table) > 1: 467 ↛ 468line 467 didn't jump to line 468 because the condition on line 467 was never true

468 self.log.warning( 

469 "diaObjectId=%s has %d current versions (validityEndMjdTai " 

470 "IS NULL); returning the first.", diaObjectId, len(table)) 

471 return table[0] 

472 

473 # ------------------------------------------------------------------ 

474 # DiaSource / DiaForcedSource loaders 

475 # 

476 # The production PPDB will permit cone searches only on DiaObject. The 

477 # encouraged way to retrieve sources for a region is therefore object- 

478 # first: cone-search DiaObjects (`load_*_for_region`), then load their 

479 # sources by diaObjectId. `load_sources`/`load_forced_sources` require a 

480 # diaObjectId and take no spatial argument. Direct, object-unrestricted 

481 # source-table access is confined to the explicit, prototype-only 

482 # `load_*_by_cone` and `load_*_by_time_window` methods. 

483 # ------------------------------------------------------------------ 

484 def _source_filter_clauses(self, bands, mjd_begin, mjd_end): 

485 """Build the time/band filter clauses shared by the source loaders.""" 

486 clauses = [] 

487 if mjd_begin is not None: 

488 clauses.append(f"midpointMjdTai >= {float(mjd_begin)}") 

489 if mjd_end is not None: 

490 clauses.append(f"midpointMjdTai < {float(mjd_end)}") 

491 if bands: 

492 clauses.append(self._band_clause(bands)) 

493 return clauses 

494 

495 def _run_source_query(self, table_name, clauses, columns, limit): 

496 """Run a single time-ordered SELECT against a source table.""" 

497 where = (" WHERE " + " AND ".join(clauses)) if clauses else "" 

498 adql = (f"{self._select_clause(columns, limit)} FROM ppdb.{table_name}" 

499 f"{where} ORDER BY midpointMjdTai") 

500 order_table = table_name if columns is None else None 

501 return self._finish(self.run_query(adql, order_table=order_table), 

502 limit, table_name) 

503 

504 def _load_source_like(self, table_name, diaObjectId, extra_clauses, 

505 columns, limit, id_chunk_size): 

506 """Dispatch a source load by one id or many ids. 

507 

508 ``diaObjectId`` is required: source tables are loaded by object. An 

509 iterable (including the empty list a region search may produce) takes 

510 the chunked ``IN`` path; a scalar takes a single ``=`` query. 

511 """ 

512 if diaObjectId is None: 

513 raise ValueError( 

514 "diaObjectId is required; DiaSources and DiaForcedSources are " 

515 "loaded by object. Use load_*_for_region for a spatial " 

516 "region, or the prototype-only load_*_by_cone / " 

517 "load_*_by_time_window methods for direct source-table access.") 

518 if isinstance(diaObjectId, numbers.Integral): 

519 clauses = [f"diaObjectId = {int(diaObjectId)}"] + list(extra_clauses) 

520 return self._run_source_query(table_name, clauses, columns, limit) 

521 return self._load_sources_for_ids(table_name, diaObjectId, 

522 extra_clauses, columns, limit, 

523 id_chunk_size) 

524 

525 def _load_sources_for_ids(self, table_name, ids, extra_clauses, columns, 

526 limit, id_chunk_size): 

527 """Load sources for many diaObjectIds via chunked ``IN`` queries. 

528 

529 A cone search can return more diaObjectIds than fit in a single 

530 ``diaObjectId IN (...)`` clause, so the ids are split into chunks 

531 of ``id_chunk_size`` and the per-chunk results are concatenated and 

532 re-sorted by time. 

533 ``limit`` caps the total rows returned across all chunks. 

534 """ 

535 if id_chunk_size < 1: 535 ↛ 536line 535 didn't jump to line 536 because the condition on line 535 was never true

536 raise ValueError("id_chunk_size must be >= 1") 

537 ids = list(dict.fromkeys(int(i) for i in ids)) 

538 if not ids or (limit is not None and limit <= 0): 

539 # Empty result, but issue a query so columns/dtypes are populated. 

540 return self._run_source_query(table_name, ["1 = 0"], columns, None) 

541 

542 order_table = table_name if columns is None else None 

543 pages = [] 

544 total = 0 

545 truncated = False 

546 for start in range(0, len(ids), id_chunk_size): 

547 if limit is not None and total >= limit: 547 ↛ 548line 547 didn't jump to line 548 because the condition on line 547 was never true

548 truncated = True 

549 break 

550 chunk = ids[start:start + id_chunk_size] 

551 remaining = None if limit is None else limit - total 

552 id_list = ", ".join(str(i) for i in chunk) 

553 clauses = [f"diaObjectId IN ({id_list})"] + list(extra_clauses) 

554 where = " WHERE " + " AND ".join(clauses) 

555 adql = (f"{self._select_clause(columns, remaining)} " 

556 f"FROM ppdb.{table_name}{where} ORDER BY midpointMjdTai") 

557 page = self.run_query(adql, order_table=order_table) 

558 total += len(page) 

559 pages.append(page) 

560 

561 table = (pages[0] if len(pages) == 1 

562 else astropy.table.vstack(pages, metadata_conflicts="silent")) 

563 if len(table) and "midpointMjdTai" in table.colnames: 563 ↛ 565line 563 didn't jump to line 565 because the condition on line 563 was always true

564 table.sort("midpointMjdTai") 

565 if limit is not None and (truncated or total == limit): 565 ↛ 566line 565 didn't jump to line 566 because the condition on line 565 was never true

566 self.log.warning( 

567 "%s query reached the row limit (%d) while loading sources for " 

568 "%d objects; results are likely truncated. Raise limit or " 

569 "narrow the query.", table_name, limit, len(ids)) 

570 return table 

571 

572 def _direct_source_query(self, table_name, region_method, kind, clauses, 

573 columns, limit): 

574 """Warn, then run a prototype-only direct source-table query. 

575 

576 ``load_sources``/``load_forced_sources`` require a diaObjectId; these 

577 direct (object-unrestricted) source-table searches will not exist in 

578 the production PPDB, so callers reach them only through the explicit 

579 ``load_*_by_cone`` / ``load_*_by_time_window`` methods, which route 

580 here and emit this warning. 

581 """ 

582 self.log.warning( 

583 "Direct %s %s search is a prototype-only capability and will not " 

584 "be available in the production PPDB; prefer %s().", 

585 table_name, kind, region_method) 

586 return self._run_source_query(table_name, clauses, columns, limit) 

587 

588 def _load_by_cone(self, table_name, region_method, *, ra, dec, radius, 

589 exposure, padding, bands, mjd_begin, mjd_end, columns, 

590 limit): 

591 """Direct cone search on a source table (prototype-only).""" 

592 cone = self._resolve_cone(ra, dec, radius, exposure, padding) 

593 if cone is None: 

594 raise ValueError("A cone search requires ra/dec/radius or an " 

595 "exposure.") 

596 clauses = [self._cone_clause(*cone)] 

597 clauses += self._source_filter_clauses(bands, mjd_begin, mjd_end) 

598 return self._direct_source_query(table_name, region_method, "cone", 

599 clauses, columns, limit) 

600 

601 def _load_by_time_window(self, table_name, region_method, *, mjd_begin, 

602 mjd_end, bands, columns, limit): 

603 """Direct time-window scan of a source table (prototype-only).""" 

604 if mjd_begin is None and mjd_end is None: 

605 raise ValueError("A time-window search requires mjd_begin and/or " 

606 "mjd_end.") 

607 clauses = self._source_filter_clauses(bands, mjd_begin, mjd_end) 

608 return self._direct_source_query(table_name, region_method, 

609 "time-window", clauses, columns, 

610 limit) 

611 

612 def load_sources(self, *, diaObjectId, mjd_begin=None, mjd_end=None, 

613 bands=None, columns=None, limit=DEFAULT_ROW_LIMIT, 

614 id_chunk_size=1000): 

615 """Load DiaSources for a given object (or objects), by time and band. 

616 

617 DiaSources are loaded by object: ``diaObjectId`` is required. This 

618 loader has no spatial argument, because the production PPDB does not 

619 permit searches on the source tables. To load sources for a region, 

620 use `load_sources_for_region` (object-first); for direct, object- 

621 unrestricted prototype-only access, use `load_sources_by_cone` or 

622 `load_sources_by_time_window`. 

623 

624 Parameters 

625 ---------- 

626 diaObjectId : `int` or iterable of `int` 

627 Restrict to sources of this object, or of any of these objects 

628 (issued as chunked ``diaObjectId IN (...)`` queries). Required. 

629 mjd_begin, mjd_end : `float`, optional 

630 Restrict to ``mjd_begin <= midpointMjdTai < mjd_end`` (TAI). 

631 bands : `list` [`str`], optional 

632 Restrict to these bands (subset of ``ugrizy``). 

633 columns : `list` [`str`], optional 

634 Columns to select; defaults to all columns. 

635 limit : `int`, optional 

636 Maximum number of rows to return; None means no limit. 

637 id_chunk_size : `int`, optional 

638 Number of diaObjectIds per ``IN`` query when an iterable is given. 

639 

640 Returns 

641 ------- 

642 sources : `astropy.table.Table` 

643 The matching DiaSources, ordered by ``midpointMjdTai``. 

644 

645 Raises 

646 ------ 

647 ValueError 

648 If ``diaObjectId`` is None. 

649 """ 

650 extra = self._source_filter_clauses(bands, mjd_begin, mjd_end) 

651 return self._load_source_like("DiaSource", diaObjectId, extra, 

652 columns, limit, id_chunk_size) 

653 

654 def load_forced_sources(self, *, diaObjectId, mjd_begin=None, 

655 mjd_end=None, bands=None, columns=None, 

656 limit=DEFAULT_ROW_LIMIT, id_chunk_size=1000): 

657 """Load DiaForcedSources, with the same filters as `load_sources`. 

658 

659 Like `load_sources`, ``diaObjectId`` is required and there is no 

660 spatial argument; use `load_forced_sources_for_region`, or the 

661 prototype-only `load_forced_sources_by_cone` / 

662 `load_forced_sources_by_time_window`. 

663 

664 Returns 

665 ------- 

666 forced_sources : `astropy.table.Table` 

667 The matching DiaForcedSources, ordered by ``midpointMjdTai``. 

668 

669 Raises 

670 ------ 

671 ValueError 

672 If ``diaObjectId`` is None. 

673 """ 

674 extra = self._source_filter_clauses(bands, mjd_begin, mjd_end) 

675 return self._load_source_like("DiaForcedSource", diaObjectId, extra, 

676 columns, limit, id_chunk_size) 

677 

678 def load_sources_for_region(self, *, ra=None, dec=None, radius=None, 

679 exposure=None, padding=DEFAULT_PADDING_ARCSEC, 

680 mjd_begin=None, mjd_end=None, bands=None, 

681 columns=None, limit=DEFAULT_ROW_LIMIT, 

682 object_limit=DEFAULT_ROW_LIMIT, 

683 id_chunk_size=1000): 

684 """Load DiaSources for a region, the production-compatible way. 

685 

686 Performs the two-step workflow the production PPDB requires: cone- 

687 search DiaObjects (the only table that allows spatial search), then 

688 load the DiaSources of those objects. 

689 

690 Parameters 

691 ---------- 

692 ra, dec, radius : `float`, optional 

693 Cone center and radius, in degrees; mutually exclusive with 

694 ``exposure``. The radius is applied to *object* positions. 

695 exposure : `lsst.afw.image.Exposure`, optional 

696 Load sources for objects covering this exposure's footprint. 

697 padding : `float`, optional 

698 Margin around the exposure footprint, in arcseconds. 

699 mjd_begin, mjd_end, bands, columns : optional 

700 Passed through to the source query (see `load_sources`). 

701 limit : `int`, optional 

702 Cap on returned sources; None means no limit. 

703 object_limit : `int`, optional 

704 Cap on the DiaObjects found by the cone search. 

705 id_chunk_size : `int`, optional 

706 diaObjectIds per ``IN`` query when fetching sources. 

707 

708 Returns 

709 ------- 

710 sources : `astropy.table.Table` 

711 DiaSources of the objects in the region, ordered by time. 

712 

713 Notes 

714 ----- 

715 Because this is object-first, DiaSources not associated with a 

716 returned DiaObject (e.g. unassociated or ssObject-only detections) 

717 are not included. To find those, use `load_sources_by_cone`. 

718 """ 

719 ids = self._region_object_ids( 

720 ra, dec, radius, exposure, padding, object_limit) 

721 return self.load_sources(diaObjectId=ids, mjd_begin=mjd_begin, 

722 mjd_end=mjd_end, bands=bands, columns=columns, 

723 limit=limit, id_chunk_size=id_chunk_size) 

724 

725 def load_forced_sources_for_region(self, *, ra=None, dec=None, 

726 radius=None, exposure=None, 

727 padding=DEFAULT_PADDING_ARCSEC, 

728 mjd_begin=None, mjd_end=None, bands=None, 

729 columns=None, limit=DEFAULT_ROW_LIMIT, 

730 object_limit=DEFAULT_ROW_LIMIT, 

731 id_chunk_size=1000): 

732 """Load DiaForcedSources for a region (object-first). 

733 

734 Same two-step workflow and caveats as `load_sources_for_region`. 

735 

736 Returns 

737 ------- 

738 forced_sources : `astropy.table.Table` 

739 DiaForcedSources of the objects in the region, ordered by time. 

740 """ 

741 ids = self._region_object_ids( 

742 ra, dec, radius, exposure, padding, object_limit) 

743 return self.load_forced_sources( 

744 diaObjectId=ids, mjd_begin=mjd_begin, mjd_end=mjd_end, bands=bands, 

745 columns=columns, limit=limit, id_chunk_size=id_chunk_size) 

746 

747 def _region_object_ids(self, ra, dec, radius, exposure, padding, 

748 object_limit): 

749 """Cone-search DiaObjects and return their diaObjectIds.""" 

750 objects = self.load_objects(ra=ra, dec=dec, radius=radius, 

751 exposure=exposure, padding=padding, 

752 columns=["diaObjectId"], 

753 limit=object_limit) 

754 if not len(objects): 

755 return [] 

756 return objects["diaObjectId"].tolist() 

757 

758 def load_sources_by_cone(self, *, ra=None, dec=None, radius=None, 

759 exposure=None, padding=DEFAULT_PADDING_ARCSEC, 

760 mjd_begin=None, mjd_end=None, bands=None, 

761 columns=None, limit=DEFAULT_ROW_LIMIT): 

762 """Directly cone-search the DiaSource table (prototype-only). 

763 

764 .. warning:: 

765 

766 The production PPDB will not allow spatial searches on the source 

767 tables. This works only against the prototype and is intended for 

768 debugging/validation (e.g. finding DiaSources with no associated 

769 DiaObject). For the production-compatible region workflow use 

770 `load_sources_for_region`. 

771 

772 Parameters mirror `load_sources_for_region` minus ``object_limit``. 

773 

774 Returns 

775 ------- 

776 sources : `astropy.table.Table` 

777 DiaSources within the cone, ordered by ``midpointMjdTai``. 

778 """ 

779 return self._load_by_cone( 

780 "DiaSource", "load_sources_for_region", ra=ra, dec=dec, 

781 radius=radius, exposure=exposure, padding=padding, bands=bands, 

782 mjd_begin=mjd_begin, mjd_end=mjd_end, columns=columns, limit=limit) 

783 

784 def load_forced_sources_by_cone(self, *, ra=None, dec=None, radius=None, 

785 exposure=None, 

786 padding=DEFAULT_PADDING_ARCSEC, 

787 mjd_begin=None, mjd_end=None, bands=None, 

788 columns=None, limit=DEFAULT_ROW_LIMIT): 

789 """Directly cone-search the DiaForcedSource table (prototype-only). 

790 

791 .. warning:: 

792 

793 Not supported by the production PPDB; prototype/debugging only. 

794 Prefer `load_forced_sources_for_region`. 

795 

796 Returns 

797 ------- 

798 forced_sources : `astropy.table.Table` 

799 DiaForcedSources within the cone, ordered by ``midpointMjdTai``. 

800 """ 

801 return self._load_by_cone( 

802 "DiaForcedSource", "load_forced_sources_for_region", ra=ra, 

803 dec=dec, radius=radius, exposure=exposure, padding=padding, 

804 bands=bands, mjd_begin=mjd_begin, mjd_end=mjd_end, columns=columns, 

805 limit=limit) 

806 

807 def load_sources_by_time_window(self, *, mjd_begin=None, mjd_end=None, 

808 bands=None, columns=None, 

809 limit=DEFAULT_ROW_LIMIT): 

810 """Load DiaSources in a time window, any object (prototype-only). 

811 

812 .. warning:: 

813 

814 The production PPDB will not allow object-unrestricted searches on 

815 the source tables. This works only against the prototype and is 

816 intended for debugging (e.g. inspecting everything processed in an 

817 MJD range). For the production-compatible workflow use 

818 `load_sources_for_region`. 

819 

820 Parameters 

821 ---------- 

822 mjd_begin, mjd_end : `float`, optional 

823 Time window on ``midpointMjdTai`` (TAI), as 

824 ``mjd_begin <= midpointMjdTai < mjd_end``. At least one bound is 

825 required (an unbounded source-table scan is not permitted). 

826 bands : `list` [`str`], optional 

827 Restrict to these bands (subset of ``ugrizy``). 

828 columns : `list` [`str`], optional 

829 Columns to select; defaults to all columns. 

830 limit : `int`, optional 

831 Maximum number of rows to return; None means no limit. 

832 

833 Returns 

834 ------- 

835 sources : `astropy.table.Table` 

836 DiaSources in the window, ordered by ``midpointMjdTai``. 

837 

838 Raises 

839 ------ 

840 ValueError 

841 If neither ``mjd_begin`` nor ``mjd_end`` is given. 

842 """ 

843 return self._load_by_time_window( 

844 "DiaSource", "load_sources_for_region", mjd_begin=mjd_begin, 

845 mjd_end=mjd_end, bands=bands, columns=columns, limit=limit) 

846 

847 def load_forced_sources_by_time_window(self, *, mjd_begin=None, 

848 mjd_end=None, bands=None, 

849 columns=None, 

850 limit=DEFAULT_ROW_LIMIT): 

851 """Load DiaForcedSources in a time window, any object (prototype-only). 

852 

853 .. warning:: 

854 

855 Not supported by the production PPDB; prototype/debugging only. 

856 Prefer `load_forced_sources_for_region`. At least one of 

857 ``mjd_begin`` / ``mjd_end`` is required. 

858 

859 Returns 

860 ------- 

861 forced_sources : `astropy.table.Table` 

862 DiaForcedSources in the window, ordered by ``midpointMjdTai``. 

863 

864 Raises 

865 ------ 

866 ValueError 

867 If neither ``mjd_begin`` nor ``mjd_end`` is given. 

868 """ 

869 return self._load_by_time_window( 

870 "DiaForcedSource", "load_forced_sources_for_region", 

871 mjd_begin=mjd_begin, mjd_end=mjd_end, bands=bands, columns=columns, 

872 limit=limit) 

873 

874 def load_sources_for_object(self, diaObjectId, *, columns=None, 

875 limit=DEFAULT_ROW_LIMIT): 

876 """Load all DiaSources associated with one diaObject. 

877 

878 Returns 

879 ------- 

880 sources : `astropy.table.Table` 

881 The object's DiaSources, ordered by ``midpointMjdTai``. 

882 """ 

883 return self.load_sources(diaObjectId=diaObjectId, columns=columns, 

884 limit=limit) 

885 

886 def load_forced_sources_for_object(self, diaObjectId, *, columns=None, 

887 limit=DEFAULT_ROW_LIMIT): 

888 """Load all DiaForcedSources associated with one diaObject. 

889 

890 Returns 

891 ------- 

892 forced_sources : `astropy.table.Table` 

893 The object's DiaForcedSources, ordered by ``midpointMjdTai``. 

894 """ 

895 return self.load_forced_sources(diaObjectId=diaObjectId, 

896 columns=columns, limit=limit) 

897 

898 def load_source(self, diaSourceId, *, columns=None): 

899 """Load a single DiaSource by id. 

900 

901 Raises 

902 ------ 

903 RuntimeError 

904 If the source does not exist. 

905 """ 

906 return self._load_one("DiaSource", "diaSourceId", diaSourceId, columns) 

907 

908 def load_forced_source(self, diaForcedSourceId, *, columns=None): 

909 """Load a single DiaForcedSource by id. 

910 

911 Raises 

912 ------ 

913 RuntimeError 

914 If the forced source does not exist. 

915 """ 

916 return self._load_one("DiaForcedSource", "diaForcedSourceId", 

917 diaForcedSourceId, columns) 

918 

919 def _load_one(self, table_name, id_column, id_value, columns):