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

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 treecorr 

29import treegp 

30from astropy.table import Table 

31from scipy.fft import fft2, ifft2, next_fast_len 

32from scipy.interpolate import RectBivariateSpline, RegularGridInterpolator 

33from scipy.ndimage import gaussian_filter 

34 

35import lsst.afw.geom as afwgeom 

36import lsst.afw.table 

37import lsst.pex.config as pexConfig 

38import lsst.pipe.base as pipeBase 

39from lsst.afw.cameraGeom import FOCAL_PLANE, PIXELS 

40 

41# We need to explicitly turn off multiprocessing in treecorr which is used 

42# by treegp. 

43treecorr.set_max_omp_threads(1) 

44 

45__all__ = [ 

46 "GaussianProcessesTurbulenceFitConnections", 

47 "GaussianProcessesTurbulenceFitConfig", 

48 "GaussianProcessesTurbulenceFitTask", 

49] 

50 

51 

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

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

54 

55 Parameters 

56 ---------- 

57 x : `np.ndarray` 

58 x-direction coordinates. 

59 y : `np.ndarray` 

60 y-direction coordinates. 

61 dx : `np.ndarray` 

62 x-direction residuals to be fit. 

63 dy : `np.ndarray` 

64 x-direction residuals to be fit. 

65 predx : `np.ndarray` 

66 x-direction prediction. 

67 predy : `np.ndarray` 

68 y-direction prediction. 

69 

70 Returns 

71 ------- 

72 fig : `matplotlib.pyplot.Figure` 

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

74 B-modes. 

75 """ 

76 

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

78 xie_resid, xib_resid, logr_resid = treegp.comp_eb_treecorr( 

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

80 ) 

81 

82 residualLimit = np.nanstd(dx) 

83 

84 fig, subs = plt.subplot_mosaic( 

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

86 figsize=(15, 8), 

87 layout="constrained", 

88 ) 

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

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

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

92 

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

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

95 

96 subs["residx"].scatter( 

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

98 ) 

99 subs["residy"].scatter( 

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

101 ) 

102 

103 cb = fig.colorbar( 

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

105 ) 

106 

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

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

109 

110 subs["eb"].scatter( 

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

112 ) 

113 subs["eb"].scatter( 

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

115 ) 

116 subs["eb"].legend() 

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

118 

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

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

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

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

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

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

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

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

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

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

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

130 

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

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

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

134 

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

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

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

138 

139 cb.set_label("mas") 

140 

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

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

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

144 

145 return fig 

146 

147 

148def compute_div_curl_eb(coords, field, gridResolution, smoothing=None, weights=None, deconvolve=True): 

149 """Compute the divergence, curl, and E/B Helmholtz decomposition of a 

150 scattered 2D vector field, evaluated back at the input positions. 

151 

152 The field is gridded with a noise-aware estimator: a (optionally 

153 inverse-variance weighted) bilinear deposit of the samples onto a 

154 regular grid, followed by a mask-normalized Gaussian smoothing. An 

155 exact interpolant must not be used here: it would carry the full 

156 per-sample shot noise into the maps, where it becomes spatially 

157 correlated and contaminates both E and B modes (pair-based 

158 correlation estimators never see that noise at non-zero lag, map 

159 estimators do). The smoothing also inpaints small coverage gaps 

160 (detector gaps) up to about the smoothing length. 

161 

162 The divergence and curl are central differences on the smoothed 

163 grid; they are the sources of the E and B modes of the field (the 

164 Laplacians of the scalar and pseudo-scalar potentials). For the E 

165 (curl-free) and B (divergence-free) vector components, the map is 

166 first multiplied by a smooth apodization window built from the 

167 coverage mask, so that the non-rectangular field boundary tapers to 

168 zero instead of jumping (hard edges leak E power into B and vice 

169 versa); each Fourier mode of the windowed, zero-padded map is then 

170 projected onto (E) and perpendicular to (B) its wave vector, with 

171 the k = 0 mode assigned to E. E + B reconstructs the windowed map, 

