Coverage for python/lsst/ap/association/utils.py: 17%

1# This file is part of ap_association. 

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"""Utilities for working with the APDB. 

23""" 

24__all__ = ("convertDataFrameToSdmSchema", "readSdmSchemaFile", "readSchemaFromApdb", 

25 "dropEmptyColumns", "make_empty_catalog", "getMidpointFromTimespan", 

26 "makeEmptyForcedSourceTable", "checkSdmSchemaColumns", "getRegion", "paddedRegion") 

27 

28from collections.abc import Mapping 

29import os 

30 

31from lsst.dax.apdb import Apdb, ApdbTables, schema_model 

32import felis.datamodel 

33import math 

34import numpy as np 

35import pandas as pd 

36from astropy.table import Table 

37import yaml 

38 

39from lsst.daf.butler import Timespan 

40import lsst.dax.apdb as daxApdb 

41import lsst.geom 

42import lsst.sphgeom 

43 

44 

45# The first entry in the returned mapping is for nullable columns, 

46# the second entry is for non-nullable columns. 

47_dtype_map: Mapping[felis.datamodel.DataType, tuple[str, str]] = { 

48 felis.datamodel.DataType.double: ("float64", "float64"), # Cassandra utilities need np.nan not pd.NA 

49 felis.datamodel.DataType.float: ("float32", "float32"), # Cassandra utilities need np.nan not pd.NA 

50 felis.datamodel.DataType.timestamp: ("datetime64[ms]", "datetime64[ms]"), 

51 felis.datamodel.DataType.long: ("Int64", "int64"), 

52 felis.datamodel.DataType.int: ("Int32", "int32"), 

53 felis.datamodel.DataType.short: ("Int16", "int16"), 

54 felis.datamodel.DataType.byte: ("Int8", "int8"), 

55 felis.datamodel.DataType.binary: ("object", "object"), 

56 felis.datamodel.DataType.char: ("object", "object"), 

57 felis.datamodel.DataType.text: ("object", "object"), 

58 felis.datamodel.DataType.string: ("object", "object"), 

59 felis.datamodel.DataType.unicode: ("object", "object"), 

60 felis.datamodel.DataType.boolean: ("boolean", "bool"), 

61} 

62 

63 

64def column_dtype(felis_type: felis.datamodel.DataType | daxApdb.schema_model.ExtraDataTypes, 

65 nullable=False) -> str: 

66 """Return Pandas data type for a given Felis column type. 

67 

68 Parameters 

69 ---------- 

70 felis_type : `felis.datamodel.DataType` 

71 Felis type, on of the enums defined in `felis.datamodel` module. 

72 

73 Returns 

74 ------- 

75 column_dtype : `type` or `str` 

76 Type that can be used for columns in Pandas. 

77 

78 Raises 

79 ------ 

80 TypeError 

81 Raised if type is cannot be handled. 

82 """ 

83 try: 

84 return _dtype_map[felis_type][0] if nullable else _dtype_map[felis_type][1] 

85 except KeyError: 

86 raise TypeError(f"Unexpected Felis type: {felis_type}") 

87 

88 

89def readSdmSchemaFile(schemaFile: str, 

90 schemaName: str = "ApdbSchema", 

91 ): 

92 """Read a schema file in YAML format. 

93 

94 Parameters 

95 ---------- 

96 schemaFile : `str` 

97 Fully specified path to the file to be read. 

98 schemaName : `str`, optional 

99 Name of the table of schemas to read from the file. 

100 

101 Returns 

102 ------- 

103 schemaTable : dict[str, schema_model.Table] 

104 A dict of the schemas in the given table defined in the specified file. 

105 

106 Raises 

107 ------ 

108 ValueError 

109 If the schema file can't be parsed. 

110 """ 

111 schemaFile = os.path.expandvars(schemaFile) 

112 with open(schemaFile) as yaml_stream: 

113 schemas_list = list(yaml.load_all(yaml_stream, Loader=yaml.SafeLoader)) 

