Merge pull request #51 from simonsobs/err_prop

Err prop

Author: skhrg

lat_alignment/alignment.py

@@ -14,7 +14,6 @@
 import megham.transform as mt
 import numpy as np
 import yaml
-from numpy.lib.arraysetops import isin
 from numpy.typing import NDArray
 from pqdm.processes import pqdm
 
@@ -191,40 +190,46 @@ def main():
                 bootstrap_from = cfg.get("bootstrap_from", "all")
                 logger.info("Bootstrapping from %s", bootstrap_from)
                 if isinstance(dataset, ds.DatasetPhotogrammetry):
-                    dataset, _ = da.align_photo(
+                    dataset, (aff, sft) = da.align_photo(
                         dataset,
                         reference,
                         True,
                         bootstrap_from,
                         **cfg.get("align_photo", {}),
                     )
                 else:
-                    dataset, _ = da.align_tracker(
+                    dataset, (aff, sft) = da.align_tracker(
                         dataset,
                         tracker_yamls[i],
                         bootstrap_from,
                         **cfg.get("align_tracker", {}),
                     )
+                cfrom = "opt_global"
                 if bootstrap_from == "primary":
-                    points = tf.coord_transform(
-                        dataset.points, "opt_primary", f"opt_{mirror}"
-                    )
+                    cfrom = "opt_primary"
                 elif bootstrap_from == "secondary":
-                    points = tf.coord_transform(
-                        dataset.points, "opt_secondary", f"opt_{mirror}"
-                    )
-                else:
-                    points = tf.coord_transform(
-                        dataset.points, "opt_global", f"opt_{mirror}"
-                    )
-                dataset.data_dict = {l: p for l, p in zip(dataset.labels, points)}
-            ddict = {f"{l}_{i}": p for l, p in zip(dataset.labels, dataset.points)}
+                    cfrom = "opt_secondary"
+                points = tf.coord_transform(dataset.points, cfrom, f"opt_{mirror}")
+                aff, sft = tf.affine_basis_transform(aff, sft, cfrom, f"opt_{mirror}")
+                errs = tf.err_transform(dataset.errs, aff)
+                dataset.data_dict = {
+                    l: np.array([p, e]) for l, p, e in zip(dataset.labels, points, errs)
+                }
+            ddict = {
+                f"{l}_{i}": np.array([p, e])
+                for l, p, e in zip(dataset.labels, dataset.points, dataset.errs)
+            }
             data_dict = data_dict | ddict
         dataset = datasets[0].__class__(data_dict)
         append = ""
         if "sample_every" in cfg:
             i, j = cfg["sample_every"]
-            ddict = {l: p for l, p in zip(dataset.labels[i::j], dataset.points[i::j])}
+            ddict = {
+                l: np.array([p, e])
+                for l, p, e in zip(
+                    dataset.labels[i::j], dataset.points[i::j], dataset.errs[i::j]
+                )
+            }
             dataset.data_dict = ddict
             append = f"_{i}_{j}"
 
@@ -251,9 +256,14 @@ def main():
         if cfg.get("only_adj", True):
             for panel in panels:
                 panel.measurements = panel.measurements[panel.adj_msk]
+                panel.meas_err = panel.meas_err[panel.adj_msk]
             measurements = np.vstack([panel.measurements for panel in panels])
-            data = {"TARGET" + str(i): meas for i, meas in enumerate(measurements)}
-            dataset = io.DatasetPhotogrammetry(data)
+            errs = np.vstack([panel.meas_err for panel in panels])
+            data = {
+                "TARGET" + str(i): np.array([meas, err])
+                for i, (meas, err) in enumerate(zip(measurements, errs))
+            }
+            dataset = datasets[0].__class__(data)
             dataset, _ = mir.remove_cm(
                 dataset, mirror, cfg.get("compensate", 0), **cfg.get("common_mode", {})
             )
@@ -267,16 +277,31 @@ def main():
             )
         logger.info("Found measurements for %d panels", len(panels))
         fig = mir.plot_panels(
-            panels, title_str, vmax=cfg.get("vmax", None), use_iqr=cfg.get("iqr", False)
+            panels,
+            False,
+            title_str,
+            vmax=cfg.get("vmax", None),
+            use_iqr=cfg.get("iqr", False),
         )
         fig.savefig(os.path.join(cfgdir, f"{title_str.replace(' ', '_')}{append}.png"))
+        fig = mir.plot_panels(
+            panels,
+            True,
+            title_str,
+            vmax=cfg.get("vmax", None),
+            use_iqr=cfg.get("iqr", False),
+        )
+        fig.savefig(
+            os.path.join(cfgdir, f"{title_str.replace(' ', '_')}_err{append}.png")
+        )
         res_all = np.vstack([panel.residuals for panel in panels])
         model_all = np.vstack([panel.model for panel in panels])
-        mir_out = np.hstack([model_all, res_all])
+        res_err_all = np.vstack([panel.residuals_err for panel in panels])
+        mir_out = np.hstack([model_all, res_all, res_err_all])
         np.savetxt(
             os.path.join(cfgdir, f"{title_str.replace(' ', '_')}_surface{append}.txt"),
             mir_out,
-            header="x y z x_res y_res z_res",
+            header="x y z x_res y_res z_res x_res_err y_res_err z_res_err",
         )
 
         # calc and save adjustments
@@ -341,9 +366,14 @@ def main():
                 )
                 for panel in panels:
                     panel.measurements = panel.measurements[panel.adj_msk]
