From e8851f9289127388e10972843a43823f0905d81f Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Tue, 17 Dec 2024 15:47:20 +0100
Subject: [PATCH 01/13] Make brunel able to use cable cells.

---
 example/brunel/brunel.cpp | 190 +++++++++++++++++++++++---------------
 1 file changed, 113 insertions(+), 77 deletions(-)

diff --git a/example/brunel/brunel.cpp b/example/brunel/brunel.cpp
index bfcfca9d0b..c913686ca0 100644
--- a/example/brunel/brunel.cpp
+++ b/example/brunel/brunel.cpp
@@ -14,12 +14,14 @@
 #include <arbor/domain_decomposition.hpp>
 #include <arbor/event_generator.hpp>
 #include <arbor/lif_cell.hpp>
+#include <arbor/cable_cell.hpp>
 #include <arbor/load_balance.hpp>
 #include <arbor/profile/meter_manager.hpp>
 #include <arbor/profile/profiler.hpp>
 #include <arbor/recipe.hpp>
 #include <arbor/simulation.hpp>
 #include <arbor/version.hpp>
+#include <arbor/morph/segment_tree.hpp>
 
 #include <arborenv/default_env.hpp>
 #include <arborenv/gpu_env.hpp>
@@ -46,19 +48,15 @@ struct cl_options {
     float delay = 0.1;
     float rel_inh_strength = 1;
     double poiss_lambda = 1;
-
+    // use cable cells instead of LIF
+    bool use_cc = false;
     // Simulation running parameters:
     double tfinal = 100.;
     double dt = 1;
     uint32_t group_size = 10;
     uint32_t seed = 42;
-
     // Parameters for spike output.
     std::string spike_file_output = "";
-
-    // Turn on/off profiling output for all ranks.
-    bool profile_only_zero = false;
-
     // Be more verbose with informational messages.
     bool verbose = false;
 };
@@ -70,8 +68,6 @@ std::optional<cl_options> read_options(int argc, char** argv);
 void banner(context ctx);
 
 // Add m unique connection from gids in interval [start, end) - gid.