172 i.e. the smoothed field wherever the window is 1. All outputs are 

173 only reported where the sampled window exceeds 0.9 (NaN elsewhere); 

174 even there, on a finite field the decomposition is not unique and 

175 some residual E/B mixing of field-scale modes is unavoidable. 

176 

177 Parameters 

178 ---------- 

179 coords : `numpy.ndarray` 

180 Positions of the field samples, in degrees. (n_samples, 2) 

181 field : `numpy.ndarray` 

182 Vector field values at `coords`, in mas. (n_samples, 2) 

183 gridResolution : `int` 

184 Number of grid nodes per axis of the intermediate regular grid. 

185 smoothing : `float`, optional 

186 Gaussian smoothing scale applied to the gridded field, in 

187 degrees. If not given, an adaptive default is used: 1.5 times 

188 the mean sample separation (about 15 samples per kernel, so the 

189 map shot-noise floor stays constant across sample densities), 

190 floored at one grid cell. 

191 weights : `numpy.ndarray`, optional 

192 Per-sample weights for the gridding (typically inverse total 

193 variance). Uniform if not given. (n_samples,) 

194 deconvolve : `bool`, optional 

195 Whether to deconvolve the known gridding transfer function (the 

196 bilinear deposit kernel times the Gaussian smoothing) from the 

197 windowed maps, with a Wiener regularization that caps the boost 

198 so that noise is not re-amplified. This corrects most of the 

199 attenuation of the map correlation functions at scales near the 

200 smoothing length. [default: True] 

201 

202 Returns 

203 ------- 

204 div : `numpy.ndarray` 

205 Divergence of the field at `coords`, in mas per degree. 

206 NaN outside the apodization window. (n_samples,) 

207 curl : `numpy.ndarray` 

208 Curl (z-component, a pseudo-scalar) of the field at `coords`, 

209 in mas per degree. (n_samples,) 

210 fieldE : `numpy.ndarray` 

211 E-mode (curl-free) component of the field at `coords`, in mas. 

212 (n_samples, 2) 

213 fieldB : `numpy.ndarray` 

214 B-mode (divergence-free) component of the field at `coords`, 

215 in mas. (n_samples, 2) 

