Coverage for python / lsst / meas / extensions / multiprofit / fit_coadd_psf.py: 24%

1# This file is part of meas_extensions_multiprofit. 

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 typing import ClassVar 

23 

24from lsst.afw.detection import InvalidPsfError 

25from lsst.daf.butler.formatters.parquet import astropy_to_arrow 

26import lsst.gauss2d.fit as g2f 

27from lsst.multiprofit import ( 

28 ComponentGroupConfig, 

29 FluxFractionParameterConfig, 

30 FluxParameterConfig, 

31 GaussianComponentConfig, 

32 ParameterConfig, 

33 SourceConfig, 

34) 

35from lsst.multiprofit.fitting.fit_psf import ( 

36 CatalogPsfFitter, 

37 CatalogPsfFitterConfig, 

38 CatalogPsfFitterConfigData, 

39) 

40import lsst.pex.config as pexConfig 

41import lsst.pipe.base as pipeBase 

42import lsst.pipe.tasks.fit_coadd_psf as fitCP 

43import lsst.utils.timer as utilsTimer 

44 

45from .errors import IsParentError 

46 

47 

48class MultiProFitPsfConfig(CatalogPsfFitterConfig, fitCP.CoaddPsfFitSubConfig): 

49 """Configuration for the MultiProFit Gaussian mixture PSF fitter.""" 

50 

51 fit_parents = pexConfig.Field[bool](default=False, doc="Whether to fit parent object PSFs") 

52 initialize_ellipses = pexConfig.Field[bool]( 

53 default=True, 

54 doc="Whether to initialize the ellipse parameters from the model config; if False, they " 

55 "will remain at the best-fit values for the previous source's PSF", 

56 ) 

57 prefix_column = pexConfig.Field[str](default="mpf_deepCoaddPsf_", doc="Column name prefix") 

58 

59 def setDefaults(self): 

60 super().setDefaults() 

61 self.model = SourceConfig( 

62 component_groups={ 

63 "": ComponentGroupConfig( 

64 components_gauss={ 

65 "gauss1": GaussianComponentConfig( 

66 size_x=ParameterConfig(value_initial=1.5), 

67 size_y=ParameterConfig(value_initial=1.5), 

68 fluxfrac=FluxFractionParameterConfig(value_initial=0.5), 

69 flux=FluxParameterConfig(value_initial=1.0, fixed=True), 

70 ), 

71 "gauss2": GaussianComponentConfig( 

72 size_x=ParameterConfig(value_initial=3.0), 

73 size_y=ParameterConfig(value_initial=3.0), 

74 fluxfrac=FluxFractionParameterConfig(value_initial=1.0, fixed=True), 

75 ), 

76 }, 

77 is_fractional=True, 

78 ) 

79 } 

80 ) 

81 self.flag_errors = {"no_inputs_flag": "InvalidPsfError"} 

82 self.prefix_column = "TwoGaussianPsf_" 

83 

84 

85class MultiProFitPsfTask(CatalogPsfFitter, fitCP.CoaddPsfFitSubTask): 

86 """Fit a Gaussian mixture PSF model at cataloged locations. 

87 

88 This task uses MultiProFit to fit a PSF model to the coadd PSF, 

89 evaluated at the centroid of each source in the corresponding 

90 catalog. 

91 

92 Parameters 

93 ---------- 

94 **kwargs 

95 Keyword arguments to pass to CoaddPsfFitSubTask.__init__. 

96 """ 

97 

98 ConfigClass: ClassVar = MultiProFitPsfConfig 

99 _DefaultName: ClassVar = "multiProFitPsf" 

100 

101 def __init__(self, **kwargs): 

102 errors_expected = {} if "errors_expected" not in kwargs else kwargs.pop("errors_expected") 

103 if InvalidPsfError not in errors_expected: 

104 # Cannot compute CoaddPsf at point (x, y) 

105 errors_expected[InvalidPsfError] = "no_inputs_flag" 

