359 lines
16 KiB
Python
Executable file
359 lines
16 KiB
Python
Executable file
#!/usr/bin/env python3
|
|
"""Build every browser artifact from the official FlyWire Codex v783 exports.
|
|
|
|
One run downloads (and checksum-verifies) the Codex CSV exports, then writes the sparse
|
|
connectivity arrays, viewer geometry, optic-lobe column tables, `meta.json` and
|
|
`circuit-roles.json` into the output directory. Rerunning with default options reproduces the
|
|
artifacts committed in `data/fafb-v783` byte for byte.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import gzip
|
|
import hashlib
|
|
import json
|
|
import struct
|
|
import urllib.request
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
|
|
ROOT = Path(__file__).resolve().parents[1]
|
|
SOURCE = ROOT / ".tools" / "flywire-v783"
|
|
OUTPUT = ROOT / "data" / "fafb-v783"
|
|
BASE_URL = "https://storage.googleapis.com/flywire-data/codex/data/fafb/783"
|
|
FILES = {
|
|
"classification.csv.gz": "e946b552f4056dfc977707be0674609832c3f64332a22d69dc0d9615e7aae663",
|
|
"connections.csv.gz": "d49dd692e59e153aa3c83f5257bfc0eff51247b86d7bb183386c6d1622c70fc9",
|
|
"consolidated_cell_types.csv.gz": "8aba246d71dc40361677493629972ce3883048c3d02010adc42bda22962a1a2d",
|
|
"coordinates.csv.gz": "14337121f451f98c2576cee72c24409ada5aaf7948b7c7ca8de9040296840e05",
|
|
"column_assignment.csv.gz": "bdf4ce7f62cc63493d53eefad3816ff2dfd08b190e97b35a492e0e453df2f0f6",
|
|
}
|
|
DATASET = "FlyWire FAFB Codex v783"
|
|
RETRIEVED = "2026-09-13"
|
|
DEFAULT_COMMAND_BUCKETS = 8
|
|
# Emitted into circuit-roles.json in this order, empty lists included.
|
|
CIRCUIT_ROLES = ("sensory", "motor", "descending", "kenyon", "mbon")
|
|
# The macro types of `docs/design/macros.md` section 11, in the order the contract lists them.
|
|
# One population per type, `macro_<type>`, drawn round-robin from the mushroom body output
|
|
# neurons plus the brain motor neurons -- the neurons whose input synapses the reward rule
|
|
# changes, which is what makes a learned preference for a macro the mushroom body doing what it
|
|
# does in the real fly. The order is the contract's and fixes which neurons land in which
|
|
# population, so it is never re-sorted.
|
|
MACRO_TYPES = (
|
|
"go_objective",
|
|
"go_out",
|
|
"go_warp",
|
|
"go_route",
|
|
"go_item",
|
|
"go_npc",
|
|
"go_frontier",
|
|
# `docs/design/macros.md` section 13's two errands, beside the other walks.
|
|
"go_shop",
|
|
"go_heal",
|
|
"talk",
|
|
"menu",
|
|
"next",
|
|
"yes",
|
|
"no",
|
|
"close",
|
|
"confirm",
|
|
"back",
|
|
# Section 14: one population per move slot, where `attack` was. No "best move" knowledge
|
|
# remains anywhere, so which move is used is the fly's choice and these four are what it
|
|
# learns on.
|
|
"move_1",
|
|
"move_2",
|
|
"move_3",
|
|
"move_4",
|
|
"switch",
|
|
"item",
|
|
"throw_ball",
|
|
"run",
|
|
"buy_potion",
|
|
"buy_ball",
|
|
"buy_antidote",
|
|
"buy_repel",
|
|
# The conversation at the Pokemon Center counter, which is section 13's `HEAL`.
|
|
"heal",
|
|
"leave",
|
|
)
|
|
# Neural-model signs for annotated transmitters; anything else is treated as excitatory.
|
|
NT_SIGN = {"ACH": 1, "GABA": -1, "GLUT": -1, "OCT": 1, "SER": 1, "DA": 1}
|
|
|
|
|
|
def macro_roles(mbon: list[int], motor: list[int]) -> dict[str, list[int]]:
|
|
"""`macro_<type>` populations, round-robin over the MBON and brain-motor pool.
|
|
|
|
A pure function of the two anatomical role lists, so it says the same thing whether it is
|
|
called from a full build or from `--macro-roles` over the committed artifact. The pool is the
|
|
union of the two lists sorted ascending -- neuron indices are already the sorted `root_id`
|
|
positions, so this is the dataset's own order and not a second opinion about it -- and
|
|
neuron `pool[i]` joins `MACRO_TYPES[i % len(MACRO_TYPES)]`. With 96 MBONs and 110 brain motor
|
|
neurons that is 206 neurons: six or seven per type at the thirty-one of section 14, seven or
|
|
eight at the twenty-seven of section 13, nine or ten at the twenty-two section 11 shipped.
|
|
|
|
The partition is a pure function of the type list, so **adding a type re-deals every
|
|
population** -- 206 neurons dealt twenty-seven ways is not twenty-two ways plus five. That is
|
|
what the earlier addition did too, and it is why section 11 says what it says about the
|
|
checkpoint: rates are restored by role *name*, new roles start at zero, and the neuron ids,
|
|
the edges and the kernel are untouched, so the compatibility string does not move. What a run
|
|
loses across the change is the mushroom body's learned preference for particular macros, which
|
|
is the cost of the button existing at all.
|
|
"""
|
|
pool = sorted(set(mbon) | set(motor))
|
|
roles: dict[str, list[int]] = {f"macro_{name}": [] for name in MACRO_TYPES}
|
|
for position, neuron in enumerate(pool):
|
|
roles[f"macro_{MACRO_TYPES[position % len(MACRO_TYPES)]}"].append(neuron)
|
|
return roles
|
|
|
|
|
|
def parse_args(argv: list[str] | None = None) -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
parser.add_argument(
|
|
"--command-buckets",
|
|
type=int,
|
|
default=DEFAULT_COMMAND_BUCKETS,
|
|
metavar="N",
|
|
help="split descending neurons into N round-robin `command_<k>` roles (default: %(default)s)",
|
|
)
|
|
parser.add_argument(
|
|
"--output",
|
|
type=Path,
|
|
default=OUTPUT,
|
|
metavar="DIR",
|
|
help="directory to write the artifacts into (default: %(default)s)",
|
|
)
|
|
parser.add_argument(
|
|
"--macro-roles",
|
|
action="store_true",
|
|
help=(
|
|
"rewrite only the `macro_<type>` roles of an existing circuit-roles.json in --output, "
|
|
"without downloading the Codex exports; the macro populations are a pure function of "
|
|
"the mbon and motor lists already in that file"
|
|
),
|
|
)
|
|
args = parser.parse_args(argv)
|
|
if args.command_buckets < 1:
|
|
parser.error("--command-buckets must be at least 1")
|
|
return args
|
|
|
|
|
|
def rewrite_macro_roles(output: Path) -> None:
|
|
"""`--macro-roles`: recompute the macro populations over the committed circuit roles.
|
|
|
|
The full build needs the five Codex exports (several GB) and rewrites every artifact; the
|
|
macro populations need neither, because they are derived from `mbon` and `motor` and nothing
|
|
else. This path rewrites `circuit-roles.json` in place: the five anatomical roles keep their
|
|
values and their positions, every `macro_*` role is dropped and re-derived, and the file is
|
|
written with the same separators, so the diff against the previous copy is exactly the macro
|
|
roles. `tests/test_build_flywire.py` pins that it agrees with a full run.
|
|
"""
|
|
path = output / "circuit-roles.json"
|
|
circuits = json.loads(path.read_text())
|
|
roles = {name: indices for name, indices in circuits["roles"].items() if not name.startswith("macro_")}
|
|
roles.update(macro_roles(roles["mbon"], roles["motor"]))
|
|
circuits["roles"] = roles
|
|
path.write_text(json.dumps(circuits, separators=(",", ":")))
|
|
print({name: len(indices) for name, indices in roles.items() if name.startswith("macro_")}, flush=True)
|
|
|
|
|
|
def fetch_sources() -> None:
|
|
SOURCE.mkdir(parents=True, exist_ok=True)
|
|
for name, expected in FILES.items():
|
|
path = SOURCE / name
|
|
if not path.exists():
|
|
print(f"downloading {name}", flush=True)
|
|
urllib.request.urlretrieve(f"{BASE_URL}/{name}", path)
|
|
actual = hashlib.sha256(path.read_bytes()).hexdigest()
|
|
if actual != expected:
|
|
raise RuntimeError(f"checksum mismatch for {name}: {actual}")
|
|
|
|
|
|
def write_gzip(output: Path, name: str, data: bytes) -> dict[str, int | str]:
|
|
path = output / name
|
|
with gzip.GzipFile(filename="", mode="wb", fileobj=path.open("wb"), mtime=0) as stream:
|
|
stream.write(data)
|
|
return {
|
|
"bytes": len(data),
|
|
"compressedBytes": path.stat().st_size,
|
|
"sha256": hashlib.sha256(path.read_bytes()).hexdigest(),
|
|
}
|
|
|
|
|
|
def pack_array(code: str, values: list[int] | list[float]) -> bytes:
|
|
return struct.pack(f"<{len(values)}{code}", *values)
|
|
|
|
|
|
def main(argv: list[str] | None = None) -> None:
|
|
args = parse_args(argv)
|
|
output = args.output
|
|
command_buckets = args.command_buckets
|
|
|
|
if args.macro_roles:
|
|
rewrite_macro_roles(output)
|
|
return
|
|
|
|
fetch_sources()
|
|
output.mkdir(parents=True, exist_ok=True)
|
|
|
|
classifications: dict[int, dict[str, str]] = {}
|
|
with gzip.open(SOURCE / "classification.csv.gz", "rt", newline="") as stream:
|
|
for row in csv.DictReader(stream):
|
|
classifications[int(row["root_id"])] = row
|
|
|
|
roots = sorted(classifications)
|
|
index = {root: i for i, root in enumerate(roots)}
|
|
n = len(roots)
|
|
print(f"indexed {n:,} neurons", flush=True)
|
|
|
|
primary: dict[int, str] = {}
|
|
with gzip.open(SOURCE / "consolidated_cell_types.csv.gz", "rt", newline="") as stream:
|
|
for row in csv.DictReader(stream):
|
|
primary[int(row["root_id"])] = row["primary_type"]
|
|
|
|
positions_sum = [[0.0, 0.0, 0.0, 0] for _ in range(n)]
|
|
with gzip.open(SOURCE / "coordinates.csv.gz", "rt", newline="") as stream:
|
|
for row in csv.DictReader(stream):
|
|
i = index.get(int(row["root_id"]))
|
|
if i is None:
|
|
continue
|
|
xyz = [float(value) for value in row["position"].strip("()[]").replace(",", " ").split()]
|
|
positions_sum[i][0] += xyz[0]
|
|
positions_sum[i][1] += xyz[1]
|
|
positions_sum[i][2] += xyz[2]
|
|
positions_sum[i][3] += 1
|
|
|
|
positions: list[float] = []
|
|
classes: list[int] = []
|
|
roles: dict[str, list[int]] = defaultdict(list)
|
|
circuit_roles: dict[str, list[int]] = {name: [] for name in CIRCUIT_ROLES}
|
|
for i, root in enumerate(roots):
|
|
sx, sy, sz, count = positions_sum[i]
|
|
positions.extend((sx / count, sy / count, sz / count) if count else (0.0, 0.0, 0.0))
|
|
row = classifications[root]
|
|
flow = row["flow"]
|
|
super_class = row["super_class"]
|
|
cell_class = row["class"]
|
|
subtype = row["sub_class"]
|
|
side = row["side"].lower()
|
|
cell_type = primary.get(root, "")
|
|
category = 0 if flow == "afferent" else 2 if flow == "efferent" else 1
|
|
classes.append(category)
|
|
if cell_type in {"DNa01", "DNa02"}:
|
|
roles[f"steer_{'left' if side.startswith('left') else 'right'}"].append(i)
|
|
if cell_type == "DNp09":
|
|
roles["forward"].append(i)
|
|
if cell_type == "MDN":
|
|
roles["backward"].append(i)
|
|
if subtype == "proboscis_motor_neuron":
|
|
roles["proboscis"].append(i)
|
|
if super_class == "descending":
|
|
roles["descending"].append(i)
|
|
roles[f"command_{i % command_buckets}"].append(i)
|
|
if super_class == "motor" or cell_class == "brain_motor_neuron":
|
|
roles["motor"].append(i)
|
|
if cell_class == "DAN" and cell_type.startswith("PAM"):
|
|
roles["reward_pam"].append(i)
|
|
# Anatomical circuit roles, written separately so the mushroom-body populations can be
|
|
# loaded and merged over meta.roles without reloading the whole metadata blob.
|
|
if super_class in ("sensory", "sensory_ascending"):
|
|
circuit_roles["sensory"].append(i)
|
|
if super_class == "motor" or cell_class == "brain_motor_neuron":
|
|
circuit_roles["motor"].append(i)
|
|
if super_class == "descending":
|
|
circuit_roles["descending"].append(i)
|
|
if cell_class == "Kenyon_Cell":
|
|
circuit_roles["kenyon"].append(i)
|
|
if cell_class == "MBON":
|
|
circuit_roles["mbon"].append(i)
|
|
|
|
visual_rows: list[tuple[int, int, float, float]] = []
|
|
with gzip.open(SOURCE / "column_assignment.csv.gz", "rt", newline="") as stream:
|
|
for row in csv.DictReader(stream):
|
|
if row["type"] != "L1":
|
|
continue
|
|
i = index.get(int(row["root_id"]))
|
|
if i is None:
|
|
continue
|
|
hemisphere = 0 if row["hemisphere"].lower().startswith("left") else 1
|
|
visual_rows.append((i, hemisphere, float(row["x"]), float(row["y"])))
|
|
roles["visual_l1"].append(i)
|
|
|
|
edge_map: dict[int, list[int | str]] = {}
|
|
with gzip.open(SOURCE / "connections.csv.gz", "rt", newline="") as stream:
|
|
for row_number, row in enumerate(csv.DictReader(stream), 1):
|
|
pre = index.get(int(row["pre_root_id"]))
|
|
post = index.get(int(row["post_root_id"]))
|
|
if pre is None or post is None:
|
|
continue
|
|
key = (pre << 32) | post
|
|
count = int(row["syn_count"])
|
|
nt = row["nt_type"].upper()
|
|
if key in edge_map:
|
|
edge_map[key][0] = int(edge_map[key][0]) + count
|
|
if edge_map[key][1] != nt:
|
|
edge_map[key][1] = "MIXED"
|
|
else:
|
|
edge_map[key] = [count, nt]
|
|
if row_number % 500_000 == 0:
|
|
print(f"parsed {row_number:,} connection rows", flush=True)
|
|
|
|
indptr = [0] * (n + 1)
|
|
targets: list[int] = []
|
|
weights: list[int] = []
|
|
viewer_edges: list[int] = []
|
|
for key, (count_value, nt_value) in sorted(edge_map.items()):
|
|
pre = key >> 32
|
|
post = key & 0xFFFFFFFF
|
|
count = int(count_value)
|
|
sign = NT_SIGN.get(str(nt_value), 1)
|
|
indptr[pre + 1] += 1
|
|
targets.append(post)
|
|
weights.append(max(-32767, min(32767, count * sign)))
|
|
if key % 43 == 0:
|
|
viewer_edges.extend((pre, post))
|
|
for i in range(n):
|
|
indptr[i + 1] += indptr[i]
|
|
|
|
visual_indices = [row[0] for row in visual_rows]
|
|
visual_hemisphere = [row[1] for row in visual_rows]
|
|
visual_xy = [value for row in visual_rows for value in row[2:]]
|
|
artifacts = {
|
|
"indptr.binz": write_gzip(output, "indptr.binz", pack_array("I", indptr)),
|
|
"targets.binz": write_gzip(output, "targets.binz", pack_array("I", targets)),
|
|
"weights.binz": write_gzip(output, "weights.binz", pack_array("h", weights)),
|
|
"positions.binz": write_gzip(output, "positions.binz", pack_array("f", positions)),
|
|
"classes.binz": write_gzip(output, "classes.binz", bytes(classes)),
|
|
"viewer-edges.binz": write_gzip(output, "viewer-edges.binz", pack_array("I", viewer_edges)),
|
|
"visual-indices.binz": write_gzip(output, "visual-indices.binz", pack_array("I", visual_indices)),
|
|
"visual-hemisphere.binz": write_gzip(output, "visual-hemisphere.binz", bytes(visual_hemisphere)),
|
|
"visual-xy.binz": write_gzip(output, "visual-xy.binz", pack_array("f", visual_xy)),
|
|
}
|
|
metadata = {
|
|
"schemaVersion": 1,
|
|
"dataset": DATASET,
|
|
"retrieved": RETRIEVED,
|
|
"neurons": n,
|
|
"edges": len(targets),
|
|
"sourceFiles": FILES,
|
|
"roles": dict(sorted(roles.items())),
|
|
"visual": {"population": "L1", "count": len(visual_rows)},
|
|
"weightEncoding": "signed aggregated synapse count; GABA/GLUT negative, other annotated transmitters positive",
|
|
"positionPolicy": "arithmetic mean of representative coordinate rows",
|
|
"artifacts": artifacts,
|
|
}
|
|
(output / "meta.json").write_text(json.dumps(metadata, separators=(",", ":")))
|
|
# Section 11's macro populations, appended after the five anatomical roles so that an
|
|
# existing consumer's key order is untouched and the diff against the previous artifact is
|
|
# additive.
|
|
circuit_roles.update(macro_roles(circuit_roles["mbon"], circuit_roles["motor"]))
|
|
circuits = {"dataset": DATASET, "neurons": n, "roles": circuit_roles}
|
|
(output / "circuit-roles.json").write_text(json.dumps(circuits, separators=(",", ":")))
|
|
print(f"wrote {len(targets):,} edges and {len(visual_rows):,} L1 inputs", flush=True)
|
|
print({name: len(indices) for name, indices in circuit_roles.items()}, flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|