digital-drosophila/models/male-cns-lif/credit-assignment-v1/phase4n_controls.py

#!/usr/bin/env python3
"""Phase 4N independent control trajectories; no frozen artifact writes."""
from __future__ import annotations

import argparse
import json
import resource
import time
from pathlib import Path

import numpy as np
import pandas as pd

from phase4h_audit import write_json
from phase4i import Validator
from phase4j import apply_v2, effect
from phase4l import OdorSimulator, pn_drive
from phase4m import ReplayEngine, compare_A_B, load_odor_inputs, replay_effect
from phase4n import load_config


def main():
    ap=argparse.ArgumentParser();ap.add_argument("--config",type=Path,required=True);ap.add_argument("--max-pairings",type=int,default=100)
    args=ap.parse_args();ncfg,mcfg,cfg,basepath=load_config(args.config);exp=Path(ncfg["experiment_root"])
    identity=pd.read_parquet(exp/"edge-identity/frozen-edge-identity.parquet")
    v=Validator(cfg,basepath,exp);sim=OdorSimulator(v,cfg,exp)
    panel,vectors,selection,mapping=load_odor_inputs(cfg);gloms=selection["glomeruli_order"];ms=float(np.median(panel.mapped_response_sum))
    pn={}
    for label in ("A","B"):
        oid=selection[f"ODOR_{label}"]["odor_id"]
        pn[label]=pn_drive(oid,vectors[oid],mapping,gloms,str(mcfg["protocol"]["input_mode"]),ms,cfg)[:2]
    matched=v.dan_targets(str(ncfg["matched_dan_type"]),str(ncfg["matched_compartment_verbatim"]))
    unrelated=v.dan_targets(str(ncfg["unrelated_dan_type"]),str(ncfg["unrelated_compartment_verbatim"]))
    src=identity.source_neuron_idx.to_numpy(np.int32);loc=np.full(v.n,-1,np.int32);loc[v.kc]=np.arange(len(v.kc),dtype=np.int32);loc=loc[src]
    k=np.float32(cfg["eligibility_v2"]["eta_effective_mv"]/cfg["model"]["w_syn_mv_per_anatomical_synapse"])
    bound=np.float32(1-cfg["eligibility_v2"]["max_fractional_depression"])
    points=set(map(int,ncfg["checkpoints"]));rows=[];walls=[]
    conditions={"B+matched-US":("B",matched,str(ncfg["matched_compartment_verbatim"])),
                "A+random-US-timing":("A",matched,str(ncfg["matched_compartment_verbatim"])),
                "A+unrelated-compartment-US":("A",unrelated,str(ncfg["unrelated_compartment_verbatim"]))}
    for condition,(odor,dan,compartment) in conditions.items():
        group_mask=identity.compartment.astype(str).to_numpy()==compartment
        for seed in ncfg["train_seeds"]:
            add=np.zeros(len(v.edge_ids),np.float32);E=np.ones(len(v.edge_ids),np.float32)
            rng=np.random.default_rng(seed+int(ncfg["random_us_seed_offset"]))
            np.savez_compressed(exp/"controls"/f"E-{condition}-seed-{seed}-pairings-0.npz",E=E,edge_id=v.edge_ids)
            rows.append({"condition":condition,"seed":seed,"pairings":0,"odor":odor,"DAN_compartment":compartment,
                "US_time_ms":None,"updated_edges":0,"matched_compartment_E_reduction_sum":0.0,
                "all_mask_E_reduction_sum":0.0,"brain_wall_seconds":0.0})
            for pairing in range(1,args.max_pairings+1):
                pulse=float(rng.uniform(*ncfg["random_us_window_ms"])) if condition=="A+random-US-timing" else float(mcfg["protocol"]["dan_time_ms"])
                run=sim.run_rates(*pn[odor],seed+pairing,float(mcfg["protocol"]["duration_ms"]),delta=add,dan_targets=dan,dan_pulses=[pulse])
                walls.append(run["wall_seconds"]);trace=run["kc_trace_at_dan"];elig=trace[loc]
                active=group_mask&(elig>0)
                E[active]=np.maximum(bound,E[active]*(np.float32(1.0)-k*elig[active]))
                add,_=apply_v2(v,trace,compartment,1.0,add,cfg)
                if pairing in points:
                    out=exp/"controls"/f"E-{condition}-seed-{seed}-pairings-{pairing}.npz"
                    np.savez_compressed(out,E=E,edge_id=v.edge_ids)
                    rows.append({"condition":condition,"seed":seed,"pairings":pairing,"odor":odor,"DAN_compartment":compartment,
                        "US_time_ms":pulse,"updated_edges":int(np.count_nonzero(E<1-1e-7)),"matched_compartment_E_reduction_sum":float(np.sum((1-E)[identity.compartment.astype(str).to_numpy()==str(ncfg["matched_compartment_verbatim"])])),
                        "all_mask_E_reduction_sum":float(np.sum(1-E)),"brain_wall_seconds":run["wall_seconds"]})
                if pairing % 20 == 0:
                    print(json.dumps({"condition":condition,"seed":seed,"pairing":pairing,"E_reduction_sum":float(np.sum(1-E))}),flush=True)
            print(json.dumps({"condition":condition,"seed":seed,"pairings":args.max_pairings,"E_reduction_sum":float(np.sum(1-E))}),flush=True)
    frame=pd.DataFrame(rows);frame.to_parquet(exp/"controls/independent-training-trajectories.parquet",index=False)
    # Common frozen held-out A/B KC spike trains; changes below are diagnostic, not memory claims.
    target_comp=identity.compartment.astype(str).to_numpy()==str(ncfg["matched_compartment_verbatim"])
    engine=ReplayEngine(sim,np.unique(identity.loc[target_comp,"target_neuron_idx"].to_numpy(np.int32)))
    w=identity.anatomical_weight_mv.to_numpy(np.float32);replay_rows=[]
    for condition in conditions:
        for seed,es in zip(ncfg["train_seeds"],ncfg["evaluation_seeds"]):
            for pair in [20,50,60,75,100]:
                with np.load(exp/"controls"/f"E-{condition}-seed-{seed}-pairings-{pair}.npz") as f:E=f["E"].copy()
                delta=w*(E-1);effects={}
                for odor in ("A","B"):
                    spikes=pd.read_parquet(Path(mcfg["experiment_root"])/"isolated-replay"/f"KC-spikes-ODOR_{odor}-seed-{es}.parquet")
                    base=engine.run(spikes,np.zeros(len(E),np.float32),float(mcfg["protocol"]["duration_ms"]))
                    post=engine.run(spikes,delta,float(mcfg["protocol"]["duration_ms"]))
                    effects[odor]=replay_effect(base,post,sim.dt)
                replay_rows.append({"condition":condition,"train_seed":seed,"eval_seed":es,"pairings":pair,
                    "A_specific_functional":compare_A_B(effects["A"],effects["B"]),
                    "A_MBON_spike_L1":effects["A"]["MBON_spike_L1"],"B_MBON_spike_L1":effects["B"]["MBON_spike_L1"],
                    "A_transmission_reduction_percent":effects["A"]["transmitted_reduction_percent"],
                    "B_transmission_reduction_percent":effects["B"]["transmitted_reduction_percent"]})
    rframe=pd.DataFrame(replay_rows);rframe.to_parquet(exp/"controls/isolated-heldout-control-replay.parquet",index=False)
    write_json(exp/"controls/independent-summary.json",{"conditions":list(conditions),"seeds":list(ncfg["train_seeds"]),
        "pairings_per_seed":args.max_pairings,"training_calls":len(walls),"training_wall_seconds":float(sum(walls)),
        "mean_call_seconds":float(np.mean(walls)),"peak_RAM_KiB":int(resource.getrusage(resource.RUSAGE_SELF).ru_maxrss),
        "matched_compartment_target_not_modified_by_unrelated_US":bool(frame[frame.condition=="A+unrelated-compartment-US"].matched_compartment_E_reduction_sum.eq(0).all()),
        "isolated_replay":rframe.groupby(["condition","pairings"]).agg(A_specific_seeds=("A_specific_functional","sum"),A_spike_L1=("A_MBON_spike_L1","mean"),B_spike_L1=("B_MBON_spike_L1","mean")).reset_index().to_dict("records")})


if __name__=="__main__":main()