[DRAFT] Add Hay et al 2011 unit test

Cargo.toml

@@ -13,6 +13,11 @@ tracing-subscriber = { version = "^0.2" }
 anyhow             = { version = "^1.0.0", features = ["backtrace"] }
 thiserror          = { version = "^1.0.0" }
 pretty_assertions  = { version = "^1.3.0" }
+serde_json         = { version = "^1.0" }
+serde              = { version = "^1.0", features = ["derive"] }
+
+[dev-dependencies]
+tempfile           = "3"
 
 [[bin]]
 name = "nmlcc"

dat/CMakeLists.txt

@@ -0,0 +1,9 @@
+cmake_minimum_required(VERSION 3.22)
+project(arbor-nmlcc LANGUAGES CXX)
+set(CMAKE_CXX_STANDARD 17)
+
+find_package(arbor REQUIRED)
+find_package(nlohmann_json 3.11.0 REQUIRED)
+
+add_executable(main main.cxx)
+target_link_libraries(main PUBLIC arbor::arbor arbor::arborio arbor::arborenv nlohmann_json::nlohmann_json)

dat/main.cxx

@@ -0,0 +1,196 @@
+#include <unordered_map>
+#include <vector>
+#include <string>
+#include <filesystem>
+#include <iostream>
+#include <fstream>
+#include <iterator>
+
+#include <arbor/context.hpp>
+#include <arbor/load_balance.hpp>
+#include <arbor/cable_cell.hpp>
+#include <arbor/simulation.hpp>
+#include <arbor/recipe.hpp>
+
+#include <arborio/neuroml.hpp>
+#include <arborio/cableio.hpp>
+
+#include <nlohmann/json.hpp>
+using json = nlohmann::json;
+
+arb::mechanism_catalogue compile(const std::filesystem::path& here) {
+    auto fn = here / \"local-catalogue.so\";
+    return arb::load_catalogue(fn);
+}
+
+std::string read_file(const std::filesystem::path& fn) {
+    auto fd = std::ifstream{fn};
+    return std::string(std::istreambuf_iterator<char>{fd}, {});
+}
+
+
+arb::locset on_segment(size_t seg, const std::string& frac) {
+    return arb::ls::on_components(std::stod(frac), arb::reg::named(std::to_string(seg)));
+};
+
+std::string mk_label(const std::string& pfx, size_t seg, const std::string& frac) {
+    return pfx + \"@seg_\" + std::to_string(seg) + \"_frac_\" + frac;
+}
+
+arb::cell_local_label_type mk_lid(const std::string& l) {
+     return {l, arb::lid_selection_policy::round_robin};
+}
+
+arb::cell_global_label_type mk_gid(arb::cell_gid_type c, const std::string& l) {
+     return {c, {l, arb::lid_selection_policy::round_robin}};
+}
+
+struct recipe: public arb::recipe {
+
+    recipe(const std::filesystem::path& cwd, const std::string& network): here{cwd} {
+        auto cat = compile(here);
+        gprop.default_parameters = arb::neuron_parameter_defaults;
+        gprop.catalogue.import(cat, prefix);
+
+        std::ifstream fd{(here / \"dat\" / network).replace_extension(\"json\")};
+        auto data = json::parse(fd);
+
+        gid_to_cell = data[\"gid_to_cell\"];
+        cell_to_morph = data[\"cell_to_morph\"];
+        i_clamps = data[\"i_clamps\"];
+        poisson_generators = data[\"poisson_generators\"];
+        regular_generators = data[\"regular_generators\"];
+        for (const auto& [k, v]: data[\"gid_to_inputs\"].items()) gid_to_inputs[std::stoi(k)] = v;
+        for (const auto& [k, v]: data[\"gid_to_synapses\"].items()) gid_to_synapses[std::stoi(k)] = v;
+        for (const auto& [k, v]: data[\"gid_to_detectors\"].items()) gid_to_detectors[std::stoi(k)] = v;
+        for (const auto& [k, v]: data[\"gid_to_connections\"].items()) gid_to_connections[std::stoi(k)] = v;
+        count = data[\"count\"];
+    }
+
+    std::any get_global_properties(arb::cell_kind) const override { return gprop; }
+    arb::cell_size_type num_cells() const override { return count; }
+    std::vector<arb::probe_info> get_probes(arb::cell_gid_type gid) const override { return {arb::cable_probe_membrane_voltage{arb::ls::location(0, 0.5)}}; }
+
+    arb::util::unique_any get_cell_description(arb::cell_gid_type gid) const override {
+        const auto& cid = gid_to_cell.at(gid);
+        const auto& mid = cell_to_morph.at(cid);
+
+        auto mrf = read_file((here / \"mrf\" / mid).replace_extension(\"nml\"));
+        auto acc = read_file((here / \"acc\" / cid).replace_extension(\"acc\"));
+
+        auto nml = arborio::neuroml{mrf};
+        auto morph = nml.morphology(mid, arborio::neuroml_options::allow_spherical_root);
+
+        auto label = arb::label_dict{};
+        label.set(\"all\", arb::reg::all());
+        label.import(morph->segments);
+        label.import(morph->named_segments);
+        label.import(morph->groups);
+
+        auto decor = std::get<arb::decor>(arborio::parse_component(acc)->component);
+
+        if (gid_to_inputs.count(gid)) {
+            for (const auto& [seg, frac, inp]: gid_to_inputs.at(gid)) {
+                if (i_clamps.count(inp)) {
+                    auto tag = on_segment(seg, frac);
+                    auto lbl = mk_label(\"ic_\" + inp, seg, frac);
+                    const auto& [lag, dur, amp] = i_clamps.at(inp);
+                    decor.place(tag, arb::i_clamp{lag, dur, amp}, lbl);
+                }
+            }
+        }
+
+        if (gid_to_synapses.count(gid)) {
+            for (const auto& [seg, frac, syn]: gid_to_synapses.at(gid)) {
+                auto tag = on_segment(seg, frac);
+                auto lbl = mk_label(\"syn_\" + syn, seg, frac);
+                decor.place(tag, arb::synapse{prefix + syn}, lbl);
+            }
+        }
+
+        if (gid_to_detectors.count(gid)) {
+            for (const auto& [seg, frac, thr]: gid_to_detectors.at(gid)) {
+                auto tag = on_segment(seg, frac);
+                auto lbl = mk_label(\"sd\", seg, frac);
+                decor.place(tag, arb::threshold_detector{thr}, lbl);
+            }
+        }
+
+        return arb::cable_cell{morph->morphology, decor, label};
+    }
+
+    arb::cell_kind get_cell_kind(arb::cell_gid_type) const override { return arb::cell_kind::cable; }
+
+    std::vector<arb::cell_connection> connections_on(arb::cell_gid_type gid) const override {
+        std::vector<arb::cell_connection> res;
+        if (gid_to_connections.count(gid)) {
+            for (const auto& [src, dec, syn, loc, wgt, del]: gid_to_connections.at(gid)) {
+                auto from = mk_gid(src, \"sd@\" + dec);
+                auto delay = std::max(min_delay, del);
+                auto to = mk_lid(\"syn_\" + syn + \"@\" + loc);
+                res.emplace_back(from, to, wgt, delay);
+            }
+        }
+        return res;
+    }
+
+    std::vector<arb::event_generator> event_generators(arb::cell_gid_type gid) const override {
+        std::vector<arb::event_generator> res;
+        if (gid_to_inputs.count(gid)) {
+            for (const auto& [seg, frac, inp]: gid_to_inputs.at(gid)) {
+                if (poisson_generators.count(inp)) {
+                    const auto& [syn, avg, wgt] = poisson_generators.at(inp);
+                    auto tgt = mk_label(\"syn_\" + syn, seg, frac);
+                    res.push_back(arb::poisson_generator(tgt, wgt, 0, avg, std::mt19937_64{gid}));
+                }
+                else if (regular_generators.count(inp)) {
+                    throw std::runtime_error{\"Regular generators not implemented yet.\"};
+                }
+                else if (i_clamps.count(inp)) {
+                    // OK, handled those above
+                }
+                else {
+                    // throw std::runtime_error{\"Unknown generator type.\"};
+                }
+            }
+        }
+        return res;
+    }
+
+    std::vector<std::string> gid_to_cell;
+    std::unordered_map<std::string, std::string> cell_to_morph;
+    std::unordered_map<int, std::vector<std::tuple<int, std::string, std::string>>> gid_to_inputs;
+    std::unordered_map<int, std::vector<std::tuple<int, std::string, std::string>>> gid_to_synapses;
+    std::unordered_map<int, std::vector<std::tuple<int, std::string, double>>> gid_to_detectors;
+    std::unordered_map<int, std::vector<std::tuple<int, std::string, std::string, std::string, double, double>>> gid_to_connections;
+    std::unordered_map<std::string, std::tuple<double, double, double>> i_clamps;
+    std::unordered_map<std::string, std::tuple<>> regular_generators;
+    std::unordered_map<std::string, std::tuple<std::string, double, double>> poisson_generators;
+
+    arb::cell_size_type count;
+    std::string prefix = \"local_\";
+    arb::cable_cell_global_properties gprop;
+
+    std::filesystem::path here;
+
+    double min_delay = 0.025;
+};
+
+int main(int argc, char* argv[]) {
+    double dt = 0.025;  // ms
+    double T  = 1000.0; // ms
+
+    if (argc != 2) {
+        std::cerr << \"Usage: main <network>\\n\";
+        return -42;
+    }
+    auto cwd = std::filesystem::path{argv[0]}.parent_path();
+    auto net = std::string{argv[1]};
+    auto ctx = arb::make_context();
+    auto mdl = recipe(cwd, net);
+    mdl.min_delay = dt;
+    auto ddc = arb::partition_load_balance(mdl, ctx);
+    auto sim = arb::simulation(mdl, ctx, ddc);
+    sim.set_binning_policy(arb::binning_kind::regular, dt);
+    sim.run(T, dt);
+}

dat/main.py

@@ -0,0 +1,140 @@
+#!/usr/bin/env python3
+
+import arbor as A
+
+import subprocess as sp
+from pathlib import Path
+from time import perf_counter as pc
+import sys
+import json
+
+here = Path(__file__).parent
+
+def compile(here):
+    here = here.resolve()
+    fn   = here / 'local-catalogue.so'
+    cat  = here / 'cat'
+    recompile = False
+    if fn.exists():
+        for src in cat.glob('*.mod'):
+            src = Path(src).resolve()
+            if src.stat().st_mtime > fn.stat().st_mtime:
+                recompile = True
+                break
+    else:
+        recompile = True
+    if recompile:
+        sp.run(f'arbor-build-catalogue local {cat}', shell=True, check=True)
+    return A.load_catalogue(fn)
+
+class recipe(A.recipe):
+    def __init__(self, network):
+        A.recipe.__init__(self)
+        self.seed = 42
+        self.prefix = 'local_'
+        self.props = A.neuron_cable_properties()
+        cat = compile(here)
+        self.props.catalogue.extend(cat, self.prefix)
+
+        with open((here / network).with_suffix('.json')) as fd:
+            data = json.load(fd)
+
+        self.gid_to_cell = data['gid_to_cell'].items()
+        self.gid_to_inputs = { int(k): v for k, v in data['gid_to_inputs'].items() }
+        self.gid_to_synapses = { int(k): v for k, v in data['gid_to_synapses'].items() }
+        self.gid_to_detectors = { int(k): v for k, v in data['gid_to_detectors'].items() }
+        self.gid_to_connections = { int(k): v for k, v in data['gid_to_connections'].items() }
+        self.cell_to_morph = data['cell_to_morph']
+        self.i_clamps = data['i_clamps']
+        self.poisson_generators = data['poisson_generators']
+        self.regular_generators = data['regular_generators']
+        self.count = data['count']
+
+    def num_cells(self):
+        return self.count
+
+    def cell_kind(self, _):
+        return A.cell_kind.cable
+
+    def cell_description(self, gid):
+        cid = self.gid_to_cell[gid]
+        mrf = self.cell_to_morph[cid]
+        nml = A.neuroml(f'{here}/mrf/{mrf}.nml').morphology(mrf, allow_spherical_root=True)
+        lbl = A.label_dict()
+        lbl.append(nml.segments())
+        lbl.append(nml.named_segments())
+        lbl.append(nml.groups())
+        lbl['all'] = '(all)'
+        dec = A.load_component(f'{here}/acc/{cid}.acc').component
+        dec.discretization(A.cv_policy_every_segment())
+        if gid in self.gid_to_inputs:
+            for seg, frac, inp in self.gid_to_inputs[gid]:
+                tag = f'(on-components {frac} (region \"{seg}\"))'
+                if inp in self.i_clamps:
+                    lag, dur, amp = self.i_clamps[inp]
+                    dec.place(tag, A.iclamp(lag, dur, amp), f'ic_{inp}@seg_{seg}_frac_{frac}')
+        if gid in self.gid_to_synapses:
+            for seg, frac, syn in self.gid_to_synapses[gid]:
+                tag = f'(on-components {frac} (region \"{seg}\"))'
+                dec.place(tag, A.synapse(self.prefix + syn), f'syn_{syn}@seg_{seg}_frac_{frac}')
+        if gid in self.gid_to_detectors:
+            for seg, frac, thr in self.gid_to_detectors[gid]:
+                tag = f'(on-components {frac} (region \"{seg}\"))'
+                dec.place(tag, A.threshold_detector(thr), f'sd@seg_{seg}_frac_{frac}')
+        return A.cable_cell(nml.morphology, dec, lbl)
+
+    def probes(self, _):
+        # Example: probe center of the root (likely the soma)
+        return [A.cable_probe_membrane_voltage('(location 0 0.5)')]
+
+    def global_properties(self, kind):
+        return self.props
+
+    def connections_on(self, gid):
+        res = []
+        if gid in self.gid_to_connections:
+            for src, dec, syn, loc, w, d in self.gid_to_connections[gid]:
+                conn = A.connection((src, A.cell_local_label(f'sd@{dec}', A.selection_policy.round_robin)), A.cell_local_label(f'syn_{syn}@{loc}', A.selection_policy.round_robin), w, d)
+                res.append(conn)
+        return res
+
+    def event_generators(self, gid):
+        res = []
+        if gid in self.gid_to_inputs:
+            for seg, frac, inp in self.gid_to_inputs[gid]:
+                if inp in self.poisson_generators:
+                    syn, avg, wgt = self.poisson_generators[inp]
+                    res.append(A.event_generator(f'syn_{syn}@seg_{seg}_frac_{frac}', wgt, A.poisson_schedule(0, avg, gid)))
+                elif inp in self.regular_generators:
+                    raise NotImplementedError()
+                else:
+                    pass
+        return res
+
+if len(sys.argv) != 2:
+    print('Usage: main.py <network-name>')
+    exit(-42)
+
+ctx = A.context()
+mdl = recipe(sys.argv[1])
+ddc = A.partition_load_balance(mdl, ctx)
+sim = A.simulation(mdl, ctx, ddc)
+hdl = sim.sample((0, 0), A.regular_schedule(0.1))
+
+print('Running simulation for 1s...')
+t0 = pc()
+sim.run(1000, 0.025)
+t1 = pc()
+print(f'Simulation done, took: {t1-t0:.3f}s')
+
+print('Trying to plot...')
+try:
+  import pandas as pd
+  import seaborn as sns
+
+  for data, meta in sim.samples(hdl):
+    df = pd.DataFrame({'t/ms': data[:, 0], 'U/mV': data[:, 1]})
+    sns.relplot(data=df, kind='line', x='t/ms', y='U/mV', ci=None).savefig('result.pdf')
+  print('Ok')
+except:
+  print('Failure, are seaborn and matplotlib installed?')

dat/run.sh

@@ -0,0 +1,9 @@
+#!/usr/bin/env bash
+
+set -euxo pipefail
+
+cmake -S . -B build
+cmake --build build
+cp build/main .
+arbor-build-catalogue local cat
+./main $*