Coverage for tests/test_fit_turbulence.py: 99%

179 statements  

« prev     ^ index     » next       coverage.py v7.16.2, created at 2026-10-17 09:42 +0000

1# This file is part of drp_tasks. 

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/>. 

21import unittest 

22 

23import astropy.units as u 

24import numpy as np 

25import treegp 

26from astropy.table import Table, vstack 

27 

28import lsst.afw.table as afwTable 

29import lsst.pex.config as pexConfig 

30import lsst.utils.tests 

31from lsst.drp.tasks.fit_turbulence import ( 

32 GaussianProcessesTurbulenceFitConfig, 

33 GaussianProcessesTurbulenceFitTask, 

34) 

35from lsst.obs.base import createInitialSkyWcsFromBoresight 

36from lsst.obs.base.instrument_tests import DummyCam 

37from lsst.pipe.base import AnnotatedPartialOutputsError 

38 

39 

40class TurbulenceTestDataMixin: 

41 """Shared helpers to simulate WCS and residual catalogs.""" 

42 

43 def _makeWcsCatalog(self, visit, camera): 

44 """Make a simulated WCS catalog.""" 

45 schema = afwTable.ExposureTable.makeMinimalSchema() 

46 schema.addField("visit", type="L", doc="Visit number") 

47 wcsCatalog = afwTable.ExposureCatalog(schema) 

48 

49 boresight = lsst.geom.SpherePoint(150 * lsst.geom.degrees, 2.5 * lsst.geom.degrees) 

50 

51 for detectorId in [0, 1]: 

52 detector = camera[detectorId] 

53 wcs = createInitialSkyWcsFromBoresight(boresight, 0 * lsst.geom.degrees, detector) 

54 record = wcsCatalog.addNew() 

55 record.setId(detectorId) 

56 record.setWcs(wcs) 

57 record.setBBox(detector.getBBox()) 

58 wcsCatalog["visit"] = visit 

59 return wcsCatalog 

60 

61 def _makeResidualCatalog(self, wcsCatalog, visit, kernel): 

62 """Make a simulated catalog of source positions and residuals.""" 

63 visitPositions = [] 

64 for detector in [0, 1]: 

65 bbox = wcsCatalog[detector].getBBox() 

66 pixToTP = wcsCatalog[detector].wcs.getFrameDict().getMapping("PIXELS", "IWC") 

67 

68 x = self.rng.uniform(bbox.getBeginX(), bbox.getEndX(), self.npoints) 

69 y = self.rng.uniform(bbox.getBeginY(), bbox.getEndY(), self.npoints) 

70 xTP, yTP = pixToTP.applyForward(np.array([x, y])) 

71 xSky, ySky = wcsCatalog[detector].wcs.pixelToSkyArray(x, y, degrees=True) 

72 detectorPositions = Table( 

73 { 

74 "xpix": x, 

75 "ypix": y, 

76 "xworld": xTP, 

77 "yworld": yTP, 

78 "xsky": xSky, 

79 "ysky": ySky, 

80 } 

81 ) 

82 detectorPositions["deviceName"] = detector 

83 visitPositions.append(detectorPositions) 

84 visitPositions = vstack(visitPositions) 

85 totPoints = len(visitPositions) 

86 

87 K = kernel.__call__(np.array([visitPositions["xworld"], visitPositions["yworld"]]).T) 

88 dx = self.rng.multivariate_normal(np.zeros(totPoints), K) 

89 dy = self.rng.multivariate_normal(np.zeros(totPoints), K) 

90 

91 dx += self.rng.normal(scale=self.noise, size=totPoints) 

92 dy += self.rng.normal(scale=self.noise, size=totPoints) 

93 

94 residError = np.ones(totPoints) * self.noise 

95 visitPositions["xresw"] = dx 

96 visitPositions["yresw"] = dy 

97 visitPositions["covTotalW_00"] = residError**2 

98 visitPositions["covTotalW_11"] = residError**2 

99 visitPositions["exposureName"] = str(visit) 

100 

101 return visitPositions 

102 

103 

104class FitTurbulenceLeget21TestCase(TurbulenceTestDataMixin, lsst.utils.tests.TestCase): 

105 

106 def setUp(self): 

107 

108 visit = 1234 

109 self.noise = 0.01 

110 self.npoints = 250 

111 self.rng = np.random.RandomState(12345) 

112 

113 config = GaussianProcessesTurbulenceFitConfig() 

114 # Set the max separation lower, since we only have two detectors. 

115 config.correlationSeparationMax = 0.1 

116 config.initKernel = "15**2 * AnisotropicVonKarman(invLam=array([[1.0/0.1**2,0],[0,1.0/0.1**2]]))" 

