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

141 statements  

« prev     ^ index     » next       coverage.py v7.4.1, created at 2024-02-03 13:22 +0000

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 

31from lsst.pipe.base import Struct 

32 

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

34from ...math import sigmaMad 

35 

36 

37def _stellarLocusFit(xs, ys, mags, paramDict): 

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

39 

40 Parameters 

41 ---------- 

42 xs : `numpy.ndarray` [`float`] 

43 The color on the xaxis. 

44 ys : `numpy.ndarray` [`float`] 

45 The color on the yaxis. 

46 mags : `numpy.ndarray` [`float`] 

47 The magnitude of the reference band flux (in mag). 

48 paramDict : `dict` [`str`, `float`] 

49 A dictionary of parameters for line fitting: 

50 

51 ``"xMin"`` 

52 The minimum x edge of the box to use for initial fitting (`float`). 

53 ``"xMax"`` 

54 The maximum x edge of the box to use for initial fitting (`float`). 

55 ``"yMin"`` 

56 The minimum y edge of the box to use for initial fitting (`float`). 

57 ``"yMax"`` 

58 The maximum y edge of the box to use for initial fitting (`float`). 

59 ``"mHW"`` 

60 The hardwired gradient for the fit (`float`). 

61 ``"bHw"`` 

62 The hardwired intercept of the fit (`float`). 

63 ``"nSigmaToClip1"`` 

64 The number of sigma perpendicular to the fit to clip in the initial 

65 fitting loop (`float`). This should probably be stricter than the 

66 final iteration (i.e. nSigmaToClip1 < nSigmaToClip2). 

67 ``"nSigmaToClip2"`` 

68 The number of sigma perpendicular to the fit to clip in the final 

69 fitting loop (`float`). 

70 ``"minObjectForFit"`` 

71 Minimum number of objects surviving cuts to attempt fit. If not 

72 met, return NANs for values in ``fitParams`` (`int`). 

73 

74 Returns 

75 ------- 

76 fitParams : `dict` 

77 A dictionary of the calculated fit parameters: 

78 

79 ``"bPerpMin"`` 

80 The intercept of the perpendicular line that goes through xMin 

81 (`float`). 

82 ``"bPerpMax"`` 

83 The intercept of the perpendicular line that goes through xMax 

84 (`float`). 

85 ``"mODR"`` 

86 The gradient from the final round of fitting (`float`). 

87 ``"bODR"`` 

88 The intercept from the final round of fitting (`float`). 

89 ``"mPerp"`` 

90 The gradient of the line perpendicular to the line from the final 

91 fit (`float`). 

92 ``"fitPoints"`` 

93 A boolean list indicating which points were used in the final fit 

94 (`list` [`bool`]). 

95 

96 Notes 

97 ----- 

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

99 fixed values given in ``paramDict`` and is done using an Orthogonal 

100 Distance Regression (ODR) fit to the points defined by the box with limits 

101 defined by the keys: xMin, xMax, yMin, and yMax. Once this fitting has been 

102 done a 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 We also perform clipping of points perpendicular to the fit line that have 

105 distances that deviate more than nSigmaToClip1/2 (for an initial and final 

106 iteration) from the fit. 

