Streamlining zk4up and caches in instrument.

python/lsst/ts/wep/instrument.py

@@ -536,7 +536,7 @@ def _getIntrinsicZernikesCached(
         Returns
         -------
         np.ndarray
-            The Zernike coefficients in meters, for Noll indices >= 4
+            The Zernike coefficients in meters
         """
         # Get the band enum
         band = BandLabel(band)
@@ -560,17 +560,15 @@ def _getIntrinsicZernikesCached(
         # Multiply by wavelength to get Zernikes in meters
         zkIntrinsic *= self.wavelength[band]
 
-        # Keep only Noll indices >= 4
-        zkIntrinsic = zkIntrinsic[4:]
-
-        return zkIntrinsic.copy()
+        return zkIntrinsic
 
     def getIntrinsicZernikes(
         self,
         xAngle: float,
         yAngle: float,
         band: Union[BandLabel, str] = BandLabel.REF,
         jmax: int = 66,
+        return4Up: bool = True,
     ) -> np.ndarray:
         """Return the intrinsic Zernikes associated with the optical design.
 
@@ -587,13 +585,22 @@ def getIntrinsicZernikes(
         jmax : int, optional
             The maximum Noll index of the intrinsic Zernikes.
             (the default is 66)
+        return4Up : bool, optional
+            Whether to only return the coefficients for Noll indices >= 4.
+            (the default is True)
 
         Returns
         -------
         np.ndarray
-            The Zernike coefficients in meters, for Noll indices >= 4
+            The Zernike coefficients in meters
         """
-        return self._getIntrinsicZernikesCached(xAngle, yAngle, band, jmax).copy()
+        zk = self._getIntrinsicZernikesCached(xAngle, yAngle, band, jmax).copy()
+
+        if return4Up:
+            # Keep only Noll indices >= 4
+            zk = zk[4:]
+
+        return zk
 
     @lru_cache(100)
     def _getOffAxisCoeffCached(
@@ -621,13 +628,13 @@ def _getOffAxisCoeffCached(
             The BandLabel Enum or corresponding string, specifying which
             batoid model to load. Only relevant if self.batoidModelName
             contains "{band}".
-        jmax : int, optional
+        jmax : int
             The maximum Noll index of the off-axis model Zernikes.
 
         Returns
         -------
         np.ndarray
-            The Zernike coefficients in meters, for Noll indices >= 4
+            The Zernike coefficients in meters
         """
         # Get the band enum
         band = BandLabel(band)
@@ -660,9 +667,6 @@ def _getOffAxisCoeffCached(
         # Multiply by wavelength to get Zernikes in meters
         zkIntrinsic *= self.wavelength[band]
 
-        # Keep only Noll indices >= 4
-        zkIntrinsic = zkIntrinsic[4:]
-
         return zkIntrinsic
 
     def getOffAxisCoeff(
@@ -672,6 +676,7 @@ def getOffAxisCoeff(
         defocalType: DefocalType,
         band: Union[BandLabel, str] = BandLabel.REF,
         jmax: int = 66,
+        return4Up: bool = True,
     ) -> np.ndarray:
         """Return the Zernike coefficients associated with the off-axis model.
 
@@ -691,16 +696,29 @@ def getOffAxisCoeff(
         jmax : int, optional
             The maximum Noll index of the off-axis model Zernikes.
             (the default is 66)
+        return4Up : bool, optional
+            Whether to only return the coefficients for Noll indices >= 4.
+            (the default is True)
 
         Returns
         -------
         np.ndarray
             The Zernike coefficients in meters, for Noll indices >= 4
         """
-        return self._getOffAxisCoeffCached(
-            xAngle, yAngle, defocalType, band, jmax
+        zk = self._getOffAxisCoeffCached(
+            xAngle,
+            yAngle,
+            defocalType,
+            band,
+            jmax,
         ).copy()
 
+        if return4Up:
+            # Keep only Noll indices >= 4
+            zk = zk[4:]
+
+        return zk
+
     @property
     def maskParams(self) -> dict:
         """The mask parameter dictionary."""