Prawn device updates

labscript_devices/PrawnBlaster/blacs_workers.py

@@ -15,6 +15,7 @@
 import h5py
 import numpy as np
 from blacs.tab_base_classes import Worker
+from labscript import LabscriptError
 from labscript_utils.connections import _ensure_str
 import labscript_utils.properties as properties
 
@@ -56,28 +57,85 @@ def init(self):
         self.wait_timeout = None
         self.h5_file = None
         self.started = False
-
-        self.prawnblaster = serial.Serial(self.com_port, 115200, timeout=1)
+        self.min_version = (1, 1, 0)
+
+        self.conn = serial.Serial(self.com_port, 115200, timeout=1)
         self.check_status()
 
         # configure number of pseudoclocks
-        self.prawnblaster.write(b"setnumpseudoclocks %d\r\n" % self.num_pseudoclocks)
-        assert self.prawnblaster.readline().decode() == "ok\r\n"
+        self.send_command_ok(f"setnumpseudoclocks {self.num_pseudoclocks}")
 
         # Configure pins
         for i, (out_pin, in_pin) in enumerate(zip(self.out_pins, self.in_pins)):
-            self.prawnblaster.write(b"setoutpin %d %d\r\n" % (i, out_pin))
-            assert self.prawnblaster.readline().decode() == "ok\r\n"
-            self.prawnblaster.write(b"setinpin %d %d\r\n" % (i, in_pin))
-            assert self.prawnblaster.readline().decode() == "ok\r\n"
+            self.send_command_ok(f"setoutpin {i} {out_pin}")
+            self.send_command_ok(f"setinpin {i} {in_pin}")
 
-        # Check if fast serial is available
         version, _ = self.get_version()
-        self.fast_serial = version >= (1, 1, 0)
+        print(f'Connected to version: {version}')
+        assert version >= self.min_version, f'Incompatible firmware, must be >= {self.min_version}'
+
+        if version >= (1, 2, 0):
+            board = self.get_board()
+        else:
+            board = 'pico1'
+            print(f'Version {version} too low to use pico2 firmware, consider upgrading firmware')
+
+        print(f'Connected to board: {board}')
+        if board.strip() != self.pico_board.strip():
+            raise LabscriptError(f'firmware reports {board} attached, labscript expects {self.pico_board}')
+
+        current_status = self.read_status()
+        print(f'Current status is {current_status}')
+
+    def _read_full_buffer(self):
+        '''Used to get any extra lines from device after a failed send_command'''
+
+        resp = self.conn.readlines()
+        str_resp = ''.join([st.decode() for st in resp])
+
+        return str_resp
+
+    def send_command(self, command, readlines=False):
+        '''Sends the supplied string command and checks for a response.
+
+        Automatically applies the correct termination characters.
+
+        Args:
+            command (str): Command to send. Termination and encoding is done automatically.
+            readlines (bool, optional): Use pyserial's readlines functionality to read multiple
+                response lines. Slower as it relies on timeout to terminate reading.
+
+        Returns:
+            str: String response from the PrawnBlaster
+        '''
+        command += '\r\n'
+        self.conn.write(command.encode())
+
+        if readlines:
+            str_resp = self._read_full_buffer()
+        else:
+            str_resp = self.conn.readline().decode()
 
+        return str_resp
+
+    def send_command_ok(self, command):
+        '''Sends the supplied string command and confirms 'ok' response.
+
+        Args:
+            command (str): String command to send.
+
+        Raises:
+            LabscriptError: If response is not `ok\\r\\n`
+        '''
+
+        resp = self.send_command(command)
+        if resp != 'ok\r\n':
+            # get complete error message
+            resp += self._read_full_buffer()
+            raise LabscriptError(f"Command '{command:s}' failed. Got response '{repr(resp)}'")
+
     def get_version(self):
-        self.prawnblaster.write(b"version\r\n")
-        version_str = self.prawnblaster.readline().decode()
+        version_str = self.send_command('version', readlines=True)
         assert version_str.startswith("version: ")
         version = version_str[9:].strip()
 
@@ -91,6 +149,17 @@ def get_version(self):
         assert len(version) == 3
 
         return version, overclock
+
+    def get_board(self):
+        '''Responds with pico board version.
+
+        Returns:
+            (str): Either "pico1" for a Pi Pico 1 board or "pico2" for a Pi Pico 2 board.'''
+        resp = self.send_command('board')
+        assert resp.startswith('board:'), f'Board command failed, got: {resp}'
+        pico_str = resp.split(':')[-1].strip()
+
+        return pico_str
 
     def check_status(self):
         """Checks the operational status of the PrawnBlaster.
@@ -128,8 +197,7 @@ def check_status(self):
         ):
             # Try to read out wait. For now, we're only reading out waits from
             # pseudoclock 0 since they should all be the same (requirement imposed by labscript)
-            self.prawnblaster.write(b"getwait %d %d\r\n" % (0, self.current_wait))
-            response = self.prawnblaster.readline().decode()
+            response = self.send_command(f'getwait {0}, {self.current_wait}')
             if response != "wait not yet available\r\n":
                 # Parse the response from the PrawnBlaster
                 wait_remaining = int(response)
@@ -202,8 +270,7 @@ def read_status(self):
                 - **clock-status** (int): Clock status code
         """
 
