flybrain 0.1.0__py3-none-any.whl

flybrain/__init__.py

@@ -0,0 +1,18 @@
+"""flybrain: the complete fruit fly nervous system (MaleCNS v1.0 connectome, 166,700 neurons)
+as a spiking network you can stimulate, read out and train readouts on.
+
+    from flybrain import FlyBrain
+    brain = FlyBrain(device="auto")      # downloads the brain files on first use
+    brain.stimulate(brain.cells(["LC4", "LPLC2"], side="L"), 0.8)
+    fired = brain.step()                 # indices of the neurons that spiked this 20 ms step
+"""
+from .brain import FlyBrain, cuda_available
+from .data import DATA, download, ensure_data, has_data
+from .eyes import ENCODER, Blob, Eyes, FeatureDetectors, blob_for
+from .reservoir import Readout, Trace, auc, bases_for, fit_logistic, fit_ridge, folds, project, run
+
+__version__ = "0.1.0"
+
+__all__ = ["FlyBrain", "cuda_available", "DATA", "download", "ensure_data", "has_data",
+           "ENCODER", "Blob", "Eyes", "FeatureDetectors", "blob_for",
+           "Readout", "Trace", "auc", "bases_for", "fit_logistic", "fit_ridge", "folds", "project", "run"]

flybrain/__main__.py

@@ -0,0 +1,51 @@
+"""Command line: `flybrain download`, `flybrain build`, `flybrain info` (or `python -m flybrain ...`)."""
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+from . import __version__
+from .data import DATA, FILES, RELEASE_URL, download, has_data
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(prog="flybrain", description="The MaleCNS fruit fly connectome as a spiking network.")
+    parser.add_argument("--version", action="version", version=f"flybrain {__version__}")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    fetch = sub.add_parser("download", help="fetch the prebuilt brain files (~260 MB)")
+    fetch.add_argument("--data", type=Path, default=DATA, help=f"where to put them (default {DATA}, or $FLY_DATA)")
+    fetch.add_argument("--url", default=RELEASE_URL, help="base URL of the files (or $FLYBRAIN_DATA_URL)")
+    fetch.add_argument("--force", action="store_true", help="download again even if the files are there")
+
+    build = sub.add_parser("build", help="download MaleCNS v1.0 (~1.1 GB) and build the brain files from it")
+    build.add_argument("--data", type=Path, default=DATA, help=f"data folder (default {DATA}, or $FLY_DATA)")
+
+    info = sub.add_parser("info", help="show the data folder and whether a GPU is usable")
+    info.add_argument("--data", type=Path, default=DATA)
+
+    args = parser.parse_args(argv)
+    if args.command == "download":
+        download(args.data, args.url, force=args.force)
+        print(f"brain files in {args.data}")
+    elif args.command == "build":
+        try:
+            from .build import build as build_brain
+        except ImportError as e:
+            raise SystemExit(f"building needs pandas and pyarrow ({e}): pip install \"flybrain[build]\"")
+        build_brain(args.data)
+    else:
+        from .brain import cuda_available
+        print(f"flybrain {__version__}")
+        print(f"data folder: {args.data}")
+        for name in FILES:
+            path = args.data / name
+            print(f"  {name}: {f'{path.stat().st_size / 1e6:,.0f} MB' if path.exists() else 'missing'}")
+        if not has_data(args.data):
+            print("  run `flybrain download` (or just create a FlyBrain) to fetch them")
+        gpu = "available" if cuda_available() else 'not available (pip install "flybrain[gpu]")'
+        print(f"cuda: {gpu}")
+
+
+if __name__ == "__main__":
+    main()

flybrain/brain.py