216 """ 

217 nSamples = len(coords) 

218 if weights is None: 218 ↛ 219line 218 didn't jump to line 219 because the condition on line 218 was never true

219 weights = np.ones(nSamples) 

220 xNodes = np.linspace(np.min(coords[:, 0]), np.max(coords[:, 0]), gridResolution) 

221 yNodes = np.linspace(np.min(coords[:, 1]), np.max(coords[:, 1]), gridResolution) 

222 xSpacing = xNodes[1] - xNodes[0] 

223 ySpacing = yNodes[1] - yNodes[0] 

224 

225 # Weighted bilinear (cloud-in-cell) deposit of the samples onto the 

226 # grid. Maps are indexed [iy, ix]. 

227 gx = np.clip((coords[:, 0] - xNodes[0]) / xSpacing, 0.0, gridResolution - 1.0) 

228 gy = np.clip((coords[:, 1] - yNodes[0]) / ySpacing, 0.0, gridResolution - 1.0) 

229 ix = np.clip(np.floor(gx).astype(np.intp), 0, gridResolution - 2) 

230 iy = np.clip(np.floor(gy).astype(np.intp), 0, gridResolution - 2) 

231 tx = gx - ix 

232 ty = gy - iy 

233 numeratorX = np.zeros((gridResolution, gridResolution)) 

234 numeratorY = np.zeros((gridResolution, gridResolution)) 

235 denominator = np.zeros((gridResolution, gridResolution)) 

236 for shiftX, shiftY, cellWeight in [ 

237 (0, 0, (1.0 - tx) * (1.0 - ty)), 

238 (1, 0, tx * (1.0 - ty)), 

239 (0, 1, (1.0 - tx) * ty), 

240 (1, 1, tx * ty), 

241 ]: 

242 np.add.at(numeratorX, (iy + shiftY, ix + shiftX), weights * cellWeight * field[:, 0]) 

243 np.add.at(numeratorY, (iy + shiftY, ix + shiftX), weights * cellWeight * field[:, 1]) 

244 np.add.at(denominator, (iy + shiftY, ix + shiftX), weights * cellWeight) 

245 

246 # Mask-normalized Gaussian smoothing: denoises the maps and inpaints 

247 # coverage gaps up to about the smoothing length. 

248 if smoothing is None: 248 ↛ 257line 248 didn't jump to line 257 because the condition on line 248 was always true

249 # Adaptive default: 1.5x the mean star separation puts ~15 

250 # stars (2 pi 1.5^2) in every smoothing kernel at any density, 

251 # keeping the map shot-noise floor sigma_n^2 / (rho A_kernel) 

252 # roughly constant across visits; floored at one grid cell. 

253 footprintArea = np.count_nonzero(denominator > 0.0) * xSpacing * ySpacing 

254 meanSeparation = np.sqrt(footprintArea / nSamples) 

255 sigmaCells = max(1.0, 1.5 * meanSeparation / (0.5 * (xSpacing + ySpacing))) 

256 else: 

257 sigmaCells = smoothing / (0.5 * (xSpacing + ySpacing)) 

258 numeratorX = gaussian_filter(numeratorX, sigmaCells) 

259 numeratorY = gaussian_filter(numeratorY, sigmaCells) 

260 denominatorSmooth = gaussian_filter(denominator, sigmaCells) 

261 reference = np.median(denominatorSmooth[denominator > 0.0]) 

262 covered = denominatorSmooth > 0.05 * reference 

263 with np.errstate(invalid="ignore", divide="ignore"): 

264 mapX = np.where(covered, numeratorX / denominatorSmooth, 0.0) 

265 mapY = np.where(covered, numeratorY / denominatorSmooth, 0.0) 

266 

267 dMapXdY, dMapXdX = np.gradient(mapX, yNodes, xNodes) 

268 dMapYdY, dMapYdX = np.gradient(mapY, yNodes, xNodes) 

269 divMap = dMapXdX + dMapYdY 

270 curlMap = dMapYdX - dMapXdY 

271 

272 # Smooth apodization window built from the (density-independent) 

273 # coverage mask, tapering the field boundary and remaining holes. 

274 window = np.clip(gaussian_filter(covered.astype(float), max(2.0 * sigmaCells, 2.0)), 0.0, 1.0) 

275 windowedX = window * mapX 

276 windowedY = window * mapY 

277 

278 # Helmholtz decomposition in Fourier space, on zero-padded maps. 

279 ny, nx = mapX.shape 

280 nyPad = next_fast_len(2 * ny) 

281 nxPad = next_fast_len(2 * nx) 

282 padX = np.zeros((nyPad, nxPad)) 

283 padY = np.zeros((nyPad, nxPad)) 

284 padX[:ny, :nx] = windowedX 

285 padY[:ny, :nx] = windowedY 

286 kx = 2.0 * np.pi * np.fft.fftfreq(nxPad, d=xSpacing)[np.newaxis, :] 

287 ky = 2.0 * np.pi * np.fft.fftfreq(nyPad, d=ySpacing)[:, np.newaxis] 

288 kSquared = kx**2 + ky**2 

289 kSquared[0, 0] = 1.0 

290 hatX = fft2(padX) 

291 hatY = fft2(padY) 

292 if deconvolve: 292 ↛ 311line 292 didn't jump to line 311 because the condition on line 292 was always true

293 # Known transfer function of the gridding: the bilinear deposit 

294 # is a convolution with a tent of base 2 cells (Fourier 

295 # transform sinc^2), and the smoothing is a Gaussian. The 

296 # Wiener regularization T / (T^2 + alpha^2) caps the 

297 # deconvolution boost at 1 / (2 alpha), so the shot noise 

298 # suppressed by the smoothing is not re-amplified (alpha tuned 

299 # on DP2 visit 2026010100522: 0.1-0.3 all behave, 0.15 caps the 

300 # boost at ~3.3 and recovers ~90% of the small-scale power). 

301 alpha = 0.15 

302 sigmaDeg = sigmaCells * 0.5 * (xSpacing + ySpacing) 

303 transfer = ( 

304 np.sinc(kx * xSpacing / (2.0 * np.pi)) ** 2 

305 * np.sinc(ky * ySpacing / (2.0 * np.pi)) ** 2 

306 * np.exp(-0.5 * kSquared * sigmaDeg**2) 

307 ) 

308 wiener = transfer / (transfer**2 + alpha**2) 

309 hatX = hatX * wiener 

310 hatY = hatY * wiener 

311 kDotU = kx * hatX + ky * hatY 

312 hatEX = kDotU * kx / kSquared 

313 hatEY = kDotU * ky / kSquared 

314 # The k = 0 mode has no direction; assign it to E by convention. 

315 hatEX[0, 0] = hatX[0, 0] 

316 hatEY[0, 0] = hatY[0, 0] 

317 mapEX = ifft2(hatEX).real[:ny, :nx] 

318 mapEY = ifft2(hatEY).real[:ny, :nx] 

319 if deconvolve: 319 ↛ 323line 319 didn't jump to line 323 because the condition on line 319 was always true

320 baseX = ifft2(hatX).real[:ny, :nx] 

321 baseY = ifft2(hatY).real[:ny, :nx] 

322 else: 

323 baseX = windowedX 

324 baseY = windowedY 

325 mapBX = baseX - mapEX 

326 mapBY = baseY - mapEY 

327 

328 outputs = [] 

329 for outputMap in [divMap, curlMap, mapEX, mapEY, mapBX, mapBY, window]: 

330 interpolator = RegularGridInterpolator( 

331 (yNodes, xNodes), 

332 outputMap, 

333 method="linear", 

334 bounds_error=False, 

335 fill_value=np.nan, 

336 ) 

337 outputs.append(interpolator(coords[:, ::-1])) 

338 div, curl, ex, ey, bx, by, windowAtCoords = outputs 

339 bad = ~(windowAtCoords > 0.9) 

340 for output in [div, curl, ex, ey, bx, by]: 

341 output[bad] = np.nan 

342 return div, curl, np.array([ex, ey]).T, np.array([bx, by]).T 

343 

344 

345class SingularMatrixError(pipeBase.AlgorithmError): 

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

347 algebra error.""" 