107 """ 

108 fitParams = {} 

109 # Initial subselection of points to use for the fit 

110 # Check for nans/infs 

111 goodPoints = np.isfinite(xs) & np.isfinite(ys) & np.isfinite(mags) 

112 

113 fitPoints = ( 

114 goodPoints 

115 & (xs > paramDict["xMin"]) 

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

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

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

119 ) 

120 if sum(fitPoints) < paramDict["minObjectForFit"]: 

121 fitParams = _setFitParamsNans(fitParams, fitPoints, paramDict) 

122 return fitParams 

123 

124 linear = scipyODR.polynomial(1) 

125 

126 fitData = scipyODR.Data(xs[fitPoints], ys[fitPoints]) 

127 odr = scipyODR.ODR(fitData, linear, beta0=[paramDict["bFixed"], paramDict["mFixed"]]) 

128 params = odr.run() 

129 mODR0 = float(params.beta[1]) 

130 bODR0 = float(params.beta[0]) 

131 mPerp0 = -1.0 / mODR0 

132 

133 # Loop twice over the fit and include sigma clipping of points 

134 # perpendicular to the fit line (stricter on first iteration). 

135 for nSigmaToClip in [paramDict["nSigmaToClip1"], paramDict["nSigmaToClip2"]]: 

136 # Having found the initial fit calculate perpendicular ends. 

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

138 # y limits of the fit line rather than the x ones. 

139 if np.abs(mODR0) > 1: 

140 yPerpMin = paramDict["yMin"] 

141 xPerpMin = (yPerpMin - bODR0) / mODR0 

142 yPerpMax = paramDict["yMax"] 

143 xPerpMax = (yPerpMax - bODR0) / mODR0 

144 else: 

145 yPerpMin = mODR0 * paramDict["xMin"] + bODR0 

146 xPerpMin = paramDict["xMin"] 

147 yPerpMax = mODR0 * paramDict["xMax"] + bODR0 

148 xPerpMax = paramDict["xMax"] 

149 

150 bPerpMin = yPerpMin - mPerp0 * xPerpMin 

151 bPerpMax = yPerpMax - mPerp0 * xPerpMax 

152 

153 # Use these perpendicular lines to choose the data and refit. 

154 fitPoints = (ys > mPerp0 * xs + bPerpMin) & (ys < mPerp0 * xs + bPerpMax) 

155 if sum(fitPoints) < paramDict["minObjectForFit"]: 

156 fitParams = _setFitParamsNans(fitParams, fitPoints, paramDict) 

157 return fitParams 

158 

159 p1 = np.array([0, bODR0]) 

160 p2 = np.array([-bODR0 / mODR0, 0]) 

161 if np.abs(sum(p1 - p2) < 1e12): # p1 and p2 must be different. 

162 p2 = np.array([(1.0 - bODR0 / mODR0), 1.0]) 

163 

164 # Sigma clip points based on perpendicular distance (in mmag) to 

165 # current fit. 

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

167 clippedStats = calcQuartileClippedStats(fitDists, nSigmaToClip=nSigmaToClip) 

168 allDists = np.array(perpDistance(p1, p2, zip(xs, ys))) * 1000 

169 keep = np.abs(allDists) <= clippedStats.clipValue 

170 fitPoints &= keep 

171 if sum(fitPoints) < paramDict["minObjectForFit"]: 

172 fitParams = _setFitParamsNans(fitParams, fitPoints, paramDict) 

173 return fitParams 

174 fitData = scipyODR.Data(xs[fitPoints], ys[fitPoints]) 

175 odr = scipyODR.ODR(fitData, linear, beta0=[bODR0, mODR0]) 

176 params = odr.run() 

177 mODR0 = params.beta[1] 

178 bODR0 = params.beta[0] 

179 

180 fitParams["bPerpMin"] = bPerpMin 

181 fitParams["bPerpMax"] = bPerpMax 

182 

183 fitParams["mODR"] = float(params.beta[1]) 

184 fitParams["bODR"] = float(params.beta[0]) 

185 

186 fitParams["mPerp"] = -1.0 / fitParams["mODR"] 

187 fitParams["goodPoints"] = goodPoints 

188 fitParams["fitPoints"] = fitPoints 

189 fitParams["paramDict"] = paramDict 

190 

191 return fitParams 

192 

193 

194def _setFitParamsNans(fitParams, fitPoints, paramDict): 

195 fitParams["bPerpMin"] = np.nan 

196 fitParams["bPerpMax"] = np.nan 

197 fitParams["mODR"] = np.nan 

198 fitParams["bODR"] = np.nan 

199 fitParams["mPerp"] = np.nan 

200 fitParams["goodPoints"] = np.nan 

201 fitParams["fitPoints"] = fitPoints 

202 fitParams["paramDict"] = paramDict 

203 return fitParams 

204 

205 

206def perpDistance(p1, p2, points): 

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

208 

209 Parameters 

210 ---------- 

211 p1 : `numpy.ndarray` [`float`] 

212 A point on the line. 

213 p2 : `numpy.ndarray` [`float`] 

214 Another point on the line. 

215 points : `zip` [(`float`, `float`)] 

216 The points to calculate the distance to. 

217 

218 Returns 

219 ------- 

220 dists : `numpy.ndarray` [`float`] 

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

222 product to work this out. 

223 """ 