114 schemas_list = [schema for schema in schemas_list if schema.get("name") == schemaName] 

115 if not schemas_list: 

116 raise ValueError(f"Schema file {schemaFile!r} does not define schema {schemaName!r}") 

117 elif len(schemas_list) > 1: 

118 raise ValueError(f"Schema file {schemaFile!r} defines multiple schemas {schemaName!r}") 

119 felis_schema = felis.datamodel.Schema.model_validate(schemas_list[0]) 

120 schema = daxApdb.schema_model.Schema.from_felis(felis_schema) 

121 schemaTable = {} 

122 

123 for singleTable in schema.tables: 

124 schemaTable[singleTable.name] = singleTable 

125 return schemaTable 

126 

127 

128def readSchemaFromApdb(apdb: Apdb) -> dict[str, schema_model.Table | None]: 

129 """Extract the schema from an APDB instance. 

130 

131 Parameters 

132 ---------- 

133 apdb : `lsst.dax.apdb.Apdb` 

134 AP database connection object. 

135 

136 Returns 

137 ------- 

138 schemaTable : dict[str, schema_model.Table | None] 

139 A dict of the schemas in the given table defined in the specified file. 

140 """ 

141 return {table.table_name(): apdb.tableDef(table) for table in ApdbTables} 

142 

143 

144def checkSdmSchemaColumns(apdbSchema, colNames, tableName): 

145 """Check if supplied column names exists in the schema. 

146 

147 Parameters 

148 ---------- 

149 apdbSchema : `dict` [`str`, `lsst.dax.apdb.schema_model.Table`] 

150 Schema from ``sdm_schemas`` containing the table definition to use. 

151 colNames : `list` of ``str` 

152 Names of the columns to check for in the table. 

153 tableName : `str` 

154 Name of the table in the schema to use. 

155 

156 Returns 

157 ------- 

158 missing : `list` of `str` 

159 All column names that are not in the schema 

160 """ 

161 table = apdbSchema[tableName] 

162 missing = [] 

163 

164 names = [columnDef.name for columnDef in table.columns] 

165 for col in colNames: 

166 if col not in names: 

167 missing.append(col) 

168 return missing 

169 

170 

171def convertDataFrameToSdmSchema(apdbSchema, sourceTable, tableName): 

172 """Force a table to conform to the schema defined by the APDB. 

173 

174 This method uses the table definitions in ``sdm_schemas`` to 

175 load the schema of the APDB, and does not actually connect to the APDB. 

176 

177 Parameters 

178 ---------- 

179 apdbSchema : `dict` [`str`, `lsst.dax.apdb.schema_model.Table`] 

180 Schema from ``sdm_schemas`` containing the table definition to use. 

181 sourceTable : `pandas.DataFrame` 

182 The input table to convert. 

183 tableName : `str` 

184 Name of the table in the schema to use. 

185 

186 Returns 

187 ------- 

188 `pandas.DataFrame` 

189 A table with the correct schema for the APDB and data copied from 

190 the input ``sourceTable``. 

191 """ 

192 if sourceTable.empty: 

193 make_empty_catalog(apdbSchema, tableName) 

194 table = apdbSchema[tableName] 

195 

196 data = {} 

197 nSrc = len(sourceTable) 

198 

199 for columnDef in table.columns: 

200 dtype = column_dtype(columnDef.datatype, nullable=columnDef.nullable) 

201 if columnDef.name in sourceTable.columns: 

202 data[columnDef.name] = pd.Series(sourceTable[columnDef.name], dtype=dtype, 

203 index=sourceTable.index) 

204 else: 

205 if columnDef.nullable: 

206 try: 

207 data[columnDef.name] = pd.Series([pd.NA]*nSrc, dtype=dtype, index=sourceTable.index) 

208 except TypeError: 

209 data[columnDef.name] = pd.Series([np.nan]*nSrc, dtype=dtype, index=sourceTable.index) 

210 else: 

211 data[columnDef.name] = pd.Series([0]*nSrc, dtype=dtype, index=sourceTable.index) 

212 return pd.DataFrame(data) 

213 

214 

215def convertTableToSdmSchema(apdbSchema, sourceTable, tableName): 

216 """Force a table to conform to the schema defined by the APDB. 

217 

218 This method uses the table definitions in ``sdm_schemas`` to 

219 load the schema of the APDB, and does not actually connect to the APDB. 

220 

221 Parameters 

222 ---------- 

223 apdbSchema : `dict` [`str`, `lsst.dax.apdb.schema_model.Table`] 

224 Schema from ``sdm_schemas`` containing the table definition to use. 

225 sourceTable : `astropy.table.Table` 

226 The input table to convert. 

227 tableName : `str` 

228 Name of the table in the schema to use. 

229 

230 Returns 

231 ------- 

232 `astropy.table.Table` 

233 A table with the correct schema for the APDB and data copied from 

234 the input ``sourceTable``. 

235 """ 

236 table = apdbSchema[tableName] 

237 

238 data = {} 

239 nSrc = len(sourceTable) 

240 

241 for columnDef in table.columns: 

242 dtype = column_dtype(columnDef.datatype, nullable=columnDef.nullable) 

243 if columnDef.name in sourceTable.columns: 

244 data[columnDef.name] = Table.Column(sourceTable[columnDef.name], dtype=dtype.lower()) 

245 else: 

246 if columnDef.nullable: 

247 try: 

248 data[columnDef.name] = Table.Column([pd.NA]*nSrc, dtype=object) 

249 except TypeError: 

250 data[columnDef.name] = Table.Column([pd.nan]*nSrc, dtype=dtype) 

251 else: 

252 data[columnDef.name] = Table.Column([0]*nSrc, dtype=dtype) 

253 return Table(data) 

254 

255 

256def dropEmptyColumns(apdbSchema, sourceTable, tableName): 

257 """Drop empty columns that are nullable. 

258 

259 This method uses the table definitions in ``sdm_schemas`` to 

260 load the schema of the APDB, and does not actually connect to the APDB. 

261 

262 Parameters 

263 ---------- 

264 apdbSchema : `dict` [`str`, `lsst.dax.apdb.schema_model.Table`] 

265 Schema from ``sdm_schemas`` containing the table definition to use. 

266 sourceTable : `pandas.DataFrame` 

267 The input table to remove missing data columns from. 

268 tableName : `str` 

269 Name of the table in the schema to use. 

270 """ 

271 table = apdbSchema[tableName] 

272 

273 nullableList = [columnDef.name for columnDef in table.columns if columnDef.nullable] 

274 nullColumns = sourceTable.isnull().all() 

275 nullColNames = nullColumns[nullColumns].index.tolist() 

276 dropColumns = list(set(nullColNames) & set(nullableList)) 

277 return sourceTable.drop(columns=dropColumns) 

278 

279 

280def make_empty_catalog(apdbSchema, tableName): 

281 """Make an empty catalog for a table with a given name. 

282 

283 Parameters 

284 ---------- 

285 apdbSchema : `dict` [`str`, `lsst.dax.apdb.schema_model.Table`] 

286 Schema from ``sdm_schemas`` containing the table definition to use. 

287 tableName : `str` 

288 Name of the table in the schema to use. 

289 

290 Returns 

291 ------- 

292 catalog : `pandas.DataFrame` 

293 An empty catalog. 

294 """ 

295 table = apdbSchema[tableName] 

296 

297 data = { 

298 columnDef.name: pd.Series(dtype=column_dtype(columnDef.datatype, nullable=columnDef.nullable)) 

299 for columnDef in table.columns 

300 } 

301 return pd.DataFrame(data) 

302 

303 

304def getMidpointFromTimespan(timespan, allowUnbounded=True): 

