Digital Drosophila Archive
347 lines
18 KiB
Python
347 lines
18 KiB
Python
"""Run the pinned author's create_model in a dedicated, serial Brian2 process.
|
|
|
|
This module never edits the author source or its parameters. The caller supplies
|
|
an exclusive output directory and the already validated, signed connectivity
|
|
multipliers in ``raw_count``. Global IDs are not FlyWire IDs and are mapped only
|
|
for this small-circuit diagnostic. Brian stores voltages in SI volts internally;
|
|
the returned state and weights are converted to mV in the selected dtype.
|
|
"""
|
|
from __future__ import annotations
|
|
|
|
import hashlib
|
|
import importlib.metadata
|
|
import json
|
|
import platform
|
|
import sys
|
|
import types
|
|
from pathlib import Path
|
|
|
|
import numpy as np
|
|
import pandas as pd
|
|
|
|
|
|
AUTHOR_SHA256 = "fc45837d7122c6ce2a7f3f2f23c515992e4b232aadb919efabb72337fac88e4e"
|
|
AUTHOR_COMMIT = "91bdd1e7dcf193f3e7ca5a8933497fcef63b7960"
|
|
PACKAGE_PINS = {
|
|
"brian2": "2.5.1", "numpy": "1.22.3", "pandas": "1.4.3",
|
|
"scipy": "1.8.1", "pyarrow": "11.0.0", "joblib": "1.2.0",
|
|
"cython": "0.29.33", "jinja2": "2.11.3", "markupsafe": "1.1.1",
|
|
"sympy": "1.12", "mpmath": "1.3.0", "pyparsing": "3.0.9",
|
|
"packaging": "23.1", "setuptools": "67.8.0",
|
|
"python-dateutil": "2.8.2", "pytz": "2023.3", "six": "1.16.0",
|
|
"wheel": "0.38.4",
|
|
}
|
|
|
|
|
|
def _json(path: Path, value) -> None:
|
|
with path.open("x", encoding="utf-8") as output:
|
|
json.dump(value, output, indent=2, sort_keys=True, allow_nan=False)
|
|
output.write("\n")
|
|
|
|
|
|
def _steps(value: float, dt: float, label: str, *, positive=False) -> int:
|
|
ratio = float(value) / dt
|
|
if not np.isfinite(ratio) or ratio < 0:
|
|
raise ValueError(f"{label} must be finite and nonnegative")
|
|
result = int(round(ratio))
|
|
if not np.isclose(ratio, result, rtol=0, atol=1e-9):
|
|
raise ValueError(f"{label} must be an integer number of dt steps")
|
|
if result > np.iinfo(np.int32).max or (positive and result == 0):
|
|
raise ValueError(f"{label} is outside the supported step range")
|
|
return result
|
|
|
|
|
|
def _int_vector(values, label: str, dtype=np.int64) -> np.ndarray:
|
|
source = np.asarray(values)
|
|
if source.ndim != 1 or source.dtype.kind not in "iu":
|
|
raise ValueError(f"{label} must be a one-dimensional integer array")
|
|
limits = np.iinfo(dtype)
|
|
if np.any(source < limits.min) or np.any(source > limits.max):
|
|
raise ValueError(f"{label} is outside {np.dtype(dtype).name} range")
|
|
return source.astype(dtype, copy=True)
|
|
|
|
|
|
def _environment() -> dict:
|
|
packages = {}
|
|
for name, expected in PACKAGE_PINS.items():
|
|
try:
|
|
actual = importlib.metadata.version(name)
|
|
except importlib.metadata.PackageNotFoundError:
|
|
actual = None
|
|
packages[name] = {"expected": expected, "actual": actual,
|
|
"matches": actual == expected}
|
|
return {
|
|
"python": platform.python_version(), "author_python": "3.10.11",
|
|
"python_patch_matches": platform.python_version() == "3.10.11",
|
|
"python_executable": sys.executable, "platform": platform.platform(),
|
|
"author_platform": "Windows", "packages": packages,
|
|
"all_package_pins_match": all(p["matches"] for p in packages.values()),
|
|
"binary_identical_author_environment_claimed": False,
|
|
}
|
|
|
|
|
|
def _code_text(value):
|
|
if isinstance(value, dict):
|
|
return {str(key): _code_text(item) for key, item in value.items()}
|
|
return str(value)
|
|
|
|
|
|
def _code_records(network) -> list:
|
|
records = []
|
|
for obj in network.sorted_objects:
|
|
record = {"name": obj.name, "type": type(obj).__name__,
|
|
"when": obj.when, "order": int(obj.order),
|
|
"active": bool(obj.active)}
|
|
for key in ("abstract_code", "user_code"):
|
|
if hasattr(obj, key):
|
|
record[key] = _code_text(getattr(obj, key))
|
|
record["code_objects"] = [
|
|
{"name": code.name, "type": type(code).__name__,
|
|
"generated_code": _code_text(code.code)}
|
|
for code in getattr(obj, "code_objects", ())
|
|
]
|
|
records.append(record)
|
|
return records
|
|
|
|
|
|
def run_brian(author_path: Path, io_dir: Path, nodes: np.ndarray,
|
|
edges: pd.DataFrame, config: dict, precision: str,
|
|
seed_node: int | None, external_arrivals=None) -> dict:
|
|
"""Execute an unchanged author model and return canonical end-of-tick arrays.
|
|
|
|
``external_arrivals`` is either None or a (T, N) array in mV. Every nonzero
|
|
entry gets its own SpikeGeneratorGroup source and zero-delay conditional
|
|
synaptic write, ordered after the anatomical synapses. It is not direct
|
|
Python assignment to g. ``seed_margin_mv`` may be zero for a strict-threshold
|
|
fixture; positive margins must remain representably above the threshold.
|
|
|
|
This function controls process-global Brian preferences and is not thread
|
|
safe. Use only the isolated oracle environment, serially or in subprocesses.
|
|
A failed run keeps its exclusive inputs and an error record for diagnosis.
|
|
"""
|
|
if precision not in ("float64", "float32"):
|
|
raise ValueError("precision must be float64 or float32")
|
|
dtype = np.dtype(precision)
|
|
nodes = _int_vector(nodes, "nodes", np.int32)
|
|
if not len(nodes) or np.any(np.diff(nodes.astype(np.int64)) <= 0):
|
|
raise ValueError("nodes must be nonempty, sorted, and unique")
|
|
required = ["edge_id", "source_neuron_idx", "target_neuron_idx", "raw_count"]
|
|
if not set(required).issubset(edges.columns):
|
|
raise ValueError(f"edges require columns {required}")
|
|
edge_ids = _int_vector(edges.edge_id, "edge_id")
|
|
if np.any(np.diff(edge_ids) <= 0):
|
|
raise ValueError("edge_id order must be canonical, increasing, and unique")
|
|
sources = _int_vector(edges.source_neuron_idx, "source_neuron_idx", np.int32)
|
|
targets = _int_vector(edges.target_neuron_idx, "target_neuron_idx", np.int32)
|
|
multipliers = _int_vector(edges.raw_count, "raw_count")
|
|
if not np.isin(sources, nodes).all() or not np.isin(targets, nodes).all():
|
|
raise ValueError("all edge endpoints must occur in nodes")
|
|
lookup = {int(node): i for i, node in enumerate(nodes)}
|
|
if seed_node is not None and int(seed_node) not in lookup:
|
|
raise ValueError("seed_node must occur in nodes")
|
|
|
|
dt_ms = float(config["dt_ms"])
|
|
if not np.isfinite(dt_ms) or dt_ms <= 0:
|
|
raise ValueError("dt_ms must be finite and positive")
|
|
total_steps = _steps(config["duration_ms"], dt_ms, "duration_ms", positive=True)
|
|
model = config["model"]
|
|
refractory_steps = _steps(model["refractory_ms"], dt_ms, "refractory_ms")
|
|
delay_steps = _steps(model["synaptic_delay_ms"], dt_ms, "synaptic_delay_ms")
|
|
margin_mv = float(config["seed_margin_mv"])
|
|
if not np.isfinite(margin_mv) or margin_mv < 0:
|
|
raise ValueError("seed_margin_mv must be finite and nonnegative")
|
|
if "dt_ms" in model and float(model["dt_ms"]) != dt_ms:
|
|
raise ValueError("top-level and model dt_ms disagree")
|
|
arrivals = None
|
|
if external_arrivals is not None:
|
|
arrivals = np.asarray(external_arrivals, dtype=np.float64)
|
|
if arrivals.shape != (total_steps, len(nodes)) or not np.isfinite(arrivals).all():
|
|
raise ValueError("external_arrivals must be finite with shape (T, N)")
|
|
|
|
author_path = Path(author_path).resolve(strict=True)
|
|
author_bytes = author_path.read_bytes()
|
|
author_hash = hashlib.sha256(author_bytes).hexdigest()
|
|
if author_hash != AUTHOR_SHA256:
|
|
raise ValueError(f"Author source hash mismatch: {author_hash}")
|
|
io_dir = Path(io_dir).resolve()
|
|
io_dir.mkdir(parents=True, exist_ok=False)
|
|
environment = _environment()
|
|
_json(io_dir / "environment.json", environment)
|
|
if environment["packages"]["brian2"]["actual"] != "2.5.1":
|
|
raise RuntimeError("The author oracle requires exactly Brian2 2.5.1")
|
|
|
|
import brian2 as b2
|
|
|
|
saved_preferences = {
|
|
"codegen.target": b2.prefs["codegen.target"],
|
|
"core.default_float_dtype": b2.prefs["core.default_float_dtype"],
|
|
"legacy.refractory_timing": b2.prefs["legacy.refractory_timing"],
|
|
}
|
|
saved_dt = b2.defaultclock.dt
|
|
try:
|
|
b2.start_scope()
|
|
b2.prefs["codegen.target"] = "numpy"
|
|
b2.prefs["core.default_float_dtype"] = dtype.type
|
|
b2.prefs["legacy.refractory_timing"] = False
|
|
b2.defaultclock.dt = dt_ms * b2.ms
|
|
# Compile the checked bytes themselves: no source rewrite or __pycache__
|
|
# is possible, and a subsequent file change cannot alter imported code.
|
|
author = types.ModuleType("phase4ad_pinned_author_" + author_hash[:12])
|
|
author.__file__ = str(author_path)
|
|
exec(compile(author_bytes, str(author_path), "exec"), author.__dict__)
|
|
parameters = dict(author.default_params)
|
|
checks = {}
|
|
for config_key, author_key, unit in (
|
|
("membrane_tau_ms", "t_mbr", b2.ms),
|
|
("synaptic_tau_ms", "tau", b2.ms),
|
|
("refractory_ms", "t_rfc", b2.ms),
|
|
("synaptic_delay_ms", "t_dly", b2.ms),
|
|
("resting_potential_mv", "v_0", b2.mV),
|
|
("reset_potential_mv", "v_rst", b2.mV),
|
|
("threshold_mv", "v_th", b2.mV),
|
|
("w_syn_mv_per_anatomical_synapse", "w_syn", b2.mV),
|
|
):
|
|
expected = float(model[config_key])
|
|
actual = float(parameters[author_key] / unit)
|
|
if not np.isfinite(expected) or not np.isclose(actual, expected, rtol=0, atol=1e-12):
|
|
raise ValueError(f"config {config_key}={expected} differs from author {actual}")
|
|
checks[config_key] = {"config": expected, "author": actual}
|
|
|
|
completeness = pd.DataFrame({"local_index": np.arange(len(nodes), dtype=np.int32)},
|
|
index=pd.Index(nodes, name="global_neuron_idx"))
|
|
connections = pd.DataFrame({
|
|
"Presynaptic_Index": [lookup[int(x)] for x in sources],
|
|
"Postsynaptic_Index": [lookup[int(x)] for x in targets],
|
|
"Excitatory x Connectivity": multipliers,
|
|
"edge_id": edge_ids,
|
|
"source_neuron_idx": sources, "target_neuron_idx": targets,
|
|
}).astype({"Presynaptic_Index": "int32", "Postsynaptic_Index": "int32"})
|
|
comp_path = io_dir / "completeness.csv"
|
|
con_path = io_dir / "connectivity.parquet"
|
|
completeness.to_csv(comp_path)
|
|
connections.to_parquet(con_path, index=False, engine="pyarrow")
|
|
neu, syn, spikes = author.create_model(comp_path, con_path, parameters)
|
|
if np.dtype(neu.variables["v"].dtype) != dtype or np.dtype(syn.variables["w"].dtype) != dtype:
|
|
raise AssertionError("Brian did not allocate the requested floating-point dtype")
|
|
np.testing.assert_array_equal(np.asarray(syn.i[:]), connections.Presynaptic_Index)
|
|
np.testing.assert_array_equal(np.asarray(syn.j[:]), connections.Postsynaptic_Index)
|
|
|
|
objects = [neu, syn, spikes]
|
|
seed_state = {"requested_global_node": None if seed_node is None else int(seed_node),
|
|
"margin_mv": margin_mv, "injections": 0}
|
|
if seed_node is not None:
|
|
seed_local = lookup[int(seed_node)]
|
|
|
|
def inject_seed():
|
|
if seed_state["injections"] == 0:
|
|
if float(b2.defaultclock.t / b2.ms) != 0:
|
|
raise AssertionError("Seed injection must occur at tick zero")
|
|
neu.v[seed_local] = parameters["v_th"] + margin_mv * b2.mV
|
|
applied = float(neu.v[seed_local] / b2.mV)
|
|
if margin_mv > 0 and not applied > float(parameters["v_th"] / b2.mV):
|
|
raise ValueError("The positive seed margin rounded away in the selected dtype")
|
|
seed_state["applied_mv"] = applied
|
|
seed_state["injections"] = 1
|
|
|
|
objects.append(b2.NetworkOperation(inject_seed, clock=neu.clock,
|
|
when="before_thresholds", order=0,
|
|
name="oracle_seed_injection"))
|
|
|
|
external_metadata = {"events": 0, "ordering": "after default_synapses.pre"}
|
|
if arrivals is not None:
|
|
event_ticks, event_targets = np.nonzero(arrivals)
|
|
event_values = arrivals[event_ticks, event_targets]
|
|
external_metadata["events"] = int(len(event_ticks))
|
|
if len(event_ticks):
|
|
event_ids = np.arange(len(event_ticks), dtype=np.int32)
|
|
external_source = b2.SpikeGeneratorGroup(
|
|
len(event_ticks), indices=event_ids,
|
|
times=event_ticks.astype(np.float64) * dt_ms * b2.ms,
|
|
sorted=True, clock=neu.clock, name="zz_external_sources")
|
|
external_syn = b2.Synapses(
|
|
external_source, neu, model="w : volt", on_pre="g_post += w",
|
|
delay=0 * b2.ms, name="zz_external_arrivals")
|
|
external_syn.connect(i=event_ids, j=event_targets)
|
|
external_syn.w = event_values * b2.mV
|
|
external_syn.pre.order = int(syn.pre.order) + 1
|
|
if external_syn.pre.when != syn.pre.when:
|
|
raise AssertionError("External and anatomical pathways must share a schedule slot")
|
|
objects.extend([external_source, external_syn])
|
|
external_metadata["pathway_order"] = int(external_syn.pre.order)
|
|
external_metadata["actual_weights_mv"] = np.asarray(external_syn.w[:] / b2.mV, dtype=dtype).tolist()
|
|
pd.DataFrame({"tick": event_ticks, "neuron_idx": nodes[event_targets],
|
|
"requested_mv": event_values}).to_csv(io_dir / "external-arrivals.csv", index=False)
|
|
|
|
trace = b2.StateMonitor(neu, ["v", "g", "lastspike", "not_refractory"],
|
|
record=True, when="end", name="oracle_end_state")
|
|
objects.append(trace)
|
|
network = b2.Network(*objects)
|
|
_json(io_dir / "run-spec.json", {
|
|
"author_sha256": author_hash, "author_commit": AUTHOR_COMMIT,
|
|
"precision": precision, "config": config, "parameter_checks": checks,
|
|
"dt_ms": dt_ms, "steps": total_steps, "refractory_steps": refractory_steps,
|
|
"delay_steps": delay_steps, "nodes": nodes.tolist(), "edge_ids": edge_ids.tolist(),
|
|
"seed": seed_state.copy(), "external_arrivals": external_metadata,
|
|
"preferences": {"codegen.target": "numpy", "legacy.refractory_timing": False,
|
|
"core.default_float_dtype": precision},
|
|
"weight_rule": "raw_count multiplied by unchanged author w_syn; no sign inference",
|
|
"internal_voltage_unit": "volt", "returned_voltage_unit": "mV",
|
|
})
|
|
_json(io_dir / "schedule.json", {"schedule": list(network.schedule),
|
|
"summary": str(network.scheduling_summary())})
|
|
network.run(float(config["duration_ms"]) * b2.ms, namespace={})
|
|
_json(io_dir / "generated-code.json", _code_records(network))
|
|
|
|
ticks_float = np.asarray(spikes.t[:] / b2.ms, dtype=np.float64) / dt_ms
|
|
ticks = np.rint(ticks_float).astype(np.int32)
|
|
if not np.allclose(ticks_float, ticks, rtol=0, atol=1e-6):
|
|
raise AssertionError("Spike timestamps do not lie on the configured integer clock")
|
|
event_nodes = nodes[np.asarray(spikes.i[:], dtype=np.int32)]
|
|
order = np.lexsort((event_nodes, ticks))
|
|
ticks, event_nodes = ticks[order], event_nodes[order]
|
|
counts = np.asarray(spikes.count[:], dtype=np.int32).copy()
|
|
if int(counts.sum()) != len(ticks):
|
|
raise AssertionError("SpikeMonitor counts disagree with event records")
|
|
trace_ticks = np.asarray(trace.t[:] / b2.ms, dtype=np.float64) / dt_ms
|
|
if trace_ticks.shape != (total_steps,) or not np.allclose(
|
|
trace_ticks, np.arange(total_steps), rtol=0, atol=1e-6):
|
|
raise AssertionError("End-state trace has the wrong tick range")
|
|
voltage = np.asarray(trace.v / b2.mV, dtype=dtype).T.copy()
|
|
conductance = np.asarray(trace.g / b2.mV, dtype=dtype).T.copy()
|
|
last_spike_ms = np.asarray(trace.lastspike / b2.ms, dtype=np.float64).T
|
|
seen = np.isfinite(last_spike_ms) & (last_spike_ms >= 0)
|
|
# Brian's never-spiked sentinel can be far outside int32. Convert only
|
|
# real nonnegative timestamps, then compute an end-of-tick countdown.
|
|
last_ticks = np.zeros(last_spike_ms.shape, dtype=np.int64)
|
|
last_ticks[seen] = np.rint(last_spike_ms[seen] / dt_ms).astype(np.int64)
|
|
elapsed = np.arange(total_steps, dtype=np.int64)[:, None] - last_ticks
|
|
ref = np.where(seen, np.maximum(refractory_steps - elapsed, 0), 0).astype(np.int32)
|
|
available = np.asarray(trace.not_refractory, dtype=bool).T
|
|
if not np.array_equal(available, ref == 0):
|
|
raise AssertionError("Recorded refractory flag disagrees with the integer countdown")
|
|
history = np.stack([voltage, conductance, ref.astype(dtype)], axis=1)
|
|
if not np.isfinite(history).all():
|
|
raise AssertionError("Nonfinite author state")
|
|
if seed_node is not None:
|
|
seed_state["spiked_at_zero"] = bool(np.any((ticks == 0) & (event_nodes == seed_node)))
|
|
if seed_state["injections"] != 1 or (margin_mv > 0 and not seed_state["spiked_at_zero"]):
|
|
raise AssertionError("The common initial seed event was not produced exactly once")
|
|
_json(io_dir / "seed.json", seed_state)
|
|
weights = np.asarray(syn.w[:] / b2.mV, dtype=dtype).copy()
|
|
_json(io_dir / "actual-weights.json", {
|
|
"edge_id": edge_ids.tolist(), "weights_mv": weights.tolist(),
|
|
"weights_volt": np.asarray(syn.w_[:], dtype=np.float64).tolist(),
|
|
})
|
|
if hashlib.sha256(author_path.read_bytes()).hexdigest() != author_hash:
|
|
raise AssertionError("Author source changed during execution")
|
|
return {"step": ticks, "neuron_idx": event_nodes, "counts": counts,
|
|
"v": voltage[-1].copy(), "g": conductance[-1].copy(),
|
|
"ref": ref[-1].copy(), "history": history, "weights": weights}
|
|
except Exception as error:
|
|
_json(io_dir / "error.json", {"type": type(error).__name__, "message": str(error)})
|
|
raise
|
|
finally:
|
|
for name, value in saved_preferences.items():
|
|
b2.prefs[name] = value
|
|
b2.defaultclock.dt = saved_dt
|