71 lines
3.7 KiB
Python
71 lines
3.7 KiB
Python
"""Exact float32 frozen state arithmetic plus explicitly labelled synthetic ablations."""
|
|
import numpy as np
|
|
|
|
|
|
def params(model):
|
|
f = np.float32
|
|
return {'gd': f(1-model['dt_ms']/model['synaptic_tau_ms']),
|
|
'alpha': f(model['dt_ms']/model['membrane_tau_ms']),
|
|
'rest': f(model['resting_potential_mv']), 'reset': f(model['reset_potential_mv']),
|
|
'threshold': f(model['threshold_mv']),
|
|
'ref_steps': round(model['refractory_ms']/model['dt_ms']),
|
|
'delay_steps': round(model['synaptic_delay_ms']/model['dt_ms'])}
|
|
|
|
|
|
def state(v, g, ref, arrival, p, policy='frozen'):
|
|
if policy not in ('frozen', 'reset_g_only', 'refractory_gate_only', 'reset_and_refractory'):
|
|
raise ValueError('Unknown synthetic policy: '+policy)
|
|
r = np.maximum(ref-1, 0).astype(np.int32)
|
|
available = r == 0
|
|
gi = np.add(np.multiply(g, p['gd'], dtype=np.float32), arrival, dtype=np.float32)
|
|
if policy in ('refractory_gate_only', 'reset_and_refractory'):
|
|
gi = np.where(available, gi, g).astype(np.float32)
|
|
leak = np.subtract(v, p['rest'], dtype=np.float32)
|
|
dv = np.multiply(p['alpha'], np.subtract(gi, leak, dtype=np.float32), dtype=np.float32)
|
|
vi = np.where(available, np.add(v, dv, dtype=np.float32), p['reset']).astype(np.float32)
|
|
fire = available & (vi >= p['threshold'])
|
|
vi[fire] = p['reset']; r[fire] = p['ref_steps']
|
|
if policy in ('reset_g_only', 'reset_and_refractory'): gi[fire] = np.float32(0)
|
|
return vi, gi, r, fire
|
|
|
|
|
|
def single(arrivals, p, policy='frozen'):
|
|
v = np.array([p['rest']], np.float32); g = np.zeros(1, np.float32); r = np.zeros(1, np.int32)
|
|
history = np.empty((len(arrivals), 3), np.float32); events = []
|
|
for t, a in enumerate(arrivals):
|
|
v, g, r, fire = state(v, g, r, np.array([a], np.float32), p, policy)
|
|
if fire[0]: events.append(t)
|
|
history[t] = v[0], g[0], r[0]
|
|
return np.asarray(events, np.int32), history
|
|
|
|
|
|
def cpu_motif(nodes, edges, weights, p, steps, seed):
|
|
assert np.all(np.diff(nodes) > 0), 'Canonical nodes must be sorted and unique'
|
|
assert np.all(np.diff(edges.edge_id.to_numpy()) > 0), 'Canonical edge order required'
|
|
n = len(nodes); lookup = {int(node): i for i, node in enumerate(nodes)}
|
|
src = np.array([lookup[int(x)] for x in edges.source_neuron_idx], np.int32)
|
|
dst = np.array([lookup[int(x)] for x in edges.target_neuron_idx], np.int32)
|
|
qlen = p['delay_steps']+1; queue = np.zeros((qlen, n), np.float32)
|
|
v = np.full(n, p['rest'], np.float32); g = np.zeros(n, np.float32); ref = np.zeros(n, np.int32)
|
|
counts = np.zeros(n, np.int32); history = np.empty((steps, 3, n), np.float32)
|
|
ts, ns = [], []
|
|
for t in range(steps):
|
|
slot = t % qlen
|
|
v, g, ref, fire = state(v, g, ref, queue[slot], p); queue[slot] = 0
|
|
if t == 0 and seed is not None:
|
|
i = lookup[seed]
|
|
assert not fire[i] and ref[i] == 0
|
|
fire[i] = True; v[i] = p['reset']; ref[i] = p['ref_steps']
|
|
counts += fire
|
|
active = np.flatnonzero(fire)
|
|
ts.extend([t]*len(active)); ns.extend(nodes[active].tolist())
|
|
selected = np.flatnonzero(fire[src]); selected = selected[np.argsort(dst[selected], kind='stable')]
|
|
for target in np.unique(dst[selected]):
|
|
ix = selected[dst[selected] == target]; acc = np.float32(0)
|
|
for e in ix: acc = np.float32(acc+weights[e])
|
|
outslot = (t+p['delay_steps']) % qlen
|
|
queue[outslot, target] = np.float32(queue[outslot, target]+acc)
|
|
history[t] = v, g, ref
|
|
return {'step': np.asarray(ts, np.int32), 'neuron_idx': np.asarray(ns, np.int32), 'counts': counts,
|
|
'v': v, 'g': g, 'ref': ref, 'history': history, 'queue': queue}
|