"""Hippocampal one-shot associative memory. Biological motivation (complementary learning systems): a fast, episodic, one-shot memory next to a slow statistical cortex (the spiking core). Key design — hybrid addressing (dentate-gyrus separation + CA3 completion): * The key of a memory is the last ``m`` tokens (patches), one 32-bit lane per context slot. * PRIMARY address: a hash of the FULL key (pattern separation — every distinct context gets its own slot, so capacity is real). * SECONDARY addresses: one per lane (auto-associative backup probes — a query with a corrupted context still reaches the entry through its intact lanes). * Retrieval is two-tier: a full-key match at the primary address dominates; otherwise candidates matching >= min_lanes lanes vote by lane-match count x log evidence count. This yields both capacity AND graceful degradation: corrupting 1 of 4 context tokens still retrieves the original continuation through the intact lanes. """ from __future__ import annotations import numpy as np from .hashing import LANE_SALTS, splitmix64 class Hippocampus: def __init__( self, m: int = 4, size: int = 1 << 23, min_lanes: int = 3, ) -> None: self.m = m self.size = size self.min_lanes = min_lanes # slot storage (vals are stable uint64 PATCH IDS, never indices) self.keys = np.zeros((size, m), dtype=np.uint32) self.vals = np.zeros(size, dtype=np.uint64) self.counts = np.zeros(size, dtype=np.uint8) self.used = np.zeros(size, dtype=np.bool_) self.n_writes = 0 # ---------------------------------------------------------------- utils def primary_addr(self, lane_keys: np.ndarray) -> int: """Pattern-separation hash of the FULL key -> primary slot (dentate-gyrus style: every distinct context gets its own address).""" z = (int(lane_keys[0]) ^ ((int(lane_keys[1]) << 8 | int(lane_keys[1]) >> 24) & 0xFFFFFFFF) ^ ((int(lane_keys[2]) << 16 | int(lane_keys[2]) >> 16) & 0xFFFFFFFF) ^ ((int(lane_keys[3]) << 24 | int(lane_keys[3]) >> 8) & 0xFFFFFFFF)) return splitmix64(z) % self.size def loo_addr(self, lane_keys: np.ndarray, l: int) -> int: """Leave-one-out hash of the key minus lane l -> CA3 completion slot. A query whose lane l is corrupted can still compute this address from its intact lanes and reach the stored entry.""" a = [int(lane_keys[i]) for i in range(self.m) if i != l] z = a[0] ^ (((a[1] << 13) | (a[1] >> 19)) & 0xFFFFFFFF) \ ^ (((a[2] << 26) | (a[2] >> 6)) & 0xFFFFFFFF) return splitmix64(z ^ (0x9E3779B9 + l * 0x85EBCA6B)) % self.size @staticmethod def lane_keys(patch_ids) -> np.ndarray: """(m,) uint32 lane keys for a context of m patch ids.""" out = np.empty(len(patch_ids), dtype=np.uint32) for lane, pid in enumerate(patch_ids): out[lane] = splitmix64(int(pid) ^ LANE_SALTS[lane]) >> 32 return out def lane_keys_stream(self, patch_ids: np.ndarray) -> np.ndarray: """Vectorized lane keys for a whole token stream: (T, m) uint32.""" T = len(patch_ids) out = np.empty((T, self.m), dtype=np.uint32) for lane in range(self.m): # splitmix64 on an array of uint64 ids, take high 32 bits z = patch_ids.astype(np.uint64) ^ np.uint64(LANE_SALTS[lane]) z = z + np.uint64(0x9E3779B97F4A7C15) z = (z ^ (z >> np.uint64(30))) * np.uint64(0xBF58476D1CE4E5B9) z = (z ^ (z >> np.uint64(27))) * np.uint64(0x94D049BB133111EB) z = z ^ (z >> np.uint64(31)) out[:, lane] = (z >> np.uint64(32)).astype(np.uint32) return out # ---------------------------------------------------------------- write def write(self, lane_keys: np.ndarray, val: int) -> None: """One-shot write of (context -> continuation patch id). Stores the entry at its PRIMARY (full-key hash) slot — pattern separation gives every distinct context its own address — plus best-effort copies at the m lane slots for corrupted-cue retrieval. """ # primary slot: exact-context capacity paddr = self.primary_addr(lane_keys) if not self.used[paddr]: self.keys[paddr] = lane_keys self.vals[paddr] = val self.counts[paddr] = 1 self.used[paddr] = True elif np.array_equal(self.keys[paddr], lane_keys): if self.vals[paddr] == val: if self.counts[paddr] < 255: self.counts[paddr] += 1 else: c = int(self.counts[paddr]) if c > 1: self.counts[paddr] = c - 1 else: self.vals[paddr] = val self.counts[paddr] = 1 elif int(self.counts[paddr]) <= 1: # evict unconsolidated memories only (keep consolidated ones) self.keys[paddr] = lane_keys self.vals[paddr] = val self.counts[paddr] = 1 # leave-one-out completion slots: robust to one corrupted lane for l in range(self.m): addr = self.loo_addr(lane_keys, l) if not self.used[addr]: self.keys[addr] = lane_keys self.vals[addr] = val self.counts[addr] = 1 self.used[addr] = True elif np.array_equal(self.keys[addr], lane_keys): if self.vals[addr] == val: if self.counts[addr] < 255: self.counts[addr] += 1 # else: a different key occupies this completion slot — keep it self.n_writes += 1 def write_stream(self, lane_stream: np.ndarray, vals: np.ndarray) -> None: """Bulk one-shot write over a stream (val[t] follows context of lane_stream[t]).""" for t in range(len(vals)): self.write(lane_stream[t], int(vals[t])) # ---------------------------------------------------------------- query def query(self, lane_keys: np.ndarray): """Query with the current context. Returns (patch_id, weight) or (None, 0.0). Two-tier voting: a candidate matching ALL m lanes dominates any partial match (a full-key entry must never lose to a partial match with a higher count); within the winning tier, votes scale with lane-match count x log evidence count. """ votes: dict[int, float] = {} max_match: dict[int, int] = {} # primary probe: exact full-key retrieval paddr = self.primary_addr(lane_keys) if self.used[paddr] and np.array_equal(self.keys[paddr], lane_keys): v = int(self.vals[paddr]) c = int(self.counts[paddr]) votes[v] = votes.get(v, 0.0) + self.m * (1.0 + np.log1p(c)) max_match[v] = self.m # leave-one-out probes: corrupted-cue completion (CA3) for l in range(self.m): addr = self.loo_addr(lane_keys, l) if not self.used[addr]: continue match = int(np.count_nonzero(self.keys[addr] == lane_keys)) if match >= self.min_lanes: v = int(self.vals[addr]) c = int(self.counts[addr]) w = match * (1.0 + np.log1p(c)) votes[v] = votes.get(v, 0.0) + w if match > max_match.get(v, 0): max_match[v] = match if not votes: return None, 0.0 if max(max_match.values()) == self.m: votes = {v: w for v, w in votes.items() if max_match[v] == self.m} best_v, best_w = max(votes.items(), key=lambda kv: kv[1]) return best_v, best_w # ---------------------------------------------------------------- stats def stats(self) -> dict: n_used = int(self.used.sum()) return { "slots": self.size, "used": n_used, "load_factor": n_used / self.size, "n_writes": self.n_writes, } def _slots_of(self, lane_keys: np.ndarray) -> list[int]: """All derivable slot addresses of a key: primary + leave-one-out.""" out = [self.primary_addr(lane_keys)] for l in range(self.m): out.append(self.loo_addr(lane_keys, l)) return out def top_entries(self, registry=None, min_count: int = 2, cap: int | None = None): """Deduplicated (lane_keys, val, count) list for export/eval, sorted by count descending. Dedup is EXACT: every slot an entry occupies is derivable from its key (the primary full-key hash plus the m leave-one-out hashes), so we mark all slots that actually hold this key. (Collisions may have blocked some redundant writes — those slots hold a different key and are not marked.) """ visited = np.zeros(self.size, dtype=np.bool_) items = [] idx = np.nonzero(self.used)[0] for addr in idx: if visited[addr]: continue k = self.keys[addr] c = int(self.counts[addr]) for a in self._slots_of(k): if self.used[a] and np.array_equal(self.keys[a], k): visited[a] = True if c >= min_count: items.append((k.copy(), int(self.vals[addr]), c)) items.sort(key=lambda it: it[2], reverse=True) if cap is not None: items = items[:cap] return items