69 lines
3.2 KiB
Python
69 lines
3.2 KiB
Python
"""CPU JIT для той же 64-состояний решётки Витерби VIDEO-004.
|
||
|
||
Режим выбирается вызывающим кодом до приёма пакета. При отсутствии Numba
|
||
этот модуль не подменяет результат медленным декодером после ошибки CRC.
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import numpy as np
|
||
|
||
from protocol import convolutional_fec, punctured_fec
|
||
|
||
try:
|
||
from numba import njit
|
||
except ImportError:
|
||
njit = None
|
||
|
||
|
||
PREDECESSORS = convolutional_fec._PREDECESSORS
|
||
BRANCH_SYMBOLS = convolutional_fec._BRANCH_SYMBOLS
|
||
|
||
|
||
def _trellis(soft_values: np.ndarray, information_bit_count: int,
|
||
predecessors: np.ndarray, branch_symbols: np.ndarray) -> np.ndarray:
|
||
"""Те же метрики, порядок сложения и выбор ветви 0 при равенстве."""
|
||
step_count = information_bit_count + convolutional_fec.TAIL_BIT_COUNT
|
||
received = soft_values.reshape(step_count, 2)
|
||
path_metric = np.full(convolutional_fec.STATE_COUNT, 1e18, dtype=np.float64)
|
||
path_metric[0] = 0.0
|
||
next_metric = np.empty(convolutional_fec.STATE_COUNT, dtype=np.float64)
|
||
choices = np.empty((step_count, convolutional_fec.STATE_COUNT), dtype=np.uint8)
|
||
for step in range(step_count):
|
||
received_0 = received[step, 0]
|
||
received_1 = received[step, 1]
|
||
for state in range(convolutional_fec.STATE_COUNT):
|
||
difference_00 = received_0 - branch_symbols[state, 0, 0]
|
||
difference_01 = received_1 - branch_symbols[state, 0, 1]
|
||
difference_10 = received_0 - branch_symbols[state, 1, 0]
|
||
difference_11 = received_1 - branch_symbols[state, 1, 1]
|
||
first = (path_metric[predecessors[state, 0]]
|
||
+ difference_00 ** 2 + difference_01 ** 2)
|
||
second = (path_metric[predecessors[state, 1]]
|
||
+ difference_10 ** 2 + difference_11 ** 2)
|
||
take_second = second < first
|
||
choices[step, state] = take_second
|
||
next_metric[state] = second if take_second else first
|
||
path_metric, next_metric = next_metric, path_metric
|
||
decoded = np.empty(step_count, dtype=np.uint8)
|
||
current_state = 0
|
||
for step in range(step_count - 1, -1, -1):
|
||
branch = choices[step, current_state]
|
||
decoded[step] = current_state >> (convolutional_fec.MEMORY - 1)
|
||
current_state = predecessors[current_state, branch]
|
||
return decoded[:information_bit_count]
|
||
|
||
|
||
COMPILED_TRELLIS = (njit(cache=False, nogil=True, fastmath=False)(_trellis)
|
||
if njit is not None else None)
|
||
|
||
|
||
def decode_soft(soft_values: np.ndarray, information_bit_count: int,
|
||
rate: str) -> np.ndarray:
|
||
"""Проверить вход и декодировать проколотое слово без запасного алгоритма."""
|
||
if COMPILED_TRELLIS is None:
|
||
raise RuntimeError("Numba недоступна для VIDEO-004")
|
||
restored = punctured_fec.depuncture_soft(soft_values, information_bit_count, rate)
|
||
return COMPILED_TRELLIS(restored, int(information_bit_count),
|
||
PREDECESSORS, BRANCH_SYMBOLS)
|