@@ -0,0 +1,201 @@
+"""Leaky integrate-and-fire simulation of the MaleCNS connectome.
+
+Dynamics follow ornata/fly (fly64/model.py) so results are comparable:
+    v <- exp(-dt/tau) v + gain * W @ spikes + tonic + noise + eye input
+    v >= 1 -> spike, reset to 0
+tonic/gain/noise are hand-calibrated, not measured. fly64 used tonic 0.18,
+gain 1.5, which parks every neuron at threshold (0.18 / (1 - 0.82) = 1.0) so
+the network ticks on its own. inject.py showed tonic 0.14, gain 3.0 keeps
+descending neurons quiet at rest (~1 Hz) while LC4/LPLC2 -> DNp01 and
+LC10a -> DNa02 signals still get through, ipsilaterally.
+
+batch > 1 runs that many independent flies (same wiring, own voltages and
+noise) in lock-step; on a GPU one sparse multiply serves them all, so 8 flies
+cost about as much as 1-2.
+"""
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import numba
+import numpy as np
+from scipy import sparse
+
+from .data import DATA, ensure_data
+
+
+@numba.njit(nogil=True, parallel=True)
+def _propagate(indptr, indices, weights, fired, n):
+    """Sum the outgoing weights (CSC columns) of every neuron that spiked.
+    Each thread scatters into its own buffer; buffers are summed at the end."""
+    threads = numba.get_num_threads()
+    partial = np.zeros((threads, n), np.float32)
+    chunk = (len(fired) + threads - 1) // threads
+    for t in numba.prange(threads):
+        acc = partial[t]
+        for k in range(t * chunk, min(len(fired), (t + 1) * chunk)):
+            j = fired[k]
+            for e in range(indptr[j], indptr[j + 1]):
+                acc[indices[e]] += weights[e]
+    current = np.zeros(n, np.float32)
+    for i in numba.prange(n):
+        s = np.float32(0.0)
+        for t in range(threads):
+            s += partial[t, i]
+        current[i] = s
+    return current
+
+
+def cuda_available() -> bool:
+    try:
+        import cupy
+        return cupy.cuda.runtime.getDeviceCount() > 0
+    except Exception:
+        return False
+
+
+class FlyBrain:
+    """device: "cpu" (numba), "cuda" (CuPy, NVIDIA GPU) or "auto"; defaults to
+    $FLY_DEVICE, else "cpu". Both run the same model; the noise streams differ,
+    so individual spikes differ between devices but statistics match.
+
+    batch: number of independent flies. Voltages are (n, batch). With batch 1,
+    step() returns the fired neuron indices; with batch > 1, a list of them,
+    one array per fly. Inputs broadcast: an amount can be a number (same for
+    every fly) or an array of length batch (one per fly)."""
+    dt = 0.020
+    tau = 0.100
+    gain = 3.0
+    tonic = 0.14            # calibrated at dt = 0.020; rescaled for other steps (see __init__)
+    noise_hz = 1.2
+    noise_amp = 0.22
+    eye_gain = 0.62
+
+    def __init__(self, data: Path | str | None = None, seed: int = 64, device: str | None = None, batch: int = 1,
+                 dt: float | None = None, sensory_input: bool = True, refractory: float = 0.0):
+        """data: folder with brain.npz and weights.npz (default $FLY_DATA, else ~/fly-data).
+        If they aren't there, the prebuilt brain is downloaded into it first (~260 MB, once).
+
+        dt: step length in seconds (default 0.020). tonic is rescaled so a silent
+        neuron settles at the same voltage as in the calibrated 20 ms model.
+
+        sensory_input: False removes every synapse onto sensory neurons (any superclass
+        containing "sensory"), so they fire only from noise and what you inject. With
+        True (the original model), olfactory receptor neurons excite each other into a
+        runaway loop and sit near maximum rate at rest, so odours add nothing.
+
+        refractory: seconds a neuron is held at 0 after it spikes (0 = none; at 20 ms
+        steps the step itself already caps rates at 50 Hz)."""
+        device = device or os.environ.get("FLY_DEVICE", "cpu")
+        if device == "auto":
+            device = "cuda" if cuda_available() else "cpu"
+        if device not in ("cpu", "cuda"):
+            raise ValueError(f"device must be cpu, cuda or auto, not {device!r}")
+        self.device = device
+        self.batch = int(batch)
+        if dt is not None:
+            self.dt = float(dt)
+        self.tonic = type(self).tonic * (1 - np.exp(-self.dt / self.tau)) / (1 - np.exp(-0.020 / self.tau))
+        self.refractory_steps = int(round(refractory / self.dt))
+        self.sensory_input = sensory_input
+        data = ensure_data(data)
+        meta = np.load(data / "brain.npz")
+        W = sparse.load_npz(data / "weights.npz")
+        if not sensory_input:
+            if "superclass" not in meta.files:
+                raise RuntimeError("brain.npz has no superclass; run `flybrain build`")
+            sensory = np.char.find(meta["superclass"].astype(str), "sensory") >= 0
+            W = sparse.diags((~sensory).astype(np.float32)) @ W.tocsr()   # rows = postsynaptic
+        if device == "cuda":
+            import cupy
+            from cupyx.scipy import sparse as cusparse
+            self.xp = cupy
+            self._W = cusparse.csr_matrix(W.tocsr().astype(np.float32))  # rows = postsynaptic
+        else:
+            self.xp = np
+        W = W.tocsc()
+        self.n = W.shape[0]
+        self.indptr, self.indices, self.weights = W.indptr, W.indices, W.data
+        self.visual = meta["visual"]
+        self.azimuth = meta["azimuth"]  # -1 far left ... +1 far right
+        self.cell_type = meta["cell_type"]
+        self.side = meta["side"]
+        self.positions = meta["positions"] if "positions" in meta.files else None
+        self.superclass = meta["superclass"] if "superclass" in meta.files else None
+        self.groups = {k.removeprefix("group_"): meta[k] for k in meta.files if k.startswith("group_")}
+        self._visual = self.xp.asarray(self.visual)
+        self.decay = np.float32(np.exp(-self.dt / self.tau))
+        self.reset(seed)
+
+    def reset(self, seed: int | None = None) -> None:
+        """Silence the network (all voltages 0, no spikes) and restart the noise."""
+        xp = self.xp
+        self.rng = xp.random.default_rng(seed)
+        self.v = xp.zeros((self.n, self.batch), xp.float32)
+        self.fired = xp.empty(0, xp.int64)   # flat indices into v
+        self.steps = 0
+        # step of each neuron's last spike, for the refractory period
+        self.last_spike = xp.full((self.n, self.batch), -10**6, xp.int32) if self.refractory_steps else None
+
+    def cells(self, types: list[str], side: str | None = None) -> np.ndarray:
+        """Neurons whose cell type is in `types`. A superclass name
+        ("descending_neuron", "visual_projection", ...) selects the whole class."""
+        mask = np.isin(self.cell_type, types)
+        if self.superclass is not None:
+            mask |= np.isin(self.superclass, types)
+        if side:
+            mask &= self.side == side
+        return np.flatnonzero(mask)
+
+    def _amount(self, amount):
+        """A number, or one value per fly, shaped to broadcast over v[idx]."""
+        a = self.xp.asarray(amount, dtype=self.xp.float32)
+        return a if a.ndim == 0 else a.reshape(1, -1)
+
+    def stimulate(self, idx: np.ndarray, amount) -> None:
+        """Add voltage to these neurons right now (before the next step)."""
+        self.v[self.xp.asarray(idx)] += self._amount(amount)
+
+    def synaptic_input(self, fired):
+        """Input current (n, batch) from the flat spike indices of the last step."""
+        xp, B = self.xp, self.batch
+        if self.device == "cuda":
+            spikes = xp.zeros((self.n, B), xp.float32)
+            spikes.ravel()[fired] = 1.0
+            if B == 1:
+                return (self._W @ spikes[:, 0])[:, None]
+            return self._W @ spikes
+        rows, cols = np.divmod(fired, B)
+        return np.column_stack([_propagate(self.indptr, self.indices, self.weights, rows[cols == b], self.n)
+                                for b in range(B)])
+
+    def step(self, eye_drive: np.ndarray | None = None, inject=()):
+        """Advance one step (dt, 20 ms by default). eye_drive: 0..1 per photoreceptor (len(self.visual)), or
+        (len(self.visual), batch); inject: (neuron indices, extra voltage) pairs
+        added this step. Returns the indices of the neurons that fired (NumPy):
+        one array with batch 1, else a list with one array per fly."""
+        xp, B = self.xp, self.batch
+        current = self.synaptic_input(self.fired) * self.gain
+        self.v *= self.decay
+        self.v += current + self.tonic
+        self.v += (self.rng.random((self.n, B)) < self.noise_hz * self.dt) * np.float32(self.noise_amp)
+        if eye_drive is not None:
+            drive = xp.asarray(eye_drive, dtype=xp.float32)
+            self.v[self._visual] += (drive[:, None] if drive.ndim == 1 else drive) * self.eye_gain
+        for idx, amount in inject:
+            self.v[xp.asarray(idx)] += self._amount(amount)
+        if self.refractory_steps:
+            self.v[(self.steps - self.last_spike) <= self.refractory_steps] = 0.0
+        fired = xp.flatnonzero(self.v >= 1.0)
+        self.v.ravel()[fired] = 0.0
+        if self.refractory_steps:
+            self.last_spike.ravel()[fired] = self.steps
+        self.fired = fired
+        self.steps += 1
+        flat = fired if xp is np else fired.get()
+        if B == 1:
+            return flat
+        rows, cols = np.divmod(flat, B)
+        order = np.argsort(cols, kind="stable")
+        return np.split(rows[order], np.cumsum(np.bincount(cols, minlength=B))[:-1])

flybrain/build.py

