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

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 NotPositiveDefiniteMatrixError(pipeBase.AlgorithmError): 

156 """Raised if the Gaussian Processes fit raises a not positive definite 

157 linear algebra error.""" 

158 

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

160 super().__init__( 

161 "The Gaussian Processes fit failed with a not-positive-definite linear algebra" " error." 

162 ) 

163 self._nSources = nSources 

164 

165 @property 

166 def metadata(self): 

167 return { 

168 "nSources": self._nSources, 

169 } 

170 

171 

172class GaussianProcessesTurbulenceFitConnections( 

173 pipeBase.PipelineTaskConnections, 

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

175 defaultTemplates={ 

176 "inputName": "gbdesHealpix3AstrometricFit", 

177 }, 

178): 

179 inputWcs = pipeBase.connectionTypes.Input( 

180 doc=( 

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

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

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

184 ), 

185 name="{inputName}SkyWcsCatalog", 

186 storageClass="ExposureCatalog", 

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

188 ) 

189 inputPositions = pipeBase.connectionTypes.Input( 

190 doc=( 

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

192 "plane coordinates and chisq values." 

193 ), 

194 name="{inputName}_fitStars", 

195 storageClass="ArrowAstropy", 

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

197 deferLoad=True, 

198 ) 

199 outputWcs = pipeBase.connectionTypes.Output( 

200 doc=( 

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

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

203 "for fast lookups of a detector." 

204 ), 

205 name="turbulenceCorrectedSkyWcsCatalog", 

206 storageClass="ExposureCatalog", 

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

208 ) 

209 hyperparameters = pipeBase.connectionTypes.Output( 

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

211 name="turbulence_fit_hyperparameters", 

212 storageClass="ArrowAstropy", 

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

214 ) 

215 

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

217 super().__init__(config=config) 

218 

219 if not self.config.healpix: 

220 self.dimensions.remove("healpix3") 

221 self.dimensions.add("tract") 

222 self.dimensions.add("skymap") 

223 self.inputWcs = dataclasses.replace( 

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

225 ) 

226 self.inputPositions = dataclasses.replace( 

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

228 ) 

229 

230 

231class GaussianProcessesTurbulenceFitConfig( 

232 pipeBase.PipelineTaskConfig, pipelineConnections=GaussianProcessesTurbulenceFitConnections 

233): 

234 initKernel = pexConfig.Field( 

235 dtype=str, 

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

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

238 ) 

239 initAnisotropicCorrelationLength = pexConfig.ListField( 

240 dtype=float, 

241 doc=( 

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

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

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

245 "where 0 means the correlation is isotropic." 

246 ), 

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

248 ) 

249 correlationSeparationMin = pexConfig.Field( 

250 dtype=float, 

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

252 default=0.0, 

253 optional=True, 

254 ) 

255 correlationSeparationMax = pexConfig.Field( 

256 dtype=float, 

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

258 default=0.3, 

259 optional=True, 

260 ) 

261 maxTrainingPoints = pexConfig.Field( 

262 dtype=int, 

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

264 default=10000, 

265 ) 

266 pixelSize = pexConfig.Field( 

267 dtype=float, 

268 doc="Pixel size in arcseconds.", 

269 default=0.2, 

270 ) 

271 healpix = pexConfig.Field( 

272 dtype=bool, 

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

274 default=True, 

275 ) 

276 splineDegree = pexConfig.Field( 

277 dtype=int, 

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

279 default=4, 

280 ) 

281 splineNNodes = pexConfig.Field( 

282 dtype=int, 

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

284 default=30, 

285 ) 

286 splineBuffer = pexConfig.Field( 

287 dtype=float, 

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

289 default=0.1, 

290 ) 

291 

292 

293class GaussianProcessesTurbulenceFitTask(pipeBase.PipelineTask): 

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

295 they are due to atmospheric turbulence.""" 

296 

297 ConfigClass = GaussianProcessesTurbulenceFitConfig 

298 _DefaultName = "gaussianProcessesTurbulenceFit" 

299 

300 def run(self, inputWcs, inputPositions): 

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

302 Gaussian Processes prediction from the WCS to account for atmospheric 

303 turbulence. 

304 

305 Parameters 

306 ---------- 

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

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

309 inputPositions : `astropy.table.Table` 

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

311 

312 Returns 

313 ------- 

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

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

316 Catalog with WCS after inserting the correction for atmospheric 

317 turbulence. 

318 ``hyperparameters`` : `astropy.table.Table` 

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

320 """ 

321 

322 visit = inputWcs[0]["visit"] 

323 

324 inputPositions = inputPositions.get( 

325 parameters={ 

326 "columns": [ 

327 "xworld", 

328 "yworld", 

329 "xresw", 

330 "yresw", 

331 "exposureName", 

332 "xpix", 

333 "ypix", 

334 "deviceName", 

335 "clip", 

336 "covTotalW_00", 

337 "covTotalW_11", 

338 ] 

339 } 

340 ) 