348 

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

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

351 self._nSources = nSources 

352 

353 @property 

354 def metadata(self): 

355 return { 

356 "nSources": self._nSources, 

357 } 

358 

359 

360class NotPositiveDefiniteMatrixError(pipeBase.AlgorithmError): 

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

362 linear algebra error.""" 

363 

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

365 super().__init__( 

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

367 ) 

368 self._nSources = nSources 

369 

370 @property 

371 def metadata(self): 

372 return { 

373 "nSources": self._nSources, 

374 } 

375 

376 

377class GaussianProcessesTurbulenceFitConnections( 

378 pipeBase.PipelineTaskConnections, 

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

380 defaultTemplates={ 

381 "inputName": "gbdesHealpix3AstrometricFit", 

382 }, 

383): 

384 inputWcs = pipeBase.connectionTypes.Input( 

385 doc=( 

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

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

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

389 ), 

390 name="{inputName}SkyWcsCatalog", 

391 storageClass="ExposureCatalog", 

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

393 ) 

394 inputPositions = pipeBase.connectionTypes.Input( 

395 doc=( 

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

397 "plane coordinates and chisq values." 

398 ), 

399 name="{inputName}_fitStars", 

400 storageClass="ArrowAstropy", 

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

402 deferLoad=True, 

403 ) 

404 camera = pipeBase.connectionTypes.PrerequisiteInput( 

405 doc="Input camera, used to transform source positions to the focal plane.", 

406 name="camera", 

407 storageClass="Camera", 

408 dimensions=("instrument",), 

409 isCalibration=True, 

410 ) 

411 outputWcs = pipeBase.connectionTypes.Output( 

412 doc=( 

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

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

415 "for fast lookups of a detector." 

416 ), 

417 name="turbulenceCorrectedSkyWcsCatalog", 

418 storageClass="ExposureCatalog", 

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

420 ) 

421 hyperparameters = pipeBase.connectionTypes.Output( 

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

423 name="turbulence_fit_hyperparameters", 

424 storageClass="ArrowAstropy", 

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

426 ) 

427 sourceTable = pipeBase.connectionTypes.Output( 

428 doc=( 

429 "Per-source table with positions, residuals, and Gaussian Processes predictions in tangent" 

430 " plane and detector pixel coordinates, along with the training/validation split, and" 

431 " optionally (config.saveDivCurl) the divergence, curl, and E/B decomposition of the" 

432 " residual field." 

433 ), 

434 name="turbulence_fit_source", 

435 storageClass="ArrowAstropy", 

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

437 ) 

438 

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

440 super().__init__(config=config) 

441 

442 if not self.config.saveSourceTable: 

443 self.outputs.remove("sourceTable") 

444 self.prerequisiteInputs.remove("camera") 

445 

446 if not self.config.healpix: 

447 self.dimensions.remove("healpix3") 

448 if self.config.healpix is None: 

449 extra_dimensions = [] 

450 else: 

451 extra_dimensions = ["tract", "skymap"] 

452 self.dimensions.update(extra_dimensions) 

453 self.inputWcs = dataclasses.replace( 

454 self.inputWcs, dimensions=["instrument", "visit"] + extra_dimensions 

455 ) 

456 self.inputPositions = dataclasses.replace( 

457 self.inputPositions, dimensions=["instrument", "band", "physical_filter"] + extra_dimensions 

458 ) 

459 self.outputWcs = dataclasses.replace( 

460 self.outputWcs, dimensions=["instrument", "visit"] + extra_dimensions 

461 ) 

462 self.hyperparameters = dataclasses.replace( 

463 self.hyperparameters, dimensions=["instrument", "visit"] + extra_dimensions 

464 ) 

465 if self.config.saveSourceTable: 

466 self.sourceTable = dataclasses.replace( 

467 self.sourceTable, dimensions=["instrument", "visit"] + extra_dimensions 

468 ) 

469 

470 

471class GaussianProcessesTurbulenceFitConfig( 

472 pipeBase.PipelineTaskConfig, pipelineConnections=GaussianProcessesTurbulenceFitConnections 

473): 

474 optimizer = pexConfig.ChoiceField( 

475 dtype=str, 

476 doc="Gaussian Processes method used to model the astrometric residuals.", 

477 default="anisotropic", 

478 allowed={ 

479 "anisotropic": ( 

480 "Fit the hyperparameters of a parametric kernel on the measured" 

481 " 2-point correlation function (Leget et al. 2021, A&A 650, A81)." 

482 ), 

483 "empirical-2pcf": ( 

484 "Use the measured 2-point correlation function directly as the" 

485 " kernel, with no hyperparameter fit (TURBO-GP," 

486 " RTN-129)." 

487 ), 

488 }, 

489 ) 

490 initKernel = pexConfig.Field( 

491 dtype=str, 

492 doc=( 

493 "The type of function that will be used to modeled spatial correlation." 

494 " Only used by the 'anisotropic' optimizer." 

495 ), 

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

497 ) 

498 initAnisotropicCorrelationLength = pexConfig.ListField( 

499 dtype=float, 

500 doc=( 

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

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

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

504 "where 0 means the correlation is isotropic." 

505 ), 

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

507 ) 

508 correlationSeparationMin = pexConfig.Field( 

509 dtype=float, 

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

511 default=0.0, 

512 optional=True, 

513 ) 

514 correlationSeparationMax = pexConfig.Field( 

515 dtype=float, 

516 doc=( 

517 "Maximum distance separation in degrees in the computation of the 2-point correlation function." 

518 " For the 'empirical-2pcf' optimizer, this is also the half width of the 2-point correlation" 

519 " function grid used as the kernel." 

520 ), 

521 default=0.3, 

522 optional=True, 

523 ) 

524 correlationPixelSize = pexConfig.Field( 

525 dtype=float, 

526 doc=( 

527 "Pixel size in degrees of the 2-point correlation function grid used as the kernel." 

528 " Only used by the 'empirical-2pcf' optimizer." 

529 ), 

530 default=0.00556, 

531 ) 

532 powerThreshold = pexConfig.Field( 

533 dtype=float, 

534 doc=( 

535 "Signal-to-noise threshold below which Fourier modes of the measured 2-point correlation" 

536 " function are set to zero. Only used by the 'empirical-2pcf' optimizer." 

537 ), 

538 default=2.5, 

539 ) 

540 apodize = pexConfig.Field( 

541 dtype=bool, 

542 doc=( 

543 "Whether to apodize the measured 2-point correlation function before taking its Fourier" 

544 " transform. Only used by the 'empirical-2pcf' optimizer." 

545 ), 

546 default=True, 

547 ) 

548 apodWindow = pexConfig.ChoiceField( 

549 dtype=str, 

550 doc="Apodization window function. Only used by the 'empirical-2pcf' optimizer.", 

551 default="hann", 

552 allowed={ 

553 "hann": "Hann window; gentler taper at the price of more spectral leakage.", 

554 "blackman-harris": "Blackman-Harris window, as used in Gomes et al. 2025.", 

555 }, 

556 ) 

557 apodRadius = pexConfig.Field( 

558 dtype=float, 

559 doc=( 

560 "Radius in degrees where the apodization window reaches zero. If None, the window reaches" 

561 " zero at correlationSeparationMax. Only used by the 'empirical-2pcf' optimizer." 

562 ), 

563 default=0.278, 

564 optional=True, 

565 ) 

566 apodAnisotropy = pexConfig.ChoiceField( 

567 dtype=str, 

568 doc="Anisotropy of the apodization window. Only used by the 'empirical-2pcf' optimizer.", 

569 default="auto", 

570 allowed={ 

571 "auto": ( 

572 "Measure the anisotropy of the 2-point correlation function with adaptive" 

573 " weighted second moments and use a matched elliptical window." 

574 ), 

575 "none": "Isotropic apodization window.", 

576 }, 

577 ) 

578 apodGScale = pexConfig.Field( 

579 dtype=float, 

580 doc=( 

581 "Factor multiplying the measured anisotropy of the apodization window when" 

582 " apodAnisotropy='auto'. Only used by the 'empirical-2pcf' optimizer." 

583 ), 

584 default=1.0, 

585 ) 

586 saveSourceTable = pexConfig.Field( 

587 dtype=bool, 

588 doc=( 

589 "Save the per-source table with positions, residuals, and Gaussian Processes predictions" 

590 " in tangent plane and detector pixel coordinates." 

591 ), 

592 default=False, 

593 ) 

594 saveDivCurl = pexConfig.Field( 

595 dtype=bool, 

596 doc=( 

597 "Add to the per-source table the divergence and curl of the residual field (the sources of" 

598 " its E and B modes) and its Helmholtz decomposition into E (curl-free) and B" 

599 " (divergence-free) vector components, for both the measured residuals and the Gaussian" 

600 " Processes predictions, in tangent plane and detector pixel coordinates. Requires" 

601 " saveSourceTable." 

602 ), 

603 default=False, 

604 ) 

605 divCurlGridResolution = pexConfig.Field( 

606 dtype=int, 

607 doc=( 

608 "Number of grid nodes per axis used to grid the residual field for the divergence, curl," 

609 " and E/B decomposition computation. Only used if saveDivCurl is set." 

610 ), 

611 default=300, 

612 ) 

613 divCurlDeconvolve = pexConfig.Field( 

614 dtype=bool, 

615 doc=( 

616 "Deconvolve the known gridding transfer function (bilinear deposit times Gaussian" 

617 " smoothing) from the E/B maps, with a Wiener regularization capping the noise boost." 

618 " Recovers most of the small-scale power attenuated by the smoothing. Only used if" 

619 " saveDivCurl is set." 

620 ), 

621 default=True, 

622 ) 

623 divCurlSmoothing = pexConfig.Field( 

624 dtype=float, 

625 doc=( 

626 "Gaussian smoothing scale in degrees applied to the gridded residual field before the" 

627 " divergence, curl, and E/B decomposition computation, to suppress the per-source shot" 

628 " noise that gridding would otherwise turn into correlated E/B power. If not set, an" 

629 " adaptive default of 1.5 times the mean source separation (floored at one grid cell)" 

630 " keeps the map noise floor constant across source densities. Only used if saveDivCurl" 

631 " is set." 

632 ), 

633 default=None, 

634 optional=True, 

635 ) 

636 maxTrainingPoints = pexConfig.Field( 

637 dtype=int, 

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

639 default=10000, 

640 ) 

641 pixelSize = pexConfig.Field( 

642 dtype=float, 

643 doc="Pixel size in arcseconds.", 

644 default=0.2, 

645 ) 

646 healpix = pexConfig.Field( 

647 dtype=bool, 

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

649 default=True, 

650 optional=True, 

651 ) 

652 splineDegree = pexConfig.Field( 

653 dtype=int, 

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

655 default=4, 

656 ) 

657 splineNNodes = pexConfig.Field( 

658 dtype=int, 

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

660 default=30, 

661 ) 

662 splineBuffer = pexConfig.Field( 

663 dtype=float, 

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

665 default=0.1, 

666 ) 

667 

668 def validate(self): 

669 super().validate() 

670 if self.saveDivCurl and not self.saveSourceTable: 

671 raise pexConfig.FieldValidationError( 

672 self.__class__.saveDivCurl, 

673 self, 

674 "saveDivCurl requires saveSourceTable to be set.", 

675 ) 

676 if self.optimizer == "empirical-2pcf": 

677 if self.correlationSeparationMax is None: 677 ↛ 678line 677 didn't jump to line 678 because the condition on line 677 was never true

678 raise pexConfig.FieldValidationError( 

679 self.__class__.correlationSeparationMax, 

680 self, 

681 "correlationSeparationMax must be set for the 'empirical-2pcf' optimizer.", 

682 ) 

683 if self.correlationSeparationMax / self.correlationPixelSize < 2: 683 ↛ 684line 683 didn't jump to line 684 because the condition on line 683 was never true

684 raise pexConfig.FieldValidationError( 

685 self.__class__.correlationPixelSize, 

686 self, 

687 "correlationSeparationMax must span at least 2 pixels of correlationPixelSize" 

688 " for the 'empirical-2pcf' optimizer.", 

689 ) 

690 

691 

692class GaussianProcessesTurbulenceFitTask(pipeBase.PipelineTask): 

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

694 they are due to atmospheric turbulence. 

695 

696 Two methods are available, selected with `config.optimizer`: fitting the 

697 hyperparameters of a parametric kernel on the measured 2-point correlation 

698 function (Leget et al. 2021, A&A 650, A81), or using the measured 2-point 

699 correlation function directly as the kernel (TURBO-GP, 

700 RTN-129). 

701 """ 

