diff --git a/policy/estimation/tie.yaml b/policy/estimation/tie.yaml index af31d231d..0a1c5c6dd 100644 --- a/policy/estimation/tie.yaml +++ b/policy/estimation/tie.yaml @@ -10,4 +10,3 @@ centerTol: 10.0e-9 centerBinary: True convergeTol: 10.0e-9 maskKwargs: null -saveHistory: False diff --git a/policy/estimation/wfEstimator.yaml b/policy/estimation/wfEstimator.yaml index c96810ef2..169c90d93 100644 --- a/policy/estimation/wfEstimator.yaml +++ b/policy/estimation/wfEstimator.yaml @@ -5,4 +5,8 @@ algoName: tie algoConfig: null instConfig: policy/instruments/LsstCam.yaml jmax: 22 +startWithIntrinsic: true +returnWfDev: false +return4Up: true units: um +saveHistory: false diff --git a/python/lsst/ts/wep/estimation/tie.py b/python/lsst/ts/wep/estimation/tie.py index 315d41626..b42abc44a 100644 --- a/python/lsst/ts/wep/estimation/tie.py +++ b/python/lsst/ts/wep/estimation/tie.py @@ -79,10 +79,6 @@ class TieAlgorithm(WfAlgorithm): Dictionary of mask keyword arguments to pass to mask creation. To see possibilities, see the docstring for lsst.ts.wep.imageMapper.ImageMapper.createPupilMask(). - saveHistory : bool, optional - Whether to save the algorithm history in the self.history attribute. - If True, then self.history contains information about the most recent - time the algorithm was run. """ def __init__( @@ -97,7 +93,6 @@ def __init__( centerBinary: Optional[bool] = None, convergeTol: Optional[float] = None, maskKwargs: Optional[dict] = None, - saveHistory: Optional[bool] = None, ) -> None: super().__init__( configFile=configFile, @@ -110,7 +105,6 @@ def __init__( centerBinary=centerBinary, convergeTol=convergeTol, maskKwargs=maskKwargs, - saveHistory=saveHistory, ) @property @@ -374,23 +368,38 @@ def history(self) -> dict: """The algorithm history. The history is a dictionary saving the intermediate products from - each iteration of the TIE solver. The first iteration is saved as - history[0]. - - The entry for each iteration is itself a dictionary containing + each iteration of the TIE solver. + + The initial products before the iteration begins are stored + in history[0], which is a dictionary with the keys: + - "intraInit" - the initial intrafocal image + - "extraInit" - the initial extrafocal image + - "zkStartIntra" - the starting intrafocal Zernikes + - "zkStartExtra" - the starting extrafocal Zernikes + - "zkStartMean" - the mean of the starting Zernikes. Note these + Zernikes are added to zkBest to estimate the + full OPD. + + Each iteration of the solver is then stored under indices >= 1. + The entry for each iteration is also a dictionary, containing the following keys: - "intraComp" - the compensated intrafocal image - "extraComp" - the compensated extrafocal image - "I0" - the estimate of the beam intensity on the pupil - "dIdz" - estimate of z-derivative of intensity across the pupil - - "zkComp" - the Zernikes used for image compensation + - "zkCompIntra" - Zernikes for compensating the intrafocal image + - "zkCompExtra" - Zernikes for compensating the extrafocal image - "zkResid" - the estimated residual Zernikes - - "zkBest" - the best estimate of the Zernikes after this iteration + - "zkBest" - the best cumulative estimate the wavefront residual. + - "zkSum" - the sum of zkBest and zkStartMean from history[0]. + This is the best estimate of the OPD at the end of + this iteration. - "converged" - flag indicating if Zernike estimation has converged - "caustic" - flag indicating if a caustic has been hit Note the units for all Zernikes are in meters, and the z-derivative - in dIdz is also in meters. + in dIdz is also in meters. Furthermore, all Zernikes start with Noll + index 4. """ return super().history @@ -472,12 +481,15 @@ def _fftSolve( # TODO: Implement the fft solver raise NotImplementedError("The fft solver is not yet implemented.") - def estimateZk( + def _estimateZk( self, I1: Image, I2: Image, # type: ignore[override] - jmax: int = 22, - instrument: Instrument = Instrument(), + zkStartI1: np.ndarray, + zkStartI2: np.ndarray, + jmax: int, + instrument: Instrument, + saveHistory: bool, ) -> np.ndarray: """Return the wavefront Zernike coefficients in meters. @@ -487,26 +499,36 @@ def estimateZk( An Image object containing an intra- or extra-focal donut image. I2 : Image A second image, on the opposite side of focus from I1. - jmax : int, optional + zkStartI1 : np.ndarray + The starting Zernikes for I1. + zkStartI2 : np.ndarray or None + The starting Zernikes for I2. + jmax : int The maximum Zernike Noll index to estimate. - (the default is 22) - instrument : Instrument, optional + instrument : Instrument The Instrument object associated with the DonutStamps. - (the default is the default Instrument) + saveHistory : bool + Whether to save the algorithm history in the self.history + attribute. If True, then self.history contains information + about the most recent time the algorithm was run. Returns ------- np.ndarray Zernike coefficients (for Noll indices >= 4) estimated from the images, in meters. + + Raises + ------ + RuntimeError + If the solver is not supported """ - # Validate the inputs + # Make sure we have been provided with two images if I1 is None or I2 is None: raise ValueError( "TIEAlgorithm requires a pair of intrafocal and extrafocal " "donuts to estimate Zernikes. Please provide both I1 and I2." ) - self._validateInputs(I1, I2, jmax, instrument) # Create the ImageMapper for centering and image compensation imageMapper = ImageMapper( @@ -515,19 +537,35 @@ def estimateZk( opticalModel=self.opticalModel, ) - # Get the initial intrafocal and extrafocal stamps + # Re-assign I1/I2 to intra/extra intra = I1.copy() if I1.defocalType == DefocalType.Intra else I2.copy() + zkStartIntra = ( + zkStartI1.copy() + if I1.defocalType == DefocalType.Intra + else zkStartI2.copy() + ) extra = I1.copy() if I1.defocalType == DefocalType.Extra else I2.copy() + zkStartExtra = ( + zkStartI1.copy() + if I1.defocalType == DefocalType.Extra + else zkStartI2.copy() + ) + + # Calculate the mean starting Zernikes + zkStartMean = np.mean([zkStartIntra, zkStartExtra], axis=0) # Initialize the variables for intra and extrafocal pupil masks intraPupilMask = None extraPupilMask = None - if self.saveHistory: - # Save the initial images in the history + if saveHistory: + # Save the initial images and intrinsic Zernikes in the history self._history[0] = { "intraInit": intra.image.copy(), "extraInit": extra.image.copy(), + "zkStartIntra": zkStartIntra.copy(), + "zkStartExtra": zkStartExtra.copy(), + "zkStartMean": zkStartMean.copy(), } # Initialize Zernike arrays at zero @@ -562,13 +600,13 @@ def estimateZk( zkCenter = zkComp.copy() intraCent = imageMapper.centerOnProjection( intra, - zkCenter, + zkCenter + zkStartIntra, binary=self.centerBinary, **self.maskKwargs, ) extraCent = imageMapper.centerOnProjection( extra, - zkCenter, + zkCenter + zkStartExtra, binary=self.centerBinary, **self.maskKwargs, ) @@ -576,13 +614,13 @@ def estimateZk( # Compensate images using the Zernikes intraComp = imageMapper.mapImageToPupil( intraCent, - zkComp, + zkComp + zkStartIntra, mask=intraPupilMask, **self.maskKwargs, ) extraComp = imageMapper.mapImageToPupil( extraCent, - zkComp, + zkComp + zkStartExtra, mask=extraPupilMask, **self.maskKwargs, ) @@ -625,6 +663,8 @@ def estimateZk( zkResid = self._expSolve(I0, dIdz, jmax, instrument) elif self.solver == "fft": zkResid = self._fftSolve(I0, dIdz, jmax, instrument) + else: + raise RuntimeError(f"Solver {self.solver} is unsupported") # Check for convergence # (1) The max absolute difference with the previous iteration @@ -635,12 +675,15 @@ def estimateZk( np.max(np.abs(newBest - zkBest)) < self.convergeTol ) - # Set the new best estimate + # Set the new best cumulative estimate of the residuals zkBest = newBest + # Add the starting Zernikes to the best residuals + zkSum = zkBest + zkStartMean + # Time to wrap up this iteration! # Should we save intermediate products in the algorithm history? - if self.saveHistory: + if saveHistory: # Save the images and Zernikes from this iteration self._history[i + 1] = { "recenter": bool(recenter), @@ -651,9 +694,11 @@ def estimateZk( "mask": mask.copy(), # type: ignore "I0": I0.copy(), "dIdz": dIdz.copy(), - "zkComp": zkComp.copy(), + "zkCompIntra": zkComp + zkStartIntra, + "zkCompExtra": zkComp + zkStartExtra, "zkResid": zkResid.copy(), "zkBest": zkBest.copy(), + "zkSum": zkSum.copy(), "converged": bool(converged), "caustic": bool(caustic), } @@ -661,10 +706,10 @@ def estimateZk( # If we are using the FFT solver, save the inner loop as well if self.solver == "fft": # TODO: After implementing fft, add inner loop here - self._history[i]["innerLoop"] = None + self._history[i + 1]["innerLoop"] = None # If we've hit a caustic or converged, we will stop early if caustic or converged: break - return zkBest + return zkSum diff --git a/python/lsst/ts/wep/estimation/wfAlgorithm.py b/python/lsst/ts/wep/estimation/wfAlgorithm.py index 971c1948f..2b4579141 100644 --- a/python/lsst/ts/wep/estimation/wfAlgorithm.py +++ b/python/lsst/ts/wep/estimation/wfAlgorithm.py @@ -27,7 +27,7 @@ import numpy as np from lsst.ts.wep import Image, Instrument -from lsst.ts.wep.utils import mergeConfigWithFile +from lsst.ts.wep.utils import mergeConfigWithFile, convertZernikesToPsfWidth class WfAlgorithm(ABC): @@ -39,10 +39,6 @@ class WfAlgorithm(ABC): Path to file specifying values for the other parameters. If the path starts with "policy/", it will look in the policy directory. Any explicitly passed parameters override values found in this file - saveHistory : bool, optional - Whether to save the algorithm history in the self.history attribute. - If True, then self.history contains information about the most recent - time the algorithm was run. ... @@ -51,13 +47,11 @@ class WfAlgorithm(ABC): def __init__( self, configFile: Optional[str] = None, - saveHistory: Optional[bool] = None, **kwargs: Any, ) -> None: # Merge keyword arguments with defaults from configFile params = mergeConfigWithFile( configFile, - saveHistory=saveHistory, **kwargs, ) @@ -66,7 +60,7 @@ def __init__( setattr(self, key, value) # Instantiate an empty history - self._history = {} # type: ignore + self._history = {} def __init_subclass__(self) -> None: """This is called when you subclass. @@ -80,44 +74,15 @@ def __init_subclass__(self) -> None: "Please use this to describe the contents of the history dict." ) - @property - def saveHistory(self) -> bool: - """Whether the algorithm history is saved.""" - return self._saveHistory - - @saveHistory.setter - def saveHistory(self, value: bool) -> None: - """Set boolean that determines whether algorithm history is saved. - - Parameters - ---------- - value : bool - Boolean that determines whether the algorithm history is saved. - - Raises - ------ - TypeError - If the value is not a boolean - """ - if not isinstance(value, bool): - raise TypeError("saveHistory must be a boolean.") - - self._saveHistory = value - - # If we are turning history-saving off, delete any old history - # This is to avoid confusion - if value is False: - self._history = {} - @property def history(self) -> dict: # Return the algorithm history # Note I have not written a real docstring here, so that I can force # subclasses to write a new docstring for this method - if not self._saveHistory: + if len(self._history) == 0: warnings.warn( - "saveHistory is False. If you want the history to be saved, " - "run self.config(saveHistory=True)." + "It looks like the history is empty. Perhaps you have not " + "yet run the algorithm with saveHistory=True?" ) return self._history @@ -126,8 +91,12 @@ def history(self) -> dict: def _validateInputs( I1: Image, I2: Optional[Image], - jmax: int = 28, - instrument: Instrument = Instrument(), + jmax: int, + instrument: Instrument, + startWithIntrinsic: bool, + returnWfDev: bool, + units: str, + saveHistory: bool, ) -> None: """Validate the inputs to estimateWf. @@ -139,17 +108,26 @@ def _validateInputs( A second image, on the opposite side of focus from I1. jmax : int, optional The maximum Zernike Noll index to estimate. - (the default is 28) instrument : Instrument, optional The Instrument object associated with the DonutStamps. - (the default is the default Instrument) + startWithIntrinsic : bool, optional + Whether to start the Zernike estimation process from the intrinsic + Zernikes rather than zero. + returnWfDev : bool, optional + 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 in which the Zernike amplitudes are returned. + Options are "m", "nm", "um", or "arcsecs". Raises ------ TypeError If any input is the wrong type ValueError - If I1 or I2 are not square arrays, or if jmax < 4 + If I1 or I2 are not square arrays, or if jmax < 4, + or unsupported unit """ # Validate I1 if not isinstance(I1, Image): @@ -178,13 +156,80 @@ def _validateInputs( if not isinstance(instrument, Instrument): raise TypeError("instrument must be an Instrument.") + # Validate startWithIntrinsic + if not isinstance(startWithIntrinsic, bool): + raise TypeError("startWithIntrinsic must be a bool.") + + # Validate returnWfDev + if not isinstance(returnWfDev, bool): + raise TypeError("returnWfDev must be a bool.") + + # Validate units + if not isinstance(units, str): + raise TypeError("units must be a str.") + allowedUnits = ["m", "nm", "um", "arcseconds"] + if units not in allowedUnits: + raise ValueError(f"units must be one of {allowedUnits}") + + # Validate saveHistory + if not isinstance(saveHistory, bool): + raise TypeError("saveHistory must be a bool.") + @abstractmethod - def estimateZk( + def _estimateZk( self, I1: Image, I2: Optional[Image], - jmax: int = 28, + zkStartI1: np.ndarray, + zkStartI2: Optional[np.ndarray], + jmax: int, + instrument: Instrument, + saveHistory: bool, + ) -> np.ndarray: + """Private Zernike estimation method that should be subclassed. + + Note that unlike the public method, this method MUST return the + coefficients for Noll indices >= 4 (in meters) corresponding to + the full OPD. + + Parameters + ---------- + I1 : DonutStamp + An Image object containing an intra- or extra-focal donut image. + I2 : DonutStamp or None + A second image, on the opposite side of focus from I1. + zkStartI1 : np.ndarray + The starting Zernikes for I1. + zkStartI2 : np.ndarray or None + The starting Zernikes for I2. + jmax : int + The maximum Zernike Noll index to estimate. + instrument : Instrument + The Instrument object associated with the DonutStamps. + saveHistory : bool + Whether to save the algorithm history in the self.history + attribute. If True, then self.history contains information + about the most recent time the algorithm was run. + + Returns + ------- + np.ndarray + Zernike coefficients representing the OPD in meters, + for Noll indices >= 4. + """ + ... + + def estimateZk( + self, + I1: Image, + I2: Optional[Image] = None, + jmax: int = 22, instrument: Instrument = Instrument(), + startWithIntrinsic: bool = True, + returnWfDev: bool = False, + return4Up: bool = True, + units: str = "m", + saveHistory: bool = False, ) -> np.ndarray: """Return the wavefront Zernike coefficients in meters. @@ -194,17 +239,117 @@ def estimateZk( An Image object containing an intra- or extra-focal donut image. 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 28) + (the default is 22) instrument : Instrument, optional The Instrument object associated with the DonutStamps. (the default is the default Instrument) + startWithIntrinsic : bool, optional + Whether to start the Zernike estimation process from the intrinsic + Zernikes rather than zero. (the default is True) + returnWfDev : bool, optional + 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". + (the default is "m") + saveHistory : bool, optional + Whether to save the algorithm history in the self.history + attribute. If True, then self.history contains information + about the most recent time the algorithm was run. + (the default is False) Returns ------- np.ndarray - Zernike coefficients (for Noll indices >= 4) estimated from - the image (or pair of images), in meters. + Zernike coefficients estimated from the image (or pair of images) """ - ... + # Validate the inputs + self._validateInputs( + I1, + I2, + jmax, + instrument, + startWithIntrinsic, + returnWfDev, + units, + saveHistory, + ) + + # Get the intrinsic Zernikes? + if startWithIntrinsic or returnWfDev: + zkIntrinsicI1 = instrument.getIntrinsicZernikes( + *I1.fieldAngle, + I1.bandLabel, + jmax, + ) + zkIntrinsicI2 = ( + None + if I2 is None + else instrument.getIntrinsicZernikes(*I2.fieldAngle, I2.bandLabel, jmax) + ) + + # Determine the Zernikes to start with + if startWithIntrinsic: + zkStartI1 = zkIntrinsicI1 + zkStartI2 = zkIntrinsicI2 + else: + zkStartI1 = np.zeros(jmax - 3) + zkStartI2 = np.zeros(jmax - 3) + + # Clear the algorithm history + self._history = {} + + # Estimate the Zernikes + zk = self._estimateZk( + I1=I1, + I2=I2, + zkStartI1=zkStartI1, + zkStartI2=zkStartI2, + jmax=jmax, + instrument=instrument, + saveHistory=saveHistory, + ) + + # Calculate the wavefront deviation? + if returnWfDev: + zkIntrinsic = ( + zkIntrinsicI1 + if I2 is None + else np.mean([zkIntrinsicI1, zkIntrinsicI2], axis=0) + ) + zk -= zkIntrinsic + + # Convert to desired units + if units == "m": + pass + elif units == "um": + zk *= 1e6 + elif units == "nm": + zk *= 1e9 + elif units == "arcsecs": + zk = convertZernikesToPsfWidth( + zernikes=zk, + diameter=self.instrument.diameter, + obscuration=self.instrument.obscuration, + ) + else: + 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 c9e8707a8..033bf35ef 100644 --- a/python/lsst/ts/wep/estimation/wfEstimator.py +++ b/python/lsst/ts/wep/estimation/wfEstimator.py @@ -29,7 +29,6 @@ from lsst.ts.wep.estimation.wfAlgorithmFactory import WfAlgorithmFactory from lsst.ts.wep.utils import ( configClass, - convertZernikesToPsfWidth, mergeConfigWithFile, ) @@ -63,9 +62,28 @@ class WfEstimator: If an Instrument object, that object is just used. jmax : int, optional The maximum Zernike Noll index to estimate. + startWithIntrinsic : bool, optional + Whether to start the Zernike estimation process from the intrinsic + Zernikes rather than zero. + returnWfDev : bool, optional + 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. units : str, optional Units in which the Zernike amplitudes are returned. Options are "m", "nm", "um", or "arcsecs". + saveHistory : bool, optional + Whether to save the algorithm history in the self.history + attribute. If True, then self.history contains information + about the most recent time the algorithm was run. + (the default is False) """ def __init__( @@ -75,7 +93,11 @@ def __init__( algoConfig: Union[str, dict, WfAlgorithm, None] = None, instConfig: Union[str, dict, Instrument, None] = None, jmax: Optional[int] = None, + startWithIntrinsic: Optional[bool] = None, + returnWfDev: Optional[bool] = None, + return4Up: Optional[bool] = None, units: Optional[str] = None, + saveHistory: Optional[bool] = None, ) -> None: # Merge keyword arguments with defaults from configFile params = mergeConfigWithFile( @@ -84,7 +106,11 @@ def __init__( algoConfig=algoConfig, instConfig=instConfig, jmax=jmax, + startWithIntrinsic=startWithIntrinsic, + returnWfDev=returnWfDev, + return4Up=return4Up, units=units, + saveHistory=saveHistory, ) # Set the algorithm @@ -97,7 +123,14 @@ def __init__( # Set the other parameters self.jmax = params["jmax"] + self.startWithIntrinsic = params["startWithIntrinsic"] + self.returnWfDev = params["returnWfDev"] + self.return4Up = params["return4Up"] self.units = params["units"] + self.saveHistory = params["saveHistory"] + + # Copy the history docstring + self.__class__.history.__doc__ = self.algo.__class__.history.__doc__ @property def algo(self) -> WfAlgorithm: @@ -116,12 +149,100 @@ def jmax(self) -> int: @jmax.setter def jmax(self, value: int) -> None: - """Set jmax""" + """Set jmax + + Parameters + ---------- + value : int + The maximum Zernike Noll index to estimate. + + Raises + ------ + ValueError + If jmax is < 4 + """ value = int(value) if value < 4: raise ValueError("jmax must be greater than or equal to 4.") self._jmax = value + @property + def startWithIntrinsic(self) -> bool: + """Whether to start Zernike estimation with the intrinsics.""" + return self._startWithIntrinsic + + @startWithIntrinsic.setter + def startWithIntrinsic(self, value: bool) -> None: + """Set startWithIntrinsic. + + Parameters + ---------- + value : bool + Whether to start the Zernike estimation process from the intrinsic + Zernikes rather than zero. + + Raises + ------ + TypeError + If the value is not a boolean + """ + if not isinstance(value, bool): + raise TypeError("startWithIntrinsic must be a bool.") + self._startWithIntrinsic = value + + @property + def returnWfDev(self) -> bool: + """Whether to return the wavefront deviation instead of the OPD.""" + return self._returnWfDev + + @returnWfDev.setter + def returnWfDev(self, value: bool) -> None: + """Set returnWfDev. + + Parameters + ---------- + value : 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. + + Raises + ------ + TypeError + If the value is not a boolean + """ + if not isinstance(value, bool): + 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. @@ -141,6 +262,38 @@ def units(self, value: str) -> None: ) self._units = value + @property + def saveHistory(self) -> bool: + """Whether to save the algorithm history.""" + return self._saveHistory + + @saveHistory.setter + def saveHistory(self, value: bool) -> None: + """Set saveHistory. + + Parameters + ---------- + value : bool + Whether to save the algorithm history in the self.history + attribute. If True, then self.history contains information + about the most recent time the algorithm was run. + + Raises + ------ + TypeError + If the value is not a boolean + """ + if not isinstance(value, bool): + raise TypeError("saveHistory must be a bool.") + self._saveHistory = value + + @property + def history(self) -> dict: + # Return the algorithm history + # This does not have a real docstring, because the dosctring from + # the algorithm history is added during the class __init__ above + return self.algo.history + def estimateZk( self, I1: Image, @@ -168,23 +321,14 @@ def estimateZk( ValueError If I1 and I2 are on the same side of focus. """ - # Estimate wavefront Zernike coefficients (in meters) - zk = self.algo.estimateZk(I1, I2, self.jmax, self.instrument) - - # Convert to desired units - if self.units == "m": - pass - elif self.units == "um": - zk *= 1e6 - elif self.units == "nm": - zk *= 1e9 - elif self.units == "arcsecs": - zk = convertZernikesToPsfWidth( - zernikes=zk, - diameter=self.instrument.diameter, - obscuration=self.instrument.obscuration, - ) - else: - raise RuntimeError(f"Conversion to unit '{self.units}' not supported.") - - return zk + return self.algo.estimateZk( + I1=I1, + I2=I2, + jmax=self.jmax, + 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 8071784a3..91a68316d 100644 --- a/python/lsst/ts/wep/imageMapper.py +++ b/python/lsst/ts/wep/imageMapper.py @@ -25,7 +25,7 @@ import numpy as np from lsst.ts.wep.image import Image from lsst.ts.wep.instrument import Instrument -from lsst.ts.wep.utils.enumUtils import DefocalType, PlaneType +from lsst.ts.wep.utils.enumUtils import DefocalType, PlaneType, BandLabel from lsst.ts.wep.utils.ioUtils import configClass, mergeConfigWithFile from lsst.ts.wep.utils.miscUtils import centerWithTemplate, polygonContains from lsst.ts.wep.utils.zernikeUtils import zernikeGradEval @@ -147,27 +147,31 @@ def _constructForwardMap( The Jacobian of the forward map np.ndarray The determinant of the Jacobian + + Raises + ------ + RuntimeWarning + If the optical model is not supported """ - # Get the reference Zernikes + # Get the Zernikes for the mapping if self.opticalModel == "onAxis": - zkRef = self.instrument.getIntrinsicZernikes( # type: ignore - *image.fieldAngle, - image.bandLabel, - ) - else: - zkRef = self.instrument.getOffAxisCoeff( # type: ignore + zkMap = zkCoeff + elif self.opticalModel == "offAxis": + # Get the off-axis coefficients + offAxisCoeff = self.instrument.getOffAxisCoeff( *image.fieldAngle, image.defocalType, image.bandLabel, + jmaxIntrinsic=len(zkCoeff) + 3, ) - # Create an array for the Zernike coefficients used for the mapping - size = max(zkCoeff.size, zkRef.size) - zkMap = np.zeros(size) - - # Add the initial and reference Zernikes - zkMap[: zkCoeff.size] = zkCoeff - zkMap[: zkRef.size] += zkRef + # Add these coefficients to the input coefficients + size = max(zkCoeff.size, offAxisCoeff.size) + zkMap = np.zeros(size) + zkMap[: zkCoeff.size] = zkCoeff + zkMap[: offAxisCoeff.size] += offAxisCoeff + else: + raise RuntimeError(f"Optical model {self.opticalModel} not supported") # Calculate all 1st- and 2nd-order Zernike derivatives d1Wdu = zernikeGradEval( @@ -864,7 +868,7 @@ def createPupilMask( def createImageMask( self, image: Image, - zkCoeff: np.ndarray = np.zeros(1), + zkCoeff: Optional[np.ndarray] = None, *, binary: bool = True, dilate: int = 0, @@ -879,9 +883,10 @@ def createImageMask( image : Image A stamp object containing the metadata required for constructing the mask. - zkCoeff : np.ndarray + zkCoeff : np.ndarray, optional 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) binary : bool, optional Whether to return a binary mask. If False, a fractional mask is returned instead. (the default is True) @@ -910,6 +915,13 @@ def createImageMask( elif dilate > 0 and not binary: raise ValueError("If dilate is greater than zero, binary must be True.") + if zkCoeff is None: + # Get the intrinsic Zernikes + zkCoeff = self.instrument.getIntrinsicZernikes( + *image.fieldAngle, + image.bandLabel, + ) + # Get the image grid inside the pupil uImage, vImage, inside = self._getImageGridInsidePupil(zkCoeff, image) @@ -976,7 +988,8 @@ def getProjectionSize( self, fieldAngle: Union[np.ndarray, tuple, list], defocalType: Union[DefocalType, str], - zkCoeff: np.ndarray = np.zeros(1), + bandLabel: Union[BandLabel, str] = BandLabel.REF, + zkCoeff: Optional[np.ndarray] = None, ) -> int: """Return size of the pupil projected onto the image plane (in pixels). @@ -993,10 +1006,15 @@ def getProjectionSize( defocalType : DefocalType or str Whether the image is intra- or extra-focal. Can be specified using a DefocalType Enum or the corresponding string. + bandLabel : BandLabel or str + Photometric band for the exposure. Can be specified using a + BandLabel Enum or the corresponding string. If None, BandLabel.REF + is used. The empty string "" also maps to BandLabel.REF. + (the default is BandLabel.REF) zkCoeff : np.ndarray, optional - The wavefront deviation at the pupil, represented as Zernike - coefficients in meters, for Noll indices >= 4. - (the default is zero) + 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) Returns ------- @@ -1008,8 +1026,16 @@ def getProjectionSize( image=np.zeros((0, 0)), fieldAngle=fieldAngle, defocalType=defocalType, + bandLabel=bandLabel, ) + if zkCoeff is None: + # Get the intrinsic Zernikes + zkCoeff = self.instrument.getIntrinsicZernikes( + *dummyImage.fieldAngle, + dummyImage.bandLabel, + ) + # Project the pupil onto the image plane theta = np.linspace(0, 2 * np.pi, 100) uPupil, vPupil = np.cos(theta), np.sin(theta) @@ -1032,7 +1058,7 @@ def getProjectionSize( def centerOnProjection( self, image: Image, - zkCoeff: np.ndarray = np.zeros(1), + zkCoeff: Optional[np.ndarray] = None, binary: bool = True, rMax: float = 10, **maskKwargs, @@ -1048,9 +1074,9 @@ def centerOnProjection( image : Image A stamp object containing the metadata needed for the mapping. zkCoeff : np.ndarray, optional - The wavefront deviation at the pupil, represented as Zernike - coefficients in meters, for Noll indices >= 4. - (the default is zero) + 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) binary : 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 @@ -1063,6 +1089,13 @@ def centerOnProjection( # Make a copy of the stamp stamp = image.copy() + if zkCoeff is None: + # Get the intrinsic Zernikes + zkCoeff = self.instrument.getIntrinsicZernikes( + *image.fieldAngle, + image.bandLabel, + ) + # Create the image template if binary: template = self.createImageMask(stamp, zkCoeff, binary=True, **maskKwargs) @@ -1077,7 +1110,7 @@ def centerOnProjection( def mapPupilToImage( self, image: Image, - zkCoeff: np.ndarray = np.zeros(1), + zkCoeff: Optional[np.ndarray] = None, mask: Optional[np.ndarray] = None, **maskKwargs, ) -> Image: @@ -1095,9 +1128,9 @@ def mapPupilToImage( It is assumed that mapping the pupil to the image plane is meant to model the image contained in this stamp. zkCoeff : np.ndarray, optional - The wavefront deviation at the pupil, represented as Zernike - coefficients in meters, for Noll indices >= 4. - (the default is zero) + 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) mask : np.ndarray, optional You can provide an image mask if you have already computed one. This is just to speed up computation. If not provided, the image @@ -1111,6 +1144,13 @@ def mapPupilToImage( # Make a copy of the stamp stamp = image.copy() + if zkCoeff is None: + # Get the intrinsic Zernikes + zkCoeff = self.instrument.getIntrinsicZernikes( + *image.fieldAngle, + image.bandLabel, + ) + # Get the image grid inside the pupil uImage, vImage, inside = self._getImageGridInsidePupil(zkCoeff, stamp) @@ -1149,7 +1189,7 @@ def mapPupilToImage( def mapImageToPupil( self, image: Image, - zkCoeff: np.ndarray = np.zeros(1), + zkCoeff: Optional[np.ndarray] = None, mask: Optional[np.ndarray] = None, **maskKwargs, ) -> Image: @@ -1168,7 +1208,7 @@ def mapImageToPupil( zkCoeff : np.ndarray, optional The wavefront at the pupil, represented as Zernike coefficients in meters, for Noll indices >= 4. - (the default is zero) + (the default are the intrinsic Zernikes at the donut position) mask : np.ndarray, optional You can provide a pupil mask if you have already computed one. This is just to speed up computation. If not provided, the pupil @@ -1186,6 +1226,13 @@ def mapImageToPupil( uPupil, vPupil = self.instrument.createPupilGrid() uImage, vImage = self.instrument.createImageGrid(stamp.image.shape[0]) + if zkCoeff is None: + # Get the intrinsic Zernikes + zkCoeff = self.instrument.getIntrinsicZernikes( + *image.fieldAngle, + image.bandLabel, + ) + # Construct the forward mapping uImageMap, vImageMap, jac, jacDet = self._constructForwardMap( uPupil, diff --git a/python/lsst/ts/wep/instrument.py b/python/lsst/ts/wep/instrument.py index 999543d54..0197d5f9d 100644 --- a/python/lsst/ts/wep/instrument.py +++ b/python/lsst/ts/wep/instrument.py @@ -362,7 +362,7 @@ def wavelength(self, value: Union[float, dict]) -> None: # Also clear the caches for the functions that use this value self._getIntrinsicZernikesCached.cache_clear() - self._getOffAxisCoeffCached.cache_clear() + self._getIntrinsicZernikesTACached.cache_clear() @property def batoidModelName(self) -> Union[str, None]: @@ -405,7 +405,7 @@ def batoidModelName(self, value: Optional[str]) -> None: # Also clear the caches for the functions that use this value self.getBatoidModel.cache_clear() self._getIntrinsicZernikesCached.cache_clear() - self._getOffAxisCoeffCached.cache_clear() + self._getIntrinsicZernikesTACached.cache_clear() @property def batoidOffsetOptic(self) -> Union[str, None]: @@ -603,7 +603,7 @@ def getIntrinsicZernikes( return zk @lru_cache(100) - def _getOffAxisCoeffCached( + def _getIntrinsicZernikesTACached( self, xAngle: float, yAngle: float, @@ -611,9 +611,7 @@ def _getOffAxisCoeffCached( band: Union[BandLabel, str], jmax: int, ) -> np.ndarray: - """Cached interior function for the getOffAxisCoeff method. - - We need to do this because numpy arrays are mutable. + """Cached function for batoid.zernikeTA. Parameters ---------- @@ -676,6 +674,7 @@ def getOffAxisCoeff( defocalType: DefocalType, band: Union[BandLabel, str] = BandLabel.REF, jmax: int = 66, + jmaxIntrinsic: int = 66, return4Up: bool = True, ) -> np.ndarray: """Return the Zernike coefficients associated with the off-axis model. @@ -696,6 +695,12 @@ def getOffAxisCoeff( jmax : int, optional The maximum Noll index of the off-axis model Zernikes. (the default is 66) + 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. return4Up : bool, optional Whether to only return the coefficients for Noll indices >= 4. (the default is True) @@ -705,19 +710,32 @@ def getOffAxisCoeff( np.ndarray The Zernike coefficients in meters, for Noll indices >= 4 """ - zk = self._getOffAxisCoeffCached( + # Get zernikeTA + zkTA = self._getIntrinsicZernikesTACached( xAngle, yAngle, defocalType, band, jmax, - ).copy() + ) + + # Get regular intrinsic zernikes + zk = self._getIntrinsicZernikesCached( + xAngle, + yAngle, + band, + min(jmax, jmaxIntrinsic), + ) + + # Subtract the intrinsics from zernikeTA + offAxisCoeff = zkTA.copy() + offAxisCoeff[: zk.size] -= zk if return4Up: # Keep only Noll indices >= 4 - zk = zk[4:] + offAxisCoeff = offAxisCoeff[4:] - return zk + return offAxisCoeff @property def maskParams(self) -> dict: