From b479627df97d3ee65426a2baf8b8820afb8f3081 Mon Sep 17 00:00:00 2001
From: Runner PC Cavity Lab <johannes.schabbauer@tuwien.ac.at>
Date: Fri, 7 Mar 2025 17:01:12 +0100
Subject: [PATCH] TimeBase: Added segment mode and more flexible channel names
 (for DIM3000 switch)

---
 TimeBaseAOMDriver/blacs_tabs.py        | 10 ++++----
 TimeBaseAOMDriver/blacs_workers.py     |  8 ++++--
 TimeBaseAOMDriver/labscript_devices.py | 35 +++++++++++++++++++++++---
 3 files changed, 42 insertions(+), 11 deletions(-)

diff --git a/TimeBaseAOMDriver/blacs_tabs.py b/TimeBaseAOMDriver/blacs_tabs.py
index 6b50986..d2180d3 100644
--- a/TimeBaseAOMDriver/blacs_tabs.py
+++ b/TimeBaseAOMDriver/blacs_tabs.py
@@ -16,12 +16,12 @@ class TimeBaseTab(DeviceTab):
         
         # Create DDS Output objects
         dds_prop = {}
-        for i in range(1,props["num_AOM"]+1):
-            if connection.find_child(self.device_name,f"CH-{i}") is None:
-                continue
-            dds_prop[f"CH-{i}"] = {"gate":{}}
+        channels = [ch.parent_port for ch in connection.child_list.values()]
+        channels.sort() # Make sure they are shown in alphabetical order in BLACS
+        for ch in channels:
+            dds_prop[ch] = {"gate":{}}
             for subchnl in ['freq', 'amp']:
