Coverage for tests/utils_tests.py: 16%

42 statements  

« prev     ^ index     » next       coverage.py v7.3.2, created at 2023-10-19 11:23 +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 

29import lsst.daf.base as dafBase 

30import lsst.geom 

31 

32 

33def makeDiaObjects(nObjects, exposure): 

34 """Make a test set of DiaObjects. 

35 

36 Parameters 

37 ---------- 

38 nObjects : `int` 

39 Number of objects to create. 

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

41 Exposure to create objects over. 

42 

43 Returns 

44 ------- 

45 diaObjects : `pandas.DataFrame` 

46 DiaObjects generated across the exposure. 

47 """ 

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

49 rand_x = np.random.uniform(bbox.getMinX(), bbox.getMaxX(), size=nObjects) 

50 rand_y = np.random.uniform(bbox.getMinY(), bbox.getMaxY(), size=nObjects) 

51 

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

53 

54 data = [] 

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

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

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

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

59 "radecMjdTai": midpointMjdTai, 

60 "diaObjectId": idx + 1, 

61 "pmParallaxNdata": 0, 

62 "nearbyObj1": 0, 

63 "nearbyObj2": 0, 

64 "nearbyObj3": 0, 

65 "flags": 1, 

66 "nDiaSources": 5} 

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

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

69 data.append(newObject) 

70 

71 return pd.DataFrame(data=data) 

72 

73 

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

75 """Make a test set of DiaSources. 

76 

77 Parameters 

78 ---------- 

79 nSources : `int` 

80 Number of sources to create. 

81 diaObjectIds : `numpy.ndarray` 

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

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

84 Exposure to create sources over. 

85 randomizeObjects : `bool`, optional 

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

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

88 necessary to get nSources objectIds. 

89 

90 Returns 

91 ------- 

92 diaSources : `pandas.DataFrame` 

93 DiaSources generated across the exposure. 

94 """ 

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

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

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

98 if randomizeObjects: 

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

100 else: 

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

102 

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

104 

105 data = [] 

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

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

108 # Put together the minimum values for the alert. 

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

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

111 "x": x, 

112 "y": y, 

113 "ccdVisitId": exposure.info.id, 

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

115 "diaObjectId": objId, 

116 "ssObjectId": 0, 

117 "parentDiaSourceId": 0, 

118 "diaSourceId": idx + 1, 

119 "midpointMjdTai": midpointMjdTai + 1.0 * idx, 

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

121 "psfNdata": 0, 

122 "trailNdata": 0, 

123 "dipoleNdata": 0, 

124 "flags": 1}) 

125 

126 return pd.DataFrame(data=data) 

127 

128 

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

130 """Make a test set of DiaSources. 

131 

132 Parameters 

133 ---------- 

134 nForcedSources : `int` 

135 Number of sources to create. 

136 diaObjectIds : `numpy.ndarray` 

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

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

139 Exposure to create sources over. 

140 randomizeObjects : `bool`, optional 

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

142 output catalog, otherwise just iterate through them. 

143 

144 Returns 

145 ------- 

146 diaForcedSources : `pandas.DataFrame` 

147 DiaForcedSources generated across the exposure. 

148 """ 

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

150 ccdVisitId = exposure.info.id 

151 if randomizeObjects: 

152 objectIds = diaObjectIds[np.random.randint(len(diaObjectIds), size=nForcedSources)] 

153 else: 

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

155 

156 data = [] 

157 for i, objId in enumerate(objectIds): 

158 # Put together the minimum values for the alert. 

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

160 "ccdVisitId": ccdVisitId + i, 

161 "diaObjectId": objId, 

162 "midpointMjdTai": midpointMjdTai + 1.0 * i, 

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

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

165 "flags": 0}) 

166 

167 return pd.DataFrame(data=data)