From ccd3c0a1f906544d143f0711cd8215ce41d36de0 Mon Sep 17 00:00:00 2001 From: Sebastian Strempfer Date: Thu, 2 Oct 2025 23:45:20 -0700 Subject: [PATCH] First draft --- README.md | 1 + pyproject.toml | 30 ++++++ src/team1k/Client.py | 64 +++++++++++++ src/team1k/KTestClient.py | 191 ++++++++++++++++++++++++++++++++++++++ src/team1k/__init__.py | 2 + src/team1k/server.py | 167 +++++++++++++++++++++++++++++++++ 6 files changed, 455 insertions(+) create mode 100644 README.md create mode 100644 pyproject.toml create mode 100644 src/team1k/Client.py create mode 100644 src/team1k/KTestClient.py create mode 100644 src/team1k/__init__.py create mode 100644 src/team1k/server.py diff --git a/README.md b/README.md new file mode 100644 index 0000000..9311632 --- /dev/null +++ b/README.md @@ -0,0 +1 @@ +# `team1k` \ No newline at end of file diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 0000000..4bdb54f --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,30 @@ +[build-system] +requires = ["setuptools >= 77.0.3"] +build-backend = "setuptools.build_meta" + +[project] +name = "team1k" +version = "0.0.1" +authors = [ + { name="Sebastian Strempfer", email="sstrempfer@anl.gov" }, +] +description = "Controls for the TEAM1k detector" +readme = "README.md" +requires-python = ">=3.9" +dependencies = [ + "numpy >= 1.24.0", + "pyserial >= 3.5", + "pydantic >= 2.0.0", + "typing-extensions >= 4.5.0", +] + +[project.scripts] +team1k-server = "team1k.server:main" + +[tool.setuptools.packages.find] +where = ["src"] +include = ["team1k*"] +exclude = ["tests*"] + +[tool.setuptools] +packages = ["team1k"] \ No newline at end of file diff --git a/src/team1k/Client.py b/src/team1k/Client.py new file mode 100644 index 0000000..d4c6916 --- /dev/null +++ b/src/team1k/Client.py @@ -0,0 +1,64 @@ +import socket +import time +from KTestClient import KTestClient, KTestCommand, KTestError, ExposureModes, TriggerModes + +class Client: + def __init__(self, host: str = 'localhost', port: int = 42003): + self.host = host + self.port = port + self.k_test_client = KTestClient(host, port) + + def send_command(self, command: KTestCommand, *args) -> str: + """ + Send a command to the KTestServer and return the response. + + Args: + command (KTestCommand): Command enum to send + *args: Arguments for the command + Returns: + str: Response message + """ + return self.k_test_client.send_command(command, *args) + + def reset_connection(self) -> None: + """ + Reset the connection between the k_test server and the detector. + """ + response = self.send_command(KTestCommand.RESET_CONNECTION) + time.sleep(1) # Wait for the connection to stabilize + + def start_daq(self) -> None: + """ + Start the data acquisition process. + """ + self.send_command(KTestCommand.RESTART_DAQ) + + def stop_daq(self) -> None: + """ + Stop the data acquisition process. + """ + self.send_command(KTestCommand.STOP_DAQ) + + def set_exposure_mode(self, exposure_mode: ExposureModes) -> None: + """ + Set the exposure mode for the detector. + """ + self.send_command(KTestCommand.SET_EXPOSURE_MODE, exposure_mode) + + def set_trigger_mode(self, trigger_mode: TriggerModes) -> None: + """ + Set the trigger mode for the detector. + """ + self.send_command(KTestCommand.SET_TRIGGER_MODE, trigger_mode) + + def set_integration_time(self, integration_time_ms: float) -> None: + """ + Set the integration time in milliseconds. + """ + self.send_command(KTestCommand.SET_INTEGRATION_TIME, integration_time_ms) + + def load_parameter_file(self, file_path: str) -> None: + """ + Load a parameter file for the detector. + """ + self.send_command(KTestCommand.LOAD_PARAMETER_FILE, file_path) \ No newline at end of file diff --git a/src/team1k/KTestClient.py b/src/team1k/KTestClient.py new file mode 100644 index 0000000..7d8a384 --- /dev/null +++ b/src/team1k/KTestClient.py @@ -0,0 +1,191 @@ +import socket +import time +import enum + +class KTestError(Exception): + """Custom exception for KTestClient errors.""" + pass + +class ExposureModes(enum.IntEnum): + """Exposure modes for the detector.""" + + ROLLING_SHUTTER = 0 + """Rolling shutter mode.""" + + ROLLING_SHUTTER_WITH_PAUSE = 1 + """Rolling shutter with pause mode.""" + + GLOBAL_SHUTTER = 2 + """Global shutter mode.""" + + GLOBAL_SHUTTER_CDS = 3 + """Global shutter with Correlated Double Sampling (CDS) using two images.""" + +class TriggerModes(enum.IntEnum): + """Trigger modes for the detector (more implemented but not supported by k_test server).""" + + INTERNAL_TRIGGER = 0 + """Internal trigger mode.""" + + EXTERNAL_TRIGGER = 1 + """External trigger mode.""" + +class KTestCommand(enum.Enum): + RESET_CONNECTION = "resetdetectorconnection" + """Reset the connection from k_test to the detector.""" + + # PRINT_EVERY_NTH = "printeverynth" + # """Set print frequency (requires value)""" + + TEST_MODE = "testmode" + """Enable/disable test mode (0 or 1)""" + + LOAD_PARAMETER_FILE = "parameterfile" + """Load parameter file (requires filename)""" + + READ_REGISTER = "readregister" + """Read register (requires address)""" + + WRITE_REGISTER = "writeregister" + """Write register (requires address and value)""" + + SET_ADC_CLOCK_FREQ = "setadcclockfreq" + """Set ADC clock frequency in MHz (50-100)""" + + SET_OUTPUT_DIR = "setoutputdir" + """Set output directory (requires path)""" + + ENABLE_FILE_WRITING = "enablefilewriting" + """Enable file writing""" + + DISABLE_FILE_WRITING = "disablefilewriting" + """Disable file writing""" + + SET_EXPOSURE_MODE = "setexposuremode" + """Set exposure mode (0-3)""" + + SET_TRIGGER_MODE = "settriggermode" + """Set trigger mode (0-2)""" + + SET_INTEGRATION_TIME = "setintegrationtime" + """Set integration time in ms""" + + SET_ADC_DATA_DELAY = "setadcdatadelay" + """Set ADC data delay""" + + RESTART_DAQ = "restartdaq" + """Restart DAQ""" + + STOP_DAQ = "stopdaq" + """Stop DAQ""" + +KTEST_COMMAND_ARGS = { + KTestCommand.RESET_CONNECTION: (), + # KTestCommand.PRINT_EVERY_NTH: (int,), + KTestCommand.TEST_MODE: (int,), + KTestCommand.LOAD_PARAMETER_FILE: (str,), + KTestCommand.READ_REGISTER: (int,), + KTestCommand.WRITE_REGISTER: (int, int), + KTestCommand.SET_ADC_CLOCK_FREQ: (float,), + KTestCommand.SET_OUTPUT_DIR: (str,), + KTestCommand.ENABLE_FILE_WRITING: (), + KTestCommand.DISABLE_FILE_WRITING: (), + KTestCommand.SET_EXPOSURE_MODE: (ExposureModes,), + KTestCommand.SET_TRIGGER_MODE: (TriggerModes,), + KTestCommand.SET_INTEGRATION_TIME: (float,), + KTestCommand.SET_ADC_DATA_DELAY: (int,), + KTestCommand.RESTART_DAQ: (), + KTestCommand.STOP_DAQ: (), +} + +class KTestClient: + def __init__(self, host='localhost', port=42003): + self.host = host + self.port = port + + def send_command(self, command: KTestCommand, *args) -> str: + """ + Send a command to the KTestServer and return the response. + + Args: + command (KTestCommand): Command enum to send + *args: Arguments for the command + Returns: + str: response_message + """ + if command not in KTEST_COMMAND_ARGS: + raise ValueError(f"Invalid command: {command}") + + expected_args = KTEST_COMMAND_ARGS[command] + if len(args) != len(expected_args): + raise ValueError(f"Invalid number of arguments for {command.value}. Expected {len(expected_args)}, got {len(args)}.") + + # Validate argument types + for i, (arg, expected_type) in enumerate(zip(args, expected_args)): + if not isinstance(arg, expected_type): + raise TypeError(f"Invalid argument type for argument {i+1} of {command.value}. Expected {expected_type.__name__}, got {type(arg).__name__}.") + + # Construct command string + command_str = command.value + if args: + command_str += ' ' + ' '.join(str(arg) for arg in args) + + ret_code, ret_msg = self.send_command_str(command_str) + if ret_code != 0: + raise KTestError(f"Command '{command_str}' failed with code {ret_code}: {ret_msg}") + return ret_msg + + def send_command_str(self, command) -> tuple[int, str]: + """ + Send a command to the KTestServer and return the response. + + Args: + command (str): Command string to send + + Returns: + tuple: (return_code, response_message) + """ + sock = None + try: + # Create TCP socket + sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + sock.connect((self.host, self.port)) + + # Send command + sock.send(command.encode('utf-8')) + + # Receive response + response = b"" + while True: + chunk = sock.recv(4096) + if not chunk: + break + response += chunk + + sock.close() + + # Parse response - return code is at the beginning + response_str = response.decode('utf-8') + lines = response_str.strip().split('\n') + + # Extract return code (first number in response) + return_code = None + for line in lines: + if line.strip().isdigit() or (line.strip().startswith('-') and line.strip()[1:].isdigit()): + return_code = int(line.strip()) + break + + if return_code is None: + raise KTestError(f"Invalid response from server: {response_str}") + + return return_code, response_str + except ConnectionRefusedError: + raise ConnectionRefusedError("Connection error: Server is not running or unreachable. Start the k_test server first.") + except Exception as e: + raise e + finally: + try: + if sock is not None: + sock.close() + except: + pass diff --git a/src/team1k/__init__.py b/src/team1k/__init__.py new file mode 100644 index 0000000..6b0fbbd --- /dev/null +++ b/src/team1k/__init__.py @@ -0,0 +1,2 @@ +from .Client import Client, KTestCommand, KTestError, ExposureModes, TriggerModes +__all__ = ['Client', 'KTestCommand', 'KTestError', 'ExposureModes', 'TriggerModes'] \ No newline at end of file diff --git a/src/team1k/server.py b/src/team1k/server.py new file mode 100644 index 0000000..ab67843 --- /dev/null +++ b/src/team1k/server.py @@ -0,0 +1,167 @@ +import os +import socket +import zmq +import zmq.asyncio +import asyncio +import serial +from .KTestClient import KTestClient, KTestCommand, KTestError, KTEST_COMMAND_ARGS +import psutil +import subprocess +import logging +import enum + +SW_PATH = "/home/daq-user/Team1k_ANL_UED_Camera/sw" +K_TEST_CONFIG_PATH = "/home/daq-user/Team1k_ANL_UED_Camera/configuration/parameterfile_Team1k_UED_ANL.txt" + +class Commands(enum.Enum): + INSERT = "insert" + """Insert the detector.""" + + RETRACT = "retract" + """Retract the detector.""" + + POWER_SUPPLY = "power_supply" + """Power the detector on or off.""" + + KILL = "kill" + """Kill this server.""" + +server_command_arg_types = { + Commands.INSERT: (), + Commands.RETRACT: (), + Commands.POWER_SUPPLY: (bool,), + Commands.KILL: (), +} + +command_arg_types = {**server_command_arg_types, **KTEST_COMMAND_ARGS} + +class Team1kServer: + def __init__(self, host="*", port=42003): + self.host = host + self.port = port + self.context = zmq.asyncio.Context() + self.socket: zmq.asyncio.Socket = self.context.socket(zmq.REP) # REP socket + self.socket.bind(f"tcp://{self.host}:{self.port}") + self.loop = asyncio.get_event_loop() + self._k_test_process = None + self.logger = logging.getLogger("Team1kServer") + self.logger.setLevel(logging.DEBUG) + + async def handle_request(self, request): + # Process the incoming request + command = request.strip().lower().split() + if not command or command[0] == "": + return ValueError("Empty command") + + try: + command[0] = Commands(command[0]) if command[0] in Commands._value2member_map_ else KTestCommand(command[0].upper()) + except ValueError: + return ValueError(f"Unknown command: {command[0]}") + + arg_types = command_arg_types.get(command, None) + if arg_types is None: + return ValueError(f"Unknown command: {command[0]}") + if len(arg_types) != len(command) - 1: + return ValueError(f"Expected {len(arg_types)} arguments ({arg_types}), got {len(command) - 1}") + for i, arg_type in enumerate(arg_types): + try: + command[i + 1] = arg_type(command[i + 1]) + except ValueError: + return ValueError(f"Argument {i + 1} should be of type {arg_type.__name__}, got {command[i + 1]}") + + match command: + case [Commands.INSERT]: + response = await self.move_camera(insert=True) + case [Commands.RETRACT]: + response = await self.move_camera(insert=False) + case [Commands.KILL]: + response = True + await self.shutdown() + case [Commands.POWER_SUPPLY, state]: + # Turn power supply on or off + pass + case [cmd, *args] if isinstance(cmd, KTestCommand): + try: + response = self.k_client.send_command(cmd, *args) + except KTestError as e: + response = e + except Exception as e: + response = RuntimeError(f"Unexpected error: {e}") + case _: + response = ValueError(f"Unhandled command: {command}") + + return response + + @property + def k_client(self) -> KTestClient: + if not hasattr(self, '_k_client'): + # See if k_test server is running on localhost + if not any("k_test" in p.info['name'] for p in psutil.process_iter(['name'])): + # If not, start it + self.logger.info("Starting k_test server on localhost") + k_test_cmd = ["k_test", K_TEST_CONFIG_PATH] + self.logger.debug(f"Running command: {' '.join(k_test_cmd)} in {SW_PATH}") + self._k_test_process = subprocess.Popen(k_test_cmd, cwd=SW_PATH) + + self._k_client = KTestClient("localhost") + return self._k_client + + async def move_camera(self, insert: bool): + """Move the camera bellow stage in or out.""" + self.logger.info("Inserting detector") + try: + with serial.Serial('/dev/CameraBellowStage', 9600, timeout=1) as ser: + await asyncio.sleep(1) + ser.write(b'\r') + await asyncio.sleep(1) + ser.write(b'\x03') # Send Ctrl-C to stop any ongoing movement + await asyncio.sleep(1) + + ser.write(b'rc=100\r') # Set speed to 100 + await asyncio.sleep(1) + + if insert: + ser.write(b'vm=100000\r') # Set maximum speed (we're being helped by atmospheric pressure) + await asyncio.sleep(1) + + ser.write(b'mr 2000000\r') # Move 2000000 (insert) TODO: Convert to absolute position + await asyncio.sleep(10) + else: + ser.write(b'vm=100001\r') # Set maximum speed (not sure where this value comes from) + await asyncio.sleep(1) + + ser.write(b'mr -2000000\r') # Move -2000000 (retract) TODO: Convert to absolute position + await asyncio.sleep(10) + + ser.write(b'pr p\r') # Print current position + await asyncio.sleep(1) + except serial.SerialTimeoutException as e: + self.logger.error(f"Serial timeout: {e}") + return TimeoutError(f"Server <--> Camera Bellow Stage serial timeout: {e}") + except serial.SerialException as e: + self.logger.error(f"Serial error: {e}") + return ConnectionError(f"Server <--> Camera Bellow Stage serial error: {e}") + except Exception as e: + self.logger.error(f"Unexpected error: {e}") + return RuntimeError(f"Server <--> Camera Bellow Stage unexpected error: {e}") + return True + + async def process_request(self, request): + self.logger.info(f"Received request: {request}") + response = await self.handle_request(request) + await self.socket.send_pyobj(response) + self.logger.info(f"Sent response: {response}") + + async def serve(self): + while True: + # Wait for a request + request = await self.socket.recv_string() + await self.process_request(request) + + async def shutdown(self): + if self._k_test_process: + self._k_test_process.terminate() + self._k_test_process.wait() + self.socket.close() + self.context.term() + os._exit(0)