Skip to content

WIP early binding / closures on simulator components

Marmaduke Woodman requested to merge early-bind into master

Created by: maedoc

While working on the #7, I found generating closures (nested/inner functions) a much easier way to ensure the gradient could be handled by autograd. This PR focuses on such "early binding" as it is frequently called, to avoid constantly re-referencing traits, and allow the simulator inner loop to focus on computing with a bunch of arrays. After just setting up early binding for a few components, the small simulations takes about half the time to run, partly by greatly reducing the number of calls to neotraits descriptors.

I'll paste the trace used and attach the call graphs with & without,

from pycallgraph import PyCallGraph
from pycallgraph.output import GraphvizOutput

from pycallgraph import GlobbingFilter
from pycallgraph import Config

from tvb.simulator.models.wong_wang import ReducedWongWang
from tvb.simulator.lab import *
import numpy

rww = ReducedWongWang(a=numpy.array([0.27]), w=numpy.array([1.0]), I_o=numpy.array([0.3]))
sim = simulator.Simulator(
    model=rww,
    connectivity=connectivity.Connectivity(
    	weights=numpy.zeros((20, 20)),
    	tract_lengths=numpy.zeros((20, 20)),
    	region_labels=numpy.array([str(i) for i in range(20)]),
    	centres=numpy.zeros((20, 3))),
    coupling=coupling.Linear(a=numpy.array([0.5 / 50.0])),
    integrator=integrators.EulerStochastic(dt=1, noise=noise.Additive(nsig=numpy.array([1e-5]))), 
    monitors=(monitors.TemporalAverage(period=1.),),
    simulation_length=5e1
).configure()

# don't trace setup stuff on first iterations
(time, data), = sim.run()

graphviz = GraphvizOutput(output_file='trace_master.png')
config = Config()
config.trace_filter = GlobbingFilter(exclude=[
    'pycallgraph.*',
    'numpy.*',
    'autograd.*'
])
with PyCallGraph(config=config, output=graphviz):
	(time, data), = sim.run()

For the current master, I see trace_master while after this commit, I see trace_binding

I would be curious to hear thoughts from @mihandrei and/or @i-Zaak on whether this would be worth pursuing generally in the simulator (independently of the gradient work)

Merge request reports