Class-based custom mechanisms for arbor in python and C++ (DONT TRY THIS AT HOME!)

main.py

import sys
sys.path.insert(0, '/home/llandsmeer/repos/llandsmeer/arbor/build-debug/prefix/lib/python3.10/site-packages')
sys.path.append('/home/llandsmeer/repos/llandsmeer/nmlcc-cat')

import arbor
import matplotlib.pyplot as plt

import nmlcat

cat = nmlcat.Catalogue.make()
cat.add('one')
cat.add('two')

tree = arbor.segment_tree()
tree.append(arbor.mnpos, arbor.mpoint(0, 0, 0, 1), arbor.mpoint(1, 0, 0, 1), tag=1)
labels = arbor.label_dict()
decor = arbor.decor() \
        .paint('(all)', arbor.density('one')) \
        .paint('(all)', arbor.density('two'))
cell = arbor.cable_cell(tree, decor, labels)
m = arbor.single_cell_model(cell)
m.properties.catalogue.extend(cat.load(), '')
m.probe('voltage', '(root)', frequency=10)
m.run(tfinal=1, dt=0.5)
print(m.traces[0].value)
plt.plot(m.traces[0].time, m.traces[0].value)
plt.show()

nmlcat.cpp

#include <iostream>

#include <dlfcn.h>

#include <memory>
#include <arbor/mechcat.hpp>
#include <arbor/mechanism.hpp>

#ifdef PYTHON
#include <pybind11/embed.h>
#include <pybind11/pybind11.h>
#endif

class Mechanism;

class Catalogue {
    friend class Mechanism;
    std::vector<std::shared_ptr<Mechanism>> mechs; /* shared doesn't make sense here but whatever */
    std::vector<arb_mechanism> mechs_for_arbor;
    Catalogue(int cat_id_);
    int cat_id;
    inline static std::vector<std::shared_ptr<Catalogue>> catalogues; /* we can never delete! */
    inline static int _counter_current_mech_based_on_counter_during_init = 0;
    static std::shared_ptr<Mechanism> _get_current_mech_based_on_counter_during_init() {
        int idx = _counter_current_mech_based_on_counter_during_init / 2;
        _counter_current_mech_based_on_counter_during_init += 1;
        return current->mechs[idx];
    }
    inline static Catalogue * current = 0;
public:
    void set_as_current_for_init() {
        current = this;
        _counter_current_mech_based_on_counter_during_init = 0;
    }
    static Catalogue * get_current_during_init() {
        return current;
    }
    std::shared_ptr<Mechanism> add(std::string name);
    static std::shared_ptr<Catalogue> make() {
        int cat_id = catalogues.size();
        catalogues.emplace_back(std::shared_ptr<Catalogue>(new Catalogue(cat_id)));
        return catalogues[catalogues.size()-1];
    }
    void * get_catalogue(int * n);
#ifdef PYTHON
    pybind11::object load();
#else
    arb::mechanism_catalogue load();
#endif

    Catalogue (const Catalogue&) = delete;
    Catalogue& operator= (const Catalogue&) = delete;
};

class Mechanism {
    friend class Catalogue;
    arb_mechanism_interface cpu;
    std::string name = "my_bytecode_mechanism";
    std::vector<arb_field_info> globals;
    std::vector<arb_field_info> state_vars;
    std::vector<arb_field_info> parameters;
    std::vector<arb_ion_info> ions;

    static Mechanism * current(arb_mechanism_ppack * pp) {
        /* we put a global in pp with a index to our mechanism */
        int cat_id = (int)pp->globals[0];
        int mech_id = (int)pp->globals[1];
        return Catalogue::catalogues.at(cat_id)->mechs.at(mech_id).get();
    }

    /* during initilization */
    static arb_mechanism_type global_make_type_callback() {
        return Catalogue::_get_current_mech_based_on_counter_during_init()->type();
    }
    static arb_mechanism_interface * global_make_cpu_callback() {
        return &Catalogue::_get_current_mech_based_on_counter_during_init()->cpu;
    }
    static arb_mechanism_interface * global_make_gpu_callback() { return 0; }

