Merge pull request #47 from simonsobs/site_fixes

Site fixes

Author: skhrg

lat_alignment/alignment.py

@@ -10,9 +10,11 @@
 from functools import partial
 from importlib.resources import files
 
+import matplotlib.pyplot as plt
 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
 
@@ -108,10 +110,16 @@ def main():
 
     mode = cfg.get("mode", "panel")
     cfgdir = os.path.dirname(os.path.abspath(args.config))
-    meas_file = os.path.abspath(os.path.join(cfgdir, cfg["measurement"]))
+    meas_files = cfg["measurement"]
+    if isinstance(meas_files, str):
+        meas_files = [meas_files]
+    meas_files = [os.path.abspath(os.path.join(cfgdir, meas)) for meas in meas_files]
+    tracker_yamls = cfg.get("tracker_yaml", "")
+    if isinstance(tracker_yamls, str):
+        tracker_yamls = [tracker_yamls]
     title_str = cfg["title"]
     logger.info("Begining alignment %s in %s mode", title_str, mode)
-    logger.debug("Using measurement file: %s", meas_file)
+    logger.debug("Using measurement files: %s", meas_files)
 
     dat_dir = os.path.abspath(os.path.join(cfgdir, cfg.get("data_dir", "/")))
     if "data_dir" in cfg:
@@ -122,7 +130,10 @@ def main():
         ref_path = str(files("lat_alignment.data").joinpath("reference.yaml"))
     with open(ref_path) as file:
         reference = yaml.safe_load(file)
-    dataset = io.load_data(meas_file, **cfg.get("load", {"source": "photo"}))
+    datasets = [
+        io.load_data(meas_file, **cfg.get("load", {"source": "photo"}))
+        for meas_file in meas_files
+    ]
     if "data_dir" in cfg:
         corner_path_m1 = os.path.join(dat_dir, f"primary_corners.yaml")
         adj_path_m1 = os.path.join(dat_dir, f"primary_adj.csv")
@@ -159,47 +170,73 @@ def main():
         logger.info("Aligning panels for the %s mirror", mirror)
 
         # init, fit, and plot panels
-        try:
-            if isinstance(dataset, ds.DatasetPhotogrammetry):
-                dataset, _ = da.align_photo(
-                    dataset, reference, True, mirror, **cfg.get("align_photo", {})
-                )
-            else:
-                dataset, _ = da.align_tracker(
-                    dataset, cfg["tracker_yaml"], mirror, **cfg.get("align_tracker", {})
-                )
-        except Exception as e:
-            logger.error("Failed to align to reference points, with error %s", str(e))
-            bootstrap_from = cfg.get("bootstrap_from", "all")
-            logger.info("Bootstrapping from %s", bootstrap_from)
-            if isinstance(dataset, ds.DatasetPhotogrammetry):
-                dataset, _ = da.align_photo(
-                    dataset,
-                    reference,
-                    True,
-                    bootstrap_from,
-                    **cfg.get("align_photo", {}),
-                )
-            else:
-                dataset, _ = da.align_tracker(
-                    dataset,
-                    cfg["tracker_yaml"],
-                    bootstrap_from,
-                    **cfg.get("align_tracker", {}),
-                )
-            if bootstrap_from == "primary":
-                points = tf.coord_transform(
-                    dataset.points, "opt_primary", f"opt_{mirror}"
-                )
-            elif bootstrap_from == "secondary":
-                points = tf.coord_transform(
-                    dataset.points, "opt_secondary", f"opt_{mirror}"
+        data_dict = {}
+        for i, dataset in enumerate(datasets):
+            try:
+                if isinstance(dataset, ds.DatasetPhotogrammetry):
+                    dataset, _ = da.align_photo(
+                        dataset, reference, True, mirror, **cfg.get("align_photo", {})
+                    )
+                else:
+                    dataset, _ = da.align_tracker(
+                        dataset,
+                        tracker_yamls[i],
+                        mirror,
+                        **cfg.get("align_tracker", {}),
+                    )
+            except Exception as e:
+                logger.error(
+                    "Failed to align to reference points, with error %s", str(e)
                 )
-            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)}
+                bootstrap_from = cfg.get("bootstrap_from", "all")
+                logger.info("Bootstrapping from %s", bootstrap_from)
+                if isinstance(dataset, ds.DatasetPhotogrammetry):
+                    dataset, _ = da.align_photo(
+                        dataset,
+                        reference,
+                        True,
+                        bootstrap_from,
+                        **cfg.get("align_photo", {}),
+                    )
+                else:
+                    dataset, _ = da.align_tracker(
+                        dataset,
+                        tracker_yamls[i],
+                        bootstrap_from,
+                        **cfg.get("align_tracker", {}),
+                    )
+                if bootstrap_from == "primary":
+                    points = tf.coord_transform(
+                        dataset.points, "opt_primary", f"opt_{mirror}"
+                    )
+                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)}
+            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])}
+            dataset.data_dict = ddict
+            append = f"_{i}_{j}"
+
+        fig = plt.figure()
+        ax = fig.add_subplot(projection="3d")
+        ax.scatter(
+            dataset.targets[:, 0],
+            dataset.targets[:, 1],
+            dataset.targets[:, 2],
+            marker="x",
+        )
+        plt.show()
         dataset, _ = mir.remove_cm(
             dataset, mirror, cfg.get("compensate", 0), **cfg.get("common_mode", {})
         )
