Coverage for python / lsst / drp / tasks / fit_turbulence.py: 17%

235 statements  

« prev     ^ index     » next       coverage.py v7.12.0, created at 2025-12-19 01:54 -0800

1# This file is part of drp_tasks. 

2# 

3# LSST Data Management System 

4# This product includes software developed by the 

5# LSST Project (http://www.lsst.org/). 

6# See COPYRIGHT file at the top of the source tree. 

7# 

8# This program is free software: you can redistribute it and/or modify 

9# it under the terms of the GNU General Public License as published by 

10# the Free Software Foundation, either version 3 of the License, or 

11# (at your option) any later version. 

12# 

13# This program is distributed in the hope that it will be useful, 

14# but WITHOUT ANY WARRANTY; without even the implied warranty of 

15# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

16# GNU General Public License for more details. 

17# 

18# You should have received a copy of the LSST License Statement and 

19# the GNU General Public License along with this program. If not, 

20# see <https://www.lsstcorp.org/LegalNotices/>. 

21# 

22import dataclasses 

23 

24import astropy.units as u 

25import astshim as ast 

26import matplotlib.pyplot as plt 

27import numpy as np 

28import treegp 

29from astropy.table import Table 

30from scipy.interpolate import RectBivariateSpline 

31 

32import lsst.afw.geom as afwgeom 

33import lsst.afw.table 

34import lsst.pex.config as pexConfig 

35import lsst.pipe.base as pipeBase 

36 

37__all__ = [ 

38 "GaussianProcessesTurbulenceFitConnections", 

39 "GaussianProcessesTurbulenceFitConfig", 

40 "GaussianProcessesTurbulenceFitTask", 

41] 

42 

43 

44def plot_visit(x, y, dx, dy, predx, predy): 

45 """Utility function for plotting Gaussian Processes results. 

46 

47 Parameters 

48 ---------- 

49 x : `np.ndarray` 

50 x-direction coordinates. 

51 y : `np.ndarray` 

52 y-direction coordinates. 

53 dx : `np.ndarray` 

54 x-direction residuals to be fit. 

55 dy : `np.ndarray` 

56 x-direction residuals to be fit. 

57 predx : `np.ndarray` 

58 x-direction prediction. 

59 predy : `np.ndarray` 

60 y-direction prediction. 

61 

62 Returns 

63 ------- 

64 fig : `matplotlib.pyplot.Figure` 

65 Figure showing input data, Gaussian Processes prediction, and E and 

66 B-modes. 

67 """ 

68 

69 xie, xib, logr = treegp.comp_eb_treecorr(x, y, dx, dy, rmin=20 / 3600, rmax=0.6, dlogr=0.3) 

70 xie_resid, xib_resid, logr_resid = treegp.comp_eb_treecorr( 

71 x, y, dx - predx, dy - predy, rmin=20 / 3600, rmax=0.6, dlogr=0.3 

72 ) 

73 

74 residualLimit = np.nanstd(dx) 

75 

76 fig, subs = plt.subplot_mosaic( 

77 [["dx", "predx", "residx", "eb"], ["dy", "predy", "residy", "eb"]], 

78 figsize=(15, 8), 

79 layout="constrained", 

80 ) 

81 plt.subplots_adjust(wspace=0.3, right=0.99, left=0.05) 

82 im = subs["dx"].scatter(x, y, c=dx, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1) 

83 subs["dy"].scatter(x, y, c=dy, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1) 

84 

85 subs["predx"].scatter(x, y, c=predx, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1) 

86 subs["predy"].scatter(x, y, c=predy, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1) 

87 

88 subs["residx"].scatter( 

89 x, y, c=dx - predx, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1 

90 ) 

91 subs["residy"].scatter( 

92 x, y, c=dy - predy, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1 

93 ) 

94 

