365 lines
19 KiB
Python
365 lines
19 KiB
Python
"""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() == []
|