reframr/ternary.py

import math

from .linalg import Vector, mean


def quantize_vector_absmean(
    values: Vector,
    *,
    threshold: float = 0.5,
) -> tuple[float, list[int]]:
    if not values:
        return 1.0, []

    scale = mean([abs(value) for value in values])
    if scale == 0.0:
        return 1.0, [0 for _ in values]

    quantized: list[int] = []
    for value in values:
        normalized = value / scale
        if normalized >= threshold:
            quantized.append(1)
        elif normalized <= -threshold:
            quantized.append(-1)
        else:
            quantized.append(0)
    return scale, quantized


def derive_ternary_mask_from_states(states: list[Vector]) -> tuple[float, list[int]]:
    if not states:
        return 1.0, []
    feature_count = len(states[0])
    feature_energy = [
        mean([state[feature] * state[feature] for state in states])
        for feature in range(feature_count)
    ]
    return derive_ternary_mask_from_feature_energy(feature_energy)


def derive_ternary_mask_from_feature_energy(
    feature_energy: Vector,
    *,
    threshold: float = 0.02,
) -> tuple[float, list[int]]:
    if not feature_energy:
        return 1.0, []

    rms_values = [math.sqrt(max(value, 0.0)) for value in feature_energy]
    scale = mean(rms_values)
    if scale == 0.0:
        return 1.0, [0 for _ in feature_energy]

    mask = [1 if value >= threshold * scale else 0 for value in rms_values]
    if not any(mask):
        mask = [1 for _ in feature_energy]
    return 1.0, mask


def apply_ternary_mask(values: Vector, mask: list[int], scale: float) -> Vector:
    if not mask:
        return values[:]
    return [scale * mask[index] * values[index] for index in range(len(values))]