Coverage for python/lsst/drp/tasks/fit_turbulence.py: 56%
257 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-10-09 11:09 +0000
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
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.interpolate import RectBivariateSpline
33import lsst.afw.geom as afwgeom
34import lsst.afw.table
35import lsst.pex.config as pexConfig
36import lsst.pipe.base as pipeBase
38# We need to explicitly turn off multiprocessing in treecorr which is used
39# by treegp.
40treecorr.set_max_omp_threads(1)
42__all__ = [
43 "GaussianProcessesTurbulenceFitConnections",
44 "GaussianProcessesTurbulenceFitConfig",
45 "GaussianProcessesTurbulenceFitTask",
46]
49def plot_visit(x, y, dx, dy, predx, predy):
50 """Utility function for plotting Gaussian Processes results.
52 Parameters
53 ----------
54 x : `np.ndarray`
55 x-direction coordinates.
56 y : `np.ndarray`
57 y-direction coordinates.
58 dx : `np.ndarray`
59 x-direction residuals to be fit.
60 dy : `np.ndarray`
61 x-direction residuals to be fit.
62 predx : `np.ndarray`
63 x-direction prediction.
64 predy : `np.ndarray`
65 y-direction prediction.
67 Returns
68 -------
69 fig : `matplotlib.pyplot.Figure`
70 Figure showing input data, Gaussian Processes prediction, and E and
71 B-modes.
72 """
74 xie, xib, logr = treegp.comp_eb_treecorr(x, y, dx, dy, rmin=20 / 3600, rmax=0.6, dlogr=0.3)
75 xie_resid, xib_resid, logr_resid = treegp.comp_eb_treecorr(
76 x, y, dx - predx, dy - predy, rmin=20 / 3600, rmax=0.6, dlogr=0.3
77 )
79 residualLimit = np.nanstd(dx)
81 fig, subs = plt.subplot_mosaic(
82 [["dx", "predx", "residx", "eb"], ["dy", "predy", "residy", "eb"]],
83 figsize=(15, 8),
84 layout="constrained",
85 )
86 plt.subplots_adjust(wspace=0.3, right=0.99, left=0.05)
87 im = subs["dx"].scatter(x, y, c=dx, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1)
88 subs["dy"].scatter(x, y, c=dy, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1)
90 subs["predx"].scatter(x, y, c=predx, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1)
91 subs["predy"].scatter(x, y, c=predy, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1)
93 subs["residx"].scatter(
94 x, y, c=dx - predx, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1
95 )
96 subs["residy"].scatter(
97 x, y, c=dy - predy, vmin=-residualLimit, vmax=residualLimit, cmap=plt.cm.seismic, s=1
98 )
100 cb = fig.colorbar(
101 im, ax=[subs["dx"], subs["dy"], subs["predx"], subs["predy"], subs["residx"], subs["residy"]]
102 )
104 subs["eb"].scatter(np.exp(logr) * 60, xie, c="b", label="E-mode")
105 subs["eb"].scatter(np.exp(logr) * 60, xib, c="r", label="B-mode")
107 subs["eb"].scatter(
108 np.exp(logr_resid) * 60, xie_resid, c="b", marker="+", label="E-mode after GP correction"
109 )
110 subs["eb"].scatter(
111 np.exp(logr_resid) * 60, xib_resid, c="r", marker="+", label="B-mode after GP correction"
112 )
113 subs["eb"].legend()
114 subs["eb"].grid(True)
116 subs["dx"].set_aspect("equal")
117 subs["dy"].set_aspect("equal")
118 subs["predx"].set_aspect("equal")
119 subs["predy"].set_aspect("equal")
120 subs["residx"].set_aspect("equal")
121 subs["residy"].set_aspect("equal")
122 subs["dy"].set_xlabel("x (degree)")
123 subs["predy"].set_xlabel("x (degree)")
124 subs["residy"].set_xlabel("x (degree)")
125 subs["dy"].set_ylabel("y (degree)")
126 subs["dx"].set_ylabel("y (degree)")
128 subs["dx"].set_title(r"$\delta$x")
129 subs["predx"].set_title("GP prediction")
130 subs["residx"].set_title("Residual")
132 subs["dy"].set_title(r"$\delta$y")
133 subs["predy"].set_title("GP prediction")
134 subs["residy"].set_title("Residual")
136 cb.set_label("mas")
138 subs["eb"].set_title("E and B modes")
139 subs["eb"].set_ylabel(r"$\xi_{E/B}$ (mas$^2$)")
140 subs["eb"].set_xlabel(r"$\Delta \theta$ (arcmin)")
142 return fig
145class SingularMatrixError(pipeBase.AlgorithmError):
146 """Raised if the Gaussian Processes fit raises a Singular Matrix linear
147 algebra error."""
149 def __init__(self, nSources) -> None:
150 super().__init__("The Gaussian Processes fit failed with a singular matrix linear algebra error.")
151 self._nSources = nSources
153 @property
154 def metadata(self):
155 return {
156 "nSources": self._nSources,
157 }
160class NotPositiveDefiniteMatrixError(pipeBase.AlgorithmError):
161 """Raised if the Gaussian Processes fit raises a not positive definite
162 linear algebra error."""
164 def __init__(self, nSources) -> None:
165 super().__init__(
166 "The Gaussian Processes fit failed with a not-positive-definite linear algebra error."
167 )
168 self._nSources = nSources
170 @property
171 def metadata(self):
172 return {
173 "nSources": self._nSources,
174 }
177class GaussianProcessesTurbulenceFitConnections(
178 pipeBase.PipelineTaskConnections,
179 dimensions=("instrument", "visit", "healpix3"),
180 defaultTemplates={
181 "inputName": "gbdesHealpix3AstrometricFit",
182 },
183):
184 inputWcs = pipeBase.connectionTypes.Input(
185 doc=(
186 "Per-healpix, per-visit world coordinate systems derived from the fitted model."
187 " These catalogs only contain entries for detectors with an output, and use"
188 " the detector id for the catalog id, sorted on id for fast lookups of a detector."
189 ),
190 name="{inputName}SkyWcsCatalog",
191 storageClass="ExposureCatalog",
192 dimensions=("instrument", "visit", "healpix3"),
193 )
194 inputPositions = pipeBase.connectionTypes.Input(
195 doc=(
196 "Catalog of sources used in fit, along with residuals in pixel coordinates and tangent "
197 "plane coordinates and chisq values."
198 ),
199 name="{inputName}_fitStars",
200 storageClass="ArrowAstropy",
201 dimensions=("instrument", "healpix3", "physical_filter"),
202 deferLoad=True,
203 )
204 outputWcs = pipeBase.connectionTypes.Output(
205 doc=(
206 "Per-visit world coordinate systems derived from the fitted model. These catalogs only contain "
207 "entries for detectors with an output, and use the detector id for the catalog id, sorted on id "
208 "for fast lookups of a detector."
209 ),
210 name="turbulenceCorrectedSkyWcsCatalog",
211 storageClass="ExposureCatalog",
212 dimensions=("instrument", "visit", "healpix3"),
213 )
214 hyperparameters = pipeBase.connectionTypes.Output(
215 doc="Best fit hyperparameters for the Gaussian Processes fit.",
216 name="turbulence_fit_hyperparameters",
217 storageClass="ArrowAstropy",
218 dimensions=("instrument", "visit", "healpix3"),
219 )
221 def __init__(self, *, config=None):
222 super().__init__(config=config)
224 if not self.config.healpix:
225 self.dimensions.remove("healpix3")
226 if self.config.healpix is None:
227 extra_dimensions = []
228 else:
229 extra_dimensions = ["tract", "skymap"]
230 self.dimensions.update(extra_dimensions)
231 self.inputWcs = dataclasses.replace(
232 self.inputWcs, dimensions=["instrument", "visit"] + extra_dimensions
233 )
234 self.inputPositions = dataclasses.replace(
235 self.inputPositions, dimensions=["instrument", "band", "physical_filter"] + extra_dimensions
236 )
237 self.outputWcs = dataclasses.replace(
238 self.outputWcs, dimensions=["instrument", "visit"] + extra_dimensions
239 )
240 self.hyperparameters = dataclasses.replace(
241 self.hyperparameters, dimensions=["instrument", "visit"] + extra_dimensions
242 )
245class GaussianProcessesTurbulenceFitConfig(
246 pipeBase.PipelineTaskConfig, pipelineConnections=GaussianProcessesTurbulenceFitConnections
247):
248 initKernel = pexConfig.Field(
249 dtype=str,
250 doc="The type of function that will be used to modeled spatial correlation.",
251 default="15**2 * AnisotropicVonKarman(invLam=array([[1./0.8**2,0],[0,1./0.8**2]]))",
252 )
253 initAnisotropicCorrelationLength = pexConfig.ListField(
254 dtype=float,
255 doc=(
256 "The initial parameters for fiting the anisotropic correlation length. p0[0] is equivalent of "
257 "the isotropic correlation length in degrees, and p0[1]/p0[2] are ellipticity parameters and are "
258 "mathematically equivalent to e1/e2 in weak-lensing. p0[1]/p0[2] must be in the range [-1,1], "
259 "where 0 means the correlation is isotropic."
260 ),
261 default=[1, -0.2, -0.2],
262 )
263 correlationSeparationMin = pexConfig.Field(
264 dtype=float,
265 doc="Minimum distance separation in degrees in the computation of the 2-point correlation function.",
266 default=0.0,
267 optional=True,
268 )
269 correlationSeparationMax = pexConfig.Field(
270 dtype=float,
271 doc="Maximum distance separation in degrees in the computation of the 2-point correlation function.",
272 default=0.3,
273 optional=True,
274 )
275 maxTrainingPoints = pexConfig.Field(
276 dtype=int,
277 doc="Maximum number of points to use in the Gaussian Processes training.",
278 default=10000,
279 )
280 pixelSize = pexConfig.Field(
281 dtype=float,
282 doc="Pixel size in arcseconds.",
283 default=0.2,
284 )
285 healpix = pexConfig.Field(
286 dtype=bool,
287 doc="Use input WCS calculated over healpix-based region. If false, use tract-based WCS.",
288 default=True,
289 optional=True,
290 )
291 splineDegree = pexConfig.Field(
292 dtype=int,
293 doc="Degree of the spline expressing Gaussian Processes prediction.",
294 default=4,
295 )
296 splineNNodes = pexConfig.Field(
297 dtype=int,
298 doc="Number of nodes to use for the spline expressing Gaussian Processes prediction.",
299 default=30,
300 )
301 splineBuffer = pexConfig.Field(
302 dtype=float,
303 doc="Minimum distance in degrees to extend spline map outside the detector boundary.",
304 default=0.1,
305 )
308class GaussianProcessesTurbulenceFitTask(pipeBase.PipelineTask):
309 """Run Gaussian Processes on astrometric residuals with the assumption that
310 they are due to atmospheric turbulence."""
312 ConfigClass = GaussianProcessesTurbulenceFitConfig
313 _DefaultName = "gaussianProcessesTurbulenceFit"
315 def run(self, inputWcs, inputPositions):
316 """Run Gaussian Processes on position residuals and subtract the fitted
317 Gaussian Processes prediction from the WCS to account for atmospheric
318 turbulence.
320 Parameters
321 ----------
322 inputWcs : `lsst.afw.table.ExposureCatalog`
323 Catalog with WCSs for each detector of the input exposure.
324 inputPositions : `astropy.table.Table`
325 Catalog of input positions with residuals to the current best fit.
327 Returns
328 -------
329 result : `lsst.pipe.base.Struct`
330 ``outputWcs`` : `lsst.afw.table.ExposureCatalog`
331 Catalog with WCS after inserting the correction for atmospheric
332 turbulence.
333 ``hyperparameters`` : `astropy.table.Table`
334 Table of best-fit hyperparameters in x and y-directions.
335 """
337 visit = inputWcs[0]["visit"]
339 inputPositions = inputPositions.get(
340 parameters={
341 "columns": [
342 "xworld",
343 "yworld",
344 "xresw",
345 "yresw",
346 "exposureName",
347 "xpix",
348 "ypix",
349 "deviceName",
350 "clip",
351 "covTotalW_00",
352 "covTotalW_11",
353 ]
354 }
355 )
357 visitPositions = inputPositions[
358 (inputPositions["exposureName"] == str(visit)) & ~inputPositions["clip"]
359 ]
361 gpx, gpy, trainInd, testInd, hyperparameters = self.runGP(inputWcs, visitPositions)
363 try:
364 self.evaluate(gpx, gpy, visitPositions, trainInd, testInd, inputWcs)
365 wcsWithSpline = self.addGPToWcs(gpx, gpy, inputWcs)
366 except np.linalg.LinAlgError as e:
367 if "not positive definite" in str(e):
368 error = pipeBase.AnnotatedPartialOutputsError.annotate(
369 NotPositiveDefiniteMatrixError(trainInd.sum()),
370 self,
371 log=self.log,
372 )
373 raise error from e
374 else:
375 raise
377 return pipeBase.Struct(outputWcs=wcsWithSpline, hyperparameters=hyperparameters)
379 def runGP(self, inputWcs, positions):
380 """Run Gaussian Processes in tangent plane coordinates.
382 Parameters
383 ----------
384 inputWcs : `lsst.afw.table.ExposureCatalog`
385 Catalog with WCSs for each detector of the input exposure.
386 inputPositions : `astropy.table.Table`
387 Catalog of input positions with residuals to the current best fit.
389 Returns
390 -------
391 gpx : `treegp.gp_interp.GPInterpolation`
392 Gaussian Processes interpolator for x-direction residuals.
393 gpy : `treegp.gp_interp.GPInterpolation`
394 Gaussian Processes interpolator for y-direction residuals.
395 trainInds : `numpy.ndarray`
396 Array of indices for points used in training.
397 testInds : `numpy.ndarray`
398 Array of indices for points not used in training.
399 """
400 dx = positions["xresw"]
401 dy = positions["yresw"]
402 dxErr = positions["covTotalW_00"] ** 0.5
403 dyErr = positions["covTotalW_11"] ** 0.5
405 # Get tangent plane coordinates for input points
406 allTPCoords = np.zeros((len(positions), 2))
407 for detector in inputWcs:
408 detId = detector["id"]
409 detWCS = detector.wcs
410 detInd = positions["deviceName"].astype(int) == detId
411 detectorSources = positions[detInd]
413 tangentPlaneToSky = detWCS.getFrameDict().getMapping("PIXELS", "IWC")
414 tangentPlaneCoords = tangentPlaneToSky.applyForward(
415 np.array([detectorSources["xpix"], detectorSources["ypix"]])
416 )
417 allTPCoords[detInd] = tangentPlaneCoords.T
419 # Choose a random subset for training.
420 rng = np.random.default_rng(1234)
421 nPoints = len(allTPCoords)
422 nTrain = min([nPoints, self.config.maxTrainingPoints])
423 perm = rng.permutation(np.arange(nPoints))
424 trainInds = perm[:nTrain]
425 testInds = perm[nTrain:]
427 # Solve Gaussian Processes in dx direction.
428 gpx = treegp.GPInterpolation(
429 kernel=self.config.initKernel,
430 optimizer="anisotropic",
431 normalize=True,
432 nbins=21,
433 min_sep=self.config.correlationSeparationMin,
434 max_sep=self.config.correlationSeparationMax,
435 p0=self.config.initAnisotropicCorrelationLength,
436 )
438 gpx.initialize(allTPCoords[trainInds], dx[trainInds], y_err=dxErr[trainInds])
440 # Solve Gaussian Processes in dy direction.
441 gpy = treegp.GPInterpolation(
442 kernel=self.config.initKernel,
443 optimizer="anisotropic",
444 normalize=True,
445 nbins=21,
446 min_sep=self.config.correlationSeparationMin,
447 max_sep=self.config.correlationSeparationMax,
448 p0=self.config.initAnisotropicCorrelationLength,
449 )
451 gpy.initialize(allTPCoords[trainInds], dy[trainInds], y_err=dyErr[trainInds])
453 try:
454 gpx.solve()
455 gpy.solve()
456 except np.linalg.LinAlgError as e:
457 if "Singular matrix" in str(e): 457 ↛ 464line 457 didn't jump to line 464 because the condition on line 457 was always true
458 error = pipeBase.AnnotatedPartialOutputsError.annotate(
459 SingularMatrixError(len(allTPCoords[trainInds])),
460 self,
461 log=self.log,
462 )
463 raise error from e
464 elif "not positive definite" in str(e):
465 error = pipeBase.AnnotatedPartialOutputsError.annotate(
466 NotPositiveDefiniteMatrixError(len(allTPCoords[trainInds])),
467 self,
468 log=self.log,
469 )
470 raise error from e
471 else:
472 raise
474 hyperparameters = Table(
475 {"x": np.array(gpx._optimizer._results_robust), "y": np.array(gpy._optimizer._results_robust)}
476 )
478 return gpx, gpy, trainInds, testInds, hyperparameters
480 def predict(self, gpx, gpy, inputWcs, sourceCatalog):
481 """Get the positions for sources after correction for atmospheric
482 turbulence.
484 Parameters
485 ----------
486 gpx : `treegp.gp_interp.GPInterpolation`
487 Gaussian Processes interpolator for x-direction residuals.
488 gpy : `treegp.gp_interp.GPInterpolation`
489 Gaussian Processes interpolator for y-direction residuals.
490 inputWcs : `lsst.afw.table.ExposureCatalog`
491 Catalog with WCSs for each detector of the input exposure.
492 inputPositions : `astropy.table.Table`
493 Catalog of input positions with residuals to the current best fit.
495 Returns
496 -------
497 outCat : `astropy.table.Table`
498 Catalog matching `inputPositions`, with `coord_ra` and `coord_dec`
499 columns corrected for atmospheric turbulence.
500 """
501 correctedCoordinates = np.zeros((len(sourceCatalog), 2))
502 prediction = np.zeros((len(sourceCatalog), 2))
503 allTPCoords = np.zeros((len(sourceCatalog), 2))
504 allCoords = np.zeros((len(sourceCatalog), 2))
506 for detector in inputWcs:
507 detId = detector["id"]
508 detWCS = detector.wcs
509 detInd = sourceCatalog["detector"] == detId
510 detectorSources = sourceCatalog[detInd]
512 # The Gaussian Processes is fit on the tangent plane coordinates,
513 # so we must transform points to the tangent plane, then subtract
514 # the effect of atmospheric turbulence, then transform the tangent
515 # plane coordinates to sky coordinates.
516 initialSky = detWCS.pixelToSkyArray(detectorSources["x"], detectorSources["y"])
517 allCoords[detInd] = np.array(initialSky).T
518 tangentPlaneToSky = detWCS.getFrameDict().getMapping("IWC", "SKY")
519 tangentPlaneCoords = tangentPlaneToSky.applyInverse(np.array(initialSky)).T
520 allTPCoords[detInd] = tangentPlaneCoords
522 xPred = gpx.predict(tangentPlaneCoords)
523 xPrediction = (xPred * u.mas).to(u.degree)
524 yPred = gpy.predict(tangentPlaneCoords)
525 yPrediction = (yPred * u.mas).to(u.degree)
526 prediction[detInd, 0] = xPred
527 prediction[detInd, 1] = yPred
529 correctedTangentPlaneX = tangentPlaneCoords[:, 0] * u.degree - xPrediction
530 correctedTangentPlaneY = tangentPlaneCoords[:, 1] * u.degree - yPrediction
531 correctedSkyCoords = tangentPlaneToSky.applyForward(
532 np.array([correctedTangentPlaneX, correctedTangentPlaneY])
533 )
534 correctedCoordinates[detInd] = ((correctedSkyCoords.T) * u.radian).to(u.degree).value
536 outCat = sourceCatalog.copy()
537 outCat["coord_ra"] = correctedCoordinates[:, 0]
538 outCat["coord_dec"] = correctedCoordinates[:, 1]
540 return outCat
542 def addGPToWcs(self, gpx, gpy, inputWcs):
543 """Convert Gaussian Processes prediction to a spline, and insert it in
544 the WCS for each detector.
546 Parameters
547 ----------
548 gpx : `treegp.gp_interp.GPInterpolation`
549 Gaussian Processes interpolator for x-direction residuals.
550 gpy : `treegp.gp_interp.GPInterpolation`
551 Gaussian Processes interpolator for y-direction residuals.
552 inputWcs : `lsst.afw.table.ExposureCatalog`
553 Catalog with WCSs for each detector of the input exposure.
555 Returns
556 -------
557 catalog : `lsst.afw.table.ExposureCatalog`
558 Exposure catalog with the WCS set to the existing WCS plus the
559 gaussian processes fit.
560 """
561 pixelFrame = ast.Frame(2, "Domain=PIXELS")
562 tpFrame = ast.Frame(2, "Domain=TP")
563 iwcFrame = ast.Frame(2, "Domain=IWC")
565 # Set up the schema for the output catalogs
566 schema = lsst.afw.table.ExposureTable.makeMinimalSchema()
567 schema.addField("visit", type="L", doc="Visit number")
569 catalog = lsst.afw.table.ExposureCatalog(schema)
570 catalog.resize(len(inputWcs))
571 catalog["visit"] = inputWcs["visit"]
573 for d, detectorRow in enumerate(inputWcs):
574 detId = detectorRow.getId()
575 catalog[d].setId(detId)
577 # Make a grid of points in tangent plane coordinates.
578 bbox = detectorRow.getBBox()
579 catalog[d].setBBox(bbox)
580 corners = np.array(
581 [
582 [bbox.getBeginX(), bbox.getEndX(), bbox.getEndX(), bbox.getBeginX()],
583 [bbox.getBeginY(), bbox.getBeginY(), bbox.getEndY(), bbox.getEndY()],
584 ]
585 ).astype(float)
587 initWcsRow = inputWcs.find(detId)
588 pixToTPMap = initWcsRow.wcs.getFrameDict().getMapping("PIXELS", "IWC")
589 tpToSky = initWcsRow.wcs.getFrameDict().getMapping("IWC", "SKY")
590 skyFrame = initWcsRow.wcs.getFrameDict().getFrame("SKY")
591 tangentPlaneX, tangentPlaneY = pixToTPMap.applyForward(corners)
593 xs = np.linspace(
594 tangentPlaneX.min() - self.config.splineBuffer,
595 tangentPlaneX.max() + self.config.splineBuffer,
596 self.config.splineNNodes,
597 )
598 ys = np.linspace(
599 tangentPlaneY.min() - self.config.splineBuffer,
600 tangentPlaneY.max() + self.config.splineBuffer,
601 self.config.splineNNodes,
602 )
604 xx, yy = np.meshgrid(xs, ys)
605 inArray = np.array([xx.ravel(), yy.ravel()]).T
607 # Get Gaussian Processes prediction on grid and fit spline to it.
608 xPred = (gpx.predict(inArray) * u.mas).to(u.degree).value
610 splineX = RectBivariateSpline(
611 xs,
612 ys,
613 (xx - xPred.reshape(self.config.splineNNodes, self.config.splineNNodes)).T,
614 s=0,
615 kx=self.config.splineDegree - 1,
616 ky=self.config.splineDegree - 1,
617 )
618 (tx, ty) = splineX.get_knots()
619 coeffsX = splineX.get_coeffs()
621 yPred = (gpy.predict(inArray) * u.mas).to(u.degree).value
622 splineY = RectBivariateSpline(
623 xs,
624 ys,
625 (yy - yPred.reshape(self.config.splineNNodes, self.config.splineNNodes)).T,
626 s=0,
627 kx=self.config.splineDegree - 1,
628 ky=self.config.splineDegree - 1,
629 )
630 coeffsY = splineY.get_coeffs()
632 # Turn spline into AST object and insert in new WCS.
633 splineMap = ast.SplineMap(
634 self.config.splineDegree,
635 self.config.splineDegree,
636 self.config.splineNNodes,
637 self.config.splineNNodes,
638 tx,
639 ty,
640 coeffsX,
641 coeffsY,
642 options="OutUnit=1",
643 )
645 newFrameDict = ast.FrameDict(pixelFrame)
646 newFrameDict.addFrame("PIXELS", pixToTPMap, tpFrame)
647 newFrameDict.addFrame("TP", splineMap, iwcFrame)
648 newFrameDict.addFrame("IWC", tpToSky, skyFrame)
649 outWcs = afwgeom.SkyWcs(newFrameDict)
650 catalog[d].setWcs(outWcs)
652 return catalog
654 def evaluate(self, gpx, gpy, positions, trainInd, testInd, inputWcs, makeValidationPlot=False):
655 """Calculate E and B-modes in the 2-point correlation function before
656 and after correcting for atmospheric turbulence, and validate
657 prediction on some of the test data.
659 Parameters
660 ----------
661 gpx : `treegp.gp_interp.GPInterpolation`
662 Gaussian Processes interpolator for x-direction residuals.
663 gpy : `treegp.gp_interp.GPInterpolation`
664 Gaussian Processes interpolator for y-direction residuals.
665 positions : `astropy.table.Table`
666 Catalog of input positions with residuals to the best fit.
667 trainInds : `numpy.ndarray`
668 Array of indices for points used in training.
669 testInds : `numpy.ndarray`
670 Array of indices for points not used in training.
671 inputWcs : `lsst.afw.table.ExposureCatalog`
672 Catalog with WCSs for each detector of the input exposure.
673 makeValidationPlot : `bool`, optional
674 Whether to make a plot showing the prediction on the validation
675 data.
676 """
677 dx = positions["xresw"]
678 dy = positions["yresw"]
680 # Get tangent plane coordinates for input points
681 tpCoords = np.zeros((len(positions), 2))
682 for detector in inputWcs:
683 detId = detector["id"]
684 detWCS = detector.wcs
685 detInd = positions["deviceName"].astype(int) == detId
686 detectorSources = positions[detInd]
688 tangentPlaneToSky = detWCS.getFrameDict().getMapping("PIXELS", "IWC")
689 tangentPlaneCoords = tangentPlaneToSky.applyForward(
690 np.array([detectorSources["xpix"], detectorSources["ypix"]])
691 )
692 tpCoords[detInd] = tangentPlaneCoords.T
694 # Calculate E/B modes before and after Gaussian Processes correction.
695 xPredict = gpx.predict(tpCoords[trainInd])
696 yPredict = gpy.predict(tpCoords[trainInd])
697 xie, xib, logr = treegp.comp_eb_treecorr(
698 tpCoords[trainInd, 0],
699 tpCoords[trainInd, 1],
700 dx[trainInd],
701 dy[trainInd],
702 rmin=20 / 3600,
703 rmax=0.6,
704 dlogr=0.3,
705 )
706 start, stop = np.searchsorted(np.exp(logr), [0, 15])
707 meanE = np.mean(xie[start:stop])
708 meanB = np.mean(xib[start:stop])
709 self.log.info(
710 "Original average correlation level over 0-15 arcminutes: E-mode=%0.2f, B-mode=%0.2f",
711 meanE,
712 meanB,
713 )
715 xie_resid, xib_resid, logr = treegp.comp_eb_treecorr(
716 tpCoords[trainInd, 0],
717 tpCoords[trainInd, 1],
718 dx[trainInd] - xPredict,
719 dy[trainInd] - yPredict,
720 rmin=20 / 3600,
721 rmax=0.6,
722 dlogr=0.3,
723 )
724 start, stop = np.searchsorted(np.exp(logr), [0, 15])
725 meanE_resid = np.mean(xie_resid[start:stop])
726 meanB_resid = np.mean(xib_resid[start:stop])
727 self.log.info(
728 "Correlation level after GP correction over 0-15 arcminutes: E-mode=%0.2f, B-mode=%0.2f",
729 meanE_resid,
730 meanB_resid,
731 )
733 # Predict on all test data and make a plot.
734 if makeValidationPlot:
735 print(len(testInd))
736 testInd = testInd[:50000]
737 chunkSize = 5000
738 nChunks = np.ceil(len(testInd) / chunkSize).astype(int)
739 xPredict = np.zeros(len(testInd))
740 yPredict = np.zeros(len(testInd))
741 for i in range(nChunks):
742 ind = testInd[chunkSize * i : chunkSize * (i + 1)]
743 xPredict[chunkSize * i : chunkSize * (i + 1)] = gpx.predict(tpCoords[ind])
744 yPredict[chunkSize * i : chunkSize * (i + 1)] = gpy.predict(tpCoords[ind])
745 fig = plot_visit(
746 tpCoords[testInd, 0], tpCoords[testInd, 1], dx[testInd], dy[testInd], xPredict, yPredict
747 )
748 return fig