Coverage for tests/utils_tests.py: 13%

80 statements  

« prev     ^ index     » next       coverage.py v7.5.1, created at 2024-05-09 11:35 +0000

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"""Helper functions for tests of DIA catalogs, including generating mock 

23catalogs for simulated APDB access. 

24""" 

25import datetime 

26import pandas as pd 

27import numpy as np 

28 

29from lsst.afw.cameraGeom.testUtils import DetectorWrapper 

30import lsst.afw.geom as afwGeom 

31import lsst.afw.image as afwImage 

32import lsst.daf.base as dafBase 

33import lsst.geom 

34 

35 

36def makeDiaObjects(nObjects, exposure, rng): 

37 """Make a test set of DiaObjects. 

38 

39 Parameters 

40 ---------- 

41 nObjects : `int` 

42 Number of objects to create. 

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

44 Exposure to create objects over. 

45 

46 Returns 

47 ------- 

48 diaObjects : `pandas.DataFrame` 

49 DiaObjects generated across the exposure. 

50 """ 

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

52 rand_x = rng.uniform(bbox.getMinX(), bbox.getMaxX(), size=nObjects) 

53 rand_y = rng.uniform(bbox.getMinY(), bbox.getMaxY(), size=nObjects) 

54 

55 midpointMjdTai = exposure.visitInfo.date.get(system=dafBase.DateTime.MJD) 

56 

57 data = [] 

58 for idx, (x, y) in enumerate(zip(rand_x, rand_y)): 

59 coord = exposure.wcs.pixelToSky(x, y) 

60 newObject = {"ra": coord.getRa().asDegrees(), 

61 "dec": coord.getDec().asDegrees(), 

62 "radecMjdTai": midpointMjdTai, 

63 "diaObjectId": idx + 1, 

64 "pmParallaxNdata": 0, 

65 "nearbyObj1": 0, 

66 "nearbyObj2": 0, 

67 "nearbyObj3": 0, 

68 "nDiaSources": 5} 

69 for f in ["u", "g", "r", "i", "z", "y"]: 

70 newObject["%s_psfFluxNdata" % f] = 0 

71 data.append(newObject) 

72 

73 return pd.DataFrame(data=data) 

74 

75 

76def makeDiaSources(nSources, diaObjectIds, exposure, rng, randomizeObjects=False): 

77 """Make a test set of DiaSources. 

78 

79 Parameters 

80 ---------- 

81 nSources : `int` 

82 Number of sources to create. 

83 diaObjectIds : `numpy.ndarray` 

84 Integer Ids of diaobjects to "associate" with the DiaSources. 

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

86 Exposure to create sources over. 

87 randomizeObjects : `bool`, optional 

88 If True, randomly draw from `diaObjectIds` to generate the ids in the 

89 output catalog, otherwise just iterate through them, repeating as 

90 necessary to get nSources objectIds. 

91 

92 Returns 

93 ------- 

94 diaSources : `pandas.DataFrame` 

95 DiaSources generated across the exposure. 

96 """ 

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

98 rand_x = rng.uniform(bbox.getMinX(), bbox.getMaxX(), size=nSources) 

99 rand_y = rng.uniform(bbox.getMinY(), bbox.getMaxY(), size=nSources) 

100 if randomizeObjects: 

101 objectIds = diaObjectIds[rng.randint(len(diaObjectIds), size=nSources)] 

102 else: 

103 objectIds = diaObjectIds[[i % len(diaObjectIds) for i in range(nSources)]] 

104 

105 midpointMjdTai = exposure.visitInfo.date.get(system=dafBase.DateTime.MJD) 

106 

107 data = [] 

108 for idx, (x, y, objId) in enumerate(zip(rand_x, rand_y, objectIds)): 

109 coord = exposure.wcs.pixelToSky(x, y) 

110 # Put together the minimum values for the alert. 

111 data.append({"ra": coord.getRa().asDegrees(), 

112 "dec": coord.getDec().asDegrees(), 

113 "x": x, 

114 "y": y, 

115 "visit": exposure.visitInfo.id, 

116 "detector": exposure.detector.getId(), 

117 "time_processed": datetime.datetime.now(), 

118 "diaObjectId": objId, 

119 "ssObjectId": 0, 

120 "parentDiaSourceId": 0, 

121 "diaSourceId": idx + 1, 

122 "midpointMjdTai": midpointMjdTai + 1.0 * idx, 

123 "band": exposure.getFilter().bandLabel, 

124 "psfNdata": 0, 

125 "trailNdata": 0, 

126 "dipoleNdata": 0}) 

