Lab053: achieve real-time 5 fps software video receiver

experiments/lab053_viterbi.py

@@ -0,0 +1,98 @@
+"""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()