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-07 14:40 +0000
« prev ^ index » next coverage.py v7.4.1, created at 2024-02-07 14:40 +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/>.
22from __future__ import annotations
24__all__ = ("StellarLocusFitAction",)
26from typing import cast
28import numpy as np
29import scipy.odr as scipyODR
30from lsst.pex.config import DictField
31from lsst.pipe.base import Struct
33from ...interfaces import KeyedData, KeyedDataAction, KeyedDataSchema, Scalar, Vector
34from ...math import sigmaMad
37def _stellarLocusFit(xs, ys, mags, paramDict):
38 """Make a fit to the stellar locus.
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:
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`).
74 Returns
75 -------
76 fitParams : `dict`
77 A dictionary of the calculated fit parameters:
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`]).
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)
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
124 linear = scipyODR.polynomial(1)
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
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"]
150 bPerpMin = yPerpMin - mPerp0 * xPerpMin
151 bPerpMax = yPerpMax - mPerp0 * xPerpMax
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
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])
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]
180 fitParams["bPerpMin"] = bPerpMin
181 fitParams["bPerpMax"] = bPerpMax
183 fitParams["mODR"] = float(params.beta[1])
184 fitParams["bODR"] = float(params.beta[0])
186 fitParams["mPerp"] = -1.0 / fitParams["mODR"]
187 fitParams["goodPoints"] = goodPoints
188 fitParams["fitPoints"] = fitPoints
189 fitParams["paramDict"] = paramDict
191 return fitParams
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
206def perpDistance(p1, p2, points):
207 """Calculate the perpendicular distance to a line from a point.
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.
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)
231 return dists
234def calcQuartileClippedStats(dataArray, nSigmaToClip=3.0):
235 """Calculate the quartile-based clipped statistics of a data array.
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.
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.
251 Returns
252 -------
253 result : `lsst.pipe.base.Struct`
254 The quartile-based clipped statistics with ``nSigmaToClip`` clipping.
255 Atributes are:
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))
283 return Struct(
284 median=median,
285 mean=quartileClippedMean,
286 stdDev=quartileClippedStdDev,
287 rms=quartileClippedRms,
288 clipValue=clipValue,
289 goodArray=good,
290 )
293class StellarLocusFitAction(KeyedDataAction):
294 r"""Determine Stellar Locus fit parameters from given input `Vector`\ s."""
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 )
321 def getInputSchema(self) -> KeyedDataSchema:
322 return (("x", Vector), ("y", Vector))
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
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"])
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"]
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"]
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 )
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]])
373 # Convert this to mmag.
374 dists = np.array(perpDistance(p1, p2, zip(xs[fitPoints], ys[fitPoints]))) * 1000
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"]
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"]
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
411 return fitParams