-                dds_prop[f"CH-{i}"][subchnl] = {'base_unit':self.base_units[subchnl], 
+                dds_prop[ch][subchnl] = {'base_unit':self.base_units[subchnl], 
                                                 'min':self.base_min[subchnl],
                                                 'max':self.base_max[subchnl],
                                                 'step':self.base_step[subchnl],
diff --git a/TimeBaseAOMDriver/blacs_workers.py b/TimeBaseAOMDriver/blacs_workers.py
index 8e24072..5803f29 100644
--- a/TimeBaseAOMDriver/blacs_workers.py
+++ b/TimeBaseAOMDriver/blacs_workers.py
@@ -46,7 +46,11 @@ class TimeBaseWorker(Worker):
                     if self.smart_cache[channel].get(attr) != group[channel].attrs[attr] or fresh:
                         any_changed = True
                         print(f"Programming {channel}, {attr}={value}")
-                        self.client.send(f"{channel}|{attr}:{value:.0f}\r\n".encode())
+                        if ":" in attr:
+                            # For segment mode
+                            self.client.send(f"{channel}|{attr}{value}\r\n".encode())
+                        else:
+                            self.client.send(f"{channel}|{attr}:{value:.0f}\r\n".encode())
                         self.smart_cache[channel][attr] = group[channel].attrs[attr]
                         time.sleep(0.01) # TODO: What's the best time here? Do we need to wait at all? According to the DIM-3000 manual the rate is 10 commands per second, it seems router can also handle faster times?
                 if group[channel].attrs["Sswpm"] != 0 and any_changed:
@@ -119,7 +123,7 @@ class TimeBaseWorker(Worker):
                 self.smart_cache[channel]["Sout"] = dds["gate"]
 
     def shutdown(self):
-        # for i in range(1,self.num_AOM+1):
+        # for i in self.channels:
         #     ch = f"CH-{i}"
         #     self.client.send(ch.encode() + b"|Sout:0\r\n")
         self.client.shutdown(socket.SHUT_RDWR)
diff --git a/TimeBaseAOMDriver/labscript_devices.py b/TimeBaseAOMDriver/labscript_devices.py
index 20b0441..f435151 100644
--- a/TimeBaseAOMDriver/labscript_devices.py
+++ b/TimeBaseAOMDriver/labscript_devices.py
@@ -4,7 +4,7 @@ from labscript_utils.unitconversions.detuning import detuning
 import numpy as np
 
 class TimeBaseAOM(Device):
-    description = "AOM conrolled by TimeBase AOM driver"
+    description = "AOM controlled by TimeBase AOM driver"
     allowed_children = [StaticAnalogQuantity]
     # This device is an adaption similar to the StatisDDS class 
     # from labscript, for AOM channels of the TimeBase driver
@@ -12,10 +12,14 @@ class TimeBaseAOM(Device):
     def __init__(
             self, name, parent_device, connection, 
             digital_gate=[], unit_conversion_parameters={},
+            supports_segment_mode = False,
             **kwargs):
 
         Device.__init__(self,name,parent_device,connection, **kwargs)
 
+        if supports_segment_mode:
+            self.segments = []
+
         # If the channel is instantiated with unit conversion parameters,
         # use them to calculate the right detuning. If not just convert 
         # between Hz and MHz.
@@ -83,6 +87,22 @@ class TimeBaseAOM(Device):
         set_value = int(np.round(np.log2(value/3200)))
         self.FM_deviation.constant(set_value)
 
+    def add_segment(self,start_freq,stop_freq=None,step_freq=0,step_time=3.90625e-9):
+        step_time = round(step_time/3.90625e-9)
+        if step_freq==0 and stop_freq is None:
+            stop_freq = start_freq
+        
+        # Do some error checks
+        if self.parent_device.min_freq > start_freq or start_freq > self.parent_device.max_freq:
+            raise LabscriptError(f"Start frequency of segment must be within [{self.parent_device.min_freq},{self.parent_device.max_freq}]Hz!")
+        if self.parent_device.min_freq > stop_freq or stop_freq > self.parent_device.max_freq:
+            raise LabscriptError(f"Stop frequency of segment must be within [{self.parent_device.min_freq},{self.parent_device.max_freq}]Hz!")
+        if step_freq < 0 or start_freq > self.parent_device.max_freq:
+            raise LabscriptError(f"Frequency step of segment must be within [0,{self.parent_device.max_freq}]!")
+        if step_time < 0 or step_time > 65535:
+            raise LabscriptError(f"Start frequency of segment must be within [3.90625,255.996]ns!")
+
+        self.segments.append(f"1;{round(start_freq)};{round(stop_freq)};{round(step_freq)};{step_time};0")
 
 
 class TimeBaseFreqSweepTrigger(DigitalOut):
@@ -146,12 +166,11 @@ class TimeBaseAOMDriver(Device):
     @set_passed_properties(
         property_names={
             "connection_table_properties":
-            ["hostname","port","mock","num_AOM","min_freq","max_freq","min_amp","max_amp","FM_deviations","max_step","min_step","max_sweep_time","min_sweep_time"]
+            ["hostname","port","mock","min_freq","max_freq","min_amp","max_amp","FM_deviations","max_step","min_step","max_sweep_time","min_sweep_time"]
     })
-    def __init__(self, name, hostname, port=8081, num_AOM=5, mock=False, **kwargs):
+    def __init__(self, name, hostname, port=8081, mock=False, **kwargs):
         super().__init__(name, parent_device=None, connection="None", **kwargs)
         self.BLACS_connection = hostname
-        self.num_AOM = num_AOM
 
 
     def generate_code(self, hdf5_file):
@@ -179,6 +198,14 @@ class TimeBaseAOMDriver(Device):
             else:
                 AOM_group.attrs["Sfmon"] = 0
 
+            # If we use the segment mode, write the segemnt commands
+            if hasattr(AOM,"segments"):
+                if len(AOM.segments)>20:
+                    raise LabscriptError(f"{AOM.name}: Cannot have more than 20 segments!")
+                for i,segment in enumerate(AOM.segments):
+                    AOM_group.attrs[f"Wseg:{i+1}"] = segment
+                AOM_group.attrs[f"Mseg"] = 1 if len(AOM.segemnts)>0 else 0
+
         # There is no need to call Device.generate_code(), because nothing is done in the TimeBaseAOM class.
 
         
-- 
GitLab