224 if sum(p2 - p1) == 0: 

225 raise ValueError(f"Must supply two different points for p1, p2. Got {p1}, {p2}") 

226 points = list(points) 

227 if len(points) == 0: 

228 raise ValueError("Must provied a non-empty zip() list of points.") 

229 dists = np.cross(p2 - p1, points - p1) / np.linalg.norm(p2 - p1) 

230 

231 return dists 

232 

233 

234def calcQuartileClippedStats(dataArray, nSigmaToClip=3.0): 

235 """Calculate the quartile-based clipped statistics of a data array. 

236 

237 The difference between quartiles[2] and quartiles[0] is the interquartile 

238 distance. 0.74*interquartileDistance is an estimate of standard deviation 

239 so, in the case that ``dataArray`` has an approximately Gaussian 

240 distribution, this is equivalent to nSigma clipping. 

241 

242 Parameters 

243 ---------- 

244 dataArray : `list` or `numpy.ndarray` [`float`] 

245 List or array containing the values for which the quartile-based 

246 clipped statistics are to be calculated. 

247 nSigmaToClip : `float`, optional 

248 Number of \"sigma\" outside of which to clip data when computing the 

249 statistics. 

250 

251 Returns 

252 ------- 

253 result : `lsst.pipe.base.Struct` 

254 The quartile-based clipped statistics with ``nSigmaToClip`` clipping. 

255 Atributes are: 

256 

257 ``median`` 

258 The median of the full ``dataArray`` (`float`). 

259 ``mean`` 

260 The quartile-based clipped mean (`float`). 

261 ``stdDev`` 

262 The quartile-based clipped standard deviation (`float`). 

263 ``rms`` 

264 The quartile-based clipped root-mean-squared (`float`). 

265 ``clipValue`` 

266 The value outside of which to clip the data before computing the 

267 statistics (`float`). 

268 ``goodArray`` 

269 A boolean array indicating which data points in ``dataArray`` were 

270 used in the calculation of the statistics, where `False` indicates 

271 a clipped datapoint (`numpy.ndarray` of `bool`). 

272 """ 

273 quartiles = np.percentile(dataArray, [25, 50, 75]) 

274 assert len(quartiles) == 3 

275 median = quartiles[1] 

276 interQuartileDistance = quartiles[2] - quartiles[0] 

277 clipValue = nSigmaToClip * 0.74 * interQuartileDistance 

278 good = np.abs(dataArray - median) <= clipValue 

279 quartileClippedMean = dataArray[good].mean() 

280 quartileClippedStdDev = dataArray[good].std() 

281 quartileClippedRms = np.sqrt(np.mean(dataArray[good] ** 2)) 

282 

283 return Struct( 

284 median=median, 

285 mean=quartileClippedMean, 

286 stdDev=quartileClippedStdDev, 

287 rms=quartileClippedRms, 

288 clipValue=clipValue, 

289 goodArray=good, 

290 ) 

291 

292 

293class StellarLocusFitAction(KeyedDataAction): 

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

295 

296 stellarLocusFitDict = DictField[str, float]( 

297 doc="The parameters to use for the stellar locus fit. For xMin/Max, yMin/Max, " 

298 "and m/bFixed, the default parameters are examples and are not generally useful " 

299 "for any of the fits, so should be updated in the PlotAction definition in the " 

300 "atools directory. The dict needs to contain xMin/xMax/yMin/yMax which are the " 

301 "limits of the initial point selection box for fitting the stellar locus, mFixed " 

302 "and bFixed are meant to represent the intercept and gradient of a canonical fit " 

303 "for a given dataset (and should be derived from data). They are used here as an " 

304 "initial guess for the fitting. nSigmaToClip1/2 set the number of sigma to clip " 

305 "perpendicular the fit in the first and second fit iterations after the initial " 

306 "guess and point selection fit. minObjectForFit sets a minimum number of points " 

307 "deemed suitable for inclusion in the fit in order to bother attempting the fit.", 

308 default={ 

309 "xMin": 0.1, 

310 "xMax": 0.2, 

311 "yMin": 0.1, 

312 "yMax": 0.2, 

313 "mHW": 0.5, 

314 "bHW": 0.0, 

315 "nSigmaToClip1": 3.5, 

316 "nSigmaToClip2": 5.0, 

317 "minObjectForFit": 3, 

318 }, 

319 ) 

