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

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") 

27 

28from collections.abc import Mapping 

29import os 

30 

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

32import felis.datamodel 

33import numpy as np 

34import pandas as pd 

35from astropy.table import Table 

36import yaml 

37 

38from lsst.daf.butler import Timespan 

39import lsst.dax.apdb as daxApdb 

40 

41 

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

58} 

59 

60 

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

62 nullable=False) -> str: 

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

64 

65 Parameters 

66 ---------- 

67 felis_type : `felis.datamodel.DataType` 

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

69 

70 Returns 

71 ------- 

72 column_dtype : `type` or `str` 

73 Type that can be used for columns in Pandas. 

74 

75 Raises 

76 ------ 

77 TypeError 

78 Raised if type is cannot be handled. 

79 """ 

80 try: 

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

82 except KeyError: 

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

84 

85 

86def readSdmSchemaFile(schemaFile: str, 

87 schemaName: str = "ApdbSchema", 

88 ): 

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

90 

91 Parameters 

92 ---------- 

93 schemaFile : `str` 

94 Fully specified path to the file to be read. 

95 schemaName : `str`, optional 

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

97 

98 Returns 

99 ------- 

100 schemaTable : dict[str, schema_model.Table] 

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

102 

103 Raises 

104 ------ 

105 ValueError 

106 If the schema file can't be parsed. 

107 """ 

108 schemaFile = os.path.expandvars(schemaFile) 

109 with open(schemaFile) as yaml_stream: 

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

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

112 if not schemas_list: 

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

114 elif len(schemas_list) > 1: 

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

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

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

118 schemaTable = {} 

119 

120 for singleTable in schema.tables: 

121 schemaTable[singleTable.name] = singleTable 

122 return schemaTable 

123 

124 

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

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

127 

128 Parameters 

129 ---------- 

130 apdb : `lsst.dax.apdb.Apdb` 

131 AP database connection object. 

132 

133 Returns 

134 ------- 

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

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

137 """ 

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

139 

140 

141def checkSdmSchemaColumns(apdbSchema, colNames, tableName): 

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

143 

144 Parameters 

145 ---------- 

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

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

148 colNames : `list` of ``str` 

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

150 tableName : `str` 

151 Name of the table in the schema to use. 

152 

153 Returns 

154 ------- 

155 missing : `list` of `str` 

156 All column names that are not in the schema 

157 """ 

158 table = apdbSchema[tableName] 

159 missing = [] 

160 

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

162 for col in colNames: 

163 if col not in names: 

164 missing.append(col) 

165 return missing 

166 

167 

168def convertDataFrameToSdmSchema(apdbSchema, sourceTable, tableName): 

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

170 

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

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

173 

174 Parameters 

175 ---------- 

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

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

178 sourceTable : `pandas.DataFrame` 

179 The input table to convert. 

180 tableName : `str` 

181 Name of the table in the schema to use. 

182 

183 Returns 

184 ------- 

185 `pandas.DataFrame` 

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

187 the input ``sourceTable``. 

188 """ 

189 if sourceTable.empty: 

190 make_empty_catalog(apdbSchema, tableName) 

191 table = apdbSchema[tableName] 

192 

193 data = {} 

194 nSrc = len(sourceTable) 

195 

196 for columnDef in table.columns: 

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

198 if columnDef.name in sourceTable.columns: 

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

200 index=sourceTable.index) 

201 else: 

202 if columnDef.nullable: 

203 try: 

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

205 except TypeError: 

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

207 else: 

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

209 return pd.DataFrame(data) 

210 

211 

212def convertTableToSdmSchema(apdbSchema, sourceTable, tableName): 

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

214 

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

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

217 

218 Parameters 

219 ---------- 

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

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

222 sourceTable : `astropy.table.Table` 

223 The input table to convert. 

224 tableName : `str` 

225 Name of the table in the schema to use. 

226 

227 Returns 

228 ------- 

229 `astropy.table.Table` 

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

231 the input ``sourceTable``. 

232 """ 

233 table = apdbSchema[tableName] 

234 

235 data = {} 

236 nSrc = len(sourceTable) 

237 

238 for columnDef in table.columns: 

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

240 if columnDef.name in sourceTable.columns: 

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

242 else: 

243 if columnDef.nullable: 

244 try: 

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

246 except TypeError: 

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

248 else: 

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

250 return Table(data) 

251 

252 

253def dropEmptyColumns(apdbSchema, sourceTable, tableName): 

254 """Drop empty columns that are nullable. 

255 

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

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

258 

259 Parameters 

260 ---------- 

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

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

263 sourceTable : `pandas.DataFrame` 

264 The input table to remove missing data columns from. 

265 tableName : `str` 

266 Name of the table in the schema to use. 

267 """ 

268 table = apdbSchema[tableName] 

269 

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

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

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

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

274 return sourceTable.drop(columns=dropColumns) 

275 

276 

277def make_empty_catalog(apdbSchema, tableName): 

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

279 

280 Parameters 

281 ---------- 

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

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

284 tableName : `str` 

285 Name of the table in the schema to use. 

286 

287 Returns 

288 ------- 

289 catalog : `pandas.DataFrame` 

290 An empty catalog. 

291 """ 

292 table = apdbSchema[tableName] 

293 

294 data = { 

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

296 for columnDef in table.columns 

297 } 

298 return pd.DataFrame(data) 

299 

300 

301def getMidpointFromTimespan(timespan, allowUnbounded=True): 

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

303 

304 Parameters 

305 ---------- 

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

307 A Timespan centered on the midpoint of a visit. 

308 allowUnbounded : `bool`, optional 

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

310 

311 Returns 

312 ------- 

313 midpoint : `astropy.time.Time` 

314 The midpoint of the timespan. 

315 

316 Raises 

317 ------ 

318 ValueError 

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

320 ``allowUnbounded`` is not set. 

321 ValueError 

322 Raised if the timespan is empty. 

323 """ 

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

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

326 

327 try: 

328 interval = timespan.end - timespan.begin 

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

330 except TypeError as e: 

331 if allowUnbounded: 

332 if timespan.end is not None: 

333 return timespan.end.tai 

334 elif timespan.begin is not None: 

335 return timespan.begin.tai 

336 else: 

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

338 else: 

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

340 

341 

342def makeEmptyForcedSourceTable(schema): 

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

344 

345 Returns 

346 ------- 

347 diaForcedSources : `pandas.DataFrame` 

348 Empty dataframe. 

349 """ 

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

351 return diaForcedSources