95 cb = fig.colorbar( 

96 im, ax=[subs["dx"], subs["dy"], subs["predx"], subs["predy"], subs["residx"], subs["residy"]] 

97 ) 

98 

99 subs["eb"].scatter(np.exp(logr) * 60, xie, c="b", label="E-mode") 

100 subs["eb"].scatter(np.exp(logr) * 60, xib, c="r", label="B-mode") 

101 

102 subs["eb"].scatter( 

103 np.exp(logr_resid) * 60, xie_resid, c="b", marker="+", label="E-mode after GP correction" 

104 ) 

105 subs["eb"].scatter( 

106 np.exp(logr_resid) * 60, xib_resid, c="r", marker="+", label="B-mode after GP correction" 

107 ) 

108 subs["eb"].legend() 

109 subs["eb"].grid(True) 

110 

111 subs["dx"].set_aspect("equal") 

112 subs["dy"].set_aspect("equal") 

113 subs["predx"].set_aspect("equal") 

114 subs["predy"].set_aspect("equal") 

115 subs["residx"].set_aspect("equal") 

116 subs["residy"].set_aspect("equal") 

117 subs["dy"].set_xlabel("x (degree)") 

118 subs["predy"].set_xlabel("x (degree)") 

119 subs["residy"].set_xlabel("x (degree)") 

120 subs["dy"].set_ylabel("y (degree)") 

121 subs["dx"].set_ylabel("y (degree)") 

122 

123 subs["dx"].set_title(r"$\delta$x") 

124 subs["predx"].set_title("GP prediction") 

125 subs["residx"].set_title("Residual") 

126 

127 subs["dy"].set_title(r"$\delta$y") 

128 subs["predy"].set_title("GP prediction") 

129 subs["residy"].set_title("Residual") 

130 

131 cb.set_label("mas") 

132 

133 subs["eb"].set_title("E and B modes") 

134 subs["eb"].set_ylabel(r"$\xi_{E/B}$ (mas$^2$)") 

135 subs["eb"].set_xlabel(r"$\Delta \theta$ (arcmin)") 

136 

137 return fig 

138 

139 

140class SingularMatrixError(pipeBase.AlgorithmError): 

141 """Raised if the Gaussian Processes fit raises a Singular Matrix linear 

142 algebra error.""" 

143 

144 def __init__(self, nSources) -> None: 

145 super().__init__("The Gaussian Processes fit failed with a singular matrix linear algebra error.") 

146 self._nSources = nSources 

147 

148 @property 

149 def metadata(self): 

150 return { 

151 "nSources": self._nSources, 

152 } 

153 

154 

155class GaussianProcessesTurbulenceFitConnections( 

156 pipeBase.PipelineTaskConnections, 

157 dimensions=("instrument", "visit", "healpix3"), 

158 defaultTemplates={ 

159 "inputName": "gbdesHealpix3AstrometricFit", 

160 }, 

161): 

162 inputWcs = pipeBase.connectionTypes.Input( 

163 doc=( 

164 "Per-healpix, per-visit world coordinate systems derived from the fitted model." 

165 " These catalogs only contain entries for detectors with an output, and use" 

166 " the detector id for the catalog id, sorted on id for fast lookups of a detector." 

167 ), 

168 name="{inputName}SkyWcsCatalog", 

169 storageClass="ExposureCatalog", 

170 dimensions=("instrument", "visit", "healpix3"), 

171 ) 

172 inputPositions = pipeBase.connectionTypes.Input( 

173 doc=( 

174 "Catalog of sources used in fit, along with residuals in pixel coordinates and tangent " 

175 "plane coordinates and chisq values." 

176 ), 

177 name="{inputName}_fitStars", 

178 storageClass="ArrowAstropy", 

179 dimensions=("instrument", "healpix3", "physical_filter"), 

180 deferLoad=True, 

181 ) 

