Coverage for python/lsst/dax/apdb/tests/data

1# This file is part of dax_apdb.

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.

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 <http://www.gnu.org/licenses/>.

22from __future__ import annotations

24import datetime

25import random

26from collections.abc import Iterator

27from typing import Any

29import astropy.time

30import numpy

31import pandas

33from lsst.sphgeom import LonLat, Region, UnitVector3d

36def _genPointsInRegion(region: Region, count: int) -> Iterator[LonLat]:

37 """Generate bunch of SpherePoints inside given region.

39 Parameters

40 ----------

41 region : `lsst.sphgeom.Region`

42 Spherical region.

43 count : `int`

44 Number of points to generate.

46 Notes

47 -----

48 Returned points are random but not necessarily uniformly distributed.

49 """

50 bbox = region.getBoundingBox()

51 center = bbox.getCenter()

52 center_lon = center.getLon().asRadians()

53 center_lat = center.getLat().asRadians()

54 width = bbox.getWidth().asRadians()

55 height = bbox.getHeight().asRadians()

56 while count > 0:

57 lon = random.uniform(center_lon - width / 2, center_lon + width / 2)

58 lat = random.uniform(center_lat - height / 2, center_lat + height / 2)

59 lonlat = LonLat.fromRadians(lon, lat)

60 uv3d = UnitVector3d(lonlat)

61 if region.contains(uv3d):

62 yield lonlat

63 count -= 1

66def makeObjectCatalog(

67 region: Region | LonLat, count: int, visit_time: astropy.time.Time, *, start_id: int = 1, **kwargs: Any

68) -> pandas.DataFrame:

69 """Make a catalog containing a bunch of DiaObjects inside a region.

71 Parameters

72 ----------

73 region : `lsst.sphgeom.Region` or `lsst.sphgeom.LonLat`

74 Spherical region or spherical coordinate.

75 count : `int`

76 Number of records to generate.

77 visit_time : `astropy.time.Time`

78 Time of the visit.

79 start_id : `int`

80 Starting diaObjectId.

81 **kwargs : `Any`

82 Additional columns and their values to add to catalog.

84 Returns

85 -------

86 catalog : `pandas.DataFrame`

87 Catalog of DiaObjects records.

89 Notes

90 -----

91 Returned catalog only contains three columns - ``diaObjectId`, ``ra``, and

92 ``dec`` (in degrees).

93 """

94 if isinstance(region, Region):

95 points = list(_genPointsInRegion(region, count))

96 else:

97 points = [region] * count

98 # diaObjectId=0 may be used in some code for DiaSource foreign key to mean

99 # the same as ``None``.

100 ids = numpy.arange(start_id, len(points) + start_id, dtype=numpy.int64)

101 ras = numpy.array([lonlat.getLon().asDegrees() for lonlat in points], dtype=numpy.float64)

102 decs = numpy.array([lonlat.getLat().asDegrees() for lonlat in points], dtype=numpy.float64)

103 nDiaSources = numpy.ones(len(points), dtype=numpy.int32)

104 data = dict(

105 kwargs,

106 diaObjectId=ids,

107 ra=ras,

108 dec=decs,

109 nDiaSources=nDiaSources,

110 )

111 df = pandas.DataFrame(data)

112 return df

113

114

115def makeTimestampNow(use_mjd: bool, offset_ms: int = 0) -> float | datetime.datetime:

116 """Return current timestamp in yeither MJD TAI or datetime format.

117

118 Parameters

119 ----------

120 use_mjd : `bool`

121 If True return time as MJD TAI, otherwise as datetime.

122 offset_ms : `int`, optional

123 Additional offset in milliseconds to add to current timestamp.

124

125 Returns

126 -------

127 timestamp : `float` or `datetime.datetime`

128 Resulting timestamp.

129 """

130 if use_mjd:

131 ts = astropy.time.Time.now().tai.mjd

132 if offset_ms != 0:

133 ts += offset_ms / (24 * 3600 * 1_000)

134 return ts

135 else:

136 # TODO: Note that for now we use naive datetime for time_processed, to

137 # have it consistent with ap_association, this is being replaces with

138 # MJD TAI in the new APDB schema.

139 dt = datetime.datetime.now()

140 if offset_ms != 0:

141 dt += datetime.timedelta(milliseconds=offset_ms)

142 return dt

143

144

145def _makeTimestampColumn(column: str, use_mjd: bool = True) -> str:

146 """Return column name before/after schema migration to MJD TAI."""

147 if use_mjd:

148 if column == "time_processed":

149 return "timeProcessedMjdTai"

150 elif column == "time_withdrawn":

151 return "timeWithdrawnMjdTai"

152 else:

153 return f"{column}MjdTai"

154 else:

