Tickets/DM-49286: sparse Noll Zernike mode

bin.src/img_closed_loop.py

@@ -64,5 +64,5 @@
         args.raw_seeing,
         args.imsim_log_file,
         args.wep_estimator,
-        args.max_noll_index,
+        args.noll_indices,
     )

doc/versionHistory.rst

@@ -5,6 +5,13 @@
 ##################
 Version History
 ##################
+
+-------------
+1.7.0
+-------------
+
+* Implement sparse Zernike mode, supersede max_noll_index with noll_indices.
+
 -------------
 1.6.5
 -------------

python/lsst/ts/imsim/closed_loop_task.py

@@ -47,7 +47,7 @@
     plot_fwhm_of_iters,
 )
 from lsst.ts.ofc import OFC, OFCData
-from lsst.ts.wep.utils import rotMatrix, runProgram
+from lsst.ts.wep.utils import makeDense, rotMatrix, runProgram
 
 
 class ClosedLoopTask:
@@ -66,8 +66,8 @@ def __init__(self) -> None:
         # imSim Component
         self.imsim_cmpt = None
 
-        # Maximum Noll index
-        self.max_noll_index = None
+        # Noll Zernike indices
+        self.noll_indices = None
 
         # Ra/Dec/RotAng coordinates used in the simulation.
         self.boresight_ra = None
@@ -205,10 +205,8 @@ def config_ofc_calc(self, cam_type: CamType, controller_config_file: str) -> Non
             Path(get_config_dir()) / "controllers" / controller_config_file
         )
         self.ofc_calc = OFC(ofc_data)
-        self.ofc_calc.ofc_data.znmin = 4
-        self.ofc_calc.ofc_data.zn_selected = np.arange(
-            self.ofc_calc.ofc_data.znmin, self.max_noll_index + 1
-        )
+        self.ofc_calc.ofc_data.znmin = min(self.noll_indices)
+        self.ofc_calc.ofc_data.zn_selected = np.array(self.noll_indices)
 
     def map_filter_ref_to_g(self, filter_type_name: str) -> str:
         """Map the reference filter to the G filter.
@@ -525,9 +523,14 @@ def _run_sim(
                 zk_file_name=wfs_zk_file_name,
             )
 
-            # Calculate the DOF
+            # Calculate the DOF.
+            # Need to switch to dense representation
+            # before passing on to OFC
             wfe = np.array(
-                [sensor_wfe.annular_zernike_poly for sensor_wfe in list_of_wf_err]
+                [
+                    makeDense(sensor_wfe.annular_zernike_poly, self.noll_indices)
+                    for sensor_wfe in list_of_wf_err
+                ]
             )
 
             sensor_ids = np.array(
@@ -886,7 +889,7 @@ def run_wep(
         butler = dafButler.Butler(butler_root_path)
 
         dataset_refs = butler.registry.queryDatasets(
-            datasetType="zernikeEstimateAvg", collections=[f"ts_imsim_{seq_num}"]
+            datasetType="zernikes", collections=[f"ts_imsim_{seq_num}"]
         )
 
         # Get the map for detector Id to detector name
@@ -907,19 +910,26 @@ def run_wep(
                 "visit": dataset.dataId["visit"],
             }
 
-            zer_coeff = butler.get(
-                "zernikeEstimateAvg",
+            zernikes = butler.get(
+                "zernikes",
                 dataId=data_id,
                 collections=[f"ts_imsim_{seq_num}"],
             )
-
-            sensor_wavefront_data = SensorWavefrontError(
-                num_of_zk=self.max_noll_index - self.ofc_calc.ofc_data.znmin + 1
+            # Read off which columns contain Zk values
+            zk_cols = [col for col in zernikes.columns if col.startswith("Z")]
+            # This is equivalent to self.noll_indices
+            noll_indices = [int(col[1:]) for col in zk_cols]
+            # The first row contains the detector average
+            zk_avg_row = zernikes[zernikes["label"] == "average"]
+            # Each QTable column contains units, hence we convert with astropo
+            zk_avg_microns = np.array(
+                [zk_avg_row[col].to(astropy.units.micron).value[0] for col in zk_cols]
             )
+            sensor_wavefront_data = SensorWavefrontError(noll_indices=noll_indices)
             sensor_name = det_id_map[dataset.dataId["detector"]].getName()
             sensor_wavefront_data.sensor_name = sensor_name
             sensor_wavefront_data.sensor_id = det_name_map[sensor_name].getId()
-            sensor_wavefront_data.annular_zernike_poly = zer_coeff[0]
+            sensor_wavefront_data.annular_zernike_poly = zk_avg_microns
 
             list_of_wf_err.append(sensor_wavefront_data)
 
@@ -1013,7 +1023,7 @@ def write_wep_configuration(
   calcZernikesTask:
     class: lsst.ts.wep.task.calcZernikesTask.CalcZernikesTask
     config:
-      estimateZernikes.maxNollIndex: {self.max_noll_index}
+      estimateZernikes.nollIndices: {self.noll_indices}
       python: |
         from lsst.ts.wep.task import EstimateZernikesTieTask, EstimateZernikesDanishTask
         config.estimateZernikes.retarget(EstimateZernikes{wep_estimator.value.title()}Task)
@@ -1044,7 +1054,7 @@ def run_img(
         raw_seeing: float,
         imsim_log_file: str,
         wep_estimator_method: str,
-        max_noll_index: int,
+        noll_indices: int,
     ) -> None:
         """Run the simulation of images.
 