-        self.prawnblaster.write(b"status\r\n")
-        response = self.prawnblaster.readline().decode()
+        response = self.send_command("status", readlines=True)
         match = re.match(r"run-status:(\d) clock-status:(\d)(\r\n)?", response)
         if match:
             return int(match.group(1)), int(match.group(2))
@@ -231,11 +298,9 @@ def program_manual(self, values):
             pin = int(channel.split()[1])
             pseudoclock = self.out_pins.index(pin)
             if value:
-                self.prawnblaster.write(b"go high %d\r\n" % pseudoclock)
+                self.send_command_ok(f"go high {pseudoclock}")
             else:
-                self.prawnblaster.write(b"go low %d\r\n" % pseudoclock)
-
-            assert self.prawnblaster.readline().decode() == "ok\r\n"
+                self.send_command_ok(f"go low {pseudoclock}")
 
         return values
 
@@ -309,15 +374,13 @@ def transition_to_buffered(self, device_name, h5file, initial_values, fresh):
         clock_frequency = self.device_properties["clock_frequency"]
 
         # Now set the clock details
-        self.prawnblaster.write(b"setclock %d %d\r\n" % (clock_mode, clock_frequency))
-        response = self.prawnblaster.readline().decode()
-        assert response == "ok\r\n", f"PrawnBlaster said '{response}', expected 'ok'"
+        response = self.send_command_ok(f"setclock {clock_mode} {clock_frequency}")
 
         # Program instructions
         for pseudoclock, pulse_program in enumerate(pulse_programs):
             total_inst = len(pulse_program)
             # check if it is more efficient to fully refresh
-            if not fresh and self.smart_cache[pseudoclock] is not None and self.fast_serial:
+            if not fresh and self.smart_cache[pseudoclock] is not None:
                 # get more convenient handles to smart cache arrays
                 curr_inst = self.smart_cache[pseudoclock]
 
@@ -337,17 +400,17 @@ def transition_to_buffered(self, device_name, h5file, initial_values, fresh):
                 if new_inst / total_inst > 0.1:
                     fresh = True
 
-            if (fresh or self.smart_cache[pseudoclock] is None) and self.fast_serial:
+            if (fresh or self.smart_cache[pseudoclock] is None):
                 print('binary programming')
-                self.prawnblaster.write(b"setb %d %d %d\r\n" % (pseudoclock, 0, len(pulse_program)))
-                response = self.prawnblaster.readline().decode()
+                self.conn.write(b"setb %d %d %d\r\n" % (pseudoclock, 0, len(pulse_program)))
+                response = self.conn.readline().decode()
                 assert (
                     response == "ready\r\n"
                 ), f"PrawnBlaster said '{response}', expected 'ready'"
                 program_array = np.array([pulse_program['half_period'],
                                           pulse_program['reps']], dtype='<u4').T
-                self.prawnblaster.write(program_array.tobytes())
-                response = self.prawnblaster.readline().decode()
+                self.conn.write(program_array.tobytes())
+                response = self.conn.readline().decode()
                 assert (
                     response == "ok\r\n"
                 ), f"PrawnBlaster said '{response}', expected 'ok'"
@@ -361,7 +424,7 @@ def transition_to_buffered(self, device_name, h5file, initial_values, fresh):
 
                     # Only program instructions that differ from what's in the smart cache:
                     if self.smart_cache[pseudoclock][i] != instruction:
-                        self.prawnblaster.write(
+                        self.conn.write(
                             b"set %d %d %d %d\r\n"
                             % (
                                 pseudoclock,
@@ -370,7 +433,7 @@ def transition_to_buffered(self, device_name, h5file, initial_values, fresh):
                                 instruction["reps"],
                             )
                         )
-                        response = self.prawnblaster.readline().decode()
+                        response = self.conn.readline().decode()
                         assert (
                             response == "ok\r\n"
                         ), f"PrawnBlaster said '{response}', expected 'ok'"
@@ -392,9 +455,7 @@ def start_run(self):
 
         # Start in software:
         self.logger.info("sending start")
-        self.prawnblaster.write(b"start\r\n")
-        response = self.prawnblaster.readline().decode()
-        assert response == "ok\r\n", f"PrawnBlaster said '{response}', expected 'ok'"
+        self.send_command_ok("start")
 
         # set started = True
         self.started = True
@@ -405,9 +466,7 @@ def wait_for_trigger(self):
 
         # Set to wait for trigger:
         self.logger.info("sending hwstart")
-        self.prawnblaster.write(b"hwstart\r\n")
-        response = self.prawnblaster.readline().decode()
-        assert response == "ok\r\n", f"PrawnBlaster said '{response}', expected 'ok'"
+        self.send_command_ok("hwstart")
 
         running = False
         while not running:
@@ -477,7 +536,7 @@ def transition_to_manual(self):
     def shutdown(self):
         """Cleanly shuts down the connection to the PrawnBlaster hardware."""
 
-        self.prawnblaster.close()
+        self.conn.close()
 
     def abort_buffered(self):
         """Aborts a currently running buffered execution.
@@ -489,8 +548,8 @@ def abort_buffered(self):
             # Only need to send abort signal if we have told the PrawnBlaster to wait
             # for a hardware trigger. Otherwise it's just been programmed with
             # instructions and there is nothing we need to do to abort.
-            self.prawnblaster.write(b"abort\r\n")
-            assert self.prawnblaster.readline().decode() == "ok\r\n"
+            self.conn.write(b"abort\r\n")
+            assert self.conn.readline().decode() == "ok\r\n"
             # loop until abort complete
             while self.read_status()[0] != 5:
                 time.sleep(0.5)

labscript_devices/PrawnBlaster/labscript_devices.py

@@ -144,14 +144,16 @@ class PrawnBlaster(PseudoclockDevice):
     """Minimum required length of a wait before retrigger can be detected.
     Corresponds to 4 instructions."""
     allowed_children = [_PrawnBlasterPseudoclock, _PrawnBlasterDummyPseudoclock]
-    max_instructions = 30000
-    """Maximum numaber of instructions per pseudoclock. Max is 30,000 for a single
-    pseudoclock."""
+    allowed_boards = ['pico1', 'pico2']
+    max_instructions_map = {'pico1' : 30000, 'pico2' : 60000}
+    """Maximum number of instructions per pseudoclock for each pico board. """
+    max_frequency_map = {'pico1' : 133e6, 'pico2' : 150e6}
 
     @set_passed_properties(
         property_names={
             "connection_table_properties": [
                 "com_port",
+                'pico_board',
                 "in_pins",
                 "out_pins",
                 "num_pseudoclocks",
@@ -175,6 +177,7 @@ def __init__(
         trigger_device=None,
         trigger_connection=None,
         com_port="COM1",
+        pico_board = 'pico1',
         num_pseudoclocks=1,
         out_pins=None,
         in_pins=None,
@@ -191,6 +194,7 @@ def __init__(
             name (str): python variable name to assign to the PrawnBlaster
             com_port (str): COM port assigned to the PrawnBlaster by the OS. Takes 
                 the form of `'COMd'`, where `d` is an integer.
+            pico_board (str): The version of pico board used, pico1 or pico2.
             num_pseudoclocks (int): Number of pseudoclocks to create. Ranges from 1-4.
             trigger_device (:class:`~labscript.IntermediateDevice`, optional): Device
                 that will send the hardware start trigger when using the PrawnBlaster
@@ -203,8 +207,8 @@ def __init__(
                 triggering. Must have length of at least `num_pseudoclocks`. 
                 Defaults to `[0,0,0,0]`
             clock_frequency (float, optional): Frequency of clock. Standard range
-                accepts up to 133 MHz. An experimental overclocked firmware is 
-                available that allows higher frequencies.
+                accepts up to 133 MHz for pico1, 150 MHz for pico2. An experimental overclocked 
+                firmware is available that allows higher frequencies.
             external_clock_pin (int, optional): If not `None` (the default), 
                 the PrawnBlaster uses an external clock on the provided pin. Valid
                 options are `20` and `22`. The external frequency must be defined
@@ -214,14 +218,26 @@ def __init__(
 
         """
 
+        if pico_board in self.allowed_boards:
+            self.pico_board = pico_board
+        else:
+            raise LabscriptError(f'Pico board specified not in {self.allowed_boards}')
+
         # Check number of pseudoclocks is within range
         if num_pseudoclocks < 1 or num_pseudoclocks > 4:
             raise LabscriptError(
                 f"The PrawnBlaster {name} only supports between 1 and 4 pseudoclocks"
             )
-
+
+        # Set max instructions based on board
+        self.max_instructions = self.max_instructions_map[self.pico_board]
         # Update the specs based on the number of pseudoclocks
         self.max_instructions = self.max_instructions // num_pseudoclocks
+
+        self.max_frequency = self.max_frequency_map[self.pico_board]
+
+        if clock_frequency > self.max_frequency:
+            raise ValueError(f'Clock frequency must be less than {int(self.max_frequency * 10**-6)} MHz')
         # Update the specs based on the clock frequency
         if self.clock_resolution != 2 / clock_frequency:
             factor = (2 / clock_frequency) / self.clock_resolution