305 """Safely retrieve the midpoint in TAI from a Timespan. 

306 

307 Parameters 

308 ---------- 

309 timespan : `lsst.daf.butler.Timespan` of `astropy.time.Time` 

310 A Timespan centered on the midpoint of a visit. 

311 allowUnbounded : `bool`, optional 

312 If set, return the start or end of an unbounded timespan. 

313 

314 Returns 

315 ------- 

316 midpoint : `astropy.time.Time` 

317 The midpoint of the timespan. 

318 

319 Raises 

320 ------ 

321 ValueError 

322 Raised if either the start or end of the timespan is None, and 

323 ``allowUnbounded`` is not set. 

324 ValueError 

325 Raised if the timespan is empty. 

326 """ 

327 if (timespan.begin == Timespan.EMPTY) or (timespan.begin == Timespan.EMPTY): 

328 raise ValueError("Cannot compute midpoint: EMPTY Timespan.") 

329 

330 try: 

331 interval = timespan.end - timespan.begin 

332 return (timespan.begin + interval/2).tai 

333 except TypeError as e: 

334 if allowUnbounded: 

335 if timespan.end is not None: 

336 return timespan.end.tai 

337 elif timespan.begin is not None: 

338 return timespan.begin.tai 

339 else: 

340 raise ValueError("Cannot compute midpoint: unbounded timespan.") from e 

341 else: 

342 raise ValueError("Cannot compute midpoint: unbounded timespan.") from e 

343 

344 

345def makeEmptyForcedSourceTable(schema): 

346 """Return a dataframe with the correct columns for diaForcedSources table. 

347 

348 Returns 

349 ------- 

350 diaForcedSources : `pandas.DataFrame` 

351 Empty dataframe. 

352 """ 

353 diaForcedSources = convertDataFrameToSdmSchema(schema, pd.DataFrame(), tableName="DiaForcedSource") 

354 return diaForcedSources 

355 

356 

357def getRegion(exposure): 

358 """Calculate an enveloping region for an exposure. 

359 

360 Parameters 

361 ---------- 

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

363 Exposure object with calibrated WCS. 

364 

365 Returns 

366 ------- 

367 region : `lsst.sphgeom.Region` 

368 Region enveloping an exposure. 

369 """ 

370 # Bounding box needs to be a `Box2D` not a `Box2I` for `wcs.pixelToSky()` 

371 bbox = lsst.geom.Box2D(exposure.getBBox()) 

372 wcs = exposure.getWcs() 

373 

374 region = lsst.sphgeom.ConvexPolygon([pp.getVector() for pp in wcs.pixelToSky(bbox.getCorners())]) 

375 

376 return region 

377 

378 

379def paddedRegion(region, margin): 

380 """Return a region that has been expanded by a buffer. 

381 

382 Parameters 

383 ---------- 

384 region : `lsst.sphgeom.Region` 

385 The region to pad. 

386 margin : `lsst.sphgeom.Angle` 

387 The amount by which to increase the region. 

388 

389 Returns 

390 ------- 

391 padded : `lsst.sphgeom.Region` 

392 An enlarged copy of ``region``. 

393 """ 

394 # region is almost certainly a (padded) detector bounding box. 

395 if isinstance(region, lsst.sphgeom.ConvexPolygon): 

396 # This is an ad-hoc, approximate implementation. It should be good 

397 # enough for catalog loading, but is not a general-purpose solution. 

398 center = lsst.geom.SpherePoint(region.getCentroid()) 

399 corners = [lsst.geom.SpherePoint(c) for c in region.getVertices()] 

400 # Approximate the region as a Euclidian square 

401 # geom.Angle(sphgeom.Angle) converter not pybind-wrapped??? 

402 diagonal_margin = lsst.geom.Angle(margin.asRadians() * math.sqrt(2.0)) 

403 padded = [c.offset(center.bearingTo(c), diagonal_margin) for c in corners] 

404 return lsst.sphgeom.ConvexPolygon.convexHull([c.getVector() for c in padded]) 

405 elif hasattr(region, "dilatedBy"): 

406 return region.dilatedBy(margin) 

407 else: 

408 return region.getBoundingCircle().dilatedBy(margin)