flywire-coding-cortex 0.1.0__py3-none-any.whl

flywire_coding_cortex/etl/build_circuit.py

@@ -0,0 +1,210 @@
+"""Build curated circuit.json from Codex CSVs or Zenodo Feathers."""
+from __future__ import annotations
+
+import csv
+import gzip
+import json
+from collections import Counter, defaultdict
+from pathlib import Path
+from typing import Any
+
+CORE_TYPES = {
+    "LC4": "lc4",
+    "LPLC2": "lplc2",
+    "DNp01": "gf",
+    "DNa02": "dna02",
+    "DNa01": "dna01",
+    "DNp09": "dnp09",
+    "DNg11": "dng11",
+    "MDN": "mdn",
+    "DNp02": "escw",
+    "DNp04": "escw",
+    "DNp11": "escw",
+}
+NT_SIGN = {"ACH": 1.0, "GABA": -1.0, "GLUT": -1.0, "DA": 0.5, "SER": 0.5, "OCT": 0.5}
+MAX_PARTNERS = 330
+
+
+def build_from_codex_dir(raw: Path, out: Path) -> dict[str, Any]:
+    def rows(name: str):
+        with gzip.open(raw / name, "rt", encoding="utf-8", errors="replace") as f:
+            r = csv.reader(f)
+            next(r)
+            yield from r
+
+    core: dict[str, str] = {}
+    type_of: dict[str, str] = {}
+    counts: dict[str, int] = defaultdict(int)
+    for row in rows("consolidated_cell_types.csv.gz"):
+        rid, ptype = row[0], row[1].strip()
+        role = CORE_TYPES.get(ptype)
+        if role:
+            core[rid] = role
+            type_of[rid] = ptype
+            counts[ptype] += 1
+    print("core populations:", dict(counts))
+    if not counts.get("LC4") or not counts.get("LPLC2") or not counts.get("DNp01"):
+        raise SystemExit("FATAL: missing a core population — check type names")
+
+    klass: dict[str, tuple[str, str]] = {}
+    for row in rows("classification.csv.gz"):
+        klass[row[0]] = (row[2], row[6])
+
+    pos: dict[str, tuple[float, float, float]] = {}
+    for row in rows("coordinates.csv.gz"):
+        rid = row[0]
+        if rid in pos:
+            continue
+        p = row[1].strip("[]").split()
+        if len(p) == 3:
+            pos[rid] = (float(p[0]), float(p[1]), float(p[2]))
+
+    partner_strength: dict[str, int] = defaultdict(int)
+    strength_by_role: dict[str, dict[str, int]] = defaultdict(lambda: defaultdict(int))
+    for row in rows("connections.csv.gz"):
+        pre, post, syn = row[0], row[1], int(row[3])
+        pre_core, post_core = pre in core, post in core
+        if pre_core and not post_core:
+            partner_strength[post] += syn
+            strength_by_role[core[pre]][post] += syn
+        elif post_core and not pre_core:
+            partner_strength[pre] += syn
+            strength_by_role[core[post]][pre] += syn
+
+    usable = lambda r: r in pos and r in klass
+    ranked = [rid for rid, _ in sorted(partner_strength.items(), key=lambda kv: -kv[1]) if usable(rid)]
+    partners: list[str] = []
+    seen: set[str] = set()
+
+    def take(cands: list[str], k: int) -> None:
+        n = 0
+        for r in cands:
+            if r in seen or not usable(r):
+                continue
+            seen.add(r)
+            partners.append(r)
+            n += 1
+            if n == k:
+                break
+
+    for role in ("gf", "dna01", "dna02", "dnp09", "dng11", "mdn", "escw"):
+        take([r for r, _ in sorted(strength_by_role[role].items(), key=lambda kv: -kv[1])], 10)
+    take([r for r in ranked if klass[r][0] == "ascending"], 24)
+    take([r for r in ranked if klass[r][0] == "sensory"], 16)
+    take(ranked, MAX_PARTNERS - len(partners))
+    print("partner super_classes:", dict(Counter(klass[r][0] for r in partners)))
+
+    members = list(core.keys()) + partners
+    member_idx = {rid: i for i, rid in enumerate(members)}
+    edges: list[list[float]] = []
+    nt_missing = 0
+    for row in rows("connections.csv.gz"):
+        pre, post = row[0], row[1]
+        i, j = member_idx.get(pre), member_idx.get(post)
+        if i is None or j is None:
+            continue
+        syn, nt = int(row[3]), row[4].strip().upper()
+        sign = NT_SIGN.get(nt)
+        if sign is None:
+            sign, nt_missing = 1.0, nt_missing + 1
+        edges.append([i, j, round(syn * sign, 1)])
+    print(f"circuit edges: {len(edges)} (unknown nt on {nt_missing})")
+
+    xs = [p[0] for p in pos.values()]
+    ys = [p[1] for p in pos.values()]
+    zs = [p[2] for p in pos.values()]
+    cx, cy, cz = (min(xs) + max(xs)) / 2, (min(ys) + max(ys)) / 2, (min(zs) + max(zs)) / 2
+    scale = 20.0 / max(max(xs) - min(xs), max(ys) - min(ys), max(zs) - min(zs))
+
+    def norm(p: tuple[float, float, float]) -> list[float]:
+        return [
+            round((p[0] - cx) * scale, 3),
+            round(-(p[1] - cy) * scale, 3),
+            round(-(p[2] - cz) * scale, 3),
+        ]
+
+    neurons = []
+    for rid in members:
+        sc, side = klass.get(rid, ("", ""))
+        p = norm(pos[rid]) if rid in pos else [0.0, 0.0, 0.0]
+        neurons.append(
+            {
+                "id": rid,
+                "type": type_of.get(rid, sc or "?"),
+                "role": core.get(rid, "other"),
+                "side": side,
+                "pos": p,
+            }
+        )
+    circuit = {
+        "neurons": neurons,
+        "edges": edges,
+        "source": "FlyWire Codex FAFB v783 connections.csv (syn>=5, signed by nt_type)",
+    }
+    out.parent.mkdir(parents=True, exist_ok=True)
+    out.write_text(json.dumps(circuit, separators=(",", ":")) + "\n", encoding="utf-8")
+    print(f"Wrote {out}: {len(neurons)} neurons, {len(edges)} edges")
+    return circuit
+
+
+def build_from_feather_dir(raw: Path, out: Path) -> dict[str, Any]:
+    """Build a curated circuit using proofread_connections Feather + cell types if present.
+
+    Falls back to requiring Codex cell-type CSVs in the same raw dir (downloaded
+    automatically alongside Feathers when missing).
+    """
+    import pandas as pd
+
+    conn_path = raw / "proofread_connections_783.feather"
+    if not conn_path.exists():
+        raise FileNotFoundError(f"Missing {conn_path}")
+
+    # Need cell types — pull Codex types if absent
+    types_gz = raw / "consolidated_cell_types.csv.gz"
+    class_gz = raw / "classification.csv.gz"
+    coords_gz = raw / "coordinates.csv.gz"
+    if not types_gz.exists():
+        from .download import download_file, load_manifest
+
+        base = load_manifest()["profiles"]["codex"]["base_url"].rstrip("/")
+        for name in (
+            "consolidated_cell_types.csv.gz",
+            "classification.csv.gz",
+            "coordinates.csv.gz",
+            "connections.csv.gz",
+        ):
+            dest = raw / name
+            if not dest.exists():
+                download_file(f"{base}/{name}", dest, label=name)
+
+    # Prefer Codex thresholded connections for role ETL consistency when present;
+    # otherwise synthesize a connections.csv.gz-like stream from Feather.
+    if (raw / "connections.csv.gz").exists():
+        return build_from_codex_dir(raw, out)
+
+    df = pd.read_feather(conn_path)
+    # Normalize column names across releases
+    cols = {c.lower(): c for c in df.columns}
+    pre_c = cols.get("pre_root_id") or cols.get("pre") or cols.get("from") or list(df.columns)[0]
+    post_c = cols.get("post_root_id") or cols.get("post") or cols.get("to") or list(df.columns)[1]
+    syn_c = cols.get("syn_count") or cols.get("synapse_count") or cols.get("count")
+    if syn_c is None:
+        for c in df.columns:
+            if "syn" in c.lower() or "count" in c.lower():
+                syn_c = c
+                break
+    if syn_c is None:
+        raise SystemExit(f"Cannot find synapse count column in {list(df.columns)}")
+
+    tmp_conn = raw / "connections.csv.gz"
+    subset = df[[pre_c, post_c, syn_c]].copy()
+    subset.columns = ["pre", "post", "syn"]
+    with gzip.open(tmp_conn, "wt", encoding="utf-8") as f:
+        w = csv.writer(f)
+        w.writerow(["pre_root_id", "post_root_id", "neuropil", "syn_count", "nt_type"])
+        for pre, post, syn in subset.itertuples(index=False, name=None):
+            syn_i = int(syn)
+            if syn_i < 5:
+                continue
+            w.writerow([str(pre), str(post), "", syn_i, "ACH"])
+    return build_from_codex_dir(raw, out)

