From caf67140a4b958b46ba5b90243330aaf4d271270 Mon Sep 17 00:00:00 2001 From: John Franklin Crenshaw <41785729+jfcrenshaw@users.noreply.github.com> Date: Sun, 17 Nov 2024 16:02:58 -0300 Subject: [PATCH] tickets/DM-47480: sparse Zernike estimation (#290) Implemented sparse Zernike estimation --- doc/versionHistory.rst | 11 ++ .../comCamRapidAnalysisPipeline.yaml | 2 +- .../comCamSimRapidAnalysisPipeline.yaml | 11 +- python/lsst/ts/wep/estimation/danish.py | 34 +++-- python/lsst/ts/wep/estimation/tie.py | 36 ++++- python/lsst/ts/wep/estimation/wfAlgorithm.py | 81 ++++++----- python/lsst/ts/wep/estimation/wfEstimator.py | 86 +++-------- python/lsst/ts/wep/imageMapper.py | 82 ++++++++++- python/lsst/ts/wep/instrument.py | 81 +++++------ python/lsst/ts/wep/task/calcZernikesTask.py | 16 +- .../lsst/ts/wep/task/estimateZernikesBase.py | 21 +-- python/lsst/ts/wep/utils/zernikeUtils.py | 137 ++++++++++++++++++ tests/estimation/test_tie.py | 1 + tests/estimation/test_wfAlgorithm.py | 9 +- tests/estimation/test_wfEstimator.py | 53 +++---- tests/task/test_calcZernikesTieTaskCwfs.py | 30 ++++ tests/test_imageMapper.py | 3 +- tests/test_instrument.py | 28 +++- tests/utils/test_zernikeUtils.py | 69 +++++++++ 19 files changed, 545 insertions(+), 246 deletions(-) diff --git a/doc/versionHistory.rst b/doc/versionHistory.rst index 91f0c718..756ff18c 100644 --- a/doc/versionHistory.rst +++ b/doc/versionHistory.rst @@ -6,6 +6,17 @@ Version History ################## +.. _lsst.ts.wep-13.0.0: + +------------- +13.0.0 +------------- + +* enabled sparse Zernike estimation +* removed most jmax and return4Up configs in favor of nollIndices configs +* removed return4Up from estimator WfEstimator and WfAlgorithm +* added makeSparse and makeDense to Zernike utils + .. _lsst.ts.wep-12.7.0: ------------- diff --git a/pipelines/production/comCamRapidAnalysisPipeline.yaml b/pipelines/production/comCamRapidAnalysisPipeline.yaml index 375f2900..9d13df60 100644 --- a/pipelines/production/comCamRapidAnalysisPipeline.yaml +++ b/pipelines/production/comCamRapidAnalysisPipeline.yaml @@ -17,7 +17,7 @@ tasks: calcZernikesTask: class: lsst.ts.wep.task.calcZernikesTask.CalcZernikesTask config: - estimateZernikes.maxNollIndex: 22 + estimateZernikes.nollIndices: [4, 5, 6, 7, 8, 9, 10, 11, 14, 15, 20, 21, 27, 28] estimateZernikes.convergeTol: 10.0e-9 estimateZernikes.requireConverge: True estimateZernikes.saveHistory: False diff --git a/pipelines/production/comCamSimRapidAnalysisPipeline.yaml b/pipelines/production/comCamSimRapidAnalysisPipeline.yaml index a2ef4ab4..9be9f33c 100644 --- a/pipelines/production/comCamSimRapidAnalysisPipeline.yaml +++ b/pipelines/production/comCamSimRapidAnalysisPipeline.yaml @@ -5,21 +5,24 @@ tasks: class: lsst.ts.wep.task.generateDonutDirectDetectTask.GenerateDonutDirectDetectTask config: donutSelector.useCustomMagLimit: True - donutSelector.sourceLimit: 5 + donutSelector.sourceLimit: 20 cutOutDonutsScienceSensorGroupTask: class: lsst.ts.wep.task.cutOutDonutsScienceSensorTask.CutOutDonutsScienceSensorTask config: python: | from lsst.ts.wep.task.pairTask import GroupPairer config.pairer.retarget(GroupPairer) - donutStampSize: 160 + donutStampSize: 200 initialCutoutPadding: 40 calcZernikesTask: class: lsst.ts.wep.task.calcZernikesTask.CalcZernikesTask config: - estimateZernikes.maxNollIndex: 28 + estimateZernikes.nollIndices: [4, 5, 6, 7, 8, 9, 10, 11, 14, 15, 20, 21, 27, 28] + estimateZernikes.convergeTol: 10.0e-9 + estimateZernikes.requireConverge: True estimateZernikes.saveHistory: False - estimateZernikes.maskKwargs: {'doMaskBlends': False} + estimateZernikes.maskKwargs: { "doMaskBlends": False } + donutStampSelector.maxSelect: 5 isr: class: lsst.ip.isr.IsrTask config: diff --git a/python/lsst/ts/wep/estimation/danish.py b/python/lsst/ts/wep/estimation/danish.py index 97d2498a..988a8514 100644 --- a/python/lsst/ts/wep/estimation/danish.py +++ b/python/lsst/ts/wep/estimation/danish.py @@ -124,6 +124,7 @@ def history(self) -> dict: following entries - "image" - the image that is being fit - "variance" - the background variance that was used for fitting + - "nollIndices" - Noll indices for which Zernikes are estimated - "zkStart" - the starting Zernike coefficients - "lstsqResult" - dictionary of results returned by least_squares - "zkFit" - the Zernike coefficients fit to the donut @@ -140,6 +141,7 @@ def _estimateSingleZk( self, image: Image, zkStart: np.ndarray, + nollIndices: np.ndarray, instrument: Instrument, factory: danish.DonutFactory, saveHistory: bool, @@ -152,6 +154,8 @@ def _estimateSingleZk( The ts_wep image of the donut stamp zkStart : np.ndarray The starting point for the Zernikes + nollIndices : np.ndarray + Noll indices for which to estimate Zernikes instrument : Instrument The ts_wep Instrument factory : danish.DonutFactory @@ -170,22 +174,16 @@ def _estimateSingleZk( if np.isfinite(image.blendOffsets).sum() > 0: warnings.warn("Danish is currently only setup for non-blended donuts.") - # Determine the maximum Noll index - jmax = len(zkStart) + 3 - # Create reference Zernikes by adding off-axis coefficients to zkStart - zkStart = np.pad(zkStart, (4, 0)) offAxisCoeff = instrument.getOffAxisCoeff( *image.fieldAngle, image.defocalType, image.bandLabel, - jmaxIntrinsic=jmax, - return4Up=False, + nollIndicesModel=np.arange(0, 79), + nollIndicesIntr=nollIndices, ) - size = max(zkStart.size, offAxisCoeff.size) - zkRef = np.zeros(size) - zkRef[: zkStart.size] = zkStart - zkRef[: offAxisCoeff.size] += offAxisCoeff + zkRef = offAxisCoeff.copy() + zkRef[nollIndices] += zkStart # Create the background mask if it does not exist if image.maskBackground is None: @@ -213,14 +211,14 @@ def _estimateSingleZk( model = danish.SingleDonutModel( factory, z_ref=zkRef, - z_terms=np.arange(4, jmax + 1), + z_terms=nollIndices, thx=np.deg2rad(image.fieldAngle[0]), thy=np.deg2rad(image.fieldAngle[1]), npix=img.shape[0], ) # Create the initial guess for the model parameters - x0 = [0.0, 0.0, 1.0] + [0.0] * (jmax - 3) + x0 = [0.0, 0.0, 1.0] + [0.0] * len(nollIndices) # Use scipy to optimize the parameters try: @@ -238,7 +236,7 @@ def _estimateSingleZk( zkFit = np.array(zkFit) # Add the starting zernikes back into the result - zkSum = zkFit + zkStart[4:] + zkSum = zkFit + zkStart # Flag that we didn't hit GalSimFFTSizeError galSimFFTSizeError = False @@ -258,8 +256,8 @@ def _estimateSingleZk( except GalSimFFTSizeError: # Fill dummy objects result = None - zkFit = np.full(jmax - 3, np.nan) - zkSum = np.full(jmax - 3, np.nan) + zkFit = np.full_like(zkStart, np.nan) + zkSum = np.full_like(zkStart, np.nan) if saveHistory: modelImage = np.full_like(img, np.nan) @@ -271,6 +269,7 @@ def _estimateSingleZk( hist = { "image": img.copy(), "variance": backgroundStd**2, + "nollIndices": nollIndices.copy(), "zkStart": zkStart.copy(), "lstsqResult": result, "zkFit": zkFit.copy(), @@ -290,6 +289,7 @@ def _estimateZk( I2: Optional[Image], zkStartI1: np.ndarray, zkStartI2: Optional[np.ndarray], + nollIndices: np.ndarray, instrument: Instrument, saveHistory: bool, ) -> np.ndarray: @@ -305,6 +305,8 @@ def _estimateZk( The starting Zernikes for I1 (in meters, for Noll indices >= 4) zkStartI2 : np.ndarray or None The starting Zernikes for I2 (in meters, for Noll indices >= 4) + nollIndices : np.ndarray + Noll indices for which you wish to estimate Zernike coefficients. instrument : Instrument The Instrument object associated with the DonutStamps. saveHistory : bool @@ -334,6 +336,7 @@ def _estimateZk( zk1, hist[I1.defocalType.value] = self._estimateSingleZk( I1, zkStartI1, + nollIndices, instrument, factory, saveHistory, @@ -344,6 +347,7 @@ def _estimateZk( zk2, hist[I2.defocalType.value] = self._estimateSingleZk( I2, zkStartI2, + nollIndices, instrument, factory, saveHistory, diff --git a/python/lsst/ts/wep/estimation/tie.py b/python/lsst/ts/wep/estimation/tie.py index 80058ea5..98d7cfc5 100644 --- a/python/lsst/ts/wep/estimation/tie.py +++ b/python/lsst/ts/wep/estimation/tie.py @@ -32,6 +32,7 @@ binArray, createZernikeBasis, createZernikeGradBasis, + makeSparse, ) from scipy.ndimage import gaussian_filter @@ -472,6 +473,8 @@ def history(self) -> dict: full OPD. - "pupil" - model pupil image, in the case where only one image was passed to the estimator. + - "nollIndices" - array of Noll indices corresponding to the + estimated Zernikes values. Each iteration of the solver is then stored under indices >= 1. The entry for each iteration is also a dictionary, containing @@ -595,7 +598,7 @@ def _expSolve( self, I0: np.ndarray, dIdz: np.ndarray, - jmax: int, + nollIndices: np.ndarray, instrument: Instrument, ) -> np.ndarray: """Solve the TIE directly using a Zernike expansion. @@ -606,8 +609,8 @@ def _expSolve( The beam intensity at the exit pupil dIdz : np.ndarray The z-derivative of the beam intensity across the exit pupil - jmax : int - The maximum Zernike Noll index to estimate + nollIndices : np.ndarray + Noll indices for which you wish to estimate Zernike coefficients. instrument : Instrument, optional The Instrument object associated with the DonutStamps. (the default is the default Instrument) @@ -618,8 +621,11 @@ def _expSolve( Numpy array of the Zernike coefficients estimated from the image or pair of images, in nm. """ - # Get Zernike Bases + # Get the pupil grid uPupil, vPupil = instrument.createPupilGrid() + + # Get Zernike Bases + jmax = nollIndices.max() zk = createZernikeBasis( uPupil, vPupil, @@ -633,6 +639,11 @@ def _expSolve( obscuration=instrument.obscuration, ) + # Down-select to the Zernikes we are solving for + zk = makeSparse(zk, nollIndices) + dzkdu = makeSparse(dzkdu, nollIndices) + dzkdv = makeSparse(dzkdv, nollIndices) + # Calculate quantities for the linear system # See Equations 43-45 of https://sitcomtn-111.lsst.io b = np.einsum("ab,jab->j", dIdz, zk, optimize=True) @@ -651,6 +662,7 @@ def _estimateZk( I2: Optional[Image], # type: ignore[override] zkStartI1: np.ndarray, zkStartI2: Optional[np.ndarray], + nollIndices: np.ndarray, instrument: Instrument, saveHistory: bool, ) -> np.ndarray: @@ -667,6 +679,8 @@ def _estimateZk( zkStartI2 : np.ndarray or None The starting Zernikes for I2 (in meters, for Noll indices >= 4) Can be None. + nollIndices : np.ndarray + Noll indices for which you wish to estimate Zernike coefficients. instrument : Instrument The Instrument object associated with the DonutStamps. saveHistory : bool @@ -685,9 +699,11 @@ def _estimateZk( RuntimeError If the solver is not supported """ + # Rescale instrument pixel size to account for binning if self.binning > 1: instrument = instrument.copy() instrument.pixelSize *= self.binning + # Create the ImageMapper for centering and image compensation imageMapper = ImageMapper(instConfig=instrument, opticalModel=self.opticalModel) @@ -699,6 +715,7 @@ def _estimateZk( zkStartI2, ) + # Bin images to reduce resolution? if self.binning > 1: if intra is not None: intra.image = binArray(intra.image, self.binning) @@ -731,6 +748,7 @@ def _estimateZk( "zkStartExtra": None if extra is None else zkStartExtra.copy(), "zkStartMean": zkStartMean.copy(), "pupil": None if pupil is None else pupil.copy(), + "nollIndices": nollIndices.copy(), } # Initialize Zernike arrays at zero @@ -744,7 +762,7 @@ def _estimateZk( compSequence = iter(self.compSequence) # Determine the maximum Noll index we will solve for - jmax = len(zkStartMean) + 3 + jmax = nollIndices.max() # Set the caustic and converged flags to False caustic = False @@ -760,7 +778,7 @@ def _estimateZk( # The gain scales how much of previous residual we incorporate # Everything past jmaxComp is set to zero zkComp += self.compGain * zkResid - zkComp[(jmaxComp - 3) :] = 0 + zkComp[nollIndices > jmaxComp] = 0 # Center the images recenter = (i == 0) or (np.max(np.abs(zkComp - zkCenter)) > self.centerTol) @@ -773,6 +791,7 @@ def _estimateZk( else imageMapper.centerOnProjection( intra, zkCenter + zkStartIntra, + nollIndices, isBinary=self.centerBinary, **self.maskKwargs, ) @@ -783,6 +802,7 @@ def _estimateZk( else imageMapper.centerOnProjection( extra, zkCenter + zkStartExtra, + nollIndices, isBinary=self.centerBinary, **self.maskKwargs, ) @@ -795,6 +815,7 @@ def _estimateZk( else imageMapper.mapImageToPupil( intraCent, zkComp + zkStartIntra, + nollIndices, masks=None if i == 0 else intraComp.masks, # noqa: F821 **self.maskKwargs, ) @@ -805,6 +826,7 @@ def _estimateZk( else imageMapper.mapImageToPupil( extraCent, zkComp + zkStartExtra, + nollIndices, masks=None if i == 0 else extraComp.masks, # noqa: F821 **self.maskKwargs, ) @@ -842,7 +864,7 @@ def _estimateZk( ) # Estimate the Zernikes - zkResid = self._expSolve(I0, dIdz, jmax, instrument) + zkResid = self._expSolve(I0, dIdz, nollIndices, instrument) # If estimating with a single donut, double the coefficients # This is because the simulated pupil image is aberration free diff --git a/python/lsst/ts/wep/estimation/wfAlgorithm.py b/python/lsst/ts/wep/estimation/wfAlgorithm.py index 95bc2520..dcd87ef2 100644 --- a/python/lsst/ts/wep/estimation/wfAlgorithm.py +++ b/python/lsst/ts/wep/estimation/wfAlgorithm.py @@ -22,11 +22,16 @@ __all__ = ["WfAlgorithm"] from abc import ABC, abstractmethod -from typing import Optional +from typing import Optional, Sequence import numpy as np from lsst.ts.wep import Image, Instrument -from lsst.ts.wep.utils import convertZernikesToPsfWidth +from lsst.ts.wep.utils import ( + checkNollIndices, + convertZernikesToPsfWidth, + makeDense, + makeSparse, +) class WfAlgorithm(ABC): @@ -66,8 +71,8 @@ def history(self) -> dict: def _validateInputs( self, I1: Image, - I2: Optional[Image], - jmax: int, + I2: Image | None, + nollIndices: np.ndarray, instrument: Instrument, startWithIntrinsic: bool, returnWfDev: bool, @@ -80,20 +85,22 @@ def _validateInputs( ---------- I1 : DonutStamp An Image object containing an intra- or extra-focal donut image. - I2 : DonutStamp, optional + I2 : DonutStamp or None A second image, on the opposite side of focus from I1. - jmax : int, optional - The maximum Zernike Noll index to estimate. - instrument : Instrument, optional + nollIndices : np.ndarray + List, tuple, or array of Noll indices for which you wish to + estimate Zernike coefficients. Note these values must be unique, + ascending, and >= 4. + instrument : Instrument The Instrument object associated with the DonutStamps. - startWithIntrinsic : bool, optional + startWithIntrinsic : bool Whether to start the Zernike estimation process from the intrinsic Zernikes rather than zero. - returnWfDev : bool, optional + returnWfDev : bool If False, the full OPD is returned. If True, the wavefront deviation is returned. The wavefront deviation is defined as the OPD - intrinsic Zernikes. - units : str, optional + units : str Units in which the Zernike amplitudes are returned. Options are "m", "nm", "um", or "arcsecs". @@ -128,11 +135,8 @@ def _validateInputs( if I2.defocalType == I1.defocalType: raise ValueError("I1 and I2 must be on opposite sides of focus.") - # Validate jmax - if not isinstance(jmax, int): - raise TypeError("jmax must be an integer.") - if jmax < 4: - raise ValueError("jmax must be greater than or equal to 4.") + # Validate nollIndices + checkNollIndices(nollIndices) # Validate the instrument if not isinstance(instrument, Instrument): @@ -164,6 +168,7 @@ def _estimateZk( I2: Optional[Image], zkStartI1: np.ndarray, zkStartI2: Optional[np.ndarray], + nollIndices: np.ndarray, instrument: Instrument, saveHistory: bool, ) -> np.ndarray: @@ -183,6 +188,8 @@ def _estimateZk( The starting Zernikes for I1 (in meters, for Noll indices >= 4) zkStartI2 : np.ndarray or None The starting Zernikes for I2 (in meters, for Noll indices >= 4) + nollIndices : np.ndarray + Noll indices for which you wish to estimate Zernike coefficients. instrument : Instrument The Instrument object associated with the DonutStamps. saveHistory : bool @@ -202,11 +209,10 @@ def estimateZk( self, I1: Image, I2: Optional[Image] = None, - jmax: int = 22, + nollIndices: Sequence = tuple(np.arange(4, 23)), instrument: Instrument = Instrument(), startWithIntrinsic: bool = True, returnWfDev: bool = False, - return4Up: bool = True, units: str = "m", saveHistory: bool = False, ) -> np.ndarray: @@ -219,9 +225,10 @@ def estimateZk( I2 : DonutStamp, optional A second image, on the opposite side of focus from I1. (the default is None) - jmax : int, optional - The maximum Zernike Noll index to estimate. - (the default is 22) + nollIndices : Sequence, optional + List, tuple, or array of Noll indices for which you wish to + estimate Zernike coefficients. Note these values must be unique, + ascending, and >= 4. (the default is indices 4-22) instrument : Instrument, optional The Instrument object associated with the DonutStamps. (the default is the default Instrument) @@ -232,14 +239,6 @@ def estimateZk( If False, the full OPD is returned. If True, the wavefront deviation is returned. The wavefront deviation is defined as the OPD - intrinsic Zernikes. (the default is False) - return4Up : bool, optional - If True, the returned Zernike coefficients start with - Noll index 4. If False, they follow the Galsim convention - of starting with index 0 (which is meaningless), so the - array index of the output corresponds to the Noll index. - In this case, indices 0-3 are always set to zero, because - they are not estimated by our pipeline. - (the default is True) units : str, optional Units in which the Zernike amplitudes are returned. Options are "m", "nm", "um", or "arcsecs". @@ -255,11 +254,14 @@ def estimateZk( np.ndarray Zernike coefficients estimated from the image (or pair of images) """ + # Convert nollIndices to an array + nollIndices = np.array(nollIndices) + # Validate the inputs self._validateInputs( I1, I2, - jmax, + nollIndices, instrument, startWithIntrinsic, returnWfDev, @@ -272,12 +274,16 @@ def estimateZk( zkIntrinsicI1 = instrument.getIntrinsicZernikes( *I1.fieldAngle, I1.bandLabel, - jmax, + nollIndices, ) zkIntrinsicI2 = ( None if I2 is None - else instrument.getIntrinsicZernikes(*I2.fieldAngle, I2.bandLabel, jmax) + else instrument.getIntrinsicZernikes( + *I2.fieldAngle, + I2.bandLabel, + nollIndices, + ) ) # Determine the Zernikes to start with @@ -285,8 +291,8 @@ def estimateZk( zkStartI1 = zkIntrinsicI1 zkStartI2 = zkIntrinsicI2 else: - zkStartI1 = np.zeros(jmax - 3) - zkStartI2 = np.zeros(jmax - 3) + zkStartI1 = np.zeros_like(nollIndices, dtype=float) + zkStartI2 = None if I2 is None else np.zeros_like(nollIndices, dtype=float) # Clear the algorithm history self._history = {} @@ -297,6 +303,7 @@ def estimateZk( I2=I2, zkStartI1=zkStartI1, zkStartI2=zkStartI2, + nollIndices=nollIndices, instrument=instrument, saveHistory=saveHistory, ) @@ -318,16 +325,14 @@ def estimateZk( elif units == "nm": zk *= 1e9 elif units == "arcsec": + zk = makeDense(zk, nollIndices) zk = convertZernikesToPsfWidth( zernikes=zk * 1e6, diameter=instrument.diameter, obscuration=instrument.obscuration, ) + zk = makeSparse(zk, nollIndices) else: # pragma: no cover raise RuntimeError(f"Conversion to unit '{units}' not supported.") - # Pad array so that array index = Noll index? - if not return4Up: - zk = np.pad(zk, (4, 0)) - return zk diff --git a/python/lsst/ts/wep/estimation/wfEstimator.py b/python/lsst/ts/wep/estimation/wfEstimator.py index 364bd846..8403b1ee 100644 --- a/python/lsst/ts/wep/estimation/wfEstimator.py +++ b/python/lsst/ts/wep/estimation/wfEstimator.py @@ -21,13 +21,13 @@ __all__ = ["WfEstimator"] -from typing import Optional, Union +from typing import Optional, Sequence, Union import numpy as np from lsst.ts.wep import Image, Instrument from lsst.ts.wep.estimation.wfAlgorithm import WfAlgorithm from lsst.ts.wep.estimation.wfAlgorithmFactory import WfAlgorithmFactory -from lsst.ts.wep.utils import configClass +from lsst.ts.wep.utils import checkNollIndices, configClass class WfEstimator: @@ -53,9 +53,10 @@ class WfEstimator: If a dictionary, it is assumed to hold keywords for configuration. If an Instrument object, that object is just used. (the default is "policy:instrument/LsstCam.yaml") - jmax : int, optional - The maximum Zernike Noll index to estimate. - (the default is 22) + nollIndices : Sequence, optional + List, tuple, or array of Noll indices for which you wish to + estimate Zernike coefficients. Note these values must be unique, + ascending, and >= 4. (the default is indices 4-22) startWithIntrinsic : bool, optional Whether to start the Zernike estimation process from the intrinsic Zernikes rather than zero. @@ -65,14 +66,6 @@ class WfEstimator: deviation is returned. The wavefront deviation is defined as the OPD - intrinsic Zernikes. (the default is False) - return4Up : bool, optional - If True, the returned Zernike coefficients start with - Noll index 4. If False, they follow the Galsim convention - of starting with index 0 (which is meaningless), so the - array index of the output corresponds to the Noll index. - In this case, indices 0-3 are always set to zero, because - they are not estimated by our pipeline. - (the default is True) units : str, optional Units in which the Zernike coefficients are returned. Options are "m", "um", "nm", or "arcsec". @@ -89,10 +82,9 @@ def __init__( algoName: str = "tie", algoConfig: Union[dict, WfAlgorithm, None] = None, instConfig: Union[str, dict, Instrument] = "policy:instruments/LsstCam.yaml", - jmax: int = 22, + nollIndices: Sequence = tuple(np.arange(4, 23)), startWithIntrinsic: bool = True, returnWfDev: bool = False, - return4Up: bool = True, units: str = "um", saveHistory: bool = False, ) -> None: @@ -101,10 +93,9 @@ def __init__( self.instrument = configClass(instConfig, Instrument) # Set the other parameters - self.jmax = jmax + self.nollIndices = nollIndices self.startWithIntrinsic = startWithIntrinsic self.returnWfDev = returnWfDev - self.return4Up = return4Up self.units = units self.saveHistory = saveHistory @@ -158,28 +149,24 @@ def instrument(self, value: Instrument) -> None: self._instrument = value @property - def jmax(self) -> int: - """Return the maximum Zernike Noll index that will be estimated.""" - return self._jmax + def nollIndices(self) -> np.ndarray: + """Return the Noll indices for which Zernikes are estimated.""" + return self._nollIndices - @jmax.setter - def jmax(self, value: int) -> None: - """Set jmax + @nollIndices.setter + def nollIndices(self, value: Sequence) -> None: + """Set the Noll indices for which the Zernikes are estimated. Parameters ---------- - value : int - The maximum Zernike Noll index to estimate. - - Raises - ------ - ValueError - If jmax is < 4 + nollIndices : Sequence, optional + List, tuple, or array of Noll indices for which you wish to + estimate Zernike coefficients. Note these values must be unique, + ascending, and >= 4. (the default is indices 4-22) """ - value = int(value) - if value < 4: - raise ValueError("jmax must be greater than or equal to 4.") - self._jmax = value + value = np.array(value) + checkNollIndices(value) + self._nollIndices = value @property def startWithIntrinsic(self) -> bool: @@ -230,34 +217,6 @@ def returnWfDev(self, value: bool) -> None: raise TypeError("returnWfDev must be a bool.") self._returnWfDev = value - @property - def return4Up(self) -> bool: - """Whether to return Zernikes starting with Noll index 4.""" - return self._return4Up - - @return4Up.setter - def return4Up(self, value: bool) -> None: - """Set return4Up. - - Parameters - ---------- - value : bool - If True, the returned Zernike coefficients start with - Noll index 4. If False, they follow the Galsim convention - of starting with index 0 (which is meaningless), so the - array index of the output corresponds to the Noll index. - In this case, indices 0-3 are always set to zero, because - they are not estimated by our pipeline. - - Raises - ------ - TypeError - If the value is not a boolean - """ - if not isinstance(value, bool): - raise TypeError("return4Up must be a bool.") - self._return4Up = value - @property def units(self) -> str: """Return the wavefront units. @@ -337,11 +296,10 @@ def estimateZk( return self.algo.estimateZk( I1=I1, I2=I2, - jmax=self.jmax, + nollIndices=self.nollIndices, instrument=self.instrument, startWithIntrinsic=self.startWithIntrinsic, returnWfDev=self.returnWfDev, - return4Up=self.return4Up, units=self.units, saveHistory=self.saveHistory, ) diff --git a/python/lsst/ts/wep/imageMapper.py b/python/lsst/ts/wep/imageMapper.py index 088f40e5..babf3705 100644 --- a/python/lsst/ts/wep/imageMapper.py +++ b/python/lsst/ts/wep/imageMapper.py @@ -28,7 +28,7 @@ from lsst.ts.wep.utils.enumUtils import BandLabel, DefocalType, PlaneType from lsst.ts.wep.utils.ioUtils import configClass from lsst.ts.wep.utils.miscUtils import centerWithTemplate, polygonContains -from lsst.ts.wep.utils.zernikeUtils import zernikeGradEval +from lsst.ts.wep.utils.zernikeUtils import makeDense, zernikeGradEval from scipy.interpolate import interpn from scipy.ndimage import binary_dilation, shift @@ -118,6 +118,7 @@ def _constructForwardMap( uPupil: np.ndarray, vPupil: np.ndarray, zkCoeff: np.ndarray, + nollIndices: np.ndarray | None, image: Image, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Construct the forward mapping from the pupil to the image plane. @@ -131,6 +132,11 @@ def _constructForwardMap( zkCoeff : np.ndarray The wavefront at the pupil, represented as Zernike coefficients in meters for Noll indices >= 4. + nollIndices : np.ndarray, optional + These are the Noll indices corresponding to the coefficients in + zkCoeff. If None, it is assumed zkCoeff contains coefficients + for consecutive indices starting with Noll index 4. + (the default is None) image : Image A stamp object containing the metadata required for the mapping. @@ -150,6 +156,13 @@ def _constructForwardMap( RuntimeWarning If the optical model is not supported """ + # Resolve the Noll indices + if nollIndices is None: + nollIndices = np.arange(4, len(zkCoeff) + 4) + + # Turn zkCoeff into dense array, starting at Noll index 4 + zkCoeff = makeDense(zkCoeff, nollIndices) + # Get the Zernikes for the mapping if self.opticalModel == "onAxis" or self.opticalModel == "paraxial": zkMap = zkCoeff @@ -159,7 +172,7 @@ def _constructForwardMap( *image.fieldAngle, image.defocalType, image.bandLabel, - jmaxIntrinsic=len(zkCoeff) + 3, + nollIndicesIntr=nollIndices, ) # Add these coefficients to the input coefficients @@ -328,6 +341,7 @@ def _constructInverseMap( uImage: np.ndarray, vImage: np.ndarray, zkCoeff: np.ndarray, + nollIndices: np.ndarray | None, image: Image, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]: """Construct the inverse mapping from the image plane to the pupil. @@ -341,6 +355,11 @@ def _constructInverseMap( zkCoeff : np.ndarray The wavefront at the pupil, represented as Zernike coefficients in meters for Noll indices >= 4. + nollIndices : np.ndarray, optional + These are the Noll indices corresponding to the coefficients in + zkCoeff. If None, it is assumed zkCoeff contains coefficients + for consecutive indices starting with Noll index 4. + (the default is None) image : Image A stamp object containing the metadata required for the mapping. @@ -371,6 +390,7 @@ def _constructInverseMap( uPupilTest, vPupilTest, zkCoeff, + nollIndices, image, ) @@ -411,6 +431,7 @@ def _constructInverseMap( uPupil, vPupil, zkCoeff, + nollIndices, image, ) @@ -430,6 +451,7 @@ def _constructInverseMap( uPupil, vPupil, zkCoeff, + nollIndices, image, ) @@ -458,6 +480,7 @@ def _constructInverseMap( def _getImageGridInsidePupil( self, zkCoeff: np.ndarray, + nollIndices: np.ndarray | None, image: Image, ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: """Return image grid and mask for which pixels are inside the pupil. @@ -469,6 +492,10 @@ def _getImageGridInsidePupil( zkCoeff : np.ndarray The wavefront at the pupil, represented as Zernike coefficients in meters for Noll indices >= 4. + nollIndices : np.ndarray, optional + These are the Noll indices corresponding to the coefficients in + zkCoeff. If None, it is assumed zkCoeff contains coefficients + for consecutive indices starting with Noll index 4. image : Image A stamp object containing the metadata required for the mapping. @@ -488,6 +515,7 @@ def _getImageGridInsidePupil( uPupil, vPupil, zkCoeff, + nollIndices, image, ) imageEdge = np.array([uImageEdge, vImageEdge]).T @@ -1164,6 +1192,7 @@ def createImageMasks( self, image: Image, zkCoeff: Optional[np.ndarray] = None, + nollIndices: Optional[np.ndarray] = None, *, isBinary: bool = True, dilate: int = 0, @@ -1191,6 +1220,11 @@ def createImageMasks( The wavefront at the pupil, represented as Zernike coefficients in meters, for Noll indices >= 4. (the default are the intrinsic Zernikes at the donut position) + nollIndices : np.ndarray, optional + These are the Noll indices corresponding to the coefficients in + zkCoeff. If None, it is assumed zkCoeff contains coefficients + for consecutive indices starting with Noll index 4. + (the default is None) isBinary : bool, optional Whether to return a binary mask. If False, a fractional mask is returned instead. (the default is True) @@ -1240,7 +1274,9 @@ def createImageMasks( ) # Get the image grid inside the pupil - uImage, vImage, inside = self._getImageGridInsidePupil(zkCoeff, image) + uImage, vImage, inside = self._getImageGridInsidePupil( + zkCoeff, nollIndices, image + ) # Get the inverse mapping from image plane to pupil plane if _invMap is None: @@ -1249,6 +1285,7 @@ def createImageMasks( uImage[inside], vImage[inside], zkCoeff, + nollIndices, image, ) else: @@ -1295,6 +1332,7 @@ def getProjectionSize( defocalType: Union[DefocalType, str], bandLabel: Union[BandLabel, str] = BandLabel.REF, zkCoeff: Optional[np.ndarray] = None, + nollIndices: Optional[np.ndarray] = None, ) -> int: """Return size of the pupil projected onto the image plane (in pixels). @@ -1320,6 +1358,11 @@ def getProjectionSize( The wavefront at the pupil, represented as Zernike coefficients in meters, for Noll indices >= 4. (the default are the intrinsic Zernikes at the donut position) + nollIndices : np.ndarray, optional + These are the Noll indices corresponding to the coefficients in + zkCoeff. If None, it is assumed zkCoeff contains coefficients + for consecutive indices starting with Noll index 4. + (the default is None) Returns ------- @@ -1348,6 +1391,7 @@ def getProjectionSize( uPupil, vPupil, zkCoeff, + nollIndices, dummyImage, ) @@ -1364,6 +1408,7 @@ def centerOnProjection( self, image: Image, zkCoeff: Optional[np.ndarray] = None, + nollIndices: Optional[np.ndarray] = None, isBinary: bool = True, rMax: float = 10, **maskKwargs, @@ -1384,6 +1429,11 @@ def centerOnProjection( The wavefront at the pupil, represented as Zernike coefficients in meters, for Noll indices >= 4. (the default are the intrinsic Zernikes at the donut position) + nollIndices : np.ndarray, optional + These are the Noll indices corresponding to the coefficients in + zkCoeff. If None, it is assumed zkCoeff contains coefficients + for consecutive indices starting with Noll index 4. + (the default is None) isBinary : bool, optional If True, a binary mask is used to estimate the center of the image, otherwise a forward model of the image is used. The latter will @@ -1408,12 +1458,15 @@ def centerOnProjection( self.createImageMasks( stamp, zkCoeff, + nollIndices, isBinary=True, **maskKwargs, ) template = stamp.mask.copy() else: - template = self.mapPupilToImage(stamp, zkCoeff, **maskKwargs).image + template = self.mapPupilToImage( + stamp, zkCoeff, nollIndices, **maskKwargs + ).image # Center the image stamp.image = centerWithTemplate(stamp.image, template, rMax) @@ -1424,6 +1477,7 @@ def mapPupilToImage( self, image: Image, zkCoeff: Optional[np.ndarray] = None, + nollIndices: Optional[np.ndarray] = None, masks: Optional[Tuple[np.ndarray]] = None, **maskKwargs, ) -> Image: @@ -1444,6 +1498,11 @@ def mapPupilToImage( The wavefront at the pupil, represented as Zernike coefficients in meters, for Noll indices >= 4. (the default are the intrinsic Zernikes at the donut position) + nollIndices : np.ndarray, optional + These are the Noll indices corresponding to the coefficients in + zkCoeff. If None, it is assumed zkCoeff contains coefficients + for consecutive indices starting with Noll index 4. + (the default is None) masks : np.ndarray, optional You can provide the image masks if they have already been computed. This is just to speed up computation. If not provided, the masks @@ -1465,13 +1524,18 @@ def mapPupilToImage( ) # Get the image grid inside the pupil - uImage, vImage, inside = self._getImageGridInsidePupil(zkCoeff, stamp) + uImage, vImage, inside = self._getImageGridInsidePupil( + zkCoeff, + nollIndices, + stamp, + ) # Construct the inverse mapping uPupil, vPupil, jac, jacDet = self._constructInverseMap( uImage[inside], vImage[inside], zkCoeff, + nollIndices, stamp, ) @@ -1483,6 +1547,7 @@ def mapPupilToImage( self.createImageMasks( stamp, zkCoeff, + nollIndices, **maskKwargs, _invMap=(uPupil, vPupil, jac, jacDet), ) @@ -1502,6 +1567,7 @@ def mapImageToPupil( self, image: Image, zkCoeff: Optional[np.ndarray] = None, + nollIndices: Optional[np.ndarray] = None, masks: Optional[np.ndarray] = None, **maskKwargs, ) -> Image: @@ -1521,6 +1587,11 @@ def mapImageToPupil( The wavefront at the pupil, represented as Zernike coefficients in meters, for Noll indices >= 4. (the default are the intrinsic Zernikes at the donut position) + nollIndices : np.ndarray, optional + These are the Noll indices corresponding to the coefficients in + zkCoeff. If None, it is assumed zkCoeff contains coefficients + for consecutive indices starting with Noll index 4. + (the default is None) masks : np.ndarray, optional You can provide the image masks if they have already been computed. This is just to speed up computation. If not provided, the masks @@ -1550,6 +1621,7 @@ def mapImageToPupil( uPupil, vPupil, zkCoeff, + nollIndices, stamp, ) diff --git a/python/lsst/ts/wep/instrument.py b/python/lsst/ts/wep/instrument.py index 0bffbbd3..f4d567c4 100644 --- a/python/lsst/ts/wep/instrument.py +++ b/python/lsst/ts/wep/instrument.py @@ -25,11 +25,12 @@ from functools import lru_cache from pathlib import Path -from typing import Optional, Tuple, Union +from typing import Optional, Sequence, Tuple, Union import batoid import numpy as np -from lsst.ts.wep.utils import BandLabel, DefocalType, EnumDict, mergeConfigWithFile +from lsst.ts.wep.utils.enumUtils import BandLabel, DefocalType, EnumDict +from lsst.ts.wep.utils.ioUtils import mergeConfigWithFile from scipy.optimize import minimize_scalar @@ -735,7 +736,7 @@ def _getIntrinsicZernikesCached( Returns ------- np.ndarray - The Zernike coefficients in meters + The Zernike coefficients in meters, starting with index 0 """ # Get the band enum band = BandLabel(band) @@ -773,8 +774,7 @@ def getIntrinsicZernikes( xAngle: float, yAngle: float, band: Union[BandLabel, str] = BandLabel.REF, - jmax: int = 78, - return4Up: bool = True, + nollIndices: Sequence = tuple(np.arange(4, 79)), ) -> np.ndarray: """Return the intrinsic Zernikes associated with the optical design. @@ -788,25 +788,22 @@ def getIntrinsicZernikes( The BandLabel Enum or corresponding string, specifying which batoid model to load. Only relevant if self.batoidModelName contains "{band}". (the default is BandLabel.REF) - jmax : int, optional - The maximum Noll index of the intrinsic Zernikes. - (the default is 78) - return4Up : bool, optional - Whether to only return the coefficients for Noll indices >= 4. - (the default is True) + nollIndices : np.ndarray, optional + Noll indices for which to return Zernikes. + (the default is indices 4-78) Returns ------- np.ndarray The Zernike coefficients in meters """ - zk = self._getIntrinsicZernikesCached(xAngle, yAngle, band, jmax).copy() + # Make sure this is an array + nollIndices = np.array(nollIndices) - if return4Up: - # Keep only Noll indices >= 4 - zk = zk[4:] + # Retrieve cached Zernikes + zk = self._getIntrinsicZernikesCached(xAngle, yAngle, band, nollIndices.max()) - return zk + return zk[nollIndices] @lru_cache(100) def _getIntrinsicZernikesTACached( @@ -838,7 +835,7 @@ def _getIntrinsicZernikesTACached( Returns ------- np.ndarray - The Zernike coefficients in meters + The Zernike coefficients in meters, starting with index 0 Notes ----- @@ -893,9 +890,8 @@ def getOffAxisCoeff( yAngle: float, defocalType: DefocalType, band: Union[BandLabel, str] = BandLabel.REF, - jmax: int = 78, - jmaxIntrinsic: int = 78, - return4Up: bool = True, + nollIndicesModel: Sequence = tuple(np.arange(4, 79)), + nollIndicesIntr: Sequence = tuple(np.arange(4, 79)), ) -> np.ndarray: """Return the Zernike coefficients associated with the off-axis model. @@ -912,32 +908,34 @@ def getOffAxisCoeff( The BandLabel Enum or corresponding string, specifying which batoid model to load. Only relevant if self.batoidModelName contains "{band}". (the default is BandLabel.REF) - jmax : int, optional - The maximum Noll index of the off-axis model Zernikes. - (the default is 78) - jmaxIntrinsic : int, optional - The off-axis coefficients are calculated by subtracting the - intrinsic Zernikes from batoid.zernikeTA. This value sets the - maximum Noll index of the intrinsic Zernikes that are subtracted - from batoid.zernikeTA. It is usually the jmax of the Zernikes - being estimated by the wavefront estimators. - (the default is 78) - return4Up : bool, optional - Whether to only return the coefficients for Noll indices >= 4. - (the default is True) + nollIndicesModel : np.ndarray, optional + Noll indices of Zernikes retrieved for the off-axis model. + (the default is indices 4-78) + nollIndicesIntr : np.ndarray, optional + Noll indices of Zernikes you are estimating in the TIE or + Danish. The off-axis coefficients are calculated by retrieving + coefficients from batoid.zernikeTA, and then subtracting off the + intrinsic Zernikes for Noll indices you are estimating. This is + allows you to determine whether intrinsic Zernikes are included + in wavefront estimates when using WfEstimator. + (the default is indices 4-22). Returns ------- np.ndarray The Zernike coefficients in meters, for Noll indices >= 4 """ + # Make sure these are arrays + nollIndicesModel = np.array(nollIndicesModel) + nollIndicesIntr = np.array(nollIndicesIntr) + # Get zernikeTA zkTA = self._getIntrinsicZernikesTACached( xAngle, yAngle, defocalType, band, - jmax, + nollIndicesModel.max(), ) # Get regular intrinsic zernikes @@ -945,18 +943,15 @@ def getOffAxisCoeff( xAngle, yAngle, band, - min(jmax, jmaxIntrinsic), + nollIndicesIntr.max(), ) - # Subtract the intrinsics from zernikeTA - offAxisCoeff = zkTA.copy() - offAxisCoeff[: zk.size] -= zk - - if return4Up: - # Keep only Noll indices >= 4 - offAxisCoeff = offAxisCoeff[4:] + # Subtract intrinsics from zernikeTA + offAxisCoeff = np.zeros(max(zkTA.size, zk.size), dtype=float) + offAxisCoeff[nollIndicesModel] = zkTA[nollIndicesModel] + offAxisCoeff[nollIndicesIntr] -= zk[nollIndicesIntr] - return offAxisCoeff + return offAxisCoeff[nollIndicesModel] @property def maskParams(self) -> dict: diff --git a/python/lsst/ts/wep/task/calcZernikesTask.py b/python/lsst/ts/wep/task/calcZernikesTask.py index 87eeaeba..b0990fc4 100644 --- a/python/lsst/ts/wep/task/calcZernikesTask.py +++ b/python/lsst/ts/wep/task/calcZernikesTask.py @@ -131,7 +131,7 @@ def __init__(self, **kwargs) -> None: # Create subtasks self.estimateZernikes = self.config.estimateZernikes self.makeSubtask("estimateZernikes") - self.maxNollIndex = self.estimateZernikes.config.maxNollIndex + self.nollIndices = self.estimateZernikes.config.nollIndices self.combineZernikes = self.config.combineZernikes self.makeSubtask("combineZernikes") @@ -165,7 +165,7 @@ def initZkTable(self) -> QTable: ("intra_frac_bad_pix", " QTable: table["extra_field"].unit = u.deg table["intra_centroid"].unit = u.pixel table["extra_centroid"].unit = u.pixel - for j in range(4, self.maxNollIndex + 1): + for j in self.nollIndices: table[f"Z{j}"].unit = u.nm return table @@ -215,8 +215,8 @@ def createZkTable( "label": "average", "used": True, **{ - f"Z{j}": zkCoeffCombined.combinedZernikes[j - 4] * u.micron - for j in range(4, self.maxNollIndex + 1) + f"Z{j}": zkCoeffCombined.combinedZernikes[i] * u.micron + for i, j in enumerate(self.nollIndices) }, "intra_field": np.nan, "extra_field": np.nan, @@ -248,7 +248,7 @@ def createZkTable( row["label"] = f"pair{i+1}" row["used"] = not flag row.update( - {f"Z{j}": zk[j - 4] * u.micron for j in range(4, self.maxNollIndex + 1)} + {f"Z{j}": zk[i] * u.micron for i, j in enumerate(self.nollIndices)} ) row["intra_field"] = ( (np.array(np.nan, dtype=pos2f_dtype) * u.deg) @@ -357,8 +357,8 @@ def empty(self) -> pipeBase.Struct: "DEFOCAL_TYPE", ] return pipeBase.Struct( - outputZernikesRaw=np.atleast_2d(np.full(self.maxNollIndex - 3, np.nan)), - outputZernikesAvg=np.atleast_2d(np.full(self.maxNollIndex - 3, np.nan)), + outputZernikesRaw=np.atleast_2d(np.full(len(self.nollIndices), np.nan)), + outputZernikesAvg=np.atleast_2d(np.full(len(self.nollIndices), np.nan)), zernikes=self.initZkTable(), donutQualityTable=QTable({name: [] for name in qualityTableCols}), ) diff --git a/python/lsst/ts/wep/task/estimateZernikesBase.py b/python/lsst/ts/wep/task/estimateZernikesBase.py index b52d69ac..646d314e 100644 --- a/python/lsst/ts/wep/task/estimateZernikesBase.py +++ b/python/lsst/ts/wep/task/estimateZernikesBase.py @@ -45,10 +45,13 @@ class EstimateZernikesBaseConfig(pexConfig.Config): dtype=str, optional=True, ) - maxNollIndex = pexConfig.Field( + nollIndices = pexConfig.ListField( dtype=int, - default=28, - doc="The maximum Zernike Noll index estimated.", + default=tuple(range(4, 29)), + doc="Noll indices for which you wish to estimate Zernike coefficients. " + + "Note these values must be unique, ascending, >= 4, and azimuthal pairs " + + "must be complete. For example, if nollIndices contains 5, it must also " + + "contain 6 (because 5 and 6 are the azimuthal pairs for astigmatism).", ) startWithIntrinsic = pexConfig.Field( dtype=bool, @@ -60,15 +63,6 @@ class EstimateZernikesBaseConfig(pexConfig.Config): default=False, doc="If True, returns wavefront deviation. If False, returns full OPD.", ) - return4Up = pexConfig.Field( - dtype=bool, - default=True, - doc="If True, the returned Zernike coefficients start with Noll index 4. " - + "If False, they follow the Galsim convention of starting with index 0 " - + "(which is meaningless), so the array index of the output corresponds " - + "to the Noll index. In this case, indices 0-3 are always set to zero, " - + "because they are not estimated by our pipeline.", - ) saveHistory = pexConfig.Field( dtype=bool, default=False, @@ -254,10 +248,9 @@ def run( algoName=self.wfAlgoName, algoConfig=self.wfAlgoConfig, instConfig=instrument, - jmax=self.config.maxNollIndex, + nollIndices=self.config.nollIndices, startWithIntrinsic=self.config.startWithIntrinsic, returnWfDev=self.config.returnWfDev, - return4Up=self.config.return4Up, units="um", saveHistory=self.config.saveHistory, ) diff --git a/python/lsst/ts/wep/utils/zernikeUtils.py b/python/lsst/ts/wep/utils/zernikeUtils.py index e563e122..7720399f 100644 --- a/python/lsst/ts/wep/utils/zernikeUtils.py +++ b/python/lsst/ts/wep/utils/zernikeUtils.py @@ -29,6 +29,9 @@ "getPsfGradPerZernike", "convertZernikesToPsfWidth", "getZernikeParity", + "makeSparse", + "makeDense", + "checkNollIndices", ] from typing import Optional @@ -557,3 +560,137 @@ def getZernikeParity(jmin: int = 4, jmax: int = 22, axis: str = "x"): raise ValueError("axis must be either 'x' or 'y'.") return np.array(parity) + + +def makeSparse(values: np.ndarray, nollIndices: np.ndarray) -> np.ndarray: + """Down-select the values array according to requested Noll indices. + + For example, if values=[1, -2, 3, -4], nollIndices=[4, 5, 7], this + function returns [1, -2, -4]. + + Parameters + ---------- + values : np.ndarray + Array of values to be down-selected. Selection is applied to + the first axis, which must correspond to consecutive Noll + indices, starting with Noll index 4. + nollIndices : np.ndarray + Array of Noll indices to select. Must contain only indices >= 4. + Note these values must be unique, ascending, and >= 4. + + Returns + ------- + np.ndarray + Down-selected values + """ + # Make sure these are arrays + values, nollIndices = np.array(values), np.array(nollIndices) + + # Validate indices + if np.any(nollIndices < 4): + raise ValueError("nollIndices must be >= 4.") + if not np.array_equal(nollIndices, np.sort(np.unique(nollIndices))): + raise ValueError("Values in nollIndices must be unique and ascending.") + + return values[nollIndices - 4] + + +def makeDense( + values: np.ndarray, + nollIndices: np.ndarray, + jmax: int | None = None, +) -> np.ndarray: + """Inserts sparse values into dense array of Zeroes. + + For example, if values=[1, -2, -4], nollIndices=[4, 5, 7], this + function returns [1, -2, 0, -4]. + + Parameters + ---------- + values : np.ndarray + Array of sparse values that will be inserted into the dense array. + nollIndices : np.ndarray + Array of Noll indices to select. Must contain only indices >= 4. + Note these values must be unique, ascending, and >= 4. + jmax : int or None, optional + The maximum Noll index of the dense array. If None, the max value + of nollIndices is used. (the default is None) + + Returns + ------- + np.ndarray + A dense array of values, corresponding to Noll indices 4 - jmax. + Missing values are zero. + """ + # Make sure these are arrays + values, nollIndices = np.array(values), np.array(nollIndices) + + # Validate indices and set jmax + if np.any(nollIndices < 4): + raise ValueError("nollIndices must be >= 4.") + if not np.array_equal(nollIndices, np.sort(np.unique(nollIndices))): + raise ValueError("Values in nollIndices must be unique and ascending.") + jmax = nollIndices.max() if jmax is None else jmax + if jmax < 4: + raise ValueError("jmax must be >= 4.") + + # Remove indices above jmax. + vals = values[nollIndices <= jmax] + idx = nollIndices[nollIndices <= jmax] - 4 + + # Create the dense array + dense = np.zeros_like(np.arange(4, jmax + 1), dtype=values.dtype) + + # Insert values + dense[idx] = vals + + return dense + + +def checkNollIndices(nollIndices: np.ndarray) -> None: + """Check that Noll indices meet requirements. + + Parameters + ---------- + nollIndices : np.ndarray + Array of Noll indices. + + Raises + ------ + ValueError + If nollIndices contains values less than 4, if they're not ascending + and unique, and if azimuthal pairs are not complete. + """ + # Simple checks on values + if any(nollIndices < 4): + raise ValueError("nollIndices must be >= 4.") + if not np.array_equal(nollIndices, np.sort(np.unique(nollIndices))): + raise ValueError("Values in nollIndices must be unique and ascending.") + + # Now we will make sure azimuthal pairs are complete... + + # Create grid of Noll indices from 4 to jmax, as well as az. symm. + # We select jmax that is greater than max value in Noll indices + # and is also azimuthally symmetric. This is so that once we + # downselect to indices without azimuthal symmetry, we are guaranteed + # to have an even number of indices in our grid. + grid = np.array([4]) + az = np.array([0]) + while grid[-1] < nollIndices.max() or az[-1] != 0: + grid = np.append(grid, grid[-1] + 1) + az = np.append(az, galsim.zernike.noll_to_zern(grid[-1])[1]) + + # Remove azimuthally symmetric indices + grid = grid[np.where(az != 0)] + + # Now consecutive Noll indices are azimuthal pairs + # Create mapping between these + paired_grid = grid.reshape(-1, 2) + pairs = {i: j for i, j in paired_grid} | {j: i for i, j in paired_grid} + + # Check all pairs are complete + for j in nollIndices: + if j in pairs and pairs[j] not in nollIndices: + raise ValueError( + f"Noll index {j} is missing azimuthal pair, Noll index {pairs[j]}." + ) diff --git a/tests/estimation/test_tie.py b/tests/estimation/test_tie.py index 40be33b8..fc49d5d7 100644 --- a/tests/estimation/test_tie.py +++ b/tests/estimation/test_tie.py @@ -113,6 +113,7 @@ def testSaveHistory(self): "zkStartExtra", "zkStartMean", "pupil", + "nollIndices", ], ) # All entries should all be arrays diff --git a/tests/estimation/test_wfAlgorithm.py b/tests/estimation/test_wfAlgorithm.py index 0db8e8ca..eb369cb1 100644 --- a/tests/estimation/test_wfAlgorithm.py +++ b/tests/estimation/test_wfAlgorithm.py @@ -90,7 +90,7 @@ def _estimateZk(self, *args, **kwargs) -> np.ndarray: goodSettings = { "I1": intra, "I2": extra, - "jmax": 28, + "nollIndices": np.arange(4, 23), "instrument": Instrument(), "startWithIntrinsic": True, "returnWfDev": False, @@ -150,12 +150,9 @@ def _estimateZk(self, *args, **kwargs) -> np.ndarray: with self.assertRaises(ValueError): wfAlg._validateInputs(**testSettings) - # Test bad jmax + # Test bad Noll indices testSettings = goodSettings.copy() - testSettings["jmax"] = "fake" - with self.assertRaises(TypeError): - wfAlg._validateInputs(**testSettings) - testSettings["jmax"] = 3 + testSettings["nollIndices"] = np.array([4, 5]) with self.assertRaises(ValueError): wfAlg._validateInputs(**testSettings) diff --git a/tests/estimation/test_wfEstimator.py b/tests/estimation/test_wfEstimator.py index aa8dab40..c8703f25 100644 --- a/tests/estimation/test_wfEstimator.py +++ b/tests/estimation/test_wfEstimator.py @@ -47,9 +47,15 @@ def testBadInstConfig(self): with self.assertRaises(FileNotFoundError): WfEstimator(instConfig="fake") - def testBadJmax(self): + def testBadNollIndices(self): with self.assertRaises(ValueError): - WfEstimator(jmax=2) + WfEstimator(nollIndices=[3, 4, 5]) + with self.assertRaises(ValueError): + WfEstimator(nollIndices=[4, 6, 5]) + with self.assertRaises(ValueError): + WfEstimator(nollIndices=[4, 5, 5]) + with self.assertRaises(ValueError): + WfEstimator(nollIndices=[4, 5]) def testBadStartWithIntrinsic(self): with self.assertRaises(TypeError): @@ -59,10 +65,6 @@ def testBadReturnWfDev(self): with self.assertRaises(TypeError): WfEstimator(returnWfDev="fake") - def testBadReturn4Up(self): - with self.assertRaises(TypeError): - WfEstimator(return4Up="fake") - def testBadUnits(self): with self.assertRaises(ValueError): WfEstimator(units="parsecs") @@ -71,22 +73,24 @@ def testBadSaveHistory(self): with self.assertRaises(TypeError): WfEstimator(saveHistory="fake") - def testDifferentJmax(self): + def testDifferentNollIndices(self): # Get the test data zkTrue, intra, extra = forwardModelPair() # Test every wavefront algorithm for name in WfAlgorithmName: - # Estimate with jmax=22 - wfEst = WfEstimator(algoName=name, jmax=22, units="m") - zk22 = wfEst.estimateZk(intra, extra) + # Estimate [4, 5, 6] + wfEst = WfEstimator(algoName=name, nollIndices=[4, 5, 6], units="m") + zk0 = wfEst.estimateZk(intra, extra) + self.assertEqual(len(zk0), 3) - # Estimate with jmax=28 - wfEst = WfEstimator(algoName=name, jmax=28, units="m") - zk28 = wfEst.estimateZk(intra, extra) + # Estimate with [4, 5, 6, 20, 21] + wfEst = WfEstimator(algoName=name, nollIndices=[4, 5, 6, 20, 21], units="m") + zk1 = wfEst.estimateZk(intra, extra) + self.assertEqual(len(zk1), 5) - # Make sure results are pretty similar up to Noll index 22 - self.assertLess(np.sqrt(np.sum(np.square(zk28[:-6] - zk22))), 75e-9) + # Make sure results are pretty similar for [4, 5, 6] + self.assertLess(np.sqrt(np.sum(np.square(zk1[:-2] - zk0))), 10e-9) def testStartWithIntrinsic(self): # Get the test data @@ -122,29 +126,12 @@ def testReturnWfDev(self): # Make sure that OPD = wf dev + intrinsics zkInt = wfEst.instrument.getIntrinsicZernikes( *intra.fieldAngle, - jmax=len(opd) + 3, + nollIndices=np.arange(opd.size) + 4, ) # Make sure the results are identical self.assertTrue(np.allclose(opd, wfDev + zkInt)) - def testReturn4Up(self): - # Get the test data - zkTrue, intra, extra = forwardModelPair() - - # Test every wavefront algorithm - for name in WfAlgorithmName: - # Get estimate starting with Noll index 4 - wfEst = WfEstimator(algoName=name, return4Up=True) - zk4up = wfEst.estimateZk(intra, extra) - - # Get estimate starting with Noll index 0 - wfEst = WfEstimator(algoName=name, return4Up=False) - zk0up = wfEst.estimateZk(intra, extra) - - # Make sure the results are identical - self.assertTrue(np.allclose(zk4up, zk0up[4:])) - def testUnits(self): # Get the test data zkTrue, intra, extra = forwardModelPair() diff --git a/tests/task/test_calcZernikesTieTaskCwfs.py b/tests/task/test_calcZernikesTieTaskCwfs.py index 19e8f351..422a8369 100644 --- a/tests/task/test_calcZernikesTieTaskCwfs.py +++ b/tests/task/test_calcZernikesTieTaskCwfs.py @@ -304,3 +304,33 @@ def testRequireConverge(self): # Everything should be NaN because we did not converge self.assertTrue(np.isnan(zernikes).all()) + + def testNollIndices(self): + # Load the stamps + donutStampDir = os.path.join(self.testDataDir, "donutImg", "donutStamps") + donutStampsExtra = DonutStamps.readFits( + os.path.join(donutStampDir, "R04_SW0_donutStamps.fits") + ) + donutStampsIntra = DonutStamps.readFits( + os.path.join(donutStampDir, "R04_SW1_donutStamps.fits") + ) + + # Estimate Zernikes 4, 5, 6 + self.task.config.estimateZernikes.nollIndices = [4, 5, 6] + zk0 = self.task.estimateZernikes.run( + donutStampsExtra, donutStampsIntra + ).zernikes[0] + + # Estimate Zernikes 4, 5, 6, 20, 21 + self.task.config.estimateZernikes.nollIndices = [4, 5, 6, 20, 21] + zk1 = self.task.estimateZernikes.run( + donutStampsExtra, donutStampsIntra + ).zernikes[0] + + # Check lengths + self.assertEqual(len(zk0), 3) + self.assertEqual(len(zk1), 5) + + # Check that 4, 5, 6 are independent of 20, 21 at less + self.assertLess(np.sqrt(np.sum(np.square(zk1[:-2] - zk0))), 0.020) + # self.assertTrue(np.all(np.abs(zk1[:3] - zk0) < 0.035)) diff --git a/tests/test_imageMapper.py b/tests/test_imageMapper.py index 37f8bda1..f3593540 100644 --- a/tests/test_imageMapper.py +++ b/tests/test_imageMapper.py @@ -674,7 +674,8 @@ def maxPercent(**kwargs): uImage, vImage, *_ = mapper._constructForwardMap( uPupil, vPupil, - inst.getIntrinsicZernikes(*angle, band, jmax=28), + inst.getIntrinsicZernikes(*angle, band), + None, Image(np.zeros((1, 1)), angle, dfType, band), ) diff --git a/tests/test_instrument.py b/tests/test_instrument.py index 439ff985..22fe07ae 100644 --- a/tests/test_instrument.py +++ b/tests/test_instrument.py @@ -108,8 +108,12 @@ def testGetIntrinsicZernikes(self): inst = Instrument() # First check the shape - self.assertEqual(inst.getIntrinsicZernikes(0, 0, jmax=66).shape, (63,)) - self.assertEqual(inst.getIntrinsicZernikes(1, 1.1, jmax=22).shape, (19,)) + self.assertEqual( + inst.getIntrinsicZernikes(0, 0, nollIndices=np.arange(4, 67)).shape, (63,) + ) + self.assertEqual( + inst.getIntrinsicZernikes(1, 1.1, nollIndices=np.arange(4, 23)).shape, (19,) + ) # Now check that in-place changes don't impact the cache intrZk = inst.getIntrinsicZernikes(1, 1) @@ -121,13 +125,23 @@ def testGetOffAxisCoeff(self): inst = Instrument() # First check the shape - self.assertEqual(inst.getOffAxisCoeff(0, 0, "intra", jmax=66).shape, (63,)) - self.assertEqual(inst.getOffAxisCoeff(1, 1.1, "extra", jmax=22).shape, (19,)) + self.assertEqual( + inst.getOffAxisCoeff( + 0, 0, "intra", nollIndicesModel=np.arange(4, 67) + ).shape, + (63,), + ) + self.assertEqual( + inst.getOffAxisCoeff( + 1, 1.1, "extra", nollIndicesModel=np.arange(4, 23) + ).shape, + (19,), + ) # Now check that in-place changes don't impact the cache - intrZk = inst.getOffAxisCoeff(0, 0, "intra") - intrZk *= 3.14159 - close = np.isclose(inst.getOffAxisCoeff(0, 0, "intra"), intrZk, atol=0) + coeff = inst.getOffAxisCoeff(0, 0, "intra") + coeff *= 3.14159 + close = np.isclose(inst.getOffAxisCoeff(0, 0, "intra"), coeff, atol=0) self.assertTrue(np.all(~close)) def testBadMaskParams(self): diff --git a/tests/utils/test_zernikeUtils.py b/tests/utils/test_zernikeUtils.py index da9df9cb..ec25488a 100644 --- a/tests/utils/test_zernikeUtils.py +++ b/tests/utils/test_zernikeUtils.py @@ -24,12 +24,15 @@ import numpy as np from lsst.ts.wep.utils import ( + checkNollIndices, convertZernikesToPsfWidth, createZernikeBasis, createZernikeGradBasis, getModulePath, getPsfGradPerZernike, getZernikeParity, + makeDense, + makeSparse, zernikeEval, zernikeFit, zernikeGradEval, @@ -186,6 +189,72 @@ def testJmin(self): yParity4 = getZernikeParity(jmin=4, axis="y") self.assertTrue(np.allclose(yParity0[4:], yParity4)) + def testSparseDense(self): + vals = np.arange(4, 23) + indices = np.array([4, 5, 6, 17]) + + # Get sparse values + sparse = makeSparse(vals, indices) + np.testing.assert_array_equal(sparse, indices) + + # Make dense without explicit jmax + dense = makeDense(sparse, indices) + np.testing.assert_array_equal(dense[indices - 4], indices) + self.assertEqual(dense[-1], indices[-1]) + + # Make dense with explicit jmax + dense = makeDense(sparse, indices, 22) + np.testing.assert_array_equal(dense[indices - 4], vals[indices - 4]) + self.assertEqual(len(dense), len(vals)) + self.assertEqual(dense[-1], 0) + + # Test fully dense does nothing + np.testing.assert_array_equal(vals, makeSparse(vals, vals)) + np.testing.assert_array_equal(vals, makeDense(vals, vals)) + + # Test that these functions work with floats + floats = np.full_like(vals, 1e-9, dtype=float) + np.testing.assert_array_equal(floats, makeSparse(floats, vals)) + np.testing.assert_array_equal(floats, makeDense(floats, vals)) + + # Test bad indices + for func in [makeSparse, makeDense]: + with self.assertRaises(ValueError): + func([1, 2, 3], [3, 4, 5]) + with self.assertRaises(ValueError): + func([1, 2, 3], [4, 6, 5]) + with self.assertRaises(ValueError): + func([1, 2, 3], [4, 6, 6]) + + def testCheckNollIndices(self): + # These values should all pass + checkNollIndices(np.array([4])) + checkNollIndices(np.array([4, 5, 6])) + checkNollIndices(np.array([20, 21])) + checkNollIndices(np.array([11, 20, 21, 22])) + + # The rest should fail... + + # < 4 + with self.assertRaises(ValueError): + checkNollIndices(np.array([3, 4, 5, 6])) + # not unique + with self.assertRaises(ValueError): + checkNollIndices(np.array([4, 5, 6, 6])) + # not ascending + with self.assertRaises(ValueError): + checkNollIndices(np.array([4, 6, 5])) + + # missing azimuthal pairs + with self.assertRaises(ValueError): + checkNollIndices(np.array([4, 5])) + with self.assertRaises(ValueError): + checkNollIndices(np.array([4, 5, 11])) + with self.assertRaises(ValueError): + checkNollIndices(np.array([4, 5, 20, 21])) + with self.assertRaises(ValueError): + checkNollIndices(np.array([4, 5, 6, 20])) + if __name__ == "__main__": # Do the unit test