allen_full_matrix.py

import urllib.request, json, urllib.parse, time, csv
from pathlib import Path

MAX_DEPTH = 6
MIN_DEPTH = 3
SLEEP = 0.15
OUT_DIR = Path(".")

def allen_get(model, params, retries=3):
    url = (f"https://api.brain-map.org/api/v2/data/{model}/query.json?"
           + urllib.parse.urlencode(params))
    for attempt in range(retries):
        try:
            req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
            with urllib.request.urlopen(req, timeout=20) as r:
                return json.loads(r.read())
        except Exception as e:
            if attempt < retries - 1:
                time.sleep(2 ** attempt)
            else:
                return {"success": False, "num_rows": 0, "msg": [], "error": str(e)}

def main():
    print("Step 1: Fetching structure tree...")
    resp = allen_get("Structure", {
        "criteria": f"[graph_id$eq1][depth$gte{MIN_DEPTH}][depth$lte{MAX_DEPTH}]",
        "num_rows": 2000,
        "only": "id,acronym,name,depth",
        "order[]": "depth asc",
    })
    structures = resp["msg"]
    struct_map = {s["id"]: s["acronym"] for s in structures}
    struct_ids = [s["id"] for s in structures]
    print(f"  Found {len(structures)} structures")

    # Save structure list
    with open("allen_structures_full.csv", "w", newline="") as f:
        w = csv.DictWriter(f, fieldnames=["id","acronym","name","depth"])
        w.writeheader()
        for s in structures:
            w.writerow(s)
    print(f"  Saved allen_structures_full.csv")

    print("\nStep 2: Finding structures with injection experiments...")
    src_to_exps = {}
    for i, sid in enumerate(struct_ids):
        resp = allen_get("ProjectionStructureUnionize", {
            "criteria": (f"[structure_id$eq{sid}]"
                         f"[is_injection$eqtrue]"
                         f"[normalized_projection_volume$gt0]"),
            "num_rows": 3,
            "only": "section_data_set_id",
        })
        if resp["success"] and resp["num_rows"] > 0:
            exps = list({r["section_data_set_id"] for r in resp["msg"]})
            src_to_exps[sid] = exps[:3]
        if (i + 1) % 50 == 0:
            print(f"  {i+1}/{len(struct_ids)} checked, {len(src_to_exps)} with experiments")
        time.sleep(SLEEP)

    print(f"\n  Sources with experiments: {len(src_to_exps)}")
    print(f"  Estimated pairs: {len(src_to_exps) * len(struct_ids):,}")

    print("\nStep 3: Building connectivity matrix...")
    out_path = OUT_DIR / "allen_connectivity_full.csv"
    fieldnames = ["src_id","src_name","tgt_id","tgt_name","conn_strength","n_exp"]
    total = nonzero = 0

    with open(out_path, "w", newline="") as f:
        writer = csv.DictWriter(f, fieldnames=fieldnames)
        writer.writeheader()

        for src_idx, (src_id, exps) in enumerate(src_to_exps.items()):
            src_name = struct_map.get(src_id, str(src_id))
            batch = []

            for tgt_id in struct_ids:
                if tgt_id == src_id:
                    continue
                tgt_name = struct_map.get(tgt_id, str(tgt_id))
                densities = []

                for exp_id in exps:
                    resp = allen_get("ProjectionStructureUnionize", {
                        "criteria": (f"[section_data_set_id$eq{exp_id}]"
                                     f"[structure_id$eq{tgt_id}]"
                                     f"[is_injection$eqfalse]"),
                        "num_rows": 1,
                        "only": "normalized_projection_volume",
                    })
                    if resp["success"] and resp["num_rows"] > 0:
                        d = resp["msg"][0].get("normalized_projection_volume", 0.0) or 0.0
                        densities.append(d)
                    time.sleep(SLEEP)

                if densities:
                    mean_d = sum(densities) / len(densities)
                    batch.append({
                        "src_id": src_id, "src_name": src_name,
                        "tgt_id": tgt_id, "tgt_name": tgt_name,
                        "conn_strength": round(mean_d, 6),
                        "n_exp": len(densities),
                    })
                    total += 1
                    if mean_d > 0.001:
                        nonzero += 1

            writer.writerows(batch)
            f.flush()

            if (src_idx + 1) % 5 == 0:
                pct = (src_idx + 1) / len(src_to_exps) * 100
                print(f"  [{pct:5.1f}%] {src_idx+1}/{len(src_to_exps)} src  "
                      f"pairs={total:,}  nonzero={nonzero:,}  last={src_name}")

    print(f"\nDone: {total:,} pairs → {out_path}")
    print(f"Non-zero: {nonzero:,} ({nonzero/max(total,1)*100:.1f}%)")

if __name__ == "__main__":
    main()