From 03e048170771188d85deb83dc98bc62ae18af782 Mon Sep 17 00:00:00 2001
From: Johannes Schabbauer <johannes.schabbauer@tuwien.ac.at>
Date: Wed, 12 Feb 2025 13:52:19 +0100
Subject: [PATCH] SpectrumAWG: Added zmq image receiver socket, the can get the
 Tweezer image during the shot to reprogram the sequence of the AWG output.
 Tested with dummy devices, not with real hardware yet.

---
 SpectrumAWG/blacs_workers.py     | 24 +++++++++++++++++++++++-
 SpectrumAWG/labscript_devices.py |  9 +++++++--
 2 files changed, 30 insertions(+), 3 deletions(-)

diff --git a/SpectrumAWG/blacs_workers.py b/SpectrumAWG/blacs_workers.py
index 9be404a..e4ff744 100644
--- a/SpectrumAWG/blacs_workers.py
+++ b/SpectrumAWG/blacs_workers.py
@@ -2,7 +2,10 @@ import labscript_utils.h5_lock
 import h5py
 from blacs.tab_base_classes import Worker
 from . import SpectrumCard
-import numpy as np
+try:
+    from pythonlib.conditional_tweezers.OccupationReceiver import OccupationReceiver
+except ImportError:
+    print("Custom Class to reprogram Tweezers not found...")
 
 class SpectrumAWGWorker(Worker):
     def init(self):
@@ -36,6 +39,11 @@ class SpectrumAWGWorker(Worker):
         # Initialize memory for smart programming
         # Keys: hash of instructions, Values: position in memory
         self.smart_cache = {}
+
+        self.occupations = []
+        if getattr(self,"occupation_receiver_port",0):
+            print(f"Creat occupation receiver socket at port {self.occupation_receiver_port}")
+            self.occupation_receiver = OccupationReceiver(self.occupation_receiver_port, self.occupations, self.AWG)
     
     def program_manual(self, values):
         if values is None:
@@ -63,6 +71,7 @@ class SpectrumAWGWorker(Worker):
         write_setup_and_start = False
         with h5py.File(h5_file,'r') as f:
             group = f[f"devices/{device_name}"]
+            function_conditional_programming = group.attrs.get("function_conditional_programming", "")
             for ch in self.channels:
                 if fresh or len(self.smart_cache)+len(group[ch].attrs) > self.memory_segments:
                     # Reset smart programming and start writing memory from the beginning
@@ -116,6 +125,19 @@ class SpectrumAWGWorker(Worker):
             self.AWG.card_write_setup() # TODO: Do we have to call that every shot or just once after the initialization?
             self.AWG.card_start()
             self.AWG.card_enable_trigger()
+
+        if function_conditional_programming:
+            print("Start conditional programming of AWG...")
+            if not (len(instruction)-1)%3 == 0:
+                raise RuntimeError("Current implementation only supports multitone for conditional programming")
+            num_samples = int(instruction[0])
+            num_tones = (len(instruction)-1)//3
+            freq = instruction[1:num_tones+1]
+            ampl = instruction[num_tones+1:2*num_tones+1]
+            phase= instruction[2*num_tones+1:]
+            self.occupation_receiver.last_tweezer_params(index+1, len(self.smart_cache), freq, ampl, phase, num_samples, self.sample_rate)
+            self.occupation_receiver.set_function(function_conditional_programming)
+
         return initial_values
 
     def transition_to_manual(self):
diff --git a/SpectrumAWG/labscript_devices.py b/SpectrumAWG/labscript_devices.py
index 4dd1524..b58c58f 100644
--- a/SpectrumAWG/labscript_devices.py
+++ b/SpectrumAWG/labscript_devices.py
@@ -121,11 +121,11 @@ class SpectrumAWG(Device):
     allowed_children = [AWGOutput]
 
     @set_passed_properties(
-        property_names={"connection_table_properties": ["device_path","timeout","external_clock_rate","sample_rate","memory_segments"],
+        property_names={"connection_table_properties": ["device_path","timeout","external_clock_rate","sample_rate","memory_segments","occupation_receiver_port"],
                         "device_properties": []                
         }
         )
-    def __init__(self, name, device_path, sample_rate, external_clock_rate=None, timeout=5000, channel_mode="seq", memory_segments=2**16, **kwargs):
+    def __init__(self, name, device_path, sample_rate, external_clock_rate=None, timeout=5000, channel_mode="seq", memory_segments=2**16, occupation_receiver_port=None, **kwargs):
         """ Create SpectrumAWG instance.
         
         Parameters
@@ -154,6 +154,9 @@ class SpectrumAWG(Device):
         internal_memory = 2**32 # 4GB
         self.max_sample_size = internal_memory//2//memory_segments
 
+    def set_conditional_reprogramming(self, function_name):
+        self.function_conditional_programming = function_name
+
     def do_checks(self):
         if len(self.child_devices)>1:
             raise NotImplementedError("This code can just handle 1 output channel for now.")
@@ -170,6 +173,8 @@ class SpectrumAWG(Device):
             group.require_group(output.connection)
             group[output.connection].require_group("labels")
 
+            group.attrs["function_conditional_programming"] = getattr(self,"function_conditional_programming","")
+
             for i,t in enumerate(np.sort(change_times)):
                 group[output.connection].attrs[str(i)] = output.instructions[t][:-1]
                 if output.instructions[t][-1] is not None:
-- 
GitLab