Implement Python Code Accelerations

[JamesPHoughton]

It would be good to use some of python's capabilities for improving runtime speeds, such as cython or theano.

[JamesPHoughton]

This should be implemented in a way that allows graceful fallback if cython or theano are not installed. Ie, at initialization, check for them, and if found, set a flag, and use one set of functions that takes advantage of the optimization. If not, have a fallback set of functions...

[JamesPHoughton]

Numba could also be a consideration...

[JamesPHoughton]

Looks like support for classes in numba is still 'in the pipeline': numba/numba#1336

[JamesPHoughton]

Spent some time working out how to use Numba for integration. Its a pretty restricted set of operations that you can use, so it isn't quite as simple as sticking a @jit decorator on things, as they imply. It can't handle classes, or even list comprehensions, so there's a pretty fundamental rewrite necessary.

Assume you have a model file (formatted without classes) called test_model.py:

def stock():
    return state['stock']

def stock_init():
    return 2

def dstock_dt():
    return flow()

def constant():
    return constant.value
constant.value = .1

def flow():
    return constant() * stock()

initial_time = 0

dt = .1

We can apply the @jit decorator the old fashioned way to each of the functions:

import numpy as np
from numba import jit

import test_model as components

for element in filter(lambda n: not n.startswith('__'), dir(components)):
    e = getattr(components, element)
    print e
    if hasattr(e, '__call__'):
        print element
        setattr(components, element, jit(e))

Here are some helper functions/constructs. Note that we have a dictionary of derivative functions that matches the state dictionary.

_stocknames = [name[:-5] for name in dir(components) if name[-5:] == '_init']
_dfuncs = {name: getattr(components, 'd%s_dt'%name) for name in _stocknames}
_initfuncs = {name: getattr(components, '%s_init'%name) for name in _stocknames}

def reset_state():
        """Sets the model state to the state described in the model file. """
        components.t = components.initial_time #set the initial time
        components.state = {}
        retry_flag = False
        for key in _stocknames:
            try:
                components.state[key] = _initfuncs[key]() #set the initial state
            except TypeError:
                retry_flag = True
        if retry_flag:
            self.reset_state() #potential for infinite loop!
reset_state()

We can define a step and integrate function using numba @jit. I'm integrating right into the state dictionary, which saves unpackaging and repackaging as we had to do with the old scipy version. (It would be nice if numba could handle the dictionary comprehension, but alas).

@jit
def step(ddt, state, dt):
    outdict = {}
    for key in ddt:
        outdict[key] = ddt[key]()*dt + state[key]
    return outdict

@jit
def integrate(ddt, timesteps):
    outputs = range(len(timesteps))
    for i, t2 in enumerate(timesteps):
        components.state = step(ddt, components.state, t2-components.t)
        components.t = t2
        outputs[i] = components.state
    return outputs

Now if we run the integrator:

ts = range(0,100)
%time integrate(_dfuncs, ts)
CPU times: user 359 ms, sys: 9.06 ms, total: 368 ms
Wall time: 367 ms

The first time the functions are called, the numba compiler takes some time to work on them. The next time is ~90 times faster:

ts = range(100,200)
%time integrate(_dfuncs, ts)
CPU times: user 3.94 ms, sys: 1.02 ms, total: 4.95 ms
Wall time: 4.21 ms

Need to do a speed comparison for the same test optimized for using the scipy integrate that we're currently doing, as that's the correct speedtest benchmark.

Have to see if we can work out a caching arrangement that works with numba, as the pattern suggested in #15 (comment) probably wont work.

[JamesPHoughton]

So, turns out that numba is really just slowing things down. Without the JIT decorator, the speed is actually better. This is somewhat annoying. Not sure what to do now...

ts = range(0,100)
%time integrate(_dfuncs, ts)
CPU times: user 168 µs, sys: 61 µs, total: 229 µs
Wall time: 177 µs

This is essentially because numba can't handle thinking of functions as object, so the idea of iterating through a list of functions to get derivative values for each of the stocks is limiting. Instead of truly converting to a low level piece of code, most of it is still being run in python, but we now have the overhead of the typing system.

If we want to use numba in a real way, including on the integrator, need to explicitly write out all the function calls.

It may still be useful in large models, if you essentially only use it on the model functions themselves.

[datnamer]

Great package!

Fyi, Numba has JIT classes now. Might be able to use a callable class.

Also, I would really encourage you to open a issue/ feauture request in the numba repo. The dev team is responsive and would appreciate the feedback (and I would love to see numba integrated).

[enekomartinmartinez]

It would be interesting to start working on this after finishing moving the py_backend to numpy #328
I have opened several issues about possible performance improvements and updated the projects tab so we can track the current status of each one easily.