182 outputWcs = pipeBase.connectionTypes.Output( 

183 doc=( 

184 "Per-visit world coordinate systems derived from the fitted model. These catalogs only contain " 

185 "entries for detectors with an output, and use the detector id for the catalog id, sorted on id " 

186 "for fast lookups of a detector." 

187 ), 

188 name="turbulenceCorrectedSkyWcsCatalog", 

189 storageClass="ExposureCatalog", 

190 dimensions=("instrument", "visit", "healpix3"), 

191 ) 

192 hyperparameters = pipeBase.connectionTypes.Output( 

193 doc="Best fit hyperparameters for the Gaussian Processes fit.", 

194 name="turbulence_fit_hyperparameters", 

195 storageClass="ArrowAstropy", 

196 dimensions=("instrument", "visit", "healpix3"), 

197 ) 

198 

199 def __init__(self, *, config=None): 

200 super().__init__(config=config) 

201 

202 if not self.config.healpix: 

203 self.dimensions.remove("healpix3") 

204 self.dimensions.add("tract") 

205 self.dimensions.add("skymap") 

206 self.inputWcs = dataclasses.replace( 

207 self.inputWcs, dimensions=("instrument", "visit", "tract", "skymap") 

208 ) 

209 self.inputPositions = dataclasses.replace( 

210 self.inputPositions, dimensions=("instrument", "tract", "skymap", "band", "physical_filter") 

211 ) 

212 

213 

214class GaussianProcessesTurbulenceFitConfig( 

215 pipeBase.PipelineTaskConfig, pipelineConnections=GaussianProcessesTurbulenceFitConnections 

216): 

217 initKernel = pexConfig.Field( 

218 dtype=str, 

219 doc="The type of function that will be used to modeled spatial correlation.", 

220 default="15**2 * AnisotropicVonKarman(invLam=array([[1./0.8**2,0],[0,1./0.8**2]]))", 

221 ) 

222 initAnisotropicCorrelationLength = pexConfig.ListField( 

223 dtype=float, 

224 doc=( 

225 "The initial parameters for fiting the anisotropic correlation length. p0[0] is equivalent of " 

226 "the isotropic correlation length in degrees, and p0[1]/p0[2] are ellipticity parameters and are " 

227 "mathematically equivalent to e1/e2 in weak-lensing. p0[1]/p0[2] must be in the range [-1,1], " 

228 "where 0 means the correlation is isotropic." 

229 ), 

230 default=[1, -0.2, -0.2], 

231 ) 

232 correlationSeparationMin = pexConfig.Field( 

233 dtype=float, 

234 doc="Minimum distance separation in degrees in the computation of the 2-point correlation function.", 

235 default=0.0, 

236 optional=True, 

237 ) 

238 correlationSeparationMax = pexConfig.Field( 

239 dtype=float, 

240 doc="Maximum distance separation in degrees in the computation of the 2-point correlation function.", 

241 default=0.3, 

242 optional=True, 

243 ) 

244 maxTrainingPoints = pexConfig.Field( 

245 dtype=int, 

246 doc="Maximum number of points to use in the Gaussian Processes training.", 

247 default=10000, 

248 ) 

249 pixelSize = pexConfig.Field( 

250 dtype=float, 

251 doc="Pixel size in arcseconds.", 

252 default=0.2, 

253 ) 

254 healpix = pexConfig.Field( 

255 dtype=bool, 

256 doc="Use input WCS calculated over healpix-based region. If false, use tract-based WCS.", 

257 default=True, 

258 ) 

259 splineDegree = pexConfig.Field( 

260 dtype=int, 

261 doc="Degree of the spline expressing Gaussian Processes prediction.", 

262 default=4, 

263 ) 

264 splineNNodes = pexConfig.Field( 

265 dtype=int, 

266 doc="Number of nodes to use for the spline expressing Gaussian Processes prediction.", 

267 default=30, 

268 ) 

269 splineBuffer = pexConfig.Field( 

270 dtype=float, 

271 doc="Minimum distance in degrees to extend spline map outside the detector boundary.", 

272 default=0.1, 

273 ) 

