Files

370 lines
20 KiB
Python

#!/usr/bin/env python3
"""Phase 4H stage 1: read-only MaleCNS mushroom-body learning-circuit audit."""
from __future__ import annotations
import argparse
import hashlib
import json
import re
import shutil
import subprocess
import sys
from collections import Counter
from datetime import datetime, timezone
from pathlib import Path
from typing import Any
import numpy as np
import pandas as pd
import pyarrow.parquet as pq
import yaml
def sha256(path: Path, chunk: int = 8 << 20) -> str:
h = hashlib.sha256()
with path.open("rb") as handle:
while data := handle.read(chunk):
h.update(data)
return h.hexdigest()
def write_json(path: Path, value: Any) -> None:
path.parent.mkdir(parents=True, exist_ok=True)
path.write_text(json.dumps(value, indent=2, sort_keys=True, default=json_default) + "\n", encoding="utf-8")
def json_default(value: Any) -> Any:
if isinstance(value, np.generic):
return value.item()
if isinstance(value, Path):
return str(value)
raise TypeError(type(value).__name__)
def parse_compartment(instance: Any) -> tuple[str | None, str]:
"""Accept only one explicit, non-empty parenthesized token; preserve it verbatim."""
if instance is None or pd.isna(instance):
return None, "instance_null"
text = str(instance)
matches = re.findall(r"\(([^()]*)\)", text)
nonempty = [item.strip() for item in matches if item.strip()]
if len(matches) == 1 and len(nonempty) == 1:
return nonempty[0], "explicit_single_parenthesized_token"
if not matches:
return None, "no_parenthesized_token"
return None, "ambiguous_parenthesized_tokens"
def verify_checksum(root: Path, control: Path) -> dict[str, Any]:
checked = 0
failures: list[dict[str, str]] = []
for line in control.read_text(encoding="utf-8").splitlines():
if not line.strip():
continue
expected, rel = line.split(None, 1)
path = root / rel.strip().lstrip("*")
checked += 1
if not path.is_file():
failures.append({"path": str(path), "reason": "missing"})
else:
actual = sha256(path)
if actual != expected:
failures.append({"path": str(path), "reason": "sha256_mismatch", "actual": actual, "expected": expected})
return {"control": str(control), "root": str(root), "checked": checked, "failures": failures, "passed": not failures}
def frozen_controls(cfg: dict[str, Any]) -> list[tuple[Path, Path]]:
p = cfg["paths"]
roots = [
(Path(p["reference_root"]), Path(p["reference_root"]) / "checksums.sha256"),
(Path(p["runtime_root"]), Path(p["runtime_root"]) / "checksums.sha256"),
(Path(p["io_root"]), Path(p["io_root"]) / "checksums.sha256"),
(Path(p["phase4c_root"]), Path(p["phase4c_root"]) / "manifests/checksums.sha256"),
(Path(p["phase4d_root"]), Path(p["phase4d_root"]) / "manifests/checksums.sha256"),
(Path(p["phase4e_root"]), Path(p["phase4e_root"]) / "manifests/checksums.sha256"),
(Path(p["phase4f_root"]), Path(p["phase4f_root"]) / "manifests/checksums.sha256"),
(Path(p["phase4f_model_root"]), Path(p["phase4f_model_root"]) / "checksums.sha256"),
(Path(p["phase4g_root"]), Path(p["phase4g_root"]) / "manifests/checksums.sha256"),
(Path(p["phase4g_model_root"]), Path(p["phase4g_model_root"]) / "checksums.sha256"),
(Path(p["phase4g_learned_root"]), Path(p["phase4g_learned_root"]) / "checksums.sha256"),
]
return roots
def run_frozen_audit(cfg: dict[str, Any]) -> dict[str, Any]:
controls = frozen_controls(cfg)
missing = [str(control) for _, control in controls if not control.is_file()]
if missing:
raise FileNotFoundError(f"frozen checksum controls missing: {missing}")
checks = [verify_checksum(root, control) for root, control in controls]
result = {"all_passed": all(check["passed"] for check in checks), "controls": checks}
if not result["all_passed"]:
raise RuntimeError("frozen artifact checksum verification failed")
return result
def graph_paths(cfg: dict[str, Any]) -> dict[str, Path]:
root = Path(cfg["paths"]["connectome_root"]).resolve()
allowed = Path("/data/digital-drosophila/datasets/male-cns-v1.0/processed/neuron-only").resolve()
if root != allowed:
raise RuntimeError(f"only allowed connectome input is {allowed}")
paths = {
"metadata": root / "neuron-index/unified-metadata.parquet",
"indptr": root / "gpu/csr-indptr.int64.npy",
"indices": root / "gpu/csr-indices.int32.npy",
"weights": root / "gpu/csr-weights.int64.npy",
}
missing = [str(path) for path in paths.values() if not path.is_file()]
if missing:
raise FileNotFoundError(missing)
return paths
def edge_frame(edge_ids: np.ndarray, sources: np.ndarray, indices: np.ndarray, weights: np.ndarray,
metadata: pd.DataFrame, prefix: str = "") -> pd.DataFrame:
src = sources[edge_ids]
dst = indices[edge_ids]
out = pd.DataFrame({
"edge_id": edge_ids.astype(np.int64),
"source_neuron_idx": src.astype(np.int32),
"source_body_id": metadata.loc[src, "body_id"].to_numpy(dtype=np.uint64),
"target_neuron_idx": dst.astype(np.int32),
"target_body_id": metadata.loc[dst, "body_id"].to_numpy(dtype=np.uint64),
"raw_synapse_count": weights[edge_ids].astype(np.int64),
})
if prefix:
out["mapping"] = prefix
return out
def tree_manifest(root: Path, excludes: set[Path]) -> dict[str, Any]:
records = []
for path in sorted(item for item in root.rglob("*") if item.is_file() and item not in excludes):
records.append({"path": str(path.relative_to(root)), "bytes": path.stat().st_size, "sha256": sha256(path)})
return {"root": str(root), "generated_utc": datetime.now(timezone.utc).isoformat(), "files": records,
"file_count": len(records), "total_bytes": sum(item["bytes"] for item in records)}
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--config", type=Path, required=True)
args = parser.parse_args()
cfg = yaml.safe_load(args.config.read_text(encoding="utf-8"))
exp = Path(cfg["paths"]["experiment_root"])
model = Path(cfg["paths"]["model_root"])
for sub in ("audit", "mappings", "plastic-mask", "manifests", "regression"):
(exp / sub).mkdir(parents=True, exist_ok=True)
model.mkdir(parents=True, exist_ok=True)
# This is deliberately completed before reading circuit data or writing audit products.
frozen_before = run_frozen_audit(cfg)
write_json(exp / "regression/frozen-checksums-before-stage1.json", frozen_before)
paths = graph_paths(cfg)
metadata = pq.read_table(paths["metadata"]).to_pandas()
required = {"neuron_idx", "body_id", "annotation_class", "annotation_type", "annotation_instance",
"annotation_superclass", "annotation_subclass", "side", "neurotransmitter", "nt_consensus_nt"}
missing_columns = sorted(required - set(metadata.columns))
if missing_columns:
raise RuntimeError(f"metadata columns missing: {missing_columns}")
metadata = metadata.sort_values("neuron_idx", kind="stable").reset_index(drop=True)
if not np.array_equal(metadata["neuron_idx"].to_numpy(), np.arange(len(metadata), dtype=np.int32)):
raise RuntimeError("metadata neuron_idx is not contiguous and row-aligned")
indptr = np.load(paths["indptr"], mmap_mode="r")
indices = np.load(paths["indices"], mmap_mode="r")
weights = np.load(paths["weights"], mmap_mode="r")
expected = cfg["expected_graph"]
graph_actual = {"neurons": int(len(indptr) - 1), "directed_edges": int(len(indices)),
"synaptic_contact_weight_sum": int(weights.sum(dtype=np.int64))}
if graph_actual != expected:
raise RuntimeError(f"graph invariant mismatch: expected={expected}, actual={graph_actual}")
classes = metadata["annotation_class"].fillna("").astype(str).to_numpy()
types = metadata["annotation_type"].fillna("").astype(str).to_numpy()
is_kc = classes == cfg["audit"]["exact_class_values"]["KC"]
is_mbon = classes == cfg["audit"]["exact_class_values"]["MBON"]
is_dan = classes == cfg["audit"]["exact_class_values"]["DAN"]
is_pam = is_dan & np.char.startswith(types.astype(str), cfg["audit"]["pam_type_prefix"])
is_ppl1 = is_dan & np.char.startswith(types.astype(str), cfg["audit"]["ppl1_type_prefix"])
is_other_dan = is_dan & ~is_pam & ~is_ppl1
populations = np.full(len(metadata), "other", dtype=object)
populations[is_kc] = "KC"
populations[is_mbon] = "MBON"
populations[is_pam] = "PAM"
populations[is_ppl1] = "PPL1"
populations[is_other_dan] = "other_DAN"
compartment_values: list[str | None] = []
compartment_status: list[str] = []
for value in metadata["annotation_instance"]:
compartment, status = parse_compartment(value)
compartment_values.append(compartment)
compartment_status.append(status)
metadata["phase4h_population"] = populations
metadata["mb_compartment_verbatim"] = compartment_values
metadata["mb_compartment_mapping_status"] = compartment_status
selected = is_kc | is_mbon | is_dan
population_columns = [
"neuron_idx", "body_id", "phase4h_population", "annotation_superclass", "annotation_class",
"annotation_subclass", "annotation_type", "annotation_instance", "side", "neurotransmitter",
"nt_consensus_nt", "mb_compartment_verbatim", "mb_compartment_mapping_status",
]
metadata.loc[selected, population_columns].to_parquet(exp / "audit/mb-populations.parquet", index=False)
# Vectorized CSR edge source materialization is audit-only and never replaces the frozen CSR.
sources = np.repeat(np.arange(graph_actual["neurons"], dtype=np.int32), np.diff(indptr).astype(np.int64))
targets = np.asarray(indices)
kc_mbon_ids = np.flatnonzero(is_kc[sources] & is_mbon[targets]).astype(np.int64)
dan_kc_ids = np.flatnonzero(is_dan[sources] & is_kc[targets]).astype(np.int64)
dan_mbon_ids = np.flatnonzero(is_dan[sources] & is_mbon[targets]).astype(np.int64)
mbon_out_ids = np.flatnonzero(is_mbon[sources]).astype(np.int64)
kc_mbon = edge_frame(kc_mbon_ids, sources, targets, weights, metadata)
dan_kc = edge_frame(dan_kc_ids, sources, targets, weights, metadata)
dan_mbon = edge_frame(dan_mbon_ids, sources, targets, weights, metadata)
mbon_downstream = edge_frame(mbon_out_ids, sources, targets, weights, metadata)
for frame in (kc_mbon, dan_kc, dan_mbon, mbon_downstream):
frame["source_annotation_type"] = metadata.loc[frame["source_neuron_idx"], "annotation_type"].to_numpy()
frame["source_annotation_instance"] = metadata.loc[frame["source_neuron_idx"], "annotation_instance"].to_numpy()
frame["source_compartment_verbatim"] = metadata.loc[frame["source_neuron_idx"], "mb_compartment_verbatim"].to_numpy()
frame["target_annotation_class"] = metadata.loc[frame["target_neuron_idx"], "annotation_class"].to_numpy()
frame["target_annotation_type"] = metadata.loc[frame["target_neuron_idx"], "annotation_type"].to_numpy()
frame["target_annotation_instance"] = metadata.loc[frame["target_neuron_idx"], "annotation_instance"].to_numpy()
frame["target_compartment_verbatim"] = metadata.loc[frame["target_neuron_idx"], "mb_compartment_verbatim"].to_numpy()
dan_mbon["exact_compartment_match"] = (
dan_mbon["source_compartment_verbatim"].notna()
& dan_mbon["target_compartment_verbatim"].notna()
& (dan_mbon["source_compartment_verbatim"] == dan_mbon["target_compartment_verbatim"])
)
confirmed_dan_mbon = dan_mbon[dan_mbon["exact_compartment_match"]].copy()
confirmed_dan_mbon.to_parquet(exp / "mappings/dan-to-mbon-confirmed.parquet", index=False)
dan_mbon.to_parquet(exp / "mappings/dan-to-mbon-all-actual-edges.parquet", index=False)
dan_kc.to_parquet(exp / "mappings/dan-to-kc-actual-edges.parquet", index=False)
kc_mbon.to_parquet(exp / "mappings/kc-to-mbon-actual-edges.parquet", index=False)
mbon_downstream.to_parquet(exp / "mappings/mbon-to-downstream-actual-edges.parquet", index=False)
reliable_mbon_nodes = set(confirmed_dan_mbon["target_neuron_idx"].astype(int).tolist())
candidate = kc_mbon[kc_mbon["target_neuron_idx"].isin(reliable_mbon_nodes)].copy()
candidate["candidate_basis"] = "actual_KC_to_MBON_edge; target_MBON_has_actual_DAN_edge_with_exact_verbatim_compartment_match"
candidate["plasticity_polarity"] = int(cfg["plasticity"]["default_confirmed_compartment_polarity"])
candidate.to_parquet(exp / "plastic-mask/plastic-edge-candidates.parquet", index=False)
candidate_ids = candidate["edge_id"].to_numpy(dtype=np.int64)
np.save(exp / "plastic-mask/plastic-edge-ids.int64.npy", candidate_ids, allow_pickle=False)
packed = np.zeros(graph_actual["directed_edges"], dtype=np.bool_)
packed[candidate_ids] = True
np.save(exp / "plastic-mask/plastic-edge-mask.bool.npy", packed, allow_pickle=False)
reliable_compartments = sorted(set(candidate["target_compartment_verbatim"].dropna().astype(str)))
polarity_rows = []
for compartment in reliable_compartments:
polarity_rows.append({
"compartment_verbatim": compartment,
"polarity_compartment": int(cfg["plasticity"]["default_confirmed_compartment_polarity"]),
"rule_family": cfg["plasticity"]["first_version_rule_family"],
"interpretation": "model hypothesis: coincident KC and compartmental DAN activity depresses eligible KC-to-MBON synapses",
"not_interpreted_as": "PAM-positive/PPL1-negative cluster rule",
"provenance": ";".join(cfg["plasticity"]["polarity_provenance"]),
})
polarity = pd.DataFrame(polarity_rows, columns=["compartment_verbatim", "polarity_compartment", "rule_family",
"interpretation", "not_interpreted_as", "provenance"])
polarity.to_parquet(exp / "mappings/compartment-plasticity-polarity.parquet", index=False)
dan_nodes = metadata.loc[is_dan, ["neuron_idx", "body_id", "phase4h_population", "annotation_type",
"annotation_instance", "mb_compartment_verbatim", "mb_compartment_mapping_status"]].copy()
dan_nodes.to_parquet(exp / "mappings/dan-to-compartment.parquet", index=False)
mbon_nodes = metadata.loc[is_mbon, ["neuron_idx", "body_id", "annotation_type", "annotation_instance",
"neurotransmitter", "nt_consensus_nt", "mb_compartment_verbatim",
"mb_compartment_mapping_status"]].copy()
mbon_nodes.to_parquet(exp / "mappings/mbon-to-compartment.parquet", index=False)
unknown = {
"DAN_without_reliable_explicit_compartment": int(dan_nodes["mb_compartment_verbatim"].isna().sum()),
"MBON_without_reliable_explicit_compartment": int(mbon_nodes["mb_compartment_verbatim"].isna().sum()),
"actual_DAN_to_MBON_edges_not_exactly_compartment_matched": int((~dan_mbon["exact_compartment_match"]).sum()),
"KC_to_MBON_edges_excluded_from_candidate_mask": int(len(kc_mbon) - len(candidate)),
"KC_to_MBON_target_MBONs_without_confirmed_DAN_match": int(kc_mbon.loc[~kc_mbon["target_neuron_idx"].isin(reliable_mbon_nodes), "target_neuron_idx"].nunique()),
}
unknown["reported_relationship_unknown_total_non_disjoint"] = int(sum(unknown.values()))
population_counts = {
"KC": int(is_kc.sum()), "MBON": int(is_mbon.sum()), "PAM": int(is_pam.sum()),
"PPL1": int(is_ppl1.sum()), "other_DAN": int(is_other_dan.sum()), "DAN_total": int(is_dan.sum()),
}
audit_pass = bool(
frozen_before["all_passed"]
and len(reliable_compartments) > 0
and len(candidate) > 0
and len(confirmed_dan_mbon) > 0
and np.all(weights[candidate_ids] > 0)
)
report = {
"phase": "4H-stage1",
"status": "PASS-AUDIT" if audit_pass else "FAIL-INSUFFICIENT-BIOLOGICAL-MAPPING",
"stage2_authorized_by_audit": audit_pass,
"no_weight_update_or_conditioning_performed": True,
"graph": graph_actual,
"population_counts": population_counts,
"reliable_mb_compartments": {"count": len(reliable_compartments), "verbatim_values": reliable_compartments},
"KC_to_MBON": {"edges": int(len(kc_mbon)), "synapse_weight": int(kc_mbon["raw_synapse_count"].sum())},
"DAN_to_KC": {"edges": int(len(dan_kc)), "synapse_weight": int(dan_kc["raw_synapse_count"].sum())},
"DAN_to_MBON": {
"actual_edges": int(len(dan_mbon)), "actual_synapse_weight": int(dan_mbon["raw_synapse_count"].sum()),
"confirmed_exact_compartment_edge_mappings": int(len(confirmed_dan_mbon)),
"confirmed_unique_DAN_MBON_pairs": int(confirmed_dan_mbon[["source_neuron_idx", "target_neuron_idx"]].drop_duplicates().shape[0]),
},
"DAN_to_compartment": {
"confirmed_neurons": int(dan_nodes["mb_compartment_verbatim"].notna().sum()),
"unknown_neurons": int(dan_nodes["mb_compartment_verbatim"].isna().sum()),
"unique_explicit_compartments": int(dan_nodes["mb_compartment_verbatim"].dropna().nunique()),
},
"MBON_to_downstream": {"edges": int(len(mbon_downstream)), "synapse_weight": int(mbon_downstream["raw_synapse_count"].sum()),
"unique_targets": int(mbon_downstream["target_neuron_idx"].nunique())},
"unknown": unknown,
"plastic_edge_candidate_mask": {
"edges": int(len(candidate)),
"fraction_of_all_edges": float(len(candidate) / graph_actual["directed_edges"]),
"percent_of_all_edges": float(100.0 * len(candidate) / graph_actual["directed_edges"]),
"synapse_weight": int(candidate["raw_synapse_count"].sum()),
"full_edge_boolean_mask_sha256": sha256(exp / "plastic-mask/plastic-edge-mask.bool.npy"),
"edge_ids_sha256": sha256(exp / "plastic-mask/plastic-edge-ids.int64.npy"),
},
"compartment_polarity": polarity.to_dict(orient="records"),
"mapping_policy": cfg["audit"],
"biological_basis_conclusion": (
"sufficient for a conservative, explicitly compartment-matched conditioning prototype"
if audit_pass else "insufficient; conditioning must not start"
),
"limitations": [
"The neuron-level CSR has no synapse-level ROI/location, so candidate membership is conservative and does not prove that every KC-MBON contact lies inside the parsed compartment.",
"Parenthesized annotation_instance tokens are preserved verbatim; missing or ambiguous tokens are not inferred from names.",
"Direct DAN-to-MBON connectivity plus exact compartment-string equality is required; extrasynaptic dopamine is not inferred.",
"polarity=-1 is an explicit anti-Hebbian/LTD model hypothesis with cited provenance, not a PAM/PPL1 valence rule.",
],
"input_sha256": {name: sha256(path) for name, path in paths.items()},
"frozen_checksum_audit": frozen_before,
}
write_json(exp / "audit/stage1-audit-report.json", report)
shutil.copy2(args.config, exp / "audit/phase4h-config.yaml")
shutil.copy2(Path(__file__).resolve(), exp / "audit/phase4h_audit.py")
write_json(model / "STAGE1_AUDIT_GATE.json", {
"status": report["status"], "stage2_authorized": audit_pass,
"audit_report": str(exp / "audit/stage1-audit-report.json"),
"audit_report_sha256": sha256(exp / "audit/stage1-audit-report.json"),
})
# Seal the stage-1 products without touching any pre-4H artifact.
manifest_path = exp / "manifests/stage1-manifest.json"
checksum_path = exp / "manifests/stage1-checksums.sha256"
print(json.dumps(report, indent=2, sort_keys=True, default=json_default))
sys.stdout.flush()
manifest = tree_manifest(exp, {manifest_path, checksum_path})
write_json(manifest_path, manifest)
targets = sorted(path for path in exp.rglob("*") if path.is_file() and path != checksum_path)
checksum_path.write_text("".join(f"{sha256(path)} {path.relative_to(exp)}\n" for path in targets), encoding="utf-8")
if __name__ == "__main__":
main()