Digital Drosophila Archive
254 lines
14 KiB
Python
254 lines
14 KiB
Python
"""CPU-only, target-local shadow of the frozen author LIF + legacy APL schedule.
|
|
|
|
Targets generate their own threshold events. Recorded whole-graph events are
|
|
used ONLY as upstream, fixed-delay inputs; recorded target spikes are never
|
|
forced into the target state. APL h is an observed external component trace.
|
|
This is a conditional shadow reconstruction, not independent whole-brain or
|
|
independent APL validation. Nothing in this module imports Brian or Torch,
|
|
runs a scientific trial, reads a file, or writes a file.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
from collections.abc import Mapping
|
|
|
|
import numpy as np
|
|
|
|
|
|
N_GLOBAL = 166700
|
|
DELAY_STEPS = 18
|
|
REF_STEPS = 22
|
|
COEFFICIENT_KEYS = ("gd", "c0", "cg", "cv", "threshold", "reset", "seed_volt")
|
|
|
|
|
|
def _require(condition, message):
|
|
if not condition:
|
|
raise ValueError(message)
|
|
|
|
|
|
def _vector(value, dtype, name, size=None, finite=False):
|
|
value = np.asarray(value)
|
|
_require(value.dtype == np.dtype(dtype) and value.ndim == 1, f"{name} must be a one-dimensional {dtype} array")
|
|
_require(size is None or len(value) == size, f"{name} has the wrong length")
|
|
if finite:
|
|
_require(np.isfinite(value).all(), f"{name} contains a nonfinite value")
|
|
return value
|
|
|
|
|
|
def _coefficients(coeff):
|
|
_require(isinstance(coeff, Mapping), "coeff must be a mapping")
|
|
result = {}
|
|
for name in COEFFICIENT_KEYS:
|
|
_require(name in coeff, f"Missing coefficient: {name}")
|
|
value = coeff[name]
|
|
_require(not isinstance(value, (bool, np.bool_)) and isinstance(value, (int, float, np.integer, np.float64)),
|
|
f"{name} must be a double-precision scalar, not a float32 or nonnumeric value")
|
|
result[name] = np.float64(value)
|
|
_require(np.isfinite(result[name]), f"Nonfinite coefficient: {name}")
|
|
_require(result["seed_volt"] > result["threshold"], "Common-request seed voltage must exceed the strict threshold")
|
|
return result
|
|
|
|
|
|
class IncomingIndex:
|
|
"""Reusable immutable target-filtered anatomical edges in original order.
|
|
|
|
src_global:int32, dst_local:int32, edge_id:int64, weights:float64 (SI volts).
|
|
All vectors have the same length. Sources and original global edge IDs
|
|
must retain canonical increasing source/edge order. Edges are NOT removed
|
|
because their weight is zero, nor reordered by target. The caller must
|
|
supply the already-audited effective weights (all fast APL source edges
|
|
zeroed in a separate copy); this object never edits its input arrays.
|
|
"""
|
|
|
|
def __init__(self, src_global, dst_local, edge_id, weights, target_count):
|
|
_require(isinstance(target_count, (int, np.integer)) and not isinstance(target_count, (bool, np.bool_)) and
|
|
0 < target_count <= N_GLOBAL, "target_count must be a positive in-range integer")
|
|
self.target_count = int(target_count)
|
|
source = _vector(src_global, "int32", "src_global")
|
|
count = len(source)
|
|
destination = _vector(dst_local, "int32", "dst_local", count)
|
|
ids = _vector(edge_id, "int64", "edge_id", count)
|
|
weight = _vector(weights, "float64", "weights", count, finite=True)
|
|
_require(np.all((source >= 0) & (source < N_GLOBAL)), "An incoming source is outside the global neuron range")
|
|
_require(np.all((destination >= 0) & (destination < target_count)), "Incoming destination is outside the target-local range")
|
|
_require(np.all(ids >= 0) and np.all(ids[1:] > ids[:-1]), "Original edge IDs must be unique and canonically increasing")
|
|
_require(np.all(source[1:] >= source[:-1]), "Incoming sources must remain in original canonical order")
|
|
self.src_global = np.array(source, dtype=np.int32, copy=True, order="C")
|
|
self.dst_local = np.array(destination, dtype=np.int32, copy=True, order="C")
|
|
self.edge_id = np.array(ids, dtype=np.int64, copy=True, order="C")
|
|
self.weights = np.array(weight, dtype=np.float64, copy=True, order="C")
|
|
self.sources, starts = np.unique(self.src_global, return_index=True)
|
|
self.starts = np.asarray(starts, dtype=np.int64)
|
|
self.stops = np.concatenate((self.starts[1:], np.array([count], np.int64))) if count else np.empty(0, np.int64)
|
|
for value in (self.src_global, self.dst_local, self.edge_id, self.weights, self.sources, self.starts, self.stops):
|
|
value.setflags(write=False)
|
|
|
|
def due(self, global_spikes):
|
|
"""Return local edge rows in original queue order; no target sort/sum."""
|
|
fired = _vector(global_spikes, "int32", "global_spikes")
|
|
_require(np.all((fired >= 0) & (fired < N_GLOBAL)) and np.all(fired[1:] > fired[:-1]),
|
|
"Presynaptic spikes must be sorted, unique, in-range globals")
|
|
if not len(fired) or not len(self.sources):
|
|
return np.empty(0, np.int64)
|
|
positions = np.searchsorted(self.sources, fired)
|
|
present = positions < len(self.sources)
|
|
present[present] &= self.sources[positions[present]] == fired[present]
|
|
positions = positions[present]
|
|
if not len(positions):
|
|
return np.empty(0, np.int64)
|
|
# Each presynaptic group occupies a contiguous source-ordered block in
|
|
# this filtered graph. Concatenation preserves every target's original
|
|
# per-edge accumulation order, including zero-weight events.
|
|
return np.concatenate([np.arange(self.starts[pos], self.stops[pos], dtype=np.int64) for pos in positions])
|
|
|
|
|
|
class _EventIndex:
|
|
def __init__(self, record, steps, name):
|
|
_require(isinstance(record, Mapping), f"{name} must be a mapping")
|
|
time_keys = [key for key in ("tick", "step") if key in record]
|
|
_require(len(time_keys) == 1 and "neuron_idx" in record, f"{name} needs exactly one tick/step vector and neuron_idx")
|
|
ticks = _vector(record[time_keys[0]], "int32", name + ".tick")
|
|
nodes = _vector(record["neuron_idx"], "int32", name + ".neuron_idx", len(ticks))
|
|
_require(np.all((ticks >= 0) & (ticks < steps)), f"{name} has an event outside the trial")
|
|
_require(np.all((nodes >= 0) & (nodes < N_GLOBAL)), f"{name} contains an out-of-range global neuron")
|
|
_require(np.all((ticks[1:] > ticks[:-1]) | ((ticks[1:] == ticks[:-1]) & (nodes[1:] > nodes[:-1]))),
|
|
f"{name} must be sorted and unique by tick/global neuron")
|
|
self.nodes = np.array(nodes, copy=True, order="C")
|
|
counts = np.bincount(ticks, minlength=steps)
|
|
self.pointer = np.concatenate((np.zeros(1, np.int64), np.cumsum(counts, dtype=np.int64)))
|
|
self.nodes.setflags(write=False)
|
|
self.pointer.setflags(write=False)
|
|
|
|
def at(self, tick):
|
|
if tick < 0:
|
|
return np.empty(0, np.int32)
|
|
return self.nodes[self.pointer[tick]:self.pointer[tick + 1]]
|
|
|
|
|
|
def _state_step(v, g, ref, request, coeff, capture):
|
|
"""Separate NumPy array operations correspond to explicit CUDA RN ops."""
|
|
next_ref = np.maximum(ref - np.int32(1), np.int32(0))
|
|
available = next_ref == 0
|
|
next_v, next_g = v.copy(), g.copy()
|
|
oldg, oldv = g[available], v[available]
|
|
next_g[available] = np.multiply(coeff["gd"], oldg)
|
|
driven = np.multiply(coeff["cg"], oldg)
|
|
leaked = np.multiply(coeff["cv"], oldv)
|
|
combined = np.add(driven, leaked)
|
|
next_v[available] = np.add(coeff["c0"], combined)
|
|
result = {}
|
|
if capture:
|
|
result.update(old_v=v.copy(), old_g=g.copy(), ref_before=ref.copy(),
|
|
available=available.copy(), request=request.copy(),
|
|
request_accepted=(request & available), request_rejected=(request & ~available),
|
|
v_integrated=next_v.copy(), g_decayed=next_g.copy())
|
|
next_v[request & available] = coeff["seed_volt"]
|
|
if capture:
|
|
result["v_threshold"] = next_v.copy()
|
|
fired = available & (next_v > coeff["threshold"])
|
|
next_v[fired] = coeff["reset"]
|
|
next_g[fired] = np.float64(0.0)
|
|
next_ref[fired] = np.int32(REF_STEPS)
|
|
_require(np.isfinite(next_v).all() and np.isfinite(next_g).all(), "State step produced a nonfinite value")
|
|
result.update(v=next_v, g=next_g, ref=next_ref, fired=fired)
|
|
if capture:
|
|
result.update(v_after_reset=next_v.copy(), g_after_reset=next_g.copy(), ref_after_reset=next_ref.copy())
|
|
return result
|
|
|
|
|
|
def state_step(v, g, ref, request, coeff, *, capture=True):
|
|
"""One group/request/threshold/reset step; no arrivals or APL yet.
|
|
|
|
Inputs are never modified. v/g are float64 SI, ref is int32 in [0,22],
|
|
request is a bool vector. With capture=True this returns the corresponding
|
|
pre-threshold/intermediate callback fields as well as v/g/ref/fired.
|
|
"""
|
|
voltage = _vector(v, "float64", "v", finite=True)
|
|
n = len(voltage)
|
|
conductance = _vector(g, "float64", "g", n, finite=True)
|
|
refractory = _vector(ref, "int32", "ref", n)
|
|
requested = _vector(request, "bool", "request", n)
|
|
_require(n > 0 and np.all((refractory >= 0) & (refractory <= REF_STEPS)), "Invalid target count/refractory state")
|
|
_require(isinstance(capture, bool), "capture must be Boolean")
|
|
return _state_step(voltage, conductance, refractory, requested, _coefficients(coeff), capture)
|
|
|
|
|
|
def replay(initial_v, initial_g, coeff, steps, target_globals, requests, events,
|
|
incoming: IncomingIndex, h32, callback=None):
|
|
"""Reconstruct targets with bounded working state and sparse output events.
|
|
|
|
initial_v/initial_g: float64 SI vectors, initially nonrefractory.
|
|
target_globals: sorted unique int32 target IDs.
|
|
requests/events: mappings with tick (or step) and neuron_idx int32 vectors;
|
|
requests outside the target subset are ignored, not reassigned.
|
|
h32: float32 (steps, target_count), ALREADY mapped to target_globals, in mV.
|
|
Full 4064-KC input is valid when the target list is those same 4064 KCs.
|
|
callback(tick, stage): optional copies of old states, available/request
|
|
masks, integrated/threshold voltages, reset state, canonical due edges,
|
|
accepted arrivals, pre-APL states, h32/delta_si, and final v/g/ref. It may
|
|
raise to stop at the first caller-detected mismatch. Modifying callback
|
|
arrays cannot affect this replay or its input arrays.
|
|
|
|
No dense history is stored. Memory is O(incoming edges + logged events +
|
|
steps + targets + current due edges + predicted sparse events), excluding
|
|
caller-owned h32 and anything the callback elects to retain.
|
|
"""
|
|
_require(isinstance(steps, (int, np.integer)) and not isinstance(steps, (bool, np.bool_)) and
|
|
0 < steps <= np.iinfo(np.int32).max, "steps must be a positive int32-range integer")
|
|
steps = int(steps)
|
|
targets = _vector(target_globals, "int32", "target_globals")
|
|
n = len(targets)
|
|
_require(0 < n <= N_GLOBAL and np.all((targets >= 0) & (targets < N_GLOBAL)) and np.all(targets[1:] > targets[:-1]),
|
|
"Targets must be sorted unique in-range globals")
|
|
targets = targets.copy()
|
|
v = _vector(initial_v, "float64", "initial_v", n, finite=True).copy()
|
|
g = _vector(initial_g, "float64", "initial_g", n, finite=True).copy()
|
|
constants = _coefficients(coeff)
|
|
_require(isinstance(incoming, IncomingIndex) and incoming.target_count == n, "Incoming index target count mismatch")
|
|
_require(callback is None or callable(callback), "callback must be callable or None")
|
|
inhibition = np.asarray(h32)
|
|
_require(inhibition.dtype == np.dtype("float32") and inhibition.shape == (steps, n), "h32 must be float32 (steps, targets), already target-mapped")
|
|
for tick in range(steps):
|
|
_require(np.isfinite(inhibition[tick]).all() and np.all(inhibition[tick] >= 0), "APL inhibition must be finite and nonnegative")
|
|
request_index = _EventIndex(requests, steps, "requests")
|
|
event_index = _EventIndex(events, steps, "events")
|
|
ref = np.zeros(n, np.int32)
|
|
counts = np.zeros(n, np.int32)
|
|
first_spike = np.full(n, -1, np.int32)
|
|
predicted_steps, predicted_nodes = [], []
|
|
capture = callback is not None
|
|
for tick in range(steps):
|
|
requested_globals = request_index.at(tick)
|
|
local = np.searchsorted(targets, requested_globals)
|
|
present = local < n
|
|
present[present] &= targets[local[present]] == requested_globals[present]
|
|
request = np.zeros(n, dtype=bool)
|
|
request[local[present]] = True
|
|
stage = _state_step(v, g, ref, request, constants, capture)
|
|
v, g, ref, fired = stage["v"], stage["g"], stage["ref"], stage["fired"]
|
|
counts += fired.astype(np.int32)
|
|
first_spike[fired & (first_spike < 0)] = np.int32(tick)
|
|
fired_local = np.flatnonzero(fired)
|
|
if len(fired_local):
|
|
predicted_steps.append(np.full(len(fired_local), tick, np.int32))
|
|
predicted_nodes.append(targets[fired_local].copy())
|
|
due = incoming.due(event_index.at(tick - DELAY_STEPS))
|
|
accepted = ref[incoming.dst_local[due]] == 0
|
|
if len(due):
|
|
accepted_edges = due[accepted]
|
|
# add.at is deliberately unbuffered: start at existing g and add
|
|
# each accepted edge in original order, NEVER sum weights first.
|
|
np.add.at(g, incoming.dst_local[accepted_edges], incoming.weights[accepted_edges])
|
|
if capture:
|
|
stage.update(due_edge_indices=due.copy(), due_edge_id=incoming.edge_id[due].copy(),
|
|
arrival_accepted=accepted.copy(), v_pre_apl=v.copy(), g_pre_apl=g.copy(), ref_pre_apl=ref.copy())
|
|
delta_si = inhibition[tick].astype(np.float64) * np.float64(.001)
|
|
g = np.subtract(g, delta_si)
|
|
_require(np.isfinite(v).all() and np.isfinite(g).all(), "Arrival/APL stage produced a nonfinite value")
|
|
if capture:
|
|
stage.update(h32=inhibition[tick].copy(), delta_si=delta_si.copy(), v=v.copy(), g=g.copy(), ref=ref.copy(), fired=fired.copy())
|
|
callback(tick, stage)
|
|
return {"v": v.copy(), "g": g.copy(), "ref": ref.copy(), "counts": counts,
|
|
"first_spike": first_spike, "target_globals": targets,
|
|
"step": np.concatenate(predicted_steps) if predicted_steps else np.empty(0, np.int32),
|
|
"neuron_idx": np.concatenate(predicted_nodes) if predicted_nodes else np.empty(0, np.int32)}
|