274 

275 

276class GaussianProcessesTurbulenceFitTask(pipeBase.PipelineTask): 

277 """Run Gaussian Processes on astrometric residuals with the assumption that 

278 they are due to atmospheric turbulence.""" 

279 

280 ConfigClass = GaussianProcessesTurbulenceFitConfig 

281 _DefaultName = "gaussianProcessesTurbulenceFit" 

282 

283 def run(self, inputWcs, inputPositions): 

284 """Run Gaussian Processes on position residuals and subtract the fitted 

285 Gaussian Processes prediction from the WCS to account for atmospheric 

286 turbulence. 

287 

288 Parameters 

289 ---------- 

290 inputWcs : `lsst.afw.table.ExposureCatalog` 

291 Catalog with WCSs for each detector of the input exposure. 

292 inputPositions : `astropy.table.Table` 

293 Catalog of input positions with residuals to the current best fit. 

294 

295 Returns 

296 ------- 

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

298 ``outputWcs`` : `lsst.afw.table.ExposureCatalog` 

299 Catalog with WCS after inserting the correction for atmospheric 

300 turbulence. 

301 ``hyperparameters`` : `astropy.table.Table` 

302 Table of best-fit hyperparameters in x and y-directions. 

303 """ 

304 

305 visit = inputWcs[0]["visit"] 

306 

307 inputPositions = inputPositions.get( 

308 parameters={ 

309 "columns": [ 

310 "xworld", 

311 "yworld", 

312 "xresw", 

313 "yresw", 

314 "exposureName", 

315 "xpix", 

316 "ypix", 

317 "deviceName", 

318 "clip", 

319 "covTotalW_00", 

320 "covTotalW_11", 

321 ] 

322 } 

323 ) 

324 

325 visitPositions = inputPositions[ 

326 (inputPositions["exposureName"] == str(visit)) & ~inputPositions["clip"] 

327 ] 

328 

329 gpx, gpy, trainInd, testInd, hyperparameters = self.runGP(inputWcs, visitPositions) 

330 

331 self.evaluate(gpx, gpy, visitPositions, trainInd, testInd, inputWcs) 

332 

333 wcsWithSpline = self.addGPToWcs(gpx, gpy, inputWcs) 

334 

335 return pipeBase.Struct(outputWcs=wcsWithSpline, hyperparameters=hyperparameters) 

336 

337 def runGP(self, inputWcs, positions): 

338 """Run Gaussian Processes in tangent plane coordinates. 

339 

340 Parameters 

341 ---------- 

342 inputWcs : `lsst.afw.table.ExposureCatalog` 

343 Catalog with WCSs for each detector of the input exposure. 

344 inputPositions : `astropy.table.Table` 

345 Catalog of input positions with residuals to the current best fit. 

346 

347 Returns 

348 ------- 

349 gpx : `treegp.gp_interp.GPInterpolation` 

350 Gaussian Processes interpolator for x-direction residuals. 

351 gpy : `treegp.gp_interp.GPInterpolation` 

352 Gaussian Processes interpolator for y-direction residuals. 

353 trainInds : `numpy.ndarray` 

354 Array of indices for points used in training. 

355 testInds : `numpy.ndarray` 

356 Array of indices for points not used in training. 

357 """ 

358 dx = positions["xresw"] 

359 dy = positions["yresw"] 

360 dxErr = positions["covTotalW_00"] ** 0.5 

361 dyErr = positions["covTotalW_11"] ** 0.5 

362 

363 # Get tangent plane coordinates for input points 

364 allTPCoords = np.zeros((len(positions), 2)) 

365 for detector in inputWcs: 

366 detId = detector["id"] 

367 detWCS = detector.wcs 

368 detInd = positions["deviceName"].astype(int) == detId 

369 detectorSources = positions[detInd] 

