palimseste/learner.py

"""PALIMPSESTE — Learning as O(1) writing (Axiome 3) + encoding helpers.

Axiome 3 — Apprendre est une ecriture, pas une mise a jour
    Observing ``(x, y)`` in context ``c``:
        M <- M ∪ {(a = bind(x, c), v = y, w = 1)}
    Cost: O(1) amortized. One insertion. That is what "training continuously
    on massive data" means here: it is just writing into a table.

This module provides:

  - :class:`Encoder`  : turns raw observations (ints, floats, strings, tuples)
    into hypervectors so the substrate can ingest arbitrary modalities. Uses
    the standard VSA recipe: random atomic item-vectors + binding/bundling to
    compose structures (Plate 1995; Kanerva 1997).
  - :class:`Learner`  : the O(1) write primitive. Given ``(x, y, ctx)`` (each
    already an HV or encodable), it computes ``a = bind(x, ctx)`` and appends
    ``(a, v=y, w=weight)`` to ``M``. It also offers a "supervised predict"
    convenience: query with ``bind(x, ctx)`` and return ``Phi``'s reconstruction.

There is deliberately *no* gradient, no epoch, no batch, no optimizer state.
The only mutable state is ``M`` itself.
"""

from __future__ import annotations

from dataclasses import dataclass, field
import numpy as np

from .hv import HV, bind, bundle, random_hv, similarity
from .memory import Memory, Trace
from .phi import Phi

__all__ = ["Encoder", "Learner", "Prediction"]


@dataclass
class Prediction:
    """Result of a supervised query."""

    query: HV
    value: HV | None
    confidence: float
    n_matches: int


# --------------------------------------------------------------------- Encoder
@dataclass
class Encoder:
    """Encode raw values into hypervectors.

    Maintains a stable atomic-symbol table (random HVs per distinct atom) so
    that the same token always maps to the same vector. Composition rules:

      - ``encode_int(k)``     : role-bound int (uses a per-position role HV)
      - ``encode_float(x)``   : scalar -> HV via level/bucket coding
      - ``encode_str(s)``     : hash-to-HV (stable)
      - ``encode_sequence(xs)``: bind each item with its positional role,
        then bundle (order-sensitive set representation)
      - ``encode_set(xs)``    : pure bundle (order-insensitive)

    All atomic vectors are ``D``-dimensional and lazily memoized.
    """

    D: int
    rng: np.random.Generator = field(default_factory=np.random.default_rng)
    _atoms: dict[object, HV] = field(default_factory=dict)
    _roles: dict[int, HV] = field(default_factory=dict)
    _levels: list[HV] | None = None
    _n_levels: int = 256

    # ------------------------------------------------------------------ atoms
    def _atom(self, key: object) -> HV:
        v = self._atoms.get(key)
        if v is None:
            # Deterministic: seed from the key's hash so the same key always
            # produces the same HV, regardless of creation order or rng state.
            # This is critical for save/load consistency: a model loaded from
            # disk must produce identical atom HVs as the original.
            import hashlib
            h = hashlib.sha256(repr(key).encode("utf-8")).digest()
            seed = int.from_bytes(h[:4], "little")
            v = random_hv(self.D, rng=np.random.default_rng(seed))
            self._atoms[key] = v
        return v

    def _role(self, pos: int) -> HV:
        v = self._roles.get(pos)
        if v is None:
            # Deterministic: seed from the position so role(pos) is always the
            # same HV regardless of creation order.
            import hashlib
            h = hashlib.sha256(f"role:{pos}".encode("utf-8")).digest()
            seed = int.from_bytes(h[:4], "little")
            v = random_hv(self.D, rng=np.random.default_rng(seed))
            self._roles[pos] = v
        return v

    # ----------------------------------------------------------- primitive encoders
    def encode_int(self, k: int) -> HV:
        return self._atom(("i", int(k)))

    def encode_str(self, s: str) -> HV:
        return self._atom(("s", s))

    def encode_bool(self, b: bool) -> HV:
        return self._atom(("b", bool(b)))

    def encode_float(self, x: float, lo: float = 0.0, hi: float = 1.0) -> HV:
        """Level-code a scalar in ``[lo, hi]`` into an HV.

        Adjacent levels are *similar* (correlated), distant levels are
        quasi-orthogonal — the classic scalar-encoding trick for VSAs so
        that ``encode_float(0.41)`` resembles ``encode_float(0.42)``.
        """
        if hi <= lo:
            raise ValueError("hi must be > lo")
        x = min(max(x, lo), hi)
        frac = (x - lo) / (hi - lo)
        levels = self._levels_for()
        idx = int(round(frac * (self._n_levels - 1)))
        return levels[idx]

    def _levels_for(self) -> list[HV]:
        if self._levels is None:
            # Correlated ladder via a bounded random walk: start from a random
            # base HV and flip ``bits_per_step`` random positions at each level.
            # Sized for an *adjacent*-level similarity of ~0.95 so that:
            #   - consecutive levels are near-identical (good local resolution),
            #   - the walk saturates to quasi-orthogonal within ~30 levels
            #     (so distant scalars decorrelate, as VSA scalar coding demands).
            base = random_hv(self.D, rng=self.rng)
            signs = self._signs(base)
            levels = [base]
            # target adjacent sim 0.95  ->  hamming = D*(1-0.95)/2 = D*0.025
            bits_per_step = max(1, int(round(self.D * 0.025)))
            cur = signs.copy()
            for _ in range(self._n_levels - 1):
                flip = self.rng.choice(self.D, size=bits_per_step, replace=False)
                cur = cur.copy()
                cur[flip] = -cur[flip]
                levels.append(self._pack(cur))
            self._levels = levels
        return self._levels

    @staticmethod
    def _signs(h: HV) -> np.ndarray:
        from .hv import bits_to_signs

        return bits_to_signs(h)

    @staticmethod
    def _pack(signs: np.ndarray) -> HV:
        from .hv import signs_to_bits

        return signs_to_bits(signs)

    # ----------------------------------------------------------- composition
    def encode_sequence(self, xs: list[HV]) -> HV:
        """Order-sensitive bundle: ``bundle_i bind(x_i, role_i)``.

        Deterministic (same inputs => same output) so that the same context
        always maps to the same address — essential for associative retrieval.
        """
        if not xs:
            return random_hv(self.D, rng=self.rng)  # non-empty default
        bound = [bind(x, self._role(i)) for i, x in enumerate(xs)]
        return bundle(bound, rng=self.rng, deterministic=True)

    def encode_set(self, xs: list[HV]) -> HV:
        """Order-insensitive bundle: ``bundle_i x_i``.

        Deterministic (same inputs => same output).
        """
        if not xs:
            return random_hv(self.D, rng=self.rng)
        return bundle(xs, rng=self.rng, deterministic=True)

    def encode_kv(self, pairs: list[tuple[HV, HV]]) -> HV:
        """Bind key/value pairs, then bundle (record/struct encoding).

        Deterministic (same inputs => same output).
        """
        if not pairs:
            return random_hv(self.D, rng=self.rng)
        bound = [bind(k, v) for k, v in pairs]
        return bundle(bound, rng=self.rng, deterministic=True)