@@ -0,0 +1,204 @@
+"""Build a simulatable network from the MaleCNS v1.0 connectome.
+
+Same recipe as ornata/fly (fly64/data.py): all superclass-annotated neurons,
+synapse-count weights, sign from predicted neurotransmitter, per-neuron input
+normalization. Also records where each photoreceptor sits in the eye so a
+1-D "scene" (opponent to the left/right, near/far) can be projected onto it.
+
+Output: <DATA>/brain.npz  (weights as CSR, neuron groups, eye azimuths)
+
+    flybrain build [--data DIR]      (needs pip install "flybrain[build]")
+"""
+from __future__ import annotations
+
+import json
+import re
+import urllib.request
+import zipfile
+import xml.etree.ElementTree as ET
+from pathlib import Path
+
+import numpy as np
+import pyarrow.feather as feather
+from scipy import sparse
+
+from .data import DATA
+
+BUCKET = "https://storage.googleapis.com/flyem-male-cns/v1.0/connectome-data/flat-connectome"
+SOURCES = {
+    "body-annotations-male-cns-v1.0-minconf-0.5.feather": f"{BUCKET}/body-annotations-male-cns-v1.0-minconf-0.5.feather",
+    "body-neurotransmitters-male-cns-v1.0.feather": f"{BUCKET}/body-neurotransmitters-male-cns-v1.0.feather",
+    "connectome-weights-male-cns-v1.0-minconf-0.5.feather": f"{BUCKET}/connectome-weights-male-cns-v1.0-minconf-0.5.feather",
+    "optic-columns.xlsx": "https://raw.githubusercontent.com/flyconnectome/2025malecns/"
+                          "67767d2233657983993ff6c2be48e836a935863c/supplemental_data/optic-column-type-assignments-v1.0.xlsx",
+}
+
+
+def download_all(raw: Path) -> None:
+    """Fetch the MaleCNS v1.0 tables (~1.1 GB) into `raw` unless already present."""
+    raw.mkdir(parents=True, exist_ok=True)
+    for name, url in SOURCES.items():
+        target = raw / name
+        if target.exists():
+            continue
+        partial = target.with_suffix(target.suffix + ".part")
+        print(f"downloading {name}")
+
+        def progress(blocks, block_size, total):
+            if total > 0:
+                print(f"\r  {min(blocks * block_size, total) / 1e6:,.0f} / {total / 1e6:,.0f} MB", end="")
+
+        urllib.request.urlretrieve(url, partial, progress)
+        print()
+        partial.replace(target)
+INHIBITORY = "gaba|glutamate|histamine"
+
+# Descending neurons (brain -> body commands) used as motor read-outs.
+# Types follow the literature: DNa02 steering, DNp01 giant-fiber escape,
+# DNg100 forward walking (as in fly64), MDN = moonwalker (backward walking).
+# punch/kick have no fly equivalent; these picks are arbitrary: DNg11 (a
+# descending neuron with no known fighting role) and pIP10 (the male
+# courtship-song command neuron). Nothing we tested drives either of them.
+MOTOR_TYPES = {
+    "forward": ["DNg100"],
+    "steer": ["DNa02"],
+    "escape": ["DNp01"],
+    "backward": ["MDN"],
+    "punch": ["DNg11"],
+    "kick": ["pIP10"],
+}
+
+
+def optic_columns(path: Path) -> dict[int, tuple[str, int, int]]:
+    """body id -> (eye side, h1, h2) hex column coordinates, from the xlsx."""
+    ns = {"s": "http://schemas.openxmlformats.org/spreadsheetml/2006/main"}
+    result = {}
+    with zipfile.ZipFile(path) as archive:
+        strings = ["".join(e.itertext()) for e in ET.fromstring(archive.read("xl/sharedStrings.xml"))]
+        for sheet in (1, 2):
+            root = ET.fromstring(archive.read(f"xl/worksheets/sheet{sheet}.xml"))
+            for row in root.findall("s:sheetData/s:row", ns)[1:]:
+                cells = {}
+                for c in row:
+                    v = c.find("s:v", ns)
+                    if v is not None:
+                        cells[re.sub(r"\d", "", c.attrib["r"])] = strings[int(v.text)] if c.get("t") == "s" else v.text
+                match = re.fullmatch(r"ME_([LR])_col_(\d+)_(\d+)", cells.get("A", ""))
+                if not match:
+                    continue
+                side, h1, h2 = match.groups()
+                for col in ("B", "C", "E"):
+                    try:
+                        body = int(cells.get(col, -1))
+                    except ValueError:
+                        continue
+                    if body > 0:
+                        result[body] = (side, int(h1), int(h2))
+    return result
+
+
+def build(data: Path | str = DATA) -> None:
+    """Download MaleCNS v1.0 into <data>/raw and write weights.npz, brain.npz and brain.json to `data`."""
+    data = Path(data)
+    raw = data / "raw"
+    download_all(raw)
+    ann = feather.read_table(raw / "body-annotations-male-cns-v1.0-minconf-0.5.feather").to_pandas()
+    nt = feather.read_table(raw / "body-neurotransmitters-male-cns-v1.0.feather",
+                            columns=["body", "consensus_nt"]).to_pandas()
+
+    ann = ann.loc[ann["superclass"].notna() & ann["superclass"].ne("")]
+    ann = ann.drop_duplicates("bodyId").sort_values("bodyId").set_index("bodyId")
+    ids = ann.index.to_numpy(np.int64)
+    n = len(ids)
+    print(f"{n:,} neurons")
+
+    labels = nt.drop_duplicates("body").set_index("body").reindex(ids)["consensus_nt"]
+    sign = np.where(labels.fillna("unclear").str.lower().str.contains(INHIBITORY), -1.0, 1.0).astype(np.float32)
+
+    edges = feather.read_table(raw /"connectome-weights-male-cns-v1.0-minconf-0.5.feather",
+                               columns=["body_pre", "body_post", "weight"], memory_map=True)
+    pre_parts, post_parts, w_parts = [], [], []
+    for i, batch in enumerate(edges.to_batches(max_chunksize=4_000_000), 1):
+        pre_id = batch.column(0).to_numpy(zero_copy_only=False)
+        post_id = batch.column(1).to_numpy(zero_copy_only=False)
+        pre = np.minimum(np.searchsorted(ids, pre_id), n - 1)
+        post = np.minimum(np.searchsorted(ids, post_id), n - 1)
+        ok = (ids[pre] == pre_id) & (ids[post] == post_id)
+        pre_parts.append(pre[ok].astype(np.int32))
+        post_parts.append(post[ok].astype(np.int32))
+        w_parts.append(batch.column(2).to_numpy(zero_copy_only=False)[ok].astype(np.float32))
+        print(f"\r  scanned {min(i * 4_000_000, edges.num_rows):,} / {edges.num_rows:,} edge rows", end="")
+    print()
+    del edges
+    pre, post, w = (np.concatenate(p) for p in (pre_parts, post_parts, w_parts))
+    del pre_parts, post_parts, w_parts
+    w *= sign[pre]
+    incoming = np.bincount(post, weights=np.abs(w), minlength=n).astype(np.float32)
+    w /= np.maximum(incoming[post], 1.0)
+    W = sparse.csr_matrix((w, (post, pre)), shape=(n, n), dtype=np.float32)
+    print(f"{W.nnz:,} connections")
+
+    cell_type = ann["flywireType"].fillna(ann["type"]).fillna("").astype(str)
+    side = ann["somaSide"].fillna(ann["rootSide"]).fillna("").astype(str).str.upper()
+    instance = ann["instance"].fillna("").astype(str).str.upper()
+
+    # Soma (or soma-tract) position of each neuron, in EM voxels; NaN if unknown.
+    positions = np.full((n, 3), np.nan, np.float32)
+    for i, (soma, to_soma) in enumerate(zip(ann["somaLocation"], ann["tosomaLocation"])):
+        loc = soma if isinstance(soma, (list, np.ndarray)) and len(soma) == 3 else to_soma
+        if isinstance(loc, (list, np.ndarray)) and len(loc) == 3:
+            positions[i] = loc
+
+    def pick(types, want_side=None):
+        mask = cell_type.isin(types)
+        if want_side:
+            mask &= (side == want_side) | instance.str.contains(f"_{want_side}")
+        return np.flatnonzero(mask.to_numpy()).astype(np.int32)
+
+    groups = {}
+    for name, types in MOTOR_TYPES.items():
+        groups[f"{name}_L"] = pick(types, "L")
+        groups[f"{name}_R"] = pick(types, "R")
+
+    # Photoreceptors and their azimuth in the eye (-1 = far left ... +1 = far right).
+    visual = pick(["R1-6", "R7", "R8"])
+    columns = optic_columns(raw / "optic-columns.xlsx")
+    known = np.array([i for i, b in enumerate(ids) if int(b) in columns], np.int32)
+    to_known = abs(W[known][:, visual]).tocsc()  # R1-6 -> strongest column-assigned partner
+    h1_max = max(h1 for _, h1, _ in columns.values())
+    azimuth = np.full(len(visual), np.nan, np.float32)
+    eye = np.array([side.iat[v] for v in visual])
+    for k, neuron in enumerate(visual):
+        loc = columns.get(int(ids[neuron]))
+        if loc is None:
+            a, b = to_known.indptr[k:k + 2]
+            if b > a:
+                loc = columns[int(ids[known[to_known.indices[a + np.argmax(to_known.data[a:b])]]])]
+        if loc is not None:
+            eye_side, h1, _ = loc
+            eye[k] = eye_side
+            # h1 runs front->back within an eye (approximate, like fly64).
+            frac = (h1 - 1) / max(h1_max - 1, 1)
+            azimuth[k] = -(0.06 + 0.94 * frac) if eye_side == "L" else (0.06 + 0.94 * frac)
+    missing = np.isnan(azimuth)
+    azimuth[missing] = np.where(eye[missing] == "L", -0.5, 0.5)  # unplaced: mid-eye guess
+    print(f"{len(visual):,} photoreceptors ({missing.sum()} without a column estimate)")
+
+    for name, idx in groups.items():
+        print(f"  {name:11s} {len(idx):3d} neurons  {sorted(set(cell_type.iloc[idx]))}")
+    if any(len(g) == 0 for g in groups.values()):
+        raise RuntimeError("a motor group resolved to zero neurons; check MOTOR_TYPES")
+
+    sparse.save_npz(data / "weights.npz", W, compressed=False)
+    np.savez(data / "brain.npz", ids=ids, visual=visual, azimuth=azimuth,
+             cell_type=cell_type.to_numpy().astype(str), side=side.to_numpy().astype(str),
+             positions=positions, superclass=ann["superclass"].to_numpy().astype(str),
+             **{f"group_{k}": v for k, v in groups.items()})
+    (data / "brain.json").write_text(json.dumps(
+        {"neurons": n, "connections": int(W.nnz), "photoreceptors": int(len(visual)),
+         "groups": {k: int(len(v)) for k, v in groups.items()}}, indent=2))
+    print(f"saved to {data}")
+
+
+if __name__ == "__main__":
+    build()

