from pathlib import Path
import subprocess as sp
import site
import sys
import re
have_timing = False
have_stats = False
have_venv = True
cur_version = [0, 10, 0]
nxt_version = [0, 11, 0]
cur_version_str = f"{cur_version[0]}.{cur_version[1]}.{cur_version[2]}"
nxt_version_str = "nxt_version[0]}.{nxt_version[1]}.{nxt_version[2]}"
here = Path(__file__).parent
if have_venv:
env = here / ".env"
if not Path(env).exists():
print(f"No .env found, setting up {here}/.env")
sp.run(
f'bash -c "/usr/bin/env python3 -mvenv {env} && source {env}/activate && pip3 install arbor=={cur_version_str} matplotlib numpy pandas cbor2"',
shell=True,
check=True,
)
version = re.match(r"(\d+\.\d+)\..+", sys.version).group(1)
site_packages = here / ".env" / "lib" / f"python{version}" / "site-packages"
old_path = list(sys.path)
site.addsitedir(site_packages)
sys.real_prefix = sys.prefix
sys.prefix = here
sys.path[:0] = [item for item in list(sys.path) if item not in old_path] + sys.path
from collections import defaultdict
import arbor as A
from arbor import units as U
import pandas as pd
import matplotlib.pyplot as plt
from time import perf_counter as pc
from cbor2 import load as load_data
from math import ceil
ver = re.match(r"(\d+)\.(\d+)\.(\d+)(-\w+)?", A.__version__)
if ver:
mj, mn, pt, sf = ver.groups()
ver = [int(mj), int(mn), int(pt)]
assert (
cur_version <= ver <= nxt_version
), f"Arbor {cur_version_str} <= version <= {nxt_version_str} is required, got {A.__version__}"
else:
print(f"Couldn't parse version {A.__version__}")
exit(-42)
def load_morphology(path):
sfx = path.suffix
if sfx == ".swc":
try:
return A.load_swc_arbor(path).morphology
except:
pass
try:
return A.load_swc_neuron(path).morphology
except:
raise RuntimeError(
f"Could load {path} neither as NEURON nor Arbor flavour."
)
elif sfx == ".asc":
return A.load_asc(path).morphology
elif sfx == ".nml":
nml = A.neuroml(path)
if len(nml.morphology_ids()) == 1:
return nml.morphology(nml.morphology_ids()[0]).morphology
else:
raise RuntimeError(f"NML file {path} contains multiple morphologies.")
else:
raise RuntimeError(f"Unknown morphology file type {sfx}")
class Timing:
def __init__(self):
self.timings = defaultdict(lambda: 0.0)
self.times = defaultdict(lambda: 0.0)
self.children = defaultdict(set)
def tic(self, key):
self.timings[key] -= pc()
def toc(self, key):
self.timings[key] += pc()
def show_times(self, root, prefix):
lbl = f"{' '*prefix}* {root}"
print(f"{lbl:<37}{self.times[root]:0.3f}")
for child in self.children[root]:
self.show_times(child, prefix + 2)
def report(self):
for path, time in self.timings.items():
last = "Total"
self.times[last] += time
for k in path.split("/"):
self.children[last].add(k)
last = k
self.times[k] += time
print(
"""
Timings
==========
"""
)
self.show_times("Total", 0)
class TimingNull:
def __init__(self):
pass
def tic(self, _):
pass
def toc(self, _):
pass
def report(self):
pass
if have_timing:
timing = Timing()
else:
timing = TimingNull()
def open_sim():
with open(here / "dat/sim.cbor", "rb") as fd:
return load_data(fd)
def close_sim(raw):
del raw
def read_int_dict(raw, key):
res = raw[key]
assert isinstance(res, dict)
return res
def read_dict(raw, key):
res = raw[key]
assert isinstance(res, dict)
return res
def read_array(raw, key):
res = raw[key]
assert isinstance(res, list)
return res
def read_int(raw, key):
res = raw[key]
assert isinstance(res, int)
return res
def read_float(raw, key):
res = raw[key]
assert isinstance(res, float)
return res
class recipe(A.recipe):
def __init__(self):
A.recipe.__init__(self)
data = open_sim()
self.N = read_int(data, "size")
self.T = read_float(data, "time")
self.dt = read_float(data, "time_step")
self.threshold = read_float(data, "spike_threshold")
self.gid_to_bio = read_int_dict(data, "cell_bio_ids")
self.gid_to_lif = read_int_dict(data, "gid_to_lif")
self.gid_to_vrt = read_int_dict(data, "gid_to_vrt")
self.mid_to_mrf = read_array(data, "morphology")
self.cid_to_acc = read_array(data, "decoration")
self.gid_to_kid = read_array(data, "cell_kind")
self.gid_to_meta = read_array(data, "metadata")
self.gid_to_inc = read_int_dict(data, "incoming_connections")
self.gid_to_syn = read_int_dict(data, "synapses")
self.gid_to_icp = read_int_dict(data, "current_clamps")
self.gid_to_prb = read_int_dict(data, "probes")
self.kind_to_count = read_array(data, "count_by_kind")
self.cable_props = A.neuron_cable_properties()
properties = read_dict(data, "cable_cell_globals")
self.cable_props.catalogue.extend(A.allen_catalogue(), "")
if properties:
self.cable_props.set_property(
Vm=properties["v_init"] * U.mV, tempK=properties["celsius"] * U.Celsius
)
max_extent = read_float(data, "max_cv_length")
if max_extent:
self.cv_policy = A.cv_policy_max_extent(max_extent)
else:
self.cv_policy = A.cv_policy_single()
close_sim(data)
self.cable_data = {}
def cell_kind(self, gid):
kind = self.gid_to_kid[gid]
if kind == 0:
return A.cell_kind.cable
elif kind == 1:
return A.cell_kind.lif
elif kind == 2:
return A.cell_kind.spike_source
else:
raise RuntimeError("Unknown cell kind")
def num_cells(self):
return self.N
def connections_on(self, gid):
if not gid in self.gid_to_inc:
return []
return [
A.connection((src, f"src-{lbl}"), f"syn-{tgt}", w, max(d, self.dt) * U.ms)
for src, lbl, tgt, w, d in self.gid_to_inc[gid]
]
def global_properties(self, kind):
if kind == A.cell_kind.cable:
return self.cable_props
return None
def cell_description(self, gid):
kind = self.gid_to_kid[gid]
if kind == 0:
return self.make_cable_cell(gid)
elif kind == 1:
return self.make_lif_cell(gid)
elif kind == 2:
return self.make_vrt_cell(gid)
else:
raise RuntimeError("Unknown cell kind")
def probes(self, gid):
res = []
if gid in self.gid_to_prb:
kind = self.cell_kind(gid)
for loc, var, tag in self.gid_to_prb[gid]:
tag = f"probe-{tag}"
if kind == A.cell_kind.cable:
loc = f'(on-components 0.5 (region "{loc}"))'
if var == "voltage":
res.append(A.cable_probe_membrane_voltage(loc, tag))
elif var.endswith("i"):
res.append(
A.cable_probe_ion_int_concentration(loc, var[:-1], tag)
)
elif var.endswith("o"):
res.append(
A.cable_probe_ion_ext_concentration(loc, var[:-1], tag)
)
else:
print(f"[UNSUPPORTED] Skipping cable probe: {var}.")
elif kind == A.cell_kind.lif:
if var == "voltage":
res.append(A.lif_probe_voltage(tag))
else:
raise RuntimeError(f"Probing var={var} not yet implemented")
else:
raise RuntimeError(f"Probing cell of kind={kind} not implemented")
return res
def make_cable_cell(self, gid):
mrf, dec = self.load_cable_data(gid)
lbl = A.label_dict().add_swc_tags()
dec.place(
"(location 0 0.5)", A.threshold_detector(self.threshold * U.mV), "src-0"
)
if gid in self.gid_to_syn:
for location, synapse, params, tag in self.gid_to_syn[gid]:
dec.place(location, A.synapse(synapse, **params), f"syn-{tag}")
if gid in self.gid_to_icp:
for loc, delay, duration, amplitude, tag in self.gid_to_icp[gid]:
dec.place(
loc,
A.iclamp(
tstart=delay * U.ms,
duration=duration * U.ms,
current=amplitude * U.nA,
),
f"ic-{tag}",
)
dec.discretization(self.cv_policy)
return A.cable_cell(mrf, dec, lbl)
def make_lif_cell(self, gid):
cell = A.lif_cell("src-0", "syn-0")
data = self.gid_to_lif[gid]
cell.C_m = 0.6 * data["cm"] * U.pF
cell.tau_m = data["tau"] * U.ms
cell.E_L = data["U_neutral"] * U.mV
cell.E_R = data["U_reset"] * U.mV
cell.V_m = data["U_0"] * U.mV
cell.V_th = data["U_th"] * U.mV
cell.t_ref = data["t_ref"] * U.ms
return cell
def make_vrt_cell(self, gid):
return A.spike_source_cell(
"src-0", A.explicit_schedule([t * U.ms for t in self.gid_to_vrt[gid]])
)
def load_cable_data(self, gid):
mid, cid = self.gid_to_bio[gid]
if not gid in self.cable_data:
timing.tic("build/simulation/io")
mrf = load_morphology(here / "mrf" / self.mid_to_mrf[mid])
dec = A.load_component(here / "acc" / self.cid_to_acc[cid]).component
self.cable_data[gid] = (mrf, dec)
timing.toc("build/simulation/io")
mrf, dec = self.cable_data[gid]
return mrf, A.decor(dec)
timing.tic("build/recipe")
rec = recipe()
timing.toc("build/recipe")
timing.tic("build/simulation")
ctx = A.context()
hints = {}
for kind, tag in zip(
[A.cell_kind.cable, A.cell_kind.lif, A.cell_kind.spike_source], [0, 1, 2]
):
if tag in rec.kind_to_count:
hints[kind] = A.partition_hint(
cpu_group_size=int(ceil(rec.kind_to_count[tag] / ctx.threads))
)
ddc = A.partition_load_balance(rec, ctx, hints)
sim = A.simulation(rec, context=ctx, domains=ddc)
timing.toc("build/simulation")
timing.tic("build/sampling")
sim.record(A.spike_recording.all)
schedule = A.regular_schedule(tstart=0 * U.ms, dt=10 * rec.dt * U.ms)
handles = {
(gid, tag): sim.sample((gid, f"probe-{tag}"), schedule=schedule)
for gid, prbs in rec.gid_to_prb.items()
for _, _, tag in prbs
}
timing.toc("build/sampling")
timing.tic("run")
sim.run(rec.T * U.ms, rec.dt * U.ms)
timing.toc("run")
timing.tic("output/spikes")
spikes = sim.spikes()
df = pd.DataFrame(
{
"time": spikes["time"],
"gid": spikes["source"]["gid"],
"lid": spikes["source"]["index"],
}
)
df["kind"] = df["gid"].map(lambda i: rec.gid_to_kid[i])
df["population"] = df["gid"].map(lambda i: rec.gid_to_meta[i]["population"])
df["type"] = df["gid"].map(lambda i: rec.gid_to_meta[i]["type_id"])
df.to_csv(here / "out" / "spikes.csv")
timing.toc("output/spikes")
timing.tic("output/samples")
for (gid, tag), handle in handles.items():
dfs = []
for data, meta in sim.samples(handle):
if isinstance(meta, list):
columns = list(map(str, meta))
else:
columns = [str(meta)]
assert data.shape[1] == len(columns) + 1
dfs.append(pd.DataFrame(data=data[:, 1:], columns=columns, index=data[:, 0]))
if not dfs:
print(f"[WARN] No data collected for tag '{tag}' on cell {gid}")
continue
df = pd.concat(dfs, axis=1)
df.index.name = "t/ms"
df.to_csv(here / "out" / f"gid_{gid}-tag_{tag}.csv")
fg, ax = plt.subplots()
df.plot(ax=ax)
fg.savefig(here / "out" / f"gid_{gid}-tag_{tag}.pdf")
timing.toc("output/samples")
if have_stats:
timing.tic("output/statistics")
N = rec.num_cells()
cells = defaultdict(lambda: defaultdict(lambda: 0))
spike = defaultdict(lambda: defaultdict(lambda: 0))
conns = defaultdict(lambda: 0)
for gid in range(N):
meta = rec.gid_to_meta[gid]
pop = meta["population"]
kind = meta["type_id"]
cells[pop][kind] += 1
cells[pop][-1] += 1
for conn in rec.connections_on(gid):
src = rec.gid_to_meta[conn.source.gid]["population"]
conns[(src, pop)] += 1
for (gid, _), _ in spikes:
meta = rec.gid_to_meta[gid]
pop = meta["population"]
kind = meta["type_id"]
spike[pop][kind] += 1
spike[pop][-1] += 1
C = sum(conns.values())
timing.toc("output/statistics")
print(
f"""
Statistics
==========
* Cells {N:>13}"""
)
for pop, kinds in cells.items():
print(f" * {pop:<20} {kinds[-1]:>13}")
for kind, num in kinds.items():
if kind == -1:
continue
print(f" * {kind:<18} {num:>13}")
print(f"* Connections {C:>13}")
for (src, tgt), num in conns.items():
print(f" * {src:<8} -> {tgt:<8} {num:>13}")
print(f"* Spikes {len(spikes):>13}")
for pop, kinds in spike.items():
print(f" * {pop:<20} {kinds[-1]:>13}")
for kind, num in kinds.items():
if kind == -1:
continue
print(f" * {kind:<18} {num:>13}")
timing.report()