Lab051: accelerate 500 ksymbol/s streaming receiver
This commit is contained in:
LittleSam129
364
experiments/lab051_streaming_receiver.py
Normal file
364
experiments/lab051_streaming_receiver.py
Normal file
@@ -0,0 +1,364 @@
|
||||
"""Lab051: общий front-end и эквивалентный timing на 500 ксимв/с.
|
||||
|
||||
PHY, FEC, CRC и порог первого каскада заимствованы без изменения правил.
|
||||
Подмены старых функций действуют только внутри последовательного push.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from collections import Counter
|
||||
from contextlib import contextmanager
|
||||
from functools import lru_cache
|
||||
import math
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
from scipy.signal import convolve, fftconvolve, oaconvolve, firwin, resample_poly
|
||||
|
||||
from experiments import lab046b_continuous_iq as iq
|
||||
from experiments import lab046c_offline_recovery as offline
|
||||
from experiments import lab050_streaming_receiver as previous
|
||||
|
||||
|
||||
phy = previous.phy
|
||||
lab045 = phy.lab045
|
||||
lab043 = lab045.lab044.lab043
|
||||
ORIGINAL_REFINE = lab043._refine_frame_start
|
||||
ORIGINAL_TIMING = lab043._preamble_timing_offset
|
||||
ORIGINAL_MATCHED = lab043._matched_waveform
|
||||
ORIGINAL_VITERBI = lab045.conv.viterbi_decode_soft
|
||||
|
||||
|
||||
@lru_cache(maxsize=4)
|
||||
def timing_kernel(samples_per_symbol: int) -> np.ndarray:
|
||||
"""Те же символы преамбулы на исходной сетке отсчётов."""
|
||||
expected = iq.radio.bpsk_modulate(lab045.lab044.preamble_bits())
|
||||
kernel = np.zeros((expected.size - 1) * samples_per_symbol + 1, np.complex64)
|
||||
kernel[::samples_per_symbol] = np.conj(expected)
|
||||
return kernel[::-1].copy()
|
||||
|
||||
|
||||
def fast_preamble_timing(matched: np.ndarray, nominal: int, expected_symbols: np.ndarray) -> int:
|
||||
"""Все прежние 1601 смещения; отрицательные/усечённые окна имеют score=0.
|
||||
|
||||
Корреляция с разреженной преамбулой эквивалентна vdot каждой выборки.
|
||||
Энергии получаются разностями сумм на каждой из 16 фаз. Диапазон не сужен.
|
||||
"""
|
||||
samples_per_symbol = lab045.SAMPLES_PER_SYMBOL
|
||||
offsets = np.arange(-800, 801)
|
||||
starts = nominal + offsets
|
||||
count = len(expected_symbols)
|
||||
# Тело не участвует ни в одном из прежних 1601 timing-окон.
|
||||
matched = matched[:nominal + 800 + (count - 1) * samples_per_symbol + 1]
|
||||
valid = (starts >= 0) & (starts + (count - 1) * samples_per_symbol < matched.size)
|
||||
scores = np.zeros(offsets.size, np.float64)
|
||||
if np.any(valid):
|
||||
kernel = timing_kernel(samples_per_symbol)
|
||||
correlations = fftconvolve(matched, kernel, mode="valid")
|
||||
padding = (-matched.size) % samples_per_symbol
|
||||
power = np.pad(np.abs(matched) ** 2, (0, padding)).reshape(-1, samples_per_symbol)
|
||||
cumulative = np.concatenate((np.zeros((1, samples_per_symbol)),
|
||||
np.cumsum(power, axis=0, dtype=np.float64))).ravel()
|
||||
selected = starts[valid]
|
||||
energy = cumulative[selected + count * samples_per_symbol] - cumulative[selected]
|
||||
scores[valid] = np.abs(correlations[selected]) / np.sqrt(count * np.maximum(energy, 0) + 1e-30)
|
||||
largest = np.partition(scores, -2)[-2:]
|
||||
if largest[-1] > 0 and largest[-1] - largest[-2] <= 32 * np.finfo(np.float32).eps * largest[-1]:
|
||||
# FFT меняет последние разряды. На неразделимых пиках сохраняем
|
||||
# исходное суммирование и его tie→первый offset, включая чистый тон.
|
||||
return ORIGINAL_TIMING(matched, nominal, expected_symbols)
|
||||
return int(offsets[int(np.argmax(scores))])
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def carrier_rotation(length: int) -> np.ndarray:
|
||||
"""Кэш переносов не очищается при каждом входном блоке старым адаптером."""
|
||||
indexes = np.arange(length)
|
||||
return np.exp(-2j * np.pi * lab043.CARRIER_OFFSET_HZ * indexes / lab045.SAMPLE_RATE_HZ).astype(np.complex64)
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def dc_response(length: int, samples_per_symbol: int) -> np.ndarray:
|
||||
"""Отклик локального окна постоянной составляющей после переноса/RRC."""
|
||||
rotation = carrier_rotation(length)
|
||||
return oaconvolve(rotation, lab043._matched_filter_taps(), mode="full")
|
||||
|
||||
|
||||
@lru_cache(maxsize=4)
|
||||
def decimation_taps(decimation: int) -> np.ndarray:
|
||||
"""Штатные коэффициенты resample_poly(up=1, window=('kaiser', 8))."""
|
||||
return firwin(20 * decimation + 1, 1 / decimation, window=("kaiser", 8)).astype(np.float32)
|
||||
|
||||
|
||||
@lru_cache(maxsize=8)
|
||||
def detection_rotation(length: int, decimation: int) -> np.ndarray:
|
||||
return np.exp(-2j * np.pi * lab045.CARRIER_OFFSET_HZ * np.arange(length)
|
||||
/ (lab045.SAMPLE_RATE_HZ / decimation)).astype(np.complex64)
|
||||
|
||||
|
||||
class DecimationCache:
|
||||
"""Повторно использовать внутреннюю часть FIR, пересчитать только края.
|
||||
|
||||
resample_poly имеет поддержку ±10 выходных отсчётов. После сдвига окна
|
||||
неизменная внутренняя область даёт ровно те же значения; левый нулевой
|
||||
край и новое продолжение считаются с прежними коэффициентами заново.
|
||||
"""
|
||||
|
||||
def __init__(self, decimation: int):
|
||||
self.decimation = decimation
|
||||
self.base = self.input_count = 0
|
||||
self.values = np.empty(0, np.complex64)
|
||||
self.reused_samples = self.filtered_input_samples = 0
|
||||
|
||||
def get(self, samples: np.ndarray, base: int) -> np.ndarray:
|
||||
stride = self.decimation
|
||||
count = math.ceil(samples.size / stride)
|
||||
shift = (base - self.base) // stride
|
||||
low = max(10, -shift + 10)
|
||||
high = min(count - 10, self.values.size - 11 - shift)
|
||||
|
||||
def filtered(array):
|
||||
self.filtered_input_samples += array.size
|
||||
return resample_poly(array, 1, stride, window=decimation_taps(stride))
|
||||
|
||||
if not self.values.size or (base - self.base) % stride or high <= low:
|
||||
output = filtered(samples)
|
||||
else:
|
||||
output = np.empty(count, np.complex64)
|
||||
output[low:high] = self.values[low + shift:high + shift]
|
||||
self.reused_samples += high - low
|
||||
prefix_end = min(samples.size, (low + 11) * stride)
|
||||
output[:low] = filtered(samples[:prefix_end])[:low]
|
||||
suffix_start = (high - 11) * stride
|
||||
output[high:] = filtered(samples[suffix_start:])[11:11 + count - high]
|
||||
self.values = output
|
||||
self.base, self.input_count = base, samples.size
|
||||
return output
|
||||
|
||||
|
||||
def fast_viterbi(soft_values: np.ndarray, information_bit_count: int) -> np.ndarray:
|
||||
"""Та же евклидова метрика и tie→branch 0, с готовыми метриками ветвей."""
|
||||
conv = lab045.conv
|
||||
if not isinstance(information_bit_count, (int, np.integer)) or information_bit_count <= 0:
|
||||
return ORIGINAL_VITERBI(soft_values, information_bit_count)
|
||||
values = np.asarray(soft_values, dtype=np.float64)
|
||||
if (values.ndim != 1 or values.size != conv.encoded_bit_count(information_bit_count)
|
||||
or not np.all(np.isfinite(values))):
|
||||
return ORIGINAL_VITERBI(soft_values, information_bit_count)
|
||||
step_count = int(information_bit_count) + conv.TAIL_BIT_COUNT
|
||||
received = values.reshape(step_count, 2)
|
||||
# Всего четыре различных пары BPSK-символов. Порядок арифметики исходный.
|
||||
negative = (received + 1.0) ** 2
|
||||
positive = (received - 1.0) ** 2
|
||||
metrics = np.column_stack((negative[:, 0] + negative[:, 1], negative[:, 0] + positive[:, 1],
|
||||
positive[:, 0] + negative[:, 1], positive[:, 0] + positive[:, 1]))
|
||||
symbols = conv._BRANCH_SYMBOLS
|
||||
codes = (2 * (symbols[:, :, 0] > 0) + (symbols[:, :, 1] > 0)).astype(np.intp)
|
||||
metric_zero = metrics[:, codes[:, 0]]
|
||||
metric_one = metrics[:, codes[:, 1]]
|
||||
predecessor_zero = conv._PREDECESSORS[:, 0]
|
||||
predecessor_one = conv._PREDECESSORS[:, 1]
|
||||
path = np.full(conv.STATE_COUNT, 1e18)
|
||||
path[0] = 0.0
|
||||
choices = np.empty((step_count, conv.STATE_COUNT), np.uint8)
|
||||
for step in range(step_count):
|
||||
zero = path[predecessor_zero] + metric_zero[step]
|
||||
one = path[predecessor_one] + metric_one[step]
|
||||
choices[step] = one < zero
|
||||
path = np.minimum(zero, one)
|
||||
decoded = np.empty(step_count, np.uint8)
|
||||
state = 0
|
||||
for step in range(step_count - 1, -1, -1):
|
||||
decoded[step] = state >> (conv.MEMORY - 1)
|
||||
state = int(conv._PREDECESSORS[state, choices[step, state]])
|
||||
assert state == 0
|
||||
return decoded[:information_bit_count]
|
||||
|
||||
|
||||
class BlockFrontEnd:
|
||||
"""Один ограниченный общий участок; кандидаты читают его представления.
|
||||
|
||||
Конечные края и локальное вычитание среднего восстановлены отдельно:
|
||||
нельзя без этого подменять фильтрацию изолированного окна непрерывной.
|
||||
"""
|
||||
|
||||
def __init__(self, samples: np.ndarray, owned_end: int, statistics: Counter,
|
||||
*, decimator: DecimationCache | None = None, base: int = 0):
|
||||
self.samples = samples
|
||||
self.statistics = statistics
|
||||
self.decimator, self.base = decimator, base
|
||||
self.reference = lab045.lab044._preamble_reference_samples()
|
||||
self.reference_length = (16 + lab043.PILOT_SYMBOL_COUNT + 10) * lab045.SAMPLES_PER_SYMBOL
|
||||
self.probe_symbols = lab045.PREAMBLE_BIT_COUNT + lab045.PHY_HEADER_BIT_COUNT
|
||||
self.probe_waveform = (self.probe_symbols + 42) * lab045.SAMPLES_PER_SYMBOL
|
||||
self.front_length = min(samples.size, owned_end + self.probe_waveform + offline.REFINEMENT_MARGIN_SAMPLES)
|
||||
self.owned_end = owned_end
|
||||
self.prefix_correlation = None
|
||||
self.changes = None
|
||||
self.matched = None
|
||||
self.baseband = None
|
||||
self.soft_cache = {}
|
||||
|
||||
def coarse_scores(self, samples: np.ndarray, reference: np.ndarray) -> np.ndarray:
|
||||
"""Сохранить DC всего старого окна, считать корреляцию только owned+1."""
|
||||
decimation = lab043.DETECTION_DECIMATION
|
||||
needed = math.ceil(self.owned_end / decimation) + reference.size + 1
|
||||
raw = (resample_poly(samples, 1, decimation, window=decimation_taps(decimation)) if self.decimator is None
|
||||
else self.decimator.get(samples, self.base))
|
||||
mean = np.mean(raw, dtype=np.complex128).astype(np.complex64)
|
||||
values = raw[:needed].copy()
|
||||
values -= mean
|
||||
values *= detection_rotation(values.size, decimation)
|
||||
if values.size < reference.size:
|
||||
return np.empty(0, np.float32)
|
||||
scores = lab043._normalized_sliding_correlation(values, reference).astype(np.float32)
|
||||
scores[~np.isfinite(scores)] = 0
|
||||
np.clip(scores, 0, 1, out=scores)
|
||||
return scores
|
||||
|
||||
def local_start(self, array: np.ndarray) -> int:
|
||||
return (array.ctypes.data - self.samples.ctypes.data) // self.samples.itemsize
|
||||
|
||||
def refine(self, samples: np.ndarray, reference: np.ndarray, coarse_start: int) -> int:
|
||||
start = self.local_start(samples)
|
||||
low = max(0, coarse_start - 2000)
|
||||
high = min(samples.size, coarse_start + self.reference_length + 2000)
|
||||
if high - low < self.reference_length or start < 0:
|
||||
return ORIGINAL_REFINE(samples, reference, coarse_start)
|
||||
if self.changes is None:
|
||||
self.changes = np.empty(self.front_length, np.int64)
|
||||
self.changes[0] = 0
|
||||
np.cumsum(self.samples[1:self.front_length] != self.samples[:self.front_length-1],
|
||||
dtype=np.int64, out=self.changes[1:])
|
||||
if self.changes[start + samples.size - 1] == self.changes[start]:
|
||||
# После вычитания локального среднего точно нулевой сигнал.
|
||||
# Общая FFT может оставить шум округления вместо нулевого tie.
|
||||
return low
|
||||
if self.prefix_correlation is None:
|
||||
centered = self.reference[:self.reference_length] - np.mean(self.reference[:self.reference_length])
|
||||
self.centered_reference_sum = np.sum(np.conj(centered), dtype=np.complex128)
|
||||
self.prefix_correlation = fftconvolve(
|
||||
self.samples[:self.front_length], np.conj(centered[::-1]), mode="valid")
|
||||
self.statistics["shared_refinement_correlations"] += 1
|
||||
first = start + low
|
||||
last = start + high - self.reference_length + 1
|
||||
correlation = self.prefix_correlation[first:last]
|
||||
if not correlation.size:
|
||||
return ORIGINAL_REFINE(samples, reference, coarse_start)
|
||||
# Та же DC-поправка, что у корреляции (samples - mean(probe)).
|
||||
correlation = correlation - np.mean(samples) * self.centered_reference_sum
|
||||
magnitudes = np.abs(correlation)
|
||||
if magnitudes.size > 1:
|
||||
largest = np.partition(magnitudes, -2)[-2:]
|
||||
if largest[-1] - largest[-2] <= 32*np.finfo(np.float32).eps*largest[-1]:
|
||||
self.statistics["refinement_tie_fallbacks"] += 1
|
||||
return ORIGINAL_REFINE(samples, reference, coarse_start)
|
||||
return low + int(np.argmax(magnitudes))
|
||||
|
||||
def matched_window(self, samples: np.ndarray, start: int, length: int):
|
||||
absolute = self.local_start(samples) + start
|
||||
taps = lab043._matched_filter_taps()
|
||||
edge = taps.size - 1
|
||||
if absolute < 0 or absolute + length + edge > self.front_length:
|
||||
self.statistics["local_body_matched_calls"] += 1
|
||||
segment = samples[start:start + length] - np.mean(samples)
|
||||
baseband = segment * carrier_rotation(length)
|
||||
return oaconvolve(baseband, taps, mode="full"), 16 * lab045.SAMPLES_PER_SYMBOL + edge
|
||||
if self.matched is None:
|
||||
self.baseband = self.samples[:self.front_length] * carrier_rotation(self.front_length)
|
||||
self.matched = oaconvolve(self.baseband, taps, mode="full")
|
||||
self.statistics["shared_matched_calls"] += 1
|
||||
matched = self.matched[absolute:absolute + length + edge].copy()
|
||||
# Восстановить нулевые края исходного локального окна, без соседнего IQ.
|
||||
matched[:edge] = convolve(self.baseband[absolute:absolute + edge], taps, method="direct")[:edge]
|
||||
matched[-edge:] = convolve(self.baseband[absolute + length - edge:absolute + length], taps,
|
||||
method="direct")[edge:]
|
||||
phase = np.complex64(np.exp(2j * np.pi * lab043.CARRIER_OFFSET_HZ * absolute / lab045.SAMPLE_RATE_HZ))
|
||||
matched *= phase
|
||||
mean = np.mean(samples)
|
||||
matched -= mean * dc_response(length, lab045.SAMPLES_PER_SYMBOL)
|
||||
nominal = 16 * lab045.SAMPLES_PER_SYMBOL + edge
|
||||
return matched, nominal
|
||||
|
||||
def timing(self, matched, nominal, expected):
|
||||
"""Локальный FFT оказался дешевле дополнительных общих timing-массивов."""
|
||||
return fast_preamble_timing(matched, nominal, expected)
|
||||
|
||||
def extract(self, samples: np.ndarray, frame_start: int, symbol_count: int) -> iq.SoftWindow:
|
||||
key = (self.local_start(samples) + frame_start, symbol_count)
|
||||
# Усечение зависит и от доступности текущего probe, а не только от start.
|
||||
length = (symbol_count + 42) * lab045.SAMPLES_PER_SYMBOL
|
||||
if frame_start < 0 or frame_start + length > samples.size:
|
||||
return iq.SoftWindow(False, np.empty(0), 0, 0, 0, {}, "TRUNCATED_IQ_WINDOW")
|
||||
if key in self.soft_cache:
|
||||
self.statistics["soft_window_cache_hits"] += 1
|
||||
return self.soft_cache[key]
|
||||
segment = samples[frame_start:frame_start + length]
|
||||
matched, nominal = lab043._matched_waveform(segment, 0, length)
|
||||
bits = lab045.lab044.preamble_bits()
|
||||
expected = iq.radio.bpsk_modulate(bits)
|
||||
timing = lab043._preamble_timing_offset(matched, nominal, expected)
|
||||
carrier = lab043._evaluate_symbols(matched, nominal, timing, bits, lab045.PREAMBLE_BIT_COUNT)
|
||||
if carrier["abs_h"] == 0:
|
||||
result = iq.SoftWindow(False, np.empty(0), timing, 0, nominal, carrier, "ZERO_GAIN")
|
||||
else:
|
||||
symbols = matched[nominal + timing:nominal + timing + symbol_count * lab045.SAMPLES_PER_SYMBOL:
|
||||
lab045.SAMPLES_PER_SYMBOL].astype(np.complex128)
|
||||
if symbols.size != symbol_count:
|
||||
result = iq.SoftWindow(False, np.empty(0), timing, 0, nominal, carrier, "TRUNCATED_SYMBOLS")
|
||||
else:
|
||||
symbol_time = (np.arange(symbol_count, dtype=np.float64) + 0.5) / lab045.SYMBOL_RATE
|
||||
normalized = symbols * np.exp(-2j * math.pi * carrier["cfo_hz"] * symbol_time)
|
||||
normalized /= carrier["abs_h"] * np.exp(1j * carrier["phase_rad"])
|
||||
result = iq.SoftWindow(True, normalized.real, int(timing),
|
||||
int((frame_start + nominal + timing) % lab045.SAMPLES_PER_SYMBOL),
|
||||
nominal, carrier, None)
|
||||
self.soft_cache[key] = result
|
||||
self.statistics["unique_soft_windows"] += 1
|
||||
return result
|
||||
|
||||
|
||||
@contextmanager
|
||||
def optimized_context(front: BlockFrontEnd):
|
||||
targets = ((lab043, "_refine_frame_start", front.refine),
|
||||
(lab043, "_matched_waveform", front.matched_window),
|
||||
(lab043, "_preamble_timing_offset", front.timing),
|
||||
(iq, "_extract_soft_window", front.extract),
|
||||
(lab045.conv, "viterbi_decode_soft", fast_viterbi),
|
||||
(offline, "_correlation_scores", front.coarse_scores))
|
||||
originals = [(owner, name, getattr(owner, name)) for owner, name, unused in targets]
|
||||
try:
|
||||
for owner, name, function in targets:
|
||||
setattr(owner, name, function)
|
||||
yield
|
||||
finally:
|
||||
for owner, name, function in reversed(originals):
|
||||
setattr(owner, name, function)
|
||||
|
||||
|
||||
class StreamingReceiver(previous.StreamingReceiver):
|
||||
"""Совместимый API Lab050, ускорение только для рабочей точки 500 ксимв/с."""
|
||||
|
||||
def __init__(self, symbol_rate: int = 500_000, *, block_samples: int = 65_536):
|
||||
if symbol_rate != 500_000:
|
||||
raise ValueError("Lab051 исследует только 500 ксимв/с")
|
||||
super().__init__(symbol_rate, block_samples=block_samples)
|
||||
self.front_statistics = Counter()
|
||||
self.decimator = DecimationCache(self.decimation)
|
||||
|
||||
def _scan(self, owned_end: int) -> list[previous.ReceivedPacket]:
|
||||
front = BlockFrontEnd(self.buffer, owned_end - self.base_sample, self.front_statistics,
|
||||
decimator=self.decimator, base=self.base_sample)
|
||||
with optimized_context(front):
|
||||
return super()._scan(owned_end)
|
||||
|
||||
def finish(self):
|
||||
result = super().finish()
|
||||
self.decimator.values = np.empty(0, np.complex64)
|
||||
return result
|
||||
|
||||
|
||||
def run_functional_tests() -> None:
|
||||
receiver = StreamingReceiver()
|
||||
assert receiver.mode.samples_per_symbol == 16
|
||||
assert receiver.finish() == []
|
||||
Reference in New Issue
Block a user