Coverage for python / lsst / drp / tasks / fit_turbulence.py: 16%
221 statements
« prev ^ index » next coverage.py v7.12.0, created at 2025-12-12 01:55 -0800
1# This file is part of drp_tasks.
2#
3# LSST Data Management System
4# This product includes software developed by the
5# LSST Project (http://www.lsst.org/).
6# See COPYRIGHT file at the top of the source tree.
7#
8# This program is free software: you can redistribute it and/or modify
9# it under the terms of the GNU General Public License as published by
10# the Free Software Foundation, either version 3 of the License, or
11# (at your option) any later version.
12#
13# This program is distributed in the hope that it will be useful,
14# but WITHOUT ANY WARRANTY; without even the implied warranty of
15# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
16# GNU General Public License for more details.
17#
18# You should have received a copy of the LSST License Statement and
19# the GNU General Public License along with this program. If not,
20# see <https://www.lsstcorp.org/LegalNotices/>.
21#
22import dataclasses
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
32import lsst.afw.geom as afwgeom
33import lsst.afw.table
34import lsst.pex.config as pexConfig
35import lsst.pipe.base as pipeBase
37__all__ = [
38 "GaussianProcessesTurbulenceFitConnections",
39 "GaussianProcessesTurbulenceFitConfig",
40 "GaussianProcessesTurbulenceFitTask",
41]
44def plot_visit(x, y, dx, dy, predx, predy):
45 """Utility function for plotting Gaussian Processes results.
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.
62 Returns
63 -------
64 fig : `matplotlib.pyplot.Figure`
65 Figure showing input data, Gaussian Processes prediction, and E and
66 B-modes.
67 """
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 )
74 residualLimit = np.nanstd(dx)
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)
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)
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 )
95 cb = fig.colorbar(
96 im, ax=[subs["dx"], subs["dy"], subs["predx"], subs["predy"], subs["residx"], subs["residy"]]
97 )
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")
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)
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)")
123 subs["dx"].set_title(r"$\delta$x")
124 subs["predx"].set_title("GP prediction")
125 subs["residx"].set_title("Residual")
127 subs["dy"].set_title(r"$\delta$y")
128 subs["predy"].set_title("GP prediction")
129 subs["residy"].set_title("Residual")
131 cb.set_label("mas")
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)")
137 return fig
140class GaussianProcessesTurbulenceFitConnections(
141 pipeBase.PipelineTaskConnections,
142 dimensions=("instrument", "visit", "healpix3"),
143 defaultTemplates={
144 "inputName": "gbdesHealpix3AstrometricFit",
145 },
146):
147 inputWcs = pipeBase.connectionTypes.Input(
148 doc=(
149 "Per-healpix, per-visit world coordinate systems derived from the fitted model."
150 " These catalogs only contain entries for detectors with an output, and use"
151 " the detector id for the catalog id, sorted on id for fast lookups of a detector."
152 ),
153 name="{inputName}SkyWcsCatalog",
154 storageClass="ExposureCatalog",
155 dimensions=("instrument", "visit", "healpix3"),
156 )
157 inputPositions = pipeBase.connectionTypes.Input(
158 doc=(
159 "Catalog of sources used in fit, along with residuals in pixel coordinates and tangent "
160 "plane coordinates and chisq values."
161 ),
162 name="{inputName}_fitStars",
163 storageClass="ArrowAstropy",
164 dimensions=("instrument", "healpix3", "physical_filter"),
165 deferLoad=True,
166 )
167 outputWcs = pipeBase.connectionTypes.Output(
168 doc=(
169 "Per-visit world coordinate systems derived from the fitted model. These catalogs only contain "
170 "entries for detectors with an output, and use the detector id for the catalog id, sorted on id "
171 "for fast lookups of a detector."
172 ),
173 name="turbulenceCorrectedSkyWcsCatalog",
174 storageClass="ExposureCatalog",
175 dimensions=("instrument", "visit", "healpix3"),
176 )
177 hyperparameters = pipeBase.connectionTypes.Output(
178 doc="Best fit hyperparameters for the Gaussian Processes fit.",
179 name="turbulence_fit_hyperparameters",
180 storageClass="ArrowAstropy",
181 dimensions=("instrument", "visit", "healpix3"),
182 )
184 def __init__(self, *, config=None):
185 super().__init__(config=config)
187 if not self.config.healpix:
188 self.dimensions.remove("healpix3")
189 self.dimensions.add("tract")
190 self.dimensions.add("skymap")
191 self.inputWcs = dataclasses.replace(
192 self.inputWcs, dimensions=("instrument", "visit", "tract", "skymap")
193 )
194 self.inputPositions = dataclasses.replace(
195 self.inputPositions, dimensions=("instrument", "tract", "skymap", "band", "physical_filter")
196 )
199class GaussianProcessesTurbulenceFitConfig(
200 pipeBase.PipelineTaskConfig, pipelineConnections=GaussianProcessesTurbulenceFitConnections
201):
202 initKernel = pexConfig.Field(
203 dtype=str,
204 doc="The type of function that will be used to modeled spatial correlation.",
205 default="15**2 * AnisotropicVonKarman(invLam=array([[1./0.8**2,0],[0,1./0.8**2]]))",
206 )
207 initAnisotropicCorrelationLength = pexConfig.ListField(
208 dtype=float,
209 doc=(
210 "The initial parameters for fiting the anisotropic correlation length. p0[0] is equivalent of "
211 "the isotropic correlation length in degrees, and p0[1]/p0[2] are ellipticity parameters and are "
212 "mathematically equivalent to e1/e2 in weak-lensing. p0[1]/p0[2] must be in the range [-1,1], "
213 "where 0 means the correlation is isotropic."
214 ),
215 default=[1, -0.2, -0.2],
216 )
217 correlationSeparationMin = pexConfig.Field(
218 dtype=float,
219 doc="Minimum distance separation in degrees in the computation of the 2-point correlation function.",
220 default=0.0,
221 optional=True,
222 )
223 correlationSeparationMax = pexConfig.Field(
224 dtype=float,
225 doc="Maximum distance separation in degrees in the computation of the 2-point correlation function.",
226 default=0.3,
227 optional=True,
228 )
229 maxTrainingPoints = pexConfig.Field(
230 dtype=int,
231 doc="Maximum number of points to use in the Gaussian Processes training.",
232 default=10000,
233 )
234 pixelSize = pexConfig.Field(
235 dtype=float,
236 doc="Pixel size in arcseconds.",
237 default=0.2,
238 )
239 healpix = pexConfig.Field(
240 dtype=bool,
241 doc="Use input WCS calculated over healpix-based region. If false, use tract-based WCS.",
242 default=True,
243 )
244 splineDegree = pexConfig.Field(
245 dtype=int,
246 doc="Degree of the spline expressing Gaussian Processes prediction.",
247 default=4,
248 )
249 splineNNodes = pexConfig.Field(
250 dtype=int,
251 doc="Number of nodes to use for the spline expressing Gaussian Processes prediction.",
252 default=30,
253 )
256class GaussianProcessesTurbulenceFitTask(pipeBase.PipelineTask):
257 """Run Gaussian Processes on astrometric residuals with the assumption that
258 they are due to atmospheric turbulence."""
260 ConfigClass = GaussianProcessesTurbulenceFitConfig
261 _DefaultName = "gaussianProcessesTurbulenceFit"
263 def run(self, inputWcs, inputPositions):
264 """Run Gaussian Processes on position residuals and subtract the fitted
265 Gaussian Processes prediction from the WCS to account for atmospheric
266 turbulence.
268 Parameters
269 ----------
270 inputWcs : `lsst.afw.table.ExposureCatalog`
271 Catalog with WCSs for each detector of the input exposure.
272 inputPositions : `astropy.table.Table`
273 Catalog of input positions with residuals to the current best fit.
275 Returns
276 -------
277 result : `lsst.pipe.base.Struct`
278 ``outputWcs`` : `lsst.afw.table.ExposureCatalog`
279 Catalog with WCS after inserting the correction for atmospheric
280 turbulence.
281 ``hyperparameters`` : `astropy.table.Table`
282 Table of best-fit hyperparameters in x and y-directions.
283 """
285 visit = inputWcs[0]["visit"]
287 inputPositions = inputPositions.get(
288 parameters={
289 "columns": [
290 "xworld",
291 "yworld",
292 "xresw",
293 "yresw",
294 "exposureName",
295 "xpix",
296 "ypix",
297 "deviceName",
298 "clip",
299 "covTotalW_00",
300 "covTotalW_11",
301 ]
302 }
303 )
305 visitPositions = inputPositions[
306 (inputPositions["exposureName"] == str(visit)) & ~inputPositions["clip"]
307 ]
309 gpx, gpy, trainInd, testInd, hyperparameters = self.runGP(inputWcs, visitPositions)
311 self.evaluate(gpx, gpy, visitPositions, trainInd, testInd, inputWcs)
313 wcsWithSpline = self.addGPToWcs(gpx, gpy, inputWcs)
315 return pipeBase.Struct(outputWcs=wcsWithSpline, hyperparameters=hyperparameters)
317 def runGP(self, inputWcs, positions):
318 """Run Gaussian Processes in tangent plane coordinates.
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 gpx : `treegp.gp_interp.GPInterpolation`
330 Gaussian Processes interpolator for x-direction residuals.
331 gpy : `treegp.gp_interp.GPInterpolation`
332 Gaussian Processes interpolator for y-direction residuals.
333 trainInds : `numpy.ndarray`
334 Array of indices for points used in training.
335 testInds : `numpy.ndarray`
336 Array of indices for points not used in training.
337 """
338 dx = positions["xresw"]
339 dy = positions["yresw"]
340 dxErr = positions["covTotalW_00"] ** 0.5
341 dyErr = positions["covTotalW_11"] ** 0.5
343 # Get tangent plane coordinates for input points
344 allTPCoords = np.zeros((len(positions), 2))
345 for detector in inputWcs:
346 detId = detector["id"]
347 detWCS = detector.wcs
348 detInd = positions["deviceName"].astype(int) == detId
349 detectorSources = positions[detInd]
351 tangentPlaneToSky = detWCS.getFrameDict().getMapping("PIXELS", "IWC")
352 tangentPlaneCoords = tangentPlaneToSky.applyForward(
353 np.array([detectorSources["xpix"], detectorSources["ypix"]])
354 )
355 allTPCoords[detInd] = tangentPlaneCoords.T
357 # Choose a random subset for training.
358 rng = np.random.default_rng(1234)
359 nPoints = len(allTPCoords)
360 nTrain = min([nPoints, self.config.maxTrainingPoints])
361 perm = rng.permutation(np.arange(nPoints))
362 trainInds = perm[:nTrain]
363 testInds = perm[nTrain:]
365 # Solve Gaussian Processes in dx direction.
366 gpx = treegp.GPInterpolation(
367 kernel=self.config.initKernel,
368 optimizer="anisotropic",
369 normalize=True,
370 nbins=21,
371 min_sep=self.config.correlationSeparationMin,
372 max_sep=self.config.correlationSeparationMax,
373 p0=self.config.initAnisotropicCorrelationLength,
374 )
376 gpx.initialize(allTPCoords[trainInds], dx[trainInds], y_err=dxErr[trainInds])
377 gpx.solve()
379 # Solve Gaussian Processes in dy direction.
380 gpy = treegp.GPInterpolation(
381 kernel=self.config.initKernel,
382 optimizer="anisotropic",
383 normalize=True,
384 nbins=21,
385 min_sep=self.config.correlationSeparationMin,
386 max_sep=self.config.correlationSeparationMax,
387 p0=self.config.initAnisotropicCorrelationLength,
388 )
390 gpy.initialize(allTPCoords[trainInds], dy[trainInds], y_err=dyErr[trainInds])
391 gpy.solve()
393 hyperparameters = Table(
394 {"x": np.array(gpx._optimizer._results_robust), "y": np.array(gpy._optimizer._results_robust)}
395 )
397 return gpx, gpy, trainInds, testInds, hyperparameters
399 def predict(self, gpx, gpy, inputWcs, sourceCatalog):
400 """Get the positions for sources after correction for atmospheric
401 turbulence.
403 Parameters
404 ----------
405 gpx : `treegp.gp_interp.GPInterpolation`
406 Gaussian Processes interpolator for x-direction residuals.
407 gpy : `treegp.gp_interp.GPInterpolation`
408 Gaussian Processes interpolator for y-direction residuals.
409 inputWcs : `lsst.afw.table.ExposureCatalog`
410 Catalog with WCSs for each detector of the input exposure.
411 inputPositions : `astropy.table.Table`
412 Catalog of input positions with residuals to the current best fit.
414 Returns
415 -------
416 outCat : `astropy.table.Table`
417 Catalog matching `inputPositions`, with `coord_ra` and `coord_dec`
418 columns corrected for atmospheric turbulence.
419 """
420 correctedCoordinates = np.zeros((len(sourceCatalog), 2))
421 prediction = np.zeros((len(sourceCatalog), 2))
422 allTPCoords = np.zeros((len(sourceCatalog), 2))
423 allCoords = np.zeros((len(sourceCatalog), 2))
425 for detector in inputWcs:
426 detId = detector["id"]
427 detWCS = detector.wcs
428 detInd = sourceCatalog["detector"] == detId
429 detectorSources = sourceCatalog[detInd]
431 # The Gaussian Processes is fit on the tangent plane coordinates,
432 # so we must transform points to the tangent plane, then subtract
433 # the effect of atmospheric turbulence, then transform the tangent
434 # plane coordinates to sky coordinates.
435 initialSky = detWCS.pixelToSkyArray(detectorSources["x"], detectorSources["y"])
436 allCoords[detInd] = np.array(initialSky).T
437 tangentPlaneToSky = detWCS.getFrameDict().getMapping("IWC", "SKY")
438 tangentPlaneCoords = tangentPlaneToSky.applyInverse(np.array(initialSky)).T
439 allTPCoords[detInd] = tangentPlaneCoords
441 xPred = gpx.predict(tangentPlaneCoords)
442 xPrediction = (xPred * u.mas).to(u.degree)
443 yPred = gpy.predict(tangentPlaneCoords)
444 yPrediction = (yPred * u.mas).to(u.degree)
445 prediction[detInd, 0] = xPred
446 prediction[detInd, 1] = yPred
448 correctedTangentPlaneX = tangentPlaneCoords[:, 0] * u.degree - xPrediction
449 correctedTangentPlaneY = tangentPlaneCoords[:, 1] * u.degree - yPrediction
450 correctedSkyCoords = tangentPlaneToSky.applyForward(
451 np.array([correctedTangentPlaneX, correctedTangentPlaneY])
452 )
453 correctedCoordinates[detInd] = ((correctedSkyCoords.T) * u.radian).to(u.degree).value
455 outCat = sourceCatalog.copy()
456 outCat["coord_ra"] = correctedCoordinates[:, 0]
457 outCat["coord_dec"] = correctedCoordinates[:, 1]
459 return outCat
461 def addGPToWcs(self, gpx, gpy, inputWcs):
462 """Convert Gaussian Processes prediction to a spline, and insert it in
463 the WCS for each detector.
465 Parameters
466 ----------
467 gpx : `treegp.gp_interp.GPInterpolation`
468 Gaussian Processes interpolator for x-direction residuals.
469 gpy : `treegp.gp_interp.GPInterpolation`
470 Gaussian Processes interpolator for y-direction residuals.
471 inputWcs : `lsst.afw.table.ExposureCatalog`
472 Catalog with WCSs for each detector of the input exposure.
474 Returns
475 -------
476 catalog : `lsst.afw.table.ExposureCatalog`
477 Exposure catalog with the WCS set to the existing WCS plus the
478 gaussian processes fit.
479 """
480 pixelFrame = ast.Frame(2, "Domain=PIXELS")
481 tpFrame = ast.Frame(2, "Domain=TP")
482 iwcFrame = ast.Frame(2, "Domain=IWC")
484 # Set up the schema for the output catalogs
485 schema = lsst.afw.table.ExposureTable.makeMinimalSchema()
486 schema.addField("visit", type="L", doc="Visit number")
488 catalog = lsst.afw.table.ExposureCatalog(schema)
489 catalog.resize(len(inputWcs))
490 catalog["visit"] = inputWcs["visit"]
492 for d, detectorRow in enumerate(inputWcs):
493 detId = detectorRow.getId()
494 catalog[d].setId(detId)
496 # Make a grid of points in tangent plane coordinates.
497 bbox = detectorRow.getBBox()
498 catalog[d].setBBox(bbox)
499 corners = np.array(
500 [
501 [bbox.getBeginX(), bbox.getEndX(), bbox.getEndX(), bbox.getBeginX()],
502 [bbox.getBeginY(), bbox.getBeginY(), bbox.getEndY(), bbox.getEndY()],
503 ]
504 ).astype(float)
506 initWcsRow = inputWcs.find(detId)
507 pixToTPMap = initWcsRow.wcs.getFrameDict().getMapping("PIXELS", "IWC")
508 tpToSky = initWcsRow.wcs.getFrameDict().getMapping("IWC", "SKY")
509 skyFrame = initWcsRow.wcs.getFrameDict().getFrame("SKY")
510 tangentPlaneX, tangentPlaneY = pixToTPMap.applyForward(corners)
512 xs = np.linspace(tangentPlaneX.min(), tangentPlaneX.max(), self.config.splineNNodes)
513 ys = np.linspace(tangentPlaneY.min(), tangentPlaneY.max(), self.config.splineNNodes)
515 xx, yy = np.meshgrid(xs, ys)
516 inArray = np.array([xx.ravel(), yy.ravel()]).T
518 # Get Gaussian Processes prediction on grid and fit spline to it.
519 xPred = (gpx.predict(inArray) * u.mas).to(u.degree).value
521 splineX = RectBivariateSpline(
522 xs,
523 ys,
524 (xx - xPred.reshape(self.config.splineNNodes, self.config.splineNNodes)).T,
525 s=0,
526 kx=self.config.splineDegree - 1,
527 ky=self.config.splineDegree - 1,
528 )
529 (tx, ty) = splineX.get_knots()
530 coeffsX = splineX.get_coeffs()
532 yPred = (gpy.predict(inArray) * u.mas).to(u.degree).value
533 splineY = RectBivariateSpline(
534 xs,
535 ys,
536 (yy - yPred.reshape(self.config.splineNNodes, self.config.splineNNodes)).T,
537 s=0,
538 kx=self.config.splineDegree - 1,
539 ky=self.config.splineDegree - 1,
540 )
541 coeffsY = splineY.get_coeffs()
543 # Turn spline into AST object and insert in new WCS.
544 splineMap = ast.SplineMap(
545 self.config.splineDegree,
546 self.config.splineDegree,
547 self.config.splineNNodes,
548 self.config.splineNNodes,
549 tx,
550 ty,
551 coeffsX,
552 coeffsY,
553 options="OutUnit=1",
554 )
556 newFrameDict = ast.FrameDict(pixelFrame)
557 newFrameDict.addFrame("PIXELS", pixToTPMap, tpFrame)
558 newFrameDict.addFrame("TP", splineMap, iwcFrame)
559 newFrameDict.addFrame("IWC", tpToSky, skyFrame)
560 outWcs = afwgeom.SkyWcs(newFrameDict)
561 catalog[d].setWcs(outWcs)
563 return catalog
565 def evaluate(self, gpx, gpy, positions, trainInd, testInd, inputWcs, makeValidationPlot=False):
566 """Calculate E and B-modes in the 2-point correlation function before
567 and after correcting for atmospheric turbulence, and validate
568 prediction on some of the test data.
570 Parameters
571 ----------
572 gpx : `treegp.gp_interp.GPInterpolation`
573 Gaussian Processes interpolator for x-direction residuals.
574 gpy : `treegp.gp_interp.GPInterpolation`
575 Gaussian Processes interpolator for y-direction residuals.
576 positions : `astropy.table.Table`
577 Catalog of input positions with residuals to the best fit.
578 trainInds : `numpy.ndarray`
579 Array of indices for points used in training.
580 testInds : `numpy.ndarray`
581 Array of indices for points not used in training.
582 inputWcs : `lsst.afw.table.ExposureCatalog`
583 Catalog with WCSs for each detector of the input exposure.
584 makeValidationPlot : `bool`, optional
585 Whether to make a plot showing the prediction on the validation
586 data.
587 """
588 dx = positions["xresw"]
589 dy = positions["yresw"]
591 # Get tangent plane coordinates for input points
592 tpCoords = np.zeros((len(positions), 2))
593 for detector in inputWcs:
594 detId = detector["id"]
595 detWCS = detector.wcs
596 detInd = positions["deviceName"].astype(int) == detId
597 detectorSources = positions[detInd]
599 tangentPlaneToSky = detWCS.getFrameDict().getMapping("PIXELS", "IWC")
600 tangentPlaneCoords = tangentPlaneToSky.applyForward(
601 np.array([detectorSources["xpix"], detectorSources["ypix"]])
602 )
603 tpCoords[detInd] = tangentPlaneCoords.T
605 # Calculate E/B modes before and after Gaussian Processes correction.
606 xPredict = gpx.predict(tpCoords[trainInd])
607 yPredict = gpy.predict(tpCoords[trainInd])
608 xie, xib, logr = treegp.comp_eb_treecorr(
609 tpCoords[trainInd, 0],
610 tpCoords[trainInd, 1],
611 dx[trainInd],
612 dy[trainInd],
613 rmin=20 / 3600,
614 rmax=0.6,
615 dlogr=0.3,
616 )
617 start, stop = np.searchsorted(np.exp(logr), [0, 15])
618 meanE = np.mean(xie[start:stop])
619 meanB = np.mean(xib[start:stop])
620 self.log.info(
621 "Original average correlation level over 0-15 arcminutes: E-mode=%0.2f, B-mode=%0.2f",
622 meanE,
623 meanB,
624 )
626 xie_resid, xib_resid, logr = treegp.comp_eb_treecorr(
627 tpCoords[trainInd, 0],
628 tpCoords[trainInd, 1],
629 dx[trainInd] - xPredict,
630 dy[trainInd] - yPredict,
631 rmin=20 / 3600,
632 rmax=0.6,
633 dlogr=0.3,
634 )
635 start, stop = np.searchsorted(np.exp(logr), [0, 15])
636 meanE_resid = np.mean(xie_resid[start:stop])
637 meanB_resid = np.mean(xib_resid[start:stop])
638 self.log.info(
639 "Correlation level after GP correction over 0-15 arcminutes: E-mode=%0.2f, B-mode=%0.2f",
640 meanE_resid,
641 meanB_resid,
642 )
644 # Predict on all test data and make a plot.
645 if makeValidationPlot:
646 print(len(testInd))
647 testInd = testInd[:50000]
648 chunkSize = 5000
649 nChunks = np.ceil(len(testInd) / chunkSize).astype(int)
650 xPredict = np.zeros(len(testInd))
651 yPredict = np.zeros(len(testInd))
652 for i in range(nChunks):
653 ind = testInd[chunkSize * i : chunkSize * (i + 1)]
654 xPredict[chunkSize * i : chunkSize * (i + 1)] = gpx.predict(tpCoords[ind])
655 yPredict[chunkSize * i : chunkSize * (i + 1)] = gpy.predict(tpCoords[ind])
656 fig = plot_visit(
657 tpCoords[testInd, 0], tpCoords[testInd, 1], dx[testInd], dy[testInd], xPredict, yPredict
658 )
659 return fig