370 

371 tangentPlaneToSky = detWCS.getFrameDict().getMapping("PIXELS", "IWC") 

372 tangentPlaneCoords = tangentPlaneToSky.applyForward( 

373 np.array([detectorSources["xpix"], detectorSources["ypix"]]) 

374 ) 

375 allTPCoords[detInd] = tangentPlaneCoords.T 

376 

377 # Choose a random subset for training. 

378 rng = np.random.default_rng(1234) 

379 nPoints = len(allTPCoords) 

380 nTrain = min([nPoints, self.config.maxTrainingPoints]) 

381 perm = rng.permutation(np.arange(nPoints)) 

382 trainInds = perm[:nTrain] 

383 testInds = perm[nTrain:] 

384 

385 # Solve Gaussian Processes in dx direction. 

386 gpx = treegp.GPInterpolation( 

387 kernel=self.config.initKernel, 

388 optimizer="anisotropic", 

389 normalize=True, 

390 nbins=21, 

391 min_sep=self.config.correlationSeparationMin, 

392 max_sep=self.config.correlationSeparationMax, 

393 p0=self.config.initAnisotropicCorrelationLength, 

394 ) 

395 

396 gpx.initialize(allTPCoords[trainInds], dx[trainInds], y_err=dxErr[trainInds]) 

397 

398 # Solve Gaussian Processes in dy direction. 

399 gpy = treegp.GPInterpolation( 

400 kernel=self.config.initKernel, 

401 optimizer="anisotropic", 

402 normalize=True, 

403 nbins=21, 

404 min_sep=self.config.correlationSeparationMin, 

405 max_sep=self.config.correlationSeparationMax, 

406 p0=self.config.initAnisotropicCorrelationLength, 

407 ) 

408 

409 gpy.initialize(allTPCoords[trainInds], dy[trainInds], y_err=dyErr[trainInds]) 

410 

411 try: 

412 gpx.solve() 

413 gpy.solve() 

414 except np.linalg.LinAlgError as e: 

415 if "Singular matrix" in str(e): 

416 error = pipeBase.AnnotatedPartialOutputsError.annotate( 

417 SingularMatrixError(len(allTPCoords[trainInds])), 

418 self, 

419 log=self.log, 

420 ) 

421 raise error from e 

422 else: 

423 raise 

424 

425 hyperparameters = Table( 

426 {"x": np.array(gpx._optimizer._results_robust), "y": np.array(gpy._optimizer._results_robust)} 

427 ) 

428 

429 return gpx, gpy, trainInds, testInds, hyperparameters 

430 

431 def predict(self, gpx, gpy, inputWcs, sourceCatalog): 

432 """Get the positions for sources after correction for atmospheric 

433 turbulence. 

434 

435 Parameters 

436 ---------- 

437 gpx : `treegp.gp_interp.GPInterpolation` 

438 Gaussian Processes interpolator for x-direction residuals. 

439 gpy : `treegp.gp_interp.GPInterpolation` 

440 Gaussian Processes interpolator for y-direction residuals. 

441 inputWcs : `lsst.afw.table.ExposureCatalog` 

442 Catalog with WCSs for each detector of the input exposure. 

443 inputPositions : `astropy.table.Table` 

444 Catalog of input positions with residuals to the current best fit. 

445 

446 Returns 

447 ------- 

448 outCat : `astropy.table.Table` 

449 Catalog matching `inputPositions`, with `coord_ra` and `coord_dec` 

450 columns corrected for atmospheric turbulence. 

451 """ 

452 correctedCoordinates = np.zeros((len(sourceCatalog), 2)) 

453 prediction = np.zeros((len(sourceCatalog), 2)) 

454 allTPCoords = np.zeros((len(sourceCatalog), 2)) 

455 allCoords = np.zeros((len(sourceCatalog), 2)) 

456 

457 for detector in inputWcs: 

458 detId = detector["id"] 