@@ -1100,8 +1110,8 @@ def run_img(
         wep_estimator_method : str
             Specify the method used to calculate Zernikes in ts_wep.
             Options are "tie" or "danish".
-        max_noll_index : int
-            Maximum Noll index to calculate Zernikes.
+        noll_indices : list
+            Noll indices used in Zernike estimation.
         """
         cam_type = CamType(inst)
         wep_estimator = WepEstimator(wep_estimator_method)
@@ -1112,7 +1122,7 @@ def run_img(
         self.boresight_ra = boresight[0]
         self.boresight_dec = boresight[1]
         self.boresight_rot_ang = rot_cam_in_deg
-        self.max_noll_index = max_noll_index
+        self.noll_indices = noll_indices
         # Remap the reference filter to g
         filter_type_name = self.map_filter_ref_to_g(filter_type_name)
 
@@ -1133,9 +1143,7 @@ def run_img(
             cam_type, obs_metadata, path_sky_file=path_sky_file, star_mag=star_mag
         )
         self.config_ofc_calc(cam_type, controller_config_file)
-        self.imsim_cmpt = ImsimCmpt(
-            num_of_zk=self.max_noll_index - self.ofc_calc.ofc_data.znmin + 1
-        )
+        self.imsim_cmpt = ImsimCmpt(noll_indices=self.noll_indices)
 
         # If path_sky_file using default OPD positions write this to disk
         # so that the Butler can load it later
@@ -1402,12 +1410,13 @@ def set_default_parser(parser: ArgumentParser) -> ArgumentParser:
             default=1,
             help="Number of processor to run imSim and DM pipeline. (default: 1)",
         )
-
         parser.add_argument(
-            "--max_noll_index",
+            "--noll_indices",
             type=int,
-            default=28,
-            help="Maximum Noll index to calculate Zernikes. (default: 28)",
+            nargs="+",
+            default=[4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 20, 21, 22, 27, 28],
+            help="List of Noll Zernike indices used in the simulation.\
+            (default: sparse selection as in ts_wep defaults, i.e. 4-16, 20-22, 27-28)",
         )
 
         parser.add_argument(

python/lsst/ts/imsim/imsim_cmpt.py

@@ -60,21 +60,21 @@ class ImsimCmpt:
         OPD file path.
     opd_metr : lsst.ts.imsim.OpdMetrology
         OPD metrology.
-    num_of_zk : int
-        Number of Zernikes.
+    noll_indices : list
+        Noll indices of used Zernikes.
     num_of_dof : int
         Number of degrees of freedom.
     dof_in_um : numpy.ndarray
         Degrees of freedom in microns.
     """
 
-    def __init__(self, num_of_zk: int) -> None:
+    def __init__(self, noll_indices: list) -> None:
         """Initialize the ImsimCmpt class.
 
         Parameters
         ----------
-        num_of_zk : int
-            Number of Zernikes.
+        noll_indices : list
+            Noll indices of used Zernikes.
         """
         # Output directories
         self._output_dir = None
@@ -84,8 +84,8 @@ def __init__(self, num_of_zk: int) -> None:
         self.opd_file_path = None
         self.opd_metr = OpdMetrology()
 
-        # Specify number of Zernikes
-        self.num_of_zk = num_of_zk
+        # Specify Noll indices of Zernikes
+        self.noll_indices = noll_indices
 
         # AOS Degrees of Freedom
         self.num_of_dof = 50
@@ -619,10 +619,8 @@ def _write_opd_zk_file(
 
         file_path = os.path.join(self.output_img_dir, zk_file_name)
         opd_data = self._map_opd_to_zk(rot_opd_in_deg, num_opd)
-        file_txt = (
-            "# The followings are OPD in rotation angle of %.2f degree in nm from z4 to z22:\n"
-            % rot_opd_in_deg
-        )
+        file_txt = f"# The following are OPD in rotation angle of {rot_opd_in_deg:.2f} degrees in nm,\
+\n# for Noll indices {self.noll_indices}:\n"
         for sensor_id, opd_zk in zip(self.opd_metr.sensor_ids, opd_data):
             zk_str = f"{sensor_id}: {opd_zk}\n"
             file_txt += zk_str
@@ -651,7 +649,7 @@ def _map_opd_to_zk(self, rot_opd_in_deg: float, num_opd: int) -> np.ndarray:
 
         # Map the OPD to the Zk basis and do the collection
         # Get the number of OPD locations by looking at length of fieldX
-        opd_data = np.zeros((num_opd, self.num_of_zk))
+        opd_data = np.zeros((num_opd, len(self.noll_indices)))
         for idx in range(num_opd):
             opd = fits.getdata(self.opd_file_path, idx)
 
@@ -675,9 +673,11 @@ def _map_opd_to_zk(self, rot_opd_in_deg: float, num_opd: int) -> np.ndarray:
             # also fits up to zk28 by default
             zk = self.opd_metr.get_zk_from_opd(opd_map=opd_rot, zk_terms=28)[0]
 
-            # Only need to collect z4 to num_of_zk
-            init_idx = 3
-            opd_data[idx, :] = zk[init_idx : init_idx + self.num_of_zk]
+            # Only need to collect noll indices
+            # noll_indices starts from 1,
+            # whereas zk from 0
+            opd_idx = np.array(self.noll_indices) - 1
+            opd_data[idx, :] = zk[opd_idx]
 
         return opd_data
 
@@ -710,7 +710,7 @@ def _write_opd_pssn_file(self, pssn_file_name: str, num_opd: int) -> None:
 
         # Write to file
         file_path = os.path.join(self.output_img_dir, pssn_file_name)
-        header = "The followings are PSSN and FWHM (in arcsec) data. The final number is the GQ value."
+        header = "The following are PSSN and FWHM (in arcsec) data. The final number is the GQ value."
         np.savetxt(file_path, data, header=header)
 
     def _calc_pssn_opd(self, num_opd: int) -> tuple[list[float], float]:
@@ -807,7 +807,7 @@ def map_opd_data_to_list_of_wf_err(
 
         list_of_wf_err = []
         for sensor_id, sensor_name in zip(sensor_id_list, sensor_name_list):
-            sensor_wavefront_data = SensorWavefrontError(num_of_zk=self.num_of_zk)
+            sensor_wavefront_data = SensorWavefrontError(noll_indices=self.noll_indices)
             sensor_wavefront_data.sensor_id = sensor_id
             sensor_wavefront_data.sensor_name = sensor_name
             # imSim outputs OPD in nanometers so need to change to microns
@@ -956,12 +956,14 @@ def reorder_and_save_wf_err_file(
             if sensor_name in name_list_in_wf_err_map:
                 wf_err = wf_err_map[sensor_name]
             else:
-                wf_err = np.zeros(self.num_of_zk)
+                wf_err = np.zeros(len(self.noll_indices))
             reordered_wf_err_map[sensor_name] = wf_err
 
         # Save the file
         file_path = os.path.join(self.output_img_dir, zk_file_name)
-        file_txt = "# The followings are ZK in um from z4 to z22:\n"
+        file_txt = (
+            f"# The following are ZK in um, \n# for Noll indices {self.noll_indices}:\n"
+        )
         for key, val in reordered_wf_err_map.items():
             zk_str = f"{camera[key].getId()}: {val}\n"
             file_txt += zk_str
@@ -1013,17 +1015,18 @@ def _get_wf_err_values_and_stack_to_matrix(
         wf_err_map : dict
             Calculated wavefront error. The dictionary key [str] is the
             abbreviated sensor name (e.g. R22_S11). The dictionary item
-            [numpy.ndarray] is the averaged wavefront error (z4-z22) in um.
+            [numpy.ndarray] is the averaged wavefront error for
+            chosen noll indices in um.
 
         Returns
         -------
         numpy.ndarray
-            Wavefront errors as a matrix. The column is z4-z22 in um. The row
-            is the individual sensor. The order is the same as the input of
-            wfErrMap.
+            Wavefront errors as a matrix. The column corresponds to
+            chosen noll indices in um. Each row is an individual sensor.
+            The order is the same as the input of wfErrMap.
         """
 
-        value_matrix = np.empty((0, self.num_of_zk))
+        value_matrix = np.empty((0, len(self.noll_indices)))
         for wf_err in wf_err_map.values():
             value_matrix = np.vstack((value_matrix, wf_err))
 
@@ -1044,5 +1047,5 @@ def save_dof_in_um_file_for_next_iter(
         """
 
         file_path = os.path.join(self.output_dir, dof_in_um_file_name)
-        header = "The followings are the DOF in um:"
+        header = "The following are the DOF in um:"
         np.savetxt(file_path, np.transpose(self.dof_in_um), header=header)

python/lsst/ts/imsim/utils/sensor_wavefront_error.py

@@ -27,13 +27,14 @@
 class SensorWavefrontError(object):
     """Contains the wavefront errors for a single sensor."""
 
-    def __init__(self, num_of_zk: int = 19) -> None:
+    def __init__(self, noll_indices: list = range(4, 23)) -> None:
         """Constructs a sensor wavefront error.
 
         Parameters
         ----------
-        num_of_zk : int, optional
-            Number of annular Zernike polynomials. (the default is 19.)
+        noll_indices : int, optional
+            Noll indices of annular Zernike polynomials.
+            (the default is a contiguous space from 4 to 22.)
         """
 
         # Sensor Id
@@ -42,11 +43,11 @@ def __init__(self, num_of_zk: int = 19) -> None:
         # Sensor Name
         self.sensor_name = "R99_S99"
 
-        # Number of zk
-        self.num_of_zk = int(num_of_zk)
+        # Noll indices of stored Zernike polynomials
+        self.noll_indices = noll_indices
 
         # Annular Zernike polynomials (zk)
-        self._annular_zernike_poly = np.zeros(self.num_of_zk)
+        self._annular_zernike_poly = np.zeros_like(self.noll_indices)
 
     @property
     def sensor_id(self) -> int:
@@ -90,9 +91,10 @@ def annular_zernike_poly(self, new_annular_zernike_poly: np.ndarray) -> None:
             annular_zernike_poly must be an array of self.num_of_zk floats.
         """
 
-        if len(new_annular_zernike_poly) != self.num_of_zk:
+        if len(new_annular_zernike_poly) != len(self.noll_indices):
             raise ValueError(
-                "annular_zernike_poly must be an array of %d floats." % self.num_of_zk
+                "annular_zernike_poly must be an array of %d floats."
+                % len(self.noll_indices)
             )
         self._annular_zernike_poly = np.array(new_annular_zernike_poly)
 

python/lsst/ts/imsim/utils/utility.py

@@ -146,7 +146,7 @@ def make_dir(new_dir: str, exist_ok: bool = True) -> None:
 
 
 def get_zk_from_file(zk_file_path: str) -> dict[int, ndarray]:
-    """Get the zk (z4-z22) from file.
+    """Get the zk from file.
 
     Parameters
     ----------
@@ -156,7 +156,7 @@ def get_zk_from_file(zk_file_path: str) -> dict[int, ndarray]:
     Returns
     -------
     numpy.ndarray
-        zk matrix. The colunm is z4-z22. The raw is each data point.
+        zk matrix. The columns are stored Zernikes. The raw is each data point.
     """
 
     with open(zk_file_path, "r") as file: