Implementation of has-mapping sensors (#368)

* Initial file writing for has-mapping

* cleanup

* Fix autogen plotting

* Fix live plot units

CHANGELOG.md

@@ -6,6 +6,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/).
 
 ### Added
 - support for sensors without `has-measure-trigger` trait
+- support for sensors that implement `has-mapping` trait
 
 ### Changed
 - Update for pint version of WrightTools (3.4.0)

yaqc_cmds/_plot.py

@@ -95,8 +95,6 @@ def create_settings(self):
 
     def on_channels_changed(self):
         new = list(sensors.get_channels_dict())
-        if "ingaas" in new:
-            new.remove("ingaas")
         self.channel.set_allowed_values(new)
 
     def on_data_file_created(self):
@@ -111,7 +109,7 @@ def on_axis_updated(self):
         with somatic._wt5.data_container as data:
             axis = data[self.axis.read()]
             units = axis.attrs.get("units")
-            units = list(wt.units.get_valid_conversions(units))
+            units = [units] + list(wt.units.get_valid_conversions(units))
             self.axis_units.set_allowed_values(units)
 
     def on_data_file_written(self):
@@ -126,7 +124,7 @@ def on_data_file_written(self):
             axis = data[self.axis.read()]
             limits = list(
                 wt.units.convert(
-                    [np.min(axis.full), np.max(axis.full)], axis.attrs.get("units"), x_units
+                    [np.nanmin(axis.full), np.nanmax(axis.full)], axis.attrs.get("units"), x_units
                 )
             )
             channel = data[self.channel.read()]
@@ -149,7 +147,8 @@ def on_data_file_written(self):
             try:
                 self.plot_widget.set_xlim(min(limits), max(limits))
                 self.plot_widget.set_ylim(np.min(channel), np.max(channel))
-            except Exception:
+            except Exception as e:
+                print(e)
                 pass
 
     def on_sensors_changed(self):

yaqc_cmds/sensors/_sensors.py

@@ -24,11 +24,8 @@ class Data(QtCore.QMutex):
     def __init__(self):
         QtCore.QMutex.__init__(self)
         self.WaitCondition = QtCore.QWaitCondition()
-        self.shape = (1,)
-        self.size = 1
         self.channels = {}
         self.signed = []
-        self.map = None
 
     def read(self):
         return self.channels
@@ -39,15 +36,12 @@ def write(self, channels):
         self.WaitCondition.wakeAll()
         self.unlock()
 
-    def write_properties(self, shape, channels, signed=False, map=None):
+    def write_properties(self, channels, signed=False):
         self.lock()
-        self.shape = shape
-        self.size = np.prod(shape)
         self.channels = channels
         self.signed = signed
         if not signed:
             self.signed = [False] * len(self.channels)
-        self.map = map
         self.WaitCondition.wakeAll()
         self.unlock()
 
@@ -66,17 +60,6 @@ def __init__(self, *args, **kwargs):
         self.Widget = kwargs.pop("Widget")
         self.data = Data()
         self.active = True
-        # shape
-        if "shape" in args[1]:
-            self.shape = args[1].pop("shape")
-        else:
-            self.shape = (1,)
-        # map
-        if "has_map" in args[1]:
-            self.has_map = args[1].pop("has_map")
-        else:
-            self.has_map = False
-        self.has_map = False  # turning this feature off for now  --Blaise 2020-09-25
         self.measure_time = pc.Number(initial_value=np.nan, display=True, decimals=3)
         super().__init__(*args, **kwargs)
         self.settings_updated.emit()
@@ -100,7 +83,7 @@ def give_widget(self, widget):
         self.gui.create_frame(widget)
 
     def initialize(self):
-        self.wait_until_still()
+        # self.wait_until_still()
         self.freerun.updated.connect(self.on_freerun_updated)
         self.update_ui.emit()
         self.driver.update_ui.connect(self.on_driver_update_ui)
@@ -141,7 +124,6 @@ def __init__(self, sensor, yaqd_port):
         self.busy = sensor.busy
         self.freerun = sensor.freerun
         self.data = sensor.data
-        self.shape = sensor.shape
         self.measure_time = sensor.measure_time
         self.thread = sensor.thread
 
@@ -168,8 +150,13 @@ def measure(self):
                     time.sleep(0.01)
             out = self.client.get_measured()
             del out["measurement_id"]
+            try:
+                del out["mapping_id"]
+            except KeyError:
+                # No mapping
+                pass
             signed = [False for _ in out]
-            self.data.write_properties(self.shape, out, signed)
+            self.data.write_properties(out, signed)
             self.busy.write(False)
         self.measure_time.write(timer.interval)
         self.update_ui.emit()

yaqc_cmds/somatic/_wt5.py

@@ -91,27 +91,43 @@ def create_data(path, headers, destinations, axes, constants, hardware, sensors)
         channel_shapes = {}
         channel_units = {}
 
+        transform_extras = []
+
         for sensor in sensors:
             # TODO allow sensors to be inactive
-            # TODO allow multi-D sensors
             sensor.active = True
 
