Digital Drosophila Archive
90 lines
3.7 KiB
Python
90 lines
3.7 KiB
Python
#!/usr/bin/env python3
|
|
"""Snapshot the official MaleCNS retinotopy T-bar annotations at body level."""
|
|
from __future__ import annotations
|
|
|
|
import argparse
|
|
import json
|
|
import time
|
|
from collections import defaultdict
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
from neuroglancer.read_precomputed_annotations import AnnotationReader
|
|
|
|
|
|
def main() -> None:
|
|
ap = argparse.ArgumentParser()
|
|
ap.add_argument("--url", required=True)
|
|
ap.add_argument("--output", type=Path, required=True)
|
|
ap.add_argument("--summary", type=Path, required=True)
|
|
args = ap.parse_args()
|
|
|
|
reader = AnnotationReader(args.url)
|
|
prop_ids = [p.id for p in reader.properties]
|
|
expected = ["hex1", "hex2", "area_to_count", "r", "theta", "v", "h", "ht", "vic", "area_fit", "body", "roi"]
|
|
if prop_ids != expected:
|
|
raise RuntimeError(f"official schema changed: {prop_ids}")
|
|
|
|
# Values are accumulated with area_to_count as the official contribution
|
|
# weight. This is a lossless body-level sufficient statistic for the
|
|
# coordinate means used by Phase 4F; no body IDs are inferred.
|
|
names = ["hex1", "hex2", "r", "theta", "v", "h", "ht", "vic", "area_fit"]
|
|
accum: dict[int, dict[str, float]] = defaultdict(lambda: defaultdict(float))
|
|
roi_sets: dict[int, set[int]] = defaultdict(set)
|
|
t0 = time.perf_counter()
|
|
annotations = 0
|
|
for ann in reader.get_within_spatial_bounds(max_parallelism=128):
|
|
p = ann.props
|
|
body = int(p[10])
|
|
weight = float(p[2])
|
|
row = accum[body]
|
|
row["weight_sum"] += weight
|
|
row["annotation_count"] += 1.0
|
|
for i, name in enumerate(names):
|
|
value = float(p[i if i < 2 else i + 1])
|
|
if np.isfinite(value):
|
|
row[f"{name}_weighted_sum"] += weight * value
|
|
row[f"{name}_finite_weight"] += weight
|
|
roi_sets[body].add(int(p[11]))
|
|
annotations += 1
|
|
if annotations % 250_000 == 0:
|
|
print(json.dumps({"annotations": annotations, "bodies": len(accum), "wall_seconds": time.perf_counter() - t0}), flush=True)
|
|
|
|
bodies = np.asarray(sorted(accum), dtype=np.uint64)
|
|
cols: dict[str, pa.Array] = {
|
|
"body_id": pa.array(bodies, type=pa.uint64()),
|
|
"tbar_annotation_count": pa.array([int(accum[int(b)]["annotation_count"]) for b in bodies], type=pa.int32()),
|
|
"visual_input_contribution_sum": pa.array([accum[int(b)]["weight_sum"] for b in bodies], type=pa.float64()),
|
|
}
|
|
for name in names:
|
|
values = []
|
|
for body in bodies:
|
|
row = accum[int(body)]
|
|
den = row[f"{name}_finite_weight"]
|
|
values.append(row[f"{name}_weighted_sum"] / den if den else None)
|
|
cols[f"tbar_{name}_weighted_mean"] = pa.array(values, type=pa.float64())
|
|
cols["tbar_roi_codes"] = pa.array([sorted(roi_sets[int(b)]) for b in bodies], type=pa.list_(pa.uint8()))
|
|
table = pa.table(cols)
|
|
args.output.parent.mkdir(parents=True, exist_ok=True)
|
|
pq.write_table(table, args.output, compression="zstd", compression_level=9)
|
|
summary = {
|
|
"source_url": args.url,
|
|
"official_property_schema": prop_ids,
|
|
"coordinate_space": reader.metadata["dimensions"],
|
|
"lower_bound": reader.lower_bound,
|
|
"upper_bound": reader.upper_bound,
|
|
"annotation_count": annotations,
|
|
"unique_body_count": len(bodies),
|
|
"aggregation": "body-level area_to_count-weighted means; official body property used verbatim",
|
|
"wall_seconds": time.perf_counter() - t0,
|
|
"output": str(args.output),
|
|
}
|
|
args.summary.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
|
|
print(json.dumps(summary), flush=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|