459 detWCS = detector.wcs 

460 detInd = sourceCatalog["detector"] == detId 

461 detectorSources = sourceCatalog[detInd] 

462 

463 # The Gaussian Processes is fit on the tangent plane coordinates, 

464 # so we must transform points to the tangent plane, then subtract 

465 # the effect of atmospheric turbulence, then transform the tangent 

466 # plane coordinates to sky coordinates. 

467 initialSky = detWCS.pixelToSkyArray(detectorSources["x"], detectorSources["y"]) 

468 allCoords[detInd] = np.array(initialSky).T 

469 tangentPlaneToSky = detWCS.getFrameDict().getMapping("IWC", "SKY") 

470 tangentPlaneCoords = tangentPlaneToSky.applyInverse(np.array(initialSky)).T 

471 allTPCoords[detInd] = tangentPlaneCoords 

472 

473 xPred = gpx.predict(tangentPlaneCoords) 

474 xPrediction = (xPred * u.mas).to(u.degree) 

475 yPred = gpy.predict(tangentPlaneCoords) 

476 yPrediction = (yPred * u.mas).to(u.degree) 

477 prediction[detInd, 0] = xPred 

478 prediction[detInd, 1] = yPred 

479 

480 correctedTangentPlaneX = tangentPlaneCoords[:, 0] * u.degree - xPrediction 

481 correctedTangentPlaneY = tangentPlaneCoords[:, 1] * u.degree - yPrediction 

482 correctedSkyCoords = tangentPlaneToSky.applyForward( 

483 np.array([correctedTangentPlaneX, correctedTangentPlaneY]) 

484 ) 

485 correctedCoordinates[detInd] = ((correctedSkyCoords.T) * u.radian).to(u.degree).value 

486 

487 outCat = sourceCatalog.copy() 

488 outCat["coord_ra"] = correctedCoordinates[:, 0] 

489 outCat["coord_dec"] = correctedCoordinates[:, 1] 

490 

491 return outCat 

492 

493 def addGPToWcs(self, gpx, gpy, inputWcs): 

494 """Convert Gaussian Processes prediction to a spline, and insert it in 

495 the WCS for each detector. 

496 

497 Parameters 

498 ---------- 

499 gpx : `treegp.gp_interp.GPInterpolation` 

500 Gaussian Processes interpolator for x-direction residuals. 

501 gpy : `treegp.gp_interp.GPInterpolation` 

502 Gaussian Processes interpolator for y-direction residuals. 

503 inputWcs : `lsst.afw.table.ExposureCatalog` 

504 Catalog with WCSs for each detector of the input exposure. 

505 

506 Returns 

507 ------- 

508 catalog : `lsst.afw.table.ExposureCatalog` 

509 Exposure catalog with the WCS set to the existing WCS plus the 

510 gaussian processes fit. 

511 """ 

512 pixelFrame = ast.Frame(2, "Domain=PIXELS") 

513 tpFrame = ast.Frame(2, "Domain=TP") 

514 iwcFrame = ast.Frame(2, "Domain=IWC") 

515 

516 # Set up the schema for the output catalogs 

517 schema = lsst.afw.table.ExposureTable.makeMinimalSchema() 

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

519 

520 catalog = lsst.afw.table.ExposureCatalog(schema) 

521 catalog.resize(len(inputWcs)) 

522 catalog["visit"] = inputWcs["visit"] 

523 

524 for d, detectorRow in enumerate(inputWcs): 

525 detId = detectorRow.getId() 

526 catalog[d].setId(detId) 

527 

528 # Make a grid of points in tangent plane coordinates. 

529 bbox = detectorRow.getBBox() 

530 catalog[d].setBBox(bbox) 

531 corners = np.array( 

532 [ 

533 [bbox.getBeginX(), bbox.getEndX(), bbox.getEndX(), bbox.getBeginX()], 

534 [bbox.getBeginY(), bbox.getBeginY(), bbox.getEndY(), bbox.getEndY()], 

535 ] 

536 ).astype(float) 