    /* during integration: */
    static void global_init_callback(arb_mechanism_ppack * pp) { current(pp)->init(pp); }
    static void global_compute_currents_callback(arb_mechanism_ppack * pp) { current(pp)->compute_currents(pp); }
    static void global_advance_state_callback(arb_mechanism_ppack * pp) { current(pp)->advance_state(pp); }
    static void global_write_ions_callback(arb_mechanism_ppack * pp) { current(pp)->write_ions(pp); }
    static void global_apply_events_callback(arb_mechanism_ppack * pp, arb_deliverable_event_stream * e) { current(pp)->apply_events(pp, e); }
    static void global_post_event_callback(arb_mechanism_ppack * pp) { current(pp)->post_event(pp); }

    arb_mechanism_type type() {
        arb_mechanism_type type;
        type.abi_version = ARB_MECH_ABI_VERSION;
        type.fingerprint = "<placeholder>";
        type.name = name.c_str();
        type.kind = arb_mechanism_kind_density;
        type.is_linear = false;
        type.has_post_events = false;
        type.globals = globals.data();
        type.n_globals = globals.size();
        type.state_vars = state_vars.data();
        type.n_state_vars = state_vars.size();
        type.parameters = parameters.data();
        type.n_parameters = parameters.size();
        type.ions = ions.data();
        type.n_ions = ions.size();
        type.n_random_variables = 0;
        return type;
    }
    Mechanism(int cat_id, int mech_id, std::string _name = "") : name(_name) {
        cpu.partition_width = 1;
        cpu.backend = arb_backend_kind_cpu;
        cpu.alignment = 8;
        cpu.init_mechanism = global_init_callback;
        cpu.compute_currents = global_compute_currents_callback;
        cpu.apply_events = global_apply_events_callback;
        cpu.advance_state = global_advance_state_callback;
        cpu.write_ions = global_write_ions_callback;
        cpu.post_event = global_post_event_callback;
        /* do ugly thing with globals */
        arb_value_type arb_cat_id = cat_id;
        arb_value_type arb_mech_id = mech_id;
        globals.emplace_back("cat_id", "", arb_cat_id, arb_cat_id, arb_cat_id);
        globals.emplace_back("mech_id", "", arb_mech_id, arb_mech_id, arb_mech_id);
    }

    arb_mechanism make() {
        arb_mechanism result;
        result.type = global_make_type_callback;
        result.i_cpu = global_make_cpu_callback;
        result.i_gpu = global_make_gpu_callback;
        return result;
    }
public:
    void init(arb_mechanism_ppack * pp) {
        std::cout << name << "::init()" << std::endl;
        (void)pp;
    }

    void compute_currents(arb_mechanism_ppack * pp) {
        std::cout << name << "::compute_currents()" << std::endl;
        (void)pp;
    }

    void advance_state(arb_mechanism_ppack * pp) {
        std::cout << name << "::advance_state()" << std::endl;
        (void)pp;
    }

    void write_ions(arb_mechanism_ppack * pp) {
        std::cout << name << "::write_ions()" << std::endl;
        (void)pp;
    }

    void apply_events(arb_mechanism_ppack * pp, arb_deliverable_event_stream * e) {
        std::cout << name << "::apply_events()" << std::endl;
        (void)pp;
        (void)e;
    }

    void post_event(arb_mechanism_ppack * pp) {
        std::cout << name << "::post_event()" << std::endl;
        (void)pp;
    }

    Mechanism (const Mechanism&) = delete;
    Mechanism& operator= (const Mechanism&) = delete;
};

Catalogue::Catalogue(int cat_id_) : cat_id(cat_id_) {
}

