Digital Drosophila Archive
284 lines
16 KiB
Python
284 lines
16 KiB
Python
"""Read existing records and replay PN targets only; no GPU or counterfactual."""
|
|
from functools import lru_cache
|
|
import time
|
|
import traceback
|
|
import numpy as np
|
|
import pandas as pd
|
|
import pyarrow as pa
|
|
import pyarrow.parquet as pq
|
|
from context import (ROOT, EXP, AH_MODEL, AG, COEFF, GROUPS, read, save, digest,
|
|
check, bit_equal, load_module)
|
|
from pn_observer import PNObserver
|
|
from apl_sources import decompose_apl_sources
|
|
|
|
|
|
def first_ticks(event_tick, event_node, targets):
|
|
first = np.full(166700, np.iinfo(np.int32).max, np.int32)
|
|
np.minimum.at(first, event_node, event_tick)
|
|
first[first == np.iinfo(np.int32).max] = -1
|
|
return first[targets]
|
|
|
|
|
|
def temporal_relation(source_tick, target_tick):
|
|
if source_tick < 0 or target_tick < 0:
|
|
return 'unobserved_source' if source_tick < 0 and target_tick >= 0 else (
|
|
'unobserved_target' if source_tick >= 0 else 'both_unobserved')
|
|
return 'before' if source_tick < target_tick else ('same_tick' if source_tick == target_tick else 'after')
|
|
|
|
|
|
def min_observed(values):
|
|
values = np.asarray(values)
|
|
values = values[values >= 0]
|
|
return int(values.min()) if len(values) else -1
|
|
|
|
|
|
class EventSink:
|
|
"""Streaming accepted antecedents, never silently truncate."""
|
|
def __init__(self, path):
|
|
ints64 = ('edge_id', 'incoming_edge_row')
|
|
ints32 = ('source_neuron_idx', 'target_neuron_idx', 'source_spike_tick', 'arrival_tick',
|
|
'weight_sign', 'ref_before', 'ref_at_arrival')
|
|
bools = ('target_is_design', 'accepted', 'available_at_integration')
|
|
floats = ('weight_SI', 'old_v_SI', 'old_g_SI', 'v_integrated_SI', 'v_prethreshold_SI',
|
|
'threshold_margin_SI', 'g_before_arrivals_SI')
|
|
self.schema = pa.schema([(k, t) for names, t in ((ints64, pa.int64()), (ints32, pa.int32()),
|
|
(bools, pa.bool_()), (floats, pa.float64()), (('source_group', 'record_scope'), pa.string())) for k in names])
|
|
assert not path.exists()
|
|
self.writer = pq.ParquetWriter(path, self.schema, compression='zstd')
|
|
self.rows = 0
|
|
|
|
def __call__(self, rows):
|
|
self.writer.write_table(pa.Table.from_pylist(rows, schema=self.schema))
|
|
self.rows += len(rows)
|
|
|
|
def close(self):
|
|
self.writer.close()
|
|
|
|
|
|
@lru_cache(maxsize=2)
|
|
def graph_view(graph_name):
|
|
core = load_module('sealed_ah_replay_ai_' + graph_name, AH_MODEL / 'replay_core.py')
|
|
folder = AG / 'graphs' / graph_name
|
|
nodes = np.load(folder / 'nodes.npy', mmap_mode='r')
|
|
pn = np.load(AG / 'apl-bundle/pn_globals.npy')
|
|
apl = np.load(AG / 'apl-bundle/apl_globals.npy')
|
|
targets, positions, _ = np.intersect1d(nodes, pn, assume_unique=True, return_indices=True)
|
|
assert len(targets) > 0, 'Empty PN fixture subset: stop and report nonapplicable; never fabricate targets'
|
|
lookup = np.full(len(nodes), -1, np.int32)
|
|
lookup[positions] = np.arange(len(targets), dtype=np.int32)
|
|
dst = np.load(folder / 'dst.npy', mmap_mode='r')
|
|
selected = np.flatnonzero(lookup[dst] >= 0)
|
|
src = nodes[np.load(folder / 'src.npy', mmap_mode='r')[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
|
|
raw = np.load(ROOT / 'datasets/male-cns-v1.0/processed/neuron-only/gpu/csr-weights.int64.npy', mmap_mode='r')[ids]
|
|
assert raw.dtype == np.int64 and np.all(raw > 0)
|
|
incoming = core.IncomingIndex(src, lookup[dst[selected]], ids, weights, len(targets))
|
|
groups = pd.read_parquet(EXP / 'protocol/neuron-source-groups.parquet').audit_source_group.to_numpy()
|
|
codes = np.array([GROUPS.index(x) for x in groups], np.int32)[src]
|
|
audit = dict(graph=graph_name, PN_targets=len(targets), incoming_edges=len(ids), contacts=int(raw.sum()),
|
|
zeroed_fast_APL_edges=int(np.isin(src, apl).sum()), graded_APL_to_PN=False,
|
|
groups={g: dict(edges=int((codes == i).sum()), contacts=int(raw[codes == i].sum())) for i, g in enumerate(GROUPS)},
|
|
source_graph_unchanged=True, effective_fast_APL_zero_mask_preserved=True)
|
|
return core, incoming, raw, codes, targets, positions, audit
|
|
|
|
|
|
@lru_cache(maxsize=2)
|
|
def pn_kc_edges(graph_name):
|
|
folder = AG / 'graphs' / graph_name
|
|
nodes = np.load(folder / 'nodes.npy', mmap_mode='r')
|
|
pn = np.load(AG / 'apl-bundle/pn_globals.npy')
|
|
kc = np.load(AG / 'apl-bundle/kc_globals.npy')
|
|
src = nodes[np.load(folder / 'src.npy', mmap_mode='r')]
|
|
dst = nodes[np.load(folder / 'dst.npy', mmap_mode='r')]
|
|
picked = np.flatnonzero(np.isin(src, pn) & np.isin(dst, kc))
|
|
return src[picked], np.searchsorted(kc, dst[picked]), np.load(folder / 'weights.npy', mmap_mode='r')[picked]
|
|
|
|
|
|
def pn_kc_before_first(graph_name, state, kc_first, design):
|
|
"""All PN edge arrivals strictly before first KC spike; no refractory yet."""
|
|
src, dst, weights = pn_kc_edges(graph_name)
|
|
rows = {'neuron_idx': np.load(AG / 'apl-bundle/kc_globals.npy').copy(), 'first_spike_tick': kc_first.copy()}
|
|
for name in ('design_ALPN', 'non_design_ALPN'):
|
|
rows[name + '_strict_pre_first_edge_events'] = np.zeros(len(kc_first), np.int64)
|
|
rows[name + '_strict_pre_first_positive_V'] = np.zeros(len(kc_first), np.float64)
|
|
rows[name + '_strict_pre_first_negative_abs_V'] = np.zeros(len(kc_first), np.float64)
|
|
# Fixed graph edges only. Before the first spike, initially zero ref has never been reset.
|
|
for source in np.unique(src):
|
|
edges = np.flatnonzero((src == source) & (kc_first[dst] >= 0))
|
|
spikes = state['step'][state['neuron_idx'] == source]
|
|
count = np.searchsorted(spikes, kc_first[dst[edges]] - 18, side='left')
|
|
label = 'design_ALPN' if source in design else 'non_design_ALPN'
|
|
np.add.at(rows[label + '_strict_pre_first_edge_events'], dst[edges], count)
|
|
np.add.at(rows[label + '_strict_pre_first_positive_V'], dst[edges], count * np.maximum(weights[edges], 0))
|
|
np.add.at(rows[label + '_strict_pre_first_negative_abs_V'], dst[edges], count * np.maximum(-weights[edges], 0))
|
|
return rows
|
|
|
|
|
|
def audit_trial(row, population):
|
|
start = time.perf_counter()
|
|
out = EXP / 'trials' / (row['case'] + '-r' + str(row['repeat']))
|
|
out.mkdir(exist_ok=False)
|
|
source, case = ROOT / row['source'], AG / 'cases' / row['case']
|
|
spec = read(case / 'spec.json')
|
|
graph_name = 'fixtures' if row['kind'] == 'fixture' else 'full'
|
|
core, incoming, raw, codes, targets, positions, 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'}
|
|
with np.load(case / 'requests.npz', allow_pickle=False) as data:
|
|
req = {k: data[k].copy() for k in data.files}
|
|
rt = pd.read_csv(source / 'requests.csv')
|
|
assert np.array_equal(req['tick'], rt.tick) and np.array_equal(req['neuron_idx'], rt.neuron_idx)
|
|
assert (rt.attempted == 1).all() and rt.accepted.isin([0, 1]).all()
|
|
assert np.array_equal(rt.rejected, 1 - rt.accepted)
|
|
if row['kind'] == 'calibration':
|
|
with np.load(case / 'pn-drive.npz', allow_pickle=False) as data:
|
|
design = np.sort(data['targets'].copy())
|
|
assert design.dtype == np.int32 and data['rates_hz'].dtype == np.float32
|
|
assert len(design) == spec['target_count'] and (data['rates_hz'] > 0).all()
|
|
else:
|
|
design = np.intersect1d(req['neuron_idx'], targets)
|
|
assert np.isin(design, targets).all() and len(np.unique(design)) == len(design)
|
|
expected_mask = np.isin(state['neuron_idx'], targets)
|
|
expected = {k: state[k][expected_mask] for k 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, dtype in (('v', 'float64'), ('g', 'float64'), ('ref', 'int32')):
|
|
histories[key] = np.load(source / f'{key}-history.{dtype}.npy', mmap_mode='r')[:, positions]
|
|
histories['pre_' + key] = np.load(source / f'preAPL-{key}.{dtype}.npy', mmap_mode='r')[:, positions]
|
|
sink = EventSink(out / 'accepted-prefirst-edge-events.parquet')
|
|
observer = PNObserver(targets, design, incoming, codes, GROUPS, read(COEFF), spec['steps'],
|
|
max_first_event_rows=1000000, event_sink=sink)
|
|
request_lookup = np.full(166700, -1, np.int32)
|
|
request_lookup[targets] = np.arange(len(targets), dtype=np.int32)
|
|
def callback(tick, stage):
|
|
for key in ('v', 'g', 'ref'):
|
|
if histories:
|
|
assert bit_equal(stage[key], histories[key][tick]), 'Fixture post state mismatch'
|
|
assert bit_equal(stage[key + '_pre_apl'], histories['pre_' + key][tick]), 'Fixture pre state mismatch'
|
|
indices = np.flatnonzero((req['tick'] == tick) & np.isin(req['neuron_idx'], targets))
|
|
if len(indices):
|
|
local = request_lookup[req['neuron_idx'][indices]]
|
|
assert np.array_equal(stage['request_accepted'][local], rt.accepted.to_numpy()[indices].astype(bool))
|
|
observer(tick, stage)
|
|
try:
|
|
result = core.replay(initial_v, initial_g, read(COEFF), spec['steps'], targets, req,
|
|
dict(step=state['step'], neuron_idx=state['neuron_idx']), incoming,
|
|
np.zeros((spec['steps'], len(targets)), np.float32), callback)
|
|
finally:
|
|
sink.close()
|
|
exact = {k: bit_equal(result[k], state[k][positions]) for k in ('v', 'g', 'ref', 'counts')}
|
|
exact.update({k: bit_equal(result[k], expected[k]) for k in ('step', 'neuron_idx')})
|
|
assert all(exact.values()), 'PN replay diverged: ' + str(exact)
|
|
data, observer_report = observer.finish(result)
|
|
assert sink.rows == observer_report['first_arrival_record_count']
|
|
table = pd.DataFrame(data)
|
|
table.insert(1, 'body_id', population.iloc[targets].body_id.to_numpy())
|
|
table.insert(2, 'annotation_type', population.iloc[targets].annotation_type.to_numpy())
|
|
table.to_parquet(out / 'per-pn.parquet', index=False)
|
|
edges = pd.DataFrame(observer.edge_summary())
|
|
if len(edges):
|
|
edge_positions = np.searchsorted(incoming.edge_id, edges.edge_id.to_numpy())
|
|
assert np.array_equal(incoming.edge_id[edge_positions], edges.edge_id.to_numpy())
|
|
edges['raw_contact_count'] = raw[edge_positions]
|
|
for prefix in ('source', 'target'):
|
|
global_ids = edges[prefix + '_neuron_idx'].to_numpy()
|
|
edges[prefix + '_body_id'] = population.iloc[global_ids].body_id.to_numpy()
|
|
edges[prefix + '_type'] = population.iloc[global_ids].annotation_type.to_numpy()
|
|
edges.to_parquet(out / 'prefirst-edge-summary.parquet', index=False)
|
|
with (out / 'replay-final.npz').open('xb') as f:
|
|
np.savez_compressed(f, **result)
|
|
kc = np.load(AG / 'apl-bundle/kc_globals.npy')
|
|
pn = np.load(AG / 'apl-bundle/pn_globals.npy')
|
|
first_kc = first_ticks(state['step'], state['neuron_idx'], kc)
|
|
first_pn = first_ticks(state['step'], state['neuron_idx'], pn)
|
|
first_nondesign = min_observed(first_pn[~np.isin(pn, design)])
|
|
first_anykc = min_observed(first_kc)
|
|
component = read(source / 'apl-component.json')
|
|
env = component['environment']
|
|
assert env['cuda_matmul_allow_tf32'] is False and env['float32_matmul_precision'] == 'highest'
|
|
d32 = np.array([int(component['decay_f32_bits'], 16)], np.uint32).view(np.float32)[0]
|
|
q32 = np.array([int(component['gain_f32_bits'], 16)], np.uint32).view(np.float32)[0]
|
|
apl = decompose_apl_sources(np.load(AG / 'apl-bundle/input_matrix.npy', mmap_mode='r'),
|
|
np.load(AG / 'apl-bundle/output_matrix.npy', mmap_mode='r'), pn, kc, design,
|
|
state['step'], state['neuron_idx'], np.load(source / 'apl-drive.float32.npy', mmap_mode='r'),
|
|
np.load(source / 'apl-state.float32.npy', mmap_mode='r'),
|
|
np.load(source / 'apl-inhibition.float32.npy', mmap_mode='r'),
|
|
decay32=d32, gain32=q32, standard_fp32=True, first_kc_tick=first_kc)
|
|
assert apl['summary']['pass_all']
|
|
save(out / 'apl-source-summary.json', apl['summary'])
|
|
pd.DataFrame(apl['tick_rows']).to_parquet(out / 'apl-per-tick.parquet', index=False)
|
|
per_kc = pd.DataFrame(apl['per_kc'])
|
|
per_kc.insert(1, 'body_id', population.iloc[kc].body_id.to_numpy())
|
|
per_kc['KC_present_in_source_graph'] = np.isin(kc, np.load(spec['graph_dir'] + '/nodes.npy'))
|
|
per_kc['first_non_design_PN_tick'] = first_nondesign
|
|
per_kc['first_non_design_relation'] = [temporal_relation(first_nondesign, int(t)) for t in first_kc]
|
|
for key, value in pn_kc_before_first(graph_name, state, first_kc, design).items():
|
|
if key not in ('neuron_idx', 'first_spike_tick'):
|
|
per_kc[key] = value
|
|
per_kc.to_parquet(out / 'per-kc-apl-sources.parquet', index=False)
|
|
report = dict(case=row['case'], repeat=row['repeat'], kind=row['kind'], pass_all=True, exact=exact,
|
|
every_tick_fixture_pre_post_exact=bool(histories), CAL_per_tick_independent_state_available=False,
|
|
steps=spec['steps'], duration_ms=spec['duration_ms'], mode=spec.get('mode'), seed=spec.get('seed'),
|
|
odor_id=spec.get('odor_id'), graph=graph_audit, observer=observer_report,
|
|
first_non_design_PN_tick=first_nondesign, first_any_KC_tick=first_anykc,
|
|
first_non_design_vs_first_KC=temporal_relation(first_nondesign, first_anykc),
|
|
fixture_design_definition='Unique requested PN IDs for numerical fixture only; not odor design' if histories else None,
|
|
APL_source_bound_pass=apl['summary']['pass_all'], candidate_generated=False, new_brain_simulation=False,
|
|
CPU_audit_wall_seconds=time.perf_counter() - start,
|
|
input_sha256={f: digest(source / f) for f in ('state.npz', 'requests.csv', 'summary.json',
|
|
'apl-component.json', 'apl-drive.float32.npy', 'apl-state.float32.npy', 'apl-inhibition.float32.npy')})
|
|
save(out / 'summary.json', report)
|
|
save(out / 'output-hashes.json', [dict(path=p.name, sha256=digest(p), bytes=p.stat().st_size)
|
|
for p in sorted(out.iterdir()) if p.is_file()])
|
|
return report
|
|
|
|
|
|
def main():
|
|
check()
|
|
population = pd.read_parquet(EXP / 'protocol/neuron-source-groups.parquet')
|
|
inventory = read(EXP / 'protocol/inventory.json')
|
|
receipts = []
|
|
for row in inventory:
|
|
try:
|
|
report = audit_trial(row, population)
|
|
receipts.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 PN / APL BOUND PASS',
|
|
round(report['CPU_audit_wall_seconds'], 3), 'CPU sec', 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=receipts, new_brain_simulation=False))
|
|
raise
|
|
repeats = []
|
|
for row in inventory:
|
|
if row['repeat'] != 1:
|
|
continue
|
|
a = EXP / 'trials' / (row['case'] + '-r0')
|
|
b = EXP / 'trials' / (row['case'] + '-r1')
|
|
checks = {}
|
|
for f in ('per-pn.parquet', 'prefirst-edge-summary.parquet', 'apl-per-tick.parquet', 'per-kc-apl-sources.parquet'):
|
|
checks[f] = pd.read_parquet(a / f).equals(pd.read_parquet(b / f))
|
|
checks['antecedent_events'] = digest(a / 'accepted-prefirst-edge-events.parquet') == digest(b / 'accepted-prefirst-edge-events.parquet')
|
|
checks['apl-summary'] = read(a / 'apl-source-summary.json') == read(b / 'apl-source-summary.json')
|
|
with np.load(a / 'replay-final.npz') as aa, np.load(b / 'replay-final.npz') as bb:
|
|
checks['replay'] = set(aa.files) == set(bb.files) and all(bit_equal(aa[k], bb[k]) for k in aa.files)
|
|
assert all(checks.values()), 'Repeated accounting differs'
|
|
repeats.append(dict(case=row['case'], pass_all=True, checks=checks))
|
|
assert len(receipts) == 80 and len(repeats) == 12
|
|
save(EXP / 'audit/completed.json', dict(pass_all=True, trials=receipts, repeat_checks=repeats,
|
|
meaning='Conditional PN exactness and bounded APL attribution only; not a learning or causal gate'))
|
|
print('PASS 80 conditional audits and 12 repeats', flush=True)
|
|
|
|
|
|
if __name__ == '__main__':
|
|
main()
|