Theano and Numpy speed comparison for Lorenz Attractor

theano_numpy_diffeq.py

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


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

def fprime_lorenz(y,*args):
    sigma, rho, beta = args
    yprime = T.zeros_like(y)
#     yprime = theano.shared(np.zeros(10))
    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

def fprime_lorenz_numpy(y,*args):

    sigma, rho, beta = args
    yprime = np.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

y = T.vector('y')
h = T.scalar('h')
i = T.iscalar('i')
out = rungekuttastep(h,y,fprime_lorenz,10.,28.,8./3.)

result, updates = theano.scan(fn=lambda y, h: rungekuttastep(h,y,fprime_lorenz,10.,28.,8./3.) ,
                                outputs_info =[{'initial':y}],
                                non_sequences=h,
                                n_steps=i)
t0 = time.time()
f = theano.function([h,y,i],result)
print("Compilation time: {:.5f} seconds".format(time.time() - t0))

n_iter = 1000

t0 = time.time()
y = np.arange(3)
for i in np.arange(n_iter):
    y = rungekuttastep(0.001,y,fprime_lorenz_numpy,10.,28.,8./3.)
print("Time in calculation (numpy): {:.5f}".format(time.time()-t0))

t0 = time.time()
result = f(0.001,np.arange(3),n_iter)
print("Time in calculation (theano): {:.5f}".format(time.time()-t0))

print("Same array? {}".format(np.allclose(y, result[-1])))