Coverage for python/lsst/analysis/tools/actions/keyedData/stellarLocusFit.py: 14%

98 statements  

« prev     ^ index     » next       coverage.py v7.2.6, created at 2023-05-24 02:36 -0700

1# This file is part of analysis_tools. 

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 

22from __future__ import annotations 

23 

24__all__ = ("StellarLocusFitAction",) 

25 

26from typing import cast 

27 

28import numpy as np 

29import scipy.odr as scipyODR 

30from lsst.pex.config import DictField 

31 

32from ...interfaces import KeyedData, KeyedDataAction, KeyedDataSchema, Scalar, Vector 

33from ...statistics import sigmaMad 

34 

35 

36def stellarLocusFit(xs, ys, paramDict): 

37 """Make a fit to the stellar locus. 

38 

39 Parameters 

40 ---------- 

41 xs : `numpy.ndarray` 

42 The color on the xaxis 

43 ys : `numpy.ndarray` 

44 The color on the yaxis 

45 paramDict : lsst.pex.config.dictField.Dict 

46 A dictionary of parameters for line fitting 

47 xMin : `float` 

48 The minimum x edge of the box to use for initial fitting 

49 xMax : `float` 

50 The maximum x edge of the box to use for initial fitting 

51 yMin : `float` 

52 The minimum y edge of the box to use for initial fitting 

53 yMax : `float` 

54 The maximum y edge of the box to use for initial fitting 

55 mHW : `float` 

56 The hardwired gradient for the fit 

57 bHW : `float` 

58 The hardwired intercept of the fit 

59 

60 Returns 

61 ------- 

62 paramsOut : `dict` 

63 A dictionary of the calculated fit parameters 

64 xMin : `float` 

65 The minimum x edge of the box to use for initial fitting 

66 xMax : `float` 

67 The maximum x edge of the box to use for initial fitting 

68 yMin : `float` 

69 The minimum y edge of the box to use for initial fitting 

70 yMax : `float` 

71 The maximum y edge of the box to use for initial fitting 

72 mHW : `float` 

73 The hardwired gradient for the fit 

74 bHW : `float` 

75 The hardwired intercept of the fit 

76 mODR : `float` 

77 The gradient calculated by the ODR fit 

78 bODR : `float` 

79 The intercept calculated by the ODR fit 

80 yBoxMin : `float` 

81 The y value of the fitted line at xMin 

82 yBoxMax : `float` 

83 The y value of the fitted line at xMax 

84 bPerpMin : `float` 

85 The intercept of the perpendicular line that goes through xMin 

86 bPerpMax : `float` 

87 The intercept of the perpendicular line that goes through xMax 

88 mODR2 : `float` 

89 The gradient from the second round of fitting 

90 bODR2 : `float` 

91 The intercept from the second round of fitting 

92 mPerp : `float` 

93 The gradient of the line perpendicular to the line from the 

94 second fit 

95 

96 Notes 

97 ----- 

98 The code does two rounds of fitting, the first is initiated using the 

99 hardwired values given in the `paramDict` parameter and is done using 

100 an Orthogonal Distance Regression fit to the points defined by the 

101 box of xMin, xMax, yMin and yMax. Once this fitting has been done a 

102 perpendicular bisector is calculated at either end of the line and 

103 only points that fall within these lines are used to recalculate the fit. 

104 """ 

105 # Points to use for the fit 

106 fitPoints = np.where( 

107 (xs > paramDict["xMin"]) 

108 & (xs < paramDict["xMax"]) 

109 & (ys > paramDict["yMin"]) 

110 & (ys < paramDict["yMax"]) 

111 )[0] 

112 

113 linear = scipyODR.polynomial(1) 

114 

115 data = scipyODR.Data(xs[fitPoints], ys[fitPoints]) 

116 odr = scipyODR.ODR(data, linear, beta0=[paramDict["bHW"], paramDict["mHW"]]) 

117 params = odr.run() 

118 mODR = float(params.beta[1]) 

119 bODR = float(params.beta[0]) 

120 

121 paramsOut = { 

122 "xMin": paramDict["xMin"], 

123 "xMax": paramDict["xMax"], 

124 "yMin": paramDict["yMin"], 

125 "yMax": paramDict["yMax"], 

126 "mHW": paramDict["mHW"], 

127 "bHW": paramDict["bHW"], 

128 "mODR": mODR, 

129 "bODR": bODR, 

130 } 

