Coverage for tests/utils_tests.py: 13%

79 statements  

« prev     ^ index     » next       coverage.py v7.5.0, created at 2024-04-25 11:25 -0700

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

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 = np.random.uniform(bbox.getMinX(), bbox.getMaxX(), size=nObjects) 

53 rand_y = np.random.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 "flags": 1, 

69 "nDiaSources": 5} 

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

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

72 data.append(newObject) 

73 

74 return pd.DataFrame(data=data) 

75 

76 

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

78 """Make a test set of DiaSources. 

79 

80 Parameters 

81 ---------- 

82 nSources : `int` 

83 Number of sources to create. 

84 diaObjectIds : `numpy.ndarray` 

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

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

87 Exposure to create sources over. 

88 randomizeObjects : `bool`, optional 

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

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

91 necessary to get nSources objectIds. 

92 

93 Returns 

94 ------- 

95 diaSources : `pandas.DataFrame` 

96 DiaSources generated across the exposure. 

97 """ 

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

99 rand_x = np.random.uniform(bbox.getMinX(), bbox.getMaxX(), size=nSources) 

100 rand_y = np.random.uniform(bbox.getMinY(), bbox.getMaxY(), size=nSources) 

101 if randomizeObjects: 

102 objectIds = diaObjectIds[np.random.randint(len(diaObjectIds), size=nSources)] 

103 else: 

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

105 

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

107 

108 data = [] 

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

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

111 # Put together the minimum values for the alert. 

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

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

114 "x": x, 

115 "y": y, 

116 "ccdVisitId": exposure.info.id, 

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 "flags": 1}) 

128 

129 return pd.DataFrame(data=data) 

130 

131 

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

133 """Make a test set of DiaSources. 

134 

135 Parameters 

136 ---------- 

137 nForcedSources : `int` 

138 Number of sources to create. 

139 diaObjectIds : `numpy.ndarray` 

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

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

142 Exposure to create sources over. 

143 randomizeObjects : `bool`, optional 

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

145 output catalog, otherwise just iterate through them. 

146 

147 Returns 

148 ------- 

149 diaForcedSources : `pandas.DataFrame` 

150 DiaForcedSources generated across the exposure. 

151 """ 

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

153 ccdVisitId = exposure.info.id 

154 if randomizeObjects: 

155 objectIds = diaObjectIds[np.random.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 "ccdVisitId": ccdVisitId + i, 

164 "diaObjectId": objId, 

165 "midpointMjdTai": midpointMjdTai + 1.0 * i, 

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

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

168 "flags": 0}) 

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