speed up glif models with numba jit

[russelljjarvis]

I am trying to speed up glif models with numba jit function decorator.

If I succeed can I make a pull request? https://towardsdatascience.com/speed-up-your-algorithms-part-2-numba-293e554c5cc1

[dyf]

Hi @russelljjarvis. Great question. That sounds exciting! That said, there is some ongoing work to make the GLIF models available in the NEST simulator. Some folks have prototyped that in this repo and shown it to be much more performant than the reference implementation here in allensdk:

https://github.com/AllenInstitute/GLIF2NEST

They are going to start working on integrating this into NEST proper soon, although I don't have a timeframe for you. Once that gets into a stable NEST release, we'll most likely deprecate the implementation here.

If you still want to give optimizing this implementation a go, I'm all for it. I'll just want you to add some regression tests.

[dyf]

Here's a link to the NEST PR, for reference: https://github.com/nest/nest-simulator/pull/1126

[russelljjarvis]

I still feel that using numba @jit decorators, would be worthwhile. I imagine NEST is not a fast and easy pip install, but allensdk, and numba are easy pip installs. So I think jit could better manage the tradeoff between package management burden, and code speeds up for many users.

[russelljjarvis]

I just made a pull request that has a skeleton code of how it would work. It's untested. https://github.com/AllenInstitute/AllenSDK/pull/378

In my experience, you can't jit every function but many you can, and the speedup is often significant.

[russelljjarvis]

Actually, I am wrong. JIT actually slows down the code, but that's not the first time I have seen this happen. JIT speeds up some functions but slows down others, I have not learned how to tell the difference. I think the problem is it can't speed up methods that use dictionaries.

I think it would be possible to get a speed up by systematically removing some of the calls to jit, but I am unsure which ones, and I don't have time to try.

https://stackoverflow.com/questions/52419402/is-it-possible-that-numba-jit-slows-down-my-gcd-execution

Please ignore both pull requests.

Regular time:
0.0019969940185546875
0.002218008041381836
0.0021698474884033203
0.0025441646575927734

Using numba jit optimization over multiple trials:

0.1391608715057373
0.11669373512268066
0.10923027992248535
0.10755181312561035
0.11692500114440918
0.10854673385620117
0.12166714668273926
0.10976099967956543

The test code (stored in a file called test_glif.py)

from allensdk.api.queries.glif_api import GlifApi
from allensdk.core.cell_types_cache import CellTypesCache
import allensdk.core.json_utilities as json_utilities

neuronal_model_id = 566302806

# download model metadata
glif_api = GlifApi()
nm = glif_api.get_neuronal_models_by_id([neuronal_model_id])[0]

# download the model configuration file
nc = glif_api.get_neuron_configs([neuronal_model_id])[neuronal_model_id]
neuron_config = glif_api.get_neuron_configs([neuronal_model_id])
json_utilities.write('neuron_config.json', neuron_config)

# download information about the cell
ctc = CellTypesCache()
ctc.get_ephys_data(nm['specimen_id'], file_name='stimulus.nwb')
ctc.get_ephys_sweeps(nm['specimen_id'], file_name='ephys_sweeps.json')
import allensdk.core.json_utilities as json_utilities
from allensdk.model.glif.glif_neuron import GlifNeuron

# initialize the neuron
neuron_config = json_utilities.read('neuron_config.json')['566302806']
neuron = GlifNeuron.from_dict(neuron_config)

# make a short square pulse. stimulus units should be in Amps.
stimulus = [ 0.0 ] * 100 + [ 10e-9 ] * 100 + [ 0.0 ] * 100

# important! set the neuron's dt value for your stimulus in seconds
neuron.dt = 5e-6

# simulate the neuron
import time
b4 = time.time()
output = neuron.run(stimulus)
after = time.time()
print(after-b4)