# ---------------------------------------------------------------------- Learner
@dataclass
class Learner:
    """The O(1) supervised-write primitive.

    ``learn(x, y, ctx)`` appends ``(bind(x, ctx), y, w)`` to ``M``.
    ``predict(x, ctx)`` queries ``Phi(bind(x, ctx))``.

    The Learner is a thin orchestration over :class:`Memory` + :class:`Phi`;
    all heavy lifting (storage, retrieval, decay) lives in those modules.
    """

    mem: Memory
    phi: Phi
    rng: np.random.Generator = field(default_factory=np.random.default_rng)

    def learn(
        self,
        x: HV,
        y: HV,
        ctx: HV | None = None,
        weight: float = 1.0,
        tag: str | None = None,
    ) -> Trace:
        """Axiome 3: append ``(a=bind(x,ctx), v=y, w=weight)`` to ``M``. O(1)."""
        if ctx is None:
            ctx = self._identity_ctx()
        address = bind(x, ctx)
        return self.mem.write(address, y, weight=weight, tag=tag)

    def predict(self, x: HV, ctx: HV | None = None) -> Prediction:
        """Query ``Phi(bind(x, ctx))`` and report confidence."""
        if ctx is None:
            ctx = self._identity_ctx()
        q = bind(x, ctx)
        ret = self.phi.retrieve(self.mem, q)
        val = self.phi(self.mem, q)
        # confidence: weighted-mean similarity of matches in [-1, 1] -> [0, 1]
        if ret.matches:
            ws = np.asarray(ret.weights)
            sims = np.asarray(ret.sims)
            conf = float(np.clip((ws * sims).sum() / ws.sum(), -1.0, 1.0))
            conf = (conf + 1.0) / 2.0
        else:
            conf = 0.0
        return Prediction(query=q, value=val, confidence=conf, n_matches=ret.n_matches)

    def _identity_ctx(self) -> HV:
        """A stable all-+1 context (binding with it is identity)."""
        from .hv import constant_hv

        return constant_hv(self.mem.D, value=+1)

    # --------------------------------------------------- batch convenience
    def learn_many(
        self, pairs: list[tuple[HV, HV]], ctx: HV | None = None
    ) -> list[Trace]:
        """Append many ``(x, y)`` pairs. Each is O(1); the loop is O(n)."""
        return [self.learn(x, y, ctx=ctx) for x, y in pairs]