-// We exclude gid because we don't want self-loops. Returns number of
-// actually _added_ connections.
 void add_subset(cell_gid_type gid,
                 cell_gid_type start, cell_gid_type end,
                 unsigned m,
@@ -80,20 +76,30 @@ void add_subset(cell_gid_type gid,
                 std::vector<cell_connection>& conns);
 
 /*
-   A Brunel network consists of nexc excitatory LIF neurons and ninh inhibitory LIF neurons.
-   Each neuron in the network receives in_degree_prop * nexc excitatory connections
-   chosen randomly, in_degree_prop * ninh inhibitory connections and next (external) Poisson connections.
-   All the connections have the same delay. The strenght of excitatory and Poisson connections is given by
-   parameter weight, whereas the strength of inhibitory connections is rel_inh_strength * weight.
-   Poisson neurons all spike independently with expected number of spikes given by parameter poiss_lambda.
-   Because of the refractory period, the activity is mostly driven by Poisson neurons and
+   A Brunel network consists of nexc excitatory LIF neurons and ninh inhibitory
+   LIF neurons. Each neuron in the network receives in_degree_prop * nexc
+   excitatory connections chosen randomly, in_degree_prop * ninh inhibitory
+   connections and next (external) Poisson connections. All the connections have
+   the same delay. The strenght of excitatory and Poisson connections is given
+   by parameter weight, whereas the strength of inhibitory connections is
+   rel_inh_strength * weight. Poisson neurons all spike independently with
+   expected number of spikes given by parameter poiss_lambda. Because of the
+   refractory period, the activity is mostly driven by Poisson neurons and
    recurrent connections have a small effect.
  */
 class brunel_recipe: public recipe {
 public:
-    brunel_recipe(cell_size_type nexc, cell_size_type ninh, cell_size_type next, double in_degree_prop,
-                  float weight, float delay, float rel_inh_strength, double poiss_lambda, int seed = 42):
-        ncells_exc_(nexc), ncells_inh_(ninh), delay_(delay), seed_(seed) {
+    brunel_recipe(cell_size_type nexc,
+                  cell_size_type ninh,
+                  cell_size_type next,
+                  double in_degree_prop,
+                  float weight,
+                  float delay,
+                  float rel_inh_strength,
+                  double poiss_lambda,
+                  int seed=42,
+                  bool use_cc=false):
+        ncells_exc_(nexc), ncells_inh_(ninh), delay_(delay), seed_(seed), use_cable_cells(use_cc) {
         // Make sure that in_degree_prop in the interval (0, 1]
         if (in_degree_prop <= 0.0 || in_degree_prop > 1.0) {
             throw std::out_of_range("The proportion of incoming connections should be in the interval (0, 1].");
@@ -108,6 +114,19 @@ class brunel_recipe: public recipe {
         // each cell receives next incoming Poisson sources with mean rate poiss_lambda, which is equivalent
         // to a single Poisson source with mean rate next*poiss_lambda
         lambda_ = next * poiss_lambda;
+        // construct cable cell prototype
+        {
+            arb::segment_tree tree;
+            tree.append(arb::mnpos, {-1.0, 0, 0, 1.0}, {1, 0, 0, 1.0}, 1);
+
+            auto dec = arb::decor{}
+                .paint(arb::reg::tagged(1), arb::density("hh"))
+                .place(arb::ls::location(0, 0.5), arb::synapse("expsyn"), "tgt")
+                .place(arb::ls::location(0, 0.5), arb::threshold_detector(-10 * U::mV), "src");
+
+            cable = arb::cable_cell(tree, dec, {});
+        }
+        prop.default_parameters = arb::neuron_parameter_defaults;
     }
 
     cell_size_type num_cells() const override {
@@ -115,7 +134,12 @@ class brunel_recipe: public recipe {
     }
 
     cell_kind get_cell_kind(cell_gid_type gid) const override {
-        return cell_kind::lif;
+        if (use_cable_cells) {
+            return cell_kind::cable;
+        }
+        else {
+            return cell_kind::lif;
+        }
     }
 
     std::vector<cell_connection> connections_on(cell_gid_type gid) const override {
@@ -127,16 +151,16 @@ class brunel_recipe: public recipe {
     }
 
     util::unique_any get_cell_description(cell_gid_type gid) const override {
-        auto cell = lif_cell("src", "tgt");
-        cell.tau_m = 10*U::ms;
-        cell.V_th = 10*U::mV;
-        cell.C_m = 20*U::pF;
-        cell.E_L = 0*U::mV;
-        cell.V_m = 0*U::mV;
-        cell.t_ref = 2*U::ms;
-        return cell;
+        if (use_cable_cells) {
+            return cable;
+        }
+        else {
+            return lif;
+        }
     }
 
+    std::any get_global_properties(cell_kind) const override { return prop; }
+
     std::vector<event_generator> event_generators(cell_gid_type gid) const override {
         return {poisson_generator({"tgt"}, weight_ext_, 0*arb::units::ms, lambda_*arb::units::kHz, gid + seed_)};
     }
@@ -144,33 +168,40 @@ class brunel_recipe: public recipe {
 private:
     // Number of excitatory cells.
     cell_size_type ncells_exc_;
-
     // Number of inhibitory cells.
     cell_size_type ncells_inh_;
-
     // Weight of excitatory synapses.
     float weight_exc_;
-
     // Weight of inhibitory synapses.
     float weight_inh_;
-
     // Weight of external Poisson cell synapses.
     float weight_ext_;
-
     // Delay of all synapses.
     float delay_;
-
     // Number of connections that each neuron receives from excitatory population.
     int in_degree_exc_;
-
     // Number of connections that each neuron receives from inhibitory population.
     int in_degree_inh_;
-
     // Expected number of poisson spikes.
     double lambda_;
-
     // Seed used for the Poisson spikes generation.
     int seed_;
+
+    // LIF cell prototype
+    arb::lif_cell lif = {
+        .source="src",
+        .target="tgt",
+        .tau_m = 10*U::ms,
+        .V_th = 10*U::mV,
+        .C_m = 20*U::pF,
+        .E_L = 0*U::mV,
+        .V_m = 0*U::mV,
+        .t_ref = 2*U::ms,
+    };
+    // Cable cell prototype
+    arb::cable_cell cable;
+    arb::cable_cell_global_properties prop;
+    bool use_cable_cells = true;
 };
 
 int main(int argc, char** argv) {
@@ -239,14 +270,18 @@ int main(int argc, char** argv) {
 
         unsigned seed = options.seed;
 
-        brunel_recipe recipe(nexc, ninh, next, in_degree_prop, w, d, rel_inh_strength, poiss_lambda, seed);
+        brunel_recipe recipe(nexc, ninh, next, in_degree_prop, w, d, rel_inh_strength, poiss_lambda, seed, options.use_cc);
 
         partition_hint_map hints;
         hints[cell_kind::lif].cpu_group_size = group_size;
+        hints[cell_kind::cable].cpu_group_size = group_size;
+        hints[cell_kind::lif].gpu_group_size = group_size;
+        hints[cell_kind::cable].gpu_group_size = group_size;
+        auto dec = partition_load_balance(recipe, context, hints);
 
         simulation sim(recipe,
                        context,
-                       [&hints](auto& r, auto c) { return partition_load_balance(r, c, hints); });
+                       dec);
 
         // Set up spike recording.
         std::vector<arb::spike> recorded_spikes;
@@ -267,13 +302,14 @@ int main(int argc, char** argv) {
         if (spike_out) {
             spike_out << std::fixed << std::setprecision(4);
             for (auto& s: recorded_spikes) {
-                spike_out << s.source.gid << ' ' << s.time << '\n';
+                spike_out << s.source.gid << ' '
+                          << s.time << '\n';
             }
         }
 
         // output profile and diagnostic feedback
-        std::cout << profile::profiler_summary() << "\n";
-        std::cout << "\nThere were " << sim.num_spikes() << " spikes\n";
+        std::cout << profile::profiler_summary() << "\n"
+                  << "\nThere were " << sim.num_spikes() << " spikes\n";
 
         auto report = profile::make_meter_report(meters, context);
         std::cout << report;
@@ -286,8 +322,8 @@ int main(int argc, char** argv) {
     }
     catch (std::exception& e) {
         // only print errors on master
-        std::cerr << sup::mask_stream(root);
-        std::cerr << e.what() << "\n";
+        std::cerr << sup::mask_stream(root)
+                  << e.what() << "\n";
         return 1;
     }
     return 0;
@@ -310,9 +346,9 @@ void add_subset(cell_gid_type gid,
                 float weight, float delay,
                 std::vector<cell_connection>& conns) {
     // We can only add this many connections!
-    auto gid_in_range = int(gid >= start && gid < end); //
+    auto gid_in_range = int(gid >= start && gid < end);
     if (m + start + gid_in_range >= end) throw std::runtime_error("Requested too many connections from the given range of gids.");
-
+    // Exclude ourself
     std::set<cell_gid_type> seen{gid};
     std::mt19937 gen(gid + 42);
     std::uniform_int_distribution<cell_gid_type> dis(start, end - 1);
@@ -345,7 +381,7 @@ std::optional<cl_options> read_options(int argc, char** argv) {
                      "-G|--group-size        [Number of cells per cell group]\n"
                      "-S|--seed              [Seed for poisson spike generators]\n"
                      "-f|--write-spikes      [Save spikes to file]\n"
-                     "-z|--profile-rank-zero [Only output profile information for rank 0]\n"
+                     "-c|--use-cable-cells   [Use a cable cell model]\n"
                      "-v|--verbose           [Print more verbose information to stdout]\n";
 
     cl_options opt;
@@ -354,22 +390,22 @@ std::optional<cl_options> read_options(int argc, char** argv) {
     };
 
     to::option options[] = {
-            { opt.nexc,              "-n", "--n-excitatory" },
-            { opt.ninh,              "-m", "--n-inhibitory" },
-            { opt.next,              "-e", "--n-external" },
-            { opt.syn_per_cell_prop, "-p", "--in-degree-prop" },
-            { opt.weight,            "-w", "--weight" },
-            { opt.delay,             "-d", "--delay" },
-            { opt.rel_inh_strength,  "-g", "--rel-inh-w" },
-            { opt.poiss_lambda,      "-l", "--lambda" },
-            { opt.tfinal,            "-t", "--tfinal" },
-            { opt.dt,                "-s", "--dt" },
-            { opt.group_size,        "-G", "--group-size" },
-            { opt.seed,              "-S", "--seed" },
-            { opt.spike_file_output, "-f", "--write-spikes" },
-            // { to::set(opt.profile_only_zero), to::flag, "-z", "--profile-rank-zero" },
-            { to::set(opt.verbose),           to::flag, "-v", "--verbose" },
-            { to::action(help),               to::flag, to::exit, "-h", "--help" }
+            { opt.nexc,                        "-n", "--n-excitatory" },
+            { opt.ninh,                        "-m", "--n-inhibitory" },
+            { opt.next,                        "-e", "--n-external" },
+            { opt.syn_per_cell_prop,           "-p", "--in-degree-prop" },
+            { opt.weight,                      "-w", "--weight" },
+            { opt.delay,                       "-d", "--delay" },
+            { opt.rel_inh_strength,            "-g", "--rel-inh-w" },
+            { opt.poiss_lambda,                "-l", "--lambda" },
+            { opt.tfinal,                      "-t", "--tfinal" },
+            { opt.dt,                          "-s", "--dt" },
+            { opt.group_size,                  "-G", "--group-size" },
+            { opt.seed,                        "-S", "--seed" },
+            { opt.spike_file_output,           "-f", "--write-spikes" },
+            { to::set(opt.use_cc),   to::flag, "-c", "--use-cable-cells" },
+            { to::set(opt.verbose),  to::flag, "-v", "--verbose" },
+            { to::action(help),      to::flag, to::exit, "-h", "--help" }
     };
 
     if (!to::run(options, argc, argv+1)) return {};
@@ -392,20 +428,20 @@ std::optional<cl_options> read_options(int argc, char** argv) {
 }
 
 std::ostream& operator<<(std::ostream& o, const cl_options& options) {
-    o << "Simulation options:\n";
-    o << "  Excitatory cells                                           : " << options.nexc << "\n";
-    o << "  Inhibitory cells                                           : " << options.ninh << "\n";
-    o << "  Poisson connections per cell                               : " << options.next << "\n";
-    o << "  Proportion of synapses/cell from each population           : " << options.syn_per_cell_prop << "\n";
-    o << "  Weight of excitatory synapses                              : " << options.weight << "\n";
-    o << "  Relative strength of inhibitory synapses                   : " << options.rel_inh_strength << "\n";
-    o << "  Delay of all synapses                                      : " << options.delay << "\n";
-    o << "  Expected number of spikes from a single poisson cell per ms: " << options.poiss_lambda << "\n";
-    o << "\n";
-    o << "  Simulation time                                            : " << options.tfinal << "\n";
-    o << "  dt                                                         : " << options.dt << "\n";
-    o << "  Group size                                                 : " << options.group_size << "\n";
-    o << "  Seed                                                       : " << options.seed << "\n";
-    o << "  Spike file output                                          : " << options.spike_file_output << "\n";
+    o << "Simulation options:\n"
+      << "  Excitatory cells                                           : " << options.nexc << "\n"
+      << "  Inhibitory cells                                           : " << options.ninh << "\n"
+      << "  Poisson connections per cell                               : " << options.next << "\n"
+      << "  Proportion of synapses/cell from each population           : " << options.syn_per_cell_prop << "\n"
+      << "  Weight of excitatory synapses                              : " << options.weight << "\n"
+      << "  Relative strength of inhibitory synapses                   : " << options.rel_inh_strength << "\n"
+      << "  Delay of all synapses                                      : " << options.delay << "\n"
+      << "  Expected number of spikes from a single poisson cell per ms: " << options.poiss_lambda << "\n"
+      << "\n"
+      << "  Simulation time                                            : " << options.tfinal << "\n"
+      << "  dt                                                         : " << options.dt << "\n"
+      << "  Group size                                                 : " << options.group_size << "\n"
+      << "  Seed                                                       : " << options.seed << "\n"
+      << "  Spike file output                                          : " << options.spike_file_output << "\n";
     return o;
 }

From 812b2d00ddd7eecfd00579603a5ef49ab77a1d40 Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Tue, 17 Dec 2024 15:50:17 +0100
Subject: [PATCH 02/13] make lif cell an aggregate.

---
 arbor/include/arbor/lif_cell.hpp | 8 ++------
 1 file changed, 2 insertions(+), 6 deletions(-)

diff --git a/arbor/include/arbor/lif_cell.hpp b/arbor/include/arbor/lif_cell.hpp
index e9d5c2fdf5..2af3c2bf7c 100644
--- a/arbor/include/arbor/lif_cell.hpp
+++ b/arbor/include/arbor/lif_cell.hpp
@@ -12,8 +12,8 @@ using namespace U::literals;
 
 // Model parameters of leaky integrate and fire neuron model.
 struct ARB_SYMBOL_VISIBLE lif_cell {
-    cell_tag_type source; // Label of source.
-    cell_tag_type target; // Label of target.
+    cell_tag_type source = ""; // Label of source.
+    cell_tag_type target = ""; // Label of target.
 
     // Neuronal parameters.
     U::quantity tau_m = 10_ms;    // Membrane potential decaying constant [ms].
@@ -23,10 +23,6 @@ struct ARB_SYMBOL_VISIBLE lif_cell {
     U::quantity E_R   =  0_mV;    // Reset potential [mV].
     U::quantity V_m   =  0_mV;    // Initial value of the Membrane potential [mV].
     U::quantity t_ref =  2_ms;    // Refractory period [ms].
-
-    lif_cell() = default;
-    lif_cell(cell_tag_type source, cell_tag_type target): source(std::move(source)), target(std::move(target)) {}
-
 };
 
 // LIF probe metadata, to be passed to sampler callbacks. Intentionally left blank.

From 37f5c29390ecd75704ee3542b9569301055d13b2 Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Tue, 17 Dec 2024 19:43:15 +0100
Subject: [PATCH 03/13] Remove introsort for classical sort.

---
 arbor/backends/event_stream_base.hpp | 89 +++++++++++++++++-----------
 1 file changed, 55 insertions(+), 34 deletions(-)

diff --git a/arbor/backends/event_stream_base.hpp b/arbor/backends/event_stream_base.hpp
index a81254a781..315a552483 100644
--- a/arbor/backends/event_stream_base.hpp
+++ b/arbor/backends/event_stream_base.hpp
@@ -8,7 +8,6 @@
 #include "backends/event_stream_state.hpp"
 #include "event_lane.hpp"
 #include "timestep_range.hpp"
-#include "util/partition.hpp"
 
 ARB_SERDES_ENABLE_EXT(arb_deliverable_event_data, mech_index, weight);
 
@@ -63,7 +62,28 @@ struct event_stream_base {
     virtual void init() = 0;
 };
 
-struct spike_event_stream_base : event_stream_base<deliverable_event> {
+struct spike_event_stream_base: event_stream_base<deliverable_event> {
+    // Take in one event lane per cell `gid` and reorganise into one stream per
+    // synapse `mech_id`.
+    //
+    // - Due to the cell group coalescing multiple cells and their synapses into
+    //   one object, one `mech_id` can touch multiple lanes / `gid`s.
+    // - Inversely, two `mech_id`s can cover different, but overlapping sets of `gid`s
+    // - Multiple `mech_id`s can receive events from the same source
+    //
+    // Pre:
+    // - Events in `lanes[ix]` forall ix
+    //   * are sorted by time
+    //   * `ix` maps to exactly one cell in the local cell group
+    // - `divs` partitions `handles` such that the target handles for cell `ix`
+    //   are located in `handles[divs[ix]..divs[ix + 1]]`
+    // - `handles` records `(mech_id, index)` of a target s.t. `index` is the instance
+    //   with the set identified by `mech_id`, e.g. a single synapse placed on a multi-
+    //   location locset (plus the merging across cells by groups)
+    // Post:
+    // - streams[mech_id] contains a list of all events for synapse `mech_id` s.t.
+    //   * the list is sorted by (time_step, lid, time)
+    //   * the list is partitioned by `time_step` via `ev_spans`
     template<typename EventStream>
     friend void initialize(const event_lane_subrange& lanes,
                            const std::vector<target_handle>& handles,
@@ -74,54 +94,56 @@ struct spike_event_stream_base : event_stream_base<deliverable_event> {
 
         // reset streams and allocate sufficient space for temporaries
         auto n_steps = steps.size();
-        for (auto& [k, v]: streams) {
-            v.clear();
-            v.spike_counter_.clear();
-            v.spike_counter_.resize(steps.size(), 0);
-            v.spikes_.clear();
+        for (auto& [id, stream]: streams) {
+            stream.clear();
+            stream.ev_spans_.resize(steps.size() + 1, 0);
+            stream.spikes_.clear();
             // ev_data_ has been cleared during v.clear(), so we use its capacity
-            v.spikes_.reserve(v.ev_data_.capacity());
+            stream.spikes_.reserve(stream.ev_data_.capacity());
         }
 
         // loop over lanes: group events by mechanism and sort them by time
         auto cell = 0;
         for (const auto& lane: lanes) {
             auto div = divs[cell];
-            ++cell;
             arb_size_type step = 0;
             for (const auto& evt: lane) {
-                auto time = evt.time;
-                auto weight = evt.weight;
-                auto target = evt.target;
-                while(step < n_steps && time >= steps[step].t_end()) ++step;
+                step = std::lower_bound(steps.begin() + step,
+                                        steps.end(),
+                                        evt.time,
+                                        [](const auto& bucket, time_type time) { return bucket.t_end() <= time; })
+                     - steps.begin();
                 // Events coinciding with epoch's upper boundary belong to next epoch
                 if (step >= n_steps) break;
-                arb_assert(div + target < handles.size());
-                const auto& handle = handles[div + target];
+                arb_assert(div + evt.target < handles.size());
+                const auto& handle = handles[div + evt.target];
                 auto& stream = streams[handle.mech_id];
-                stream.spikes_.push_back(spike_data{step, handle.mech_index, time, weight});
-                // insertion sort with last element as pivot
-                // ordering: first w.r.t. step, within a step: mech_index, within a mech_index: time
-                auto first = stream.spikes_.begin();
-                auto last = stream.spikes_.end();
-                auto pivot = std::prev(last, 1);
-                std::rotate(std::upper_bound(first, pivot, *pivot), pivot, last);
-                // increment count in current time interval
-                stream.spike_counter_[step]++;
+                stream.spikes_.emplace_back(step, handle.mech_index, evt.time, evt.weight);
+                stream.ev_spans_[step + 1]++;
             }
+            ++cell;
         }
 
+        // TODO parallelise over streams
+        // auto tg = threading::task_group(ts.get());
         for (auto& [id, stream]: streams) {
-            // copy temporary deliverable_events into stream's ev_data_
-            stream.ev_data_.reserve(stream.spikes_.size());
-            std::transform(stream.spikes_.begin(), stream.spikes_.end(), std::back_inserter(stream.ev_data_),
-                [](auto const& e) noexcept -> arb_deliverable_event_data {
-                return {e.mech_index, e.weight}; });
-            // scan over spike_counter_ and  written to ev_spans_
-            util::make_partition(stream.ev_spans_, stream.spike_counter_);
-            // delegate to derived class init: static cast necessary to access protected init()
-            static_cast<spike_event_stream_base&>(stream).init();
+            // tg.run([&stream=stream]() {
+                // scan to partition stream
+                std::inclusive_scan(stream.ev_spans_.begin(), stream.ev_spans_.end(),
+                                    stream.ev_spans_.begin());
+                // This is made slightly faster by using pdqsort.
+                // Despite events being sorted by time in the partitions defined
+                // by the lane index here, they are not _totally_ sorted, thus
+                // sort is needed, merge not being strong enough :/
+                std::sort(stream.spikes_.begin(), stream.spikes_.end());
+                // copy temporary deliverable_events into stream's ev_data_
+                stream.ev_data_.reserve(stream.spikes_.size());
+                for (const auto& spike: stream.spikes_) stream.ev_data_.emplace_back(spike.mech_index, spike.weight);
+                // delegate to derived class init: static cast necessary to access protected init()
+                static_cast<spike_event_stream_base&>(stream).init();
+            // });
         }
+        // tg.wait();
     }
 
   protected: // members
@@ -133,7 +155,6 @@ struct spike_event_stream_base : event_stream_base<deliverable_event> {
         auto operator<=>(spike_data const&) const noexcept = default;
     };
     std::vector<spike_data> spikes_;
-    std::vector<std::size_t> spike_counter_;
 };
 
 struct sample_event_stream_base : event_stream_base<sample_event> {

From 5a82e141bb89e2086054934e07d54dd06f2ad3df Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Wed, 18 Dec 2024 10:18:58 +0100
Subject: [PATCH 04/13] Disable --- for now --- extremely slow tests.

---
 scripts/run_cpp_examples.sh | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/scripts/run_cpp_examples.sh b/scripts/run_cpp_examples.sh
index b916c23594..184497eaed 100755
--- a/scripts/run_cpp_examples.sh
+++ b/scripts/run_cpp_examples.sh
@@ -25,7 +25,15 @@ ok=0
 # List of all examples
 all_examples=(
     "bench"
-    "brunel"
+    "brunel -G 1"
+    "brunel -G 10"
+    "brunel -G 100"
+    "brunel -G 100000"
+    "brunel --n-excitatory 10000 --n-inhibitory 2500 --tfinal 10 --in-degree-prop 0.1 --dt 0.01 --lambda 1.5 -g 0.8 --delay 1.0 --weight 0.00010 -G 1 --use-cable-cells"
+    "brunel --n-excitatory 10000 --n-inhibitory 2500 --tfinal 10 --in-degree-prop 0.1 --dt 0.01 --lambda 1.5 -g 0.8 --delay 1.0 --weight 0.00010 -G 10 --use-cable-cells"
+    # TODO These need to be enabled when the group size issue is fixed.
+    # "brunel --n-excitatory 10000 --n-inhibitory 2500 --tfinal 10 --in-degree-prop 0.1 --dt 0.01 --lambda 1.5 -g 0.8 --delay 1.0 --weight 0.00010 -G 100 --use-cable-cells"
+    # "brunel --n-excitatory 10000 --n-inhibitory 2500 --tfinal 10 --in-degree-prop 0.1 --dt 0.01 --lambda 1.5 -g 0.8 --delay 1.0 --weight 0.00010 -G 100000 --use-cable-cells"
     "gap_junctions"
     "generators"
     "lfp"
@@ -49,6 +57,13 @@ skip_local=(
 expected_outputs=(
     972
     6998
+    6998
+    6998
+    6998
+    38956
+    38956
+    # 38956
+    # 38956
     "30"
     ""
     ""

From fa7127af50c15d67dbc05215a07d70d24470c21a Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Tue, 17 Dec 2024 19:43:15 +0100
Subject: [PATCH 05/13] Remove introsort for classical sort.

---
 arbor/backends/event_stream_base.hpp | 89 +++++++++++++++++-----------
 1 file changed, 55 insertions(+), 34 deletions(-)

diff --git a/arbor/backends/event_stream_base.hpp b/arbor/backends/event_stream_base.hpp
index a81254a781..315a552483 100644
--- a/arbor/backends/event_stream_base.hpp
+++ b/arbor/backends/event_stream_base.hpp
@@ -8,7 +8,6 @@
 #include "backends/event_stream_state.hpp"
 #include "event_lane.hpp"
 #include "timestep_range.hpp"
-#include "util/partition.hpp"
 
 ARB_SERDES_ENABLE_EXT(arb_deliverable_event_data, mech_index, weight);
 
@@ -63,7 +62,28 @@ struct event_stream_base {
     virtual void init() = 0;
 };
 
-struct spike_event_stream_base : event_stream_base<deliverable_event> {
+struct spike_event_stream_base: event_stream_base<deliverable_event> {
+    // Take in one event lane per cell `gid` and reorganise into one stream per
+    // synapse `mech_id`.
+    //
+    // - Due to the cell group coalescing multiple cells and their synapses into
+    //   one object, one `mech_id` can touch multiple lanes / `gid`s.
+    // - Inversely, two `mech_id`s can cover different, but overlapping sets of `gid`s
+    // - Multiple `mech_id`s can receive events from the same source
+    //
+    // Pre:
+    // - Events in `lanes[ix]` forall ix
+    //   * are sorted by time
+    //   * `ix` maps to exactly one cell in the local cell group
+    // - `divs` partitions `handles` such that the target handles for cell `ix`
+    //   are located in `handles[divs[ix]..divs[ix + 1]]`
+    // - `handles` records `(mech_id, index)` of a target s.t. `index` is the instance
+    //   with the set identified by `mech_id`, e.g. a single synapse placed on a multi-
+    //   location locset (plus the merging across cells by groups)
+    // Post:
+    // - streams[mech_id] contains a list of all events for synapse `mech_id` s.t.
+    //   * the list is sorted by (time_step, lid, time)
+    //   * the list is partitioned by `time_step` via `ev_spans`
     template<typename EventStream>
     friend void initialize(const event_lane_subrange& lanes,
                            const std::vector<target_handle>& handles,
@@ -74,54 +94,56 @@ struct spike_event_stream_base : event_stream_base<deliverable_event> {
 
         // reset streams and allocate sufficient space for temporaries
         auto n_steps = steps.size();
-        for (auto& [k, v]: streams) {
-            v.clear();
-            v.spike_counter_.clear();
-            v.spike_counter_.resize(steps.size(), 0);
-            v.spikes_.clear();
+        for (auto& [id, stream]: streams) {
+            stream.clear();
+            stream.ev_spans_.resize(steps.size() + 1, 0);
+            stream.spikes_.clear();
             // ev_data_ has been cleared during v.clear(), so we use its capacity
-            v.spikes_.reserve(v.ev_data_.capacity());
+            stream.spikes_.reserve(stream.ev_data_.capacity());
         }
 
         // loop over lanes: group events by mechanism and sort them by time
         auto cell = 0;
         for (const auto& lane: lanes) {
             auto div = divs[cell];
-            ++cell;
             arb_size_type step = 0;
             for (const auto& evt: lane) {
-                auto time = evt.time;
-                auto weight = evt.weight;
-                auto target = evt.target;
-                while(step < n_steps && time >= steps[step].t_end()) ++step;
+                step = std::lower_bound(steps.begin() + step,
+                                        steps.end(),
+                                        evt.time,
+                                        [](const auto& bucket, time_type time) { return bucket.t_end() <= time; })
+                     - steps.begin();
                 // Events coinciding with epoch's upper boundary belong to next epoch
                 if (step >= n_steps) break;
-                arb_assert(div + target < handles.size());
-                const auto& handle = handles[div + target];
+                arb_assert(div + evt.target < handles.size());
+                const auto& handle = handles[div + evt.target];
                 auto& stream = streams[handle.mech_id];
-                stream.spikes_.push_back(spike_data{step, handle.mech_index, time, weight});
-                // insertion sort with last element as pivot
-                // ordering: first w.r.t. step, within a step: mech_index, within a mech_index: time
-                auto first = stream.spikes_.begin();
-                auto last = stream.spikes_.end();
-                auto pivot = std::prev(last, 1);
-                std::rotate(std::upper_bound(first, pivot, *pivot), pivot, last);
-                // increment count in current time interval
-                stream.spike_counter_[step]++;
+                stream.spikes_.emplace_back(step, handle.mech_index, evt.time, evt.weight);
+                stream.ev_spans_[step + 1]++;
             }
+            ++cell;
         }
 
+        // TODO parallelise over streams
+        // auto tg = threading::task_group(ts.get());
         for (auto& [id, stream]: streams) {
-            // copy temporary deliverable_events into stream's ev_data_
-            stream.ev_data_.reserve(stream.spikes_.size());
-            std::transform(stream.spikes_.begin(), stream.spikes_.end(), std::back_inserter(stream.ev_data_),
-                [](auto const& e) noexcept -> arb_deliverable_event_data {
-                return {e.mech_index, e.weight}; });
-            // scan over spike_counter_ and  written to ev_spans_
-            util::make_partition(stream.ev_spans_, stream.spike_counter_);
-            // delegate to derived class init: static cast necessary to access protected init()
-            static_cast<spike_event_stream_base&>(stream).init();
+            // tg.run([&stream=stream]() {
+                // scan to partition stream
+                std::inclusive_scan(stream.ev_spans_.begin(), stream.ev_spans_.end(),
+                                    stream.ev_spans_.begin());
+                // This is made slightly faster by using pdqsort.
+                // Despite events being sorted by time in the partitions defined
+                // by the lane index here, they are not _totally_ sorted, thus
+                // sort is needed, merge not being strong enough :/
+                std::sort(stream.spikes_.begin(), stream.spikes_.end());
+                // copy temporary deliverable_events into stream's ev_data_
+                stream.ev_data_.reserve(stream.spikes_.size());
+                for (const auto& spike: stream.spikes_) stream.ev_data_.emplace_back(spike.mech_index, spike.weight);
+                // delegate to derived class init: static cast necessary to access protected init()
+                static_cast<spike_event_stream_base&>(stream).init();
+            // });
         }
+        // tg.wait();
     }
 
   protected: // members
@@ -133,7 +155,6 @@ struct spike_event_stream_base : event_stream_base<deliverable_event> {
         auto operator<=>(spike_data const&) const noexcept = default;
     };
     std::vector<spike_data> spikes_;
-    std::vector<std::size_t> spike_counter_;
 };
 
 struct sample_event_stream_base : event_stream_base<sample_event> {

From e1f4fcd8df8e54011937fa971d13f004b1a8ac80 Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Wed, 18 Dec 2024 08:25:36 +0100
Subject: [PATCH 06/13] Comment on aliasing

---
 arbor/backends/event_stream_base.hpp | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/arbor/backends/event_stream_base.hpp b/arbor/backends/event_stream_base.hpp
index 315a552483..15a00e519d 100644
--- a/arbor/backends/event_stream_base.hpp
+++ b/arbor/backends/event_stream_base.hpp
@@ -108,8 +108,7 @@ struct spike_event_stream_base: event_stream_base<deliverable_event> {
             auto div = divs[cell];
             arb_size_type step = 0;
             for (const auto& evt: lane) {
-                step = std::lower_bound(steps.begin() + step,
-                                        steps.end(),
+                step = std::lower_bound(steps.begin() + step, steps.end(),
                                         evt.time,
                                         [](const auto& bucket, time_type time) { return bucket.t_end() <= time; })
                      - steps.begin();
@@ -128,7 +127,7 @@ struct spike_event_stream_base: event_stream_base<deliverable_event> {
         // auto tg = threading::task_group(ts.get());
         for (auto& [id, stream]: streams) {
             // tg.run([&stream=stream]() {
-                // scan to partition stream
+                // scan to partition stream (aliasing is explicitly allowed)
                 std::inclusive_scan(stream.ev_spans_.begin(), stream.ev_spans_.end(),
                                     stream.ev_spans_.begin());
                 // This is made slightly faster by using pdqsort.

From 8eb35cf7a938f92f460065422b86d17d276f2848 Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Mon, 6 Jan 2025 14:53:18 +0100
Subject: [PATCH 07/13] Clean-up

---
 arbor/backends/event_stream_base.hpp | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/arbor/backends/event_stream_base.hpp b/arbor/backends/event_stream_base.hpp
index a25777f248..d7a0d4b872 100644
--- a/arbor/backends/event_stream_base.hpp
+++ b/arbor/backends/event_stream_base.hpp
@@ -108,10 +108,6 @@ struct spike_event_stream_base: event_stream_base<deliverable_event> {
             auto div = divs[cell];
             arb_size_type step = 0;
             for (const auto& evt: lane) {
-                auto time = evt.time;
-                auto weight = evt.weight;
-                auto target = evt.target;
-                while(step < n_steps && time >= steps[step].t_end()) ++step;
                 step = std::lower_bound(steps.begin() + step,
                                         steps.end(),
                                         evt.time,

From 00f8dfc63f60ddd7d82e79ac6792ac8952f6f864 Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Mon, 6 Jan 2025 15:48:38 +0100
Subject: [PATCH 08/13] Test?

---
 arbor/backends/event_stream_base.hpp | 1 +
 1 file changed, 1 insertion(+)

diff --git a/arbor/backends/event_stream_base.hpp b/arbor/backends/event_stream_base.hpp
index d7a0d4b872..f518d6b233 100644
--- a/arbor/backends/event_stream_base.hpp
+++ b/arbor/backends/event_stream_base.hpp
@@ -8,6 +8,7 @@
 #include "backends/event_stream_state.hpp"
 #include "event_lane.hpp"
 #include "timestep_range.hpp"
+#include "util/partition.hpp"
 
 ARB_SERDES_ENABLE_EXT(arb_deliverable_event_data, mech_index, weight);
 

From 2cd0f6293133304c686ea36cc1b8789d6ba378b0 Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Mon, 6 Jan 2025 17:10:39 +0100
Subject: [PATCH 09/13] Tweak ctors

---
 arbor/backends/event.hpp | 12 +++++++-----
 1 file changed, 7 insertions(+), 5 deletions(-)

diff --git a/arbor/backends/event.hpp b/arbor/backends/event.hpp
index 7fd9eb8491..e062dbccb2 100644
--- a/arbor/backends/event.hpp
+++ b/arbor/backends/event.hpp
@@ -39,17 +39,19 @@ struct deliverable_event {
     target_handle handle;
 
     deliverable_event() = default;
-    constexpr deliverable_event(time_type time, target_handle handle, float weight) noexcept:
-        time(time), weight(weight), handle(handle) {}
+    constexpr deliverable_event(const time_type time,
+                                target_handle handle,
+                                const float weight) noexcept:
+        time(time), weight(weight), handle(std::move(handle)) {}
 
     ARB_SERDES_ENABLE(deliverable_event, time, weight, handle);
 };
 
 // Subset of event information required for mechanism delivery.
 struct deliverable_event_data {
-    cell_local_size_type mech_index; // same as target_handle::mech_index
-    float weight;
-    deliverable_event_data(cell_local_size_type idx, float w):
+    cell_local_size_type mech_index = 0; // same as target_handle::mech_index
+    float weight = 0;
+    deliverable_event_data(const cell_local_size_type idx, const float w) noexcept:
         mech_index(idx),
         weight(w) {}
     ARB_SERDES_ENABLE(deliverable_event_data,

From e546d5d313431c9f04305d7ba4e4a76dc6b8e87f Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Tue, 7 Jan 2025 08:29:50 +0100
Subject: [PATCH 10/13] C++20 support do be weird?

---
 arbor/backends/event.hpp             | 6 +++---
 arbor/backends/event_stream_base.hpp | 4 ++--
 2 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/arbor/backends/event.hpp b/arbor/backends/event.hpp
index e062dbccb2..e7abdc6d45 100644
--- a/arbor/backends/event.hpp
+++ b/arbor/backends/event.hpp
@@ -39,9 +39,9 @@ struct deliverable_event {
     target_handle handle;
 
     deliverable_event() = default;
-    constexpr deliverable_event(const time_type time,
+    constexpr deliverable_event(time_type time,
                                 target_handle handle,
-                                const float weight) noexcept:
+                                float weight) noexcept:
         time(time), weight(weight), handle(std::move(handle)) {}
 
     ARB_SERDES_ENABLE(deliverable_event, time, weight, handle);
@@ -51,7 +51,7 @@ struct deliverable_event {
 struct deliverable_event_data {
     cell_local_size_type mech_index = 0; // same as target_handle::mech_index
     float weight = 0;
-    deliverable_event_data(const cell_local_size_type idx, const float w) noexcept:
+    deliverable_event_data(cell_local_size_type idx, float w) noexcept:
         mech_index(idx),
         weight(w) {}
     ARB_SERDES_ENABLE(deliverable_event_data,
diff --git a/arbor/backends/event_stream_base.hpp b/arbor/backends/event_stream_base.hpp
index f518d6b233..0b3fb940a6 100644
--- a/arbor/backends/event_stream_base.hpp
+++ b/arbor/backends/event_stream_base.hpp
@@ -104,7 +104,7 @@ struct spike_event_stream_base: event_stream_base<deliverable_event> {
         }
 
         // loop over lanes: group events by mechanism and sort them by time
-        auto cell = 0;
+        arb_size_type cell = 0;
         for (const auto& lane: lanes) {
             auto div = divs[cell];
             arb_size_type step = 0;
@@ -119,7 +119,7 @@ struct spike_event_stream_base: event_stream_base<deliverable_event> {
                 arb_assert(div + evt.target < handles.size());
                 const auto& handle = handles[div + evt.target];
                 auto& stream = streams[handle.mech_id];
-                stream.spikes_.emplace_back(step, handle.mech_index, evt.time, evt.weight);
+                stream.spikes_.push_back(spike_data(step, handle.mech_index, evt.time, evt.weight));
                 stream.ev_spans_[step + 1]++;
             }
             ++cell;

From 78d8cc34a6b74f45e801cb9de48344fbb8bc21c5 Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Tue, 7 Jan 2025 09:06:09 +0100
Subject: [PATCH 11/13] appease older compilers, i guess?

---
 arbor/backends/event_stream_base.hpp | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/arbor/backends/event_stream_base.hpp b/arbor/backends/event_stream_base.hpp
index 0b3fb940a6..0e62d6ddf4 100644
--- a/arbor/backends/event_stream_base.hpp
+++ b/arbor/backends/event_stream_base.hpp
@@ -119,7 +119,7 @@ struct spike_event_stream_base: event_stream_base<deliverable_event> {
                 arb_assert(div + evt.target < handles.size());
                 const auto& handle = handles[div + evt.target];
                 auto& stream = streams[handle.mech_id];
-                stream.spikes_.push_back(spike_data(step, handle.mech_index, evt.time, evt.weight));
+                stream.spikes_.emplace_back({step, handle.mech_index, evt.time, evt.weight});
                 stream.ev_spans_[step + 1]++;
             }
             ++cell;
@@ -139,7 +139,7 @@ struct spike_event_stream_base: event_stream_base<deliverable_event> {
                 std::sort(stream.spikes_.begin(), stream.spikes_.end());
                 // copy temporary deliverable_events into stream's ev_data_
                 stream.ev_data_.reserve(stream.spikes_.size());
-                for (const auto& spike: stream.spikes_) stream.ev_data_.emplace_back(spike.mech_index, spike.weight);
+                for (const auto& spike: stream.spikes_) stream.ev_data_.emplace_back({spike.mech_index, spike.weight});
                 // delegate to derived class init: static cast necessary to access protected init()
                 static_cast<spike_event_stream_base&>(stream).init();
             // });

From d2008263a0f8e0510e66fb30f51ce3708e490fde Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Tue, 7 Jan 2025 10:20:52 +0100
Subject: [PATCH 12/13] Name the types

---
 arbor/backends/event_stream_base.hpp | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/arbor/backends/event_stream_base.hpp b/arbor/backends/event_stream_base.hpp
index 0e62d6ddf4..9c9122f778 100644
--- a/arbor/backends/event_stream_base.hpp
+++ b/arbor/backends/event_stream_base.hpp
@@ -8,7 +8,6 @@
 #include "backends/event_stream_state.hpp"
 #include "event_lane.hpp"
 #include "timestep_range.hpp"
-#include "util/partition.hpp"
 
 ARB_SERDES_ENABLE_EXT(arb_deliverable_event_data, mech_index, weight);
 
@@ -119,7 +118,7 @@ struct spike_event_stream_base: event_stream_base<deliverable_event> {
                 arb_assert(div + evt.target < handles.size());
                 const auto& handle = handles[div + evt.target];
                 auto& stream = streams[handle.mech_id];
-                stream.spikes_.emplace_back({step, handle.mech_index, evt.time, evt.weight});
+                stream.spikes_.emplace_back(spike_data{step, handle.mech_index, evt.time, evt.weight});
                 stream.ev_spans_[step + 1]++;
             }
             ++cell;
@@ -139,7 +138,7 @@ struct spike_event_stream_base: event_stream_base<deliverable_event> {
                 std::sort(stream.spikes_.begin(), stream.spikes_.end());
                 // copy temporary deliverable_events into stream's ev_data_
                 stream.ev_data_.reserve(stream.spikes_.size());
-                for (const auto& spike: stream.spikes_) stream.ev_data_.emplace_back({spike.mech_index, spike.weight});
+                for (const auto& spike: stream.spikes_) stream.ev_data_.emplace_back(event_data_type{spike.mech_index, spike.weight});
                 // delegate to derived class init: static cast necessary to access protected init()
                 static_cast<spike_event_stream_base&>(stream).init();
             // });

From a029a33616efad1e0496b1c80a3335ffbde69130 Mon Sep 17 00:00:00 2001
From: Thorsten Hater <24411438+thorstenhater@users.noreply.github.com>
Date: Tue, 7 Jan 2025 11:30:05 +0100
Subject: [PATCH 13/13] Enable once too-slow tests.

---
 scripts/run_cpp_examples.sh | 9 ++++-----
 1 file changed, 4 insertions(+), 5 deletions(-)

diff --git a/scripts/run_cpp_examples.sh b/scripts/run_cpp_examples.sh
index 184497eaed..7ae744b02b 100755
--- a/scripts/run_cpp_examples.sh
+++ b/scripts/run_cpp_examples.sh
@@ -31,9 +31,8 @@ all_examples=(
     "brunel -G 100000"
     "brunel --n-excitatory 10000 --n-inhibitory 2500 --tfinal 10 --in-degree-prop 0.1 --dt 0.01 --lambda 1.5 -g 0.8 --delay 1.0 --weight 0.00010 -G 1 --use-cable-cells"
     "brunel --n-excitatory 10000 --n-inhibitory 2500 --tfinal 10 --in-degree-prop 0.1 --dt 0.01 --lambda 1.5 -g 0.8 --delay 1.0 --weight 0.00010 -G 10 --use-cable-cells"
-    # TODO These need to be enabled when the group size issue is fixed.
-    # "brunel --n-excitatory 10000 --n-inhibitory 2500 --tfinal 10 --in-degree-prop 0.1 --dt 0.01 --lambda 1.5 -g 0.8 --delay 1.0 --weight 0.00010 -G 100 --use-cable-cells"
-    # "brunel --n-excitatory 10000 --n-inhibitory 2500 --tfinal 10 --in-degree-prop 0.1 --dt 0.01 --lambda 1.5 -g 0.8 --delay 1.0 --weight 0.00010 -G 100000 --use-cable-cells"
+    "brunel --n-excitatory 10000 --n-inhibitory 2500 --tfinal 10 --in-degree-prop 0.1 --dt 0.01 --lambda 1.5 -g 0.8 --delay 1.0 --weight 0.00010 -G 100 --use-cable-cells"
+    "brunel --n-excitatory 10000 --n-inhibitory 2500 --tfinal 10 --in-degree-prop 0.1 --dt 0.01 --lambda 1.5 -g 0.8 --delay 1.0 --weight 0.00010 -G 100000 --use-cable-cells"
     "gap_junctions"
     "generators"
     "lfp"
@@ -62,8 +61,8 @@ expected_outputs=(
     6998
     38956
     38956
-    # 38956
-    # 38956
+    38956
+    38956
     "30"
     ""
     ""