"""ThinkRF integration for the RIA toolkit.""" from typing import Any, Dict, Optional import numpy as np try: from pyrf.devices.thinkrf import WSA except ImportError as exc: # pragma: no cover - optional dependency raise ImportError( "pyrf is required to use the ThinkRF integration. " "Install with: pip install ria-toolkit-oss[thinkrf]" ) from exc except SyntaxError as exc: # pragma: no cover - Python 2/3 compatibility issue import sys from pathlib import Path # pyrf ships with Python 2 syntax - try to auto-fix it print("\033[93mWARNING: pyrf has Python 2 syntax. Attempting automatic fix...\033[0m") try: from lib2to3.refactor import RefactoringTool, get_fixers_from_package import pyrf thinkrf_path = Path(pyrf.__file__).resolve().parent / "devices" / "thinkrf.py" print(f"Fixing: {thinkrf_path}") fixers = get_fixers_from_package("lib2to3.fixes") tool = RefactoringTool(fixers) tool.refactor_file(str(thinkrf_path), write=True) print("\033[92m✅ Fixed pyrf for Python 3. Please restart Python/reload the module.\033[0m") print("Or run: python -m ria_toolkit_oss.sdr.thinkrf_fix") sys.exit(1) # Exit so user can reload except Exception as fix_exc: print(f"\033[91m❌ Auto-fix failed: {fix_exc}\033[0m") print("Manual fix: Run `python scripts/fix_pyrf_python3.py` from ria-toolkit-oss directory") raise exc from ria_toolkit_oss.sdr.sdr import SDR, SDRParameterError class ThinkRF(SDR): """SDR adapter for ThinkRF analyzers using the PyRF API.""" BASE_SAMPLE_RATE = 125_000_000 SUPPORTED_DECIMATIONS = (1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024) MAX_ONBOARD_SAMPLES = 33_500_000 # Confirmed: 512 packets @ dec 1 = 33.5M samples (268ms) DEFAULT_SPP = 65504 # VRT packet size (samples per packet) def __init__(self, identifier: Optional[str] = None): super().__init__() if identifier is None: raise SDRParameterError("ThinkRF requires an IP address or hostname identifier") self.identifier = identifier try: self.radio = WSA() self.radio.connect(identifier) self.radio.request_read_perm() print(f"Connected to ThinkRF at [{identifier}].") except Exception as exc: print(f"Failed to connect to ThinkRF at [{identifier}].") raise exc self.configure_frontend() self._last_context: Optional[Any] = None def configure_frontend( self, *, rfe_mode: str = "ZIF", attenuation: int = 0, gain_profile: str = "HIGH", trigger_config: Optional[Dict[str, Any]] = None, samples_per_packet: int = 65504, packets_per_block: int = 1, capture_mode: str = "block", stream_id: int = 1, min_stream_decimation: int = 16, ) -> None: """Persist settings applied during the next RX initialisation. ``capture_mode`` selects between buffered ``"block"`` captures that use the analyser's onboard RAM and ``"stream"`` captures that push data over GigE in real time. Streaming requires a sufficiently large decimation to keep within the link budget; ``min_stream_decimation`` forms the lower bound. """ mode = capture_mode.lower() if mode not in {"block", "stream"}: raise SDRParameterError("capture_mode must be either 'block' or 'stream'") self._rfe_mode = rfe_mode self._attenuation = int(max(0, min(attenuation, 30))) self._gain_profile = gain_profile.upper() self._trigger_config = trigger_config self._samples_per_packet = int(samples_per_packet) self._packets_per_block = max(1, int(packets_per_block)) self._capture_mode = mode self._stream_id = int(stream_id) self._min_stream_decimation = max(1, int(min_stream_decimation)) self._streaming_active = False def init_rx( self, sample_rate: int | float, center_frequency: int | float, gain: int, channel: int, gain_mode: Optional[str] = "absolute", decimation: Optional[int] = None, ): if channel not in (0, None): raise SDRParameterError("ThinkRF supports only channel 0 for RX.") stream_mode = getattr(self, "_capture_mode", "block") == "stream" actual_decimation, _ = self.set_rx_sample_rate( sample_rate=sample_rate, decimation=decimation, stream_mode=stream_mode ) self.radio.reset() self.radio.scpiset(":SYSTEM:FLUSH") try: self.radio.scpiset(":TRACE:STREAM:STOP") except Exception: pass self.radio.rfe_mode(self._rfe_mode) self.set_rx_center_frequency(center_frequency=center_frequency) self.set_rx_gain(gain=gain, gain_mode=gain_mode, actual_decimation=actual_decimation) self.radio.decimation(actual_decimation) if stream_mode: self.radio.scpiset(f":SENSE:DECIMATION {actual_decimation}") trigger = self._trigger_config or self._default_trigger(center_frequency) self.radio.trigger(trigger) self.radio.scpiset(f":TRACE:SPP {self._samples_per_packet}") if stream_mode: self._streaming_active = False else: print( f"ThinkRF: Configuring block capture - SPP={self._samples_per_packet}, PPB={self._packets_per_block}" ) self.radio.scpiset(f":TRACE:BLOCK:PACKETS {self._packets_per_block}") self.radio.scpiset(":TRACE:BLOCK:DATA?") self.rx_buffer_size = self._samples_per_packet self.rx_channel = 0 self._rx_initialized = True self._tx_initialized = False def set_rx_sample_rate(self, sample_rate, decimation, stream_mode): """ Set the sample rate of the receiver. Not callable during recording; ThinkRF requires stream stop/restart to change sample rate. """ # Enforce sample rate / decimation # Note: decimation parameter takes precedence if provided actual_decimation, actual_sample_rate = self.enforce_sample_rate(sample_rate, decimation) if stream_mode and actual_decimation < self._min_stream_decimation: enforced = self._min_stream_decimation print( "Requested ThinkRF sample rate exceeds typical GigE throughput; " f"enforcing decimation {enforced} for streaming." ) actual_decimation = enforced actual_sample_rate = self.BASE_SAMPLE_RATE / actual_decimation self._decimation = actual_decimation self.rx_sample_rate = actual_sample_rate print(f"ThinkRF RX Sample Rate = {actual_sample_rate}") return actual_decimation, actual_sample_rate def set_rx_center_frequency(self, center_frequency): """ Set the center frequency of the receiver. Callable during streaming. """ with self._param_lock: self.radio.freq(int(center_frequency)) self.rx_center_frequency = self.radio.freq print(f"ThinkRF RX Center Frequency = {self.radio.freq}") def set_rx_gain(self, gain, gain_mode, actual_decimation): attenuation = self._attenuation if gain is None else int(gain) # gain attenuation = max(0, min(attenuation, 30)) self.radio.attenuator(attenuation) gain_profile = self._gain_profile if gain_mode and isinstance(gain_mode, str) and gain_mode.upper() in {"LOW", "MEDIUM", "HIGH", "VLOW"}: gain_profile = gain_mode.upper() self.radio.gain(gain_profile.lower()) # WSA.gain() expects lowercase self.rx_gain = { "attenuation_dB": attenuation, "profile": gain_profile, "decimation": actual_decimation, "rfe_mode": self._rfe_mode, "spp": self._samples_per_packet, "ppb": self._packets_per_block, } def _stream_rx(self, callback): if not self._rx_initialized: raise RuntimeError("RX was not initialized. init_rx() must be called before _stream_rx() or record().") self._enable_rx = True packets_processed = 0 stream_mode = getattr(self, "_capture_mode", "block") == "stream" if stream_mode and not self._streaming_active: try: self.radio.scpiset(f":TRACE:STREAM:START {self._stream_id}") self._streaming_active = True except Exception as exc: print(f"Failed to start ThinkRF stream: {exc}") return print("ThinkRF Starting RX...") while self._enable_rx: packet = self._safe_read(stream_mode, packets_processed) if packet is None: # No more packets available if not stream_mode and packets_processed >= self._packets_per_block: # Finished reading block break continue if packet.is_context_packet(): self._last_context = packet continue if not packet.is_data_packet(): # Unknown packet type - skip continue metadata = metadata = self._extract_metadata(packet) complex_buffer = self._extract_iq(packet) if complex_buffer is None: continue # Send packet data to callback (accumulation handled by parent) callback(buffer=complex_buffer, metadata=metadata) packets_processed += 1 # In block mode, stop after receiving all packets in the block if not stream_mode and packets_processed >= self._packets_per_block: # Got all packets for this block break print("ThinkRF RX Completed.") if stream_mode and self._streaming_active: self._stop_stream() self.radio.scpiset(":SYSTEM:FLUSH") def _safe_read(self, stream_mode, packets_processed): packet = None try: packet = self.radio.read() except Exception as e: # In block mode, reaching end of block can cause exceptions if not stream_mode and packets_processed > 0: # We got some packets in block mode, so finish gracefully print(f"ThinkRF: Block read complete ({packets_processed} packets received)") else: print(f"ThinkRF read error: {e}") return packet def _extract_iq(self, packet): # packet.data is an iterable IQData object that yields (I, Q) tuples # Convert to numpy array: collect all [I, Q] pairs try: iq_pairs = list(packet.data) if not iq_pairs: return None iq_array = np.array(iq_pairs, dtype=np.float32) return (iq_array[:, 0] + 1j * iq_array[:, 1]).astype(np.complex64) except Exception as e: print(f"Error extracting IQ from packet.data: {e}") return None def _extract_metadata(self, packet): if not hasattr(packet, "fields"): return None metadata = packet.fields if metadata.get("sample_loss"): print("\033[93mWarning: ThinkRF sample overflow detected\033[0m") return metadata def _stop_stream(self): try: self.radio.scpiset(":TRACE:STREAM:STOP") except Exception: pass self._streaming_active = False def init_tx( self, sample_rate: int | float, center_frequency: int | float, gain: int, channel: int, gain_mode: Optional[str] = "absolute", ): raise NotImplementedError("ThinkRF devices do not support transmit operations") def _stream_tx(self, callback): raise NotImplementedError("ThinkRF devices do not support transmit operations") def set_clock_source(self, source): raise NotImplementedError("ThinkRF clock configuration is not implemented") def close(self): try: self.radio.scpiset(":TRACE:STREAM:STOP") except Exception: # pragma: no cover - best effort cleanup pass try: self.radio.scpiset(":SYSTEM:FLUSH") except Exception: pass try: self.radio.disconnect() finally: self._enable_rx = False self._enable_tx = False print(f"Disconnected from ThinkRF at [{self.identifier}].") def supports_bias_tee(self) -> bool: return False def set_bias_tee(self, enable: bool): # pragma: no cover - interface compliance raise NotImplementedError("ThinkRF radios do not expose a controllable bias-tee") def _derive_decimation(self, target_sample_rate: int | float) -> int: """ Derive decimation from target sample rate. Always rounds DOWN decimation (UP sample rate) to meet or exceed user's requested rate. Example: 30 MS/s requested → dec 4 (31.25 MS/s), NOT dec 8 (15.625 MS/s) """ if not target_sample_rate: return 1 requested = float(target_sample_rate) if requested >= self.BASE_SAMPLE_RATE: return 1 desired_decimation = self.BASE_SAMPLE_RATE / requested # Round DOWN decimation (UP sample rate) to meet or exceed requested rate # Find largest decimation that gives sample rate >= requested valid_decimations = [d for d in self.SUPPORTED_DECIMATIONS if d <= desired_decimation] if valid_decimations: # Use largest valid decimation (gives sample rate >= requested) best = max(valid_decimations) else: # Requested rate too low, use minimum decimation (max sample rate) best = self.SUPPORTED_DECIMATIONS[0] return int(best) def enforce_sample_rate( self, requested_sample_rate: int | float, decimation: Optional[int] = None ) -> tuple[int, float]: """ Enforce valid sample rate and decimation. If decimation is provided, it takes precedence. Otherwise, derive decimation from requested sample rate. Returns: (decimation, actual_sample_rate) """ if decimation is not None: # Decimation provided - validate and use it if decimation not in self.SUPPORTED_DECIMATIONS: # Round to nearest supported decimation = min(self.SUPPORTED_DECIMATIONS, key=lambda d: abs(d - decimation)) print(f"ThinkRF: Requested decimation not supported. Using decimation={decimation}") else: # Derive from sample rate decimation = self._derive_decimation(requested_sample_rate) actual_sample_rate = self.BASE_SAMPLE_RATE / decimation if abs(actual_sample_rate - requested_sample_rate) > 1e3: # More than 1 kHz difference print(f"ThinkRF: Requested {requested_sample_rate/1e6:.2f} MS/s → \ Using decimation={decimation} ({actual_sample_rate/1e6:.2f} MS/s)") return decimation, actual_sample_rate def calculate_spp_ppb(self, num_samples: int, spp: Optional[int] = None) -> tuple[int, int]: """ Calculate optimal SPP (samples per packet) and PPB (packets per block). Strategy: - Maximize SPP (use DEFAULT_SPP) unless num_samples < DEFAULT_SPP - Calculate PPB to get as close as possible to num_samples - Actual captured samples = SPP * PPB (may exceed num_samples slightly) Args: num_samples: Desired number of samples spp: Override SPP (for advanced users, not recommended) Returns: (spp, ppb) """ if spp is not None: # User override - use as-is actual_spp = max(1, int(spp)) else: # Maximize SPP unless samples requested is smaller if num_samples < self.DEFAULT_SPP: actual_spp = num_samples else: actual_spp = self.DEFAULT_SPP # Calculate PPB to get close to num_samples ppb = max(1, int(np.ceil(num_samples / actual_spp))) actual_samples = actual_spp * ppb if actual_samples != num_samples: print( f"ThinkRF: Requested {num_samples} samples → Capturing {actual_samples} (SPP={actual_spp}, PPB={ppb})" ) return actual_spp, ppb def check_ram_limit(self, num_samples: int, decimation: int) -> None: """ Check if requested capture exceeds onboard RAM limits. Raises warning if exceeds MAX_ONBOARD_SAMPLES at low decimations. For decimation 1 or 2, block captures are limited by onboard RAM. """ if decimation <= 2 and num_samples > self.MAX_ONBOARD_SAMPLES: raise SDRParameterError( f"ThinkRF: Cannot capture {num_samples} samples at decimation {decimation}. " f"Onboard RAM limit is ~{self.MAX_ONBOARD_SAMPLES} samples for dec 1/2. " f"Either reduce num_samples or use stream mode (increase decimation to >=4)." ) def _default_trigger(self, center_frequency: int | float) -> Dict[str, Any]: span = 40_000_000 half = span // 2 return { "type": "NONE", "fstart": int(center_frequency) - half, "fstop": int(center_frequency) + half, "amplitude": -100, } def supports_dynamic_updates(self) -> dict: return {"center_frequency": True, "sample_rate": False, "gain": False}