155 return column

156

157

158def makeSourceCatalog(

159 objects: pandas.DataFrame,

160 visit_time: astropy.time.Time,

161 start_id: int = 0,

162 visit: int = 1,

163 detector: int = 1,

164 use_mjd: bool = True,

165) -> pandas.DataFrame:

166 """Make a catalog containing a bunch of DiaSources associated with the

167 input DiaObjects.

168

169 Parameters

170 ----------

171 objects : `pandas.DataFrame`

172 Catalog of DiaObject records.

173 visit_time : `astropy.time.Time`

174 Time of the visit.

175 start_id : `int`

176 Starting value for ``diaObjectId``.

177 visit, detector : `int`

178 Value for ``visit`` and ``detector`` fields.

179 use_mjd : `bool`

180 If True use MJD TAI for timestamp columns.

181

182 Returns

183 -------

184 catalog : `pandas.DataFrame`

185 Catalog of DiaSource records.

186

187 Notes

188 -----

189 Returned catalog only contains small number of columns needed for tests.

190 """

191 nrows = len(objects)

192 midpointMjdTai = visit_time.mjd

193 df = pandas.DataFrame(

194 {

195 "diaSourceId": numpy.arange(start_id, start_id + nrows, dtype=numpy.int64),

196 "diaObjectId": objects["diaObjectId"],

197 "visit": numpy.full(nrows, visit, dtype=numpy.int64),

198 "detector": numpy.full(nrows, detector, dtype=numpy.int16),

199 "parentDiaSourceId": 0,

200 "ra": objects["ra"],

201 "dec": objects["dec"],

202 "midpointMjdTai": numpy.full(nrows, midpointMjdTai, dtype=numpy.float64),

203 "flags": numpy.full(nrows, 0, dtype=numpy.int64),

204 "ssObjectId": pandas.NA,

205 _makeTimestampColumn("time_processed", use_mjd): makeTimestampNow(use_mjd),

206 }

207 )

208 return df

209

210

211def makeForcedSourceCatalog(

212 objects: pandas.DataFrame,

213 visit_time: astropy.time.Time,

214 visit: int = 1,

215 detector: int = 1,

216 use_mjd: bool = True,

217) -> pandas.DataFrame:

218 """Make a catalog containing a bunch of DiaForcedSources associated with

219 the input DiaObjects.

220

221 Parameters

222 ----------

223 objects : `pandas.DataFrame`

224 Catalog of DiaObject records.

225 visit_time : `astropy.time.Time`

226 Time of the visit.

227 visit, detector : `int`

228 Value for ``visit`` and ``detector`` fields.

229 use_mjd : `bool`

230 If True use MJD TAI for timestamp columns.

231

232 Returns

233 -------

234 catalog : `pandas.DataFrame`

235 Catalog of DiaForcedSource records.

236

237 Notes

238 -----

239 Returned catalog only contains small number of columns needed for tests.

240 """

241 nrows = len(objects)

242 midpointMjdTai = visit_time.mjd

243 df = pandas.DataFrame(

244 {

245 "diaObjectId": objects["diaObjectId"],

246 "visit": numpy.full(nrows, visit, dtype=numpy.int64),

247 "detector": numpy.full(nrows, detector, dtype=numpy.int16),

248 "ra": objects["ra"],

249 "dec": objects["dec"],

250 "midpointMjdTai": numpy.full(nrows, midpointMjdTai, dtype=numpy.float64),

251 "flags": numpy.full(nrows, 0, dtype=numpy.int64),

252 _makeTimestampColumn("time_processed", use_mjd): makeTimestampNow(use_mjd),

253 }

254 )

255 return df

256

257

258def makeSSObjectCatalog(count: int, start_id: int = 1, flags: int = 0) -> pandas.DataFrame:

259 """Make a catalog containing a bunch of SSObjects.

260

261 Parameters

262 ----------

263 count : `int`

264 Number of records to generate.

265 startID : `int`

266 Initial SSObject ID.

267 flags : `int`

268 Value for ``flags`` column.

269

270 Returns

271 -------

272 catalog : `pandas.DataFrame`

273 Catalog of SSObjects records.

274

275 Notes

276 -----

277 Returned catalog only contains three columns - ``ssObjectId`, ``arc``,

278 and ``flags``.

279 """

280 ids = numpy.arange(start_id, count + start_id, dtype=numpy.int64)

281 arc = numpy.full(count, 0.001, dtype=numpy.float32)

282 flags_array = numpy.full(count, flags, dtype=numpy.int64)

283 df = pandas.DataFrame({"ssObjectId": ids, "arc": arc, "flags": flags_array})

284 return df

Coverage for python/lsst/dax/apdb/tests/data_factory.py: 18%

69 statements