#!/usr/bin/env python3 # -*- coding: utf-8 -*- """ This script can be used to talk to SingleQuantum Retina Nerve. This project is licensed under the terms of the MIT license. Copyright (c) 2023 Single Quantum B. V. """ import json import struct import threading import time import uuid from collections import OrderedDict from ctypes import Structure, c_uint8, c_uint32 from urllib.parse import urlencode, urlsplit, urlunsplit from urllib.request import Request, urlopen from typing import Any, Dict, List, Optional, Tuple, Union import requests import websockets from websockets.sync.client import connect # The time (ms) of a CU (channel unit) cycle. CU_INTTIME = 10 COUNTS_MESSAGE_SIZE = 9 def unpack_counts_header(offset: int, buffer: bytes) -> Dict[str, int]: """Unpack the websocket counts header.""" time, int_size = struct.unpack_from(" Dict[str, int]: """Unpack one channel-unit counts record from a binary payload.""" status = buffer[offset] channel_number = struct.unpack_from(" Dict[str, int]: """Convert this structure into a plain dictionary.""" return { "cuStatus": self.cuStatus, "channel_number": self.channel_number, "counts": self.counts, } class JsonRpc(object): """This class takes the `api_url` and then can be used to send standard HTTP requests with `request` or json rpc requests with `jsonrpc`. """ def __init__(self, api_url: str, jsonrpc_version: str = "2.0"): """Initialize a JsonRpc class. Parameters ---------- api_url : str The URL of the api endpoint. jsonrpc_version : str The JSON RPC version this endpoint uses. """ self.api_url = api_url self.jsonrpc_version = jsonrpc_version def request(self, params: Optional[list] = None, payload: Optional[Dict] = None): """Perform a GET HTTP request, if given a payload a POST. Parameters ---------- params : list, optional A list of parameters added to the `self.api_url`. payload : dict, optional If provided sends `payload` as JSON with a POST. Returns ------- dict Returns the result as a dict converted from the json received. Raises ------ AssertionError Raises an assertion error when the HTTP response code is not 200. """ headers = {} request_data = None headers["Accept"] = "application/json" target = self.api_url if params: target += "?" + urlencode(params, doseq=True, safe="/") if payload: headers["Content-Type"] = "application/json; charset=UTF-8" request_data = json.dumps(payload).encode() http_request = Request(target, data=request_data, headers=headers) http_response = urlopen(http_request, timeout=10) content_charset = "utf-8" # default if hasattr(http_response, "status"): # Python3 only assert http_response.status == 200, ( "Got HTTP " + str(http_response.status) + " with " + http_response.read().decode(content_charset) ) content_charset = http_response.headers.get_content_charset(content_charset) body = http_response.read().decode(content_charset) return json.loads(body) def jsonrpc(self, method: str, **params): """Makes a JSON RPC request to the `self.api_url`. Parameters ---------- method : str The name of the method you want to call. params : dict, optional Keyword arguments that are sent as JSON-RPC parameters. Returns ------- dict with keys as str The result of the function as a dictionary. Raises ------ AssertionError If the response is malformed. ValueError If the endpoint returns an ``error`` message. """ identifier = str(uuid.uuid4()) payload = { "method": method, "params": params or [], "jsonrpc": self.jsonrpc_version, "id": identifier, } response_data = self.request(payload=payload) assert response_data["jsonrpc"], "No jsonrpc response" assert response_data["id"] == identifier, "Incorrect identifier in response" if "result" in response_data: return response_data["result"] elif "error" in response_data: if "message" in response_data["error"]: raise ValueError(response_data["error"]["message"]) else: return None def synchronized_method(method, *args, **kws): outer_lock = threading.Lock() lock_name = "__" + method.__name__ + "_lock" + "__" def sync_method(self, *args, **kws): with outer_lock: if not hasattr(self, lock_name): setattr(self, lock_name, threading.Lock()) lock = getattr(self, lock_name) with lock: return method(self, *args, **kws) sync_method.__name__ = method.__name__ sync_method.__doc__ = method.__doc__ sync_method.__module__ = method.__module__ return sync_method class WebSocketThread(threading.Thread): """ Background thread that receives counts over the websocket stream. """ def __init__(self, websocket_url: str): threading.Thread.__init__( self, name=f"Retina Websocket Thread, for: {websocket_url}" ) self.lock = threading.Lock() self.rlock = threading.RLock() self.websocket_url = websocket_url self.websocket = connect(self.websocket_url) self.default_buffer_size = 5 self.buffer_size = self.default_buffer_size self.message_buffer = OrderedDict() self.latest_message_buffer = {} self.update_interval = 50 # ms self.shutdown = False @synchronized_method def close(self): """Close the websocket.""" self.shutdown = True self.websocket.close() def get_messages_no_delay( self, channels: Optional[list[int]] = None, attribute: Optional[str] = None ): """Return the latest messages currently in the websocket buffer.""" if channels: latest_messages = [ self.latest_message_buffer.get(channel, None) for channel in channels ] if attribute: return [ msg.get(attribute, None) if msg else None for msg in latest_messages ] return latest_messages return None @synchronized_method def get_n_messages( self, number: int, channels: Optional[list[int]] = None, attribute: Optional[str] = None, ) -> list[Any]: """ Returns N number of messages. Parameters ---------- number : int The number of messages to return. channels : list, optional The channels to filter messages by. attribute : str, optional The attribute to extract from each message. Returns ------- list A list of N messages, for the provided channels. If an attribute is provided, a list of only that attribute is returned. """ if attribute: if attribute not in ["intTime", "status", "channel_number", "counts"]: raise ValueError(f"Attribute {attribute} is not a valid attribute.") needed_buffer_size = round(number + (1000 / self.update_interval)) if needed_buffer_size > self.buffer_size: self.buffer_size = needed_buffer_size starting_timestamp = None timeout = 10 # seconds start_time = time.time() while not starting_timestamp: # Awaiting the first message, this should take 1 times the integration time. time.sleep(self.update_interval / 1000) starting_timestamp = next(reversed(self.message_buffer.keys()), None) if time.time() - start_time > timeout: raise TimeoutError("Timeout while waiting for the first message.") measuring_time = self.update_interval * number # ms end_timestamp = round((starting_timestamp + measuring_time)) # ms # We wait here till we have passed our desired timestamp. waiting_for_timestamp = True while waiting_for_timestamp: current_timestamp = next(reversed(self.message_buffer.keys())) waiting_for_timestamp = not (current_timestamp >= end_timestamp + 0.1) time.sleep(0.01) relevant_messages = [ msg for ts, msg in self.message_buffer.items() if starting_timestamp <= ts < end_timestamp + 0.005 ] messages = relevant_messages if channels: for idx, msg_list in enumerate(messages): msg_list_new = [None] * len(channels) for i, channel in enumerate(channels): channel_messages = [ msg for msg in msg_list if msg["channel_number"] == channel ] msg_list_new[i] = ( channel_messages[0] if len(channel_messages) > 0 else {"channel_number": channel} ) messages[idx] = msg_list_new if attribute: messages = [ sorted(msg_list, key=lambda msg: msg["channel_number"]) for msg_list in messages ] messages = [[el.get(attribute, None) for el in msg] for msg in messages] self.buffer_size = self.default_buffer_size return messages @synchronized_method def set_update_interval(self, update_interval: int) -> None: """Change the update interval in milliseconds.""" self.update_interval = update_interval self.message_buffer.clear() def add_message(self, timestamp: int, msg: Dict): """Adds any received messages to the message buffer.""" # Direct access buffer, this can give you the latest message that was received without any delay. self.latest_message_buffer[msg["channel_number"]] = msg if timestamp not in self.message_buffer: if len(self.message_buffer) >= self.buffer_size: self.message_buffer.popitem(last=False) # Remove oldest timestamp self.message_buffer[timestamp] = [] self.message_buffer[timestamp].append(msg) def decode_counts_websocket_msg( self, msg: bytes, int_time: int ) -> List[Dict]: """Decode packed websocket counts payload into message dictionaries.""" return [ {"intTime": int_time, **unpack_counts_data(offset, msg)} for offset in range(0, len(msg), COUNTS_MESSAGE_SIZE) ] def run(self): """Main loop collecting and processing the websocket stream messages.""" failcount = 0 while self.shutdown is False: try: message = self.websocket.recv(timeout=None) datatype_byte = message[0] ws_messages = [] if datatype_byte == 0: # Counts messages counts_header = unpack_counts_header(0, message[1:13]) ws_messages = self.decode_counts_websocket_msg( message[13:], counts_header["intTime"] ) with self.lock: for msg in ws_messages: self.add_message(counts_header["time"], msg) if datatype_byte == 1: # IV data messages are ignored. continue failcount = 0 except Exception as e: if isinstance(e, websockets.ConnectionClosedOK): print("Retina counts stream connection was closed.") self.close() failcount += 1 if failcount > 10: self.close() raise Exception( f"Something went wrong while collecting messages: {e}" ) class NerveController(JsonRpc): """ This class can send requests to a SQ Nerve via the JSON RPC protocol. Set the `domain` of the SQ Nerve in the initialization. SQ Nerve is the latest control software for the Single Quantum Retina as of January 2026. """ def __init__(self, domain: str): """Initialize a NerveController class Parameters ---------- domain : str The domain the nerve controller can be accessed on. """ res = urlsplit(domain) # We assume http if only an IP is given. scheme = "http" if not res.scheme else res.scheme self.netloc = domain if not res.scheme else res.netloc self.ws_url = urlunsplit(("ws", self.netloc, "/counts", "", "")) self.ws_thread = None self.ws_thread = WebSocketThread(self.ws_url) self.ws_thread.daemon = True self.ws_thread.start() api_url = urlunsplit((scheme, self.netloc, "/api", "", "")) super(NerveController, self).__init__(api_url) # Get settings already on init self.getSettings() int_time = self.getIntTime() self.ws_thread.set_update_interval(int_time) channel_data = self.getChannels() self._channel_numbers = sorted([int(key) for key in channel_data.keys()]) def __enter__(self): return self def __exit__(self, exc_type, exc_val, exc_tb): self.close() def close(self): if self.ws_thread: self.ws_thread.close() self.ws_thread.join() def getChannels(self) -> Dict[str, Dict[str, Any]]: """ Returns a dictionary with the channels in the system, with the following structure:: { channel_number: { *channel_parameters*, }, ... } """ return self.jsonrpc("getChannels") def getChannelInformation( self, parameter_name: str, channels: Optional[List[int]] = None ) -> List[Any]: """Gets the channel information for the parameter "parameter_name". Parameters ---------- parameter_name : str The name of the parameter to get. channels : list, optional A list of channels given by their channel numbers. Returns ------- list a list containing the retrieved values for the requested quantity. """ channels_data = self.getChannels() channel_numbers = ( channels if channels else sorted([int(key) for key in channels_data.keys()]) ) return [ channels_data[str(channel_number)].get(parameter_name, None) if str(channel_number) in channels_data else None for channel_number in channel_numbers ] def setChannelUnitsValue( self, parameter_name: str, value_or_values: Union[Any, List[Any]], channels: Optional[List[int]] = None, ): """Sets the parameter 'parameter_name' to 'value_or_values' for all given channels. If no channels are provided it will set it for all channels. Parameters ---------- parameter_name : str, optional The name of the parameter to set. value_or_values : list or float The value or list of values to set for the parameter. If a single value is provided all channels will be set to that value. If a list is provided it is assumed to be sorted by channel number. channels : list, optional A list of channel numbers to set the parameter for. If this is not provided it will select all of them. Returns ------- updated_settings: dict The new and updated settings. """ if channels is None: channels = self._channel_numbers if isinstance(value_or_values, list): if len(value_or_values) != len(channels): raise ValueError( f"The amount of values provided: {len(value_or_values)} does not match the amount of channels: {len(channels)}" ) else: value_or_values = [value_or_values] * len(channels) settings_update = { "channels": { str(ch): {parameter_name: value_or_values[idx]} for idx, ch in enumerate(channels) } } return self.setSettings(**settings_update) def getTriggerV(self, channels: Optional[List[int]] = None): """Gets the trigger level for each channel. Parameters ---------- channels : list, optional A list of channel numbers to get the trigger level for. If this is not provided it will select all of them. Returns ------- list a list containing the trigger level for each requested channel """ return self.getChannelInformation("triggerV", channels=channels) def setTriggerV( self, value_or_values: List[Union[float, int]], channels: Optional[List[int]] = None, ): """Sets the trigger level for the counters for each channel (all selected channels the same value). The trigger voltage is in Volts and must be in the range (-10, 10). Parameters ---------- value_or_values : float or list The trigger level to set (in V) for all channels. Supported range: (-10, 10)V. If a list is provided it is assumed to be sorted by channel number. If a single value is provided all channels will be set to that value. channels : list, optional A list of channel numbers to set the parameter for. If this is not provided it will select all of them. Returns ------- updated_settings: dict The new and updated settings. """ return self.setChannelUnitsValue("triggerV", value_or_values, channels=channels) def getBiasI(self, channels: Optional[List[int]] = None): """Gets the bias current for each channel. Parameters ---------- channels : list, optional A list of channel numbers to get the trigger level for. If this is not provided it will select all of them. Returns ------- list a list containing the bias current for each requested channel """ return self.getChannelInformation("biasI", channels=channels) def setBiasI( self, value_or_values: List[Union[float, int]], channels: Optional[List[int]] = None, ): """Sets the bias current level for each channel (all selected channels the same value). Parameters ---------- value_or_values : float or list The bias current to set (in A). If a list is provided it is assumed to be sorted by channel number. If a single value is provided all channels will be set to that value. channels : list, optional A list of channel numbers to set the parameter for. If this is not provided it will select all of them. Returns ------- updated_settings: dict The new and updated settings. """ return self.setChannelUnitsValue("biasI", value_or_values, channels=channels) def restartSoftware(self, delay: int = 0): """Restarts the software. Parameters ---------- delay : int Optionally a delay can be passed in seconds. """ return self.jsonrpc("reboot", delay=delay) def getSettings(self) -> Dict: """Gets the current Retina driver settings.""" return self.jsonrpc("getSettings") def setSettings(self, **params: Any) -> Dict: """Set new Retina driver settings and overwrite only provided fields.""" return self.jsonrpc("setSettings", **params) def getIntTime(self) -> int: """Get the integration time in milliseconds.""" settings = self.getSettings() return settings["backend"]["intTime"] def setIntTime(self, intTime: int) -> Dict: """Set integration time in milliseconds.""" int_time_to_set = round(intTime / 10) * 10 self.ws_thread.set_update_interval(int_time_to_set) return self.setSettings(backend={"intTime": int_time_to_set}) def setStreamState(self, state: bool) -> Dict: """Turn the counts stream on or off.""" return self.setSettings(backend={"isRunning": state}) def getLog(self, lines: int = 1000) -> List[Dict]: """Returns the current log object. Parameters ---------- lines : int The number of lines/entries to get from the log Returns ------- log: list[dict] The log structed as a list of dictionary objects. """ return self.jsonrpc("getLog", lines=lines) def getTemperatureData(self) -> Dict: """Returns all the stored temperature data.""" return self.jsonrpc("getTemperature") def getTemperatures(self) -> Tuple[float, Optional[float]]: """Return the latest measured temperatures for sensor 1 and sensor 2.""" temperature_data = self.getTemperatureData() if len(temperature_data["1"]["temperature"]) == 0: return None, None if temperature_data["secondSensor"]: return temperature_data["1"]["temperature"][-1], temperature_data["2"][ "temperature" ][-1] return temperature_data["1"]["temperature"][-1], None def startIv( self, biasIStart: float, biasIStop: float, biasIStep: float, intTime: float, channels: Optional[List[int]] = None, ): """Start a IV measurement on the selected channels or if not provided all of them. Parameters ---------- biasIStart : float The current to start (in uA). biasIStop : float The current to stop (in uA) biasIStep : float The step size of the sweep (in uA). intTime : float The integration time (in ms) of a single step. channels : list, optional A list of channels given as either their rank or their location. The location of a channel is given as 'mcuId.cuId'. If this is not provided it will select all of them. Returns ------- dict The the new updated settings you have send to the server. """ if channels is None: channels = self._channel_numbers IVSweepConfiguration = { "ranks": channels, "start": biasIStart, "stop": biasIStop, "step": biasIStep, "intTime": intTime, } return self.jsonrpc("startIV", **IVSweepConfiguration) def stopIv(self) -> str: """Stop the current IV measurement that is running. Returns ------- string Success if the IV sweep was stopped successfully. """ return self.jsonrpc("stopIV") def getIvData(self, raw: bool = False) -> Dict: """ Get the data of the latest IV sweep. Parameters ---------- raw : bool If ``True``, return the raw driver payload under ``data``. Returns ------- dict A dictionary which either contains the raw data from the driver if raw is True, or a processed dictionary with lists of biasI, counts, monitorV and integration_time per channel. The keys in this dictionary are the channel numbers. """ response = requests.get(f"http://{self.netloc}/download/iv?filetype=json") json_data = response.json() if "data" not in json_data: raise ValueError("No IV data found in the response.") if raw: return json_data["data"] iv_data = {} for channel_number in json_data["data"].keys(): iv_data[channel_number] = {} channel_data = json_data["data"][channel_number]["data"] integration_time = json_data["data"][channel_number]["intTime"] iv_data[channel_number]["biasI"] = [float(i["biasI"]) for i in channel_data] iv_data[channel_number]["counts"] = [int(c["counts"]) for c in channel_data] iv_data[channel_number]["monitorV"] = [ float(v["monitorV"]) for v in channel_data ] iv_data[channel_number]["integration_time"] = integration_time return iv_data def getNMessages( self, n: int = 1, channels: Optional[List[int]] = None, attribute: Optional[str] = None, ) -> List[Any]: """ Retrieves a N number of messages from the driver. Parameters ---------- n : int Number of integration intervals to collect. channels : list, optional Channel numbers to include. If not provided, all channels are used. attribute: str, optional Optional field to extract from each message. Valid options are: ``intTime``, ``status``, ``channel_number``, ``counts``. Returns ------- list List of N lists of requested attribute for each channel, or a list of N lists of the message dictionaries. """ self.setStreamState(True) # Ensure stream is on to get the latest messages if channels is None: channels = self._channel_numbers return self.ws_thread.get_n_messages( number=n, channels=channels, attribute=attribute ) def getMessageNoDelay( self, channels: Optional[List[int]] = None, attribute: Optional[str] = None ): """ Directly returns the latest received messages in the websocket without any delay. WARNING: channel unit messages can be send/received later/earlier then others, so your result can be slightly desynced. Parameters ---------- channels : list, optional Channel numbers to include. If not provided, all channels are used. attribute: str, optional Attribute can be passed here to only return a list with the value of this attribute. Options: mcuId, cuId, cuStatus, monitorV, biasI, counts, intSize, rank, time. Returns ------- list List of requested attribute for each channel, or a list of the message dictionaries. """ if channels is None: channels = self._channel_numbers return self.ws_thread.get_messages_no_delay( channels=channels, attribute=attribute ) def getCounts(self, channels: Optional[List[int]] = None): """Get the current counts measurement. The amount of counts during the current integration time. Parameters ---------- channels : list, optional Channel numbers to include. If not provided, all channels are used. Returns ------- list List of current counts for each channel. """ counts_messages = self.getNMessages(channels=channels, n=1, attribute="counts") return counts_messages[0] def getNCounts( self, n: int = 10, channels: Optional[List[int]] = None ): """ Collects an n number of counts, for all the given channels. Parameters ---------- n : int The number of counts to collect. channels : list, optional Channel numbers to include. If not provided, all channels are used. Returns ------- list List of length n with counts for each channel. """ counts_messages = self.getNMessages( channels=channels, n=n, attribute="counts", ) return counts_messages def transformToArray(self, iv_data: Dict, quantity: str): """Transform your iv data or counts data to an array for `quantity` Parameters ---------- iv_data: dict Dictionary which is the iv data from getIvData or sweepIv or the counts data from getCounts or collectCounts quantity: str This is the name of quantity: biasI, counts, or monitorV for iv data counts, time, monitorV for count data Returns ------- result : list[list] Two-dimensional array where the first row is ``biasI`` and each following row contains ``quantity`` values for one channel. """ channel_numbers = list(iv_data.keys()) channel_numbers = sorted([int(ch_number) for ch_number in channel_numbers]) res = [iv_data[str(channel_numbers[0])]["biasI"]] + [ [] for _ in channel_numbers ] for i, channel_number in enumerate(channel_numbers): res[i + 1] = iv_data[str(channel_number)][quantity] return res def getIvHistory(self): """ Gets the IV history for all channels and returns a list of lists with the bias current in the first list and the monitor voltage for each CU as the next. Returns ------- result: list[list] the first list is the bias current the other lists are the monitor voltages (V) of the channels sorted by rank. """ ivData = self.getIvData() if len(ivData) == 0: return [] return self.transformToArray(ivData, "monitorV") def getIcHistory(self): """ Gets the IC history for all channels and returns a list of lists with the bias current in the first list and counts for each CU in following lists. Returns ------- result: list[list] the first list is the bias current the other lists are the counts of the channels sorted by rank. """ ivData = self.getIvData() if len(ivData) == 0: return [] return self.transformToArray(ivData, "counts") if __name__ == "__main__": import os websq_domain = os.environ.get("WEBSQ_DOMAIN", "http://192.168.7.3/") sq = NerveController(websq_domain)