117 config.splineNNodes = 60 

118 self.task = GaussianProcessesTurbulenceFitTask(config=config) 

119 

120 kernel_skl = treegp.eval_kernel(self.task.config.initKernel) 

121 

122 # Make a simulated WCS catalog and position catalog. 

123 self.wcsCatalog = self._makeWcsCatalog(visit, DummyCam().getCamera()) 

124 self.positions = self._makeResidualCatalog(self.wcsCatalog, visit, kernel_skl) 

125 

126 # Run the Gaussian Processes and predict at the input positions. 

127 self.gpx, self.gpy, *_ = self.task.runGP(self.wcsCatalog, self.positions) 

128 

129 sourceCatalog = Table( 

130 { 

131 "x": self.positions["xpix"], 

132 "y": self.positions["ypix"], 

133 "detector": self.positions["deviceName"], 

134 } 

135 ) 

136 self.predictCat = self.task.predict(self.gpx, self.gpy, self.wcsCatalog, sourceCatalog) 

137 

138 def test_runGP(self): 

139 """Check that the predicted position is about the same as the simulated 

140 dx/dy values, modulo the scatter in the data.""" 

141 dRA = self.positions["xsky"] - self.predictCat["coord_ra"] 

142 dDec = self.positions["ysky"] - self.predictCat["coord_dec"] 

143 

144 np.testing.assert_allclose(dRA, (self.positions["xresw"] * u.mas).to(u.degree).value, rtol=1e-1) 

145 np.testing.assert_allclose(dDec, (self.positions["yresw"] * u.mas).to(u.degree).value, rtol=1e-1) 

146 

147 def test_makeWcs(self): 

148 """Check that the output WCS maps pixel positions to the same RA and 

149 Dec as the Gaussian Processes prediction. 

150 """ 

151 

152 wcsWithGP = self.task.addGPToWcs(self.gpx, self.gpy, self.wcsCatalog) 

153 for detectorRow in self.wcsCatalog: 

154 detector = detectorRow.getId() 

155 newWcs = wcsWithGP.find(detector).wcs 

156 

157 sources = self.predictCat[self.predictCat["detector"] == detector] 

158 ra, dec = newWcs.pixelToSkyArray(sources["x"], sources["y"], degrees=True) 

159 

160 np.testing.assert_allclose(ra, sources["coord_ra"]) 

161 np.testing.assert_allclose(dec, sources["coord_dec"], rtol=1e-6) 

162 

163 def test_raiseError(self): 

164 """Check that fitting insufficient data raises an 

165 AnnotatedPartialOutputsError.""" 

166 # Try fitting GP with only a small amount of data, which should fail. 

167 with self.assertRaises(AnnotatedPartialOutputsError): 

168 self.task.runGP(self.wcsCatalog, self.positions[:5]) 

169 

170 

171@unittest.skipIf(not hasattr(treegp, "empirical_2pcf"), "treegp lacks the 'empirical-2pcf' optimizer") 

172class FitTurbulenceTurboGpTestCase(TurbulenceTestDataMixin, lsst.utils.tests.TestCase): 

173 """Test the 'empirical-2pcf' (TURBO-GP, RTN-129) optimizer and the 

174 per-source output table.""" 

175 

176 def setUp(self): 

177 

178 self.visit = 1234 

179 # Realistic measurement noise (mas); near-noiseless data leave the 

180 # kernel system too ill-conditioned for a reliable solve. 

181 self.noise = 1.0 

182 self.npoints = 250 

183 self.rng = np.random.RandomState(12345) 

184 

185 config = GaussianProcessesTurbulenceFitConfig() 

186 config.optimizer = "empirical-2pcf" 

187 # The correlation length of the simulated field (0.02 degrees, 

188 # below) must sit well inside the two-detector footprint 

189 # (~0.16 degrees) for the 2-point correlation function to be 

190 # measurable, and the 2-pcf scales must match it. 

191 config.correlationSeparationMax = 0.06 

192 config.correlationPixelSize = 0.005 

193 config.apodRadius = 0.03 

194 # Force a non-empty validation set from the 500 simulated points. 

195 config.maxTrainingPoints = 400 

196 config.saveSourceTable = True 

197 self.task = GaussianProcessesTurbulenceFitTask(config=config) 

198 

199 # Simulate a smooth Gaussian field with a correlation length well 

200 # inside the footprint. A von Karman field with a correlation 

201 # length of the footprint size (as in the Leget et al. 2021 test 

202 # case) is a degenerate regime for measuring a 2-pcf, and its 