flywire_coding_cortex/etl/download.py

@@ -0,0 +1,128 @@
+"""Download profile assets into the user cache."""
+from __future__ import annotations
+
+import json
+import shutil
+import urllib.request
+from pathlib import Path
+from typing import Any
+
+from ..paths import bundled_manifest, bundled_seed_circuit, profile_dir
+
+
+def load_manifest() -> dict[str, Any]:
+    return json.loads(bundled_manifest().read_text(encoding="utf-8"))
+
+
+def list_profiles() -> dict[str, Any]:
+    return load_manifest()["profiles"]
+
+
+def download_file(url: str, dest: Path, *, label: str = "") -> None:
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    tmp = dest.with_suffix(dest.suffix + ".part")
+    print(f"Downloading {label or dest.name} …")
+    urllib.request.urlretrieve(url, tmp)  # noqa: S310 — curated manifest URLs
+    tmp.replace(dest)
+    print(f"  -> {dest} ({dest.stat().st_size:,} bytes)")
+
+
+def fetch_seed(*, force: bool = False) -> Path:
+    dest = profile_dir("seed") / "circuit.json"
+    if dest.exists() and not force:
+        return dest
+    seed = bundled_seed_circuit()
+    if seed is None:
+        raise FileNotFoundError(
+            "circuit.seed.json missing. Build it with: flywire-cortex fetch --profile codex"
+        )
+    shutil.copy2(seed, dest)
+    meta = {
+        "profile": "seed",
+        "source": "bundled circuit.seed.json",
+        "neurons": _count_neurons(dest),
+    }
+    (profile_dir("seed") / "meta.json").write_text(json.dumps(meta, indent=2) + "\n", encoding="utf-8")
+    return dest
+
+
+def fetch_codex(*, force: bool = False) -> Path:
+    from .build_circuit import build_from_codex_dir
+
+    raw = profile_dir("codex") / "raw"
+    circuit = profile_dir("codex") / "circuit.json"
+    if circuit.exists() and not force:
+        return circuit
+    manifest = load_manifest()["profiles"]["codex"]
+    base = manifest["base_url"].rstrip("/")
+    raw.mkdir(parents=True, exist_ok=True)
+    for f in manifest["files"]:
+        name = f["name"]
+        dest = raw / name
+        if dest.exists() and not force:
+            continue
+        download_file(f"{base}/{f['path']}", dest, label=name)
+    build_from_codex_dir(raw, circuit)
+    meta = {
+        "profile": "codex",
+        "source": base,
+        "neurons": _count_neurons(circuit),
+    }
+    (profile_dir("codex") / "meta.json").write_text(json.dumps(meta, indent=2) + "\n", encoding="utf-8")
+    # also refresh bundled seed for packaging workflows
+    seed_out = bundled_seed_circuit()
+    if seed_out is None:
+        seed_out = Path(__file__).resolve().parents[3] / "data" / "circuit.seed.json"
+    seed_out.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(circuit, seed_out)
+    assets = Path(__file__).resolve().parents[1] / "assets" / "circuit.seed.json"
+    assets.parent.mkdir(parents=True, exist_ok=True)
+    shutil.copy2(circuit, assets)
+    return circuit
+
+
+def fetch_full(*, force: bool = False, include_synapses: bool = False, yes: bool = False) -> Path:
+    if not yes:
+        raise SystemExit(
+            "Profile 'full' downloads ~10.6 GB. Re-run with --yes to confirm:\n"
+            "  flywire-cortex fetch --profile full --yes\n"
+            "Optional largest synapse table: add --include-synapses"
+        )
+    try:
+        import pandas as pd  # noqa: F401
+        import pyarrow  # noqa: F401
+    except ImportError as exc:
+        raise SystemExit(
+            "Profile 'full' needs: pip install 'flywire-coding-cortex[full]'"
+        ) from exc
+
+    from .build_circuit import build_from_feather_dir
+
+    raw = profile_dir("full") / "raw"
+    circuit = profile_dir("full") / "circuit.json"
+    if circuit.exists() and not force:
+        return circuit
+    manifest = load_manifest()["profiles"]["full"]
+    raw.mkdir(parents=True, exist_ok=True)
+    for f in manifest["files"]:
+        if f.get("optional") and not include_synapses:
+            print(f"Skipping optional {f['name']} (pass --include-synapses to fetch)")
+            continue
+        dest = raw / f["name"]
+        if dest.exists() and not force:
+            continue
+        download_file(f["url"], dest, label=f"{f['name']} (~{f.get('approx_bytes', 0):,} B)")
+    build_from_feather_dir(raw, circuit)
+    meta = {
+        "profile": "full",
+        "source": manifest.get("record_url"),
+        "neurons": _count_neurons(circuit),
+        "include_synapses": include_synapses,
+    }
+    (profile_dir("full") / "meta.json").write_text(json.dumps(meta, indent=2) + "\n", encoding="utf-8")
+    return circuit
+
+
+def _count_neurons(path: Path) -> int:
+    data = json.loads(path.read_text(encoding="utf-8"))
+    return len(data.get("neurons", []))