131 

132 # Having found the initial fit calculate perpendicular ends 

133 mPerp = -1.0 / mODR 

134 # When the gradient is really steep we need to use 

135 # the y limits of the box rather than the x ones 

136 

137 if np.abs(mODR) > 1: 

138 yBoxMin = paramDict["yMin"] 

139 xBoxMin = (yBoxMin - bODR) / mODR 

140 yBoxMax = paramDict["yMax"] 

141 xBoxMax = (yBoxMax - bODR) / mODR 

142 else: 

143 yBoxMin = mODR * paramDict["xMin"] + bODR 

144 xBoxMin = paramDict["xMin"] 

145 yBoxMax = mODR * paramDict["xMax"] + bODR 

146 xBoxMax = paramDict["xMax"] 

147 

148 bPerpMin = yBoxMin - mPerp * xBoxMin 

149 

150 paramsOut["yBoxMin"] = yBoxMin 

151 paramsOut["bPerpMin"] = bPerpMin 

152 

153 bPerpMax = yBoxMax - mPerp * xBoxMax 

154 

155 paramsOut["yBoxMax"] = yBoxMax 

156 paramsOut["bPerpMax"] = bPerpMax 

157 

158 # Use these perpendicular lines to chose the data and refit 

159 fitPoints = (ys > mPerp * xs + bPerpMin) & (ys < mPerp * xs + bPerpMax) 

160 data = scipyODR.Data(xs[fitPoints], ys[fitPoints]) 

161 odr = scipyODR.ODR(data, linear, beta0=[bODR, mODR]) 

162 params = odr.run() 

163 mODR = float(params.beta[1]) 

164 bODR = float(params.beta[0]) 

165 

166 paramsOut["mODR2"] = float(params.beta[1]) 

167 paramsOut["bODR2"] = float(params.beta[0]) 

168 

169 paramsOut["mPerp"] = -1.0 / paramsOut["mODR2"] 

170 

171 return paramsOut 

172 

173 

174def perpDistance(p1, p2, points): 

175 """Calculate the perpendicular distance to a line from a point. 

176 

177 Parameters 

178 ---------- 

179 p1 : `numpy.ndarray` 

180 A point on the line 

181 p2 : `numpy.ndarray` 

182 Another point on the line 

183 points : `zip` 

184 The points to calculate the distance to 

185 

186 Returns 

187 ------- 

188 dists : `list` 

189 The distances from the line to the points. Uses the cross 

190 product to work this out. 

191 """ 

192 dists = [] 

193 for point in points: 

194 point = np.array(point) 

195 distToLine = np.cross(p2 - p1, point - p1) / np.linalg.norm(p2 - p1) 

196 dists.append(distToLine) 

197 

198 return dists 

199 

200 

201class StellarLocusFitAction(KeyedDataAction): 

202 r"""Determine Stellar Locus fit parameters from given input `Vector`\ s.""" 

203 

204 stellarLocusFitDict = DictField[str, float]( 

205 doc="The parameters to use for the stellar locus fit. The default parameters are examples and are " 

206 "not useful for any of the fits. The dict needs to contain xMin/xMax/yMin/yMax which are the " 

207 "limits of the initial box for fitting the stellar locus, mHW and bHW are the initial " 

208 "intercept and gradient for the fitting.", 

209 default={"xMin": 0.1, "xMax": 0.2, "yMin": 0.1, "yMax": 0.2, "mHW": 0.5, "bHW": 0.0}, 

210 ) 

211 

212 def getInputSchema(self) -> KeyedDataSchema: 

213 return (("x", Vector), ("y", Vector)) 

214 

215 def getOutputSchema(self) -> KeyedDataSchema: 

216 value = ( 

217 (f"{self.identity or ''}_sigmaMAD", Scalar), 

218 (f"{self.identity or ''}_median", Scalar), 

219 (f"{self.identity or ''}_hardwired_sigmaMAD", Scalar), 

220 (f"{self.identity or ''}_hardwired_median", Scalar), 

221 ) 

222 return value 

223 

224 def __call__(self, data: KeyedData, **kwargs) -> KeyedData: 

225 xs = cast(Vector, data["x"]) 

226 ys = cast(Vector, data["y"]) 

227 fitParams = stellarLocusFit(xs, ys, self.stellarLocusFitDict) 

