diff --git a/Makefile b/Makefile index 0529b80..47f245b 100644 --- a/Makefile +++ b/Makefile @@ -1,11 +1,14 @@ .PHONY: help install install-all install-dev test lint docstyle typecheck check build clean \ - examples example-basic example-connectome example-navigation \ + examples example-basic example-connectome example-navigation example-gui \ + example-gui-smoke figures \ cli-list cli-simulate cli-simulate-offline cli-train cli-train-offline \ test-neuprint test-all PYTHON := .venv/bin/python PIP := .venv/bin/pip PYTEST := $(PYTHON) -m pytest +GUI_ARGS ?= +FIGURES_ARGS ?= SRC := src/flynet TESTS := tests @@ -15,7 +18,7 @@ TESTS := tests export NEUPRINT_APPLICATION_CREDENTIALS help: ## Show this help - @grep -E '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf " \033[36m%-24s\033[0m %s\n", $$1, $$2}' + @grep -hE '^[a-zA-Z_-]+:.*?## .*$$' $(MAKEFILE_LIST) | awk 'BEGIN {FS = ":.*?## "}; {printf " \033[36m%-24s\033[0m %s\n", $$1, $$2}' # ── Setup ────────────────────────────────────────────────────────────────── @@ -81,6 +84,15 @@ example-connectome: ## Run connectome simulation (offline) example-navigation: ## Run navigation task (offline, 5 episodes) $(PYTHON) examples/navigation_task.py --episodes 5 --offline +example-gui: ## Open the interactive foraging GUI (real brain; needs token or warm cache) + $(PYTHON) examples/fly_foraging_gui.py $(GUI_ARGS) + +example-gui-smoke: ## Headless foraging GUI check on the synthetic brain (no token) + $(PYTHON) examples/fly_foraging_gui.py --synthetic --smoke --smoke-frames 200 --speed 8 + +figures: ## Regenerate the README figures into docs/images (needs real brain) + $(PYTHON) examples/make_figures.py $(FIGURES_ARGS) + # ── CLI ──────────────────────────────────────────────────────────────────── cli-list: ## List available connectome datasets diff --git a/README.md b/README.md index 2b6777a..76b2818 100644 --- a/README.md +++ b/README.md @@ -89,6 +89,8 @@ Run `make help` to see all targets: example-basic Run basic network example example-connectome Run connectome simulation (offline) example-navigation Run navigation task (offline, 5 episodes) + example-gui Open the interactive foraging GUI (real brain) + example-gui-smoke Headless foraging GUI check (synthetic, no token) cli-list List available connectome datasets cli-simulate-offline Simulate with synthetic brain (no token needed) cli-simulate Simulate with real connectome (needs token) @@ -110,7 +112,7 @@ Run `make help` to see all targets: | `inhibition` | Winner-takes-all and soft lateral inhibition, dynamic thresholds | | `network` | `SpikingNetwork` -- rate-based settling dynamics, calibration, sensor-to-motor steering | | `connectome` | `ConnectomeLoader` -- neuPrint fetch, disk caching, `ConnectomeUnavailableError` on fetch failure (never a silent synthetic fallback); synthetic only via `synthetic_brain()` / CLI `--offline` | -| `learning` | `RewardHebbian` (three-factor) and `STDPTrainer` training pipelines | +| `learning` | `RewardHebbian` and `MotorRewardHebbian` (three-factor, signed rewards) plus `STDPTrainer` pipelines | | `jax_network` | `JaxNetwork` -- differentiable FLYNN-style (arXiv 2607.00025) recurrent dynamics on the connectome | | `jax_trainer` | `GradientTrainer` (DAgger-style BPTT imitation), `RLTrainer` (REINFORCE), `ExpertTeacher` | | `visualize` | Brain circuits, trajectories, learning curves, spike rasters, weight heatmaps | @@ -131,7 +133,7 @@ gradient extra: from flynet.connectome import ConnectomeLoader from flynet.network import SpikingNetwork -loader = ConnectomeLoader(dataset="male-cns:v1.0") +loader = ConnectomeLoader(dataset="male-cns:v0.9") # Small brain slice (~41 neurons around a cell type) ids, edges, motors = loader.load_or_fetch_mini(neuron_type="DNge104") @@ -223,13 +225,48 @@ The `examples/` directory contains complete scripts: | `connectome_simulation.py` | Load real connectome and query wiring | Optional | | `navigation_task.py` | Fly navigates to food using real brain wiring + Hebbian learning | Optional | | `classification.py` | Unsupervised STDP pattern learning | No | +| `fly_foraging_gui.py` | **Interactive GUI**: a fly forages for food with sight + smell while its brain learns live | Yes (or `--synthetic`) | ```bash make example-basic # or: .venv/bin/python examples/basic_network.py make example-connectome # offline connectome demo make example-navigation # offline navigation with learning +make example-gui # interactive fly, real brain, live learning ``` +### Watch the fly learn (interactive GUI) + +`fly_foraging_gui.py` opens a window in which a fly walks a 2-D arena looking +for food. It senses the world with **sight** (a forward visual field; food +inside it is resolved onto the two eyes) and **smell** (an odor plume +sampled at the two antennae). Both are mixed into the brain's two sensor +channels, the real connectome settles, and the motor difference becomes a +turn. Every hunt that ends is fed back as reward, so the fly gets better +while you watch. + +```bash +make example-gui # full real brain (~66k neurons) +make example-gui GUI_ARGS=--mini # 41-neuron mini brain, much faster +.venv/bin/python examples/fly_foraging_gui.py --synthetic # no token needed +``` + +Controls: `[` / `]` change the steps-per-frame speed (a slider does the +same), `space` pauses, `r` restarts the current hunt, `b` resets the brain +back to the connectome's own weights (handy for seeing what learning bought +you), `q` quits. Raise the speed to learn faster: more steps per second +means more hunts per second, and the green success-rate curve and the red +steps curve on the right tell you whether it is working. + +Learning uses `MotorRewardHebbian`, which reinforces only the synapses +projecting into the motor neurons, scaled by how much each step actually +closed on the food. On the full male CNS that lifts success from 71% to 83% +over 120 hunts, measured on a fixed set of 24 held-out start conditions; the +average hunt length improves much less (121 -> 117 steps) and wanders by a few +steps between runs, so trust the success rate over the step count. Keep +`--eta` small (the default `0.05`): larger values learn faster at first and +then destabilise the steering. See +[Does It Actually Learn?](#does-it-actually-learn) for the measured figure. + ## Architecture ``` @@ -243,14 +280,14 @@ flynet/ inhibition.py LateralInhibition, ThresholdManager network.py SpikingNetwork (rate settle, calibrate, turn) connectome.py ConnectomeLoader (neuPrint fetch, disk cache, loud fetch errors) - learning.py RewardHebbian, STDPTrainer, TrainingLogger + learning.py RewardHebbian, MotorRewardHebbian, STDPTrainer, TrainingLogger jax_network.py JaxNetwork (differentiable FLYNN-style dynamics) jax_trainer.py GradientTrainer, RLTrainer, ExpertTeacher visualize.py plot_brain_circuit, plot_trajectory, etc. cli.py flynet CLI entry point ``` -**Total: ~5,680 lines** of Python across 14 modules, with 62 tests. +**Total: ~5,680 lines** of Python across 14 modules, with 76 tests. ## Results @@ -265,21 +302,56 @@ Red = sensors, Blue = motors, Grey = other neurons. Edges are proportional to sy ![Trajectory](docs/images/real_trajectory.png) -The fly starts at (15, 20) and navigates toward food at (80, 80) by sensing a smell gradient. -Left: before learning (77 steps). Right: after 5 episodes of Hebbian training (50 steps). +The fly starts at (15, 20) and navigates toward food at (80, 80) by sensing a smell +gradient. Both panels replay the **same episode seed**, so the difference is the +learned weights and nothing else: 199 steps before training, 97 steps after 5 +episodes of reward-gated Hebbian plasticity. ### Learning Curve ![Learning Curve](docs/images/real_learning_curve.png) -Steps to reach food decrease as reward-gated Hebbian plasticity strengthens co-active synapses. +Steps to food over 30 navigation episodes on the mini brain. The smoothed curve +falls from about 89 to about 67 steps; single hunts stay noisy, and the early +episodes include a 144-step outlier. + +### Interactive Foraging GUI + +![Foraging GUI](docs/images/gui_foraging.png) + +`fly_foraging_gui.py` on the full male CNS (66,657 neurons). The arena shows the +odor plume, the food, the fly's visual field and its path; the panel on the right +plots hunt length and success rate while the HUD reports the live sensor +readings, the motor command and the plasticity update that was just applied. + +### Does It Actually Learn? + +![Foraging learning](docs/images/gui_foraging_learning.png) + +Left: the live training signal, which is genuinely noisy. Right: the honest check +— the untrained and the trained connectome scored on the **same 24 held-out start +conditions** after 120 hunts. Success rises from 71% to 83% and the average hunt +shortens from 121 to 117 steps. Keep `--eta` small: at `0.25` the fly learns +faster at first and then destabilises. + +### Regenerating the figures + +Every image above is produced by `examples/make_figures.py`, which also prints the +measured numbers so the captions can be checked against them: + +```bash +make figures # all five (needs the real brain) +make figures FIGURES_ARGS="--only trajectory" # just one +``` ## Tested With +Measured on `male-cns:v0.9`; regenerate with `make figures`. + | Connectome | Neurons | Edges | Result | | --- | --- | --- | --- | -| DNge104 mini brain | 41 | 40 | 100% navigation success, learning reduces path 77 -> 43 steps | -| Full male CNS | 66,729 | 228,220 | Settles in <0.1s, 1,236 neurons activate, steering works | +| DNge104 mini brain | 41 | 40 | Navigation succeeds; 199 -> 97 steps after 5 Hebbian episodes | +| Full male CNS | 66,657 | 227,941 | Foraging success 71% -> 83% over 120 hunts; one settle iteration is a 228k-nnz mat-vec (~1 ms) | ## License diff --git a/docs/images/gui_foraging.png b/docs/images/gui_foraging.png new file mode 100644 index 0000000..1c158a0 Binary files /dev/null and b/docs/images/gui_foraging.png differ diff --git a/docs/images/gui_foraging_learning.png b/docs/images/gui_foraging_learning.png new file mode 100644 index 0000000..d74ae85 Binary files /dev/null and b/docs/images/gui_foraging_learning.png differ diff --git a/docs/images/real_brain_circuit.png b/docs/images/real_brain_circuit.png index b36f3d5..de0411a 100644 Binary files a/docs/images/real_brain_circuit.png and b/docs/images/real_brain_circuit.png differ diff --git a/docs/images/real_learning_curve.png b/docs/images/real_learning_curve.png index fb9980b..2d4d41c 100644 Binary files a/docs/images/real_learning_curve.png and b/docs/images/real_learning_curve.png differ diff --git a/docs/images/real_trajectory.png b/docs/images/real_trajectory.png index 23831ca..5990b99 100644 Binary files a/docs/images/real_trajectory.png and b/docs/images/real_trajectory.png differ diff --git a/examples/connectome_simulation.py b/examples/connectome_simulation.py index 001a4d9..9b66514 100644 --- a/examples/connectome_simulation.py +++ b/examples/connectome_simulation.py @@ -29,7 +29,7 @@ def main() -> None: help="Load the cached real brain only (no network; error if not cached)", ) p.add_argument("--full", action="store_true", help="Use whole connectome") - p.add_argument("--dataset", default="male-cns:v1.0") + p.add_argument("--dataset", default="male-cns:v0.9") p.add_argument("--neuron-type", default="DNge104") p.add_argument("--max-edges", type=int, default=40) p.add_argument("--settle", type=int, default=12) diff --git a/examples/fly_foraging_gui.py b/examples/fly_foraging_gui.py new file mode 100644 index 0000000..32fe3f9 --- /dev/null +++ b/examples/fly_foraging_gui.py @@ -0,0 +1,1362 @@ +#!/usr/bin/env python3 +"""Interactive GUI: a fly forages for food while its brain keeps learning. + +The fly walks a 2-D arena driven by the **real** fly brain loaded from +neuPrint (the full ~66k-neuron male CNS by default). It has two senses: + +* **sight** -- a forward-facing visual field; food inside the field is + detected by contrast, strongest in the eye that faces it; +* **smell** -- an odor plume around the food, sampled at the two antennae. + +Both senses are mixed into the two sensor channels the brain exposes, the +brain settles, and the motor difference becomes a turn. Reward-gated +Hebbian plasticity (:class:`flynet.learning.MotorRewardHebbian`) reinforces the +co-active synapses of every successful hunt, so the fly visibly gets better +while the window is open. + +Run it:: + + python examples/fly_foraging_gui.py # full real brain + python examples/fly_foraging_gui.py --synthetic # no token needed + python examples/fly_foraging_gui.py --smoke # headless CI check + +Controls +-------- +``[`` / ``-`` slower ``]`` / ``+`` faster +``space`` pause / resume ``r`` restart the current hunt +``b`` reset the brain weights to the connectome's own +``q`` / ``esc`` quit +""" + +from __future__ import annotations + +import argparse +import time +from dataclasses import dataclass, field + +import numpy as np + +from flynet.connectome import ConnectomeLoader, ConnectomeUnavailableError +from flynet.learning import MotorRewardHebbian, TrainingLogger +from flynet.network import SpikingNetwork + +# --------------------------------------------------------------------------- +# Arena / sensing constants +# --------------------------------------------------------------------------- +ARENA = 100.0 +FOOD = (80.0, 80.0) +FOUND_RADIUS = 4.0 +MAX_STEPS = 300 +MIN_START_DIST = 25.0 +CURVE_HORIZON = 50 # initial episode axis; doubles as the run goes on + +ANTENNA_OFFSET = 2.7 # antenna tip distance from the body centre +STEP_LEN = 2.0 # walking distance per step while tracking +SEARCH_STEP = 2.5 # walking distance per step while searching +WANDER_TURN = 0.25 # random heading jitter while searching (rad) +SCAN_RATE = 1.1 # head-scan frequency (rad/step) +SCAN_SWEEP = 0.55 # head-scan amplitude (rad) +SEARCH_SWEEP = 0.10 # body turn while searching: wide arcs, not a spin + +FOV_DEG = 110.0 # total width of the visual field +SIGHT_RANGE = 35.0 # how far the fly can see food +SIGHT_GAIN = 1.8 # weight of sight vs smell in the sensor channels +SMELL_GAIN = 1.0 +SIGHT_NOISE = 0.06 # absolute sensory noise (contrast saturates at 1.0) +SMELL_NOISE = 0.012 +ODOR_THRESHOLD = 0.005 # odor strength that marks the plume as usable + +STEER_GAIN = 6.0 # motor difference -> turn angle +MAX_TURN = 0.45 # turn-angle clip per step (rad) + +SETTLE_ITERS_FULL = 6 # one settle iteration is a 228k-nnz mat-vec (~1ms) +SETTLE_ITERS_SMALL = 12 +ACTIVITY_TOPK = 256 # neurons kept per learning event +HEBBIAN_EVENTS = 24 # events per plasticity update (bounds the cost) +ROLLING_WINDOW = 20 # episodes in the rolling success/step average + +PROGRESS_GAIN = 4.0 # reward per unit of distance closed on the food +SUCCESS_BONUS = 3.0 # terminal reward for reaching the food +FAILURE_PENALTY = 0.0 # terminal reward for running out of steps +PROBE_SEED = 4242 # fixed seed so the steering probe is repeatable + +MIN_SPEED = 1 +MAX_SPEED = 16 +FLASH_FRAMES = 10 # how long the "FOUND FOOD!" banner stays up + + +# --------------------------------------------------------------------------- +# Sensing +# --------------------------------------------------------------------------- +def wrap_angle(angle: float) -> float: + """Wrap an angle into ``(-pi, pi]``. + + Args: + angle (float): Angle in radians. + + Returns: + float: The equivalent angle in ``(-pi, pi]``. + """ + return float((angle + np.pi) % (2.0 * np.pi) - np.pi) + + +def smell_at(pos: np.ndarray, food: tuple[float, float]) -> float: + """Return the odor concentration at *pos*. + + Args: + pos (np.ndarray): ``(x, y)`` position in the arena. + food (tuple[float, float]): ``(x, y)`` position of the food. + + Returns: + float: Odor strength, ``1.0`` at the food and decaying with the + squared distance. + """ + d = float(np.linalg.norm(np.asarray(food, dtype=float) - pos)) + return 2.0 / (1.0 + 0.3 * d * d) + + +def bearing_error( + pos: np.ndarray, heading: float, food: tuple[float, float] +) -> float: + """Signed angle from *heading* to the food. + + The fly's left is its forward direction rotated by +90 degrees, so a + positive result means the food lies off the fly's left side and a + negative one off its right. + + Args: + pos (np.ndarray): ``(x, y)`` position in the arena. + heading (float): Direction the fly faces, in radians. + food (tuple[float, float]): ``(x, y)`` position of the food. + + Returns: + float: Signed bearing error in ``(-pi, pi]``, positive to the left. + """ + to_food = np.asarray(food, dtype=float) - pos + return wrap_angle(float(np.arctan2(to_food[1], to_food[0])) - heading) + + +def antenna_readings( + pos: np.ndarray, heading: float, food: tuple[float, float] +) -> tuple[float, float]: + """Sample the odor plume at the two antenna tips. + + The antennae stick out at +/-45 degrees from the heading, so a plume + coming from the right biases the right antenna. + + Args: + pos (np.ndarray): ``(x, y)`` position in the arena. + heading (float): Direction the fly faces, in radians. + food (tuple[float, float]): ``(x, y)`` position of the food. + + Returns: + tuple[float, float]: ``(right, left)`` odor readings. + """ + d = np.array([np.cos(heading), np.sin(heading)]) + s = np.array([-d[1], d[0]]) # left-hand perpendicular + right_tip = pos + ANTENNA_OFFSET * (d - s) + left_tip = pos + ANTENNA_OFFSET * (d + s) + return smell_at(right_tip, food), smell_at(left_tip, food) + + +def visual_readings( + pos: np.ndarray, + head: float, + food: tuple[float, float], + fov_deg: float = FOV_DEG, + view_range: float = SIGHT_RANGE, +) -> tuple[float, float]: + """Sample the two eyes for the food inside the visual field. + + Food outside the field of view or beyond *view_range* is invisible. + Inside it, contrast falls off with distance and is shared between the + eyes according to which side of the midline the food sits on; food + straight ahead lights both eyes equally, which drives the fly forward. + A food off the fly's left brightens the left eye, and vice versa. + + Args: + pos (np.ndarray): ``(x, y)`` position in the arena. + head (float): Direction the head faces, in radians. + food (tuple[float, float]): ``(x, y)`` position of the food. + fov_deg (float): Total width of the visual field, in degrees. + view_range (float): Maximum distance at which food is visible. + + Returns: + tuple[float, float]: ``(right, left)`` contrast readings. + """ + err = bearing_error(pos, head, food) + dist = float(np.linalg.norm(np.asarray(food, dtype=float) - pos)) + half_fov = np.radians(fov_deg) / 2.0 + if dist > view_range or abs(err) > half_fov: + return 0.0, 0.0 + contrast = 1.0 / (1.0 + 0.01 * dist * dist) + # bearing_error is positive to the fly's left, so the right eye's + # share of the image shrinks as the food moves left. + right_share = 0.5 - 0.5 * err / half_fov + return contrast * right_share, contrast * (1.0 - right_share) + + +def topk_activity(trace: list[np.ndarray], k: int = ACTIVITY_TOPK) -> tuple: + """Compress a settling trace into a sparse top-k activity vector. + + The full brain settles sparsely, so keeping only the ``k`` strongest + neurons turns a 66k-element vector into a few hundred values. The + result is the ``(indices, values)`` form that + :meth:`flynet.learning.MotorRewardHebbian.update` understands. + + Args: + trace (list[np.ndarray]): Activity vectors, one per iteration. + k (int): Number of neurons to keep. + + Returns: + tuple: ``(indices, values)`` numpy arrays of equal length. + """ + if not trace: + empty_i = np.zeros(0, dtype=np.int64) + return empty_i, np.zeros(0, dtype=np.float64) + act = np.mean(trace[-2:], axis=0) + keep = int(min(k, act.size)) + idx = np.argpartition(np.abs(act), -keep)[-keep:] + return idx, act[idx] + + +def encode_readings( + smell: tuple[float, float], sight: tuple[float, float] +) -> tuple[float, float]: + """Turn the four sensory channels into the brain's two sensor channels. + + The brain's motor pair is far more sensitive to a left/right + *contrast* than to a symmetric drive: probing the real brain shows a + symmetric input produces a motor difference about eight times weaker + than an antisymmetric one of the same size. Feeding the raw levels + would therefore let a constant offset dominate the steering, so the + contrast between the two sides is what gets encoded, split so that + only the stronger side drives the network. Food straight ahead + encodes as no input at all, which walks the fly forward. + + Args: + smell (tuple[float, float]): ``(right, left)`` odor readings. + sight (tuple[float, float]): ``(right, left)`` contrast readings. + + Returns: + tuple[float, float]: ``(rho_right, rho_left)`` sensor channels. + """ + delta = SIGHT_GAIN * (sight[0] - sight[1]) + SMELL_GAIN * ( + smell[0] - smell[1] + ) + return max(delta, 0.0), max(-delta, 0.0) + + +# --------------------------------------------------------------------------- +# Episode / forager +# --------------------------------------------------------------------------- +@dataclass +class Episode: + """State of a single hunt for food. + + Attributes: + pos (np.ndarray): Current ``(x, y)`` position. + heading (float): Body direction in radians. + scan_phase (float): Head-scan oscillator phase in radians. + step (int): Steps taken so far. + found (bool): Whether the food was reached. + events (list): ``(rho_right, rho_left, activity, reward)`` learning + events, one per guided step. + path (list[np.ndarray]): Visited positions, oldest first. + mode (str): Sensory mode of the last step. + start_dist (float): Initial distance to the food. + progress (float): Reward accumulated by the guided steps. + """ + + pos: np.ndarray + heading: float + scan_phase: float = 0.0 + step: int = 0 + found: bool = False + events: list = field(default_factory=list) + path: list = field(default_factory=list) + mode: str = "search" + start_dist: float = 0.0 + progress: float = 0.0 + + +class FlyForager: + """Headless foraging simulation driven by a :class:`SpikingNetwork`. + + Owns the sensory mixing, the three behaviour modes (visual search, + odor tracking, and plume following) and the steering law. It never + touches matplotlib, so it can be unit-tested without a display. + """ + + def __init__( + self, + net: SpikingNetwork, + food: tuple[float, float] = FOOD, + arena: float = ARENA, + settle_iters: int = SETTLE_ITERS_FULL, + rng: np.random.Generator | None = None, + invert_steer: bool = False, + Minv: np.ndarray | None = None, + ) -> None: + """Build a forager on top of *net* and calibrate it. + + Args: + net (SpikingNetwork): Brain providing the steering command. + food (tuple[float, float]): ``(x, y)`` position of the food. + arena (float): Side length of the square arena. + settle_iters (int): Settling iterations per steering call. + rng (np.random.Generator | None): Random generator; a fresh one + is created when ``None``. + invert_steer (bool): Flip the turn direction, for brains whose + motor pair is wired the other way round. + Minv (np.ndarray | None): Frozen inverse calibration matrix. + Pass the initial one to keep steering comparable while the + weights are still moving; ``None`` recalibrates from *net*. + """ + self.net = net + self.food = food + self.arena = arena + self.settle_iters = settle_iters + self.rng = rng if rng is not None else np.random.default_rng(0) + self.invert_steer = invert_steer + + # One calibration, held for the whole run: turn() documents that a + # cached Minv keeps steering consistent while weights are updated. + # Recalibrating every frame would cancel out learning, because a + # 2x2 pinv always inverts whatever the current matrix is. + if Minv is None: + matrix, self.Minv = net.calibrate() + else: + matrix, _ = net.calibrate() + self.Minv = Minv + self.cond = float(np.linalg.cond(matrix)) + + self.episode = Episode( + pos=np.zeros(2), heading=0.0, path=[np.zeros(2)] + ) + self.distance = 0.0 + self.readings = (0.0, 0.0) + self.sight = (0.0, 0.0) + self.smell = (0.0, 0.0) + self.drv = 0.0 + self.step_count = 0 + self.reset() + + # -- episode control ------------------------------------------------ + def reset( + self, start: np.ndarray | None = None, heading: float | None = None + ) -> Episode: + """Start a new hunt from a random spot far enough from the food. + + Args: + start (np.ndarray | None): Explicit ``(x, y)`` start; a random + one is drawn when ``None``. + heading (float | None): Explicit body direction in radians; a + random one is drawn when ``None``. + + Returns: + Episode: The freshly reset episode. + """ + if start is None: + for _ in range(64): + cand = self.rng.uniform(4.0, self.arena - 4.0, size=2) + if np.linalg.norm(cand - np.asarray(self.food)) >= MIN_START_DIST: + start = cand + break + else: # pragma: no cover - arena too small to satisfy the rule + start = self.rng.uniform(4.0, self.arena - 4.0, size=2) + if heading is None: + heading = float(self.rng.uniform(-np.pi, np.pi)) + self.episode = Episode( + pos=np.asarray(start, dtype=float).copy(), + heading=float(heading), + scan_phase=0.0, + path=[np.asarray(start, dtype=float).copy()], + start_dist=float( + np.linalg.norm(np.asarray(start) - np.asarray(self.food)) + ), + ) + self.distance = self.episode.start_dist + self.drv = 0.0 + self.readings = (0.0, 0.0) + self.sight = (0.0, 0.0) + self.smell = (0.0, 0.0) + return self.episode + + # -- sensing -------------------------------------------------------- + def sense(self, pos: np.ndarray, heading: float) -> tuple[float, float, str]: + """Mix sight and smell into the brain's two sensor channels. + + Three regimes, chosen by what the fly can actually perceive: + + ``sight`` + the food is inside the visual field -- steer on the image; + ``smell`` + the plume is strong enough to follow -- steer on the odor + difference between the antennae; + ``search`` + nothing usable yet -- sweep the head and body while walking. + + Args: + pos (np.ndarray): ``(x, y)`` position in the arena. + heading (float): Body direction in radians. + + Returns: + tuple: ``(rho_right, rho_left, mode)``. + """ + ep = self.episode + ep.scan_phase += SCAN_RATE + head = heading + SCAN_SWEEP * float(np.sin(ep.scan_phase)) + + smell_r, smell_l = antenna_readings(pos, heading, self.food) + sight_r, sight_l = visual_readings(pos, head, self.food) + # Real senses are noisy, and the noise is what gives the learner + # something to improve on: early on the fly steers noisily, and + # reward-gated plasticity sharpens the motor projection. The + # noise floor is absolute, so uncertainty never vanishes. + if SIGHT_NOISE or SMELL_NOISE: + sight_r = max(0.0, sight_r + self.rng.normal(0.0, SIGHT_NOISE)) + sight_l = max(0.0, sight_l + self.rng.normal(0.0, SIGHT_NOISE)) + smell_r = max(0.0, smell_r + self.rng.normal(0.0, SMELL_NOISE)) + smell_l = max(0.0, smell_l + self.rng.normal(0.0, SMELL_NOISE)) + self.smell = (smell_r, smell_l) + self.sight = (sight_r, sight_l) + + rho_right, rho_left = encode_readings( + (smell_r, smell_l), (sight_r, sight_l) + ) + + if sight_r + sight_l > 1e-9: + mode = "sight" + elif max(smell_r, smell_l) >= ODOR_THRESHOLD: + mode = "smell" + else: + mode = "search" + return rho_right, rho_left, mode + + # -- one simulation step -------------------------------------------- + def step(self) -> Episode: + """Advance the hunt by one step. + + Guided steps (``sight`` / ``smell``) ask the brain for a steering + command and record a learning event; ``search`` steps sweep the head + and body so the fly scans the horizon instead. + + Returns: + Episode: The current episode state. + """ + ep = self.episode + if ep.found: + return ep + + pending: list = [] + rho_right, rho_left, mode = self.sense(ep.pos, ep.heading) + self.readings = (rho_right, rho_left) + ep.mode = mode + before = self.distance + + if mode == "search": + ep.heading = wrap_angle( + ep.heading + + SEARCH_SWEEP + + self.rng.uniform(-WANDER_TURN, WANDER_TURN) + ) + advance = SEARCH_STEP + else: + self.drv, trace = self.net.turn( + rho_right, rho_left, Minv=self.Minv, + iterations=self.settle_iters, + ) + sign = -1.0 if self.invert_steer else 1.0 + turn = float( + np.clip(sign * STEER_GAIN * self.drv, -MAX_TURN, MAX_TURN) + ) + ep.heading = wrap_angle(ep.heading + turn) + advance = STEP_LEN + pending.append((rho_right, rho_left, topk_activity(trace))) + + direction = np.array([np.cos(ep.heading), np.sin(ep.heading)]) + ep.pos = np.clip(ep.pos + direction * advance, 0.0, self.arena) + ep.path.append(ep.pos.copy()) + ep.step += 1 + self.step_count += 1 + self.distance = float(np.linalg.norm(ep.pos - np.asarray(self.food))) + + # Third factor: a guided step is rewarded by the progress it made + # towards the food and punished when it wandered further away. + if mode != "search": + reward = (before - self.distance) * PROGRESS_GAIN + ep.progress += reward + rho_r, rho_l, activity = pending.pop() + ep.events.append((rho_r, rho_l, activity, reward)) + + if self.distance < FOUND_RADIUS: + ep.found = True + return ep + + def auto_steer_sign(self, steps: int = 40) -> int: + """Find the turn direction that actually closes on the food. + + Which way the motor difference has to be driven to turn towards + the food depends on how the brain's motor pair happens to be + wired, so it is measured instead of assumed: two short probes + start the fly short of the food, one with the food off its left + and one off its right, and the sign that ends up closer wins. + + The probes do not learn, so *net* is used as-is. + + Args: + steps (int): Probe length in steps. + + Returns: + int: ``-1`` when the steering must be inverted, ``1`` + otherwise. ``1`` is returned when neither sign is clearly + better, so a degenerate brain falls back to the default. + """ + food = np.asarray(self.food, dtype=float) + probes = ( + np.array([food[0] - 20.0, food[1] - 15.0]), # food off the left + np.array([food[0] - 20.0, food[1] + 15.0]), # food off the right + ) + totals: dict[bool, float] = {} + for invert in (False, True): + probe = FlyForager( + self.net, + food=self.food, + arena=self.arena, + settle_iters=self.settle_iters, + rng=np.random.default_rng(PROBE_SEED), + invert_steer=invert, + Minv=self.Minv, + ) + total = 0.0 + for start in probes: + probe.reset(start=start, heading=0.0) + for _ in range(steps): + probe.step() + if probe.episode.found: + break + total += probe.distance + totals[invert] = total + if abs(totals[False] - totals[True]) < 1e-6: + return 1 + return -1 if totals[True] < totals[False] else 1 + + # -- learning ------------------------------------------------------- + def learning_batch( + self, limit: int = HEBBIAN_EVENTS, found: bool | None = None + ) -> list: + """Return recent guided steps as signed-reward learning events. + + The last *limit* events are returned: the plasticity update costs + one sparse pass over every synapse per event, so a long hunt would + otherwise freeze the window. When *found* is given, a terminal + bonus or penalty is added to the last event so the whole hunt -- + not just its individual steps -- carries the outcome. + + Args: + limit (int): Maximum number of events to return. + found (bool | None): Outcome of the hunt; ``None`` scores the + steps on progress alone. + + Returns: + list: ``(rho_right, rho_left, activity, reward)`` tuples. + """ + if limit <= 0: + return [] + batch = self.episode.events[-limit:] + if found is None or not batch: + return list(batch) + rho_r, rho_l, activity, reward = batch[-1] + terminal = SUCCESS_BONUS if found else FAILURE_PENALTY + batch = list(batch[:-1]) + [(rho_r, rho_l, activity, + reward + terminal)] + return batch + + +# --------------------------------------------------------------------------- +# Training metrics +# --------------------------------------------------------------------------- +class TrainingStats: + """Rolling success/step statistics plus a full learning curve. + + Wraps :class:`flynet.learning.TrainingLogger` so the plotted curve is + the same history the library's learning plots consume. + """ + + def __init__(self, window: int = ROLLING_WINDOW) -> None: + """Create empty statistics. + + Args: + window (int): Number of recent episodes used for the rolling + averages. + """ + self.window = window + self.logger = TrainingLogger() + self.episodes = 0 + self.successes = 0 + self.best_steps: int | None = None + self.history: list[tuple[bool, int]] = [] + self.recent: list[tuple[bool, int]] = [] + self.steps_done = 0 + + def record(self, found: bool, steps: int) -> None: + """Record the outcome of one hunt. + + Args: + found (bool): Whether the fly reached the food. + steps (int): Steps the hunt took. + + Returns: + None: Appends to the history and the rolling window. + """ + self.episodes += 1 + self.steps_done += steps + if found: + self.successes += 1 + if self.best_steps is None or steps < self.best_steps: + self.best_steps = steps + self.history.append((found, steps)) + self.recent.append((found, steps)) + if len(self.recent) > self.window: + self.recent.pop(0) + self.logger.log_episode(1.0 if found else 0.0, steps) + + def success_rate(self) -> float: + """Return the success rate over the rolling window. + + Returns: + float: Fraction of recent hunts that found the food, ``0.0`` + before the first hunt finishes. + """ + if not self.recent: + return 0.0 + return sum(1 for found, _ in self.recent if found) / len(self.recent) + + def avg_steps(self) -> float: + """Return the mean hunt length over the rolling window. + + Returns: + float: Mean steps per hunt, ``0.0`` before the first hunt. + """ + if not self.recent: + return 0.0 + return float(np.mean([steps for _, steps in self.recent])) + + def smoothed_steps(self, window: int = 5) -> tuple[list[int], list[float]]: + """Return a moving average of hunt length for plotting. + + Args: + window (int): Width of the moving average in episodes. + + Returns: + tuple[list[int], list[float]]: Episode numbers and smoothed + steps, both shorter than the full history. + """ + episodes, steps = self.logger.get_learning_curve() + if not steps: + return [], [] + arr = np.asarray(steps, dtype=float) + if len(arr) < window: + return episodes, list(arr) + kernel = np.ones(window) / window + smoothed = np.convolve(arr, kernel, mode="valid") + return list(episodes[window - 1:]), list(smoothed) + + def success_series(self) -> tuple[list[int], list[float]]: + """Return the rolling success rate after each episode. + + Returns: + tuple[list[int], list[float]]: Episode numbers and the success + rate over the trailing window, ready to plot. + """ + xs: list[int] = [] + ys: list[float] = [] + for i, (found, _) in enumerate(self.history, start=1): + chunk = self.history[max(0, i - self.window):i] + xs.append(i) + ys.append(sum(1 for f, _ in chunk if f) / len(chunk)) + return xs, ys + + +def apply_learning( + net: SpikingNetwork, + hebbian: MotorRewardHebbian, + events: list, +) -> bool: + """Apply one reward-gated Hebbian update in place. + + Args: + net (SpikingNetwork): Brain whose weights are updated. + hebbian (MotorRewardHebbian): Plasticity rule. + events (list): Learning events from a successful hunt. + + Returns: + bool: ``True`` when at least one event was applied. + """ + if not events: + return False + hebbian.update(net, events) + return True + + +# --------------------------------------------------------------------------- +# Brain loading +# --------------------------------------------------------------------------- +def load_brain(args: argparse.Namespace) -> tuple[SpikingNetwork, str, str]: + """Load the brain requested on the command line. + + Real data is the default and the only silent path. ``--offline`` reads + the cache only, and ``--synthetic`` must be asked for by name and is + reported loudly. + + Args: + args (argparse.Namespace): Parsed command-line arguments. + + Returns: + tuple: ``(net, banner, description)`` where *banner* is the + provenance line to print and *description* names the dataset. + + Raises: + SystemExit: If the requested real brain is unavailable. + """ + if args.synthetic: + ids, edges, motors = ConnectomeLoader.synthetic_brain( + max_edges=40, seed=args.seed + ) + return ( + SpikingNetwork(ids, edges, motor_ids=motors), + "[connectome] --synthetic: using SYNTHETIC brain", + "synthetic 41-neuron test brain", + ) + + loader = ConnectomeLoader(dataset=args.dataset) + try: + if args.offline: + if args.mini: + cached = loader.load_cached_mini("DNge104") + if cached is None: + raise SystemExit(_cold_cache_message(loader, "DNge104")) + ids, edges, motors = cached + else: + cached = loader.load_cached_full() + if cached is None: + raise SystemExit(_cold_cache_message(loader, args.dataset)) + ids, edges = cached + motors = [] + elif args.mini: + ids, edges, motors = loader.load_or_fetch_mini( + neuron_type="DNge104", max_edges=40 + ) + else: + ids, edges, motors = loader.load_or_fetch_full( + min_weight=args.min_weight + ) + except ConnectomeUnavailableError as exc: + print(f"error: {exc}", file=__import__("sys").stderr) + raise SystemExit(1) from exc + + source = loader.last_source + net = SpikingNetwork(ids, edges, motor_ids=motors if motors else None) + scope = "DNge104 mini brain" if args.mini else "full connectome" + return ( + net, + f"[connectome] data source: {source} ({scope}, {args.dataset})", + f"{scope} from {args.dataset}", + ) + + +def _cold_cache_message(loader: ConnectomeLoader, what: str) -> str: + """Build the error shown when ``--offline` finds no cached brain. + + Args: + loader (ConnectomeLoader): Loader whose cache directory is empty. + what (str): The dataset or neuron type that was requested. + + Returns: + str: The error message. + """ + return ( + f"offline: no cached real brain for {what!r} in {loader.cache_dir}. " + "Run once with neuPrint access to populate the cache. flynet never " + "substitutes synthetic data here; pass --synthetic to ask for the " + "test brain explicitly." + ) + + +# --------------------------------------------------------------------------- +# GUI +# --------------------------------------------------------------------------- +class ForagingApp: + """Matplotlib front end: draws the arena and drives the forager. + + The figure is built once and only artist data is updated per frame, so + the window stays responsive while the brain settles. + """ + + def __init__( + self, + forager: FlyForager, + stats: TrainingStats, + hebbian: MotorRewardHebbian, + banner: str, + description: str, + speed: int = 1, + ) -> None: + """Create the figure and its artists. + + Args: + forager (FlyForager): Simulation to drive. + stats (TrainingStats): Metrics to display. + hebbian (MotorRewardHebbian): Plasticity rule applied to each hunt. + banner (str): Provenance line for the title. + description (str): Human-readable brain description. + speed (int): Simulation steps per animation frame. + """ + import matplotlib.pyplot as plt + from matplotlib.patches import Wedge + + self.forager = forager + self.stats = stats + self.hebbian = hebbian + self.banner = banner + self.description = description + self.speed = int(np.clip(speed, MIN_SPEED, MAX_SPEED)) + self.paused = False + self.running = True + self.started = time.time() + self.last_learn = "start" + self.slider = None + self.anim = None + self.flash_frames = 0 + + self.fig = plt.figure(figsize=(13.0, 7.2)) + self.ax = self.fig.add_axes([0.035, 0.07, 0.55, 0.86]) + self.ax_curve = self.fig.add_axes([0.655, 0.70, 0.32, 0.23]) + self.ax_hud = self.fig.add_axes([0.655, 0.06, 0.32, 0.50]) + self.ax_hud.set_axis_off() + + # -- arena ------------------------------------------------------ + arena = forager.arena + grid = np.linspace(0.0, arena, 96) + gx, gy = np.meshgrid(grid, grid) + field = 2.0 / (1.0 + 0.3 * ((gx - forager.food[0]) ** 2 + + (gy - forager.food[1]) ** 2)) + # The plume decays as 1/d^2, so plot a compressed field: the + # physics still uses smell_at(), this is display only. + shown = (field / field.max()) ** 0.18 + self.ax.imshow( + shown, origin="lower", extent=(0, arena, 0, arena), + cmap="YlGn", alpha=0.45, zorder=0, + ) + self.ax.set_xlim(0, arena) + self.ax.set_ylim(0, arena) + self.ax.set_aspect("equal") + self.ax.set_xlabel("x") + self.ax.set_ylabel("y") + self.ax.set_title(self.banner, fontsize=9) + + self.ax.plot( + forager.food[0], forager.food[1], "*", color="gold", + markersize=22, markeredgecolor="black", zorder=6, + label="food", + ) + (self.trail,) = self.ax.plot([], [], "-", color="black", alpha=0.35, + linewidth=1.0, zorder=2) + self.fov = Wedge( + (0, 0), SIGHT_RANGE, 0, 1, facecolor="dodgerblue", alpha=0.13, + edgecolor="dodgerblue", zorder=1, + ) + self.ax.add_patch(self.fov) + (self.body,) = self.ax.plot([], [], "-", color="black", + linewidth=2.5, zorder=4) + self.ax.plot([], [], "o", color="black", markersize=4, zorder=4) + self.fly_marker = self.ax.lines[-1] + (self.antenna,) = self.ax.plot([], [], "-", color="darkgreen", + linewidth=1.4, zorder=4) + self.flash = self.ax.text( + 0.5, 0.94, "", transform=self.ax.transAxes, ha="center", + fontsize=15, fontweight="bold", color="crimson", zorder=8, + ) + + # -- learning curve --------------------------------------------- + # Fixed limits keep the tick labels static, which is what makes + # blitting safe: nothing outside the returned artists ever moves. + (self.curve_raw,) = self.ax_curve.plot( + [], [], "-", color="steelblue", alpha=0.35, label="steps" + ) + (self.curve_avg,) = self.ax_curve.plot( + [], [], "-", color="crimson", linewidth=2, label="smoothed" + ) + self.ax_curve.set_xlim(0, CURVE_HORIZON) + self.ax_curve.set_ylim(0, float(MAX_STEPS)) + self.horizon = CURVE_HORIZON + self.full_redraw = False + self.ax_curve.set_xlabel("episode", fontsize=8) + self.ax_curve.set_ylabel("steps to food", fontsize=8) + self.ax_curve.tick_params(labelsize=7) + self.ax_curve.set_title("Learning curve", fontsize=10) + self.ax_rate = self.ax_curve.twinx() + (self.curve_rate,) = self.ax_rate.plot( + [], [], "-", color="seagreen", linewidth=1.5, alpha=0.9 + ) + self.ax_rate.set_ylabel("success rate", fontsize=8, color="seagreen") + self.ax_rate.tick_params(labelsize=7, colors="seagreen") + self.ax_rate.set_ylim(-0.05, 1.05) + + self.hud = self.ax_hud.text( + 0.0, 1.0, "", va="top", family="monospace", fontsize=9.5, + transform=self.ax_hud.transAxes, + ) + + # -- interaction ---------------------------------------------------- + def add_controls(self) -> None: + """Attach the speed slider and keyboard bindings to the figure. + + Safe to call without a display: the slider is just artists, so + the documentation figures reuse this to show it. + + Returns: + None: Registers matplotlib callbacks on the figure canvas. + """ + from matplotlib.widgets import Slider + + self.fig.canvas.mpl_connect("key_press_event", self.on_key) + self.fig.canvas.mpl_connect("close_event", self.on_close) + + slider_ax = self.fig.add_axes([0.70, 0.605, 0.20, 0.022]) + self.slider = Slider( + slider_ax, "steps/frame", MIN_SPEED, MAX_SPEED, + valinit=self.speed, valstep=1, + ) + self.slider.on_changed(self.on_slide) + self.fig.canvas.draw_idle() + + def on_slide(self, value: float) -> None: + """Apply a new speed from the slider. + + Args: + value (float): Requested steps per frame. + + Returns: + None: Updates the simulation rate. + """ + self.set_speed(int(value)) + + def set_speed(self, speed: int) -> None: + """Clamp and apply a new steps-per-frame value. + + Args: + speed (int): Requested steps per frame. + + Returns: + None: Updates the simulation rate. + """ + self.speed = int(np.clip(speed, MIN_SPEED, MAX_SPEED)) + if self.slider is not None and int(self.slider.val) != self.speed: + self.slider.set_val(self.speed) + + def on_key(self, event) -> None: + """Handle a keypress in the figure window. + + Args: + event: The matplotlib key event. + + Returns: + None: Changes speed, pauses, restarts or quits. + """ + key = (event.key or "").lower() + if key in ("]", "+", "="): + self.set_speed(self.speed + 1) + elif key in ("[", "-"): + self.set_speed(self.speed - 1) + elif key == " ": + self.paused = not self.paused + elif key == "r": + self.forager.reset() + elif key == "b": + self.reset_brain() + elif key in ("q", "escape"): + self.running = False + + def on_close(self, _event) -> None: + """Stop the animation when the window is closed. + + Args: + _event: The matplotlib close event. + + Returns: + None: Stops the animation loop. + """ + self.running = False + + def reset_brain(self) -> None: + """Restore the connectome's own weights, discarding learning. + + Returns: + None: Replaces the brain with an untrained copy. + """ + self.forager.net = self.forager.net.copy() + self.forager = FlyForager( + self.forager.net, + food=self.forager.food, + arena=self.forager.arena, + settle_iters=self.forager.settle_iters, + rng=self.forager.rng, + invert_steer=self.forager.invert_steer, + ) + self.last_learn = "brain reset to connectome weights" + + # -- simulation loop ------------------------------------------------- + def finish_episode(self) -> None: + """Record the finished hunt and reinforce it when it succeeded. + + Returns: + None: Applies plasticity and starts the next hunt. + """ + ep = self.forager.episode + self.stats.record(ep.found, ep.step) + batch = self.forager.learning_batch(found=ep.found) + if apply_learning(self.forager.net, self.hebbian, batch): + self.last_learn = ( + f"{len(batch)} events, " + f"reward {ep.progress:+.1f}" + f"{' +bonus' if ep.found else ' -penalty'}" + ) + self.flash.set_text("FOUND FOOD!" if ep.found else "") + self.flash_frames = FLASH_FRAMES if ep.found else 0 + self.forager.reset() + + def tick(self) -> None: + """Advance the simulation by ``speed`` steps. + + Returns: + None: Runs the forager and closes out finished episodes. + """ + if self.paused: + return + for _ in range(self.speed): + if self.forager.episode.found or \ + self.forager.episode.step >= MAX_STEPS: + self.finish_episode() + self.forager.step() + if self.forager.episode.found or \ + self.forager.episode.step >= MAX_STEPS: + self.finish_episode() + + def on_frame(self, _frame) -> list: + """Animation callback: tick, update artists, and stop when asked. + + Returns the list of artists that changed so that + :class:`matplotlib.animation.FuncAnimation` can blit just those, + which is far cheaper than redrawing the whole figure. + """ + if not self.running: + return [] + self.tick() + self.draw() + if self.full_redraw: + # The axis limits just moved, so the cached blit background is + # stale: repaint everything once, then go back to blitting. + self.full_redraw = False + self.fig.canvas.draw() + return self.dirty_artists() + + def dirty_artists(self) -> list: + """Return every artist that :meth:`draw` updates. + + Returns: + list: Artists to re-blit, in draw order. + """ + return [ + self.trail, self.fov, self.body, self.fly_marker, self.antenna, + self.flash, self.curve_raw, self.curve_avg, self.curve_rate, + self.hud, + ] + + def draw(self) -> None: + """Refresh every artist from the current simulation state. + + Returns: + None: Updates arena, learning curve and HUD. Deliberately + does not draw the canvas: the animation callback blits. + """ + forager = self.forager + ep = forager.episode + pos = ep.pos + heading = ep.heading + head = heading + SCAN_SWEEP * float(np.sin(ep.scan_phase)) + + d = np.array([np.cos(heading), np.sin(heading)]) + s = np.array([-d[1], d[0]]) + self.body.set_data([pos[0] - 3 * d[0], pos[0] + 3 * d[0]], + [pos[1] - 3 * d[1], pos[1] + 3 * d[1]]) + self.fly_marker.set_data([pos[0] + 3 * d[0]], [pos[1] + 3 * d[1]]) + self.antenna.set_data( + [pos[0] + 2 * d[0] - 3 * s[0], pos[0] + 2 * d[0] + 3 * s[0]], + [pos[1] + 2 * d[1] - 3 * s[1], pos[1] + 2 * d[1] + 3 * s[1]], + ) + self.fov.set_center((float(pos[0]), float(pos[1]))) + self.fov.set_theta1(np.degrees(head) - FOV_DEG / 2.0) + self.fov.set_theta2(np.degrees(head) + FOV_DEG / 2.0) + path = ep.path + self.trail.set_data([p[0] for p in path], [p[1] for p in path]) + + episodes, steps = self.stats.logger.get_learning_curve() + self.curve_raw.set_data(episodes, steps) + sm_x, sm_y = self.stats.smoothed_steps() + self.curve_avg.set_data(sm_x, sm_y) + rate_x, rate_y = self.stats.success_series() + self.curve_rate.set_data(rate_x, rate_y) + + # Grow the episode axis in coarse steps so the curve stays + # readable without moving tick labels every frame (which would + # invalidate blitting). + if episodes and episodes[-1] > self.horizon * 0.9: + self.horizon *= 2 + self.ax_curve.set_xlim(0, self.horizon) + self.full_redraw = True + + elapsed = max(time.time() - self.started, 1e-9) + rate = self.forager.step_count / elapsed * 60.0 + n_active = len(self.forager.net.get_activity(0.05)) + mean_w = float(self.forager.net.W.data.mean()) + lines = [ + f"brain {len(self.forager.net.ids):,} neurons", + f" {self.forager.net.W.nnz:,} synapses", + f"dataset {self.description}", + "", + f"episode {self.stats.episodes}", + f"mode {ep.mode}", + f"distance {forager.distance:5.1f} (start {ep.start_dist:.0f})", + f"motor drv {forager.drv:+.3f}", + f"sight R/L {forager.sight[0]:.3f} / {forager.sight[1]:.3f}", + f"smell R/L {forager.smell[0]:.3f} / {forager.smell[1]:.3f}", + f"neurons on {n_active:,}", + f"mean syn. {mean_w:.5f}", + "", + f"success {self.stats.success_rate():.0%} (last " + f"{len(self.stats.recent)})", + f"avg steps {self.stats.avg_steps():.1f}", + f"best steps {self.stats.best_steps if self.stats.best_steps else '-'}", + f"hunts {self.stats.successes}/{self.stats.episodes}", + f"learn {self.last_learn}", + "", + f"speed {self.speed} steps/frame" + f"{' [PAUSED]' if self.paused else ''}", + f"steps sim {self.forager.step_count:,}", + f"rate {rate:,.0f} steps/min", + "", + "[ / ] speed space pause", + "r restart hunt b reset brain", + "q quit", + ] + if self.flash_frames > 0: + self.flash_frames -= 1 + if self.flash_frames == 0: + self.flash.set_text("") + self.hud.set_text("\n".join(lines)) + + def summary(self) -> str: + """Return a one-shot text summary of the run. + + Returns: + str: Episode counts, success rate, steps simulated and best + hunt length. + """ + best = self.stats.best_steps + return ( + f"steps={self.forager.step_count} " + f"episodes={self.stats.episodes} " + f"success={self.stats.successes} " + f"({self.stats.success_rate():.0%} rolling) " + f"best_steps={best if best else '-'}" + ) + + +# --------------------------------------------------------------------------- +# Entry point +# --------------------------------------------------------------------------- +def build_parser() -> argparse.ArgumentParser: + """Build the command-line parser. + + Returns: + argparse.ArgumentParser: Parser for the foraging GUI. + """ + p = argparse.ArgumentParser( + description="Watch a fly learn to find food on the real fly brain." + ) + p.add_argument( + "--dataset", default="male-cns:v0.9", + help="neuPrint dataset (default: %(default)s)", + ) + p.add_argument( + "--min-weight", type=int, default=50, + help="minimum synapse weight for the full connectome (default: %(default)s)", + ) + p.add_argument( + "--mini", action="store_true", + help="use the 41-neuron DNge104 mini brain instead of the full connectome", + ) + p.add_argument( + "--offline", action="store_true", + help="use the cached real brain only (error if the cache is cold)", + ) + p.add_argument( + "--synthetic", action="store_true", + help="use the explicit synthetic test brain (no token needed)", + ) + p.add_argument( + "--speed", type=int, default=1, + help="simulation steps per frame at startup (default: %(default)s)", + ) + p.add_argument( + "--eta", type=float, default=0.05, + help="Hebbian learning rate (default: %(default)s)", + ) + p.add_argument( + "--invert-steer", action="store_true", + help="flip the turn direction instead of detecting it", + ) + p.add_argument( + "--steer-sign", choices=["auto", "default"], default="auto", + help="detect the turn direction (default: %(default)s)", + ) + p.add_argument( + "--seed", type=int, default=3, help="random seed (default: %(default)s)" + ) + p.add_argument( + "--smoke", action="store_true", + help="headless run of --smoke-frames frames, then exit (for CI)", + ) + p.add_argument( + "--smoke-frames", type=int, default=240, + help="frames to run in --smoke mode (default: %(default)s)", + ) + p.add_argument( + "--save", default=None, + help="save the final frame to this PNG path", + ) + p.add_argument( + "--duration", type=float, default=0.0, + help="close the window automatically after N seconds (0 = never)", + ) + return p + + +def main() -> None: + """Run the foraging GUI. + + Raises: + SystemExit: If the requested real brain cannot be loaded. + """ + import matplotlib + + args = build_parser().parse_args() + if args.smoke: + matplotlib.use("Agg") + + net, banner, description = load_brain(args) + print(banner) + print(f"[gui] {description}: {len(net.ids):,} neurons, " + f"{net.W.nnz:,} synapses") + + settle = SETTLE_ITERS_SMALL if (args.mini or args.synthetic) \ + else SETTLE_ITERS_FULL + forager = FlyForager( + net, + settle_iters=settle, + rng=np.random.default_rng(args.seed), + invert_steer=args.invert_steer, + ) + print(f"[gui] sensors={net.sensors} motors={net.motors} " + f"calibration condition number={forager.cond:.2f}") + if forager.cond > 50: + print("[gui] warning: the sensor->motor matrix is poorly conditioned; " + "steering may be weak. Try --mini, or a different dataset.") + + if args.steer_sign == "auto" and not args.invert_steer: + invert = forager.auto_steer_sign() == -1 + forager.invert_steer = invert + print(f"[gui] steering probe: turn direction " + f"{'inverted' if invert else 'as wired'}") + + stats = TrainingStats() + hebbian = MotorRewardHebbian(eta=args.eta) + app = ForagingApp( + forager, stats, hebbian, banner, description, speed=args.speed + ) + + if args.smoke: + for _ in range(args.smoke_frames): + app.tick() + app.draw() + _save_frame(app, args.save) + print(f"[gui] {app.summary()}") + return + + app.add_controls() + import matplotlib.pyplot as plt + from matplotlib.animation import FuncAnimation + + # The animation must stay referenced, or it is garbage collected before + # it ever renders a frame. + app.anim = FuncAnimation( + app.fig, app.on_frame, interval=33, blit=True, + cache_frame_data=False, repeat=False, + ) + if args.duration > 0: + _close_after(app.fig, args.duration) + print("[gui] running -- press q or close the window to stop") + plt.show() + _save_frame(app, args.save) + print(f"[gui] {app.summary()}") + + +def _save_frame(app: "ForagingApp", path: str | None) -> None: + """Write the current figure to *path* when one was requested. + + Args: + app (ForagingApp): Application whose figure is saved. + path (str | None): Destination PNG path; ignored when ``None``. + + Returns: + None: Saves the figure and reports the path. + """ + if not path: + return + app.fig.savefig(path, dpi=110, bbox_inches="tight") + print(f"[gui] saved {path}") + + +def _close_after(fig, seconds: float) -> None: + """Close *fig* automatically once *seconds* have elapsed. + + Uses the canvas timer API, falling back to a repeating callback on + older matplotlib where ``single_shot_add_callback`` does not exist. + + Args: + fig (matplotlib.figure.Figure): Figure to close. + seconds (float): Delay before closing. + + Returns: + None: Starts a one-shot timer bound to the figure's canvas. + """ + import matplotlib.pyplot as plt + + timer = fig.canvas.new_timer(interval=int(seconds * 1000)) + once = getattr(timer, "single_shot_add_callback", None) + if once is not None: + once(plt.close, fig) + else: # matplotlib < 3.9 + def _close() -> None: + timer.stop() + plt.close(fig) + + timer.add_callback(_close) + fig._flynet_timer = timer # keep a reference alive + timer.start() + + +if __name__ == "__main__": + main() diff --git a/examples/make_figures.py b/examples/make_figures.py new file mode 100644 index 0000000..6af8a80 --- /dev/null +++ b/examples/make_figures.py @@ -0,0 +1,528 @@ +#!/usr/bin/env python3 +"""Regenerate the figures in ``docs/images/``. + +The repository's documentation used to ship images with no code to +produce them, so the numbers in the README drifted away from what the +code actually does. This script is that missing generator: it writes +every figure referenced by the README and prints the measured values so +the captions can be checked against them. + +Run it after changing the connectome loader, the learning rules or the +examples:: + + .venv/bin/python examples/make_figures.py + .venv/bin/python examples/make_figures.py --only gui --outdir /tmp + +It needs the real connectome, so it requires either a warm +``~/.flynet/cache`` or ``NEUPRINT_APPLICATION_CREDENTIALS``; it never +falls back to synthetic data. +""" + +from __future__ import annotations + +import argparse +import importlib.util +import sys +from pathlib import Path + +import matplotlib + +matplotlib.use("Agg") + +import matplotlib.pyplot as plt # noqa: E402 +import numpy as np # noqa: E402 + +from flynet.connectome import ConnectomeLoader # noqa: E402 +from flynet.learning import RewardHebbian, TrainingLogger # noqa: E402 +from flynet.network import SpikingNetwork # noqa: E402 +from flynet.visualize import ( # noqa: E402 + plot_brain_circuit, + plot_learning_curve, + plot_trajectory, +) + +_HERE = Path(__file__).resolve().parent +MINI_TYPE = "DNge104" +NAV_SEED = 7 # seed for the before/after navigation comparison +NAV_TRAIN_EPISODES = 5 # episodes between the before/after pair +CURVE_EPISODES = 30 # longer run so the learning curve has a trend +NAV_EVAL_SEED = 21 +GUI_SEED = 11 +GUI_TRAIN_EPISODES = 120 +EVAL_SEED0 = 90000 # held-out hunt seeds for the before/after check +EVAL_EPISODES = 24 + + +def _load(name: str): + """Import a sibling example module by file path. + + ``examples/`` is not a package, so its modules are loaded explicitly + instead of being imported by name. + + Args: + name (str): File name of the module inside ``examples/``. + + Returns: + module: The imported module. + """ + path = _HERE / f"{name}.py" + spec = importlib.util.spec_from_file_location(f"flynet_example_{name}", path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _load_brain(loader: ConnectomeLoader, mini: bool) -> SpikingNetwork: + """Load the mini or full connectome as a :class:`SpikingNetwork`. + + Args: + loader (ConnectomeLoader): Loader to fetch through. + mini (bool): Use the DNge104 mini brain instead of the whole + connectome. + + Returns: + SpikingNetwork: The assembled brain. + """ + if mini: + ids, edges, motors = loader.load_or_fetch_mini( + neuron_type=MINI_TYPE, max_edges=40 + ) + else: + ids, edges, motors = loader.load_or_fetch_full(min_weight=50) + motors = [] + return SpikingNetwork(ids, edges, motor_ids=motors if motors else None) + + +# --------------------------------------------------------------------------- +# Figure 1: brain circuit +# --------------------------------------------------------------------------- +def figure_brain_circuit(loader: ConnectomeLoader, outdir: Path) -> str: + """Draw the wiring around the DNge104 descending neuron. + + Args: + loader (ConnectomeLoader): Loader to fetch the mini brain through. + outdir (Path): Directory to write the PNG into. + + Returns: + str: One-line summary of what was measured. + """ + net = _load_brain(loader, mini=True) + fig = plt.figure(figsize=(9, 9)) + plot_brain_circuit(net, ax=fig.add_subplot(111), + title=f"Brain Circuit ({MINI_TYPE}, real connectome)") + path = outdir / "real_brain_circuit.png" + fig.savefig(path, dpi=130, bbox_inches="tight") + plt.close(fig) + return (f"{len(net.ids)} neurons, {net.W.nnz} synapses, " + f"sensors={net.sensors}, motors={net.motors}") + + +# --------------------------------------------------------------------------- +# Figures 2 and 3: navigation before/after, and its learning curve +# --------------------------------------------------------------------------- +def _train_mini(net: SpikingNetwork, episodes: int) -> TrainingLogger: + """Run Hebbian training on the mini brain and log each episode. + + Args: + net (SpikingNetwork): Brain to train in place. + episodes (int): Number of training episodes. + + Returns: + TrainingLogger: Per-episode reward and step history. + """ + nav = _load("navigation_task") + hebbian = RewardHebbian(eta=0.6) + logger = TrainingLogger() + for ep in range(1, episodes + 1): + rng = np.random.default_rng(NAV_SEED + ep) + steps, _path, events, found = nav.play_game(net, rng) + logger.log_episode(1.0 if found else 0.0, steps or nav.MAX_STEPS) + if found and events: + hebbian.update(net, events) + return logger + + +def _steps_label(steps: int | None, found: bool, path_len: int) -> str: + """Describe a hunt in a few words for a figure title. + + Args: + steps (int | None): Steps reported by the episode, ``None`` when + the food was never reached. + found (bool): Whether the food was reached. + path_len (int): Number of positions in the recorded path. + + Returns: + str: A short label such as ``"97 steps"`` or ``"gave up (300 + steps)"``. + """ + if found: + return f"{steps} steps" + return f"gave up ({path_len - 1} steps)" + + +def _frame_axes(ax, arena: float, legend: bool = True) -> None: + """Give a navigation panel a fixed frame so panels stay comparable. + + ``plot_trajectory`` autoscales to the path, which makes a + before/after pair hard to read; this pins both panels to the same + square arena and moves the legend off the food. + + Args: + ax (matplotlib.axes.Axes): Panel to adjust. + arena (float): Side length of the square arena. + legend (bool): Whether to reposition the legend. + + Returns: + None: Adjusts the axes in place. + """ + ax.set_xlim(0, arena) + ax.set_ylim(0, arena) + ax.set_aspect("equal") + if legend and ax.get_legend() is not None: + ax.get_legend().loc = "lower right" + + +def figure_navigation(loader: ConnectomeLoader, outdir: Path) -> str: + """Draw the mini-brain navigation path before and after training. + + Both panels replay the *same* episode seed so the difference is due + to the learned weights and nothing else. + + Args: + loader (ConnectomeLoader): Loader to fetch the mini brain through. + outdir (Path): Directory to write the PNG into. + + Returns: + str: One-line summary of the measured before/after step counts. + """ + nav = _load("navigation_task") + before_net = _load_brain(loader, mini=True) + + steps_before, path_before, _, found_before = nav.play_game( + before_net, np.random.default_rng(NAV_EVAL_SEED) + ) + + after_net = before_net.copy() + logger = _train_mini(after_net, NAV_TRAIN_EPISODES) + steps_after, path_after, _, found_after = nav.play_game( + after_net, np.random.default_rng(NAV_EVAL_SEED) + ) + + fig, axes = plt.subplots(1, 2, figsize=(13, 6.4)) + for ax, path, steps, found, label in ( + (axes[0], path_before, steps_before, found_before, "Before learning"), + (axes[1], path_after, steps_after, found_after, + f"After {NAV_TRAIN_EPISODES} episodes of Hebbian training"), + ): + plot_trajectory( + [(p[0], p[1]) for p in path], nav.FOOD, tuple(nav.START), ax=ax, + title=f"{label} - {_steps_label(steps, found, len(path))}", + ) + _frame_axes(ax, nav.ARENA) + fig.suptitle("Fly Navigation (real connectome, DNge104 mini brain)") + fig.tight_layout() + path = outdir / "real_trajectory.png" + fig.savefig(path, dpi=130, bbox_inches="tight") + plt.close(fig) + + steps = logger.get_learning_curve()[1] + return (f"before={_steps_label(steps_before, found_before, len(path_before))} " + f"after={_steps_label(steps_after, found_after, len(path_after))} " + f"(training steps: {steps})") + + +def figure_learning_curve(loader: ConnectomeLoader, outdir: Path) -> str: + """Draw the mini-brain navigation learning curve. + + Args: + loader (ConnectomeLoader): Loader to fetch the mini brain through. + outdir (Path): Directory to write the PNG into. + + Returns: + str: One-line summary of the curve's first and last episodes. + """ + net = _load_brain(loader, mini=True) + logger = _train_mini(net, CURVE_EPISODES) + + fig = plt.figure(figsize=(10, 5)) + plot_learning_curve( + logger, ax=fig.add_subplot(111), + title=f"Navigation learning curve ({MINI_TYPE}, " + f"{CURVE_EPISODES} episodes)", + ) + path = outdir / "real_learning_curve.png" + fig.savefig(path, dpi=130, bbox_inches="tight") + plt.close(fig) + + steps = logger.get_learning_curve()[1] + return f"steps per episode: {steps}" + + +# --------------------------------------------------------------------------- +# Figures 4 and 5: the interactive foraging GUI +# --------------------------------------------------------------------------- +def _gui_context(loader: ConnectomeLoader): + """Load the full connectome and a foraging app on top of it. + + Args: + loader (ConnectomeLoader): Loader to fetch the connectome through. + + Returns: + tuple: ``(app, forager, rule)`` for driving and training. + """ + gui = _load("fly_foraging_gui") + net = _load_brain(loader, mini=False) + forager = gui.FlyForager( + net, rng=np.random.default_rng(GUI_SEED), + settle_iters=gui.SETTLE_ITERS_FULL, + ) + forager.invert_steer = forager.auto_steer_sign() == -1 + stats = gui.TrainingStats() + rule = gui.MotorRewardHebbian(eta=0.05) + app = gui.ForagingApp( + forager, stats, rule, + banner="[connectome] data source: real male CNS " + f"({loader.last_source})", + description=f"full connectome from {loader.dataset}", + ) + return app, forager, rule + + +def figure_gui_frame(loader: ConnectomeLoader, outdir: Path) -> str: + """Capture one frame of the running GUI on the full connectome. + + The frame is taken while the fly is actively tracking the food with + one of its senses, which is the state worth showing in the docs. + + Args: + loader (ConnectomeLoader): Loader to fetch the connectome through. + outdir (Path): Directory to write the PNG into. + + Returns: + str: One-line summary of the captured state. + """ + gui = _load("fly_foraging_gui") + app, forager, _rule = _gui_context(loader) + app.add_controls() + app.set_speed(16) + + # Wait for enough hunts that the HUD and the learning curve have + # something in them, then catch the fly while it is tracking. + for _ in range(3000): + app.tick() + if app.stats.episodes < 20: + continue + mode = forager.episode.mode + if mode in ("sight", "smell") and 12.0 < forager.distance < 45.0: + break + # Settle transient UI (the "FOUND FOOD!" banner counts down per draw) + # without advancing the simulation, so the still frame is the steady + # state a reader would see. + for _ in range(gui.FLASH_FRAMES + 1): + app.draw() + app.fig.canvas.draw() + path = outdir / "gui_foraging.png" + app.fig.savefig(path, dpi=110, bbox_inches="tight") + plt.close(app.fig) + return (f"episode {app.stats.episodes}, " + f"success {app.stats.successes}/{app.stats.episodes}, " + f"frame mode={forager.episode.mode}") + + +def _evaluate_hunts(gui, net, Minv, invert, seeds) -> tuple[int, float]: + """Score a brain on a fixed set of episode seeds. + + Uses the same held-out seeds for every brain so that before/after + numbers are comparable; a single episode is far too noisy to judge + learning on the full connectome. + + Args: + gui (module): The foraging GUI example module. + net (SpikingNetwork): Brain to score. + Minv (np.ndarray): Frozen steering calibration to reuse. + invert (bool): Whether the steering direction is flipped. + seeds (range): Episode seeds to score. + + Returns: + tuple: ``(successes, mean_steps)`` over the seeds. + """ + runner = gui.FlyForager( + net, rng=np.random.default_rng(EVAL_SEED0), Minv=Minv, + invert_steer=invert, + ) + hits, steps = 0, [] + for seed in seeds: + runner.rng = np.random.default_rng(seed) + runner.reset() + while (not runner.episode.found + and runner.episode.step < gui.MAX_STEPS): + runner.step() + if runner.episode.found: + hits += 1 + steps.append(runner.episode.step) + return hits, float(np.mean(steps)) if steps else float("nan") + + +def figure_gui_learning(loader: ConnectomeLoader, outdir: Path) -> str: + """Show that the GUI's fly learns, on training and on held-out hunts. + + The left panel is the live training signal the window shows; the right + panel is the honest check, scoring the untrained and the trained + connectome on the *same* 24 held-out episode seeds. + + Args: + loader (ConnectomeLoader): Loader to fetch the connectome through. + outdir (Path): Directory to write the PNG into. + + Returns: + str: One-line summary of the measured success rates. + """ + gui = _load("fly_foraging_gui") + _app, forager, rule = _gui_context(loader) + stats = gui.TrainingStats() + seeds = range(EVAL_SEED0, EVAL_SEED0 + EVAL_EPISODES) + + hits_before, steps_before = _evaluate_hunts( + gui, forager.net, forager.Minv, forager.invert_steer, seeds + ) + # Public API only: snapshot the connectome's own weights, then train. + untrained = forager.net.copy() + + for ep in range(GUI_TRAIN_EPISODES): + forager.rng = np.random.default_rng(GUI_SEED * 1000 + ep) + forager.reset() + while (not forager.episode.found + and forager.episode.step < gui.MAX_STEPS): + forager.step() + stats.record(forager.episode.found, forager.episode.step) + gui.apply_learning( + forager.net, rule, + forager.learning_batch(found=forager.episode.found), + ) + + hits_after, steps_after = _evaluate_hunts( + gui, forager.net, forager.Minv, forager.invert_steer, seeds + ) + del untrained + + fig, axes = plt.subplots(1, 2, figsize=(13, 5.2)) + + # -- left: the live training signal -------------------------------- + episodes, steps = stats.logger.get_learning_curve() + axes[0].plot(episodes, steps, "-", color="steelblue", alpha=0.3, + label="steps per hunt") + sm_x, sm_y = stats.smoothed_steps() + axes[0].plot(sm_x, sm_y, "-", color="crimson", linewidth=2, + label="smoothed") + rate_x, rate_y = stats.success_series() + axes[0].set_xlabel("hunt") + axes[0].set_ylabel("steps to food") + axes[0].set_ylim(0, float(gui.MAX_STEPS)) + axes[0].set_title(f"Training the fly ({GUI_TRAIN_EPISODES} hunts)") + axes[0].legend(loc="upper right", fontsize=8) + + rate_ax = axes[0].twinx() + rate_ax.plot(rate_x, rate_y, "-", color="seagreen", linewidth=1.5) + rate_ax.set_ylabel("rolling success rate", color="seagreen", fontsize=9) + rate_ax.tick_params(labelsize=8, colors="seagreen") + rate_ax.set_ylim(-0.05, 1.05) + + # -- right: the held-out comparison -------------------------------- + labels = ["untrained\nconnectome", f"after {GUI_TRAIN_EPISODES}\nhunts"] + rates = [100.0 * hits_before / EVAL_EPISODES, + 100.0 * hits_after / EVAL_EPISODES] + means = [steps_before, steps_after] + bars = axes[1].bar(labels, rates, color=["#9fb8cd", "#3f7f9f"], width=0.5) + axes[1].set_ylabel(f"success rate on {EVAL_EPISODES} held-out hunts (%)") + axes[1].set_ylim(0, 110) + axes[1].set_title("Held-out check (same 24 start conditions)") + for bar, rate, mean in zip(bars, rates, means): + axes[1].text( + bar.get_x() + bar.get_width() / 2, bar.get_height() + 2, + f"{rate:.0f}%\n{mean:.0f} steps avg", + ha="center", va="bottom", fontsize=9, + ) + + fig.suptitle("Foraging on the full real connectome: live learning") + fig.tight_layout() + path = outdir / "gui_foraging_learning.png" + fig.savefig(path, dpi=130, bbox_inches="tight") + plt.close(fig) + return (f"held-out success {hits_before}/{EVAL_EPISODES} -> " + f"{hits_after}/{EVAL_EPISODES}, " + f"avg steps {steps_before:.0f} -> {steps_after:.0f}") + + +FIGURES = { + "brain_circuit": figure_brain_circuit, + "trajectory": figure_navigation, + "curve": figure_learning_curve, + "gui": figure_gui_frame, + "gui_learning": figure_gui_learning, +} + + +def build_parser() -> argparse.ArgumentParser: + """Build the command-line parser. + + Returns: + argparse.ArgumentParser: Parser for the figure generator. + """ + p = argparse.ArgumentParser( + description="Regenerate the figures used by the README." + ) + p.add_argument( + "--outdir", default=str(_HERE.parent / "docs" / "images"), + help="where to write the PNGs (default: %(default)s)", + ) + p.add_argument( + "--dataset", default="male-cns:v0.9", + help="neuPrint dataset (default: %(default)s)", + ) + p.add_argument( + "--only", default="", + help="comma-separated subset of: " + ", ".join(FIGURES), + ) + p.add_argument( + "--list", action="store_true", help="list the figure names and exit" + ) + return p + + +def main() -> None: + """Regenerate the requested figures. + + Raises: + SystemExit: If the real connectome cannot be loaded. + """ + args = build_parser().parse_args() + if args.list: + for name in FIGURES: + print(name) + return + + wanted = ( + [s.strip() for s in args.only.split(",") if s.strip()] + if args.only else list(FIGURES) + ) + unknown = [s for s in wanted if s not in FIGURES] + if unknown: + raise SystemExit( + f"unknown figure(s): {', '.join(unknown)}; " + f"choose from {', '.join(FIGURES)}" + ) + + outdir = Path(args.outdir) + outdir.mkdir(parents=True, exist_ok=True) + loader = ConnectomeLoader(dataset=args.dataset) + + for name in wanted: + print(f"[figures] {name} ...", flush=True) + summary = FIGURES[name](loader, outdir) + print(f"[figures] {name}: {summary}", flush=True) + print(f"[figures] wrote {len(wanted)} figure(s) to {outdir}") + + +if __name__ == "__main__": + main() diff --git a/src/flynet/__init__.py b/src/flynet/__init__.py index b484cef..04be2d2 100644 --- a/src/flynet/__init__.py +++ b/src/flynet/__init__.py @@ -14,7 +14,12 @@ from flynet.stdp import STDPRule from flynet.network import SpikingNetwork from flynet.connectome import ConnectomeLoader, ConnectomeUnavailableError -from flynet.learning import RewardHebbian, STDPTrainer, TrainingLogger +from flynet.learning import ( + MotorRewardHebbian, + RewardHebbian, + STDPTrainer, + TrainingLogger, +) from flynet.inhibition import LateralInhibition # Gradient training (optional – requires jax) @@ -46,6 +51,7 @@ "SpikingNetwork", "ConnectomeLoader", "ConnectomeUnavailableError", + "MotorRewardHebbian", "RewardHebbian", "STDPTrainer", "TrainingLogger", diff --git a/src/flynet/cli.py b/src/flynet/cli.py index 4c8be17..e227f92 100644 --- a/src/flynet/cli.py +++ b/src/flynet/cli.py @@ -32,7 +32,7 @@ def main(argv: list[str] | None = None) -> None: # -- flynet simulate ---------------------------------------------------- sim = sub.add_parser("simulate", help="Run a simulation with the real or synthetic connectome") - sim.add_argument("--dataset", default="male-cns:v1.0", + sim.add_argument("--dataset", default="male-cns:v0.9", help="Connectome dataset (default: %(default)s)") sim.add_argument("--neuron-type", default="DNge104", help="Cell type for mini brain (default: %(default)s)") @@ -55,7 +55,7 @@ def main(argv: list[str] | None = None) -> None: # -- flynet train ------------------------------------------------------- tr = sub.add_parser("train", help="Train the brain (Hebbian or gradient-based)") - tr.add_argument("--dataset", default="male-cns:v1.0") + tr.add_argument("--dataset", default="male-cns:v0.9") tr.add_argument("--neuron-type", default="DNge104") tr.add_argument("--full", action="store_true") tr.add_argument("--episodes", type=int, default=2, help="Training episodes (default: %(default)s)") diff --git a/src/flynet/connectome.py b/src/flynet/connectome.py index 23dca3a..b11742a 100644 --- a/src/flynet/connectome.py +++ b/src/flynet/connectome.py @@ -52,7 +52,7 @@ class ConnectomeLoader: Attributes: server (str): Base URL of the neuPrint server. - dataset (str): Dataset identifier (e.g. ``"male-cns:v1.0"``). + dataset (str): Dataset identifier (e.g. ``"male-cns:v0.9"``). cache_dir (str): Local directory used for CSV/JSON cache files. last_source (str): Provenance of the most recent load -- ``"cache"``, ``"neuprint"``, ``"synthetic"``, or ``""``. @@ -62,8 +62,7 @@ class ConnectomeLoader: DATASETS: dict[str, str] = { "hemibrain:v1.2.1": "adult FEMALE central brain (the classic hemibrain)", - "male-cns:v0.9": "male CNS draft (brain + nerve cord + optic lobes)", - "male-cns:v1.0": "adult MALE brain + nerve cord + optic lobes (default)", + "male-cns:v0.9": "male CNS draft (brain + nerve cord + optic lobes) (default)", "manc:v1.0": "adult MALE ventral nerve cord (MANC)", "manc:v1.2.1": "adult MALE ventral nerve cord (MANC), updated", "manc:v1.2.3": "adult MALE ventral nerve cord (MANC), current", @@ -75,7 +74,7 @@ class ConnectomeLoader: def __init__( self, server: str = "https://neuprint.janelia.org", - dataset: str = "male-cns:v1.0", + dataset: str = "male-cns:v0.9", cache_dir: str | None = None, ) -> None: """Initialize the ConnectomeLoader. diff --git a/src/flynet/learning.py b/src/flynet/learning.py index 7e575a2..d625215 100644 --- a/src/flynet/learning.py +++ b/src/flynet/learning.py @@ -17,7 +17,94 @@ if TYPE_CHECKING: from flynet.network import SpikingNetwork -__all__ = ["RewardHebbian", "STDPTrainer", "TrainingLogger"] +__all__ = [ + "MotorRewardHebbian", + "RewardHebbian", + "STDPTrainer", + "TrainingLogger", +] + + +class MotorRewardHebbian: + """Reward-modulated plasticity on the motor projection. + + The steering command is the difference between the two motor neurons' + activities, and a motor's activity is + ``sum_pre a[pre] * W[pre, motor]``. The synapses that can therefore + change the command are the ones *into* the motors, and the + first-order sensitivity of the command to each of them is proportional + to the presynaptic activity. This rule reinforces exactly those + synapses, scaled by a signed reward, so a large brain learns through + the handful of synapses that actually drive its output instead of the + hundreds of thousands that do not. + + Attributes: + eta (float): Maximum absolute weight change per update. + w_max (float): Upper bound on an updated weight. + """ + + def __init__(self, eta: float = 0.05, w_max: float = 1.5) -> None: + """Initialize the motor-projection learner. + + Args: + eta (float): Maximum absolute weight change per update. + w_max (float): Upper bound on an updated weight. + """ + self.eta = eta + self.w_max = w_max + + def update(self, network: "SpikingNetwork", events: list, W: Any = None) -> None: + """Reinforce or weaken the synapses that project into the motors. + + Args: + network (SpikingNetwork): Network whose ``rows``/``cols`` COO + views and ``motors`` define the update. + events (list): ``(rho_right, rho_left, activity, reward)`` + tuples. *activity* is a dense vector or a sparse + ``(indices, values)`` pair; *reward* is signed, so steps + that moved towards the goal strengthen the active input + synapses and steps that moved away weaken them. + W (np.ndarray | None): Weight matrix to update. Default: + ``network.W``. + + Returns: + None: Updates *W* in place. + """ + if W is None: + W = network.W + rows, cols = network.rows, network.cols + + motor_cols = [network._ix[b] for b in network.motors[:2]] + mask = np.zeros(cols.shape[0], dtype=bool) + for m in motor_cols: + mask |= cols == m + if not mask.any(): + return + + pre_rows = rows[mask] + dW = np.zeros(int(mask.sum()), dtype=np.float64) + n = len(network.ids) + + for event in events: + if len(event) == 4: + _, _, act, reward = event + else: + _, _, act = event + reward = 1.0 + if isinstance(act, tuple): + idx, vals = act + a = np.zeros(n) + a[idx] = vals + else: + a = act + dW += reward * a[pre_rows] + + peak = float(np.abs(dW).max()) if dW.size else 0.0 + if peak <= 0.0: + return + W.data[mask] = np.clip( + W.data[mask] + self.eta * dW / peak, -self.w_max, self.w_max + ) class RewardHebbian: @@ -40,11 +127,21 @@ def __init__(self, eta: float = 0.6, reward_bonus: float = 6.0) -> None: def update(self, network: "SpikingNetwork", events: list, W: Any = None) -> None: """Apply Hebbian updates to the network weights. + Each event is ``(rho_right, rho_left, activity)``. A fourth + element may be supplied as an explicit signed reward, which + overrides the default ``rho_right + rho_left + reward_bonus``. + Negative rewards depress the co-active synapses, which turns the + rule into a proper three-factor learner: synapses active during a + good decision are strengthened, synapses active during a bad one + are weakened. + Args: network (SpikingNetwork): SpikingNetwork instance with ``rows`` and ``cols`` attributes (COO views of the synapse matrix). events (list): List of ``(rho_right, rho_left, activity_vector)`` - tuples from successful games. + tuples from successful games, optionally with a trailing + signed reward. *activity* is either a dense vector or a + sparse ``(indices, values)`` pair. W (np.ndarray | None): Weight matrix to update. Default: ``network.W``. @@ -55,22 +152,28 @@ def update(self, network: "SpikingNetwork", events: list, W: Any = None) -> None W = network.W rows, cols = network.rows, network.cols dW = np.zeros(rows.shape[0], dtype=np.float64) + n = len(network.ids) - for rho_right, rho_left, act in events: - reward = (rho_right + rho_left) + self.reward_bonus - if reward <= 0: - continue + for event in events: + if len(event) == 4: + rho_right, rho_left, act, reward = event + else: + rho_right, rho_left, act = event + reward = (rho_right + rho_left) + self.reward_bonus if isinstance(act, tuple): idx, vals = act - a = np.zeros(len(network.ids)) + a = np.zeros(n) a[idx] = vals else: a = act dW += reward * (a[rows] * a[cols]) - if dW.max() <= 0: + # Normalise by the largest magnitude so that an all-negative batch + # depresses instead of being skipped. + peak = float(np.abs(dW).max()) if dW.size else 0.0 + if peak <= 0.0: return - W.data += self.eta * dW / dW.max() + W.data += self.eta * dW / peak class STDPTrainer: diff --git a/src/flynet/network.py b/src/flynet/network.py index 959756c..29571b1 100644 --- a/src/flynet/network.py +++ b/src/flynet/network.py @@ -300,6 +300,7 @@ def turn( rho_left: float, W: sp.csr_matrix | None = None, Minv: np.ndarray | None = None, + iterations: int = SETTLE, ) -> tuple[float, list[np.ndarray]]: """Compute steering command for a smell pair. @@ -314,6 +315,9 @@ def turn( If ``None``, calibration is computed fresh. Pass a cached ``Minv`` to keep steering consistent across weight updates during training. + iterations (int): Settling iterations. Lower values trade settling + accuracy for speed, which matters on the full ~66k-neuron + connectome where each iteration is a sparse mat-vec. Returns: tuple[float, list[np.ndarray]]: ``(drv, trace)`` -- the motor difference and the full @@ -323,7 +327,7 @@ def turn( if Minv is None: _, Minv = self.calibrate(W=W) z = Minv @ np.array([rho_right, rho_left]) - tr = self.settle(z[0], z[1]) + tr = self.settle(z[0], z[1], iterations=iterations) m = [self._ix[b] for b in self.motors[:2]] if len(m) < 2: drv = float(tr[-1][m[0]]) diff --git a/tests/test_fly_foraging_gui.py b/tests/test_fly_foraging_gui.py new file mode 100644 index 0000000..49348c1 --- /dev/null +++ b/tests/test_fly_foraging_gui.py @@ -0,0 +1,297 @@ +"""Tests for the interactive foraging GUI example. + +The example lives in ``examples/`` and is not importable as a package, so +it is loaded by path. Only display-free pieces are exercised: the sensing +geometry, the step/reward bookkeeping and the plasticity rule. A short +headless run of ``main()`` covers the wiring end to end. +""" + +import importlib.util +import sys +from pathlib import Path + +import matplotlib +import numpy as np +import pytest + +matplotlib.use("Agg") + +from flynet.connectome import ConnectomeLoader +from flynet.learning import MotorRewardHebbian, RewardHebbian +from flynet.network import SpikingNetwork + +_EXAMPLE = ( + Path(__file__).resolve().parents[1] / "examples" / "fly_foraging_gui.py" +) + + +def _load_example(): + """Import ``examples/fly_foraging_gui.py`` by path. + + Returns: + module: The imported example module. + """ + spec = importlib.util.spec_from_file_location("fly_foraging_gui", _EXAMPLE) + module = importlib.util.module_from_spec(spec) + sys.modules["fly_foraging_gui"] = module + spec.loader.exec_module(module) + return module + + +fg = _load_example() + + +@pytest.fixture() +def small_net(): + """Return the 41-neuron synthetic brain as a SpikingNetwork.""" + ids, edges, motors = ConnectomeLoader.synthetic_brain(max_edges=40) + return SpikingNetwork(ids, edges, motor_ids=motors) + + +# --------------------------------------------------------------------------- +# sensing geometry +# --------------------------------------------------------------------------- +def test_wrap_angle_stays_in_range(): + assert fg.wrap_angle(3.0 * np.pi) == pytest.approx(-np.pi) + assert fg.wrap_angle(0.5) == pytest.approx(0.5) + assert abs(fg.wrap_angle(7.0)) <= np.pi + + +def test_smell_decays_with_distance(): + food = (50.0, 50.0) + near = fg.smell_at(np.array([50.0, 55.0]), food) + far = fg.smell_at(np.array([50.0, 70.0]), food) + assert near > far > 0.0 + assert fg.smell_at(np.array([50.0, 50.0]), food) == pytest.approx(2.0) + + +def test_visual_field_is_symmetric_when_food_is_ahead(): + # Facing +x, food directly ahead: both eyes see the same contrast. + right, left = fg.visual_readings(np.array([20.0, 50.0]), 0.0, (50.0, 50.0)) + assert right == pytest.approx(left) + assert right > 0.0 + + +def test_visual_field_favours_the_eye_the_food_is_in(): + # Facing +x in a y-up frame, the fly's left is +y, so food at +y must + # brighten the left eye and food at -y the right one. + pos = np.array([20.0, 50.0]) + right, left = fg.visual_readings(pos, 0.0, (50.0, 60.0)) + assert left > right > 0.0 + right, left = fg.visual_readings(pos, 0.0, (50.0, 40.0)) + assert right > left > 0.0 + + +def test_visual_field_ignores_food_behind_and_beyond_range(): + pos = np.array([20.0, 50.0]) + # Facing -x, so food at larger x is behind the fly. + assert fg.visual_readings(pos, np.pi, (30.0, 50.0)) == (0.0, 0.0) + # Food ahead but past the visual range. + far = (20.0 + fg.SIGHT_RANGE + 10.0, 50.0) + assert fg.visual_readings(pos, 0.0, far) == (0.0, 0.0) + + +def test_antennae_favour_the_side_the_plume_is_on(): + pos = np.array([20.0, 50.0]) + # Facing +x, its left antenna points to +y. + right, left = fg.antenna_readings(pos, 0.0, (50.0, 60.0)) + assert left > right + right, left = fg.antenna_readings(pos, 0.0, (50.0, 40.0)) + assert right > left + + +def test_topk_activity_is_sparse_and_sized(): + trace = [np.zeros(1000), np.arange(1000, dtype=float)] + idx, vals = fg.topk_activity(trace, k=10) + assert idx.size == 10 + assert vals.size == 10 + # The mean of the last two iterations is kept, and the largest + # magnitudes survive the truncation. + assert np.abs(vals).max() == pytest.approx(999.0 / 2.0) + empty_i, empty_v = fg.topk_activity([], k=10) + assert empty_i.size == 0 and empty_v.size == 0 + + +# --------------------------------------------------------------------------- +# forager +# --------------------------------------------------------------------------- +def test_forager_runs_and_records_signed_rewards(small_net): + forager = fg.FlyForager(small_net, rng=np.random.default_rng(0)) + for _ in range(40): + forager.step() + assert forager.episode.step == 40 + assert forager.step_count == 40 + for event in forager.episode.events: + assert len(event) == 4 + _, _, activity, reward = event + assert isinstance(activity, tuple) + assert np.isfinite(reward) + + +def test_forager_stays_inside_the_arena(small_net): + forager = fg.FlyForager(small_net, rng=np.random.default_rng(1)) + for _ in range(120): + forager.step() + assert np.all(forager.episode.pos >= 0.0) + assert np.all(forager.episode.pos <= forager.arena) + + +def test_search_mode_actually_travels(small_net): + """Search must cover ground, not spin on the spot.""" + forager = fg.FlyForager(small_net, rng=np.random.default_rng(2)) + forager.reset(start=np.array([10.0, 10.0])) + while forager.episode.mode != "search" and forager.episode.step < 20: + forager.step() + forager.episode.events.clear() + start = forager.episode.pos.copy() + for _ in range(25): + forager.step() + assert float(np.linalg.norm(forager.episode.pos - start)) > 10.0 + + +def test_learning_batch_bounds_length_and_adds_terminal_reward(small_net): + forager = fg.FlyForager(small_net, rng=np.random.default_rng(3)) + for _ in range(4): + forager.reset(start=np.array([10.0, 10.0])) + while not forager.episode.events and forager.episode.step < fg.MAX_STEPS: + forager.step() + if forager.episode.events: + break + assert forager.episode.events, "no guided step produced a learning event" + + batch = forager.learning_batch(limit=5, found=True) + assert 1 <= len(batch) <= 5 + plain = forager.learning_batch(limit=5) + assert batch[-1][3] == pytest.approx(plain[-1][3] + fg.SUCCESS_BONUS) + failed = forager.learning_batch(limit=5, found=False) + assert failed[-1][3] == pytest.approx( + plain[-1][3] + fg.FAILURE_PENALTY + ) + assert forager.learning_batch(limit=0) == [] + + +# --------------------------------------------------------------------------- +# learning rules +# --------------------------------------------------------------------------- +def test_motor_reward_hebbian_moves_the_motor_projection(small_net): + rows, cols = small_net.rows, small_net.cols + mask = np.zeros(cols.shape[0], dtype=bool) + for body in small_net.motors[:2]: + mask |= cols == small_net._ix[body] + assert mask.any() + + before = small_net.W.data.copy() + events = [(2.0, 0.5, (np.arange(20), np.linspace(0.2, 1.0, 20)), 1.0)] + MotorRewardHebbian(eta=0.5).update(small_net, events) + after = small_net.W.data + assert not np.allclose(before[mask], after[mask]) + # Nothing outside the motor projection may move. + assert np.allclose(before[~mask], after[~mask]) + + +def test_motor_reward_hebbian_handles_all_negative_rewards(small_net): + """A purely negative batch must depress, not be skipped.""" + mask = np.zeros(small_net.cols.shape[0], dtype=bool) + for body in small_net.motors[:2]: + mask |= small_net.cols == small_net._ix[body] + before = small_net.W.data[mask].copy() + # Realistic events: non-negative activity, negative reward. + events = [(0.5, 2.0, (np.arange(20), np.linspace(0.2, 1.0, 20)), -1.0)] + MotorRewardHebbian(eta=0.5).update(small_net, events) + after = small_net.W.data[mask] + assert not np.allclose(before, after) + # Every delta is non-positive, so the total drive must shrink. + assert after.sum() < before.sum() + + +def test_reward_hebbian_still_accepts_three_element_events(small_net): + """The pre-existing 3-tuple contract must keep working.""" + before = float(small_net.W.data.mean()) + events = [(2.0, 0.5, (np.arange(20), np.linspace(0.2, 1.0, 20)))] + RewardHebbian(eta=0.1).update(small_net, events) + assert float(small_net.W.data.mean()) != before + + +# --------------------------------------------------------------------------- +# stats and app +# --------------------------------------------------------------------------- +def test_training_stats_rolling_window(): + stats = fg.TrainingStats(window=3) + assert stats.success_rate() == 0.0 + assert stats.avg_steps() == 0.0 + for found, steps in ((True, 10), (False, 300), (True, 20)): + stats.record(found, steps) + assert stats.episodes == 3 + assert stats.successes == 2 + assert stats.best_steps == 10 + assert stats.success_rate() == pytest.approx(2 / 3) + stats.record(True, 5) + # The window drops the oldest entry, so the perfect run sticks. + assert stats.best_steps == 5 + assert len(stats.recent) == 3 + xs, ys = stats.success_series() + assert xs == [1, 2, 3, 4] + assert ys[-1] == pytest.approx(2 / 3) + + +def test_app_clamps_speed_and_advances(small_net): + forager = fg.FlyForager(small_net, rng=np.random.default_rng(4)) + app = fg.ForagingApp( + forager, fg.TrainingStats(), MotorRewardHebbian(eta=0.25), "b", "b", + speed=999, + ) + assert app.speed == fg.MAX_SPEED + app.set_speed(0) + assert app.speed == fg.MIN_SPEED + + app.set_speed(3) + before = forager.step_count + app.tick() + assert forager.step_count == before + 3 + + app.paused = True + held = forager.step_count + app.tick() + assert forager.step_count == held + + app.on_key(type("E", (), {"key": " "})()) + assert app.paused is False + app.on_key(type("E", (), {"key": "q"})()) + assert app.running is False + + +# --------------------------------------------------------------------------- +# end to end +# --------------------------------------------------------------------------- +def test_headless_main_runs_on_the_synthetic_brain(tmp_path, capsys): + out = tmp_path / "frame.png" + argv = sys.argv + sys.argv = [ + "fly_foraging_gui", "--synthetic", "--smoke", "--smoke-frames", "60", + "--speed", "8", "--save", str(out), + ] + try: + fg.main() + finally: + sys.argv = argv + captured = capsys.readouterr().out + assert "SYNTHETIC" in captured + assert "steps=" in captured + assert out.exists() and out.stat().st_size > 0 + + +def test_offline_without_cache_exits_cleanly(monkeypatch, tmp_path, capsys): + monkeypatch.setattr( + ConnectomeLoader, "cache_dir", str(tmp_path), raising=False + ) + monkeypatch.setattr( + ConnectomeLoader, "load_cached_full", lambda self: None + ) + argv = sys.argv + sys.argv = ["fly_foraging_gui", "--offline", "--smoke", "--smoke-frames", "1"] + try: + with pytest.raises(SystemExit) as exc: + fg.main() + finally: + sys.argv = argv + assert "no cached real brain" in str(exc.value)