flybrain/data.py

@@ -0,0 +1,80 @@
+"""Where the brain files live, and fetching them.
+
+FlyBrain needs two files built from the MaleCNS v1.0 connectome: weights.npz (the signed,
+normalized connection matrix) and brain.npz (cell types, sides, positions, readout groups,
+eye layout). They are too big for the package, so the first FlyBrain() downloads a prebuilt
+copy (~260 MB) into $FLY_DATA (default ~/fly-data). `flybrain build` makes the same files
+from the original MaleCNS release instead.
+"""
+from __future__ import annotations
+
+import hashlib
+import os
+import sys
+import urllib.error
+import urllib.request
+from pathlib import Path
+
+DATA = Path(os.environ.get("FLY_DATA", Path.home() / "fly-data"))
+
+RELEASE_URL = os.environ.get("FLYBRAIN_DATA_URL",
+                             "https://github.com/alextitonis/fly.ai/releases/download/brain-v1")
+
+# sha256 of the prebuilt files (166,700 neurons, 25,582,938 connections)
+FILES = {
+    "brain.npz": "cc9bd1ecd00bd703a6fa648bc6ad145c93c7c1ee53debdcc9ce0d1f4305e6aca",
+    "weights.npz": "c29919aa44069a271b1ee978abe05fa9bf6e45e4ba3e436e92b624ef1b5be40c",
+}
+
+
+def has_data(data: Path | str = DATA) -> bool:
+    """True if both brain files are in `data`."""
+    return all((Path(data) / name).exists() for name in FILES)
+
+
+def download(data: Path | str = DATA, url: str = RELEASE_URL, force: bool = False) -> Path:
+    """Fetch the prebuilt brain files into `data`, checking their sha256. Files already
+    there are kept unless force=True."""
+    data = Path(data)
+    data.mkdir(parents=True, exist_ok=True)
+    for name, expected in FILES.items():
+        target = data / name
+        if target.exists() and not force:
+            continue
+        partial = target.with_suffix(target.suffix + ".part")
+        digest = hashlib.sha256()
+        print(f"downloading {name}", file=sys.stderr)
+        try:
+            with urllib.request.urlopen(f"{url.rstrip('/')}/{name}") as response, open(partial, "wb") as out:
+                total = int(response.headers.get("Content-Length") or 0)
+                done = 0
+                while chunk := response.read(1 << 20):
+                    out.write(chunk)
+                    digest.update(chunk)
+                    done += len(chunk)
+                    if total:
+                        print(f"\r  {done / 1e6:,.0f} / {total / 1e6:,.0f} MB", end="", file=sys.stderr)
+            print(file=sys.stderr)
+        except (urllib.error.URLError, OSError) as e:
+            partial.unlink(missing_ok=True)
+            raise RuntimeError(f"could not download {name} from {url}: {e}") from e
+        if digest.hexdigest() != expected:
+            partial.unlink(missing_ok=True)
+            raise RuntimeError(f"{name} from {url} has the wrong checksum; not using it")
+        partial.replace(target)
+    return data
+
+
+def ensure_data(data: Path | str | None = None) -> Path:
+    """The data folder, downloading the prebuilt brain first if it isn't there yet."""
+    data = DATA if data is None else Path(data)
+    if has_data(data):
+        return data
+    print(f"flybrain: no brain files in {data}; downloading the prebuilt brain (~260 MB, once)", file=sys.stderr)
+    try:
+        download(data)
+    except RuntimeError as e:
+        raise FileNotFoundError(
+            f"no brain files in {data} and the download failed ({e}). Build them from the MaleCNS "
+            f"release instead: pip install \"flybrain[build]\" && flybrain build --data \"{data}\"") from e
+    return data