@@ -229,13 +266,15 @@ def main():
                 cfg.get("adjuster_radius", 100),
             )
         logger.info("Found measurements for %d panels", len(panels))
-        fig = mir.plot_panels(panels, title_str, vmax=cfg.get("vmax", None))
-        fig.savefig(os.path.join(cfgdir, f"{title_str.replace(' ', '_')}.png"))
+        fig = mir.plot_panels(
+            panels, 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"))
         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])
         np.savetxt(
-            os.path.join(cfgdir, f"{title_str.replace(' ', '_')}_surface.txt"),
+            os.path.join(cfgdir, f"{title_str.replace(' ', '_')}_surface{append}.txt"),
             mir_out,
             header="x y z x_res y_res z_res",
         )
@@ -249,11 +288,15 @@ def main():
         order = np.lexsort((adjustments[:, 2], adjustments[:, 1], adjustments[:, 0]))
         adjustments = adjustments[order]
         np.savetxt(
-            os.path.join(cfgdir, f"{title_str.replace(' ', '_')}.csv"),
+            os.path.join(cfgdir, f"{title_str.replace(' ', '_')}{append}.csv"),
             adjustments,
             fmt=["%d", "%d", "%d"] + ["%.5f"] * 14,
         )
     elif mode == "optical":
+        if len(datasets) > 1 or len(tracker_yamls) > 1:
+            raise ValueError("Cannot have multiple files in optical mode")
+        dataset = datasets[0]
+        tracker_yaml = tracker_yamls[0]
         align_to = cfg["align_to"]
         if align_to not in ["primary", "secondary", "receiver", "bearing"]:
             raise ValueError(f"Invalid element specified for 'align_to': {align_to}")
@@ -264,7 +307,7 @@ def main():
             )
         else:
             dataset, _ = da.align_tracker(
-                dataset, cfg["tracker_yaml"], "all", **cfg.get("align_tracker", {})
+                dataset, tracker_yaml, "all", **cfg.get("align_tracker", {})
             )
 
         # Load data and compute the transformation to align with the model

lat_alignment/data/reference_va.yaml

