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

56 statements  

« prev     ^ index     » next       coverage.py v6.4.4, created at 2022-08-18 12:39 -0700

1from __future__ import annotations 

2 

3__all__ = ("StellarLocusFitAction",) 

4 

5from typing import cast 

6 

7import numpy as np 

8from lsst.pex.config import DictField 

9from scipy.stats import median_absolute_deviation as sigmaMad 

10 

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

12from ..plot.plotUtils import perpDistance, stellarLocusFit 

13 

14 

15class StellarLocusFitAction(KeyedDataAction): 

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

17 

18 stellarLocusFitDict = DictField[str, float]( 

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

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

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

22 "intercept and gradient for the fitting.", 

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

24 ) 

25 

26 def getInputSchema(self) -> KeyedDataSchema: 

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

28 

29 def getOutputSchema(self) -> KeyedDataSchema: 

30 value = ( 

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

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

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

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

35 ) 

36 return value 

37 

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

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

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

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

42 fitPoints = np.where( 

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

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

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

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

47 )[0] 

48 

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

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

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

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

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

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

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

56 

57 else: 

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

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

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

61 ysFitLine = np.array( 

62 [ 

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

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

65 ] 

66 ) 

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

68 ysFitLine2 = np.array( 

69 [ 

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

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

72 ] 

73 ) 

74 

75 # Calculate the distances to that line 

76 # Need two points to characterise the lines we want 

77 # to get the distances to 

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

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

80 

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

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

83 

84 distsHW = perpDistance(p1HW, p2HW, zip(xs[fitPoints], ys[fitPoints])) 

85 dists = perpDistance(p1, p2, zip(xs[fitPoints], ys[fitPoints])) 

86 

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

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

89 # perpendicular lines that intersect at the box edges 

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

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

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

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

94 else: 

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

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

97 

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

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

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

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

102 else: 

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

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

105 

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

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

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

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

110 

111 return fitParams # type: ignore