"""MUSHROOM BODY — память нормального тоннеля и детектор новизны. Схема взята из коннектома напрямую и почти не изменена: * **PN → KC.** Около 2000 клеток Кеньона на полушарие; каждая получает вход всего от ~6 проекционных нейронов, выбранных почти случайно (в FlyWire и hemibrain среднее число «когтей» ≈ 5–7). Это случайная разрежённая проекция, поднимающая размерность примерно в 40 раз. * **APL.** Один гигантский тормозный нейрон собирает активность всех KC и возвращает торможение всем сразу. Результат — «победитель забирает всё»: одновременно активны около 5 % клеток Кеньона. * **KC → MBON.** Выход считывается немногими выходными нейронами. Синапсы KC→MBON **депрессируются** при повторном предъявлении стимула, поэтому знакомый запах даёт слабый ответ, а новый — сильный. Именно так устроена «familiarity suppression» у MBON-α′3. Почему это точно наша задача. Разметки нет, и почти все данные — пустой тоннель. Геометрический слой честно сообщает обо всём, что торчит в габарит, и вместе с препятствиями выдаёт кабельные лотки, ниши, гермозатворы, кромки платформ и стрелочные приводы. Грибовидное тело обучается на пустых проездах **без единой метки** и после этого молча гасит всё, что тоннель уже показывал, оставляя громкими только незнакомые формы. Хранится не образ, а битовый код: поиск знакомости — это индексация массива, поэтому стадия стоит микросекунды и легко укладывается в реальное время. """ from __future__ import annotations from dataclasses import dataclass from pathlib import Path import numpy as np from .lobula import Candidate # имена признаков-«проекционных нейронов»; порядок фиксирован и входит в артефакт FEATURES = ( "log_d", "lat", "abs_lat", "h", "log_w", "log_h", "aspect", "log_rays", "fill", "ang_w", "ang_h", "log_gap", "shadow", "inten", "el", "floor_def", "rings", "cols", "log_depth", "elongation", "containment", "h_min", ) def describe(c: Candidate) -> np.ndarray: """Вектор признаков кандидата — активность «проекционных нейронов». Сознательно смешаны две группы: форма и угловой размер (обобщаются на любой тоннель) и положение в сечении (именно оно делает лоток лотком). Первая даёт переносимость, вторая — способность запомнить штатную обстановку. """ d = max(c.d, 1.0) w = max(c.width, 0.02) hh = max(c.height, 0.02) area = max((c.bbox[1] - c.bbox[0] + 1) * (c.bbox[3] - c.bbox[2] + 1), 1) return np.array([ np.log(d), c.u, abs(c.u), c.h, np.log(w), np.log(hh), np.log(w / hh), np.log(max(c.n_rays, 1)), c.n_rays / area, np.degrees(w / d), np.degrees(hh / d), np.log1p(max(c.gap, 0.0)), c.shadow, np.log1p(max(c.inten, 0.0)), c.el_deg, np.log1p(max(c.floor_deficit, 0.0)), np.log(max(c.n_rings, 1)), np.log(max(c.n_cols, 1)), np.log1p(max(c.depth, 0.0)), np.log(max(c.depth, 0.05) / max(w, hh)), # вытянутость вдоль пути c.containment, c.h_min, # есть ли опора снизу ], dtype=np.float32) @dataclass class MushroomBodyConfig: """Параметры памяти. `claws = 6` взято прямо из коннектома: клетка Кеньона получает в среднем 5–7 входов от проекционных нейронов. Это значение и оказалось лучшим в подборе (`tools/tune_memory.py`): 4 когтя дают AUC 0.76, 8 — 0.89, 10 — 0.87, шесть — 0.89…0.90. Число клеток Кеньона, напротив, увеличено против мушиных ~2000. Причина инженерная: муха за жизнь встречает сотни запахов, а нам нужно запомнить десятки тысяч видов тоннельной обстановки. При 2000 клетках и 5 % активных память насыщается после нескольких тысяч примеров — подавлены все синапсы, и новым не выглядит уже ничто, включая настоящее препятствие. Разрежённость поэтому тоже снижена: активны 50 клеток из 50 000. """ n_kc: int = 50_000 # клеток Кеньона claws: int = 6 # входов на клетку Кеньона (из коннектома) sparsity: float = 0.001 # доля активных после торможения APL seed: int = 20240617 class MushroomBody: """Разрежённое кодирование + выученная знакомость.""" def __init__(self, cfg: MushroomBodyConfig | None = None, n_pn: int | None = None): self.cfg = cfg or MushroomBodyConfig() self.n_pn = n_pn or len(FEATURES) self.n_active = max(1, int(round(self.cfg.n_kc * self.cfg.sparsity))) rng = np.random.default_rng(self.cfg.seed) # разрежённая случайная проекция PN → KC: по `claws` входов на клетку idx = np.stack([rng.choice(self.n_pn, size=self.cfg.claws, replace=False) for _ in range(self.cfg.n_kc)]) sign = rng.choice((-1.0, 1.0), size=idx.shape) w = np.zeros((self.cfg.n_kc, self.n_pn), np.float32) np.put_along_axis(w, idx, sign.astype(np.float32), axis=1) self.W = w # нормировка входа (аналог дивизивной нормировки на общую интенсивность) self.mean = np.zeros(self.n_pn, np.float32) self.scale = np.ones(self.n_pn, np.float32) # вес синапсов KC → MBON: 1 — «не видели», 0 — «совсем привычно» self.w_mbon = np.ones(self.cfg.n_kc, np.float32) self.n_seen = 0 # ------------------------------------------------------------------ кодирование def fit_normalizer(self, X: np.ndarray) -> None: self.mean = X.mean(0).astype(np.float32) s = X.std(0).astype(np.float32) self.scale = np.where(s > 1e-6, s, 1.0).astype(np.float32) def encode(self, X: np.ndarray, device: str | None = None) -> np.ndarray: """Вектора признаков → индексы активных клеток Кеньона, (N, n_active). `device="cuda"` переносит проекцию и отбор на GPU: обучение на сотнях тысяч примеров так идёт в десятки раз быстрее. В инференсе кандидатов единицы, и путь через CPU заведомо дешевле запуска ядра. """ X = np.atleast_2d(np.asarray(X, np.float32)) if X.shape[1] != self.n_pn: raise ValueError( f"память обучена на {self.n_pn} признаках, а дескриптор даёт " f"{X.shape[1]}: набор признаков менялся — переобучите память " f"(tools/train_mushroom_body.py)") k = self.n_active # матрица откликов (N × n_kc) при большой выборке не помещается в память # целиком, поэтому кодирование идёт порциями фиксированного объёма chunk = max(1, int(2 ** 26 // max(self.cfg.n_kc, 1))) target_dev = device if target_dev == "auto": from .device import get_device target_dev = get_device("auto") if target_dev and target_dev != "cpu": try: import torch out = np.empty((X.shape[0], k), np.int64) with torch.no_grad(): m = torch.as_tensor(self.mean, device=target_dev) s = torch.as_tensor(self.scale, device=target_dev) w = torch.as_tensor(self.W, device=target_dev).T.contiguous() for i in range(0, X.shape[0], chunk): t = torch.as_tensor(X[i:i + chunk], device=target_dev) y = ((t - m) / s) @ w out[i:i + chunk] = torch.topk(y, k, dim=1).indices.cpu().numpy() return out except Exception as e: from .device import notify_cuda_error notify_cuda_error(e) # Fallback на CPU ниже if X.shape[0] <= chunk: z = (X - self.mean) / self.scale y = z @ self.W.T # (N, n_kc) # APL: глобальное торможение оставляет только сильнейшие return np.argpartition(-y, k - 1, axis=1)[:, :k] out = np.empty((X.shape[0], k), np.int64) for i in range(0, X.shape[0], chunk): z = (X[i:i + chunk] - self.mean) / self.scale y = z @ self.W.T out[i:i + chunk] = np.argpartition(-y, k - 1, axis=1)[:, :k] return out # ------------------------------------------------------------------ обучение def auto_rate(self, n_samples: int, target: float = 1.5) -> float: """Темп депрессии, согласованный с размером обучающей выборки. Каждый пример депрессирует `n_active` клеток, поэтому на одну клетку Кеньона в среднем приходится `n · n_active / n_kc` попаданий. Если темп не уменьшать вместе с ростом выборки, после нескольких десятков тысяч примеров подавлены все синапсы и новым не выглядит уже ничто. Темп выбирается так, чтобы «типичная» клетка ослабла примерно в e^target раз. Тогда шкала новизны отражает **частоту** обстановки: постоянно встречающиеся лотки и ниши уходят в ноль, а редкая форма остаётся заметной. Это и есть смысл familiarity suppression, а не «видел хоть раз». """ mean_hits = max(n_samples * self.n_active / self.cfg.n_kc, 1e-6) return float(np.clip(target / mean_hits, 1e-4, 0.5)) def learn(self, X: np.ndarray, rate: float = 0.35, device: str | None = None) -> None: """Депрессия синапсов KC→MBON на предъявленных примерах. Каждое предъявление умножает вес активных синапсов на (1 − rate), что для набора сразу эквивалентно возведению в степень по числу попаданий, — поэтому цикл по примерам не нужен. """ target_dev = device or "cpu" # без явной просьбы — процессор, как раньше if target_dev == "auto": from .device import get_device target_dev = get_device("auto") if target_dev.startswith("cuda"): try: import torch dev = torch.device(target_dev) k = self.n_active chunk = max(1, int(2 ** 26 // max(self.cfg.n_kc, 1))) m = torch.as_tensor(self.mean, device=dev) s = torch.as_tensor(self.scale, device=dev) w = torch.as_tensor(self.W, device=dev).T.contiguous() act_list = [] with torch.no_grad(): for i in range(0, X.shape[0], chunk): t = torch.as_tensor(X[i:i + chunk], device=dev) y = ((t - m) / s) @ w top_idx = torch.topk(y, k, dim=1).indices act_list.append(top_idx) act_t = torch.cat(act_list, dim=0) if len(act_list) > 1 else act_list[0] cnt_t = torch.bincount(act_t.reshape(-1), minlength=self.cfg.n_kc) w_mbon_t = torch.as_tensor(self.w_mbon, device=dev) w_mbon_t *= torch.pow(torch.tensor(1.0 - rate, device=dev), cnt_t.float()) self.w_mbon = w_mbon_t.cpu().numpy() self.n_seen += X.shape[0] return except Exception as e: from .device import notify_cuda_error notify_cuda_error(e) act = self.encode(X, device="cpu") cnt = np.bincount(act.ravel(), minlength=self.cfg.n_kc) self.w_mbon *= np.power(1.0 - rate, cnt).astype(np.float32) self.n_seen += len(act) def novelty(self, X: np.ndarray) -> np.ndarray: """Ответ MBON: 1 — совершенно ново, 0 — вполне привычно.""" act = self.encode(X) return self.w_mbon[act].mean(axis=1) def novelty_of(self, c: Candidate) -> float: return float(self.novelty(describe(c)[None, :])[0]) def annotate(self, cands: list[Candidate]) -> list[Candidate]: if not cands: return cands X = np.stack([describe(c) for c in cands]) nov = self.novelty(X) for c, n in zip(cands, nov): c.novelty = float(n) return cands # ------------------------------------------------------------------ хранение def save(self, path: str | Path) -> None: np.savez_compressed(path, w_mbon=self.w_mbon, mean=self.mean, scale=self.scale, n_kc=self.cfg.n_kc, claws=self.cfg.claws, sparsity=self.cfg.sparsity, seed=self.cfg.seed, n_pn=self.n_pn, n_seen=self.n_seen, features=np.array(FEATURES)) @staticmethod def load(path: str | Path) -> "MushroomBody": d = np.load(path, allow_pickle=False) cfg = MushroomBodyConfig(n_kc=int(d["n_kc"]), claws=int(d["claws"]), sparsity=float(d["sparsity"]), seed=int(d["seed"])) mb = MushroomBody(cfg, n_pn=int(d["n_pn"])) mb.w_mbon = d["w_mbon"].astype(np.float32) mb.mean = d["mean"].astype(np.float32) mb.scale = d["scale"].astype(np.float32) mb.n_seen = int(d["n_seen"]) return mb # Признаки для привыкания. Отличие от `FEATURES` принципиальное: здесь нет # ничего, что меняется при подъезде к одному и тому же предмету. Ни дальности, # ни числа лучей, ни углового размера, ни интенсивности (она падает как 1/R²), # ни контраста в метрах (он растёт с дальностью). Остаются физические свойства # формы: где она стоит в сечении, какая она и есть ли под ней опора. # # Иначе механизм не работает вовсе: код одного и того же кронштейна, снятого с # 90 и с 40 м, разъезжается, повторы не узнаются, а подъезжающий предмет, # наоборот, каждый кадр выглядит новой формой и привыкает сам к себе. Замерено: # на полном наборе признаков за двенадцать разных мест новизна падала с 1.000 # до 0.981 — то есть ни на что. # # У мухи ту же роль играет инвариантность к концентрации запаха: APL нормирует # общий уровень возбуждения, и один и тот же запах слабее или сильнее даёт # почти один и тот же набор активных клеток Кеньона. HAB_FEATURES = ("lat", "abs_lat", "h", "h_min", "log_w", "log_h", "aspect", "log_depth", "elongation", "containment") def describe_shape(c: Candidate) -> np.ndarray: """Вектор формы, не зависящий от дальности до неё.""" w = max(c.width, 0.02) hh = max(c.height, 0.02) return np.array([ c.u, abs(c.u), c.h, c.h_min, np.log(w), np.log(hh), np.log(w / hh), np.log1p(max(c.depth, 0.0)), np.log(max(c.depth, 0.05) / max(w, hh)), c.containment, ], dtype=np.float32) @dataclass class HabituationConfig: """Параметры привыкания внутри одного проезда. `n_kc` намеренно маленькое — 4000 против 50 000 у долговременной памяти, то есть примерно как у самой мухи. Задача здесь обратная: не различить десятки тысяч видов обстановки, а **обобщить** — чтобы седьмой кронштейн гасился первыми шестью, хотя он и не в точности такой же. Чем меньше популяция и чем шире разрежённость, тем сильнее перекрываются коды похожих форм. """ n_kc: int = 4000 claws: int = 4 sparsity: float = 0.02 # 80 активных клеток seed: int = 20260214 rate: float = 0.25 # депрессия за одно НОВОЕ место place_m: float = 5.0 # насколько разнести места, чтобы считать их разными same_frac: float = 0.30 # доля клеток «с этого же места» — значит, то же самое read_q: float = 0.75 # квантиль отсчёта MBON: доля клеток, которых мало quant: float = 0.5 # шаг огрубления признаков, в долях разброса recover_m: float = 800.0 # путь, за который привыкание сходит наполовину warmup: int = 60 # кандидатов на разогрев нормировки norm_n: int = 300 # после стольких нормировка замирает min_novelty: float = 0.05 # ниже этого привыкание не опускает class Habituation: """Кратковременное привыкание: гасит то, что повторяется в РАЗНЫХ местах пути. Обученная память решает задачу «этот тоннель я уже видел». На новом участке она бесполезна по определению: новизна у всего максимальна, улика копится беспрепятственно, и ложных тревог становится 24.3 на километр вместо 9.1. А приватный тест — это как раз новый участок. Но даже на совершенно незнакомой линии есть признак, не требующий ни единого прошлого проезда. **Тоннельная обстановка повторяется вдоль пути, а посторонний предмет — нет.** Кабельный кронштейн, рама крепи, стык тюбингов встречаются каждые несколько метров в одном и том же виде; упавший предмет лежит в одном месте. Поэтому привыкание считается не по времени и не по числу кадров, а по **числу разных точек пути**, где встретилась эта форма. Разница принципиальна. Настоящий предмет виден сто кадров подряд, но всё это время он стоит в одной точке мира: его код депрессируется один раз и остаётся новым до конца подъезда. Кронштейны за те же сто кадров дают десятки разных точек и гаснут. Биологически это та же схема, только с другим учителем. Синапсы KC→MBON депрессируются при повторном предъявлении — это и есть familiarity suppression MBON-α′3. Момент депрессии здесь разрешает координата пути из центрального комплекса: у мухи дофаминергические PPL1/PAM точно так же решают, когда пластичности быть, и их вход зависит от состояния животного, а не только от стимула. Привыкание постепенно сходит (`recover_m`): выехав с однообразного перегона на станцию, система снова смотрит в оба. """ def __init__(self, cfg: HabituationConfig | None = None): self.cfg = cfg or HabituationConfig() self.enc = MushroomBody( MushroomBodyConfig(n_kc=self.cfg.n_kc, claws=self.cfg.claws, sparsity=self.cfg.sparsity, seed=self.cfg.seed), n_pn=len(HAB_FEATURES)) n = self.cfg.n_kc self.w = np.ones(n, np.float32) # где в последний раз эта клетка Кеньона участвовала в наблюдении self.last_s = np.full(n, -1e9, np.float32) self.places = 0 self._n = 0 self._sum = np.zeros(len(HAB_FEATURES), np.float64) self._sq = np.zeros(len(HAB_FEATURES), np.float64) @property def level(self) -> float: """Насколько привыкла популяция: 0 — всё ново, 1 — всё знакомо.""" return float(1.0 - self.w.mean()) # ------------------------------------------------------------------ такт def advance(self, ds: float) -> None: """Забывание по пройденному пути.""" if ds <= 0 or self.cfg.recover_m <= 0: return k = 1.0 - 0.5 ** (ds / self.cfg.recover_m) self.w += (1.0 - self.w) * k def _fit(self, X: np.ndarray) -> None: """Дивизивная нормировка по накопленной статистике проезда. После `norm_n` кандидатов замирает, и это не мелочь. Пока нормировка плывёт, вместе с ней плывёт и код: у неподвижного предмета, чьи признаки формы вообще не меняются, набор активных клеток обновлялся настолько, что предмет засчитывался как двадцать шесть разных мест и гасил сам себя. С замороженной нормировкой одна и та же форма даёт один и тот же код, и правило «то же место — не повтор» наконец работает. Порог низкий намеренно. При 4000 в бэге `roundT_doubleT` (всего около 570 кандидатов за проезд) нормировка не замирала вовсе: привыкание осталось на 0.20, медианная новизна кандидата 1.00 и ложных треков ровно столько же, сколько без привыкания. Трёхсот кандидатов на десять признаков для оценки среднего и разброса достаточно с запасом. """ self._n += X.shape[0] if self._n > self.cfg.norm_n: return self._sum += X.sum(0) self._sq += (X.astype(np.float64) ** 2).sum(0) m = self._sum / self._n v = np.maximum(self._sq / self._n - m * m, 1e-12) self.enc.mean = m.astype(np.float32) self.enc.scale = np.sqrt(v).astype(np.float32) def _coarse(self, X: np.ndarray) -> np.ndarray: """Огрубить признаки до шага `quant` разброса — до кодирования. Без этого механизм не работает вовсе. Два кронштейна, снятые в разных местах, отличаются на сантиметры, случайная проекция разносит их в разные наборы клеток, и второй не узнаёт первого. Замер: при ёмкости, достаточной чтобы не насыщаться (20 000 клеток), привыкание давало ровно ноль — 7.5 ложных трека на км против тех же 7.5 без него. Весь видимый эффект маленькой популяции был не узнаванием повторов, а насыщением, которое одинаково давит и обстановку, и предмет. Огрубление до половины разброса заставляет близкие формы давать буквально один и тот же код. Это coarse coding, и у мухи оно тоже есть: клетка Кеньона реагирует на широкую область пространства запахов, а не на точку. """ q = self.cfg.quant if q <= 0: return X z = (X - self.enc.mean) / self.enc.scale return (np.round(z / q) * q) * self.enc.scale + self.enc.mean def update(self, cands: list[Candidate], s_path: float, ds: float) -> None: """Приглушить знакомое и запомнить то, что встретилось в новом месте. `s_path` — пройденный путь, м; положение формы в мире считается как `s_path + c.d`, поэтому неподвижный предмет остаётся в одной точке, сколько бы кадров к нему ни ехали. Оценка собственного движения копит около 3 % ошибки, то есть за подъезд с сотни метров точка уползает на два-три метра — отсюда и `place_m`: меньший разнос считается тем же местом, больший — новым. """ self.advance(ds) if not cands: return X = np.stack([describe_shape(c) for c in cands]) self._fit(X) if self._n < self.cfg.warmup: return # нормировка ещё не осмысленна act = self.enc.encode(self._coarse(X)) # Отсчёт берётся КВАНТИЛЕМ, а не средним, и это решает главную проблему # механизма. Коды похожих форм перекрываются частично: кронштейн, гасясь # сам, задевает часть клеток предмета. При среднем половина подавленных # клеток уже вдвое снижает новизну предмета — замерено, что новизна # вставленного человека падала с 0.76 до 0.20, а обнаружение на 70–90 м # со 100 % до 58 %. Квантиль 0.75 требует, чтобы подавлено было не менее # трёх четвертей клеток кода: частичное перекрытие форму уже не гасит, # а настоящий повтор гасит полностью. fam = np.quantile(self.w[act], self.cfg.read_q, axis=1) lo = self.cfg.min_novelty for c, f in zip(cands, fam): # Берётся минимум, а не произведение: два канала знакомости, и # достаточно одного, чтобы приглушить. Перемножение давило бы # предмет дважды за одно и то же. c.novelty = float(max(min(c.novelty, f), lo)) keep = 1.0 - self.cfg.rate for i, c in enumerate(cands): s_obj = np.float32(s_path + c.d) a = act[i] same = np.abs(self.last_s[a] - s_obj) <= self.cfg.place_m # Отметка ставится всегда, в том числе клеткам, впервые попавшим в # код: иначе подъезжающий предмет каждый кадр вербует свежие клетки # и депрессирует сам себя. self.last_s[a] = s_obj if same.mean() >= self.cfg.same_frac: continue # то же самое место — это не повтор self.w[a] *= keep self.places += 1 class NoMemory: """Заглушка: пока память не обучена, всё считается новым.""" n_seen = 0 @staticmethod def annotate(cands: list[Candidate]) -> list[Candidate]: return cands @staticmethod def novelty_of(_c: Candidate) -> float: return 1.0