341 

342 visitPositions = inputPositions[ 

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

344 ] 

345 

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

347 

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

349 

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

351 

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

353 

354 def runGP(self, inputWcs, positions): 

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

356 

357 Parameters 

358 ---------- 

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

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

361 inputPositions : `astropy.table.Table` 

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

363 

364 Returns 

365 ------- 

366 gpx : `treegp.gp_interp.GPInterpolation` 

367 Gaussian Processes interpolator for x-direction residuals. 

368 gpy : `treegp.gp_interp.GPInterpolation` 

369 Gaussian Processes interpolator for y-direction residuals. 

370 trainInds : `numpy.ndarray` 

371 Array of indices for points used in training. 

372 testInds : `numpy.ndarray` 

373 Array of indices for points not used in training. 

374 """ 

375 dx = positions["xresw"] 

376 dy = positions["yresw"] 

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

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

379 

380 # Get tangent plane coordinates for input points 

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

382 for detector in inputWcs: 

383 detId = detector["id"] 

384 detWCS = detector.wcs 

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

386 detectorSources = positions[detInd] 

387 

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

389 tangentPlaneCoords = tangentPlaneToSky.applyForward( 

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

391 ) 

392 allTPCoords[detInd] = tangentPlaneCoords.T 

393 

394 # Choose a random subset for training. 

395 rng = np.random.default_rng(1234) 

396 nPoints = len(allTPCoords) 

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

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

399 trainInds = perm[:nTrain] 

400 testInds = perm[nTrain:] 

401 

402 # Solve Gaussian Processes in dx direction. 

403 gpx = treegp.GPInterpolation( 

404 kernel=self.config.initKernel, 

405 optimizer="anisotropic", 

406 normalize=True, 

407 nbins=21, 

408 min_sep=self.config.correlationSeparationMin, 

409 max_sep=self.config.correlationSeparationMax, 

410 p0=self.config.initAnisotropicCorrelationLength, 

411 ) 

412 

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

414 

415 # Solve Gaussian Processes in dy direction. 

416 gpy = treegp.GPInterpolation( 

417 kernel=self.config.initKernel, 

418 optimizer="anisotropic", 

419 normalize=True, 

420 nbins=21, 

421 min_sep=self.config.correlationSeparationMin, 

422 max_sep=self.config.correlationSeparationMax, 

423 p0=self.config.initAnisotropicCorrelationLength, 

424 ) 

425 

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

427 

428 try: 

429 gpx.solve() 

430 gpy.solve() 

431 except np.linalg.LinAlgError as e: 

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

433 error = pipeBase.AnnotatedPartialOutputsError.annotate( 

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

435 self, 

436 log=self.log, 

437 ) 

438 raise error from e 

439 elif "not positive definite" in str(e): 

440 error = pipeBase.AnnotatedPartialOutputsError.annotate( 

441 NotPositiveDefiniteMatrixError(len(allTPCoords[trainInds])), 

442 self, 

443 log=self.log, 

444 ) 

445 raise error from e 

446 else: 

447 raise 

448 

449 hyperparameters = Table( 

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

451 ) 

452 

453 return gpx, gpy, trainInds, testInds, hyperparameters 

454 

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

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

457 turbulence. 

458 

459 Parameters 

460 ---------- 

461 gpx : `treegp.gp_interp.GPInterpolation` 

462 Gaussian Processes interpolator for x-direction residuals. 

463 gpy : `treegp.gp_interp.GPInterpolation` 

464 Gaussian Processes interpolator for y-direction residuals. 

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

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

467 inputPositions : `astropy.table.Table` 

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

469 

470 Returns 

471 ------- 

472 outCat : `astropy.table.Table` 

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

474 columns corrected for atmospheric turbulence. 

475 """ 

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

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

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

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

480 

481 for detector in inputWcs: 

482 detId = detector["id"] 

483 detWCS = detector.wcs 

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

485 detectorSources = sourceCatalog[detInd] 

486 

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

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

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

490 # plane coordinates to sky coordinates. 

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

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

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

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

495 allTPCoords[detInd] = tangentPlaneCoords 

496 

497 xPred = gpx.predict(tangentPlaneCoords) 

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

499 yPred = gpy.predict(tangentPlaneCoords) 

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

501 prediction[detInd, 0] = xPred 

502 prediction[detInd, 1] = yPred 

503 

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

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

506 correctedSkyCoords = tangentPlaneToSky.applyForward( 

507 np.array([correctedTangentPlaneX, correctedTangentPlaneY]) 

508 ) 

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

510 

511 outCat = sourceCatalog.copy() 

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

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

514 

515 return outCat 

516 

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

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

519 the WCS for each detector. 

520 

521 Parameters 

522 ---------- 

523 gpx : `treegp.gp_interp.GPInterpolation` 