127 

128 return pd.DataFrame(data=data) 

129 

130 

131def makeDiaForcedSources(nForcedSources, diaObjectIds, exposure, rng, randomizeObjects=False): 

132 """Make a test set of DiaSources. 

133 

134 Parameters 

135 ---------- 

136 nForcedSources : `int` 

137 Number of sources to create. 

138 diaObjectIds : `numpy.ndarray` 

139 Integer Ids of diaobjects to "associate" with the DiaSources. 

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

141 Exposure to create sources over. 

142 randomizeObjects : `bool`, optional 

143 If True, randomly draw from `diaObjectIds` to generate the ids in the 

144 output catalog, otherwise just iterate through them. 

145 

146 Returns 

147 ------- 

148 diaForcedSources : `pandas.DataFrame` 

149 DiaForcedSources generated across the exposure. 

150 """ 

151 midpointMjdTai = exposure.visitInfo.date.get(system=dafBase.DateTime.MJD) 

152 visit = exposure.visitInfo.id 

153 detector = exposure.detector.getId() 

154 if randomizeObjects: 

155 objectIds = diaObjectIds[rng.randint(len(diaObjectIds), size=nForcedSources)] 

156 else: 

157 objectIds = diaObjectIds[[i % len(diaObjectIds) for i in range(nForcedSources)]] 

158 

159 data = [] 

160 for i, objId in enumerate(objectIds): 

161 # Put together the minimum values for the alert. 

162 data.append({"diaForcedSourceId": i + 1, 

163 "visit": visit + i, 

164 "detector": detector, 

165 "diaObjectId": objId, 

166 "midpointMjdTai": midpointMjdTai + 1.0 * i, 

167 "time_processed": datetime.datetime.now(), 

168 "band": exposure.getFilter().bandLabel}) 

169 

170 return pd.DataFrame(data=data) 

171 

172 

173def makeExposure(flipX=False, flipY=False): 

174 """Create an exposure and flip the x or y (or both) coordinates. 

175 

176 Returns bounding boxes that are right or left handed around the bounding 

177 polygon. 

178 

179 Parameters 

180 ---------- 

181 flipX : `bool` 

182 Flip the x coordinate in the WCS. 

183 flipY : `bool` 

184 Flip the y coordinate in the WCS. 

185 

186 Returns 

187 ------- 

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

189 Exposure with a valid bounding box and wcs. 

190 """ 

191 metadata = dafBase.PropertySet() 

192 

193 metadata.set("SIMPLE", "T") 

194 metadata.set("BITPIX", -32) 

195 metadata.set("NAXIS", 2) 

196 metadata.set("NAXIS1", 1024) 

197 metadata.set("NAXIS2", 1153) 

198 metadata.set("RADECSYS", 'FK5') 

199 metadata.set("EQUINOX", 2000.) 

200 

201 metadata.setDouble("CRVAL1", 215.604025685476) 

202 metadata.setDouble("CRVAL2", 53.1595451514076) 

203 metadata.setDouble("CRPIX1", 1109.99981456774) 

204 metadata.setDouble("CRPIX2", 560.018167811613) 

205 metadata.set("CTYPE1", 'RA---SIN') 

206 metadata.set("CTYPE2", 'DEC--SIN') 

207 

208 xFlip = 1 

209 if flipX: 

210 xFlip = -1 

211 yFlip = 1 

212 if flipY: 

213 yFlip = -1 

214 metadata.setDouble("CD1_1", xFlip * 5.10808596133527E-05) 

215 metadata.setDouble("CD1_2", yFlip * 1.85579539217196E-07) 

216 metadata.setDouble("CD2_2", yFlip * -5.10281493481982E-05) 

217 metadata.setDouble("CD2_1", xFlip * -8.27440751733828E-07) 

218 

219 wcs = afwGeom.makeSkyWcs(metadata) 

220 exposure = afwImage.makeExposure( 

221 afwImage.makeMaskedImageFromArrays(np.ones((1024, 1153))), wcs) 

222 detector = DetectorWrapper(id=23, bbox=exposure.getBBox()).detector 

223 visit = afwImage.VisitInfo( 

224 exposureTime=200., 

225 date=dafBase.DateTime("2014-05-13T17:00:00.000000000", 

226 dafBase.DateTime.Timescale.TAI)) 

227 exposure.info.id = 1234 

228 exposure.setDetector(detector) 

229 exposure.info.setVisitInfo(visit) 

230 exposure.setFilter(afwImage.FilterLabel(band='g')) 

231 

232 return exposure