LAT pointing model and ffp updates

sotodlib/coords/fp_containers.py

@@ -173,14 +173,18 @@ def dist(self):
             return np.linalg.norm(self.diff, axis=1)
         return np.linalg.norm(self.diff[:, :2], axis=1)
 
+    @property
+    def padded(self):
+        return self.tot_weight is None 
+
     @classmethod
     def empty(cls, template, stream_id, wafer_slot, n_aman):
         if template is None:
             raise TypeError("template must be an instance of Template, not None")
         full_fp = np.full(template.fp.shape + (n_aman,), np.nan)
-        tot_weight = np.zeros(len(template.det_ids))
+        tot_weight = np.zeros((len(template.det_ids), 2))
         avg_fp = np.full_like(template.fp, np.nan)
-        weight = np.zeros(len(template.det_ids))
+        weight = np.zeros((len(template.det_ids), 2))
         transformed = template.fp.copy()
         center = template.center.copy()
         center_transformed = template.center.copy()
@@ -219,6 +223,9 @@ def map_by_det_id(self, aman):
         srt = np.argsort(aman.det_info.det_id[msk])
         xi = aman.pointing.xi[msk][srt][mapping]
         eta = aman.pointing.eta[msk][srt][mapping]
+        r2 = np.nan + np.zeros_like(eta)
+        if "R2" in aman.pointing:
+            r2 = aman.pointing.R2[msk][srt][mapping]
         if "polarization" in aman:
             # name of field just a placeholder for now
             gamma = aman.polarization.polang[msk][srt][mapping]
@@ -227,12 +234,13 @@ def map_by_det_id(self, aman):
         else:
             gamma = np.full(len(xi), np.nan)
         fp = np.column_stack((xi, eta, gamma))
-        return fp, template_msk
+        return fp, r2, template_msk
 
     def add_fp(self, i, fp, weights, template_msk):
         if self.full_fp is None or self.tot_weight is None:
             raise ValueError("full_fp or tot_weight not initialized")
-        self.full_fp[template_msk, :, i] = fp * weights[..., None]
+        self.full_fp[template_msk, :, i] = fp * weights[:, 0][..., None]
+        weights = np.nan_to_num(weights)
         self.tot_weight[template_msk] += weights
 
     def save(self, f, db_info, group):
@@ -242,9 +250,10 @@ def save(self, f, db_info, group):
             ("xi", np.float32),
             ("eta", np.float32),
             ("gamma", np.float32),
+            ("padded", np.bool_),
         ]
         fpout = np.fromiter(
-            zip(self.det_ids, *(self.transformed.T)), dtype=outdt, count=ndets
+            zip(self.det_ids, *(self.transformed.T), np.ones(len(self.det_ids), dtype=bool)*(self.padded)), dtype=outdt, count=ndets
         )
         write_dataset(
             metadata.ResultSet.from_friend(fpout),
@@ -269,6 +278,7 @@ def save(self, f, db_info, group):
             ("eta_m", np.float32),
             ("gamma_m", np.float32),
             ("weights", np.float32),
+            ("r2", np.float32),
             ("n_point", np.int8),
             ("n_gamma", np.int8),
         ]
@@ -277,7 +287,7 @@ def save(self, f, db_info, group):
                 self.det_ids,
                 *(self.transformed.T),
                 *(self.avg_fp.T),
-                self.weights,
+                *(self.weights.T),
                 self.n_point,
                 self.n_gamma,
             ),
@@ -305,6 +315,8 @@ def save(self, f, db_info, group):
     def load(cls, group):
         stream_id = group.name.split("/")[-1]
         fp_full = read_dataset(group.file, f"{group.name}/focal_plane_full")
+        if fp_full.keys is None:
+            raise ValueError("fp_full somehow has no keys")
         det_ids = fp_full["dets:det_id"]
         avg_fp = np.column_stack(
             (
@@ -313,7 +325,10 @@ def load(cls, group):
                 np.array(fp_full["gamma_m"]),
             )
         )
