From c2cefa4458e188c79e77122b3056326d59404637 Mon Sep 17 00:00:00 2001
From: John Franklin Crenshaw <jfcrenshaw@gmail.com>
Date: Thu, 18 Jan 2024 09:00:56 -0800
Subject: [PATCH] Streamlining zk4up and caches in instrument.

---
 python/lsst/ts/wep/instrument.py | 46 ++++++++++++++++++++++----------
 1 file changed, 32 insertions(+), 14 deletions(-)

diff --git a/python/lsst/ts/wep/instrument.py b/python/lsst/ts/wep/instrument.py
index 6d854580f..6b7783a1d 100644
--- a/python/lsst/ts/wep/instrument.py
+++ b/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,10 +560,7 @@ 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,
@@ -571,6 +568,7 @@ def getIntrinsicZernikes(
         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."""