Theano and Numpy speed comparison for Lorenz Attractor

theano_numpy_ODE.py

import theano 
import theano.tensor as T 
import numpy
import h5py 
import os 
import time


def rungekuttastep(h, y, fprime, *theta):
    k1 = h*fprime(y,*theta)
    k2 = h*fprime(y + 0.5*k1,*theta)
    k3 = h*fprime(y + 0.5*k2,*theta)
    k4 = h*fprime(y + k3,*theta)
    y_np1 = y + (1./6.)*k1 + (1./3.)*k2 + (1./3.)*k3 + (1./6.)*k4
    return y_np1


class Measurement(object):
    def __init__(self, name):
        self.name = name
        return

    def __enter__(self):
        self.t_start = time.time()
        return

    def __exit__(self, exc_type, exc_val, exc_tb):
        t_end = time.time()
        print("{} took {:.3f} seconds".format(self.name, t_end - self.t_start))
        return


class LorenzAttractor(object):
    theta = (10., 28., 8./3.)

    @staticmethod
    def fprime_theano(y, *theta):
        sigma, rho, beta = theta
        yprime = T.zeros_like(y)
        yprime = T.set_subtensor(yprime[0],sigma*(y[1] - y[0]) )
        yprime = T.set_subtensor(yprime[1], y[0]*(rho - y[2]) - y[1])
        yprime = T.set_subtensor(yprime[2],y[0]*y[1] - beta*y[2] )
        return yprime

    @staticmethod
    def fprime_numpy(y,*theta):
        sigma, rho, beta = theta
        yprime = numpy.zeros(y.shape[0])
        yprime[0] = sigma*(y[1] - y[0])
        yprime[1] = y[0]*(rho - y[2]) - y[1]
        yprime[2] = y[0]*y[1] - beta*y[2]
        return yprime


results = {}

with Measurement("compilation"):
    y = T.vector('y')
    h = T.scalar('h')
    i = T.iscalar('i')


    out = rungekuttastep(h, y, LorenzAttractor.fprime_theano, *LorenzAttractor.theta)

    result, updates = theano.scan(fn=lambda y, h: rungekuttastep(h, y, LorenzAttractor.fprime_theano, *LorenzAttractor.theta),
                                outputs_info =[{'initial':y}],
                                non_sequences=h,
                                n_steps=i)
    f = theano.function([h,y,i], result)

n_iter = 1000

with Measurement("numpy-integration"):
    y = numpy.arange(3)
    for i in numpy.arange(n_iter):
        y = rungekuttastep(0.001, y, LorenzAttractor.fprime_numpy, *LorenzAttractor.theta)
    results["numpy"] = y

with Measurement("theano-integration"):
    results["theano"] = f(0.001, numpy.arange(3), n_iter)[-1]


solution = results["numpy"]
for key,sln in results.items():
    if not numpy.allclose(sln, solution):
        print("Solution by method '{}' does not match with numpy-solution.".format(key))
print("All Done.")