flywire-gnn 0.1.0__tar.gz

flywire_gnn-0.1.0/PKG-INFO

@@ -0,0 +1,114 @@
+Metadata-Version: 2.4
+Name: flywire-gnn
+Version: 0.1.0
+Summary: The FlyWire FAFB Drosophila connectome as a ready-to-train PyTorch Geometric dataset, with baseline models and reproducible benchmarks.
+Author: Omkar
+License: MIT
+Project-URL: FlyWire, https://flywire.ai
+Project-URL: Zenodo, https://doi.org/10.5281/zenodo.10676866
+Keywords: connectome,flywire,drosophila,gnn,graph-neural-network,neuroscience,pytorch,benchmark
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.24
+Requires-Dist: pandas>=2.0
+Requires-Dist: pyarrow>=14.0
+Requires-Dist: torch>=2.0
+Requires-Dist: torch_geometric>=2.4
+Requires-Dist: scikit-learn>=1.3
+Provides-Extra: hub
+Requires-Dist: huggingface_hub>=0.20; extra == "hub"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+
+# flywire-gnn
+
+**The complete fruit-fly brain connectome (FlyWire FAFB v783) as a ready-to-train PyTorch Geometric dataset** — 139,255 proofread neurons, 2.7M directed synaptic connections, cell-type labels on every neuron, deterministic stratified splits, and three reproducible baselines. One `pip install`, one class, no authentication, no left-over data-wrangling.
+
+## Install
+
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cpu
+pip install torch_geometric
+pip install git+https://github.com/omkar-dhakane/flywire-gnn.git   # or: pip install -e .
+```
+
+Tested with Python 3.13, torch 2.14.0+cpu, torch_geometric 2.8.0.post1, pyarrow 22.0, numpy 2.3.1. Older `torch>=2.0` / `torch_geometric>=2.4` should work; the benchmark numbers above are from the tested versions.
+
+## Quickstart (the whole thing)
+
+```python
+from flywire_gnn import FlyWireFAFB
+
+ds = FlyWireFAFB()                    # downloads 852MB + 1.1MB + 27MB once, caches locally
+data = ds.data                        # torch_geometric.data.Data
+train_mask, val_mask, test_mask = ds.splits()   # deterministic, stratified, seed 42
+print(data)      # Data(x=[139255, 174], edge_index=[2, 2700513], edge_attr=[2700513, 7], y=[139255])
+```
+
+Everything is cached in `~/.cache/flywire_gnn` — subsequent loads are instant.
+
+## The dataset
+
+| | |
+|---|---|
+| Nodes | 139,255 proofread neurons (whole adult female fly brain) |
+| Edges | 2,700,513 unique directed pairs at ≥5 synapses (15,091,983 at ≥1) |
+| Node features (`x`, 174-dim) | wiring structure: degrees/synapse totals, per-neuropil in/out profiles across 79 brain regions, in/out neurotransmitter profiles |
+| Edge attributes (`edge_attr`, 7-dim) | `log1p(syn_count)` + 6 synapse-weighted neurotransmitter probabilities |
+| Labels (`y`) | `super_class`: 9 classes covering **100% of nodes** (also available: `cell_class`, `cell_sub_class`, `cell_type`) |
+| Task | node classification from wiring structure |
+| Splits | 70/15/15 train/val/test, stratified, seed 42 |
+
+Data sources (all public, **no auth needed**):
+- Connectivity: FlyWire Whole-brain Connectome Connectivity Data v783 — [Zenodo, CC-BY-4.0](https://doi.org/10.5281/zenodo.10676866)
+- Annotation: [Schlegel et al. 2024](https://doi.org/10.1038/s41586-024-07686-5), Supplementary Data 5
+- Packaged copies (parquet, pair-level): [huggingface.co/datasets/SLOP011/flywire-fafb-connectome](https://huggingface.co/datasets/SLOP011/flywire-fafb-connectome)
+
+## Leaderboard — node classification on `super_class`
+
+Real runs, full-batch, CPU, seed 42, `min_synapses=5`, hidden 128 (see `flywire_gnn/train.py`):
+
+| Model | Test accuracy | Macro-F1 | Best val acc (epoch) | Epochs | Wall time |
+|---|---|---|---|---|---|
+| MLP (features only) | **0.9851** | 0.7115 | 0.9860 (200) | 200 | 3.5 min |
+| GraphSAGE | 0.9812 | **0.7563** | 0.9820 (150) | 150 | 13 min |
+| GCN | 0.9166 | 0.4702 | 0.9202 (180) | 200 | 23.5 min |
+
+**What the benchmark shows**: wiring-profile features alone nearly saturate accuracy (98.5%) — most of a cell's coarse class is readable directly from its projection pattern. Mean-aggregated message passing (GraphSAGE) roughly matches features and wins on the rare classes (best macro-F1); GCN's symmetric normalization over-smooths and trails. If your architecture can't beat 0.9851 accuracy *and* 0.7563 macro-F1 on these exact splits, it isn't adding anything over a feature baseline.
+
+### Reproduce
+
+```bash
+python -m flywire_gnn.train --models mlp,gcn,sage          # all three
+python -m flywire_gnn.train --model sage --epochs 150 --seed 42
+python -m flywire_gnn.train --model sage --labels cell_class   # harder multi-class task
+python -m flywire_gnn.train --model sage --min-synapses 1       # dense 15M-edge graph
+```
+
+## Design choices (and why)
+
+- **`min_synapses=5`** is the published FAFB "connection" threshold (same default as Codex). Pass `min_synapses=1` for the full-density graph.
+- **Node features come from the ≥1-synapse wiring** (a neuron's total projection profile); edges are filtered by `min_synapses`. So features are the same regardless of the edge threshold you benchmark.
+- **Stratified splits on labeled nodes only**: with `super_class`, all 139,255 nodes are labeled. Deterministic via `seed=42` (NumPy `RandomState` + sklearn stratified split).
+- The raw 9.5 GB per-synapse file (`flywire_synapses_783.feather`) is **not** needed: the 852 MB pair×neuropil table is sufficient.
+
+## Not in v1 (deliberately)
+
+Mesh/skeleton loading (use `fafbseg` + meshparty), other connectomes (MANC/MAOL/MCNS/BANC), hosted leaderboard server, spiking neural simulation, per-synapse link prediction. The package is intentionally complete at this scope.
+
+## Tests
+
+```bash
+pytest tests/ -v
+```
+
+Covers: graph assembly counts on a synthetic graph, feature finiteness, split disjointness/determinism, model forward + backward smoke tests, and (if the real cache is present) the real 139,255-node / 2,700,513-edge invariants.
+
+## License & citation
+
+Code: MIT. Data: **CC-BY-4.0**. Using the dataset means citing:
+
+1. **Dorkenwald et al.** 2024. *Neuronal wiring diagram of an adult brain.* Nature. [doi:10.1038/s41586-024-07558-y](https://doi.org/10.1038/s41586-024-07558-y)
+2. **Schlegel et al.** 2024. *Whole-brain annotation and multi-connectome cell typing of Drosophila.* Nature. [doi:10.1038/s41586-024-07686-5](https://doi.org/10.1038/s41586-024-07686-5)
+
+Not affiliated with the FlyWire Consortium. Interactive exploration: [codex.flywire.ai](https://codex.flywire.ai) — analysis in Python: [navis](https://github.com/navis-org/navis) / [fafbseg](https://github.com/navis-org/fafbseg-py).

flywire_gnn-0.1.0/README.md

@@ -0,0 +1,92 @@
+# flywire-gnn
+
+**The complete fruit-fly brain connectome (FlyWire FAFB v783) as a ready-to-train PyTorch Geometric dataset** — 139,255 proofread neurons, 2.7M directed synaptic connections, cell-type labels on every neuron, deterministic stratified splits, and three reproducible baselines. One `pip install`, one class, no authentication, no left-over data-wrangling.
+
+## Install
+
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cpu
+pip install torch_geometric
+pip install git+https://github.com/omkar-dhakane/flywire-gnn.git   # or: pip install -e .
+```
+
+Tested with Python 3.13, torch 2.14.0+cpu, torch_geometric 2.8.0.post1, pyarrow 22.0, numpy 2.3.1. Older `torch>=2.0` / `torch_geometric>=2.4` should work; the benchmark numbers above are from the tested versions.
+
+## Quickstart (the whole thing)
+
+```python
+from flywire_gnn import FlyWireFAFB
+
+ds = FlyWireFAFB()                    # downloads 852MB + 1.1MB + 27MB once, caches locally
+data = ds.data                        # torch_geometric.data.Data
+train_mask, val_mask, test_mask = ds.splits()   # deterministic, stratified, seed 42
+print(data)      # Data(x=[139255, 174], edge_index=[2, 2700513], edge_attr=[2700513, 7], y=[139255])
+```
+
+Everything is cached in `~/.cache/flywire_gnn` — subsequent loads are instant.
+
+## The dataset
+
+| | |
+|---|---|
+| Nodes | 139,255 proofread neurons (whole adult female fly brain) |
+| Edges | 2,700,513 unique directed pairs at ≥5 synapses (15,091,983 at ≥1) |
+| Node features (`x`, 174-dim) | wiring structure: degrees/synapse totals, per-neuropil in/out profiles across 79 brain regions, in/out neurotransmitter profiles |
+| Edge attributes (`edge_attr`, 7-dim) | `log1p(syn_count)` + 6 synapse-weighted neurotransmitter probabilities |
+| Labels (`y`) | `super_class`: 9 classes covering **100% of nodes** (also available: `cell_class`, `cell_sub_class`, `cell_type`) |
+| Task | node classification from wiring structure |
+| Splits | 70/15/15 train/val/test, stratified, seed 42 |
+
+Data sources (all public, **no auth needed**):
+- Connectivity: FlyWire Whole-brain Connectome Connectivity Data v783 — [Zenodo, CC-BY-4.0](https://doi.org/10.5281/zenodo.10676866)
+- Annotation: [Schlegel et al. 2024](https://doi.org/10.1038/s41586-024-07686-5), Supplementary Data 5
+- Packaged copies (parquet, pair-level): [huggingface.co/datasets/SLOP011/flywire-fafb-connectome](https://huggingface.co/datasets/SLOP011/flywire-fafb-connectome)
+
+## Leaderboard — node classification on `super_class`
+
+Real runs, full-batch, CPU, seed 42, `min_synapses=5`, hidden 128 (see `flywire_gnn/train.py`):
+
+| Model | Test accuracy | Macro-F1 | Best val acc (epoch) | Epochs | Wall time |
+|---|---|---|---|---|---|
+| MLP (features only) | **0.9851** | 0.7115 | 0.9860 (200) | 200 | 3.5 min |
+| GraphSAGE | 0.9812 | **0.7563** | 0.9820 (150) | 150 | 13 min |
+| GCN | 0.9166 | 0.4702 | 0.9202 (180) | 200 | 23.5 min |
+
+**What the benchmark shows**: wiring-profile features alone nearly saturate accuracy (98.5%) — most of a cell's coarse class is readable directly from its projection pattern. Mean-aggregated message passing (GraphSAGE) roughly matches features and wins on the rare classes (best macro-F1); GCN's symmetric normalization over-smooths and trails. If your architecture can't beat 0.9851 accuracy *and* 0.7563 macro-F1 on these exact splits, it isn't adding anything over a feature baseline.
+
+### Reproduce
+
+```bash
+python -m flywire_gnn.train --models mlp,gcn,sage          # all three
+python -m flywire_gnn.train --model sage --epochs 150 --seed 42
+python -m flywire_gnn.train --model sage --labels cell_class   # harder multi-class task
+python -m flywire_gnn.train --model sage --min-synapses 1       # dense 15M-edge graph
+```
+
+## Design choices (and why)
+
+- **`min_synapses=5`** is the published FAFB "connection" threshold (same default as Codex). Pass `min_synapses=1` for the full-density graph.
+- **Node features come from the ≥1-synapse wiring** (a neuron's total projection profile); edges are filtered by `min_synapses`. So features are the same regardless of the edge threshold you benchmark.
+- **Stratified splits on labeled nodes only**: with `super_class`, all 139,255 nodes are labeled. Deterministic via `seed=42` (NumPy `RandomState` + sklearn stratified split).
+- The raw 9.5 GB per-synapse file (`flywire_synapses_783.feather`) is **not** needed: the 852 MB pair×neuropil table is sufficient.
+
+## Not in v1 (deliberately)
+
+Mesh/skeleton loading (use `fafbseg` + meshparty), other connectomes (MANC/MAOL/MCNS/BANC), hosted leaderboard server, spiking neural simulation, per-synapse link prediction. The package is intentionally complete at this scope.
+
+## Tests
+
+```bash
+pytest tests/ -v
+```
+
+Covers: graph assembly counts on a synthetic graph, feature finiteness, split disjointness/determinism, model forward + backward smoke tests, and (if the real cache is present) the real 139,255-node / 2,700,513-edge invariants.
+
+## License & citation
+
+Code: MIT. Data: **CC-BY-4.0**. Using the dataset means citing:
+
+1. **Dorkenwald et al.** 2024. *Neuronal wiring diagram of an adult brain.* Nature. [doi:10.1038/s41586-024-07558-y](https://doi.org/10.1038/s41586-024-07558-y)
+2. **Schlegel et al.** 2024. *Whole-brain annotation and multi-connectome cell typing of Drosophila.* Nature. [doi:10.1038/s41586-024-07686-5](https://doi.org/10.1038/s41586-024-07686-5)
+
+Not affiliated with the FlyWire Consortium. Interactive exploration: [codex.flywire.ai](https://codex.flywire.ai) — analysis in Python: [navis](https://github.com/navis-org/navis) / [fafbseg](https://github.com/navis-org/fafbseg-py).

flywire_gnn-0.1.0/flywire_gnn/__init__.py

@@ -0,0 +1,6 @@
+from .dataset import FlyWireFAFB
+from .models import GCN, GraphSAGE, MLP
+
+__version__ = "0.1.0"
+
+__all__ = ["FlyWireFAFB", "MLP", "GCN", "GraphSAGE", "__version__"]

flywire_gnn-0.1.0/flywire_gnn/dataset.py

@@ -0,0 +1,441 @@
+"""The FlyWire FAFB (v783) Drosophila connectome as a PyTorch Geometric dataset.
+
+Data sources (all public, no authentication required):
+- Connectivity: FlyWire Whole-brain Connectome Connectivity Data v783
+  (Zenodo, CC-BY-4.0, doi:10.5281/zenodo.10676866), Dorkenwald et al., Nature 2024
+- Cell annotations: Supplementary Data 5 of Schlegel et al., Nature 2024
+  (doi:10.1038/s41586-024-07686-5)
+
+Graph layout:
+- x:          float32 [N, 174] node features derived from wiring
+- edge_index: int64   [2, E] directed edges (pre -> post), pair-level, filtered
+              to synapses >= min_synapses (default 5, the published FAFB
+              "connection" threshold)
+- edge_attr:  float32 [E, 7]: [log1p(syn_count), 6 weighted NT probabilities]
+- y:          int64 [N] label index (-1 = unlabeled)
+"""
+
+from __future__ import annotations
+
+import hashlib
+import os
+import time
+import urllib.request
+from pathlib import Path
+from typing import Dict, Optional, Tuple
+
+import numpy as np
+import pandas as pd
+import pyarrow as pa
+import pyarrow.compute as pc
+import pyarrow.feather
+import torch
+
+ZENODO_DOI = "10.5281/zenodo.10676866"
+
+CONNECTIONS_URL = (
+    "https://zenodo.org/records/10676866/files/"
+    "proofread_connections_783.feather?download=1"
+)
+CONNECTIONS_MD5 = "f48f972d262323a102aed49af1396b8a"
+CONNECTIONS_SIZE = 852_022_274
+
+ROOT_IDS_URL = (
+    "https://zenodo.org/records/10676866/files/"
+    "proofread_root_ids_783.npy?download=1"
+)
+ROOT_IDS_MD5 = "e0e6c19732fd8c7a4e39a2d170105421"
+ROOT_IDS_SIZE = 1_114_168
+
+LABELS_URL = (
+    "https://media.springernature.com/original/springer-static/esm/"
+    "art%3A10.1038%2Fs41586-024-07686-5/MediaObjects/"
+    "41586_2024_7686_MOESM5_ESM.tsv"
+)
+LABELS_SIZE = 27_015_208
+
+LABEL_COLUMNS = ("super_class", "cell_class", "cell_sub_class", "cell_type")
+
+NT_AVG_COLUMNS = ("gaba_avg", "ach_avg", "glut_avg", "oct_avg", "ser_avg", "da_avg")
+NT_NAMES = ("gaba", "ach", "glut", "oct", "ser", "da")
+
+DEFAULT_SEED = 42
+DEFAULT_RATIOS = (0.7, 0.15, 0.15)
+DEFAULT_MIN_SYNAPSES = 5
+
+
+def _default_cache_dir() -> Path:
+    env = os.environ.get("FLYWIRE_GNN_CACHE")
+    if env:
+        return Path(env)
+    return Path.home() / ".cache" / "flywire_gnn"
+
+
+def _md5(path: Path, chunk_size: int = 1 << 22) -> str:
+    h = hashlib.md5()
+    with open(path, "rb") as f:
+        for chunk in iter(lambda: f.read(chunk_size), b""):
+            h.update(chunk)
+    return h.hexdigest()
+
+
+def _download(url: str, dest: Path, expected_size: Optional[int] = None) -> None:
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dest.with_suffix(dest.suffix + ".part")
+    for attempt in range(5):
+        existing = tmp.stat().st_size if tmp.exists() else 0
+        headers = {"Range": f"bytes={existing}-"} if existing else {}
+        req = urllib.request.Request(url, headers=headers)
+        mode = "ab" if existing else "wb"
+        try:
+            with urllib.request.urlopen(req, timeout=180) as r, open(tmp, mode) as f:
+                while chunk := r.read(1 << 22):
+                    f.write(chunk)
+            if expected_size is None or tmp.stat().st_size >= expected_size:
+                tmp.replace(dest)
+                return
+        except Exception as e:
+            if attempt == 4:
+                raise RuntimeError(f"download failed after 5 attempts: {url}") from e
+            time.sleep(2 * (attempt + 1))
+    raise RuntimeError(f"download failed: {url}")
+
+
+def aggregate(connections_path: Path) -> Dict[str, object]:
+    """Memory-safe aggregation of the pair-by-neuropil connection table.
+
+    Returns numpy/pandas frames:
+      pairs:      dict of columns for unique (pre, post) pairs with
+                  syn_count and NT weighted sums (all at >=1 synapse)
+      neuro_pre / neuro_post: (root_id, neuropil, syn) frames
+      node_out / node_in:     per-root syn totals and NT weighted sums
+      neuropils:  sorted list of all neuropil names
+    """
+    tbl = pa.feather.read_table(connections_path, memory_map=True)
+    tbl = tbl.filter(
+        pc.not_equal(pc.field("pre_pt_root_id"), pc.field("post_pt_root_id"))
+    )
+    syn_f = tbl.column("syn_count").cast(pa.float64())
+    w_cols = []
+    for c in NT_AVG_COLUMNS:
+        w = c.replace("_avg", "_w")
+        tbl = tbl.append_column(w, pc.multiply(tbl.column(c), syn_f))
+        w_cols.append(w)
+
+    pairs_t = tbl.group_by(["pre_pt_root_id", "post_pt_root_id"]).aggregate(
+        [("syn_count", "sum")] + [(w, "sum") for w in w_cols]
+    )
+    pairs = {
+        c: pairs_t.column(c).combine_chunks().to_numpy(zero_copy_only=False)
+        for c in pairs_t.column_names
+    }
+
+    def as_frame(t):
+        return t.to_pandas()
+
+    neuro_pre = as_frame(
+        tbl.group_by(["pre_pt_root_id", "neuropil"]).aggregate([("syn_count", "sum")])
+    ).rename(columns={"pre_pt_root_id": "root_id", "syn_count_sum": "syn"})
+
+    neuro_post = as_frame(
+        tbl.group_by(["post_pt_root_id", "neuropil"]).aggregate([("syn_count", "sum")])
+    ).rename(columns={"post_pt_root_id": "root_id", "syn_count_sum": "syn"})
+
+    node_out = as_frame(
+        tbl.group_by("pre_pt_root_id").aggregate(
+            [("syn_count", "sum")] + [(w, "sum") for w in w_cols]
+        )
+    ).rename(columns={"pre_pt_root_id": "root_id", "syn_count_sum": "syn_total"})
+    node_in = as_frame(
+        tbl.group_by("post_pt_root_id").aggregate(
+            [("syn_count", "sum")] + [(w, "sum") for w in w_cols]
+        )
+    ).rename(columns={"post_pt_root_id": "root_id", "syn_count_sum": "syn_total"})
+
+    neuropils = sorted(
+        set(neuro_pre["neuropil"].unique()) | set(neuro_post["neuropil"].unique())
+    )
+
+    return {
+        "pairs": pairs,
+        "neuro_pre": neuro_pre,
+        "neuro_post": neuro_post,
+        "node_out": node_out,
+        "node_in": node_in,
+        "neuropils": neuropils,
+    }
+
+
+def assemble(
+    agg: Dict[str, object],
+    node_ids: np.ndarray,
+    labels_df: pd.DataFrame,
+    label_col: str,
+    min_synapses: int,
+) -> Dict[str, object]:
+    """Build graph tensors from aggregated frames. Pure numpy/torch, testable."""
+    node_ids = np.sort(np.asarray(node_ids, dtype=np.int64))
+    n = len(node_ids)
+
+    def map_ids(arr: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        arr = np.asarray(arr, dtype=np.int64)
+        pos = np.searchsorted(node_ids, arr)
+        pos_c = np.clip(pos, 0, n - 1)
+        ok = node_ids[pos_c] == arr
+        return pos_c, ok
+
+    pairs = agg["pairs"]
+    pre = pairs["pre_pt_root_id"]
+    post = pairs["post_pt_root_id"]
+    syn = pairs["syn_count_sum"].astype(np.int64)
+    wsum_names = [f"{s}_w_sum" for s in NT_NAMES]
+
+    upos, uok = map_ids(pre)
+    vpos, vok = map_ids(post)
+    valid = uok & vok
+    dropped_edges = int((~valid).sum())
+
+    keep = (syn >= min_synapses) & valid
+    u = upos[keep]
+    v = vpos[keep]
+    s = syn[keep]
+    wsum = np.stack([pairs[w][keep] for w in wsum_names], axis=1).astype(np.float64)
+    nt_edge = (wsum / np.maximum(s, 1)[:, None]).astype(np.float32)
+
+    edge_index = torch.from_numpy(np.stack([u, v])).long()
+    edge_attr = torch.from_numpy(
+        np.column_stack([np.log1p(s.astype(np.float64)), nt_edge])
+    ).float()
+
+    upos_all, _ = map_ids(pre[valid])
+    vpos_all, _ = map_ids(post[valid])
+    s_all = syn[valid]
+
+    x = np.zeros((n, 4 + 2 * len(agg["neuropils"]) + 12), dtype=np.float32)
+    x[:, 0] = np.log1p(np.bincount(upos_all, minlength=n).astype(np.float32))
+    x[:, 1] = np.log1p(np.bincount(vpos_all, minlength=n).astype(np.float32))
+    x[:, 2] = np.log1p(
+        np.bincount(upos_all, weights=s_all, minlength=n).astype(np.float32)
+    )
+    x[:, 3] = np.log1p(
+        np.bincount(vpos_all, weights=s_all, minlength=n).astype(np.float32)
+    )
+
+    neuropils = agg["neuropils"]
+    n_n = len(neuropils)
+    np_name2idx = {name: i for i, name in enumerate(neuropils)}
+
+    def neuro_matrix(df: pd.DataFrame, col0: int) -> None:
+        rpos, ok = map_ids(df["root_id"].to_numpy())
+        codes = df["neuropil"].map(np_name2idx).to_numpy(dtype=np.int64)
+        vals = df["syn"].to_numpy(dtype=np.float64)
+        np.add.at(x, (rpos[ok], col0 + codes[ok]), np.log1p(vals[ok]))
+
+    neuro_matrix(agg["neuro_pre"], 4)
+    neuro_matrix(agg["neuro_post"], 4 + n_n)
+
+    col_nt = 4 + 2 * n_n
+
+    def nt_profiles(df: pd.DataFrame, col0: int) -> None:
+        rpos, ok = map_ids(df["root_id"].to_numpy())
+        tot = df["syn_total"].to_numpy(dtype=np.float64)
+        for j, name in enumerate(NT_NAMES):
+            w = df[f"{name}_w_sum"].to_numpy(dtype=np.float64)
+            profile = np.where(tot > 0, w / np.maximum(tot, 1), 0.0)
+            x[rpos[ok], col0 + j] = profile[ok].astype(np.float32)
+
+    nt_profiles(agg["node_out"], col_nt)
+    nt_profiles(agg["node_in"], col_nt + 6)
+
+    y = np.full(n, -1, dtype=np.int64)
+    label_names: list = []
+    if labels_df is not None and label_col in labels_df.columns:
+        sub = labels_df[["root_id", label_col]].dropna().copy()
+        labels = sub[label_col].astype(str)
+        label_names = sorted(labels.unique())
+        enc = {name: i for i, name in enumerate(label_names)}
+        rpos, ok = map_ids(sub["root_id"].to_numpy())
+        y[rpos[ok]] = labels.map(enc).to_numpy()[ok]
+
+    return {
+        "x": torch.from_numpy(x),
+        "edge_index": edge_index,
+        "edge_attr": edge_attr,
+        "y": torch.from_numpy(y),
+        "node_ids": torch.from_numpy(node_ids),
+        "neuropils": neuropils,
+        "label_col": label_col,
+        "label_names": label_names,
+        "meta": {
+            "num_nodes": n,
+            "num_edges": int(edge_index.shape[1]),
+            "dropped_edges_unmapped_root": dropped_edges,
+            "min_synapses": min_synapses,
+            "feature_dim": int(x.shape[1]),
+            "labeled_nodes": int((y >= 0).sum()),
+        },
+    }
+
+
+def make_splits(
+    y: torch.Tensor,
+    seed: int = DEFAULT_SEED,
+    ratios: Tuple[float, float, float] = DEFAULT_RATIOS,
+    stratify: bool = True,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    y_np = y.numpy()
+    n = len(y_np)
+    labeled = np.where(y_np >= 0)[0]
+    labeled.sort()
+    if stratify and len(labeled) > 0:
+        from sklearn.model_selection import train_test_split
+
+        try:
+            train_idx, rest = train_test_split(
+                labeled,
+                train_size=ratios[0],
+                random_state=seed,
+                stratify=y_np[labeled],
+            )
+            rel = ratios[2] / (ratios[1] + ratios[2])
+            val_idx, test_idx = train_test_split(
+                np.sort(rest), test_size=rel, random_state=seed, stratify=y_np[np.sort(rest)]
+            )
+        except ValueError:
+            rng = np.random.RandomState(seed)
+            idx = labeled.copy()
+            rng.shuffle(idx)
+            k1 = int(len(idx) * ratios[0])
+            k2 = int(len(idx) * (ratios[0] + ratios[1]))
+            train_idx, val_idx, test_idx = idx[:k1], idx[k1:k2], idx[k2:]
+    else:
+        rng = np.random.RandomState(seed)
+        idx = rng.permutation(n if len(labeled) == 0 else labeled.copy())
+        k1 = int(len(idx) * ratios[0])
+        k2 = int(len(idx) * (ratios[0] + ratios[1]))
+        train_idx, val_idx, test_idx = idx[:k1], idx[k1:k2], idx[k2:]
+
+    def mask(idx) -> torch.Tensor:
+        m = np.zeros(n, dtype=bool)
+        m[np.asarray(idx, dtype=np.int64)] = True
+        return torch.from_numpy(m)
+
+    return mask(train_idx), mask(val_idx), mask(test_idx)
+
+
+class FlyWireFAFB:
+    """FlyWire FAFB v783 connectome as a ready-to-train graph dataset.
+
+    Parameters
+    ----------
+    root:
+        Cache directory (default ~/.cache/flywire_gnn or $FLYWIRE_GNN_CACHE).
+    labels:
+        Annotation column to use for node classification:
+        'super_class' (default), 'cell_class', 'cell_sub_class', 'cell_type'.
+    min_synapses:
+        Minimum synapse count for a directed edge between two neurons.
+        Default 5 (the published FAFB 'connection' threshold). Set 1 for the
+        densest graph.
+    seed:
+        Seed for the train/val/test split (default 42).
+    download:
+        Download the source files if they are not already in `root`.
+    """
+
+    def __init__(
+        self,
+        root: Optional[str] = None,
+        labels: str = "super_class",
+        min_synapses: int = DEFAULT_MIN_SYNAPSES,
+        seed: int = DEFAULT_SEED,
+        download: bool = True,
+    ) -> None:
+        if labels not in LABEL_COLUMNS:
+            raise ValueError(f"labels must be one of {LABEL_COLUMNS}, got {labels!r}")
+        self.root = Path(root) if root else _default_cache_dir()
+        self.root.mkdir(parents=True, exist_ok=True)
+        self.labels_col = labels
+        self.min_synapses = int(min_synapses)
+        self.seed = seed
+
+        self._conn = self.root / "proofread_connections_783.feather"
+        self._ids = self.root / "proofread_root_ids_783.npy"
+        self._labels = self.root / "annotations_v783.tsv"
+        self._cache_path = (
+            self.root / f"processed_v783_{labels}_syn{self.min_synapses}.pt"
+        )
+
+        if self._cache_path.exists():
+            self._cache = torch.load(self._cache_path, weights_only=True)
+        else:
+            if not download:
+                raise FileNotFoundError(
+                    f"no processed cache at {self._cache_path} and download=False"
+                )
+            self._ensure_raw()
+            self._cache = self._build()
+            torch.save(self._cache, self._cache_path)
+
+    def _ensure_raw(self) -> None:
+        if not self._conn.exists():
+            _download(CONNECTIONS_URL, self._conn, CONNECTIONS_SIZE)
+        if os.path.getsize(self._conn) != CONNECTIONS_SIZE or (
+            _md5(self._conn) != CONNECTIONS_MD5
+        ):
+            raise RuntimeError("connection file failed integrity check")
+        if not self._ids.exists():
+            _download(ROOT_IDS_URL, self._ids, ROOT_IDS_SIZE)
+        if _md5(self._ids) != ROOT_IDS_MD5:
+            raise RuntimeError("root-id file failed integrity check")
+        if not self._labels.exists():
+            _download(LABELS_URL, self._labels, LABELS_SIZE)
+
+    def _build(self) -> Dict[str, object]:
+        agg = aggregate(self._conn)
+        node_ids = np.load(self._ids)
+        labels_df = pd.read_csv(self._labels, sep="\t", low_memory=False)
+        out = assemble(
+            agg,
+            node_ids=node_ids,
+            labels_df=labels_df,
+            label_col=self.labels_col,
+            min_synapses=self.min_synapses,
+        )
+        return out
+
+    @property
+    def data(self):
+        from torch_geometric.data import Data
+
+        c = self._cache
+        d = Data(x=c["x"], edge_index=c["edge_index"], edge_attr=c["edge_attr"], y=c["y"])
+        d.node_ids = c["node_ids"]
+        d.num_classes = len(c["label_names"])
+        return d
+
+    def splits(
+        self, seed: Optional[int] = None, ratios: Tuple[float, float, float] = DEFAULT_RATIOS
+    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        return make_splits(self._cache["y"], seed=seed or self.seed, ratios=ratios)
+
+    def summary(self) -> str:
+        m = self._cache["meta"]
+        return (
+            "FlyWire FAFB v783\n"
+            f"  nodes:            {m['num_nodes']:,}\n"
+            f"  edges:            {m['num_edges']:,} (min_synapses={m['min_synapses']})\n"
+            f"  dropped edges:    {m['dropped_edges_unmapped_root']:,}\n"
+            f"  feature dim:      {m['feature_dim']}\n"
+            f"  labeled nodes:    {m['labeled_nodes']:,} ({m['labeled_nodes']/m['num_nodes']*100:.1f}%)\n"
+            f"  task labels:      {self._cache['label_col']} "
+            f"({len(self._cache['label_names'])} classes)\n"
+            f"  source:           {ZENODO_DOI}"
+        )
+
+    def __repr__(self) -> str:
+        return (
+            f"FlyWireFAFB(labels={self.labels_col!r}, min_synapses={self.min_synapses}, "
+            f"nodes={self._cache['meta']['num_nodes']}, "
+            f"edges={self._cache['meta']['num_edges']})"
+        )

flywire_gnn-0.1.0/flywire_gnn/models.py

@@ -0,0 +1,57 @@
+"""Baseline models for FlyWire FAFB node classification."""
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+class MLP(nn.Module):
+    def __init__(self, in_dim, hidden=128, out_dim=None, num_layers=3, dropout=0.5):
+        super().__init__()
+        out_dim = out_dim or hidden
+        dims = [in_dim] + [hidden] * (num_layers - 1) + [out_dim]
+        self.layers = nn.ModuleList(
+            nn.Linear(dims[i], dims[i + 1]) for i in range(len(dims) - 1)
+        )
+        self.dropout = dropout
+
+    def forward(self, x, edge_index=None):
+        for layer in self.layers[:-1]:
+            x = F.dropout(F.relu(layer(x)), self.dropout, training=self.training)
+        return self.layers[-1](x)
+
+
+class GCN(nn.Module):
+    def __init__(self, in_dim, hidden=128, out_dim=None, num_layers=3, dropout=0.5):
+        super().__init__()
+        from torch_geometric.nn import GCNConv
+
+        out_dim = out_dim or hidden
+        dims = [in_dim] + [hidden] * (num_layers - 1) + [out_dim]
+        self.convs = nn.ModuleList(
+            GCNConv(dims[i], dims[i + 1]) for i in range(len(dims) - 1)
+        )
+        self.dropout = dropout
+
+    def forward(self, x, edge_index):
+        for conv in self.convs[:-1]:
+            x = F.dropout(F.relu(conv(x, edge_index)), self.dropout, training=self.training)
+        return self.convs[-1](x, edge_index)
+
+
+class GraphSAGE(nn.Module):
+    def __init__(self, in_dim, hidden=128, out_dim=None, num_layers=3, dropout=0.5):
+        super().__init__()
+        from torch_geometric.nn import SAGEConv
+
+        out_dim = out_dim or hidden
+        dims = [in_dim] + [hidden] * (num_layers - 1) + [out_dim]
+        self.convs = nn.ModuleList(
+            SAGEConv(dims[i], dims[i + 1]) for i in range(len(dims) - 1)
+        )
+        self.dropout = dropout
+
+    def forward(self, x, edge_index):
+        for conv in self.convs[:-1]:
+            x = F.dropout(F.relu(conv(x, edge_index)), self.dropout, training=self.training)
+        return self.convs[-1](x, edge_index)

flywire_gnn-0.1.0/flywire_gnn/train.py

@@ -0,0 +1,149 @@
+"""Train a baseline model on FlyWire FAFB node classification."""
+
+from __future__ import annotations
+
+import argparse
+import json
+import time
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from .dataset import FlyWireFAFB
+from .models import GCN, GraphSAGE, MLP
+
+MODELS = {"mlp": MLP, "gcn": GCN, "sage": GraphSAGE}
+
+
+def evaluate(model, data, mask, model_name):
+    model.eval()
+    with torch.no_grad():
+        out = model(data.x, data.edge_index)
+        pred = out[mask].argmax(dim=1)
+        y = data.y[mask]
+    acc = (pred == y).float().mean().item()
+    correct = pred == y
+    f1_sum, f1_n = 0.0, 0
+    for c in range(out.shape[1]):
+        c_pred = pred == c
+        c_true = y == c
+        tp = (c_pred & c_true & (y >= 0)).sum().item()
+        fp = (c_pred & ~c_true & (y >= 0)).sum().item()
+        fn = (~c_pred & c_true).sum().item()
+        if c_true.sum().item() == 0:
+            continue
+        f1_sum += 2 * tp / max(2 * tp + fp + fn, 1)
+        f1_n += 1
+    return acc, f1_sum / max(f1_n, 1)
+
+
+def train_one(
+    model_name,
+    data,
+    masks,
+    hidden=128,
+    epochs=200,
+    lr=0.01,
+    weight_decay=5e-4,
+    seed=42,
+    device="cpu",
+    log_every=25,
+):
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    device = torch.device(device)
+    data = data.to(device)
+    train_mask, val_mask, test_mask = [m.to(device) for m in masks]
+
+    model = MODELS[model_name](
+        data.x.shape[1], hidden=hidden, out_dim=int(data.y.max().item()) + 1
+    ).to(device)
+    opt = torch.optim.Adam(model.parameters(), lr=lr, weight_decay=weight_decay)
+
+    best_val, best_state, best_epoch = -1.0, None, -1
+    t0 = time.time()
+    for epoch in range(1, epochs + 1):
+        model.train()
+        opt.zero_grad()
+        out = model(data.x, data.edge_index)
+        loss = F.cross_entropy(out[train_mask], data.y[train_mask])
+        loss.backward()
+        opt.step()
+        if epoch % 10 == 0 or epoch == epochs:
+            acc_v, _ = evaluate(model, data, val_mask, model_name)
+            improved = acc_v > best_val
+            if improved:
+                best_val = acc_v
+                best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()}
+                best_epoch = epoch
+            if log_every and epoch % log_every == 0:
+                print(
+                    f"  epoch {epoch:4d}/{epochs}  loss={loss.item():.4f}  "
+                    f"val_acc={acc_v:.4f}  best_val={best_val:.4f}",
+                    flush=True,
+                )
+    model.load_state_dict(best_state)
+    acc_t, f1_t = evaluate(model, data, test_mask, model_name)
+    return {
+        "model": model_name,
+        "test_acc": round(acc_t, 4),
+        "test_macro_f1": round(f1_t, 4),
+        "best_val_acc": round(best_val, 4),
+        "best_epoch": best_epoch,
+        "epochs": epochs,
+        "seed": seed,
+        "min_synapses": None,
+        "wall_time_s": round(time.time() - t0, 1),
+    }
+
+
+def main():
+    ap = argparse.ArgumentParser(description="FlyWire FAFB baseline training")
+    ap.add_argument("--model", choices=list(MODELS), default="gcn")
+    ap.add_argument("--models", default=None,
+                    help="comma-separated list to run several, e.g. mlp,gcn,sage")
+    ap.add_argument("--labels", default="super_class",
+                    choices=["super_class", "cell_class", "cell_sub_class", "cell_type"])
+    ap.add_argument("--min-synapses", type=int, default=5)
+    ap.add_argument("--epochs", type=int, default=200)
+    ap.add_argument("--hidden", type=int, default=128)
+    ap.add_argument("--lr", type=float, default=0.01)
+    ap.add_argument("--weight-decay", type=float, default=5e-4)
+    ap.add_argument("--seed", type=int, default=42)
+    ap.add_argument("--device", default="cpu")
+    ap.add_argument("--root", default=None, help="dataset cache dir")
+    args = ap.parse_args()
+
+    names = args.models.split(",") if args.models else [args.model]
+
+    print(f"Loading FlyWire FAFB v783 (labels={args.labels}, min_synapses={args.min_synapses})...")
+    ds = FlyWireFAFB(root=args.root, labels=args.labels, min_synapses=args.min_synapses)
+    data = ds.data
+    masks = ds.splits(seed=args.seed)
+    print(ds.summary())
+    n_train = int(masks[0].sum())
+    n_val = int(masks[1].sum())
+    n_test = int(masks[2].sum())
+    print(f"  splits: train={n_train:,} val={n_val:,} test={n_test:,}")
+
+    results = []
+    for name in names:
+        print(f"\n=== {name.upper()} (hidden={args.hidden}, epochs={args.epochs}, seed={args.seed}) ===")
+        r = train_one(
+            name, data, masks, hidden=args.hidden, epochs=args.epochs,
+            lr=args.lr, weight_decay=args.weight_decay, seed=args.seed,
+            device=args.device, log_every=25,
+        )
+        r["min_synapses"] = args.min_synapses
+        results.append(r)
+        print(f"RESULT {json.dumps(r)}")
+
+    print("\n=== SUMMARY ===")
+    for r in results:
+        print(f"{r['model']:>5}  acc={r['test_acc']:.4f}  macro_f1={r['test_macro_f1']:.4f}  "
+              f"val={r['best_val_acc']:.4f}  epochs_used={r['best_epoch']}")
+
+
+if __name__ == "__main__":
+    main()

flywire_gnn-0.1.0/flywire_gnn.egg-info/PKG-INFO

@@ -0,0 +1,114 @@
+Metadata-Version: 2.4
+Name: flywire-gnn
+Version: 0.1.0
+Summary: The FlyWire FAFB Drosophila connectome as a ready-to-train PyTorch Geometric dataset, with baseline models and reproducible benchmarks.
+Author: Omkar
+License: MIT
+Project-URL: FlyWire, https://flywire.ai
+Project-URL: Zenodo, https://doi.org/10.5281/zenodo.10676866
+Keywords: connectome,flywire,drosophila,gnn,graph-neural-network,neuroscience,pytorch,benchmark
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+Requires-Dist: numpy>=1.24
+Requires-Dist: pandas>=2.0
+Requires-Dist: pyarrow>=14.0
+Requires-Dist: torch>=2.0
+Requires-Dist: torch_geometric>=2.4
+Requires-Dist: scikit-learn>=1.3
+Provides-Extra: hub
+Requires-Dist: huggingface_hub>=0.20; extra == "hub"
+Provides-Extra: dev
+Requires-Dist: pytest>=7.0; extra == "dev"
+
+# flywire-gnn
+
+**The complete fruit-fly brain connectome (FlyWire FAFB v783) as a ready-to-train PyTorch Geometric dataset** — 139,255 proofread neurons, 2.7M directed synaptic connections, cell-type labels on every neuron, deterministic stratified splits, and three reproducible baselines. One `pip install`, one class, no authentication, no left-over data-wrangling.
+
+## Install
+
+```bash
+pip install torch --index-url https://download.pytorch.org/whl/cpu
+pip install torch_geometric
+pip install git+https://github.com/omkar-dhakane/flywire-gnn.git   # or: pip install -e .
+```
+
+Tested with Python 3.13, torch 2.14.0+cpu, torch_geometric 2.8.0.post1, pyarrow 22.0, numpy 2.3.1. Older `torch>=2.0` / `torch_geometric>=2.4` should work; the benchmark numbers above are from the tested versions.
+
+## Quickstart (the whole thing)
+
+```python
+from flywire_gnn import FlyWireFAFB
+
+ds = FlyWireFAFB()                    # downloads 852MB + 1.1MB + 27MB once, caches locally
+data = ds.data                        # torch_geometric.data.Data
+train_mask, val_mask, test_mask = ds.splits()   # deterministic, stratified, seed 42
+print(data)      # Data(x=[139255, 174], edge_index=[2, 2700513], edge_attr=[2700513, 7], y=[139255])
+```
+
+Everything is cached in `~/.cache/flywire_gnn` — subsequent loads are instant.
+
+## The dataset
+
+| | |
+|---|---|
+| Nodes | 139,255 proofread neurons (whole adult female fly brain) |
+| Edges | 2,700,513 unique directed pairs at ≥5 synapses (15,091,983 at ≥1) |
+| Node features (`x`, 174-dim) | wiring structure: degrees/synapse totals, per-neuropil in/out profiles across 79 brain regions, in/out neurotransmitter profiles |
+| Edge attributes (`edge_attr`, 7-dim) | `log1p(syn_count)` + 6 synapse-weighted neurotransmitter probabilities |
+| Labels (`y`) | `super_class`: 9 classes covering **100% of nodes** (also available: `cell_class`, `cell_sub_class`, `cell_type`) |
+| Task | node classification from wiring structure |
+| Splits | 70/15/15 train/val/test, stratified, seed 42 |
+
+Data sources (all public, **no auth needed**):
+- Connectivity: FlyWire Whole-brain Connectome Connectivity Data v783 — [Zenodo, CC-BY-4.0](https://doi.org/10.5281/zenodo.10676866)
+- Annotation: [Schlegel et al. 2024](https://doi.org/10.1038/s41586-024-07686-5), Supplementary Data 5
+- Packaged copies (parquet, pair-level): [huggingface.co/datasets/SLOP011/flywire-fafb-connectome](https://huggingface.co/datasets/SLOP011/flywire-fafb-connectome)
+
+## Leaderboard — node classification on `super_class`
+
+Real runs, full-batch, CPU, seed 42, `min_synapses=5`, hidden 128 (see `flywire_gnn/train.py`):
+
+| Model | Test accuracy | Macro-F1 | Best val acc (epoch) | Epochs | Wall time |
+|---|---|---|---|---|---|
+| MLP (features only) | **0.9851** | 0.7115 | 0.9860 (200) | 200 | 3.5 min |
+| GraphSAGE | 0.9812 | **0.7563** | 0.9820 (150) | 150 | 13 min |
+| GCN | 0.9166 | 0.4702 | 0.9202 (180) | 200 | 23.5 min |
+
+**What the benchmark shows**: wiring-profile features alone nearly saturate accuracy (98.5%) — most of a cell's coarse class is readable directly from its projection pattern. Mean-aggregated message passing (GraphSAGE) roughly matches features and wins on the rare classes (best macro-F1); GCN's symmetric normalization over-smooths and trails. If your architecture can't beat 0.9851 accuracy *and* 0.7563 macro-F1 on these exact splits, it isn't adding anything over a feature baseline.
+
+### Reproduce
+
+```bash
+python -m flywire_gnn.train --models mlp,gcn,sage          # all three
+python -m flywire_gnn.train --model sage --epochs 150 --seed 42
+python -m flywire_gnn.train --model sage --labels cell_class   # harder multi-class task
+python -m flywire_gnn.train --model sage --min-synapses 1       # dense 15M-edge graph
+```
+
+## Design choices (and why)
+
+- **`min_synapses=5`** is the published FAFB "connection" threshold (same default as Codex). Pass `min_synapses=1` for the full-density graph.
+- **Node features come from the ≥1-synapse wiring** (a neuron's total projection profile); edges are filtered by `min_synapses`. So features are the same regardless of the edge threshold you benchmark.
+- **Stratified splits on labeled nodes only**: with `super_class`, all 139,255 nodes are labeled. Deterministic via `seed=42` (NumPy `RandomState` + sklearn stratified split).
+- The raw 9.5 GB per-synapse file (`flywire_synapses_783.feather`) is **not** needed: the 852 MB pair×neuropil table is sufficient.
+
+## Not in v1 (deliberately)
+
+Mesh/skeleton loading (use `fafbseg` + meshparty), other connectomes (MANC/MAOL/MCNS/BANC), hosted leaderboard server, spiking neural simulation, per-synapse link prediction. The package is intentionally complete at this scope.
+
+## Tests
+
+```bash
+pytest tests/ -v
+```
+
+Covers: graph assembly counts on a synthetic graph, feature finiteness, split disjointness/determinism, model forward + backward smoke tests, and (if the real cache is present) the real 139,255-node / 2,700,513-edge invariants.
+
+## License & citation
+
+Code: MIT. Data: **CC-BY-4.0**. Using the dataset means citing:
+
+1. **Dorkenwald et al.** 2024. *Neuronal wiring diagram of an adult brain.* Nature. [doi:10.1038/s41586-024-07558-y](https://doi.org/10.1038/s41586-024-07558-y)
+2. **Schlegel et al.** 2024. *Whole-brain annotation and multi-connectome cell typing of Drosophila.* Nature. [doi:10.1038/s41586-024-07686-5](https://doi.org/10.1038/s41586-024-07686-5)
+
+Not affiliated with the FlyWire Consortium. Interactive exploration: [codex.flywire.ai](https://codex.flywire.ai) — analysis in Python: [navis](https://github.com/navis-org/navis) / [fafbseg](https://github.com/navis-org/fafbseg-py).

flywire_gnn-0.1.0/flywire_gnn.egg-info/SOURCES.txt

@@ -0,0 +1,12 @@
+README.md
+pyproject.toml
+flywire_gnn/__init__.py
+flywire_gnn/dataset.py
+flywire_gnn/models.py
+flywire_gnn/train.py
+flywire_gnn.egg-info/PKG-INFO
+flywire_gnn.egg-info/SOURCES.txt
+flywire_gnn.egg-info/dependency_links.txt
+flywire_gnn.egg-info/requires.txt
+flywire_gnn.egg-info/top_level.txt
+tests/test_pipeline.py

flywire_gnn-0.1.0/flywire_gnn.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

flywire_gnn-0.1.0/flywire_gnn.egg-info/requires.txt

@@ -0,0 +1,12 @@
+numpy>=1.24
+pandas>=2.0
+pyarrow>=14.0
+torch>=2.0
+torch_geometric>=2.4
+scikit-learn>=1.3
+
+[dev]
+pytest>=7.0
+
+[hub]
+huggingface_hub>=0.20

flywire_gnn-0.1.0/flywire_gnn.egg-info/top_level.txt

@@ -0,0 +1 @@
+flywire_gnn

flywire_gnn-0.1.0/pyproject.toml

@@ -0,0 +1,35 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "flywire-gnn"
+version = "0.1.0"
+description = "The FlyWire FAFB Drosophila connectome as a ready-to-train PyTorch Geometric dataset, with baseline models and reproducible benchmarks."
+readme = "README.md"
+requires-python = ">=3.10"
+license = { text = "MIT" }
+authors = [{ name = "Omkar" }]
+keywords = ["connectome", "flywire", "drosophila", "gnn", "graph-neural-network", "neuroscience", "pytorch", "benchmark"]
+dependencies = [
+    "numpy>=1.24",
+    "pandas>=2.0",
+    "pyarrow>=14.0",
+    "torch>=2.0",
+    "torch_geometric>=2.4",
+    "scikit-learn>=1.3",
+]
+
+[project.optional-dependencies]
+hub = ["huggingface_hub>=0.20"]
+dev = ["pytest>=7.0"]
+
+[project.urls]
+FlyWire = "https://flywire.ai"
+Zenodo = "https://doi.org/10.5281/zenodo.10676866"
+
+[tool.setuptools.packages.find]
+include = ["flywire_gnn*"]
+
+[tool.pytest.ini_options]
+testpaths = ["tests"]

flywire_gnn-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

flywire_gnn-0.1.0/tests/test_pipeline.py

@@ -0,0 +1,140 @@
+"""Tests for the flywire_gnn pipeline.
+
+Synthetic-data tests run anywhere. The real-data smoke test is skipped unless
+a processed cache is present (i.e. after a build on the developer machine).
+"""
+
+from pathlib import Path
+
+import numpy as np
+import pandas as pd
+import pytest
+import torch
+
+from flywire_gnn import FlyWireFAFB, GraphSAGE, MLP
+from flywire_gnn.dataset import assemble, make_splits
+
+
+def tiny_agg():
+    pairs = {
+        "pre_pt_root_id": np.array([1, 1, 2, 3, 4, 5], dtype=np.int64),
+        "post_pt_root_id": np.array([2, 3, 3, 1, 5, 4], dtype=np.int64),
+        "syn_count_sum": np.array([7, 3, 9, 2, 6, 11], dtype=np.int64),
+        "gaba_w_sum": np.array([0.5, 0.1, 8.0, 0.4, 5.0, 10.0]),
+        "ach_w_sum": np.array([6.0, 2.5, 0.5, 1.5, 0.5, 0.5]),
+        "glut_w_sum": np.array([0.3, 0.2, 0.2, 0.0, 0.4, 0.3]),
+        "oct_w_sum": np.zeros(6),
+        "ser_w_sum": np.zeros(6),
+        "da_w_sum": np.zeros(6),
+    }
+    neuro_pre = pd.DataFrame(
+        {
+            "root_id": np.array([1, 1, 2, 3, 4, 5], dtype=np.int64),
+            "neuropil": ["AL_L", "ME_R", "AL_L", "ME_R", "AL_L", "ME_R"],
+            "syn": np.array([7, 3, 9, 2, 6, 11], dtype=np.int64),
+        }
+    )
+    neuro_post = pd.DataFrame(
+        {
+            "root_id": np.array([2, 3, 3, 1, 5, 4], dtype=np.int64),
+            "neuropil": ["AL_L", "AL_L", "ME_R", "AL_L", "ME_R", "AL_L"],
+            "syn": np.array([7, 3, 9, 2, 6, 11], dtype=np.int64),
+        }
+    )
+    node_out = pd.DataFrame(
+        {
+            "root_id": np.array([1, 2, 3, 4, 5], dtype=np.int64),
+            "syn_total": np.array([10.0, 9.0, 2.0, 6.0, 11.0]),
+            "gaba_w_sum": np.array([0.6, 8.0, 0.4, 5.0, 10.0]),
+            "ach_w_sum": np.array([8.5, 0.5, 1.5, 0.5, 0.5]),
+            "glut_w_sum": np.array([0.5, 0.2, 0.0, 0.4, 0.3]),
+            "oct_w_sum": np.zeros(5),
+            "ser_w_sum": np.zeros(5),
+            "da_w_sum": np.zeros(5),
+        }
+    )
+    node_in = node_out.copy()
+    labels_df = pd.DataFrame(
+        {
+            "root_id": np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10], dtype=np.int64),
+            "super_class": ["a", "b", "a", "b", "a", "b", "a", "b", "a", "b"],
+        }
+    )
+    return {
+        "pairs": pairs,
+        "neuro_pre": neuro_pre,
+        "neuro_post": neuro_post,
+        "node_out": node_out,
+        "node_in": node_in,
+        "neuropils": ["AL_L", "ME_R"],
+    }, labels_df
+
+
+def test_assemble_counts():
+    agg, labels_df = tiny_agg()
+    node_ids = np.arange(1, 11, dtype=np.int64)
+    out = assemble(agg, node_ids, labels_df, "super_class", min_synapses=5)
+    assert out["meta"]["num_nodes"] == 10
+    # pairs with syn >= 5: (1,2,7), (2,3,9), (4,5,6), (5,4,11)
+    assert out["meta"]["num_edges"] == 4
+    assert out["meta"]["feature_dim"] == 4 + 2 * 2 + 12
+    assert out["meta"]["labeled_nodes"] == 10
+    assert out["meta"]["dropped_edges_unmapped_root"] == 0
+    assert out["x"].shape == (10, 20)
+    assert out["edge_index"].dtype == torch.int64
+
+
+def test_features_finite():
+    agg, labels_df = tiny_agg()
+    node_ids = np.arange(1, 11, dtype=np.int64)
+    out = assemble(agg, node_ids, labels_df, "super_class", min_synapses=5)
+    assert torch.isfinite(out["x"]).all()
+    assert torch.isfinite(out["edge_attr"]).all()
+
+
+def test_splits_disjoint_and_deterministic():
+    y = torch.tensor([0, 1] * 50, dtype=torch.int64)
+    tm, vm, te = make_splits(y, seed=42)
+    assert not (tm & vm).any().item()
+    assert not (tm & te).any().item()
+    assert not (vm & te).any().item()
+    assert (tm | vm | te).all().item()
+    tm2, vm2, te2 = make_splits(y, seed=42)
+    assert (tm == tm2).all().item()
+    assert (vm == vm2).all().item()
+    assert (te == te2).all().item()
+
+
+def test_models_forward_small_graph():
+    torch.manual_seed(0)
+    n, e, f, c = 10, 40, 20, 3
+    x = torch.randn(n, f)
+    y = torch.randint(0, c, (n,))
+    edge_index = torch.randint(0, n, (2, e))
+    for model in [MLP(f, hidden=16, out_dim=c), GraphSAGE(f, hidden=16, out_dim=c)]:
+        out = model(x, edge_index)
+        assert out.shape == (n, c)
+        assert torch.isfinite(out).all()
+        loss = torch.nn.functional.cross_entropy(out, y)
+        loss.backward()
+        assert torch.isfinite(loss)
+
+
+@pytest.mark.skipif(
+    not (Path("data") / "processed_v783_super_class_syn5.pt").exists(),
+    reason="real FAFB cache not built on this machine",
+)
+def test_real_cache_counts():
+    ds = FlyWireFAFB(root="data", labels="super_class", min_synapses=5, download=False)
+    m = ds._cache["meta"]
+    assert m["num_nodes"] == 139255
+    assert m["num_edges"] == 2700513
+    assert m["labeled_nodes"] == 139255
+    assert m["feature_dim"] == 174
+    assert len(ds._cache["label_names"]) == 9
+    tm, vm, te = ds.splits()
+    n = m["num_nodes"]
+    assert int(tm.sum()) + int(vm.sum()) + int(te.sum()) == n
+    assert not (tm & vm).any().item()
+    assert not (tm & te).any().item()
+    assert not (vm & te).any().item()