Digital Drosophila Archive
206 lines
12 KiB
Python
206 lines
12 KiB
Python
"""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()
|