-        weights = fp_full["weights"]
+        # For backwards compatibility
+        weights = np.array(fp_full["weights"])
+        if "r2" in fp_full.keys:
+            weights = np.column_stack((weights, np.array(fp_full["r2"])))
         transformed = np.column_stack(
             (
                 np.array(fp_full["xi_t"]),
@@ -367,6 +382,11 @@ def __post_init__(self):
             self.center, self.transform_fullcm.affine, self.transform_fullcm.shift
         )
 
+    @property
+    def num_fps(self):
+        fps = [fp for fp in self.focal_planes if fp.tot_weight is not None]
+        return len(fps)
+
     @classmethod
     def from_pointing_cfg(cls, pointing_cfg):
         name = pointing_cfg["tube_slot"]
@@ -544,9 +564,9 @@ def plot_ufm(focal_plane, plot_dir):
 
 def plot_ot(ot, plot_dir):
     fig, (ax1, ax2) = plt.subplots(1, 2, constrained_layout=True)
-    dists = [fp.dist[fp.isfinite] * 180 * 60 * 60 / np.pi for fp in ot.focal_planes]
-    xis = [fp.transformed[fp.isfinite, 0] for fp in ot.focal_planes]
-    etas = [fp.transformed[fp.isfinite, 1] for fp in ot.focal_planes]
+    dists = [fp.dist[fp.isfinite] * 180 * 60 * 60 / np.pi for fp in ot.focal_planes if fp.tot_weight is not None]
+    xis = [fp.transformed[fp.isfinite, 0] for fp in ot.focal_planes if fp.tot_weight is not None]
+    etas = [fp.transformed[fp.isfinite, 1] for fp in ot.focal_planes if fp.tot_weight is not None]
 
     # Plot the radial dist
     r = np.sqrt(np.hstack(xis) ** 2 + np.hstack(etas) ** 2)
@@ -633,21 +653,23 @@ def plot_receiver(receiver, plot_dir):
     xlims, ylims = receiver.lims
 
     fig, axs_all = plt.subplots(
-        4, 3, sharex="col", sharey="row", constrained_layout=True
+        len(valid_ids), 3, sharex="col", sharey="row", constrained_layout=True
     )
     axs = axs_all.flat
     axs[0].set_title("Xi")
     axs[1].set_title("Eta")
     axs[2].set_title("Gamma")
     cf = None
-    for i in range(4):
+    for i in range(len(valid_ids)):
         for j in range(3):
             axs[3 * i + j].set_aspect("equal")
             axs[3 * i + j].set_xlim(xlims)
             axs[3 * i + j].set_ylim(ylims)
         for fp in receiver.focal_planes:
+            if fp.tot_weight is None:
+                continue
             msk = (fp.template.id_strs == valid_ids[i]) * fp.isfinite
-            if np.sum(msk) == 0:
+            if np.sum(msk) < 3:
                 continue
             diff = fp.diff * 180 * 60 * 60 / np.pi
             cf = axs[3 * i + 0].tricontourf(

sotodlib/coords/pointing_model.py

@@ -72,28 +72,33 @@ def apply_pointing_model(tod, pointing_model=None, ancil=None,
       AxisManager: the corrected boresight.
 
     """
-    if pointing_model is None and 'pointing_model' not in tod:
-        logger.warning('No pointing_model found -- applying basic model.')
-        assert wrap in (None, 'boresight'), \
-            'When using naive pointing model, wrap=... not supported'
+    if pointing_model is None and "pointing_model" not in tod:
+        logger.warning("No pointing_model found -- applying basic model.")
+        assert wrap in (
+            None,
+            "boresight",
+        ), "When using naive pointing model, wrap=... not supported"
         return apply_basic_pointing_model(tod)
 
-    pointing_model = _valid_arg(pointing_model, 'pointing_model',
-                                src=tod)
-    ancil = _valid_arg(ancil, 'ancil', src=tod)
+    pointing_model = _valid_arg(pointing_model, "pointing_model", src=tod)
+    ancil = _valid_arg(ancil, "ancil", src=tod)
 
     # Encoder values, to radians.
-    vers = pointing_model['version']
-    tel_type = vers.split('_')[0]
-    if tel_type == 'sat':
+    if pointing_model is None:
+        raise ValueError("No pointing_model specified")
+    if "version" not in pointing_model:
+        raise ValueError("Pointing model version not specified")
+    vers = pointing_model["version"]
+    tel_type = vers.split("_")[0]
+    if tel_type == "sat":
         boresight = apply_pointing_model_sat(vers, pointing_model, tod, ancil)
-    elif tel_type == 'lat':
-        boresight = apply_pointing_model_lat(vers, pointing_model, tod, ancil)
+    elif tel_type == "lat":
+        boresight = apply_pointing_model_lat(vers, pointing_model, ancil)
     else:
         raise ValueError(f'Unimplemented pointing model "{vers}"')
 
     if wrap is None:
-        wrap = 'boresight'
+        wrap = "boresight"
     if wrap is not False:
         if wrap in tod._fields:
             del tod[wrap]
@@ -115,8 +120,109 @@ def apply_pointing_model_sat(vers, params, tod, ancil):
         raise ValueError(f'Unimplemented pointing model "{vers}"')
 
 
-def apply_pointing_model_lat(vers, tod, pointing_model, ancil):
-    raise ValueError(f'Unimplemented pointing model "{vers}"')
+def apply_pointing_model_lat(vers, params, ancil):
+    az, el, roll = _get_lat_enc_radians(ancil)
+
+    if vers == 'lat_naive':
+        return _new_boresight(ancil.samps, az=az, el=el, roll=roll)
+
+    if vers == "lat_v1":
+        az1, el1, roll1 = model_lat_v1(params, az, el, roll)
+        return _new_boresight(ancil.samps, az=az1, el=el1, roll=roll1)
+
+    else:
+        raise ValueError(f'Unimplemented pointing model "{vers}"')
+
+
+#
+# LAT model(s)
+#       
+def model_lat_v1(params, az, el, roll):
+    """Applies pointing model to (az, el, roll).
+
+    Args:
+      params: AxisManager (or dict) of pointing parameters.
+      az, el, roll: naive horizon coordinates, in radians, of the
+        boresight.
+
+    The implemented model parameters are all in Radians:
+    - enc_offset_{az, el, cr}: Encoder offsets in Radians.
+      Sign convention: True = Encoder + Offset
+    - cr_center_{xi,eta}0: The (xi,eta) coordinate in the LATR-centered 
+      focal plane that remains fixed under corotation. 
+    - el_axis_center_{xi,eta}0: The (xi,eta) coordinate in the CR-centered
+      focal plane that appears fixed when the elevation structure is rotated 
+      about its axis.
+    - mir_center_{xi,eta}0: The (xi,eta) coordinate in the El-structure-centered
+      focal plane about which any mirror misalignment rotates.
+
+    """
+    _p = dict(param_defaults['lat_v1'])
+    if isinstance(params, dict):
+        _p.update(params)
+    else:
+        _p.update({k: params[k] for k in params._fields.keys()})
+    params, _p = _p, None   
+
+    for k, v in params.items():
+        if k == 'version':
+            continue
+        if k not in param_defaults['lat_v1'] and v != 0.:
+            raise ValueError(f'Handling of model param "{k}" is not implemented.')
+
+    # Here we reconstruct the naive corotator value before applying 
+    # the elevation encoder offset.
+    cr = el - roll - np.deg2rad(60)    
+
+    az_orig = az.copy()
+    az = az + params['enc_offset_az']
+    el = el + params['enc_offset_el'] 
+    cr = cr + params['enc_offset_cr']
+
+    # Lonlat rotation with az and el encoder offsets included.
+    q_lonlat = quat.rotation_lonlat(-1 * az, el)
+
+    # Rotation that takes a vector in elevation-hub centered 
+    # coordinates to final boresight centered coordinates. 
+    # Accounts for any mirror-related offsets.
+    q_mir_center = ~quat.rotation_xieta(
+        params['mir_center_xi0'],
+        params['mir_center_eta0']
+    )
+
+    # Elevation component of roll motion
+    q_el_roll = quat.euler(2, el - np.deg2rad(60))
+
+    # Rotation that takes a vector in telescope's elevation-hub centered 
+    # coordinates to mirror-centered coordinates
+    q_el_axis_center = ~quat.rotation_xieta(
+        params['el_axis_center_xi0'],
+        params['el_axis_center_eta0']
+    )
+
+    # Corotator component of roll motion
+    q_cr_roll = quat.euler(2, -1 * cr)
+
+    # Rotation that takes a vector in LATR/focal plane-centered coordinates to  
+    # corotator-centered coordinates.
+    q_cr_center = ~quat.rotation_xieta(
+        params['cr_center_xi0'],
+        params['cr_center_eta0']
+    )
+
+    # Horizon Coordiantes
+    q_hs = (
+        q_lonlat * q_mir_center
+        * q_el_roll * q_el_axis_center
+        * q_cr_roll * q_cr_center
+    )
+    new_az, el, roll = quat.decompose_lonlat(q_hs)* np.array([-1, 1, 1])[..., None]
+
+    # Make corrected az as close as possible to the input az.    
+    change = ((new_az - az_orig) + np.pi) % (2 * np.pi) - np.pi
+    az = az_orig + change
+
+    return az, el, roll
 
 
 #
@@ -126,19 +232,6 @@ def apply_pointing_model_lat(vers, tod, pointing_model, ancil):
 # sat_v1: you can expand v1, as long as new params don't do anything
 # if their value is zero (and that should be the registered default).
 
-defaults_sat_v1 = {
-    'enc_offset_az': 0.,
-    'enc_offset_el': 0.,
-    'enc_offset_boresight': 0.,
-    'fp_offset_xi0': 0.,
-    'fp_offset_eta0': 0.,
-    'fp_rot_xi0': 0.,
-    'fp_rot_eta0': 0.,
-    'az_rot': 0.,
-    'base_tilt_cos': 0.,
-    'base_tilt_sin': 0.,
-}
-
 def model_sat_v1(params, az, el, roll):
     """Applies pointing model to (az, el, roll).
 
@@ -161,7 +254,7 @@ def model_sat_v1(params, az, el, roll):
       - az_rot: Dimensionless parameter describing a linear dependence of Az on El.
 
     """
-    _p = dict(defaults_sat_v1)
+    _p = dict(param_defaults['sat_v1'])
     if isinstance(params, dict):
         _p.update(params)
     else:
@@ -171,7 +264,7 @@ def model_sat_v1(params, az, el, roll):
     for k, v in params.items():
         if k == 'version':
             continue
-        if k not in defaults_sat_v1 and v != 0.:
+        if k not in param_defaults['sat_v1'] and v != 0.:
             raise ValueError(f'Handling of model param "{k}" is not implemented.')
 
     # Construct offsetted encoders.
@@ -205,6 +298,31 @@ def model_sat_v1(params, az, el, roll):
 
 
 # Support functions
+param_defaults={
+    'lat_v1' : {
+        'enc_offset_az': 0,
+        'enc_offset_el': 0,
+        'enc_offset_cr': 0,
+        'el_axis_center_xi0': 0,
+        'el_axis_center_eta0': 0,
+        'cr_center_xi0': 0,
+        'cr_center_eta0': 0,
+        'mir_center_xi0': 0,
+        'mir_center_eta0': 0,
+    },
+    'sat_v1' : {
+        'enc_offset_az': 0.,
+        'enc_offset_el': 0.,
+        'enc_offset_boresight': 0.,
+        'fp_offset_xi0': 0.,
+        'fp_offset_eta0': 0.,
+        'fp_rot_xi0': 0.,
+        'fp_rot_eta0': 0.,
+        'az_rot': 0.,
+        'base_tilt_cos': 0.,
+        'base_tilt_sin': 0.
+    }
+}
 
 def _new_boresight(samps, az=None, el=None, roll=None):
     boresight = core.AxisManager(samps)
@@ -219,6 +337,15 @@ def _get_sat_enc_radians(ancil):
             ancil.el_enc * DEG,
             -ancil.boresight_enc * DEG)
 
+def _get_lat_enc_radians(ancil):
+    az = ancil.az_enc * DEG
+    el = ancil.el_enc * DEG
+    if "roll_enc" in ancil:
+        roll = ancil.roll_enc * DEG
+    else:
+        roll = (ancil.el_enc - ancil.corotator_enc - 60) * DEG
+    return (az, el, roll)
+
 def get_base_tilt_q(c, s):
     """Returns the quaternion rotation that applies base tilt, taking
     vectors in the platforms horizon coordinates to vectors in the