106 CatalogPsfFitter.__init__(self, errors_expected=errors_expected) 

107 fitCP.CoaddPsfFitSubTask.__init__(self, **kwargs) 

108 

109 def check_source(self, source, config): 

110 if ( 

111 config 

112 and hasattr(config, "fit_parents") 

113 and not config.fit_parents 

114 and (source["parent"] == 0) 

115 and (source["deblend_nChild"] > 1) 

116 ): 

117 raise IsParentError( 

118 f"{source['id']=} is a parent with nChild={source['deblend_nChild']}" f" and will be skipped" 

119 ) 

120 

121 def initialize_model( 

122 self, 

123 model: g2f.ModelD, 

124 config_data: CatalogPsfFitterConfigData, 

125 limits_x: g2f.LimitsD = None, 

126 limits_y: g2f.LimitsD = None, 

127 ) -> None: 

128 """Initialize a ModelD for a single source row. 

129 

130 Parameters 

131 ---------- 

132 model 

133 The model object to initialize. 

134 config_data 

135 The fitter config with cached data. 

136 limits_x 

137 Hard limits for the source's x centroid. 

138 limits_y 

139 Hard limits for the source's y centroid. 

140 """ 

141 n_rows, n_cols = model.data[0].image.data.shape 

142 cen_x, cen_y = n_cols / 2.0, n_rows / 2.0 

143 centroids = set() 

144 if limits_x is None: 

145 limits_x = g2f.LimitsD(0, n_cols) 

146 if limits_y is None: 

147 limits_y = g2f.LimitsD(0, n_rows) 

148 

149 for component, config_comp in zip( 

150 config_data.components.values(), config_data.component_configs.values() 

151 ): 

152 centroid = component.centroid 

153 if centroid not in centroids: 

154 centroid.x_param.value = cen_x 

155 centroid.x_param.limits = limits_x 

156 centroid.y_param.value = cen_y 

157 centroid.y_param.limits = limits_y 

158 centroids.add(centroid) 

159 

160 if self.config.initialize_ellipses: 

161 ellipse = component.ellipse 

162 ellipse.size_x_param.limits = limits_x 

163 ellipse.size_x = config_comp.size_x.value_initial 

164 ellipse.size_y_param.limits = limits_y 

165 ellipse.size_y = config_comp.size_y.value_initial 

166 ellipse.rho = config_comp.rho.value_initial 

167 

168 @utilsTimer.timeMethod 

169 def run( 

170 self, 

171 catexp: fitCP.CatalogExposurePsf, 

172 **kwargs, 

173 ) -> pipeBase.Struct: 

174 """Run the MultiProFit PSF task on a catalog-exposure pair. 

175 

176 Parameters 

177 ---------- 

178 catexp 

179 An exposure to fit a model PSF at the position of all 

180 sources in the corresponding catalog. 

181 **kwargs 

182 Additional keyword arguments to pass to self.fit. 

183 

184 Returns 

185 ------- 

186 catalog 

187 A table with fit parameters for the PSF model at the location 

188 of each source. 

189 """ 

190 is_parent_name = IsParentError.column_name() 

191 if not self.config.fit_parents: 

192 if is_parent_name not in self.errors_expected: 

193 self.errors_expected[IsParentError] = is_parent_name 

194 if "IsParentError" not in self.config.flag_errors.values(): 

195 self.config._frozen = False 

196 self.config.flag_errors[is_parent_name] = "IsParentError" 

197 self.config._frozen = True 

198 elif is_parent_name in self.errors_expected: 

199 del self.errors_expected[is_parent_name] 

200 if is_parent_name in self.config.flag_errors.keys(): 

201 self.config._frozen = False 

202 del self.config.flag_errors[is_parent_name] 

203 self.config._frozen = True 

204 

205 config_data = CatalogPsfFitterConfigData(config=self.config) 

206 catalog = self.fit(catexp=catexp, config_data=config_data, **kwargs) 

207 return pipeBase.Struct(output=astropy_to_arrow(catalog))