320 

321 def getInputSchema(self) -> KeyedDataSchema: 

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

323 

324 def getOutputSchema(self) -> KeyedDataSchema: 

325 value = ( 

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

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

328 ) 

329 return value 

330 

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

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

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

334 mags = cast(Vector, data["mag"]) 

335 

336 fitParams = _stellarLocusFit(xs, ys, mags, self.stellarLocusFitDict) 

337 # Bail out if there were not enough points to fit. 

338 for value in fitParams.values(): 

339 if isinstance(value, float): 

340 if np.isnan(value): 

341 fitParams[f"{self.identity or ''}_sigmaMAD"] = np.nan 

342 fitParams[f"{self.identity or ''}_median"] = np.nan 

343 return fitParams 

344 fitPoints = fitParams["fitPoints"] 

345 

346 if np.fabs(self.stellarLocusFitDict["mFixed"]) > 1: 

347 ysFitLineFixed = np.array([self.stellarLocusFitDict["yMin"], self.stellarLocusFitDict["yMax"]]) 

348 xsFitLineFixed = (ysFitLineFixed - self.stellarLocusFitDict["bFixed"]) / self.stellarLocusFitDict[ 

349 "mFixed" 

350 ] 

351 ysFitLine = np.array([self.stellarLocusFitDict["yMin"], self.stellarLocusFitDict["yMax"]]) 

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

353 

354 else: 

355 xsFitLineFixed = np.array([self.stellarLocusFitDict["xMin"], self.stellarLocusFitDict["xMax"]]) 

356 ysFitLineFixed = ( 

357 self.stellarLocusFitDict["mFixed"] * xsFitLineFixed + self.stellarLocusFitDict["bFixed"] 

358 ) 

359 xsFitLine = [self.stellarLocusFitDict["xMin"], self.stellarLocusFitDict["xMax"]] 

360 ysFitLine = np.array( 

361 [ 

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

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

364 ] 

365 ) 

366 

367 # Calculate the distances to that line. 

368 # Need two points to characterize the lines we want to get the 

369 # distances to. 

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

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

372 

373 # Convert this to mmag. 

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

375 

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

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

378 # perpendicular lines that intersect at the box edges. 

379 if np.fabs(self.stellarLocusFitDict["mFixed"]) > 1: 

380 xMid = (self.stellarLocusFitDict["yMin"] - fitParams["bODR"]) / fitParams["mODR"] 

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

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

383 else: 

384 xs = np.array( 

385 [ 

386 self.stellarLocusFitDict["xMin"] - 0.2, 

387 self.stellarLocusFitDict["xMin"], 

388 self.stellarLocusFitDict["xMin"] + 0.2, 

389 ] 

390 ) 

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

392 

393 if np.fabs(self.stellarLocusFitDict["mFixed"]) > 1: 

394 xMid = (self.stellarLocusFitDict["yMax"] - fitParams["bODR"]) / fitParams["mODR"] 

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

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

397 else: 

398 xs = np.array( 

399 [ 

400 self.stellarLocusFitDict["xMax"] - 0.2, 

401 self.stellarLocusFitDict["xMax"], 

402 self.stellarLocusFitDict["xMax"] + 0.2, 

403 ] 

404 ) 

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

406 

407 fit_sigma, fit_med = (sigmaMad(dists), np.median(dists)) if len(dists) else (np.nan, np.nan) 

408 fitParams[f"{self.identity or ''}_sigmaMAD"] = fit_sigma 

409 fitParams[f"{self.identity or ''}_median"] = fit_med 

410 

411 return fitParams