228 fitPoints = np.where( 

229 (xs > fitParams["xMin"]) # type: ignore 

230 & (xs < fitParams["xMax"]) # type: ignore 

231 & (ys > fitParams["yMin"]) # type: ignore 

232 & (ys < fitParams["yMax"]) # type: ignore 

233 )[0] 

234 

235 if np.fabs(fitParams["mHW"]) > 1: 

236 ysFitLineHW = np.array([fitParams["yMin"], fitParams["yMax"]]) 

237 xsFitLineHW = (ysFitLineHW - fitParams["bHW"]) / fitParams["mHW"] 

238 ysFitLine = np.array([fitParams["yMin"], fitParams["yMax"]]) 

239 xsFitLine = (ysFitLine - fitParams["bODR"]) / fitParams["mODR"] 

240 ysFitLine2 = np.array([fitParams["yMin"], fitParams["yMax"]]) 

241 xsFitLine2 = (ysFitLine2 - fitParams["bODR2"]) / fitParams["mODR2"] 

242 

243 else: 

244 xsFitLineHW = np.array([fitParams["xMin"], fitParams["xMax"]]) 

245 ysFitLineHW = fitParams["mHW"] * xsFitLineHW + fitParams["bHW"] 

246 xsFitLine = [fitParams["xMin"], fitParams["xMax"]] 

247 ysFitLine = np.array( 

248 [ 

249 fitParams["mODR"] * xsFitLine[0] + fitParams["bODR"], 

250 fitParams["mODR"] * xsFitLine[1] + fitParams["bODR"], 

251 ] 

252 ) 

253 xsFitLine2 = [fitParams["xMin"], fitParams["xMax"]] 

254 ysFitLine2 = np.array( 

255 [ 

256 fitParams["mODR2"] * xsFitLine2[0] + fitParams["bODR2"], 

257 fitParams["mODR2"] * xsFitLine2[1] + fitParams["bODR2"], 

258 ] 

259 ) 

260 

261 # Calculate the distances to that line 

262 # Need two points to characterise the lines we want 

263 # to get the distances to 

264 p1 = np.array([xsFitLine[0], ysFitLine[0]]) 

265 p2 = np.array([xsFitLine[1], ysFitLine[1]]) 

266 

267 p1HW = np.array([xsFitLine[0], ysFitLineHW[0]]) 

268 p2HW = np.array([xsFitLine[1], ysFitLineHW[1]]) 

269 

270 # Convert this to mmag 

271 distsHW = np.array(perpDistance(p1HW, p2HW, zip(xs[fitPoints], ys[fitPoints]))) * 1000 

272 dists = np.array(perpDistance(p1, p2, zip(xs[fitPoints], ys[fitPoints]))) * 1000 

273 

274 # Now we have the information for the perpendicular line we 

275 # can use it to calculate the points at the ends of the 

276 # perpendicular lines that intersect at the box edges 

277 if np.fabs(fitParams["mHW"]) > 1: 

278 xMid = (fitParams["yMin"] - fitParams["bODR2"]) / fitParams["mODR2"] 

279 xs = np.array([xMid - 0.5, xMid, xMid + 0.5]) 

280 ys = fitParams["mPerp"] * xs + fitParams["bPerpMin"] 

281 else: 

282 xs = np.array([fitParams["xMin"] - 0.2, fitParams["xMin"], fitParams["xMin"] + 0.2]) 

283 ys = xs * fitParams["mPerp"] + fitParams["bPerpMin"] 

284 

285 if np.fabs(fitParams["mHW"]) > 1: 

286 xMid = (fitParams["yMax"] - fitParams["bODR2"]) / fitParams["mODR2"] 

287 xs = np.array([xMid - 0.5, xMid, xMid + 0.5]) 

288 ys = fitParams["mPerp"] * xs + fitParams["bPerpMax"] 

289 else: 

290 xs = np.array([fitParams["xMax"] - 0.2, fitParams["xMax"], fitParams["xMax"] + 0.2]) 

291 ys = xs * fitParams["mPerp"] + fitParams["bPerpMax"] 

292 

293 fitParams[f"{self.identity or ''}_sigmaMAD"] = sigmaMad(dists) 

294 fitParams[f"{self.identity or ''}_median"] = np.median(dists) 

295 fitParams[f"{self.identity or ''}_hardwired_sigmaMAD"] = sigmaMad(distsHW) 

296 fitParams[f"{self.identity or ''}_hardwired_median"] = np.median(distsHW) 

297 

298 return fitParams # type: ignore