99 lines
3.8 KiB
Python
99 lines
3.8 KiB
Python
"""Lab053: те же переходы soft Viterbi, циклы компилируются только для CPU."""
|
||
|
||
from __future__ import annotations
|
||
|
||
from pathlib import Path
|
||
import sys
|
||
import time
|
||
|
||
import numpy as np
|
||
|
||
if __package__ in (None, ""):
|
||
sys.path.insert(0, str(Path(__file__).resolve().parents[1]))
|
||
|
||
from experiments import lab051_streaming_receiver as previous
|
||
|
||
runtime = Path(__file__).resolve().parents[1] / ".lab053_cpu_runtime/packages"
|
||
if runtime.is_dir():
|
||
sys.path.insert(0, str(runtime))
|
||
try:
|
||
from numba import njit
|
||
except ImportError:
|
||
njit = None
|
||
|
||
conv = previous.lab045.conv
|
||
symbols = conv._BRANCH_SYMBOLS
|
||
CODES = (2 * (symbols[:, :, 0] > 0) + (symbols[:, :, 1] > 0)).astype(np.intp)
|
||
PREDECESSORS = conv._PREDECESSORS.copy()
|
||
CODES.flags.writeable = PREDECESSORS.flags.writeable = False
|
||
|
||
|
||
def trellis_decode(metrics, predecessors, codes):
|
||
"""Порядок сложений исходный; строгое < сохраняет tie в ветвь 0.
|
||
|
||
Не объединяем несколько шагов: это могло бы изменить решение при
|
||
равных округлённых суммах. Старт/финиш 0, 64 состояния, хвост 6 бит.
|
||
"""
|
||
step_count = metrics.shape[0]
|
||
path = np.full(64, 1e18, dtype=np.float64)
|
||
path[0] = 0.0
|
||
next_path = np.empty(64, np.float64)
|
||
choices = np.empty((step_count, 64), np.uint8)
|
||
for step in range(step_count):
|
||
for state in range(64):
|
||
zero = path[predecessors[state, 0]] + metrics[step, codes[state, 0]]
|
||
one = path[predecessors[state, 1]] + metrics[step, codes[state, 1]]
|
||
take_one = one < zero
|
||
choices[step, state] = take_one
|
||
next_path[state] = one if take_one else zero
|
||
path, next_path = next_path, path
|
||
decoded = np.empty(step_count, np.uint8)
|
||
state = 0
|
||
for step in range(step_count - 1, -1, -1):
|
||
decoded[step] = state >> 5
|
||
state = predecessors[state, choices[step, state]]
|
||
assert state == 0
|
||
return decoded[:-6]
|
||
|
||
|
||
compiled_trellis = njit(cache=True, nogil=True, fastmath=False)(trellis_decode) if njit is not None else None
|
||
|
||
|
||
def viterbi_decode_soft(soft_values, information_bit_count):
|
||
if compiled_trellis is None:
|
||
return previous.fast_viterbi(soft_values, information_bit_count)
|
||
if not isinstance(information_bit_count, (int, np.integer)) or information_bit_count <= 0:
|
||
return previous.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 previous.ORIGINAL_VITERBI(soft_values, information_bit_count)
|
||
received = values.reshape(int(information_bit_count) + 6, 2)
|
||
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]))
|
||
return compiled_trellis(metrics, PREDECESSORS, CODES)
|
||
|
||
|
||
def warmup() -> dict:
|
||
started = time.perf_counter()
|
||
values = np.ones(conv.encoded_bit_count(8), np.float64)
|
||
assert np.array_equal(viterbi_decode_soft(values, 8), previous.fast_viterbi(values, 8))
|
||
return {"backend": "NUMBA_CPU" if compiled_trellis is not None else "NUMPY_LAB051",
|
||
"warmup_seconds": time.perf_counter() - started, "fastmath": False, "gpu_used": False}
|
||
|
||
|
||
def run_functional_tests() -> None:
|
||
values = np.zeros(conv.encoded_bit_count(8), np.float64)
|
||
assert np.array_equal(viterbi_decode_soft(values, 8), previous.fast_viterbi(values, 8))
|
||
|
||
|
||
def main() -> None:
|
||
print(warmup())
|
||
run_functional_tests()
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|