PicoScope MCP Server

"""Device manager for PicoScope oscilloscopes.""" from typing import Optional, Any import ctypes import numpy as np # Try to import PicoSDK, but handle gracefully if not installed try: from picosdk.discover import find_all_units from picosdk.ps5000a import ps5000a as ps from picosdk.functions import adc2mV, assert_pico_ok, mV2adc PICOSDK_AVAILABLE = True except Exception as e: PICOSDK_AVAILABLE = False print(f"Warning: PicoSDK not available: {e}") print("Server will run but device operations will fail until PicoSDK is installed.") from .models import DeviceInfo, ChannelConfig, TriggerConfig, CaptureData class PicoScopeManager: """Manager for PicoScope device operations.""" def __init__(self): """Initialize the device manager.""" self.current_device: Optional[Any] = None self.device_info: Optional[DeviceInfo] = None self.channel_configs: dict[str, ChannelConfig] = {} self.chandle: Optional[ctypes.c_int16] = None self.status: dict[str, Any] = {} def discover_devices(self) -> list[dict[str, Any]]: """Discover all connected PicoScope devices. Returns: List of discovered devices with their information. """ if not PICOSDK_AVAILABLE: return [] try: devices = find_all_units() device_list = [] for scope in devices: device_list.append( { "info": str(scope.info), "variant": getattr(scope, "variant", "unknown"), } ) scope.close() return device_list except Exception as e: return [] def connect(self, device_index: int = 0, serial: Optional[str] = None) -> bool: """Connect to a PicoScope 5000A device. Args: device_index: Index of device to connect to (0 for first device). serial: Optional serial number string for specific device. Returns: True if connection successful, False otherwise. """ if not PICOSDK_AVAILABLE: return False try: # Create handle for PS5000A self.chandle = ctypes.c_int16() # Set resolution (12-bit default) resolution = ps.PS5000A_DEVICE_RESOLUTION["PS5000A_DR_12BIT"] # Convert serial to bytes if provided serial_bytes = serial.encode() if serial else None # Open the device self.status["openunit"] = ps.ps5000aOpenUnit( ctypes.byref(self.chandle), serial_bytes, resolution ) # Handle power status issues try: assert_pico_ok(self.status["openunit"]) except: power_status = self.status["openunit"] # Power status codes: 286 = USB3.0 powered, 282 = needs external power if power_status in [286, 282]: self.status["changePowerSource"] = ps.ps5000aChangePowerSource( self.chandle, power_status ) assert_pico_ok(self.status["changePowerSource"]) else: raise # Get device info info_string = ctypes.create_string_buffer(256) required_size = ctypes.c_int16(256) # Get variant info self.status["getUnitInfo"] = ps.ps5000aGetUnitInfo( self.chandle, ctypes.byref(info_string), ctypes.c_int16(256), ctypes.byref(required_size), ps.PICO_INFO["PICO_VARIANT_INFO"] ) variant = info_string.value.decode('utf-8') # Get batch and serial self.status["getUnitInfo"] = ps.ps5000aGetUnitInfo( self.chandle, ctypes.byref(info_string), ctypes.c_int16(256), ctypes.byref(required_size), ps.PICO_INFO["PICO_BATCH_AND_SERIAL"] ) batch_serial = info_string.value.decode('utf-8') # Store device information self.device_info = DeviceInfo( handle=self.chandle.value, model="PS5000A", serial=batch_serial, variant=variant, batch_and_serial=batch_serial, max_adc_value=32767, # 15-bit for PS5000A min_adc_value=-32767, num_channels=4, # PS5000A has 4 channels ) return True except Exception as e: self.chandle = None self.device_info = None return False def disconnect(self) -> bool: """Disconnect from current device. Returns: True if disconnection successful, False otherwise. """ try: if self.chandle: # Stop any running capture ps.ps5000aStop(self.chandle) # Close the unit ps.ps5000aCloseUnit(self.chandle) self.chandle = None self.current_device = None self.device_info = None self.channel_configs.clear() self.status.clear() return True except Exception as e: return False def is_connected(self) -> bool: """Check if a device is currently connected. Returns: True if connected, False otherwise. """ return self.current_device is not None def get_info(self) -> Optional[DeviceInfo]: """Get information about the current device. Returns: DeviceInfo object or None if not connected. """ return self.device_info def configure_channel(self, config: ChannelConfig) -> bool: """Configure a channel. Args: config: Channel configuration. Returns: True if successful, False otherwise. """ if not self.is_connected(): return False try: # Map channel letter to PS5000A channel constant channel_map = { "A": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_A"], "B": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_B"], "C": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_C"], "D": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_D"], } if config.channel not in channel_map: return False # Map coupling type coupling = ( ps.PS5000A_COUPLING["PS5000A_AC"] if config.coupling.value == "AC" else ps.PS5000A_COUPLING["PS5000A_DC"] ) # Map voltage range to closest available range range_map = { 0.02: "PS5000A_20MV", 0.05: "PS5000A_50MV", 0.1: "PS5000A_100MV", 0.2: "PS5000A_200MV", 0.5: "PS5000A_500MV", 1.0: "PS5000A_1V", 2.0: "PS5000A_2V", 5.0: "PS5000A_5V", 10.0: "PS5000A_10V", 20.0: "PS5000A_20V", } # Find closest range closest_range = min(range_map.keys(), key=lambda x: abs(x - config.voltage_range)) voltage_range = ps.PS5000A_RANGE[range_map[closest_range]] # Convert analog offset to ADC counts if self.device_info: analog_offset_adc = mV2adc( config.analog_offset * 1000, # V to mV voltage_range, self.device_info.max_adc_value ) else: analog_offset_adc = 0 # Set the channel self.status[f"setCh{config.channel}"] = ps.ps5000aSetChannel( self.chandle, channel_map[config.channel], 1 if config.enabled else 0, coupling, voltage_range, analog_offset_adc ) assert_pico_ok(self.status[f"setCh{config.channel}"]) # Store configuration self.channel_configs[config.channel] = config return True except Exception as e: return False def set_trigger(self, config: TriggerConfig) -> bool: """Set up trigger. Args: config: Trigger configuration. Returns: True if successful, False otherwise. """ if not self.is_connected(): return False try: # Map source to channel source_map = { "A": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_A"], "B": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_B"], "C": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_C"], "D": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_D"], "External": ps.PS5000A_CHANNEL["PS5000A_EXTERNAL"], } if config.source not in source_map: return False # Map trigger direction direction_map = { "Rising": ps.PS5000A_THRESHOLD_DIRECTION["PS5000A_RISING"], "Falling": ps.PS5000A_THRESHOLD_DIRECTION["PS5000A_FALLING"], "Rising_Or_Falling": ps.PS5000A_THRESHOLD_DIRECTION[ "PS5000A_RISING_OR_FALLING" ], } direction = direction_map[config.direction.value] # Convert threshold from mV to ADC counts # Use the configured range for the source channel if available if config.source in self.channel_configs: ch_config = self.channel_configs[config.source] range_map = { 0.02: "PS5000A_20MV", 0.05: "PS5000A_50MV", 0.1: "PS5000A_100MV", 0.2: "PS5000A_200MV", 0.5: "PS5000A_500MV", 1.0: "PS5000A_1V", 2.0: "PS5000A_2V", 5.0: "PS5000A_5V", 10.0: "PS5000A_10V", 20.0: "PS5000A_20V", } closest_range = min( range_map.keys(), key=lambda x: abs(x - ch_config.voltage_range) ) voltage_range = ps.PS5000A_RANGE[range_map[closest_range]] else: voltage_range = ps.PS5000A_RANGE["PS5000A_2V"] # Default threshold_adc = mV2adc( config.threshold_mv, voltage_range, self.device_info.max_adc_value if self.device_info else 32767, ) # Set simple trigger self.status["trigger"] = ps.ps5000aSetSimpleTrigger( self.chandle, 1, # Enable trigger source_map[config.source], threshold_adc, direction, 0, # Delay (samples) config.auto_trigger_ms, # Auto-trigger timeout ) assert_pico_ok(self.status["trigger"]) return True except Exception as e: return False def capture_block( self, pre_trigger: int, post_trigger: int ) -> Optional[dict[str, CaptureData]]: """Capture a block of data. Args: pre_trigger: Number of samples before trigger. post_trigger: Number of samples after trigger. Returns: Dictionary of channel data or None if failed. """ if not self.is_connected(): return None try: total_samples = pre_trigger + post_trigger # Set up buffers for enabled channels buffers = {} channel_map = { "A": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_A"], "B": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_B"], "C": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_C"], "D": ps.PS5000A_CHANNEL["PS5000A_CHANNEL_D"], } for ch_name, ch_config in self.channel_configs.items(): if ch_config.enabled: # Create buffer buffer = (ctypes.c_int16 * total_samples)() buffers[ch_name] = buffer # Set data buffer self.status[f"setDataBuffer{ch_name}"] = ps.ps5000aSetDataBuffer( self.chandle, channel_map[ch_name], ctypes.byref(buffer), total_samples, 0, # Segment index ps.PS5000A_RATIO_MODE["PS5000A_RATIO_MODE_NONE"], ) assert_pico_ok(self.status[f"setDataBuffer{ch_name}"]) # Get timebase - using timebase 0 for fastest sampling timebase = 0 time_interval_ns = ctypes.c_float() max_samples = ctypes.c_int32() self.status["getTimebase"] = ps.ps5000aGetTimebase2( self.chandle, timebase, total_samples, ctypes.byref(time_interval_ns), ctypes.byref(max_samples), 0, # Segment index ) # If timebase 0 doesn't work, try higher values while self.status["getTimebase"] != 0 and timebase < 100: timebase += 1 self.status["getTimebase"] = ps.ps5000aGetTimebase2( self.chandle, timebase, total_samples, ctypes.byref(time_interval_ns), ctypes.byref(max_samples), 0, ) assert_pico_ok(self.status["getTimebase"]) # Run block capture self.status["runBlock"] = ps.ps5000aRunBlock( self.chandle, pre_trigger, post_trigger, timebase, None, # Time indisposed (not used) 0, # Segment index None, # Callback (not used, we'll poll) None, # Callback parameter ) assert_pico_ok(self.status["runBlock"]) # Wait for capture to complete ready = ctypes.c_int16(0) check = ctypes.c_int16(0) while ready.value == check.value: self.status["isReady"] = ps.ps5000aIsReady(self.chandle, ctypes.byref(ready)) # Get values overflow = ctypes.c_int16() c_total_samples = ctypes.c_int32(total_samples) self.status["getValues"] = ps.ps5000aGetValues( self.chandle, 0, # Start index ctypes.byref(c_total_samples), 1, # Downsampling ratio ps.PS5000A_RATIO_MODE["PS5000A_RATIO_MODE_NONE"], 0, # Segment index ctypes.byref(overflow), ) assert_pico_ok(self.status["getValues"]) # Convert ADC values to mV and create CaptureData objects result = {} time_values = np.arange(total_samples) * (time_interval_ns.value / 1e9) # Convert to seconds for ch_name, buffer in buffers.items(): ch_config = self.channel_configs[ch_name] # Get voltage range range_map = { 0.02: "PS5000A_20MV", 0.05: "PS5000A_50MV", 0.1: "PS5000A_100MV", 0.2: "PS5000A_200MV", 0.5: "PS5000A_500MV", 1.0: "PS5000A_1V", 2.0: "PS5000A_2V", 5.0: "PS5000A_5V", 10.0: "PS5000A_10V", 20.0: "PS5000A_20V", } closest_range = min(range_map.keys(), key=lambda x: abs(x - ch_config.voltage_range)) voltage_range = ps.PS5000A_RANGE[range_map[closest_range]] # Convert to numpy array adc_array = np.array(buffer) # Convert ADC to mV voltage_mv = adc2mV( adc_array, voltage_range, self.device_info.max_adc_value if self.device_info else 32767, ) result[ch_name] = CaptureData( channel=ch_name, time_values=time_values.tolist(), voltage_values=voltage_mv.tolist(), sample_interval_ns=int(time_interval_ns.value), num_samples=total_samples, ) return result except Exception as e: return None def start_streaming( self, sample_interval_ns: int, buffer_size: int ) -> bool: """Start streaming mode. Args: sample_interval_ns: Sample interval in nanoseconds. buffer_size: Buffer size for streaming. Returns: True if successful, False otherwise. """ if not self.is_connected(): return False # TODO: Implement streaming mode # This will need to set up buffers and call ps*RunStreaming return True def stop_streaming(self) -> bool: """Stop streaming mode. Returns: True if successful, False otherwise. """ if not self.is_connected(): return False # TODO: Implement streaming stop # Call ps*Stop return True # Global instance of the device manager device_manager = PicoScopeManager()