| """36-state Viterbi selection over the Q2 symbolic lattice. |
| |
| This is a direct port of Phase 3 of the C engine. Each graph block owns a |
| ``6 x 6`` coarse/fine lattice; candidate errors are reduced to a free-energy |
| per bin, the Sieve marginals become a log-prior bonus, and an exact DP finds |
| the minimum-cost state sequence with a Manhattan transition penalty. |
| |
| Both the 576->36 bin reduction and the path->candidate mapping are fully |
| vectorised by reshaping the ``(24, 24)`` candidate grid into ``(6, 4, 6, 4)``. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import numpy as np |
|
|
| from .config import HPCConfig |
| from .factored import CAND_TO_QUHIT, N_CAND_M, TOTAL_CANDIDATES |
|
|
| N_STATES = 36 |
|
|
| _BIN_CAND_MAP = ( |
| np.arange(TOTAL_CANDIDATES, dtype=np.int64) |
| .reshape(6, 4, 6, 4) |
| .transpose(0, 2, 1, 3) |
| .reshape(N_STATES, 16) |
| ) |
|
|
|
|
| def _median_upper(values: np.ndarray) -> float: |
| arr = np.asarray(values) |
| if arr.size == 0: |
| return 0.0 |
| k = arr.size // 2 |
| return float(np.partition(arr, k)[k]) |
|
|
|
|
| def robust_temperature( |
| candidate_errors: np.ndarray, cfg: HPCConfig |
| ) -> float: |
| n_blocks = candidate_errors.shape[0] |
| if n_blocks <= 0: |
| return 1.0e-10 |
| step = max(1, n_blocks // cfg.temperature_sample_blocks) |
| medians = [] |
| for b in range(0, n_blocks, step): |
| row = candidate_errors[b] |
| finite = row[np.isfinite(row)] |
| if finite.size == 0: |
| continue |
| medians.append(_median_upper(np.maximum(0.0, row - finite.min()))) |
| if len(medians) >= cfg.temperature_sample_blocks: |
| break |
| t = _median_upper(np.asarray(medians)) |
| return t if np.isfinite(t) and t > 1.0e-10 else 1.0e-10 |
|
|
|
|
| def _bin_geometry(): |
| di = np.arange(TOTAL_CANDIDATES) // N_CAND_M |
| mi = np.arange(TOTAL_CANDIDATES) % N_CAND_M |
| qi_d = CAND_TO_QUHIT[di] |
| qi_m = CAND_TO_QUHIT[mi] |
| return qi_d, qi_m, qi_d * 6 + qi_m |
|
|
|
|
| def viterbi_select( |
| candidate_errors: np.ndarray, |
| coarse_marg: np.ndarray, |
| fine_marg: np.ndarray, |
| temperature: float, |
| stride: int, |
| graph_blocks: int, |
| cfg: HPCConfig, |
| ) -> tuple[np.ndarray, np.ndarray]: |
| n_blocks = candidate_errors.shape[0] |
| temp = max(float(temperature), 1.0e-12) |
|
|
| |
| ce_f32 = np.asarray(candidate_errors, dtype=np.float32) |
| ce_view = ce_f32.reshape(n_blocks, 6, 4, 6, 4) |
| blk_bin_min = np.min(ce_view, axis=(2, 4)).reshape(n_blocks, N_STATES).astype(np.float64) |
|
|
| ce_bins = ce_view.transpose(0, 1, 3, 2, 4).reshape(n_blocks, N_STATES, 16) |
| blk_bin_arg = np.argmin(ce_bins, axis=2) |
| blk_bin_cand = _BIN_CAND_MAP[np.arange(N_STATES)[None, :], blk_bin_arg] |
| blk_bin_cand = np.where(np.isfinite(blk_bin_min), blk_bin_cand, -1) |
|
|
| pad = graph_blocks * stride - n_blocks |
| if pad > 0: |
| pad_arr = np.full((pad, 6, 4, 6, 4), np.inf, dtype=np.float32) |
| grouped = np.concatenate([ce_view, pad_arr], axis=0).reshape(graph_blocks, stride, 6, 4, 6, 4) |
| else: |
| grouped = ce_view.reshape(graph_blocks, stride, 6, 4, 6, 4) |
|
|
| group_min_4d = np.min(grouped, axis=(1, 3, 5)) |
| group_min = group_min_4d.reshape(graph_blocks, N_STATES).astype(np.float64) |
| finite = np.isfinite(grouped) |
| counts = np.sum(finite, axis=(1, 3, 5)).reshape(graph_blocks, N_STATES) |
| valid = (counts > 0) & np.isfinite(group_min) |
| gm_b = np.where( np.isfinite(group_min_4d), group_min_4d, 0.0)[:, None, :, None, :, None] |
| shifted = np.where(finite, (grouped - gm_b) / np.float32(2.0 * temp), np.inf) |
| sums = np.sum(np.exp(-shifted), axis=(1, 3, 5), dtype=np.float64).reshape(graph_blocks, N_STATES) |
| mean_exp = np.maximum(sums / np.maximum(counts, 1), 1.0e-300) |
| bin_err = np.where(valid, group_min - 2.0 * temp * np.log(mean_exp), np.inf) |
|
|
| rep_idx = np.arange(graph_blocks) * stride |
| bin_cand = blk_bin_cand[rep_idx] |
|
|
| c_tot = np.sum(coarse_marg, axis=1, keepdims=True) |
| f_tot = np.sum(fine_marg, axis=1, keepdims=True) |
| pc = np.where(c_tot > 1.0e-30, coarse_marg / np.maximum(c_tot, 1.0e-30), 1.0 / 6.0) |
| pf = np.where(f_tot > 1.0e-30, fine_marg / np.maximum(f_tot, 1.0e-30), 1.0 / 6.0) |
| log_prior = ( |
| np.log(pc[:, :, None] + 1.0e-30) + np.log(pf[:, None, :] + 1.0e-30) |
| ).reshape(graph_blocks, N_STATES) |
|
|
| finite_err = bin_err[np.isfinite(bin_err)] |
| scale_err = float(finite_err.mean()) if finite_err.size else 1.0e-10 |
| scale_err = max(scale_err, 1.0e-20) |
|
|
| local = np.where( |
| np.isfinite(bin_err), |
| bin_err - cfg.viterbi_beta * scale_err * log_prior, |
| np.inf, |
| ) |
|
|
| state_d = np.arange(N_STATES) // 6 |
| state_m = np.arange(N_STATES) % 6 |
| trans = ( |
| cfg.viterbi_alpha |
| * scale_err |
| * (np.abs(state_d[None, :] - state_d[:, None]) + np.abs(state_m[None, :] - state_m[:, None])) |
| ) |
|
|
| dp = np.full((graph_blocks, N_STATES), np.inf, dtype=np.float64) |
| back = np.zeros((graph_blocks, N_STATES), dtype=np.int64) |
| dp[0] = local[0] |
| s_range = np.arange(N_STATES) |
| for gi in range(1, graph_blocks): |
| totals = dp[gi - 1][:, None] + trans |
| best_sp = np.argmin(totals, axis=0) |
| best_val = totals[best_sp, s_range] |
| dp[gi] = np.where(np.isfinite(best_val), best_val + local[gi], np.inf) |
| back[gi] = best_sp |
|
|
| if not np.any(np.isfinite(dp[-1])): |
| path = np.zeros(graph_blocks, dtype=np.int64) |
| else: |
| path = np.zeros(graph_blocks, dtype=np.int64) |
| path[-1] = int(np.argmin(dp[-1])) |
| for gi in range(graph_blocks - 2, -1, -1): |
| path[gi] = back[gi + 1, path[gi + 1]] |
|
|
| |
| blk_gi = np.arange(n_blocks) // stride |
| chosen_state = path[blk_gi] |
| chosen_cand = blk_bin_cand[np.arange(n_blocks), chosen_state] |
| global_best_cand = np.argmin(candidate_errors, axis=1).astype(np.int64) |
| best_candidate = np.where(chosen_cand >= 0, chosen_cand, global_best_cand) |
|
|
| cur_err = candidate_errors[np.arange(n_blocks), best_candidate] |
| g_err = candidate_errors[np.arange(n_blocks), global_best_cand] |
| best_candidate = np.where( |
| g_err < cur_err * cfg.greedy_override_ratio, global_best_cand, best_candidate |
| ) |
|
|
| return best_candidate, path |
|
|