flybrain/eyes.py

@@ -0,0 +1,124 @@
+"""What the fly "sees".
+
+Two routes into the brain:
+* Eyes: a 1-D panorama projected onto the 6,006 photoreceptors. Each has an
+  azimuth (-1 = far left, +1 = far right) estimated from the MaleCNS
+  optic-column tables. Kept for completeness: in a spiking model this signal
+  fades at the lamina (see sweep.py).
+* FeatureDetectors: drive the fly's own visual projection neuron types
+  directly, on the side where things are. This is the route that works.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import numpy as np
+
+BACKGROUND = 0.9
+
+
+@dataclass
+class Blob:
+    center: float      # azimuth, -1..1
+    half_width: float  # azimuth units
+    darkness: float    # 0 = invisible, 1 = black
+
+
+def render(azimuth: np.ndarray, blobs: list[Blob]) -> np.ndarray:
+    lum = np.full(len(azimuth), BACKGROUND, np.float32)
+    for b in blobs:
+        inside = np.abs(azimuth - b.center) <= b.half_width
+        lum[inside] = np.minimum(lum[inside], BACKGROUND * (1 - b.darkness))
+    return lum
+
+
+class Eyes:
+    def __init__(self, azimuth: np.ndarray):
+        self.azimuth = azimuth
+        self.previous: np.ndarray | None = None
+
+    def drive(self, blobs: list[Blob]) -> np.ndarray:
+        lum = render(self.azimuth, blobs)
+        change = np.zeros_like(lum) if self.previous is None else np.abs(lum - self.previous)
+        self.previous = lum
+        return np.clip(0.45 * lum + 1.6 * change, 0, 1)
+
+
+# Encoder parameters (the hand-set defaults). Each can also be an array with one value per
+# fly for a batched brain, so several encoders run side by side (sshfighter/replay.py).
+ENCODER = {
+    "loom_gain": 10.0,    # angular growth per game frame -> LPLC2 voltage
+    "loom_size": 0.0,     # angular size -> LPLC2 voltage (looming neurons also respond to size)
+    "chase_base": 0.6,    # LC10a voltage whenever the opponent is visible...
+    "chase_gain": 0.2,    # ...plus this much per unit of angular size
+    "threat_max": 0.8,    # LC4 voltage when the opponent attacks at close range
+    "shot_gain": 10.0,    # projectile angular growth -> LPLC1 voltage
+    "cap": 0.8,           # most voltage any channel adds per step
+}
+
+# Which identified neuron types each channel drives. The picks are ours, guided by
+# the literature and checked by stimulating each type (all four reach descending
+# neurons): LPLC2 = looming, LC4 = fast looming and escape (strongly drives
+# DNp01/02/04), LPLC1 = small approaching objects, LC10a = the target a male chases.
+CHANNELS = {"loom": ["LPLC2"], "threat": ["LC4"], "shot": ["LPLC1"], "chase": ["LC10a"]}
+
+
+class FeatureDetectors:
+    """Shortcut past the lamina, which a spiking model can't relay (its neurons
+    are graded in real flies). As in Eon's embodied fly, this visual front end
+    is a model; everything downstream of these neurons is the connectome.
+    """
+
+    def __init__(self, brain, **encoder):
+        unknown = set(encoder) - set(ENCODER)
+        if unknown:
+            raise ValueError(f"unknown encoder parameters: {sorted(unknown)}")
+        self.p = {k: np.asarray(encoder.get(k, v), np.float32) for k, v in ENCODER.items()}
+        self.cells = {ch: {s: brain.cells(types, s) for s in "LR"} for ch, types in CHANNELS.items()}
+        self.previous: dict = {}
+        self.last = {f"{ch}{s}": 0.0 for ch in CHANNELS for s in "LR"}   # for display: mean over flies
+
+    @property
+    def loom(self):
+        return self.cells["loom"]
+
+    @property
+    def chase(self):
+        return self.cells["chase"]
+
+    def inject(self, opp=None, shots=(), threat: float = 0.0) -> list:
+        """opp: (dx, size) of the opponent, or None; shots: hostile projectiles as
+        (stable key, dx, size); threat: 0..1, how hard the opponent is attacking
+        right now. dx is in screen units from the fly. Call once per game frame.
+        Amounts are numbers, or one per fly when encoder parameters are arrays."""
+        p = self.p
+        drive = {key: np.float32(0.0) for key in self.last}
+        seen = {}
+
+        def angle_and_growth(key, dx, size):
+            angle = size / max(abs(dx), 8.0)
+            seen[key] = angle
+            return angle, max(0.0, angle - self.previous.get(key, angle))
+
+        if opp is not None:
+            dx, size = opp
+            s = "L" if dx < 0 else "R"
+            angle, growth = angle_and_growth("opp", dx, size)
+            drive[f"loom{s}"] = np.clip(growth * p["loom_gain"] + angle * p["loom_size"], 0, p["cap"])
+            drive[f"chase{s}"] = np.clip(p["chase_base"] + p["chase_gain"] * angle, 0, p["cap"])
+            drive[f"threat{s}"] = p["threat_max"] * np.float32(np.clip(threat, 0, 1))
+        for key, dx, size in shots:
+            s = "L" if dx < 0 else "R"
+            _, growth = angle_and_growth(key, dx, size)
+            drive[f"shot{s}"] = np.maximum(drive[f"shot{s}"], np.clip(growth * p["shot_gain"], 0, p["cap"]))
+        self.previous = seen
+        self.last = {key: float(np.mean(amount)) for key, amount in drive.items()}
+        return [(self.cells[key[:-1]][key[-1]], amount) for key, amount in drive.items() if np.any(amount > 0)]
+
+
+def blob_for(dx: float, size: float, darkness: float) -> Blob:
+    """An object `dx` world units to the side (screen coordinates) of the fly."""
+    distance = max(abs(dx), 8.0)
+    return Blob(center=float(np.clip(dx / 110.0, -1, 1)),
+                half_width=float(np.clip(size / distance * 0.5, 0.03, 0.7)),
+                darkness=darkness)

flybrain/reservoir.py

@@ -0,0 +1,275 @@
+"""Generic reservoir computing on the fly connectome.
+
+    input -> encoder -> fly brain (frozen) -> trace -> trained readout -> output
+
+`flybrain.FlyBrain` is never trained: its weights come straight from the connectome.
+This module is everything task-agnostic around it:
+
+  * `Trace` turns the neurons that fire each step into a decaying feature vector,
+    for any neuron population you name (a cell type, a `brain.groups[...]` set,
+    or your own index array).
+  * `run` steps the brain over a sequence of inputs and collects a `Trace`,
+    so the whole loop is one call from a notebook.
+  * `Readout` fits a linear (ridge) or logistic PCA readout from that activity to
+    your own labels, model-selected by cross-validation, and predicts on new
+    activity.
+
+Nothing here knows about SSH Fighter, or any other task. `sshfighter/reservoir.py`
+is a worked example of using this module for one game; `flyreservoir_example.py`
+at the repo root is a minimal one with synthetic data, runnable without any game
+or recordings.
+
+Encoders (turning task input into neuron drive) are necessarily task-specific --
+you write them by picking `brain.cells([...types], side=...)` and passing
+`(idx, amount)` pairs to `brain.step(inject=...)`. `flybrain/eyes.py` is a worked
+example of a visual encoder for SSH Fighter.
+"""
+from __future__ import annotations
+
+from pathlib import Path
+
+import numpy as np
+from scipy.optimize import minimize
+from scipy.special import expit
+
+__all__ = ["Trace", "run", "bases_for", "project", "fit_logistic", "fit_ridge", "auc", "folds", "Readout"]
+
+
+# ---- trace: turns spikes into features ---------------------------------------------------
+
+class Trace:
+    """Exponentially decaying spike trace of a neuron population.
+
+    Pick the population by `types` (a list of MaleCNS cell types, via `brain.cells`),
+    `group` (a name from `brain.groups`), or your own `idx` array. `side` restricts
+    `types` to "L" or "R" (see `FlyBrain.cells`).
+
+    With a batched brain (`batch > 1`, several flies stepped together), `aggregate`
+    controls what `features()` returns: "mean" (default) is one vector, the flies'
+    average trace -- the natural choice when the flies vote on one shared decision;
+    "batch" keeps one column per fly, shape (n_neurons, batch), for independent
+    per-fly readouts.
+    """
+
+    def __init__(self, brain, types: list[str] | None = None, group: str | None = None,
+                 idx: np.ndarray | None = None, side: str | None = None, tau: float = 0.1,
+                 aggregate: str = "mean"):
+        if sum(x is not None for x in (types, group, idx)) != 1:
+            raise ValueError("pass exactly one of types, group, idx")
+        if idx is None:
+            idx = brain.groups[group] if group is not None else brain.cells(types, side=side)
+        if aggregate not in ("mean", "batch"):
+            raise ValueError('aggregate must be "mean" or "batch"')
+        self.idx = np.asarray(idx)
+        self.slot = np.full(brain.n, -1, np.int64)
+        self.slot[self.idx] = np.arange(len(self.idx))
+        self.aggregate = aggregate
+        self.batch = brain.batch
+        width = len(self.idx)
+        self.trace = np.zeros(width if aggregate == "mean" else (width, self.batch), np.float32)
+        self.decay = np.float32(np.exp(-brain.dt / tau))
+
+    def observe(self, fired) -> np.ndarray:
+        """`fired`: whatever `FlyBrain.step()` returned -- one array of spike indices
+        (batch 1), or a list of them, one per fly. Returns `features()`."""
+        self.trace *= self.decay
+        flies = fired if isinstance(fired, list) else [fired]
+        if self.aggregate == "mean":
+            for f in flies:
+                slots = self.slot[f]
+                self.trace[slots[slots >= 0]] += 1.0 / len(flies)
+        else:
+            for b, f in enumerate(flies):
+                slots = self.slot[f]
+                self.trace[slots[slots >= 0], b] += 1.0
+        return self.features()
+
+    def features(self) -> np.ndarray:
+        return self.trace.copy()
+
+    def reset(self) -> None:
+        self.trace[...] = 0
+
+
+def run(brain, steps: int, encode=None, trace: Trace | None = None, eye_drive=None) -> np.ndarray:
+    """Step the brain `steps` times and collect activity, so a whole task can run
+    from one call: `activity = flybrain.run(brain, len(inputs), encode=...)`.
+
+    `encode(t)`, if given, returns the `inject` list for step `t` (see
+    `FlyBrain.step`) -- typically `encoder.inject(...)` for some task-specific
+    encoder built on `brain.cells(...)`. `eye_drive`, if given, is either an
+    array (same every step) or `eye_drive(t)`. `trace` defaults to a `Trace`
+    over every neuron; pass your own to watch a specific population.
+
+    Returns activity stacked over time: `(steps, n_features)`, or
+    `(steps, n_features, batch)` if `trace.aggregate == "batch"`.
+    """
+    if trace is None:
+        trace = Trace(brain, idx=np.arange(brain.n))
+    out = []
+    for t in range(steps):
+        inject = encode(t) if encode is not None else ()
+        drive = eye_drive(t) if callable(eye_drive) else eye_drive
+        fired = brain.step(eye_drive=drive, inject=inject)
+        out.append(trace.observe(fired))
+    return np.stack(out)
+
+
+# ---- PCA basis shared by training and inference ------------------------------------------
+
+def bases_for(X: np.ndarray, ks) -> dict:
+    """For each k in `ks`: (mean, projection, scale) of the top-k principal
+    components of X's columns. k=None just standardises every column (no PCA)."""
+    mu = X.mean(0)
+    out, vt = {}, None
+    for k in ks:
+        if k is None:
+            P = np.eye(X.shape[1], dtype=np.float32)
+        else:
+            if vt is None:
+                vt = np.linalg.svd(X - mu, full_matrices=False)[2]
+            P = vt[:min(k, len(vt))].T.astype(np.float32)
+        sd = ((X - mu) @ P).std(0) + 1e-6
+        out[k] = (mu, P, sd)
+    return out
+
+
+def project(basis, X: np.ndarray) -> np.ndarray:
+    mu, P, sd = basis
+    return ((X - mu) @ P) / sd
+
+
+# ---- fitting: linear (ridge) and logistic readouts ----------------------------------------
+
+def fit_logistic(Z: np.ndarray, y: np.ndarray, lam: float):
+    """L2-regularised logistic regression, y in {0, 1}. Returns (weights, bias)."""
+    def objective(w):
+        z = Z @ w[:-1] + w[-1]
+        loss = np.mean(np.logaddexp(0, z) - y * z) + 0.5 * lam * w[:-1] @ w[:-1]
+        r = (expit(z) - y) / len(y)
+        return loss, np.append(Z.T @ r + lam * w[:-1], r.sum())
+
+    w = minimize(objective, np.zeros(Z.shape[1] + 1), jac=True, method="L-BFGS-B").x
+    return w[:-1], w[-1]
+
+
+def fit_ridge(Z: np.ndarray, y: np.ndarray, lam: float):
+    """Ridge regression. y is (n,) for one output or (n, outputs) for several
+    (e.g. one column per class, fit jointly). Returns (weights, bias): weights
+    is (features,) or (outputs, features) to match y; bias matches in the same way."""
+    single = y.ndim == 1
+    y2 = y[:, None] if single else y
+    zm, ym = Z.mean(0), y2.mean(0)
+    Zc = Z - zm
+    W = np.linalg.solve(Zc.T @ Zc + lam * len(y2) * np.eye(Z.shape[1]), Zc.T @ (y2 - ym))
+    b = ym - zm @ W
+    return (W[:, 0], float(b[0])) if single else (W.T, b)
+
+
+def auc(y: np.ndarray, s: np.ndarray) -> float:
+    """Area under the ROC curve. nan if y has only one class."""
+    pos, neg = y == 1, y == 0
+    if pos.sum() == 0 or neg.sum() == 0:
+        return float("nan")
+    ranks = np.empty(len(s))
+    ranks[np.argsort(s, kind="stable")] = np.arange(1, len(s) + 1)
+    return (ranks[pos].sum() - pos.sum() * (pos.sum() + 1) / 2) / (pos.sum() * neg.sum())
+
+
+def folds(n: int, groups: np.ndarray | None = None, k: int = 5):
+    """Cross-validation splits over n samples, as (train_idx, test_idx) pairs.
+    With `groups` (one label per sample, e.g. a match or recording id), this is
+    leave-one-group-out. Without it, plain k-fold."""
+    if groups is not None:
+        groups = np.asarray(groups)
+        for g in np.unique(groups):
+            test = np.flatnonzero(groups == g)
+            train = np.flatnonzero(groups != g)
+            if len(train):
+                yield train, test
+    else:
+        idx = np.arange(n)
+        for part in np.array_split(idx, min(k, n)):
+            if len(part) == 0 or len(part) == n:
+                continue
+            test = part
+            train = np.setdiff1d(idx, test, assume_unique=True)
+            yield train, test
+
+
+class Readout:
+    """A PCA + linear readout: activity -> label, with the PCA rank and L2
+    strength picked by cross-validation.
+
+    `kind="ridge"` regresses onto a continuous or multi-column target (mean
+    squared error). `kind="logistic"` classifies a binary target in {0, 1}
+    (held-out AUC). Both compress activity to its top principal components
+    first (see `bases_for`) -- this is what makes it a *reservoir* readout: the
+    brain does the nonlinear mixing, the readout only has to be linear on that.
+    """
+
+    def __init__(self, kind: str, basis, w, b, cv_score: float, components: int | None, lam: float):
+        self.kind, self.basis, self.w, self.b = kind, basis, w, b
+        self.cv_score, self.components, self.lam = cv_score, components, lam
+
+    @classmethod
+    def fit(cls, X: np.ndarray, y: np.ndarray, kind: str = "ridge", groups: np.ndarray | None = None,
+            components=(5, 20, 60), lambdas=(1e-2, 1e-1, 1.0, 10.0), verbose: bool = False) -> "Readout":
+        """X: (n_samples, n_features) activity, e.g. from `Trace.features()` stacked
+        over time. y: (n_samples,) labels -- 0/1 for `kind="logistic"`, numeric
+        (or (n_samples, outputs)) for `kind="ridge"`. `groups`, if given, makes
+        cross-validation leave-one-group-out (e.g. one group per recorded episode)
+        instead of plain k-fold, so the score isn't inflated by correlated samples
+        from the same episode landing in both train and test.
+
+        Every (k, lambda) in the grid is scored by cross-validated held-out AUC
+        (logistic) or negative MSE (ridge); the best is refit on all of X, y.
+        """
+        if kind not in ("ridge", "logistic"):
+            raise ValueError('kind must be "ridge" or "logistic"')
+        fit_fn = fit_logistic if kind == "logistic" else fit_ridge
+        grid = [(k, lam) for k in components for lam in lambdas]
+        scores = {g: [] for g in grid}
+        for train, test in folds(len(X), groups):
+            Xtr, ytr, Xte, yte = X[train], y[train], X[test], y[test]
+            if kind == "logistic" and (ytr.min() == ytr.max() or yte.min() == yte.max()):
+                continue
+            bases = bases_for(Xtr, components)
+            for k, lam in grid:
+                w, b = fit_fn(project(bases[k], Xtr), ytr, lam)
+                pred = project(bases[k], Xte) @ w.T + b if w.ndim > 1 else project(bases[k], Xte) @ w + b
+                score = auc(yte, expit(pred)) if kind == "logistic" else -np.mean((pred - yte) ** 2)
+                if not np.isnan(score):
+                    scores[(k, lam)].append(score)
+        avg = {g: (float(np.mean(s)) if s else float("-inf")) for g, s in scores.items()}
+        k, lam = max(avg, key=avg.get)
+        basis = bases_for(X, [k])[k]
+        w, b = fit_fn(project(basis, X), y, lam)
+        if verbose:
+            print(f"readout: {k} components, lambda {lam:g}, "
+                  f"cross-validated {'AUC' if kind == 'logistic' else 'neg-MSE'} {avg[(k, lam)]:.3f}")
+        return cls(kind, basis, w, b, avg[(k, lam)], k, lam)
+
+    def predict(self, x: np.ndarray) -> np.ndarray | float:
+        """x: one sample (n_features,) or a batch (n_samples, n_features). Returns
+        the raw ridge output, or the logistic probability, in the same shape."""
+        single = x.ndim == 1
+        z = project(self.basis, x[None] if single else x)
+        raw = z @ self.w.T + self.b if self.w.ndim > 1 else z @ self.w + self.b
+        out = expit(raw) if self.kind == "logistic" else raw
+        return out[0] if single else out
+
+    def save(self, path: str | Path) -> None:
+        np.savez(path, kind=self.kind, mu=self.basis[0], P=self.basis[1], sd=self.basis[2],
+                 w=self.w, b=np.asarray(self.b), cv_score=self.cv_score,
+                 components=self.components if self.components is not None else -1, lam=self.lam)
+
+    @classmethod
+    def load(cls, path: str | Path) -> "Readout":
+        d = np.load(path, allow_pickle=False)
+        basis = (d["mu"], d["P"], d["sd"])
+        b = d["b"]
+        b = float(b) if b.ndim == 0 else b
+        components = int(d["components"])
+        return cls(str(d["kind"]), basis, d["w"], b, float(d["cv_score"]),
+                    None if components < 0 else components, float(d["lam"]))