203 # near-singular covariance makes the seeded multivariate_normal 

204 # draw platform-dependent (the SVD factorization of a singular 

205 # matrix differs between BLAS builds), so fit-quality statistics 

206 # were unstable across machines. 

207 kernel_skl = treegp.eval_kernel( 

208 "15**2 * AnisotropicRBF(invLam=array([[1.0/0.02**2,0],[0,1.0/0.02**2]]))" 

209 ) 

210 

211 self.camera = DummyCam().getCamera() 

212 self.wcsCatalog = self._makeWcsCatalog(self.visit, self.camera) 

213 self.positions = self._makeResidualCatalog(self.wcsCatalog, self.visit, kernel_skl) 

214 

215 ( 

216 self.gpx, 

217 self.gpy, 

218 self.trainInds, 

219 self.testInds, 

220 self.hyperparameters, 

221 self.allTPCoords, 

222 ) = self.task.runGP(self.wcsCatalog, self.positions) 

223 

224 def test_runGP(self): 

225 """Check that the empirical kernel is used and that the prediction 

226 catches most of the correlated signal.""" 

227 self.assertIsInstance(self.gpx.kernel, treegp.EmpiricalCorrelationKernel) 

228 self.assertEqual(len(self.gpx.kernel.theta), 0) 

229 

230 # Thresholds calibrated by rerunning this exact setup (treegp 

231 # only) over 20 seeds plus Cholesky-draw variants: the train 

232 # residual ratio stays below 0.07 and the validation ratio below 

233 # 0.19, while a broken fit gives about 1. 

234 dx = np.array(self.positions["xresw"]) 

235 predTrain = self.gpx.predict(self.allTPCoords[self.trainInds]) 

236 residTrain = dx[self.trainInds] - predTrain 

237 self.assertLess(np.std(residTrain), 0.3 * np.std(dx[self.trainInds])) 

238 

239 # The validation points are not used in the fit, but the prediction 

240 # should still remove most of the correlated signal. 

241 predTest = self.gpx.predict(self.allTPCoords[self.testInds]) 

242 residTest = dx[self.testInds] - predTest 

243 self.assertLess(np.std(residTest), 0.5 * np.std(dx[self.testInds])) 

244 

245 def test_hyperparameters(self): 

246 """Check the diagnostics stored in place of the hyperparameters.""" 

247 self.assertEqual(self.hyperparameters.colnames, ["x", "y"]) 

248 # [xi0, g1 measured, g2 measured, g1 applied, g2 applied, 

249 # pixel size, number of grid pixels] 

250 self.assertEqual(len(self.hyperparameters["x"]), 7) 

251 self.assertGreater(self.hyperparameters["x"][0], 0.0) 

252 self.assertEqual(self.hyperparameters["x"][5], self.task.config.correlationPixelSize) 

253 # npix = 2 * ceil(correlationSeparationMax / correlationPixelSize) 

254 self.assertEqual(self.hyperparameters["x"][6], 24) 

255 

256 def test_sourceTable(self): 

257 """Check the per-source output table.""" 

258 sourceTable = self.task.makeSourceTable( 

259 self.gpx, 

260 self.gpy, 

261 self.positions, 

262 self.allTPCoords, 

263 self.trainInds, 

264 self.wcsCatalog, 

265 self.camera, 

266 self.visit, 

267 ) 

268 

269 expectedColumns = [ 

270 "visit", 

271 "detector", 

272 "xTP", 

273 "yTP", 

274 "dxTP", 

275 "dyTP", 

276 "dxTPGP", 

277 "dyTPGP", 

278 "isTraining", 

279 "xPix", 

280 "yPix", 

281 "fpX", 

282 "fpY", 

283 "dxPix", 

284 "dyPix", 

285 "dxPixGP", 

286 "dyPixGP", 

287 ] 

288 self.assertEqual(sourceTable.colnames, expectedColumns) 

289 self.assertEqual(len(sourceTable), len(self.positions)) 

290 self.assertEqual(np.sum(sourceTable["isTraining"]), self.task.config.maxTrainingPoints) 

291 np.testing.assert_array_equal(sourceTable["visit"], self.visit) 

292 np.testing.assert_allclose(sourceTable["dxTP"], self.positions["xresw"]) 

293 self.assertTrue(np.all(np.isfinite(sourceTable["dxTPGP"]))) 

294 self.assertTrue(np.all(np.isfinite(sourceTable["fpX"]))) 

295 

296 # The pixel-frame offsets must have the same amplitude as the 

297 # tangent plane offsets divided by the local plate scale. 

298 for detector in self.wcsCatalog: 

299 detId = detector["id"] 

