Digital Drosophila Archive
137 lines
9.6 KiB
Python
137 lines
9.6 KiB
Python
"""Preregister all AK trials before credit/MBON outcomes; never train."""
|
|
from datetime import datetime, timezone
|
|
import platform
|
|
import shutil
|
|
import subprocess
|
|
import sys
|
|
import numpy as np
|
|
import pandas as pd
|
|
import yaml
|
|
from context import (ROOT, MODEL, EXP, AK, H, AG, META, CORE, COEFF, CFG, PINS,
|
|
read, save, digest, sources, integrity, write_npz)
|
|
|
|
|
|
def main():
|
|
assert not EXP.exists() or (EXP.is_dir() and not any(EXP.iterdir())), 'Populated experiment exists: preserve it, do not regenerate'
|
|
assert shutil.disk_usage(EXP.parent).free > CFG['output_budget_bytes'] + 10 * 1024**3
|
|
for relative, expected in PINS.items():
|
|
assert digest(ROOT / relative) == expected, relative
|
|
assert read(AK / 'audit/final-status.json')['status'] == 'PASS-FROZEN-D18-INDEPENDENT-VALIDATION'
|
|
assert not CFG['training'] and not CFG['candidate_generation']
|
|
jc = yaml.safe_load((ROOT / 'models/male-cns-lif/mb-physiology-v1/phase4j.yaml').read_text())
|
|
j = jc['eligibility_v2']
|
|
assert j['kc_trace_tau_ms'] == CFG['kc_trace_tau_ms'] and j['trace_cap'] == CFG['trace_cap']
|
|
assert j['max_fractional_depression'] == CFG['conditional_depression_fraction'] == .25
|
|
assert jc['direct_kc']['dan_time_ms'] / CFG['dt_ms'] == CFG['virtual_gate_tick']
|
|
inventory = read(AK / 'protocol/validation-cases.json')
|
|
assert len(inventory) == CFG['primary_trials']
|
|
seeds = sorted({r['seed'] for r in inventory})
|
|
assert len(seeds) == 5
|
|
planned = []
|
|
for row in inventory:
|
|
for repeat in (0, 1) if row['seed'] == seeds[0] else (0,):
|
|
planned.append(dict(row, repeat=repeat))
|
|
assert len(planned) == CFG['primary_trials'] + CFG['technical_repeats']
|
|
inputs = dict(PINS)
|
|
h_manifest = {line.split(' ', 1)[1]: line.split(' ', 1)[0]
|
|
for line in (H / 'manifests/checksums.sha256').read_text().splitlines()}
|
|
static_h = ('plastic-mask/plastic-edge-candidates.parquet', 'plastic-mask/plastic-edge-ids.int64.npy',
|
|
'plastic-mask/plastic-edge-mask.bool.npy', 'mappings/dan-to-mbon-confirmed.parquet',
|
|
'mappings/mbon-to-compartment.parquet', 'mappings/kc-to-mbon-actual-edges.parquet')
|
|
for rel in static_h:
|
|
assert digest(H / rel) == h_manifest[rel]
|
|
inputs[(H / rel).relative_to(ROOT).as_posix()] = h_manifest[rel]
|
|
for row in planned:
|
|
case = AK / 'cases' / row['cohort'] / row['name']
|
|
trial = AK / 'gpu' / row['cohort'] / (row['name'] + '-r' + str(row['repeat']))
|
|
for p in (case / 'spec.json', case / 'requests.npz', trial / 'state.npz', trial / 'summary.json'):
|
|
inputs[p.relative_to(ROOT).as_posix()] = digest(p)
|
|
paths = [COEFF] + list((AG / 'graphs/full').glob('*.npy')) + list((AG / 'apl-bundle').glob('*.npy'))
|
|
raw_path = ROOT / 'datasets/male-cns-v1.0/processed/neuron-only/gpu/csr-weights.int64.npy'
|
|
for p in paths + [raw_path]:
|
|
inputs[p.relative_to(ROOT).as_posix()] = digest(p)
|
|
for name in ('protocol', 'trials', 'audit', 'regression', 'logs', 'pairs'):
|
|
(EXP / name).mkdir(parents=True, exist_ok=False)
|
|
test = subprocess.run([sys.executable, '-B', '-m', 'unittest', 'discover', '-v', '-s', str(MODEL)],
|
|
cwd=MODEL, text=True, capture_output=True)
|
|
with (EXP / 'regression/unit-tests.log').open('x', encoding='utf-8') as f:
|
|
f.write(test.stdout + test.stderr)
|
|
assert test.returncode == 0, 'Unit tests failed; preserve incomplete experiment'
|
|
assert 'skipped=' not in test.stderr, 'Server arithmetic comparison must not be skipped'
|
|
save(EXP / 'protocol/preregistered.json', dict(timestamp_utc=datetime.now(timezone.utc).isoformat(),
|
|
sources=sources(), config=CFG, inputs=[dict(path=k, sha256=v) for k, v in sorted(inputs.items())],
|
|
parent_commit=subprocess.check_output(['git', '-C', str(ROOT), 'rev-parse', 'HEAD'], text=True).strip(),
|
|
planned=planned, prior_AK_representation_known=True, before_new_credit_and_MBON_outcomes=True,
|
|
training=False, learning_achieved=False, virtual_gate_is_dopamine_measurement=False,
|
|
MBON_recurrent_and_all_other_presynaptic_events_clamped=True,
|
|
expected_functional_effect_not_a_preregistered_learning_gate=True,
|
|
python=sys.version, platform=platform.platform(), numpy=np.__version__))
|
|
save(EXP / 'regression/old-integrity-before.json', integrity())
|
|
meta = pd.read_parquet(META).sort_values('neuron_idx').reset_index(drop=True)
|
|
assert np.array_equal(meta.neuron_idx, np.arange(CFG['neurons'], dtype=np.int32))
|
|
kc = np.flatnonzero(meta.annotation_class.eq('Kenyon_Cell')).astype(np.int32)
|
|
mbon = np.flatnonzero(meta.annotation_class.eq('MBON')).astype(np.int32)
|
|
apl = np.flatnonzero(meta.annotation_type.eq('APL')).astype(np.int32)
|
|
assert len(kc) == 4064 and len(mbon) == 97
|
|
assert np.array_equal(kc, np.load(AG / 'apl-bundle/kc_globals.npy'))
|
|
assert np.array_equal(apl, np.load(AG / 'apl-bundle/apl_globals.npy'))
|
|
folder = AG / 'graphs/full'
|
|
arrays = {k: np.load(folder / (k + '.npy'), mmap_mode='r') for k in ('nodes', 'src', 'dst', 'edge_id', 'weights')}
|
|
assert np.array_equal(arrays['nodes'], np.arange(CFG['neurons'], dtype=np.int32))
|
|
assert np.array_equal(arrays['edge_id'], np.arange(CFG['edges'], dtype=np.int64))
|
|
raw = np.load(raw_path, mmap_mode='r')
|
|
candidate = pd.read_parquet(H / static_h[0])
|
|
ids = np.load(H / static_h[1]); mask = np.load(H / static_h[2])
|
|
assert ids.dtype == np.int64 and ids.shape == (CFG['candidate_edges'],)
|
|
assert np.array_equal(candidate.edge_id, ids) and np.array_equal(np.flatnonzero(mask), ids)
|
|
for prefix, key in (('source', 'src'), ('target', 'dst')):
|
|
assert np.array_equal(candidate[prefix + '_neuron_idx'], arrays[key][ids])
|
|
assert np.array_equal(candidate[prefix + '_body_id'], meta.iloc[arrays[key][ids]].body_id)
|
|
assert np.array_equal(candidate.raw_synapse_count, raw[ids])
|
|
assert np.isin(candidate.source_neuron_idx, kc).all() and np.isin(candidate.target_neuron_idx, mbon).all()
|
|
assert (candidate.plasticity_polarity == -1).all() and candidate.target_compartment_verbatim.notna().all()
|
|
confirmed = pd.read_parquet(H / static_h[3])
|
|
assert confirmed.exact_compartment_match.all()
|
|
assert np.array_equal(confirmed.source_compartment_verbatim, confirmed.target_compartment_verbatim)
|
|
for row in confirmed.itertuples():
|
|
assert arrays['src'][row.edge_id] == row.source_neuron_idx and arrays['dst'][row.edge_id] == row.target_neuron_idx
|
|
assert raw[row.edge_id] == row.raw_synapse_count and meta.iloc[row.source_neuron_idx].annotation_class == 'DAN'
|
|
assert set(candidate.target_neuron_idx).issubset(set(confirmed.target_neuron_idx))
|
|
compartments = pd.read_parquet(H / static_h[4]).sort_values('neuron_idx').reset_index(drop=True)
|
|
assert np.array_equal(compartments.neuron_idx, mbon)
|
|
compartments.to_parquet(EXP / 'protocol/mbon-metadata.parquet', index=False)
|
|
candidate.to_parquet(EXP / 'protocol/candidate-static.parquet', index=False)
|
|
lookup = np.full(CFG['neurons'], -1, np.int32); lookup[mbon] = np.arange(len(mbon), dtype=np.int32)
|
|
selected = np.flatnonzero(lookup[arrays['dst']] >= 0)
|
|
src, dst = arrays['src'][selected], lookup[arrays['dst'][selected]]
|
|
weights = arrays['weights'][selected].copy()
|
|
assert weights.dtype == np.float64 and raw.dtype == np.int64
|
|
weights[np.isin(src, apl)] = 0.0
|
|
is_candidate = mask[selected]
|
|
assert int(is_candidate.sum()) == CFG['candidate_edges'] and (weights[is_candidate] > 0).all()
|
|
groups = np.full(len(selected), CFG['groups'].index('other'), np.int32)
|
|
for name, chosen in (('nonplastic_KC', np.isin(src, kc)), ('DAN', meta.iloc[src].annotation_class.eq('DAN').to_numpy()),
|
|
('MBON', np.isin(src, mbon)), ('APL', np.isin(src, apl)), ('plastic_KC', is_candidate)):
|
|
groups[chosen] = CFG['groups'].index(name)
|
|
kcall = pd.read_parquet(H / static_h[5])
|
|
assert np.array_equal(selected[np.isin(src, kc)], kcall.edge_id.to_numpy())
|
|
assert np.array_equal(raw[kcall.edge_id.to_numpy()], kcall.raw_synapse_count.to_numpy())
|
|
write_npz(EXP / 'protocol/incoming.npz', src_global=src, dst_local=dst, edge_id=selected.astype(np.int64),
|
|
weights=weights, raw_contacts=raw[selected], groups=groups, is_candidate=is_candidate,
|
|
targets=mbon, kc_globals=kc, edge_source_local=np.searchsorted(kc, candidate.source_neuron_idx).astype(np.int32))
|
|
save(EXP / 'protocol/graph-audit.json', dict(neurons=len(meta), all_edges=CFG['edges'], KC=len(kc), MBON=len(mbon),
|
|
all_MBON_incoming_edges=len(selected), all_KC_MBON_edges=len(kcall), candidate_edges=len(ids),
|
|
candidate_targets=int(candidate.target_neuron_idx.nunique()), confirmed_DAN_MBON_edges=len(confirmed),
|
|
compartments={str(c): dict(edges=len(g), targets=int(g.target_neuron_idx.nunique()), sources=int(g.source_neuron_idx.nunique()))
|
|
for c, g in candidate.groupby('target_compartment_verbatim')},
|
|
fast_APL_incoming_edges_zeroed=int(np.isin(src, apl).sum()), groups={name: int((groups == i).sum()) for i, name in enumerate(CFG['groups'])},
|
|
mask_all_edges_identity_verified=True, mapping_basis='Frozen explicit annotation token + real DAN-to-MBON edge; not synapse ROI',
|
|
feedback_neuron_class_not_inferred=True, weight_units='float64 SI volts; not measured conductance'))
|
|
save(EXP / 'protocol/prepared-hashes.json', [dict(path=p.relative_to(EXP).as_posix(), bytes=p.stat().st_size, sha256=digest(p))
|
|
for p in sorted((EXP / 'protocol').iterdir()) if p.is_file() and p.name != 'prepared-hashes.json'])
|
|
print('PREREGISTERED 120 primary +24 repeats; static edge identities validated; no new trajectory or training', flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|