"""Reconstruct existing KC targets; never launch a brain simulation or CUDA.""" from functools import lru_cache import time import traceback import numpy as np import pandas as pd from context import ROOT, EXP, AG, META, COEFF, read, save, digest, check from replay_core import IncomingIndex, replay from accounting import DriveAccounting, bit_equal from populations import GROUPS @lru_cache(maxsize=2) def graph_view(graph_name): folder = AG / 'graphs' / graph_name nodes = np.load(folder / 'nodes.npy', mmap_mode='r') all_kc = np.load(AG / 'apl-bundle/kc_globals.npy', mmap_mode='r') apl = np.load(AG / 'apl-bundle/apl_globals.npy', mmap_mode='r') targets, node_positions, apl_columns = np.intersect1d(nodes, all_kc, assume_unique=True, return_indices=True) lookup = np.full(len(nodes), -1, np.int32) lookup[node_positions] = np.arange(len(targets), dtype=np.int32) destination = np.load(folder / 'dst.npy', mmap_mode='r') selected = np.flatnonzero(lookup[destination] >= 0) src_local = np.load(folder / 'src.npy', mmap_mode='r')[selected] src = nodes[src_local] dst = lookup[destination[selected]] ids = np.load(folder / 'edge_id.npy', mmap_mode='r')[selected] weights = np.load(folder / 'weights.npy', mmap_mode='r')[selected].copy() weights[np.isin(src, apl)] = 0.0 original = np.load(ROOT / 'datasets/male-cns-v1.0/processed/neuron-only/gpu/csr-weights.int64.npy', mmap_mode='r') raw = original[ids] assert raw.dtype == np.int64 and np.all(raw > 0) incoming = IncomingIndex(src, dst, ids, weights, len(targets)) group_frame = pd.read_parquet(EXP / 'protocol/neuron-source-groups.parquet') labels = group_frame.audit_source_group.to_numpy() group_codes = np.array([GROUPS.index(x) for x in labels], np.int32)[src] audit = dict(graph_name=graph_name, graph_nodes=len(nodes), graph_edges=len(destination), KC_targets=len(targets), incoming_edges=len(ids), incoming_contacts=int(raw.sum()), zeroed_APL_incoming_edges=int(np.isin(src, apl).sum()), groups={name: dict(edges=int(np.count_nonzero(group_codes == gi)), contacts=int(raw[group_codes == gi].sum())) for gi, name in enumerate(GROUPS)}, source_graph=str(folder), conditional_target_filter_only=True, source_graph_modified=False) return incoming, raw, group_codes, targets, node_positions, apl_columns, audit def input_audit(case, source, spec, state, population): with np.load(case / 'requests.npz', allow_pickle=False) as data: req = {k: data[k].copy() for k in data.files} table = pd.read_csv(source / 'requests.csv') assert len(table) == len(req['tick']) assert np.array_equal(table.tick.to_numpy(), req['tick']) assert np.array_equal(table.neuron_idx.to_numpy(), req['neuron_idx']) assert (table.attempted == 1).all() and table.accepted.isin([0, 1]).all() assert np.array_equal(table.rejected.to_numpy(), 1 - table.accepted.to_numpy()) request_keys = req['tick'].astype(np.int64) * 166700 + req['neuron_idx'] event_keys = state['step'].astype(np.int64) * 166700 + state['neuron_idx'] accepted = np.isin(request_keys, event_keys) assert np.array_equal(accepted, table.accepted.to_numpy().astype(bool)) labels = population.audit_source_group.to_numpy() all_pn = np.flatnonzero(labels == 'ALPN').astype(np.int32) observed_count = np.bincount(state['neuron_idx'], minlength=166700) requested_count = np.bincount(req['neuron_idx'], minlength=166700) accepted_count = np.bincount(req['neuron_idx'][accepted], minlength=166700) targets = np.empty(0, np.int32) rates = np.empty(0, np.float32) if spec['kind'] == 'calibration': with np.load(case / 'pn-drive.npz', allow_pickle=False) as data: targets, rates = data['targets'].copy(), data['rates_hz'].copy() assert targets.dtype == np.int32 and rates.dtype == np.float32 assert targets.shape == rates.shape == (spec['target_count'],) assert len(np.unique(targets)) == len(targets) and np.isfinite(rates).all() and (rates > 0).all() assert np.isin(targets, all_pn).all() and np.isin(req['neuron_idx'], targets).all() rate_map = dict(zip(targets.tolist(), rates.tolist())) pn_table = population.iloc[all_pn][['neuron_idx', 'body_id', 'annotation_type']].copy() pn_table['design_target'] = np.isin(all_pn, targets) pn_table['design_rate_hz'] = [rate_map.get(int(i), 0.0) for i in all_pn] pn_table['requests'] = requested_count[all_pn] pn_table['accepted_requests'] = accepted_count[all_pn] pn_table['rejected_requests'] = requested_count[all_pn] - accepted_count[all_pn] pn_table['actual_spikes'] = observed_count[all_pn] pn_table['spikes_not_coincident_with_accepted_request'] = observed_count[all_pn] - accepted_count[all_pn] assert (pn_table.spikes_not_coincident_with_accepted_request >= 0).all() info = dict(requests=len(request_keys), accepted=int(accepted.sum()), rejected=int((~accepted).sum()), actual_spikes_by_source_group={name: int(observed_count[labels == name].sum()) for name in GROUPS}, all_ALPN=len(all_pn), odor_design_target_count=len(targets), PN_ever_requested=int(np.count_nonzero(requested_count[all_pn])), PN_ever_accepted_request=int(np.count_nonzero(accepted_count[all_pn])), PN_actual_active=int(np.count_nonzero(observed_count[all_pn])), design_target_PN_spikes=int(observed_count[targets].sum()), non_design_ALPN_spikes=int(observed_count[np.setdiff1d(all_pn, targets)].sum()), accepted_request_is_not_proof_spike_was_caused_exclusively_by_request=True, actual_requested_events_not_forced_to_equal_across_odors=True) return req, info, pn_table def audit_trial(row, population): start = time.perf_counter() name = row['case'] + '-r' + str(row['repeat']) out = EXP / 'trials' / name out.mkdir(exist_ok=False) source = ROOT / row['source'] case = AG / 'cases' / row['case'] spec = read(case / 'spec.json') graph_name = 'fixtures' if row['kind'] == 'fixture' else 'full' incoming, raw, groups, targets, positions, apl_columns, graph_audit = graph_view(graph_name) with np.load(source / 'state.npz', allow_pickle=False) as data: state = {k: data[k].copy() for k in data.files} assert set(state) == {'step', 'neuron_idx', 'counts', 'v', 'g', 'ref'} assert all(np.isfinite(v).all() for v in state.values()) req, input_report, pn_table = input_audit(case, source, spec, state, population) event_mask = np.isin(state['neuron_idx'], targets) expected = {key: state[key][event_mask] for key in ('step', 'neuron_idx')} initial_v = np.full(len(targets), spec['model']['resting_potential_mv'], np.float64) * np.float64(.001) initial_g = np.zeros(len(targets), np.float64) if (case / 'arrays.npz').exists(): with np.load(case / 'arrays.npz', allow_pickle=False) as data: initial_v = data['initial_v_mv'][positions] * np.float64(.001) initial_g = data['initial_g_mv'][positions] * np.float64(.001) histories = {} if row['kind'] == 'fixture': for key, suffix in (('v', 'float64'), ('g', 'float64'), ('ref', 'int32')): histories[key] = np.load(source / f'{key}-history.{suffix}.npy', mmap_mode='r')[:, positions] histories['pre_' + key] = np.load(source / f'preAPL-{key}.{suffix}.npy', mmap_mode='r')[:, positions] h = np.load(source / 'apl-inhibition.float32.npy', mmap_mode='r') assert h.shape == (spec['steps'], 4064) if not np.array_equal(apl_columns, np.arange(4064)): h = h[:, apl_columns] coeff = read(COEFF) accounting = DriveAccounting(targets, incoming, raw, groups, GROUPS, expected, coeff, state['counts'][positions], histories) result = replay(initial_v, initial_g, coeff, spec['steps'], targets, req, dict(step=state['step'], neuron_idx=state['neuron_idx']), incoming, h, accounting) exact = {key: bit_equal(result[key], state[key][positions]) for key in ('v', 'g', 'ref', 'counts')} exact.update({key: bit_equal(result[key], expected[key]) for key in ('step', 'neuron_idx')}) assert all(exact.values()), 'Reconstruction diverged: ' + str(exact) assert accounting.callback_count == spec['steps'] data, report = accounting.finish(result['g']) assert bit_equal(data['first_spike_tick'], result['first_spike']) per_kc = pd.DataFrame(data) per_kc.insert(1, 'body_id', population.iloc[targets].body_id.to_numpy()) per_kc.to_parquet(out / 'per-kc.parquet', index=False) pn_table.to_parquet(out / 'per-pn.parquet', index=False) pd.DataFrame(accounting.tick_rows).to_parquet(out / 'per-tick.parquet', index=False) with (out / 'replay-final.npz').open('xb') as f: np.savez_compressed(f, **result) late_counts = np.bincount(state['neuron_idx'][state['step'] >= spec['steps'] - 18], minlength=166700) pending = late_counts[incoming.src_global] report.update(case=row['case'], repeat=row['repeat'], kind=row['kind'], pass_all=True, exact=exact, full_KC_sparse_events_and_final_state_exact=True, every_tick_fixture_pre_post_state_exact=bool(histories), graph=graph_audit, input=input_report, steps=spec['steps'], duration_ms=spec['duration_ms'], KC_active=int(np.count_nonzero(result['counts'])), KC_spikes=int(result['counts'].sum()), pending_after_record_window=dict(edge_events=int(pending.sum()), contacts=int(np.dot(pending, raw)), nominal_abs_SI=float(np.dot(pending, np.abs(incoming.weights)))), mode=spec.get('mode'), odor_id=spec.get('odor_id'), seed=spec.get('seed'), CPU_audit_wall_seconds=time.perf_counter() - start, original_brain_time_and_wall_not_rebenchmarked=True, training=False, new_brain_simulation=False, candidate_generated=False, input_sha256={filename: digest(source / filename) for filename in ('state.npz', 'apl-inhibition.float32.npy', 'requests.csv', 'summary.json')}) save(out / 'summary.json', report) return report def main(): check() prepared = read(EXP / 'protocol/prepared-hashes.json') for row in prepared: assert digest(EXP / row['path']) == row['sha256'] population = pd.read_parquet(EXP / 'protocol/neuron-source-groups.parquet') inventory = read(EXP / 'protocol/inventory.json') reports = [] for row in inventory: try: report = audit_trial(row, population) reports.append(dict(case=row['case'], repeat=row['repeat'], pass_all=report['pass_all'])) assert sum(p.stat().st_size for p in EXP.rglob('*') if p.is_file()) < 2 * 1024 ** 3 print(row['case'], row['repeat'], 'EXACT', report['KC_active'], 'KC active; CPU audit seconds', round(report['CPU_audit_wall_seconds'], 3), flush=True) except Exception as error: save(EXP / 'audit/failure.json', dict(case=row, error=type(error).__name__, message=str(error), traceback=traceback.format_exc(), completed=reports, new_brain_simulation=False)) raise repeat_checks = [] for row in inventory: if row['repeat'] != 1: continue first = EXP / 'trials' / (row['case'] + '-r0') repeat = EXP / 'trials' / (row['case'] + '-r1') checks = {} for filename in ('per-kc.parquet', 'per-pn.parquet', 'per-tick.parquet'): a, b = pd.read_parquet(first / filename), pd.read_parquet(repeat / filename) checks[filename] = a.equals(b) with np.load(first / 'replay-final.npz') as a, np.load(repeat / 'replay-final.npz') as b: checks['replay-final'] = set(a.files) == set(b.files) and all(bit_equal(a[k], b[k]) for k in a.files) assert all(checks.values()), 'Repeat accounting differs' repeat_checks.append(dict(case=row['case'], pass_all=True, checks=checks)) assert len(reports) == 80 and len(repeat_checks) == 12 save(EXP / 'audit/completed.json', dict(pass_all=True, trials=reports, repeat_checks=repeat_checks, meaning='Exact conditional KC reconstruction and reproducible accounting only', candidate_generated=False)) print('PASS 80/80 conditional reconstructions; 12/12 accounting repeats', flush=True) if __name__ == '__main__': main()