flywire_coding_cortex/install_skills.py

@@ -0,0 +1,74 @@
+"""Copy bundled skill into Cursor / OpenClaw / Claude skill directories."""
+from __future__ import annotations
+
+import json
+import shutil
+from pathlib import Path
+from typing import Any
+
+from .paths import bundled_skill_dir, home
+
+
+def _cursor_skills() -> Path:
+    return Path.home() / ".cursor" / "skills" / "flywire-coding-cortex"
+
+
+def _openclaw_skills() -> list[Path]:
+    h = Path.home()
+    return [
+        h / ".openclaw" / "workspace" / "skills" / "flywire-coding-cortex",
+        h / ".openclaw" / "skills" / "flywire-coding-cortex",
+    ]
+
+
+def _claude_skills() -> Path:
+    return Path.home() / ".claude" / "skills" / "flywire-coding-cortex"
+
+
+def _copy_skill(dest: Path) -> str:
+    src = bundled_skill_dir()
+    dest.parent.mkdir(parents=True, exist_ok=True)
+    if dest.exists():
+        shutil.rmtree(dest)
+    shutil.copytree(src, dest)
+    return str(dest)
+
+
+def mcp_snippet() -> dict[str, Any]:
+    return {
+        "mcpServers": {
+            "flywire-coding-cortex": {
+                "command": "flywire-cortex",
+                "args": ["mcp"],
+            }
+        }
+    }
+
+
+def install_all(targets: list[str] | None = None) -> dict[str, Any]:
+    targets = targets or ["cursor", "openclaw", "claude"]
+    result: dict[str, Any] = {"installed": [], "skipped": [], "mcp": mcp_snippet(), "home": str(home())}
+    if "cursor" in targets:
+        result["installed"].append({"target": "cursor", "path": _copy_skill(_cursor_skills())})
+    if "openclaw" in targets:
+        # prefer workspace path; also try top-level skills
+        paths = _openclaw_skills()
+        primary = paths[0]
+        result["installed"].append({"target": "openclaw", "path": _copy_skill(primary)})
+        # second path as optional mirror if parent exists
+        if paths[1].parent.exists() or paths[1].parent.parent.exists():
+            try:
+                paths[1].parent.mkdir(parents=True, exist_ok=True)
+                result["installed"].append({"target": "openclaw-alt", "path": _copy_skill(paths[1])})
+            except OSError as exc:
+                result["skipped"].append({"target": "openclaw-alt", "error": str(exc)})
+    if "claude" in targets:
+        result["installed"].append({"target": "claude", "path": _copy_skill(_claude_skills())})
+    tip = home() / "mcp.snippet.json"
+    tip.write_text(json.dumps(mcp_snippet(), indent=2) + "\n", encoding="utf-8")
+    result["mcp_snippet_file"] = str(tip)
+    result["hint"] = (
+        "Add the mcp snippet to Cursor MCP settings, then restart the agent session. "
+        "Ensure `flywire-cortex` is on PATH (pip/uv install)."
+    )
+    return result