Coverage for python/lsst/drp/tasks/fit_turbulence.py: 58%
251 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-08 02:12 -0700
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 self.evaluate(gpx, gpy, visitPositions, trainInd, testInd, inputWcs)
365 wcsWithSpline = self.addGPToWcs(gpx, gpy, inputWcs)
367 return pipeBase.Struct(outputWcs=wcsWithSpline, hyperparameters=hyperparameters)
369 def runGP(self, inputWcs, positions):
370 """Run Gaussian Processes in tangent plane coordinates.
372 Parameters
373 ----------
374 inputWcs : `lsst.afw.table.ExposureCatalog`
375 Catalog with WCSs for each detector of the input exposure.
376 inputPositions : `astropy.table.Table`
377 Catalog of input positions with residuals to the current best fit.
379 Returns
380 -------
381 gpx : `treegp.gp_interp.GPInterpolation`
382 Gaussian Processes interpolator for x-direction residuals.
383 gpy : `treegp.gp_interp.GPInterpolation`
384 Gaussian Processes interpolator for y-direction residuals.
385 trainInds : `numpy.ndarray`
386 Array of indices for points used in training.
387 testInds : `numpy.ndarray`
388 Array of indices for points not used in training.
389 """
390 dx = positions["xresw"]
391 dy = positions["yresw"]
392 dxErr = positions["covTotalW_00"] ** 0.5
393 dyErr = positions["covTotalW_11"] ** 0.5
395 # Get tangent plane coordinates for input points
396 allTPCoords = np.zeros((len(positions), 2))
397 for detector in inputWcs:
398 detId = detector["id"]
399 detWCS = detector.wcs
400 detInd = positions["deviceName"].astype(int) == detId
401 detectorSources = positions[detInd]
403 tangentPlaneToSky = detWCS.getFrameDict().getMapping("PIXELS", "IWC")
404 tangentPlaneCoords = tangentPlaneToSky.applyForward(
405 np.array([detectorSources["xpix"], detectorSources["ypix"]])
406 )
407 allTPCoords[detInd] = tangentPlaneCoords.T
409 # Choose a random subset for training.
410 rng = np.random.default_rng(1234)
411 nPoints = len(allTPCoords)
412 nTrain = min([nPoints, self.config.maxTrainingPoints])
413 perm = rng.permutation(np.arange(nPoints))
414 trainInds = perm[:nTrain]
415 testInds = perm[nTrain:]
417 # Solve Gaussian Processes in dx direction.
418 gpx = treegp.GPInterpolation(
419 kernel=self.config.initKernel,
420 optimizer="anisotropic",
421 normalize=True,
422 nbins=21,
423 min_sep=self.config.correlationSeparationMin,
424 max_sep=self.config.correlationSeparationMax,
425 p0=self.config.initAnisotropicCorrelationLength,
426 )
428 gpx.initialize(allTPCoords[trainInds], dx[trainInds], y_err=dxErr[trainInds])
430 # Solve Gaussian Processes in dy direction.
431 gpy = treegp.GPInterpolation(
432 kernel=self.config.initKernel,
433 optimizer="anisotropic",
434 normalize=True,
435 nbins=21,
436 min_sep=self.config.correlationSeparationMin,
437 max_sep=self.config.correlationSeparationMax,
438 p0=self.config.initAnisotropicCorrelationLength,
439 )
441 gpy.initialize(allTPCoords[trainInds], dy[trainInds], y_err=dyErr[trainInds])
443 try:
444 gpx.solve()
445 gpy.solve()
446 except np.linalg.LinAlgError as e:
447 if "Singular matrix" in str(e): 447 ↛ 454line 447 didn't jump to line 454 because the condition on line 447 was always true
448 error = pipeBase.AnnotatedPartialOutputsError.annotate(
449 SingularMatrixError(len(allTPCoords[trainInds])),
450 self,
451 log=self.log,
452 )
453 raise error from e
454 elif "not positive definite" in str(e):
455 error = pipeBase.AnnotatedPartialOutputsError.annotate(
456 NotPositiveDefiniteMatrixError(len(allTPCoords[trainInds])),
457 self,
458 log=self.log,
459 )
460 raise error from e
461 else:
462 raise
464 hyperparameters = Table(
465 {"x": np.array(gpx._optimizer._results_robust), "y": np.array(gpy._optimizer._results_robust)}
466 )
468 return gpx, gpy, trainInds, testInds, hyperparameters
470 def predict(self, gpx, gpy, inputWcs, sourceCatalog):
471 """Get the positions for sources after correction for atmospheric
472 turbulence.
474 Parameters
475 ----------
476 gpx : `treegp.gp_interp.GPInterpolation`
477 Gaussian Processes interpolator for x-direction residuals.
478 gpy : `treegp.gp_interp.GPInterpolation`
479 Gaussian Processes interpolator for y-direction residuals.
480 inputWcs : `lsst.afw.table.ExposureCatalog`
481 Catalog with WCSs for each detector of the input exposure.
482 inputPositions : `astropy.table.Table`
483 Catalog of input positions with residuals to the current best fit.
485 Returns
486 -------
487 outCat : `astropy.table.Table`
488 Catalog matching `inputPositions`, with `coord_ra` and `coord_dec`
489 columns corrected for atmospheric turbulence.
490 """
491 correctedCoordinates = np.zeros((len(sourceCatalog), 2))
492 prediction = np.zeros((len(sourceCatalog), 2))
493 allTPCoords = np.zeros((len(sourceCatalog), 2))
494 allCoords = np.zeros((len(sourceCatalog), 2))
496 for detector in inputWcs:
497 detId = detector["id"]
498 detWCS = detector.wcs
499 detInd = sourceCatalog["detector"] == detId
500 detectorSources = sourceCatalog[detInd]
502 # The Gaussian Processes is fit on the tangent plane coordinates,
503 # so we must transform points to the tangent plane, then subtract
504 # the effect of atmospheric turbulence, then transform the tangent
505 # plane coordinates to sky coordinates.
506 initialSky = detWCS.pixelToSkyArray(detectorSources["x"], detectorSources["y"])
507 allCoords[detInd] = np.array(initialSky).T
508 tangentPlaneToSky = detWCS.getFrameDict().getMapping("IWC", "SKY")
509 tangentPlaneCoords = tangentPlaneToSky.applyInverse(np.array(initialSky)).T
510 allTPCoords[detInd] = tangentPlaneCoords
512 xPred = gpx.predict(tangentPlaneCoords)
513 xPrediction = (xPred * u.mas).to(u.degree)
514 yPred = gpy.predict(tangentPlaneCoords)
515 yPrediction = (yPred * u.mas).to(u.degree)
516 prediction[detInd, 0] = xPred
517 prediction[detInd, 1] = yPred
519 correctedTangentPlaneX = tangentPlaneCoords[:, 0] * u.degree - xPrediction
520 correctedTangentPlaneY = tangentPlaneCoords[:, 1] * u.degree - yPrediction
521 correctedSkyCoords = tangentPlaneToSky.applyForward(
522 np.array([correctedTangentPlaneX, correctedTangentPlaneY])
523 )
524 correctedCoordinates[detInd] = ((correctedSkyCoords.T) * u.radian).to(u.degree).value
526 outCat = sourceCatalog.copy()
527 outCat["coord_ra"] = correctedCoordinates[:, 0]
528 outCat["coord_dec"] = correctedCoordinates[:, 1]
530 return outCat
532 def addGPToWcs(self, gpx, gpy, inputWcs):
533 """Convert Gaussian Processes prediction to a spline, and insert it in
534 the WCS for each detector.
536 Parameters
537 ----------
538 gpx : `treegp.gp_interp.GPInterpolation`
539 Gaussian Processes interpolator for x-direction residuals.
540 gpy : `treegp.gp_interp.GPInterpolation`
541 Gaussian Processes interpolator for y-direction residuals.
542 inputWcs : `lsst.afw.table.ExposureCatalog`
543 Catalog with WCSs for each detector of the input exposure.
545 Returns
546 -------
547 catalog : `lsst.afw.table.ExposureCatalog`
548 Exposure catalog with the WCS set to the existing WCS plus the
549 gaussian processes fit.
550 """
551 pixelFrame = ast.Frame(2, "Domain=PIXELS")
552 tpFrame = ast.Frame(2, "Domain=TP")
553 iwcFrame = ast.Frame(2, "Domain=IWC")
555 # Set up the schema for the output catalogs
556 schema = lsst.afw.table.ExposureTable.makeMinimalSchema()
557 schema.addField("visit", type="L", doc="Visit number")
559 catalog = lsst.afw.table.ExposureCatalog(schema)
560 catalog.resize(len(inputWcs))
561 catalog["visit"] = inputWcs["visit"]
563 for d, detectorRow in enumerate(inputWcs):
564 detId = detectorRow.getId()
565 catalog[d].setId(detId)
567 # Make a grid of points in tangent plane coordinates.
568 bbox = detectorRow.getBBox()
569 catalog[d].setBBox(bbox)
570 corners = np.array(
571 [
572 [bbox.getBeginX(), bbox.getEndX(), bbox.getEndX(), bbox.getBeginX()],
573 [bbox.getBeginY(), bbox.getBeginY(), bbox.getEndY(), bbox.getEndY()],
574 ]
575 ).astype(float)
577 initWcsRow = inputWcs.find(detId)
578 pixToTPMap = initWcsRow.wcs.getFrameDict().getMapping("PIXELS", "IWC")
579 tpToSky = initWcsRow.wcs.getFrameDict().getMapping("IWC", "SKY")
580 skyFrame = initWcsRow.wcs.getFrameDict().getFrame("SKY")
581 tangentPlaneX, tangentPlaneY = pixToTPMap.applyForward(corners)
583 xs = np.linspace(
584 tangentPlaneX.min() - self.config.splineBuffer,
585 tangentPlaneX.max() + self.config.splineBuffer,
586 self.config.splineNNodes,
587 )
588 ys = np.linspace(
589 tangentPlaneY.min() - self.config.splineBuffer,
590 tangentPlaneY.max() + self.config.splineBuffer,
591 self.config.splineNNodes,
592 )
594 xx, yy = np.meshgrid(xs, ys)
595 inArray = np.array([xx.ravel(), yy.ravel()]).T
597 # Get Gaussian Processes prediction on grid and fit spline to it.
598 xPred = (gpx.predict(inArray) * u.mas).to(u.degree).value
600 splineX = RectBivariateSpline(
601 xs,
602 ys,
603 (xx - xPred.reshape(self.config.splineNNodes, self.config.splineNNodes)).T,
604 s=0,
605 kx=self.config.splineDegree - 1,
606 ky=self.config.splineDegree - 1,
607 )
608 (tx, ty) = splineX.get_knots()
609 coeffsX = splineX.get_coeffs()
611 yPred = (gpy.predict(inArray) * u.mas).to(u.degree).value
612 splineY = RectBivariateSpline(
613 xs,
614 ys,
615 (yy - yPred.reshape(self.config.splineNNodes, self.config.splineNNodes)).T,
616 s=0,
617 kx=self.config.splineDegree - 1,
618 ky=self.config.splineDegree - 1,
619 )
620 coeffsY = splineY.get_coeffs()
622 # Turn spline into AST object and insert in new WCS.
623 splineMap = ast.SplineMap(
624 self.config.splineDegree,
625 self.config.splineDegree,
626 self.config.splineNNodes,
627 self.config.splineNNodes,
628 tx,
629 ty,
630 coeffsX,
631 coeffsY,
632 options="OutUnit=1",
633 )
635 newFrameDict = ast.FrameDict(pixelFrame)
636 newFrameDict.addFrame("PIXELS", pixToTPMap, tpFrame)
637 newFrameDict.addFrame("TP", splineMap, iwcFrame)
638 newFrameDict.addFrame("IWC", tpToSky, skyFrame)
639 outWcs = afwgeom.SkyWcs(newFrameDict)
640 catalog[d].setWcs(outWcs)
642 return catalog
644 def evaluate(self, gpx, gpy, positions, trainInd, testInd, inputWcs, makeValidationPlot=False):
645 """Calculate E and B-modes in the 2-point correlation function before
646 and after correcting for atmospheric turbulence, and validate
647 prediction on some of the test data.
649 Parameters
650 ----------
651 gpx : `treegp.gp_interp.GPInterpolation`
652 Gaussian Processes interpolator for x-direction residuals.
653 gpy : `treegp.gp_interp.GPInterpolation`
654 Gaussian Processes interpolator for y-direction residuals.
655 positions : `astropy.table.Table`
656 Catalog of input positions with residuals to the best fit.
657 trainInds : `numpy.ndarray`
658 Array of indices for points used in training.
659 testInds : `numpy.ndarray`
660 Array of indices for points not used in training.
661 inputWcs : `lsst.afw.table.ExposureCatalog`
662 Catalog with WCSs for each detector of the input exposure.
663 makeValidationPlot : `bool`, optional
664 Whether to make a plot showing the prediction on the validation
665 data.
666 """
667 dx = positions["xresw"]
668 dy = positions["yresw"]
670 # Get tangent plane coordinates for input points
671 tpCoords = np.zeros((len(positions), 2))
672 for detector in inputWcs:
673 detId = detector["id"]
674 detWCS = detector.wcs
675 detInd = positions["deviceName"].astype(int) == detId
676 detectorSources = positions[detInd]
678 tangentPlaneToSky = detWCS.getFrameDict().getMapping("PIXELS", "IWC")
679 tangentPlaneCoords = tangentPlaneToSky.applyForward(
680 np.array([detectorSources["xpix"], detectorSources["ypix"]])
681 )
682 tpCoords[detInd] = tangentPlaneCoords.T
684 # Calculate E/B modes before and after Gaussian Processes correction.
685 xPredict = gpx.predict(tpCoords[trainInd])
686 yPredict = gpy.predict(tpCoords[trainInd])
687 xie, xib, logr = treegp.comp_eb_treecorr(
688 tpCoords[trainInd, 0],
689 tpCoords[trainInd, 1],
690 dx[trainInd],
691 dy[trainInd],
692 rmin=20 / 3600,
693 rmax=0.6,
694 dlogr=0.3,
695 )
696 start, stop = np.searchsorted(np.exp(logr), [0, 15])
697 meanE = np.mean(xie[start:stop])
698 meanB = np.mean(xib[start:stop])
699 self.log.info(
700 "Original average correlation level over 0-15 arcminutes: E-mode=%0.2f, B-mode=%0.2f",
701 meanE,
702 meanB,
703 )
705 xie_resid, xib_resid, logr = treegp.comp_eb_treecorr(
706 tpCoords[trainInd, 0],
707 tpCoords[trainInd, 1],
708 dx[trainInd] - xPredict,
709 dy[trainInd] - yPredict,
710 rmin=20 / 3600,
711 rmax=0.6,
712 dlogr=0.3,
713 )
714 start, stop = np.searchsorted(np.exp(logr), [0, 15])
715 meanE_resid = np.mean(xie_resid[start:stop])
716 meanB_resid = np.mean(xib_resid[start:stop])
717 self.log.info(
718 "Correlation level after GP correction over 0-15 arcminutes: E-mode=%0.2f, B-mode=%0.2f",
719 meanE_resid,
720 meanB_resid,
721 )
723 # Predict on all test data and make a plot.
724 if makeValidationPlot:
725 print(len(testInd))
726 testInd = testInd[:50000]
727 chunkSize = 5000
728 nChunks = np.ceil(len(testInd) / chunkSize).astype(int)
729 xPredict = np.zeros(len(testInd))
730 yPredict = np.zeros(len(testInd))
731 for i in range(nChunks):
732 ind = testInd[chunkSize * i : chunkSize * (i + 1)]
733 xPredict[chunkSize * i : chunkSize * (i + 1)] = gpx.predict(tpCoords[ind])
734 yPredict[chunkSize * i : chunkSize * (i + 1)] = gpy.predict(tpCoords[ind])
735 fig = plot_visit(
736 tpCoords[testInd, 0], tpCoords[testInd, 1], dx[testInd], dy[testInd], xPredict, yPredict
737 )
738 return fig