370 lines
20 KiB
Python
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()
|