537 

538 initWcsRow = inputWcs.find(detId) 

539 pixToTPMap = initWcsRow.wcs.getFrameDict().getMapping("PIXELS", "IWC") 

540 tpToSky = initWcsRow.wcs.getFrameDict().getMapping("IWC", "SKY") 

541 skyFrame = initWcsRow.wcs.getFrameDict().getFrame("SKY") 

542 tangentPlaneX, tangentPlaneY = pixToTPMap.applyForward(corners) 

543 

544 xs = np.linspace( 

545 tangentPlaneX.min() - self.config.splineBuffer, 

546 tangentPlaneX.max() + self.config.splineBuffer, 

547 self.config.splineNNodes, 

548 ) 

549 ys = np.linspace( 

550 tangentPlaneY.min() - self.config.splineBuffer, 

551 tangentPlaneY.max() + self.config.splineBuffer, 

552 self.config.splineNNodes, 

553 ) 

554 

555 xx, yy = np.meshgrid(xs, ys) 

556 inArray = np.array([xx.ravel(), yy.ravel()]).T 

557 

558 # Get Gaussian Processes prediction on grid and fit spline to it. 

559 xPred = (gpx.predict(inArray) * u.mas).to(u.degree).value 

560 

561 splineX = RectBivariateSpline( 

562 xs, 

563 ys, 

564 (xx - xPred.reshape(self.config.splineNNodes, self.config.splineNNodes)).T, 

565 s=0, 

566 kx=self.config.splineDegree - 1, 

567 ky=self.config.splineDegree - 1, 

568 ) 

569 (tx, ty) = splineX.get_knots() 

570 coeffsX = splineX.get_coeffs() 

571 

572 yPred = (gpy.predict(inArray) * u.mas).to(u.degree).value 

573 splineY = RectBivariateSpline( 

574 xs, 

575 ys, 

576 (yy - yPred.reshape(self.config.splineNNodes, self.config.splineNNodes)).T, 

577 s=0, 

578 kx=self.config.splineDegree - 1, 

579 ky=self.config.splineDegree - 1, 

580 ) 

581 coeffsY = splineY.get_coeffs() 

582 

583 # Turn spline into AST object and insert in new WCS. 

584 splineMap = ast.SplineMap( 

585 self.config.splineDegree, 

586 self.config.splineDegree, 

587 self.config.splineNNodes, 

588 self.config.splineNNodes, 

589 tx, 

590 ty, 

591 coeffsX, 

592 coeffsY, 

593 options="OutUnit=1", 

594 ) 

595 

596 newFrameDict = ast.FrameDict(pixelFrame) 

597 newFrameDict.addFrame("PIXELS", pixToTPMap, tpFrame) 

598 newFrameDict.addFrame("TP", splineMap, iwcFrame) 

599 newFrameDict.addFrame("IWC", tpToSky, skyFrame) 

600 outWcs = afwgeom.SkyWcs(newFrameDict) 

601 catalog[d].setWcs(outWcs) 

602 

603 return catalog 

604 

605 def evaluate(self, gpx, gpy, positions, trainInd, testInd, inputWcs, makeValidationPlot=False): 

606 """Calculate E and B-modes in the 2-point correlation function before 

607 and after correcting for atmospheric turbulence, and validate 

608 prediction on some of the test data. 

609 

610 Parameters 

611 ---------- 

612 gpx : `treegp.gp_interp.GPInterpolation` 

613 Gaussian Processes interpolator for x-direction residuals. 

614 gpy : `treegp.gp_interp.GPInterpolation` 

615 Gaussian Processes interpolator for y-direction residuals. 

616 positions : `astropy.table.Table` 

617 Catalog of input positions with residuals to the best fit. 

618 trainInds : `numpy.ndarray` 

619 Array of indices for points used in training. 

620 testInds : `numpy.ndarray` 

621 Array of indices for points not used in training. 

622 inputWcs : `lsst.afw.table.ExposureCatalog` 

623 Catalog with WCSs for each detector of the input exposure. 

624 makeValidationPlot : `bool`, optional 

625 Whether to make a plot showing the prediction on the validation 

626 data. 

627 """ 