+                    panel.meas_err = panel.meas_err[panel.adj_msk]
                 measurements = np.vstack([panel.measurements for panel in panels])
-                data = {"TARGET" + str(i): meas for i, meas in enumerate(measurements)}
-                meas = io.Dataset(data)
+                errs = np.vstack([panel.meas_err for panel in panels])
+                data = {
+                    "TARGET" + str(i): np.array([meas, err])
+                    for i, (meas, err) in enumerate(zip(measurements, errs))
+                }
+                meas = meas.__class__(data)
                 meas, common_mode_2 = mir.remove_cm(
                     meas,
                     "primary",
@@ -394,9 +424,14 @@ def main():
                 )
                 for panel in panels:
                     panel.measurements = panel.measurements[panel.adj_msk]
+                    panel.meas_err = panel.meas_err[panel.adj_msk]
                 measurements = np.vstack([panel.measurements for panel in panels])
-                data = {"TARGET" + str(i): meas for i, meas in enumerate(measurements)}
-                meas = io.Dataset(data)
+                errs = np.vstack([panel.meas_err for panel in panels])
+                data = {
+                    "TARGET" + str(i): np.array([meas, err])
+                    for i, (meas, err) in enumerate(zip(measurements, errs))
+                }
+                meas = meas.__class__(data)
                 meas, common_mode_2 = mir.remove_cm(
                     meas,
                     "secondary",
@@ -422,15 +457,16 @@ def main():
                 meas, alignment = da.align_photo(
                     dataset, reference, True, "bearing", **cfg.get("align_photo", {})
                 )
+                meas, cyl_fit = br.cylinder_fit(meas)
+                full_alignment = mt.compose_transform(*alignment, *cyl_fit)
             else:
-                meas, alignment = da.align_tracker(
+                full_alignment, alignment = da.align_tracker(
                     dataset,
                     cfg["tracker_yaml"],
                     "bearing",
                     **cfg.get("align_tracker", {}),
                 )
-            meas, cyl_fit = br.cylinder_fit(meas)
-            full_alignment = mt.compose_transform(*alignment, *cyl_fit)
+                logger.warning("Can't do cylinder fit on bearing with tracker data!")
             log_alignment(full_alignment, logger)
         except Exception as e:
             print(

lat_alignment/data_alignment.py

@@ -8,13 +8,13 @@
 
 import matplotlib.pyplot as plt
 import numpy as np
-from megham.transform import apply_transform, decompose_rotation, get_affine, get_rigid
+from megham.transform import apply_transform, decompose_rotation, get_rigid
 from megham.utils import make_edm
 from numpy.typing import NDArray
 
 from . import io
 from .dataset import Dataset, DatasetPhotogrammetry, DatasetReference
-from .transforms import coord_transform
+from .transforms import coord_transform, err_transform
 
 logger = logging.getLogger("lat_alignment")
 
@@ -152,12 +152,6 @@ def align_photo(
         logger.info("\t\tAssociated %s with %s", label, pname)
     if blind_search > 0:
         raise NotImplementedError("Blind search not implemented yet!")
-    # Set 12
-    # ref = [rpoint for rpoint, _ in reference[element]]
-    # ref = np.array(ref)[[True, True, False, True]]
-    # invars = ["TARGET35", "TARGET4", "TARGET484"] #, "TARGET421"]
-    # pts = [dataset[label] for label in invars]
-    # print(invars)
     if len(ref) < 4:
         logger.warning(f"Only {len(ref)} reference points found!")
         logger.warning(f"Adding reference codes")
@@ -194,8 +188,10 @@ def align_photo(
         pts_scaled = pts * scale_fac
         logger.debug("\t\tScale factor of %f applied", scale_fac)
 
-        new_rot, new_sft = get_rigid(pts_scaled[msk], ref[msk], method="mean")
-        pts_t = apply_transform(pts_scaled[msk], new_rot, new_sft)
+        new_rot, new_sft = get_rigid(
+            pts_scaled[msk].astype(np.float64), ref[msk], method="mean"
+        )
+        pts_t = apply_transform(pts_scaled[msk].astype(np.float64), new_rot, new_sft)
         if plot:
             fig = plt.figure()
             ax = fig.add_subplot(projection="3d")
@@ -227,22 +223,27 @@ def align_photo(
     if rot is None or sft is None:
         raise ValueError("Transformation is None")
 
-    coords_transformed = apply_transform(dataset.points * scale_fac, rot, sft)
+    logger.debug("\t\tShift is %s mm", str(sft))
+    logger.debug("\t\tRotation is %s deg", str(np.rad2deg(decompose_rotation(rot))))
+    scale_fac = np.eye(3) * scale_fac
+    rot @= scale_fac
+
+    coords_transformed = apply_transform(dataset.points, rot, sft)
+    errs_transformed = err_transform(dataset.errs, rot)
     labels = dataset.labels
 
     if kill_refs:
         msk = ~np.isin(dataset.labels, invars)
         labels = labels[msk]
         coords_transformed = coords_transformed[msk]
+        errs_transformed = errs_transformed[msk]
 
-    data = {label: coord for label, coord in zip(labels, coords_transformed)}
+    data = {
+        label: np.array([coord, err])
+        for label, coord, err in zip(labels, coords_transformed, errs_transformed)
+    }
     transformed = DatasetPhotogrammetry(data)
 
-    logger.debug("\t\tShift is %s mm", str(sft))
-    logger.debug("\t\tRotation is %s deg", str(np.rad2deg(decompose_rotation(rot))))
-    scale_fac = np.eye(3) * scale_fac
-    rot @= scale_fac
-
     if plot:
         fig = plt.figure()
         ax = fig.add_subplot(projection="3d")
@@ -356,17 +357,21 @@ def align_tracker(
         rms,
     )
 
-    coords_transformed = apply_transform(dataset.points * scale_fac, rot, sft)
-    labels = dataset.labels
-
-    data = {label: coord for label, coord in zip(labels, coords_transformed)}
-    transformed = Dataset(data)
-
     logger.debug("\t\tShift is %s mm", str(sft))
     logger.debug("\t\tRotation is %s deg", str(np.rad2deg(decompose_rotation(rot))))
     scale_fac = np.eye(3) * scale_fac
     rot @= scale_fac
 
+    coords_transformed = apply_transform(dataset.points, rot, sft)
+    errs_transformed = err_transform(dataset.errs, rot)
+    labels = dataset.labels
+
+    data = {
+        label: np.array([coord, err])
+        for label, coord, err in zip(labels, coords_transformed, errs_transformed)
+    }
+    transformed = Dataset(data)
+
     if plot:
         fig = plt.figure()
         ax = fig.add_subplot(projection="3d")