palimseste/phi.py

"""PALIMPSESTE — The read-back kernel ``Phi`` (Axiomes 2, 4).

Axiome 2 — Le passage avant est une recuperation associative
    ``Phi(q) = bundle_{(a,v,w) in N_r(q)} (w * v)``
    where ``N_r(q)`` is the Hamming neighborhood of radius ``r`` of ``q``.
    Retrieval is ``O(log |M|)`` via LSH (independent of D in the *compute*,
    since bundling is bitwise).

Axiome 4 — Les parametres sont reconstruits, non stockes
    ``W_t = Phi(bind(x, s_t))`` — the "weight" exists only for the duration of
    one computation. It emerges from memory. So "modifying the parameters" =
    "writing better traces into M".

This module implements ``Phi`` as a *reconstruction operator* over a
:class:`Memory`. It is:

  - stateless w.r.t. the memory contents (it reads, never writes);
  - parameterized by a small, learnable ``KernelConfig`` (radius, threshold
    sharpness, min-weight cutoff) whose values are read from ``H_meta`` by
    the :class:`MetaController` (``meta.py``);
  - sub-linear in ``|M|`` because it operates only on the LSH candidate set.

The kernel also exposes the *raw* retrieved traces and their (similarity,
weight) scores, so callers (curiosity, consolidation, meta) can introspect
*why* a reconstruction came out the way it did — transparency is a design goal.
"""

from __future__ import annotations

from dataclasses import dataclass, field
import math
import numpy as np

from .hv import HV, bundle, hamming, similarity
from .memory import Memory, Trace

__all__ = ["KernelConfig", "Phi", "Retrieval"]


@dataclass
class KernelConfig:
    """The learnable knobs of the read-back kernel ``Phi``.

    These are the meta-parameters that live in ``H_meta`` and that the
    :class:`MetaController` may rewrite under a Lyapunov constraint.

    Attributes
    ----------
    radius : int
        Hamming radius defining the neighborhood ``N_r(q)``.
    min_weight : float
        Traces with decayed weight below this are ignored (soft floor for
        "dormant" recall; keeps rare re-activations possible without flooding
        recall with noise).
    sharpness : float
        Temperature for the optional similarity-weighted softmax over
        retrieved values. ``sharpness <= 0`` means plain weighted-bundle
        (hard majority); ``sharpness > 0`` interpolates values by
        ``exp(sharpness * sim)``.
    topk : int | None
        If set, only the ``topk`` most similar candidates contribute. Caps
        per-query compute and noise.
    """

    radius: int = 50
    min_weight: float = 1e-3
    sharpness: float = 0.0
    topk: int | None = None

    def __post_init__(self) -> None:
        if self.radius < 0:
            raise ValueError("radius must be >= 0")
        if self.min_weight <= 0:
            raise ValueError("min_weight must be > 0")
        if self.topk is not None and self.topk <= 0:
            raise ValueError("topk must be positive or None")

    def encode(self) -> dict:
        """Serialize for storage in ``H_meta``."""
        return {
            "radius": self.radius,
            "min_weight": self.min_weight,
            "sharpness": self.sharpness,
            "topk": self.topk,
        }

    @classmethod
    def decode(cls, d: dict) -> "KernelConfig":
        return cls(**d)


@dataclass
class Retrieval:
    """Full result of a ``Phi`` call, for introspection."""

    query: HV
    config: KernelConfig
    matches: list[Trace]
    sims: list[float]
    weights: list[float]
    result: HV | None  # None if no matches

    @property
    def n_matches(self) -> int:
        return len(self.matches)