300 mask = sourceTable["detector"] == detId 

301 pixToTP = detector.wcs.getFrameDict().getMapping("PIXELS", "IWC") 

302 point = np.array([[sourceTable["xPix"][mask][0]], [sourceTable["yPix"][mask][0]]]) 

303 step = pixToTP.applyForward(point + np.array([[1.0], [0.0]])) - pixToTP.applyForward(point) 

304 plateScale = np.hypot(step[0, 0], step[1, 0]) 

305 

306 pixAmplitude = np.hypot(sourceTable["dxPix"][mask], sourceTable["dyPix"][mask]) * plateScale 

307 tpAmplitude = ( 

308 (np.hypot(sourceTable["dxTP"][mask], sourceTable["dyTP"][mask]) * u.mas).to(u.degree).value 

309 ) 

310 np.testing.assert_allclose(pixAmplitude, tpAmplitude, rtol=1e-2) 

311 

312 def test_divCurl(self): 

313 """Check the optional divergence/curl and E/B decomposition columns 

314 of the per-source table.""" 

315 # saveDivCurl requires saveSourceTable. 

316 badConfig = GaussianProcessesTurbulenceFitConfig() 

317 badConfig.saveDivCurl = True 

318 with self.assertRaises(pexConfig.FieldValidationError): 

319 badConfig.validate() 

320 

321 config = GaussianProcessesTurbulenceFitConfig() 

322 config.saveSourceTable = True 

323 config.saveDivCurl = True 

324 # Coarse grid: the DummyCam test catalog is sparse, and the 

325 # cloud-in-cell gridding needs a few sources per cell. 

326 config.divCurlGridResolution = 20 

327 task = GaussianProcessesTurbulenceFitTask(config=config) 

328 sourceTable = task.makeSourceTable( 

329 self.gpx, 

330 self.gpy, 

331 self.positions, 

332 self.allTPCoords, 

333 self.trainInds, 

334 self.wcsCatalog, 

335 self.camera, 

336 self.visit, 

337 ) 

338 

339 newColumns = [] 

340 for label in ["", "GP"]: 

341 newColumns += [f"divTP{label}", f"curlTP{label}"] 

342 for ebLabel in ["E", "B"]: 

343 newColumns += [f"dxTP{label}{ebLabel}", f"dyTP{label}{ebLabel}"] 

344 for label in ["", "GP"]: 

345 for ebLabel in ["E", "B"]: 

346 newColumns += [f"dxPix{label}{ebLabel}", f"dyPix{label}{ebLabel}"] 

347 for column in newColumns: 

348 self.assertIn(column, sourceTable.colnames) 

349 

350 # Most sources are inside the gridded region and get finite values; 

351 # sources near the field edge (outside the convex hull of the 

352 # gridded field) are NaN. 

353 finite = np.isfinite(sourceTable["divTP"]) 

354 self.assertGreater(np.mean(finite), 0.5) 

355 

356 # The Helmholtz decomposition reconstructs the (deconvolved) 

357 # gridded field: E + B must track the GP column coherently. This 

358 # is a sanity check, not a precision one: the catalog is sparse 

359 # (~1 source per grid cell) and the deconvolution sharpens the 

360 # map, so the pointwise scatter is intrinsically ~0.2-0.35 of 

361 # the field rms on this setup (calibrated on synthetic Gaussian 

362 # fields over many seeds; a broken decomposition gives ~1). 

363 # Tolerance set to 0.75 accordingly. 

364 for xy, ind in [("x", 0), ("y", 1)]: 

365 total = sourceTable[f"d{xy}TPGPE"] + sourceTable[f"d{xy}TPGPB"] 

366 good = np.isfinite(total) 

367 self.assertGreater(np.mean(good), 0.5) 

368 scatter = np.std(total[good] - sourceTable[f"d{xy}TPGP"][good]) 

369 self.assertLess(scatter, 0.75 * np.std(sourceTable[f"d{xy}TPGP"][good])) 

370 

371 # Pixel-frame E/B components are finite wherever the tangent plane 

372 # ones are, and have the same amplitude up to the plate scale. 

373 for column in ["dxPixE", "dyPixB", "dxPixGPE", "dyPixGPB"]: 

374 tpColumn = column.replace("Pix", "TP") 

375 good = np.isfinite(sourceTable[tpColumn]) 

376 self.assertTrue(np.all(np.isfinite(sourceTable[column][good]))) 

377 

378 

379if __name__ == "__main__": 379 ↛ 380line 379 didn't jump to line 380 because the condition on line 379 was never true

380 lsst.utils.tests.init() 

381 unittest.main()