-            # InGaAs array detector
-            if hasattr(sensor.driver.client, "get_map"):
-                map_shape = full_scan_shape + tuple(sensor.shape)
-                full_scan_shape = full_scan_shape + (1,)
-                for k, v in variable_shapes.items():
-                    variable_shapes[k] = v + (1,)
-                for k, v in channel_shapes.items():
-                    channel_shapes[k] = v + (1,)
-
-                variable_shapes["wa"] = map_shape
-                variable_units["wa"] = "nm"
-                variable_labels["wa"] = "a"
-
-                for ch in sensor.channel_names:
-                    channel_shapes[ch] = map_shape
-                    channel_units[ch] = None
+            if "has-mapping" in sensor.driver.client.traits:
+                channel_mappings = sensor.driver.client.get_channel_mappings()
+                shapes = sensor.driver.client.get_channel_shapes()
+                chs_left = set(channel_mappings.keys())
+
+                variable_units.update(sensor.driver.client.get_mapping_units())
+                channel_units.update(sensor.driver.client.get_channel_units())
+                mappings = sensor.driver.client.get_mappings()
+
+                while chs_left:
+                    chs_broad = [chs_left.pop()]
+                    maps_broad = channel_mappings[chs_broad[0]]
+                    ndim = len(shapes[chs_broad[0]])
+                    for ch in chs_left:
+                        if any(i in mappings[ch] for i in maps_broad):
+                            maps_broad.extend(channel_mappings[ch])
+                            chs_broad.append(ch)
+                            chs_left.remove(ch)
+                    for k, v in variable_shapes.items():
+                        variable_shapes[k] = v + (1,) * ndim
+                    for k, v in channel_shapes.items():
+                        channel_shapes[k] = v + (1,) * ndim
+                    variable_shapes.update(
+                        {m: full_scan_shape + mappings[m].shape for m in maps_broad}
+                    )
+                    transform_extras.extend(maps_broad)
+                    channel_shapes.update(
+                        {ch: full_scan_shape + tuple(shapes[ch]) for ch in chs_broad}
+                    )
+                    full_scan_shape += (1,) * ndim
+                # TODO channels with _no_ mapping?
 
             else:
                 for ch in sensor.channel_names:
@@ -140,8 +156,7 @@ def create_data(path, headers, destinations, axes, constants, hardware, sensors)
         # When we have some better hinting in WT itself, this should likely be reverted
         # KFS 2020-11-13
         transform = [f"{a.name}_points" if hasattr(a, "points") else a.name for a in axes]
-        if "wa" in variable_shapes:
-            transform += ["wa"]
+        transform.extend(transform_extras)
         data.transform(*transform)
 
         for ch, sh in channel_shapes.items():
@@ -163,9 +178,26 @@ def write_data(idx, hardware, sensors):
             for rec, (obj, *_) in hw.recorded.items():
                 data[rec][idx] = obj.read(data[rec].units)
         for s in sensors:
-            for ch, val in s.channels.items():
-                data[ch][idx] = val
-            if s.shape[0] > 1:
-                data["wa"][idx] = s.driver.client.get_map(data["wm"][idx])
+            if "has-mapping" in s.driver.client.traits:
+                for ch, val in s.channels.items():
+                    idx_ch = list(data[ch].shape)
+                    print(idx_ch)
+                    idx_ch = [slice(None) if i != 1 else 1 for i in idx_ch]
+                    idx_ch[: len(in_idx)] = in_idx
+                    idx_ch = tuple(idx_ch)
+                    print(data[ch].shape, idx_ch, val.shape)
+                    data[ch][idx_ch] = val
+                for var, val in s.driver.client.get_mappings().items():
+                    if var == "mapping_id":
+                        continue
+                    idx_map = list(data[var].shape)
+                    idx_map = [slice(None) if i != 1 else 1 for i in idx_map]
+                    idx_map[: len(in_idx)] = in_idx
+                    idx_map = tuple(idx_map)
+                    print(data[var].shape, idx_map, val.shape)
+                    data[var][idx_map] = val
+            else:
+                for ch, val in s.channels.items():
+                    data[ch][idx] = val
         data.flush()
         data_container.last_idx_written = in_idx

yaqc_cmds/somatic/modules/scan.py

@@ -125,12 +125,24 @@ def process(self, scan_folder):
             for channel_name in channels:
                 channel_path = data_folder / channel_name
                 output_path = data_folder
-                if data.ndim > 2:
+                orig_transform = data.axis_names
+                print(data.axis_names)
+                axes = []
+                for a in orig_transform:
+                    print(data[channel_name].shape, data[a].shape)
+                    if all(
+                        ch_sh >= a_sh
+                        for ch_sh, a_sh in zip(data[channel_name].shape, data[a].shape)
+                    ):
+                        axes.append(a)
+                data.transform(*axes)
+                print(data.axis_names)
+                if sum(a > 1 for a in data[channel_name].shape) > 2:
                     output_path = channel_path
                     channel_path.mkdir()
                 channel_index = data.channel_names.index(channel_name)
                 image_fname = channel_name
-                if data.ndim == 1:
+                if sum(a > 1 for a in data[channel_name].shape) == 1:
                     outs = wt.artists.quick1D(
                         data,
                         channel=channel_index,
@@ -152,6 +164,7 @@ def process(self, scan_folder):
                     )
                 if channel_name == main_channel:
                     outputs = outs
+                data.transform(*orig_transform)
             # get output image
             if len(outputs) == 1:
                 output_image_path = outputs[0]