flybrain-0.1.0.dist-info/METADATA

@@ -0,0 +1,87 @@
+Metadata-Version: 2.4
+Name: flybrain
+Version: 0.1.0
+Summary: The complete fruit fly nervous system (MaleCNS connectome, 166,700 neurons) as a spiking network, on CPU or GPU.
+Author: alextitonis
+License-Expression: MIT
+Project-URL: Homepage, https://flyaiworld.com
+Project-URL: Source, https://github.com/alextitonis/fly.ai
+Keywords: connectome,drosophila,spiking neural network,neuroscience,reservoir computing
+Classifier: Programming Language :: Python :: 3
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Science/Research
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy>=2.0
+Requires-Dist: scipy>=1.13
+Requires-Dist: numba>=0.61
+Provides-Extra: gpu
+Requires-Dist: cupy-cuda12x[ctk]>=13; extra == "gpu"
+Provides-Extra: build
+Requires-Dist: pandas>=2.2; extra == "build"
+Requires-Dist: pyarrow>=15; extra == "build"
+Dynamic: license-file
+
+# flybrain
+
+The complete central nervous system of an adult male fruit fly, *Drosophila melanogaster*, as a
+spiking network you can run on your own computer: **166,700 neurons and 25.6 million
+connections** from the [MaleCNS v1.0 connectome](https://male-cns.janelia.org), wired as
+electron microscopy found them.
+
+Nothing inside the brain is trained. You drive some of the fly's own neurons, step the network
+forward, and read out what its descending neurons (the brain's commands to the body) do.
+
+```sh
+pip install flybrain              # CPU (numba)
+pip install "flybrain[gpu]"       # plus CuPy for an NVIDIA GPU (CUDA 12)
+```
+
+```python
+from flybrain import FlyBrain
+
+brain = FlyBrain(device="auto")          # first run downloads the brain files (~260 MB) to ~/fly-data
+left_loom = brain.cells(["LC4", "LPLC2"], side="L")   # looming detectors, left eye
+giant_fiber = brain.cells(["DNp01"], side="L")        # the escape command neuron
+
+for step in range(50):                   # one second at 20 ms per step
+    fired = brain.step(inject=[(left_loom, 0.8)])
+    if set(giant_fiber) & set(fired):
+        print(f"left giant fiber fired at {step * brain.dt:.2f} s")
+```
+
+## What's in it
+
+* `FlyBrain`: leaky integrate-and-fire over the whole connectome. `device="cpu" | "cuda" | "auto"`,
+  `batch=8` runs 8 independent flies at once, plus `dt`, `sensory_input` and `refractory` options.
+  `brain.cells([...])` finds neurons by cell type or superclass (`"descending_neuron"`).
+* `Trace`, `run`, `Readout`: reservoir computing. Collect a spike trace of any neuron population
+  over your task, then fit a cross-validated linear or logistic PCA readout to your labels.
+* `Eyes`, `FeatureDetectors`: a visual encoder that drives the fly's visual projection neurons.
+
+## Data
+
+The brain files live in `$FLY_DATA` (default `~/fly-data`). The first `FlyBrain()` downloads them;
+you can also run it ahead of time:
+
+```sh
+flybrain download                 # prebuilt files, sha256-checked
+flybrain build                    # or build them from the MaleCNS release (~1.1 GB; pip install "flybrain[build]")
+flybrain info                     # data folder and GPU status
+```
+
+The first step on CPU is slow while numba compiles; later steps take about 12–15 ms on 24 threads.
+On an RTX 4060 a step takes 1.4 ms.
+
+## Credits and license
+
+Code: MIT. The connectome data is MaleCNS v1.0 by FlyEM (HHMI Janelia), the University of
+Cambridge, the MRC Laboratory of Molecular Biology and Google Research, used under
+[CC BY 4.0](https://male-cns.janelia.org/download/). If you use it, cite Berg, S. et al. (2026),
+*Sexual dimorphism in the complete connectome of the Drosophila male central nervous system*, *Cell*.
+The neuron model follows [Fly64](https://github.com/ornata/fly) by Jessica Paquette.
+
+Source, experiments and results: [github.com/alextitonis/fly.ai](https://github.com/alextitonis/fly.ai)

flybrain-0.1.0.dist-info/RECORD

@@ -0,0 +1,13 @@
+flybrain/__init__.py,sha256=q1U72lEWZfnOJO8-jYNytGU4CuEpdckE6-XT5YVEkBA,1002
+flybrain/__main__.py,sha256=eP9lEQhK5FyOHIlvu47TbaIo_Df53u1i6xSRws4HLx4,2426
+flybrain/brain.py,sha256=FC7P_9XTLjBlWCnjSziwQCTkm-rmYr0HgnnpYQx4mww,9406
+flybrain/build.py,sha256=NNcWdFN8fxIGuSyPIz3L07h405rd37t6iRa2LS9MHpU,9634
+flybrain/data.py,sha256=ePvdQ2AUYvNhUOnoticZ6S57Qi7YFYHD96ipdUyzHYg,3447
+flybrain/eyes.py,sha256=lfBgPXbuWDvaLysO70xcz2D6LHZi466qEGQoYR3KgX4,5455
+flybrain/reservoir.py,sha256=48h7bM5Hhc530SkoM4NfAl_PWssZGzN-eeov2xxy1o0,12687
+flybrain-0.1.0.dist-info/licenses/LICENSE,sha256=yOLVyKQYhZFHyoP0vyf7SG4_vvm__UglSFIQ5oeScRY,1068
+flybrain-0.1.0.dist-info/METADATA,sha256=YIN3o-ql7ORfaERK_e0OV_f5kXQmMBanXxsp2e_6OiQ,4044
+flybrain-0.1.0.dist-info/WHEEL,sha256=YVMoNqKzERt-wjUZwJ33xBGAwnFl-4cqbYkTtWa4itE,91
+flybrain-0.1.0.dist-info/entry_points.txt,sha256=EKDxAS8z8EdppS_-TrSIPR75w7nutTFVbEHFxrpxbY0,52
+flybrain-0.1.0.dist-info/top_level.txt,sha256=ZDXUud3CZilW1jhzGIqBix0bKnpxLJx21KRbCLwZy9I,9
+flybrain-0.1.0.dist-info/RECORD,,

flybrain-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (84.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

flybrain-0.1.0.dist-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+flybrain = flybrain.__main__:main

flybrain-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 alextitonis
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

flybrain-0.1.0.dist-info/top_level.txt

@@ -0,0 +1 @@
+flybrain