Arbor master recipe v1

master_recipe.py

import arbor

class master_recipe(arbor.recipe):
    def __init__(self, *recipes):
        arbor.recipe.__init__(self)
        self.recipes = list(recipes)
        self.setup_master_recipe()
        ### :::THIS IS THE ONLY PROBLEM:::
        ### global_properties() can not be specified up to subrecipe level
        ### which means catalogue mechanism names can not collide
        self.props = arbor.neuron_cable_properties()
        self.cat = arbor.default_catalogue()
        self.props.register(self.cat)

    def setup_master_recipe(self):
        'Setup inner structures after self.recipes update'
        self.total_cells = sum(r.num_cells() for r in self.recipes)
        self.gid_offsets = [0]
        for r in self.recipes[1:]:
            offset = self.gid_offsets[-1] + r.num_cells()
            self.gid_offsets.append(offset)

    def to_gid(self, recipe_id, cell_id):
        'Convert subrecipe index and gid to master_recipe gid'
        assert 0 <= cell_id < self.recipes[recipe_id].num_cells()
        return self.gid_offsets[recipe_id] + cell_id

    def from_gid(self, gid):
        'Convert master_recipe gid to subrecipe index and gid'
        for rid, offset in enumerate(self.gid_offsets):
            if gid >= offset:
                res = rid, gid - offset
        return res

    def num_cells(self):
        return self.total_cells

    def cell_description(self, gid):
        rid, sub_id = self.from_gid(gid)
        return self.recipes[rid].cell_description(sub_id)

    def cell_kind(self, gid):
        rid, sub_id = self.from_gid(gid)
        return self.recipes[rid].cell_kind(sub_id)

    def connections_on(self, gid):
        rid, sub_id = self.from_gid(gid)
        old_connections = self.recipes[rid].connections_on(sub_id)
        new_connections = []
        for conn in old_connections:
            gid = self.to_gid(rid, conn.source.gid)
            source = gid, conn.source.label
            new_conn = arbor.connection(source, conn.dest, conn.weight, conn.delay)
            new_connections.append(new_conn)
        return new_connections

    def event_generators(self, gid):
        rid, sub_id = self.from_gid(gid)
        return self.recipes[rid].event_generators(sub_id)

    def probes(self, gid):
        rid, sub_id = self.from_gid(gid)
        return self.recipes[rid].probes(sub_id)

    def global_properties(self, kind):
        ## THIS!!!!! NOT GLOBAL
        return self.props


#### FROM PYTHON EXAMPLES: ####

from math import sqrt

def make_cable_cell(gid):
    tree = arbor.segment_tree()
    s = tree.append(arbor.mnpos, arbor.mpoint(-12, 0, 0, 6), arbor.mpoint(0, 0, 0, 6), tag=1)
    b0 = tree.append(s, arbor.mpoint(0, 0, 0, 2), arbor.mpoint(50, 0, 0, 2), tag=3)
    b1 = tree.append(b0, arbor.mpoint(50, 0, 0, 2), arbor.mpoint(50+50/sqrt(2), 50/sqrt(2), 0, 0.5), tag=3)
    b2 = tree.append(b0, arbor.mpoint(50, 0, 0, 1), arbor.mpoint(50+50/sqrt(2), -50/sqrt(2), 0, 1), tag=3)
    labels = arbor.label_dict()
    labels['soma'] = '(tag 1)'
    labels['dend'] = '(tag 3)'
    labels['synapse_site'] = '(location 1 0.5)'
    labels['root'] = '(root)'
    decor = arbor.decor()
    decor.paint('"soma"', arbor.density('hh'))
    decor.paint('"dend"', arbor.density('pas'))
    decor.place('"synapse_site"', arbor.synapse('expsyn'), 'syn')
    decor.place('"root"', arbor.spike_detector(-10), 'detector')
    cell = arbor.cable_cell(tree, labels, decor)
    return cell


class ring_recipe(arbor.recipe):
    def __init__(self, ncells):
        arbor.recipe.__init__(self)
        self.ncells = ncells
        self.props = arbor.neuron_cable_properties()
        self.cat = arbor.default_catalogue()
        self.props.register(self.cat)
    def num_cells(self):
        return self.ncells
    def cell_description(self, gid):
        return make_cable_cell(gid)
    def cell_kind(self, gid):
        return arbor.cell_kind.cable
    def connections_on(self, gid):
        src = (gid-1)%self.ncells
        w = 0.01 # 0.01 μS on expsyn
        d = 5 # ms delay
        return [arbor.connection((src,'detector'), 'syn', w, d)]
    def event_generators(self, gid):
        if gid==0:
            sched = arbor.explicit_schedule([1]) # one event at 1 ms
            weight = 0.1 # 0.1 μS on expsyn
            return [arbor.event_generator('syn', weight, sched)]
        return []
    def probes(self, gid):
        return [arbor.cable_probe_membrane_voltage('"root"')]
    def global_properties(self, kind):
        return self.props

ncells = 4
recipe = master_recipe(
        ring_recipe(ncells),
        ring_recipe(ncells),
        )
context = arbor.context()
decomp = arbor.partition_load_balance(recipe, context)
sim = arbor.simulation(recipe, decomp, context)
sim.record(arbor.spike_recording.all)
handles = [sim.sample((gid, 0), arbor.regular_schedule(0.1)) for gid in range(ncells)]
sim.run(100)
print('Simulation finished')
print('spikes:')

for gid in range(recipe.num_cells()):
    rid, sub_id = recipe.from_gid(gid)
    gid_back = recipe.to_gid(rid, sub_id)
    print(f'{gid=} {rid=} {sub_id=} {gid_back=}')
for sp in sim.spikes():
    gid = sp[0][0]
    rid, sub_id = recipe.from_gid(gid)
    print(' ', rid, sub_id, sp)