@@ -158,12 +158,12 @@ def align_photo(
     # invars = ["TARGET35", "TARGET4", "TARGET484"] #, "TARGET421"]
     # pts = [dataset[label] for label in invars]
     # print(invars)
-    if len(ref) < 3:
+    if len(ref) < 4:
         logger.warning(f"Only {len(ref)} reference points found!")
         logger.warning(f"Adding reference codes")
         pts += found_coded
         ref += ref_coded
-    if len(ref) < 3:
+    if len(ref) < 4:
         raise ValueError(
             f"Only {len(ref)} reference points found including codes! Can't align!"
         )

lat_alignment/data_alignment.py

@@ -52,11 +52,29 @@ def _load_tracker_yaml(path: str):
     return DatasetReference(data)
 
 
-def _load_tracker_txt(path: str):
-    data = np.genfromtxt(path, usecols=(3, 4, 5), skip_header=1, dtype=str)
+def _load_tracker_txt(path: str, group_dist=0.02, group_thresh=0.02):
+    data = np.genfromtxt(
+        path, usecols=(3, 4, 5), skip_header=1, dtype=str, delimiter="\t"
+    )
     data = np.char.replace(data, ",", "").astype(float)
 
-    data = {f"TARGET_{i}": dat for i, dat in enumerate(data)}
+    to_kill = []
+    if group_dist > 0 and group_thresh > 0:
+        done = []
+        edm = make_edm(data[:, :2])
+        np.fill_diagonal(edm, np.nan)
+        for i in range(len(edm)):
+            if i in to_kill or i in done:
+                continue
+            group_idx = np.hstack(([i], np.where(edm[i] <= group_dist)[0]))
+            done += group_idx.tolist()
+            if len(group_idx) == 1:
+                continue
+            zs = data[group_idx, 2]
+            bad_zs = np.abs(zs - np.median(zs)) > group_thresh
+            to_kill += group_idx[bad_zs].tolist()
+        logger.info("\tFound and removed %d bad group points", len(to_kill))
+    data = {f"TARGET_{i}": dat for i, dat in enumerate(data) if i not in to_kill}
 
     return Dataset(data)
 
@@ -68,7 +86,7 @@ def _load_tracker_csv(path: str):
     )
 
 
-def load_tracker(path: str) -> Dataset:
+def load_tracker(path: str, group_dist=0.02, group_thresh=0.02) -> Dataset:
     """
     Load laser tracker data.
     TODO: This interface needs to be unified with `load_photo` so all code can use either datatype interchangibly
@@ -78,6 +96,14 @@ def load_tracker(path: str) -> Dataset:
     path : str
         The path to the laser tracker data.
         The type of data will be infered from the extension.
+    group_dist : float, default: 0.02
+        Distance between points in xy needed to group them for cuts.
+        Only used for `.txt` files.
+        Set to 0 to disable.
+    group_thresh : float, default: 0.02
+        Difference in z between point and the median z for a group to cut at.
+        Only used for `.txt` files.
+        Set to 0 to disable.
 
     Returns
     -------
@@ -90,7 +116,7 @@ def load_tracker(path: str) -> Dataset:
     if ext == ".yaml":
         return _load_tracker_yaml(path)
     elif ext == ".txt":
-        return _load_tracker_txt(path)
+        return _load_tracker_txt(path, group_dist, group_thresh)
     elif ext == ".csv":
         return _load_tracker_csv(path)
     raise ValueError(f"Invalid tracker data with extension {ext}")
@@ -152,12 +178,20 @@ def load_photo(
             to_kill += [labels[trg_msk][i]]
     msk = ~np.isin(labels, to_kill)
     logger.info("\tFound and removed %d doubles", len(to_kill))
-    labels, coords = labels[msk], coords[msk]
+    labels, coords, err = labels[msk], coords[msk], err[msk]
 
     if plot:
         fig = plt.figure()
         ax = fig.add_subplot(projection="3d")
-        ax.scatter(coords[:, 0], coords[:, 1], coords[:, 2], marker="x")
+        p = ax.scatter(
+            coords[:, 0],
+            coords[:, 1],
+            coords[:, 2],
+            marker="x",
+            c=err,
+            vmax=np.percentile(err, 90),
+        )
+        fig.colorbar(p)
         plt.show()
 
     data = {label: coord for label, coord in zip(labels, coords)}