524 Gaussian Processes interpolator for x-direction residuals. 

525 gpy : `treegp.gp_interp.GPInterpolation` 

526 Gaussian Processes interpolator for y-direction residuals. 

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

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

529 

530 Returns 

531 ------- 

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

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

534 gaussian processes fit. 

535 """ 

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

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

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

539 

540 # Set up the schema for the output catalogs 

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

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

543 

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

545 catalog.resize(len(inputWcs)) 

546 catalog["visit"] = inputWcs["visit"] 

547 

548 for d, detectorRow in enumerate(inputWcs): 

549 detId = detectorRow.getId() 

550 catalog[d].setId(detId) 

551 

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

553 bbox = detectorRow.getBBox() 

554 catalog[d].setBBox(bbox) 

555 corners = np.array( 

556 [ 

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

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

559 ] 

560 ).astype(float) 

561 

562 initWcsRow = inputWcs.find(detId) 

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

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

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

566 tangentPlaneX, tangentPlaneY = pixToTPMap.applyForward(corners) 

567 

568 xs = np.linspace( 

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

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

571 self.config.splineNNodes, 

572 ) 

573 ys = np.linspace( 

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

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

576 self.config.splineNNodes, 

577 ) 

578 

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

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

581 

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

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

584 

585 splineX = RectBivariateSpline( 

586 xs, 

587 ys, 

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

589 s=0, 

590 kx=self.config.splineDegree - 1, 

591 ky=self.config.splineDegree - 1, 

592 ) 

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

594 coeffsX = splineX.get_coeffs() 

595 

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

597 splineY = RectBivariateSpline( 

598 xs, 

599 ys, 

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

601 s=0, 

602 kx=self.config.splineDegree - 1, 

603 ky=self.config.splineDegree - 1, 

604 ) 

605 coeffsY = splineY.get_coeffs() 

606 

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

608 splineMap = ast.SplineMap( 

609 self.config.splineDegree, 

610 self.config.splineDegree, 

611 self.config.splineNNodes, 

612 self.config.splineNNodes, 

613 tx, 

614 ty, 

615 coeffsX, 

616 coeffsY, 

617 options="OutUnit=1", 

618 ) 

619 

620 newFrameDict = ast.FrameDict(pixelFrame) 

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

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

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

624 outWcs = afwgeom.SkyWcs(newFrameDict) 

625 catalog[d].setWcs(outWcs) 

626 

627 return catalog 

628 

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

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

631 and after correcting for atmospheric turbulence, and validate 

632 prediction on some of the test data. 

633 

634 Parameters 

635 ---------- 

636 gpx : `treegp.gp_interp.GPInterpolation` 

637 Gaussian Processes interpolator for x-direction residuals. 

638 gpy : `treegp.gp_interp.GPInterpolation` 

639 Gaussian Processes interpolator for y-direction residuals. 

640 positions : `astropy.table.Table` 

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

642 trainInds : `numpy.ndarray` 

643 Array of indices for points used in training. 

644 testInds : `numpy.ndarray` 

645 Array of indices for points not used in training. 

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

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

648 makeValidationPlot : `bool`, optional 

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

650 data. 

651 """ 

652 dx = positions["xresw"] 

653 dy = positions["yresw"] 

654 

655 # Get tangent plane coordinates for input points 

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

657 for detector in inputWcs: 

658 detId = detector["id"] 

659 detWCS = detector.wcs 

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

661 detectorSources = positions[detInd] 

662 

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

664 tangentPlaneCoords = tangentPlaneToSky.applyForward( 

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

666 ) 

667 tpCoords[detInd] = tangentPlaneCoords.T 

668 

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

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

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

672 xie, xib, logr = treegp.comp_eb_treecorr( 

673 tpCoords[trainInd, 0], 

674 tpCoords[trainInd, 1], 

675 dx[trainInd], 

676 dy[trainInd], 

677 rmin=20 / 3600, 

678 rmax=0.6, 

679 dlogr=0.3, 

680 ) 

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

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

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

684 self.log.info( 

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

686 meanE, 

687 meanB, 

688 ) 

689 

690 xie_resid, xib_resid, logr = treegp.comp_eb_treecorr( 

691 tpCoords[trainInd, 0], 

692 tpCoords[trainInd, 1], 

693 dx[trainInd] - xPredict, 

694 dy[trainInd] - yPredict, 

695 rmin=20 / 3600, 

696 rmax=0.6, 

697 dlogr=0.3, 

698 ) 

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

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

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

702 self.log.info( 

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

704 meanE_resid, 

705 meanB_resid, 

706 ) 

707 

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

709 if makeValidationPlot: 

710 print(len(testInd)) 

711 testInd = testInd[:50000] 

712 chunkSize = 5000 

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

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

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

716 for i in range(nChunks): 

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

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

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

720 fig = plot_visit( 

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

722 ) 

723 return fig