collect¶
The only notebook that talks to Codex or CAVE.
Everything downstream reads the files it writes
(Cfg.CENSUS, Cfg.CONNECTIONS, Cfg.SKELETONS, Cfg.CODEX).
- cell types and other tables from Codex
- the connections of every proofread neuron, from CAVE
- skeletons for 3D figures
class Sources:
"""Remote services and the literal ids of the products we pull from
them."""
# The only CAVE datastack this account can read;
# flywire_fafb_production returns 403.
# Frozen at materialization 783 (the 2024 public release).
STACK = "flywire_fafb_public"
# One finished SWC per root id, in the public bucket behind Codex.
# lod1 is their healed skeleton at materialization 783, which
# matches STACK; the whole set is the sk_lod1_783_healed.zip on the
# Codex download page.
# Nothing in this path names a materialization, unlike the archive
# under skeletons/fafb/archives/783/,
# so a republish for a new one would silently change the geometry.
# Confirmed reachable on 2026-09-29;
# the Codex FAQ does not document it.
SKELETON_URL = (
"https://storage.googleapis.com/flywire-data/codex/skeletons/fafb/lod1"
)
# Codex account token, a different login from CAVE, from the account
# page on codex.flywire.ai.
# It lives in the gitignored .env at the repo root;
# direnv exports it and marimo edit loads it.
# main() only needs it when Cfg.CENSUS is missing.
# Environment, not a cell, so it never lands in the saved notebook.
CODEX_ENV = "CODEX_API_TOKEN"
CODEX_HOST = "https://codex.flywire.ai/api/download_resource"
# The census; download_codex fetches any of CODEX_PRODUCTS on
# demand.
CODEX_CENSUS = "consolidated_cell_types"
# Every table download_resource serves for fafb, checked by name on
# 2026-09-29 (all gzipped CSV).
# The skeletons are not in here;
# they are one SWC per root id under SKELETON_URL.
CODEX_PRODUCTS = (
"consolidated_cell_types",
"classification",
"names",
"neurons",
"cell_stats",
"coordinates",
"labels",
"processed_labels",
"visual_neuron_types",
"column_assignment",
"connectivity_tags",
"connections_princeton",
"connections_princeton_no_threshold",
"synapse_table",
)
# Codex also serves other datasets (hemibrain, ...); this work is
# FAFB.
CODEX_DATASET = "fafb"
# The synapse set behind the Codex release (one row per synapse,
# with a size and a neuropil).
# Every connection count comes from here.
# Views are not listed by get_tables.
# On the cells compared (six LC10a) it reproduces every Codex
# connection to proofread neurons with identical counts,
# which TRANSMITTER_TABLE does not;
# counting its rows is what the release's thresholds of five and
# ten synapses refer to.
SYNAPSE_VIEW = "synapses_v3_neuropil_v6_merge_view"
# The unfiltered Buhmann predictions, with the per-transmitter
# probabilities that the view lacks.
# Only the tool_* notebooks read it;
# no connection count comes from it.
TRANSMITTER_TABLE = "synapses_nt_v1"
# Root ids per query.
# query_view silently stops at ROW_CAP rows,
# so ids go in batches and a batch that hits the cap is split in
# half and retried.
# Fifty is arbitrary; a few hundred also works for LC-sized cells,
# while one id per call is slow.
BATCH = 50
ROW_CAP = 500_000
# Seconds a query may take before it fails and is retried.
# One batch of fifty cells answers in about ten seconds,
# and the largest queries a few times that;
# five minutes is arbitrary but far above both.
QUERY_TIMEOUT = 300
pl.DataFrame(
{
"name": [_k for _k in vars(Sources) if _k.isupper()],
"value": [
str(_v) for _k, _v in vars(Sources).items() if _k.isupper()
],
}
)
shape: (12, 2)
| name | value |
|---|---|
| str | str |
| "STACK" | "flywire_fafb_public" |
| "SKELETON_URL" | "https://storage.googleapis.com… |
| "CODEX_ENV" | "CODEX_API_TOKEN" |
| "CODEX_HOST" | "https://codex.flywire.ai/api/d… |
| "CODEX_CENSUS" | "consolidated_cell_types" |
| … | … |
| "SYNAPSE_VIEW" | "synapses_v3_neuropil_v6_merge_… |
| "TRANSMITTER_TABLE" | "synapses_nt_v1" |
| "BATCH" | "50" |
| "ROW_CAP" | "500000" |
| "QUERY_TIMEOUT" | "300" |
Codex tables¶
CAVE v783 has no cell-type column. Codex ships type names, other per-cell annotations, and whole-brain connection tables as gzipped CSVs. The census is one row per cell:
\[(\mathrm{root\_id},\ \mathrm{primary\_type})\]
def codex_url(product: str, api_token: str) -> str:
"""Codex download URL for one of Sources.CODEX_PRODUCTS in
Sources.CODEX_DATASET."""
url = Sources.CODEX_HOST
url += f"?data_product={product}"
url += f"&dataset={Sources.CODEX_DATASET}"
url += f"&api_token={api_token}"
return url
'https://codex.flywire.ai/api/download_resource?data_product=consolidated_cell_types&dataset=fafb&api_token=<token>'
def download_codex(product: str, dest=None):
"""Write one Codex product to dest, by default
Cfg.CODEX/<product>.csv.gz.
Needs Sources.CODEX_ENV in the environment.
"""
token = os.environ.get(Sources.CODEX_ENV, "").strip()
if not token:
msg = f"set {Sources.CODEX_ENV} (codex.flywire.ai account token)"
raise SystemExit(msg)
if dest is None:
dest = Cfg.CODEX / f"{product}.csv.gz"
dest.parent.mkdir(parents=True, exist_ok=True)
urllib.request.urlretrieve(codex_url(product, token), dest)
return dest
census_btn = mo.ui.run_button(
label="download_codex(Sources.CODEX_CENSUS, Cfg.CENSUS)"
)
census_btn
Any other table goes to Cfg.CODEX.
connections_princeton is the whole v783 connection list without
CAVE, and synapse_table is 2.7 GB.
product_pick = mo.ui.dropdown(
Sources.CODEX_PRODUCTS, value="column_assignment", label="product"
)
product_btn = mo.ui.run_button(label="download_codex(product)")
mo.hstack([product_pick, product_btn], justify="start")
Synapses (CAVE)¶
One synapse is one contact \(s\) from a presynaptic cell to a
postsynaptic cell, with a unique id.
For a set of root ids \(R\) we want every synapse that touches \(R\):
\[
S(R) = \{\, s : s.\mathrm{pre} \in R \ \lor\
s.\mathrm{post} \in R \,\}
\]
filter_in_dict tests one column,
so this is two queries,
one on post_pt_root_id (inputs) and one on pre_pt_root_id
(outputs).
A synapse between two cells of \(R\) comes back from both,
so the two answers are merged on id before any weight is counted.
class TimeoutAdapter(HTTPAdapter):
"""Gives every request a timeout.
requests has none by default, and a connection that dies mid-query
(a laptop suspended during a long fetch) then waits forever.
"""
def send(self, request, **kwargs):
if kwargs.get("timeout") is None:
kwargs["timeout"] = Sources.QUERY_TIMEOUT
return super().send(request, **kwargs)
def cave_materialize():
"""CAVE materialize client for Sources.STACK, with a timeout on
every request.
Needs a saved CAVE token.
"""
materialize = CAVEclient(Sources.STACK).materialize
materialize.session.mount("https://", TimeoutAdapter())
return materialize
# The network demos in this section fetch one LC16 cell,
# a few thousand synapses.
one_cell = root_ids_of_type(load_census(), "LC16")[:1]
one_cell_btn = mo.ui.run_button(label="fetch one LC16 cell")
one_cell_btn
def query_synapses(m, column: str, ids: list[int]):
"""Synapse rows (id, pre, post) with `column` in ids, split in half
until each query is under Sources.ROW_CAP.
CAVE hands the rows back as a pandas frame.
"""
rows = m.query_view(
Sources.SYNAPSE_VIEW,
filter_in_dict={column: ids},
select_columns=["id", "pre_pt_root_id", "post_pt_root_id"],
)
# A full page means the server truncated the answer; there is no
# error, so the row count is the only signal.
if len(rows) >= Sources.ROW_CAP and len(ids) == 1:
msg = (
f"root id {ids[0]} has at least {Sources.ROW_CAP:,} synapses "
f"in {column}; the answer is truncated"
)
raise RuntimeError(msg)
if len(rows) >= Sources.ROW_CAP:
mid = len(ids) // 2
first = query_synapses(m, column, ids[:mid])
second = query_synapses(m, column, ids[mid:])
rows = pd.concat([first, second], ignore_index=True)
return rows[["id", "pre_pt_root_id", "post_pt_root_id"]]
mo.stop(not one_cell_btn.value)
# The client's session sends every request through TimeoutAdapter.
isinstance(
cave_materialize().session.get_adapter("https://"), TimeoutAdapter
)
mo.stop(not one_cell_btn.value)
# Inputs only: synapses onto the cell.
query_synapses(cave_materialize(), "post_pt_root_id", one_cell)
Self-contacts¶
Rows with \(\mathrm{pre} = \mathrm{post}\) look like autapses. Here they are segmentation artifacts: LC dendrites and axons sit in different neuropils, so a true self-contact is not expected. The other options were keeping them as autapses (wrong for these cells) or flagging them; we drop them.
def drop_self_edges(syn):
"""Remove synapses whose pre and post root ids are the same cell."""
return syn[syn["pre_pt_root_id"] != syn["post_pt_root_id"]]
# The third synapse is a self-contact.
drop_self_edges(
pd.DataFrame(
{
"id": [1, 2, 3],
"pre_pt_root_id": [10, 10, 10],
"post_pt_root_id": [20, 30, 10],
}
)
)
| id | pre_pt_root_id | post_pt_root_id | |
|---|---|---|---|
| 0 | 1 | 10 | 20 |
| 1 | 2 | 10 | 30 |
Connections of every proofread neuron¶
One row per ordered pair of cells where at least one end is a proofread neuron, weighted by synapse count,
\[w(i \to j) = \#\{\, s \in S : s.\mathrm{pre} = i,\ s.\mathrm{post}
= j \,\}.\]
Counting synapses is what Codex does, over the same synapse set
(Sources.SYNAPSE_VIEW),
so a weight here equals the Codex weight for two proofread neurons.
The view has a size per synapse that could weight each one;
we do not use it.
An analysis filters what it needs when it loads
(load.scan_connections),
so it does not depend on which types were fetched.
Pairs between two unproofread fragments are left out.
A batch of proofread cells gives its outputs from a query on
pre_pt_root_id and its inputs from a query on post_pt_root_id.
A pair whose presynaptic end is proofread is counted from that
cell's own batch,
so the input query keeps only the pairs whose presynaptic end is a
fragment, and every pair is counted once:
def batch_connections(
outputs: pl.DataFrame, inputs: pl.DataFrame, proofread: pl.Series
) -> pl.DataFrame:
"""Pair weights of one batch:
DataFrame[pre_pt_root_id, post_pt_root_id, weight].
outputs: synapses whose presynaptic cell is in the batch.
inputs: synapses whose postsynaptic cell is in the batch.
proofread: root ids of every proofread neuron.
Self-contacts are dropped.
"""
inputs_from_fragments = inputs.filter(
~pl.col("pre_pt_root_id").is_in(proofread.implode())
)
return (
pl.concat([outputs, inputs_from_fragments])
.filter(pl.col("pre_pt_root_id") != pl.col("post_pt_root_id"))
.group_by("pre_pt_root_id", "post_pt_root_id")
.agg(weight=pl.len().cast(pl.UInt32))
.sort("pre_pt_root_id", "post_pt_root_id")
)
# Cells 1 and 2 are proofread and the batch is cell 2.
# Cell 9 is a fragment.
# Its two synapses onto 2 are one pair of weight 2;
# the pair 1 to 2 is left to cell 1's batch;
# the self-contact 2 to 2 is dropped.
toy_outputs = pl.DataFrame(
{"pre_pt_root_id": [2, 2, 2], "post_pt_root_id": [1, 9, 2]}
)
toy_inputs = pl.DataFrame(
{"pre_pt_root_id": [1, 9, 9, 2], "post_pt_root_id": [2, 2, 2, 2]}
)
batch_connections(toy_outputs, toy_inputs, pl.Series([1, 2]))
shape: (3, 3)
| pre_pt_root_id | post_pt_root_id | weight |
|---|---|---|
| i64 | i64 | u32 |
| 2 | 1 | 1 |
| 2 | 9 | 1 |
| 9 | 2 | 2 |
About a hundred million synapses are fetched,
so fetch_connection_shard takes one batch of Sources.BATCH
cells and writes one file for it,
retrying a failed query a few times.
The demo fetches the first batch to a scratch file.
def fetch_connection_shard(
ids: list[int], proofread: pl.Series, dest: Path
) -> int:
"""Fetch the outputs and inputs of ids, write their pair weights to
dest, return the number of pairs."""
if not hasattr(CLIENTS, "materialize"):
CLIENTS.materialize = cave_materialize()
# A gateway timeout (502, 503) usually means the query was too
# heavy for the server, so the third attempt asks for ten cells at
# a time and the fourth for one.
for attempt in range(1, 6):
size = len(ids) if attempt <= 2 else 10 if attempt == 3 else 1
try:
frames = {}
for column in ("pre_pt_root_id", "post_pt_root_id"):
answers = [
query_synapses(
CLIENTS.materialize, column, ids[i : i + size]
)
for i in range(0, len(ids), size)
]
rows = pd.concat(answers, ignore_index=True)
# Built from numpy so an empty answer still has integer
# columns.
frames[column] = pl.DataFrame(
{
name: rows[name].to_numpy(dtype="int64")
for name in ("pre_pt_root_id", "post_pt_root_id")
}
)
break
except RuntimeError:
raise
except Exception as e:
# A dropped connection or a server hiccup;
# if every attempt fails,
# a rerun resumes from the batches on disk.
if attempt == 5:
raise
print(
f"retry {attempt} after {type(e).__name__}: {e}",
file=sys.stderr,
)
time.sleep(15 * attempt)
pairs = batch_connections(
frames["pre_pt_root_id"], frames["post_pt_root_id"], proofread
)
# Written under a temporary name so a fetch that dies halfway
# leaves no truncated file for the next run to skip.
part = dest.with_name(dest.name + ".part")
pairs.write_parquet(part)
part.replace(dest)
return pairs.height
mo.stop(not shard_btn.value)
_proofread = load_proofread_ids()["root_id"].sort()
fetch_connection_shard(
_proofread.head(Sources.BATCH).to_list(),
_proofread,
Path(tempfile.gettempdir()) / "connection_shard_demo.parquet",
)
collect_connections runs every batch, a few at a time,
and skips the batches whose file is already in
Cfg.CONNECTIONS_PARTS, so a stopped run resumes.
The batches take hours (about ten seconds each, 2,800 of them),
and the number of parallel workers is how hard the public CAVE
service is asked.
When every batch is on disk the files are merged,
sorted by presynaptic cell and in row groups of a million so a
filter on a cell reads little,
into Cfg.CONNECTIONS,
and the folder of batch files is deleted.
def collect_connections(workers: int = 4) -> None:
"""Fetch every proofread cell's connections in batches and merge
them into Cfg.CONNECTIONS."""
proofread = load_proofread_ids()["root_id"].sort()
ids = proofread.to_list()
batches = [
ids[i : i + Sources.BATCH] for i in range(0, len(ids), Sources.BATCH)
]
Cfg.CONNECTIONS_PARTS.mkdir(parents=True, exist_ok=True)
def shard_path(k: int) -> Path:
return Cfg.CONNECTIONS_PARTS / f"part-{k:05d}.parquet"
todo = [k for k in range(len(batches)) if not shard_path(k).exists()]
print(
f"{len(batches) - len(todo)} of {len(batches)} batches on disk; "
f"fetching {len(todo)} with {workers} workers",
flush=True,
)
failed = []
started = time.time()
with ThreadPoolExecutor(max_workers=workers) as pool:
futures = {
pool.submit(
fetch_connection_shard, batches[k], proofread, shard_path(k)
): k
for k in todo
}
for done, future in enumerate(as_completed(futures), start=1):
k = futures[future]
try:
pairs = future.result()
except Exception as e:
# Keep going; the failed batches are reported and a
# rerun fetches them.
failed.append(k)
print(
f"batch {k + 1} failed: {type(e).__name__}: {e}",
flush=True,
)
continue
minutes = (time.time() - started) / 60
print(
f"batch {k + 1}/{len(batches)}: {pairs:,} pairs; "
f"{done}/{len(todo)} fetched in {minutes:.0f} min",
flush=True,
)
if failed:
msg = (
f"{len(failed)} batches failed (first: {sorted(failed)[:10]}); "
f"rerun collect_connections() to fetch them"
)
raise SystemExit(msg)
merged = Cfg.CONNECTIONS.with_name(Cfg.CONNECTIONS.name + ".part")
(
pl.scan_parquet(Cfg.CONNECTIONS_PARTS / "part-*.parquet")
.sort("pre_pt_root_id", "post_pt_root_id")
.sink_parquet(merged, row_group_size=1_000_000)
)
merged.replace(Cfg.CONNECTIONS)
shutil.rmtree(Cfg.CONNECTIONS_PARTS)
total = pl.scan_parquet(Cfg.CONNECTIONS).select(pl.len()).collect().item()
print(f"wrote {Cfg.CONNECTIONS}: {total:,} pairs")
connections_btn = mo.ui.run_button(
label=(
"collect_connections() (hours; resumes; "
"overwrites Cfg.CONNECTIONS)"
)
)
connections_btn
Skeletons¶
A skeleton is a thin tree tracing a neuron's cable.
Codex publishes one finished SWC per root id at
Sources.SKELETON_URL;
these are their healed lod1 skeletons at materialization 783,
denser than a skeleton rebuilt from a mesh.
They do not come from CAVE materialize,
so they still work when materialize is down,
and one root id per call keeps a bad id from failing the batch.
A cell whose skeleton is missing is skipped and reported.
def skeleton_url(root_id) -> str:
"""Codex SWC URL for one root id, the only skeleton source besides
the bulk zip on their download page."""
return f"{Sources.SKELETON_URL}/{root_id}.swc"
# The URL for one cell; every other skeleton is the same path with
# its own root id.
skeleton_url(root_ids_of_type(load_census(), "LC16")[0])
'https://storage.googleapis.com/flywire-data/codex/skeletons/fafb/lod1/720575940606173630.swc'
def collect_skeletons(types=SKELETON_TYPES):
"""Write the first SKELETONS_PER_TYPE skeletons of each of types to
Cfg.SKELETONS/<root_id>.swc."""
census = load_census()
Cfg.SKELETONS.mkdir(parents=True, exist_ok=True)
written = []
for cell_type in types:
for root_id in root_ids_of_type(census, cell_type)[
:SKELETONS_PER_TYPE
]:
path = Cfg.SKELETONS / f"{root_id}.swc"
# Download under a temporary name so a fetch that dies
# halfway leaves no truncated SWC behind to be read as data.
part = path.with_name(path.name + ".part")
try:
urllib.request.urlretrieve(skeleton_url(root_id), part)
except Exception as e:
# One missing skeleton should not stop the rest.
part.unlink(missing_ok=True)
print(
f"skip {cell_type} {root_id}: {type(e).__name__}: {e}",
file=sys.stderr,
)
continue
part.replace(path)
written.append(path)
print(f"wrote {len(written)} skeletons to {Cfg.SKELETONS}")
return written
On disk¶
_paths = [
Cfg.CENSUS,
Cfg.CONNECTIONS,
Cfg.SKELETONS,
Cfg.CODEX,
]
pl.DataFrame(
{
"path": [str(_p.relative_to(Cfg.ROOT)) for _p in _paths],
"exists": [_p.exists() for _p in _paths],
"MB": [
round(_p.stat().st_size / 1e6, 2) if _p.is_file() else None
for _p in _paths
],
}
)
shape: (4, 3)
| path | exists | MB |
|---|---|---|
| str | bool | f64 |
| "data/cell_types.csv.gz" | true | 0.9 |
| "data/connections.parquet" | true | 309.97 |
| "data/skeletons" | true | null |
| "data/codex" | true | null |