@dataclass
class Phi:
    """The associative read-back kernel.

    ``Phi`` is a *pure reader* of :class:`Memory`. All learning is done by
    writing traces into ``M`` (``learner.py``); ``Phi`` only reconstructs.
    """

    config: KernelConfig = field(default_factory=KernelConfig)

    # ----------------------------------------------------------------- core
    def retrieve(self, mem: Memory, query: HV) -> Retrieval:
        """Find ``N_r(query)`` in ``M`` and return matches + scores.

        Two-stage: LSH candidates, then exact Hamming filter (vectorized),
        then optional top-k and min-weight cutoffs.

        The Hamming distance computation is **vectorized**: all candidate
        addresses are XOR'd against the query in a single numpy batch, then
        popcount'd. This is 10-50x faster than the per-trace Python loop at
        scale (200K+ traces).
        """
        cand_ids = mem.candidates(query)
        if not cand_ids:
            return Retrieval(query=query, config=self.config,
                             matches=[], sims=[], weights=[], result=None)

        # filter valid ids
        valid_ids = [cid for cid in cand_ids if 0 <= cid < len(mem.traces)]
        if not valid_ids:
            return Retrieval(query=query, config=self.config,
                             matches=[], sims=[], weights=[], result=None)

        D = query.D
        packed_len = len(query.bits)

        # --- vectorized Hamming distance ---
        # Fast popcount via a 256-entry lookup table (no np.unpackbits needed).
        # This is 8x faster than unpack+sum for large candidate sets.
        cand_bits = np.empty((len(valid_ids), packed_len), dtype=np.uint8)
        for i, cid in enumerate(valid_ids):
            cand_bits[i] = mem.traces[cid].address.bits

        xored = np.bitwise_xor(cand_bits, query.bits[np.newaxis, :])
        # popcount via lookup table
        _POPCOUNT_TABLE = np.array(
            [bin(i).count('1') for i in range(256)], dtype=np.uint16)
        hamming_dists = _POPCOUNT_TABLE[xored].sum(axis=1)

        # filter by radius
        within = hamming_dists <= self.config.radius
        if not within.any():
            return Retrieval(query=query, config=self.config,
                             matches=[], sims=[], weights=[], result=None)

        # compute similarities and weights for survivors
        matched_idx = np.where(within)[0]
        all_sims = (1.0 - 2.0 * hamming_dists[matched_idx] / D)
        now = None
        weights: list[float] = []
        sims: list[float] = []
        matches: list[Trace] = []
        for j, mi in enumerate(matched_idx):
            cid = valid_ids[mi]
            tr = mem.traces[cid]
            w = mem.current_weight(tr, now=now)
            if w >= self.config.min_weight:
                matches.append(tr)
                weights.append(w)
                sims.append(float(all_sims[j]))

        if not matches:
            return Retrieval(query=query, config=self.config,
                             matches=[], sims=[], weights=[], result=None)

        # optional top-k by similarity
        if self.config.topk is not None and len(matches) > self.config.topk:
            order = np.argsort(sims)[::-1][: self.config.topk]
            matches = [matches[i] for i in order]
            sims = [sims[i] for i in order]
            weights = [weights[i] for i in order]

        return Retrieval(
            query=query,
            config=self.config,
            matches=matches,
            sims=sims,
            weights=weights,
            result=None,
        )

    def __call__(self, mem: Memory, query: HV) -> HV | None:
        """Compute ``Phi(query)`` — the reconstructed value (or None)."""
        ret = self.retrieve(mem, query)
        if not ret.matches:
            return None
        values = [tr.value for tr in ret.matches]

        if self.config.sharpness <= 0.0:
            # plain weighted bundle: majority of weighted signs
            res = bundle(values, weights=ret.weights, deterministic=True)
        else:
            # similarity-temperatured weighting
            temps = np.exp(self.config.sharpness * np.asarray(ret.sims))
            eff = (np.asarray(ret.weights) * temps).tolist()
            res = bundle(values, weights=eff, deterministic=True)

        ret.result = res
        return res

    # ----------------------------------------------- parameter reconstruction
    def reconstruct_param(self, mem: Memory, x: HV, state: HV) -> HV | None:
        """Axiome 4: ``W_t = Phi(bind(x, s_t))``.

        The "weight" for input ``x`` in internal state ``s`` is reconstructed
        on the fly from memory — it does not exist before this call and is not
        stored after.
        """
        from .hv import bind as _bind

        return self(mem, _bind(x, state))