std::shared_ptr<Mechanism> Catalogue::add(std::string name) {
    int mech_id = mechs.size();
    auto m = std::shared_ptr<Mechanism>(new Mechanism(cat_id, mech_id, name));
    mechs.push_back(m);
    return m;
}

void * Catalogue::get_catalogue(int * n) {
    if (mechs_for_arbor.size() == 0) {
        for (auto & m : mechs) {
            mechs_for_arbor.push_back(m->make());
        }
    }
    *n = mechs_for_arbor.size();
    return mechs_for_arbor.data();
}

#ifndef PYTHON
arb::mechanism_catalogue Catalogue::load() {
    arb::mechanism_catalogue result;
    for(auto & mech : mechs) {
        auto type = mech->type();
        result.add(mech->name, type);
        result.register_implementation(mech->name, std::make_unique<arb::mechanism>(type, mech->cpu));
    }
    return result;
}
#endif

#ifdef TEST
int main() {
    /*
    auto x = Catalogue::make();
    auto m = x->add("test");
    auto res = x->build();
    */
}
#endif

extern "C" {
    void * get_catalogue(int * n) {
        Catalogue * c = Catalogue::get_current_during_init();
        if (c) {
            return c->get_catalogue(n);
        } else {
            /* this will crash */
            *n = 0;
            return 0;
        }
    }
}

#ifdef PYTHON
namespace py = pybind11;

py::object Catalogue::load() {
    Dl_info dl_info;
    dladdr((const void*)::get_catalogue, &dl_info);
    //py::scoped_interpreter guard; 
    auto arbor = py::module::import("arbor");
    set_as_current_for_init();
    return arbor.attr("load_catalogue")(dl_info.dli_fname);
}

PYBIND11_MODULE(nmlcat, m) {
    py::module arbor = py::module::import("arbor");
    py::class_<Mechanism, std::shared_ptr<Mechanism>>(m, "Mechanism")
        ;
    py::class_<Catalogue, std::shared_ptr<Catalogue>>(m, "Catalogue")
        .def_static("make", &Catalogue::make)
        .def("add", &Catalogue::add)
        .def("load", &Catalogue::load)
        ;
}

#endif

output.txt

two::init()
one::init()
two::write_ions()
one::write_ions()
two::init()
one::init()
two::apply_events()
two::compute_currents()
one::apply_events()
one::compute_currents()
two::advance_state()
one::advance_state()
two::write_ions()
one::write_ions()
two::apply_events()
two::compute_currents()
one::apply_events()
one::compute_currents()
two::advance_state()
one::advance_state()
two::write_ions()
one::write_ions()
two::apply_events()
two::compute_currents()
one::apply_events()
one::compute_currents()
two::advance_state()
one::advance_state()
two::write_ions()
one::write_ions()
[-65.0, -65.0, -65.0, -65.0, -65.0, -65.0, -65.0, -65.0, -65.0, -65.0]

test.sh

PREFIX=~/repos/llandsmeer/arbor/build-debug/prefix
SRC=nmlcat.cpp

g++ ${SRC} \
    -I "${PREFIX}/include/" \
    -std=c++20 \
    ${PREFIX}/lib/libarbor*.a \
    -DTEST \
    -g

g++ \
    -I "${PREFIX}/include/" \
    -g \
    -Wall \
    -Wextra \
    -Werror \
    -shared -fPIC \
    -std=c++20 \
    $(python3 -m pybind11 --includes) \
    ${SRC} \
    -o nmlcat$(python3-config --extension-suffix) \
    -DPYTHON \
    /home/llandsmeer/.local/lib/python3.10/site-packages/arbor/_arbor.cpython-310-x86_64-linux-gnu.so \
    /home/llandsmeer/.local/lib/python3.10/site-packages/arbor.libs/libxml2-3998bec4.so.2.9.1 \
    /home/llandsmeer/.local/lib/python3.10/site-packages/arbor.libs/liblzma-004595ca.so.5.2.2

./a.out

python3 main.py