702 

703 ConfigClass = GaussianProcessesTurbulenceFitConfig 

704 _DefaultName = "gaussianProcessesTurbulenceFit" 

705 

706 def run(self, inputWcs, inputPositions, camera=None): 

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

708 Gaussian Processes prediction from the WCS to account for atmospheric 

709 turbulence. 

710 

711 Parameters 

712 ---------- 

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

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

715 inputPositions : `astropy.table.Table` 

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

717 camera : `lsst.afw.cameraGeom.Camera`, optional 

718 Camera object, used to transform source positions to the focal 

719 plane. Only provided when `config.saveSourceTable` is set. 

720 

721 Returns 

722 ------- 

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

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

725 Catalog with WCS after inserting the correction for atmospheric 

726 turbulence. 

727 ``hyperparameters`` : `astropy.table.Table` 

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

729 ``sourceTable`` : `astropy.table.Table` 

730 Per-source table with positions, residuals, and Gaussian 

731 Processes predictions in tangent plane and detector pixel 

732 coordinates. Only set when `config.saveSourceTable` is set. 

733 """ 

734 

735 visit = inputWcs[0]["visit"] 

736 

737 columns = [ 

738 "xworld", 

739 "yworld", 

740 "xresw", 

741 "yresw", 

742 "exposureName", 

743 "xpix", 

744 "ypix", 

745 "deviceName", 

746 "clip", 

747 "covTotalW_00", 

748 "covTotalW_11", 

749 ] 

750 inputPositions = inputPositions.get(parameters={"columns": columns}) 

751 

752 visitPositions = inputPositions[ 

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

754 ] 

755 

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

757 

758 try: 

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

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

761 except np.linalg.LinAlgError as e: 

762 if "not positive definite" in str(e): 

763 error = pipeBase.AnnotatedPartialOutputsError.annotate( 

764 NotPositiveDefiniteMatrixError(trainInd.sum()), 

765 self, 

766 log=self.log, 

767 ) 

768 raise error from e 

769 else: 

770 raise 

771 

772 outputs = {"outputWcs": wcsWithSpline, "hyperparameters": hyperparameters} 

773 if self.config.saveSourceTable: 

774 outputs["sourceTable"] = self.makeSourceTable( 

775 gpx, gpy, visitPositions, allTPCoords, trainInd, inputWcs, camera, visit 

776 ) 

777 

778 return pipeBase.Struct(**outputs) 

779 

780 def runGP(self, inputWcs, positions): 

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

782 

783 Parameters 

784 ---------- 

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

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

787 inputPositions : `astropy.table.Table` 

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

789 

790 Returns 

791 ------- 

792 gpx : `treegp.gp_interp.GPInterpolation` 

793 Gaussian Processes interpolator for x-direction residuals. 

794 gpy : `treegp.gp_interp.GPInterpolation` 

795 Gaussian Processes interpolator for y-direction residuals. 

796 trainInds : `numpy.ndarray` 

797 Array of indices for points used in training. 

798 testInds : `numpy.ndarray` 

799 Array of indices for points not used in training. 

800 hyperparameters : `astropy.table.Table` 

801 Table with an ``x`` and a ``y`` column, one per residual 

802 direction. For the 'anisotropic' optimizer, they hold the 

803 best-fit kernel hyperparameters. For the 'empirical-2pcf' 

804 optimizer, they hold diagnostics of the empirical kernel, in 

805 the order [xi0, g1 measured, g2 measured, g1 applied, 

806 g2 applied, pixel size, number of grid pixels], where xi0 is 

807 the zero-lag variance in mas^2 and the g values describe the 

808 anisotropy of the apodization window (NaN if not measured). 

809 allTPCoords : `numpy.ndarray` 

810 Tangent plane coordinates of the input positions in degrees, 

811 using input WCS, without gp correction. 

812 (n_samples, 2) 

813 """ 

814 if self.config.optimizer == "empirical-2pcf" and not hasattr(treegp, "empirical_2pcf"): 814 ↛ 815line 814 didn't jump to line 815 because the condition on line 814 was never true

815 raise RuntimeError( 

816 "The version of treegp in this environment does not support the 'empirical-2pcf'" 

817 " optimizer (TURBO-GP, RTN-129); it requires the treegp branch tickets/DM-55875." 

818 ) 

819 

820 dx = positions["xresw"] 

821 dy = positions["yresw"] 

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

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

824 

825 # Get tangent plane coordinates for input points 

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

827 for detector in inputWcs: 

828 detId = detector["id"] 

829 detWCS = detector.wcs 

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

831 detectorSources = positions[detInd] 

832 

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

834 tangentPlaneCoords = tangentPlaneToSky.applyForward( 

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