628 dx = positions["xresw"] 

629 dy = positions["yresw"] 

630 

631 # Get tangent plane coordinates for input points 

632 tpCoords = np.zeros((len(positions), 2)) 

633 for detector in inputWcs: 

634 detId = detector["id"] 

635 detWCS = detector.wcs 

636 detInd = positions["deviceName"].astype(int) == detId 

637 detectorSources = positions[detInd] 

638 

639 tangentPlaneToSky = detWCS.getFrameDict().getMapping("PIXELS", "IWC") 

640 tangentPlaneCoords = tangentPlaneToSky.applyForward( 

641 np.array([detectorSources["xpix"], detectorSources["ypix"]]) 

642 ) 

643 tpCoords[detInd] = tangentPlaneCoords.T 

644 

645 # Calculate E/B modes before and after Gaussian Processes correction. 

646 xPredict = gpx.predict(tpCoords[trainInd]) 

647 yPredict = gpy.predict(tpCoords[trainInd]) 

648 xie, xib, logr = treegp.comp_eb_treecorr( 

649 tpCoords[trainInd, 0], 

650 tpCoords[trainInd, 1], 

651 dx[trainInd], 

652 dy[trainInd], 

653 rmin=20 / 3600, 

654 rmax=0.6, 

655 dlogr=0.3, 

656 ) 

657 start, stop = np.searchsorted(np.exp(logr), [0, 15]) 

658 meanE = np.mean(xie[start:stop]) 

659 meanB = np.mean(xib[start:stop]) 

660 self.log.info( 

661 "Original average correlation level over 0-15 arcminutes: E-mode=%0.2f, B-mode=%0.2f", 

662 meanE, 

663 meanB, 

664 ) 

665 

666 xie_resid, xib_resid, logr = treegp.comp_eb_treecorr( 

667 tpCoords[trainInd, 0], 

668 tpCoords[trainInd, 1], 

669 dx[trainInd] - xPredict, 

670 dy[trainInd] - yPredict, 

671 rmin=20 / 3600, 

672 rmax=0.6, 

673 dlogr=0.3, 

674 ) 

675 start, stop = np.searchsorted(np.exp(logr), [0, 15]) 

676 meanE_resid = np.mean(xie_resid[start:stop]) 

677 meanB_resid = np.mean(xib_resid[start:stop]) 

678 self.log.info( 

679 "Correlation level after GP correction over 0-15 arcminutes: E-mode=%0.2f, B-mode=%0.2f", 

680 meanE_resid, 

681 meanB_resid, 

682 ) 

683 

684 # Predict on all test data and make a plot. 

685 if makeValidationPlot: 

686 print(len(testInd)) 

687 testInd = testInd[:50000] 

688 chunkSize = 5000 

689 nChunks = np.ceil(len(testInd) / chunkSize).astype(int) 

690 xPredict = np.zeros(len(testInd)) 

691 yPredict = np.zeros(len(testInd)) 

692 for i in range(nChunks): 

693 ind = testInd[chunkSize * i : chunkSize * (i + 1)] 

694 xPredict[chunkSize * i : chunkSize * (i + 1)] = gpx.predict(tpCoords[ind]) 

695 yPredict[chunkSize * i : chunkSize * (i + 1)] = gpy.predict(tpCoords[ind]) 

696 fig = plot_visit( 

697 tpCoords[testInd, 0], tpCoords[testInd, 1], dx[testInd], dy[testInd], xPredict, yPredict 

698 ) 

699 return fig