%matplotlib inline
%load_ext autoreload
%autoreload 2

import numpy as np
import matplotlib.pyplot as plt


from pymks import MKSRegressionModel
from pymks import FastMKSRegressionModel
from pymks import FiPyCHModel
from sklearn import metrics
from sklearn.grid_search import GridSearchCV
from pymks import FastMKSRegressionModel
import numexpr

mse = metrics.mean_squared_error

np.random.seed(101)
X = np.random.random((1000, 21, 21))

from IPython.parallel import Client
engines = Client()
dview = engines[:]
dview.block = True
dview.activate()
print "engines IDs",engines.ids

dview.scatter('X', X)

%px print X.shape

%%px

from pymks import FiPyCHMode
y = FiPyCHModel().predict(X)

y = dview.gather('y')
print y.dtype
print y.shape

from sklearn.cross_validation import train_test_split

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.5, random_state=3)

Nbins = np.arange(2, 20)

from pymks import MKSRegressionModel
from sklearn import metrics
from sklearn.grid_search import GridSearchCV
tuning_parameters = [{'Nbin': Nbins}]
    
def neg_mse(a, b):
    from sklearn import metrics
    mse = metrics.mean_squared_error
    return -mse(a, b)

gridSearch = GridSearchCV(MKSRegressionModel(), tuning_parameters, cv=5,
                          score_func=neg_mse, n_jobs=-1)

out = gridSearch.fit(X_train, y_train)

print gridSearch.best_estimator_
gridSearch.score(X_test, y_test)

for params, mean_score, scores in gridSearch.grid_scores_:
    print("%1.3e (+/-%1.3e) for %r"% (mean_score, scores.std() / 2, params))

for Nbin in Nbins:
    model = MKSRegressionModel(Nbin=Nbin)
    model.fit(X_train, y_train)
    y_true, y_pred = y_test, model.predict(X_test)
    print 'Nbin: {0}, mse: {1:1.3e}'.format(Nbin, mse(y_true, y_pred))

np.random.seed(102)
a = np.random.random(1e7)
b = np.random.random(1e7)
c = np.random.random(1e7)
a.dtype

?%timeit

d = a**2 + b**2 + 2*a*b
%timeit a**2 + b**2 + 2*a*b

e = numexpr.evaluate('a**2 + b**2 + 2 * a * b')
%timeit e = numexpr.evaluate('a**2 + b**2 + 2 * a * b')

np.all(d == e)

%timeit a**2

numexpr.set_num_threads(8)
%timeit numexpr.evaluate("a**2")

Nproc = 8
for Nthread in range(1, Nproc + 1):
    print "number of threads:",Nthread
    numexpr.set_num_threads(Nthread)
    %timeit numexpr.evaluate('a**2 + b**2 + 2 * a * b')

%timeit sin(a)**2 + cos(b)**2

%timeit numexpr.evaluate('sin(a)**2 + cos(b)**2')

a = np.random.random((1024, 16384))
b_numpy = np.fft.fftn(a, axes=(0, 1))
%timeit b_numpy = np.fft.fftn(a, axes=(0, 1))

import pyfftw

input_array = pyfftw.n_byte_align_empty(a.shape, 16, 'complex128')
output_array = pyfftw.n_byte_align_empty(a.shape, 16, 'complex128')
fft_plan = pyfftw.FFTW(input_array, output_array, threads=Nproc, axes=(0, 1), direction='FFTW_FORWARD')

input_array[:] = a;
b_fftw = fft_plan()
%timeit input_array[:] = a; b_fftw = fft_plan()

print b_fftw.dtype
print b_numpy.dtype

np.allclose(b_numpy, b_fftw)

for Nthread in range(1, Nproc + 1):
    fft_plan = pyfftw.FFTW(input_array, output_array, threads=Nthread, axes=(0, 1), direction='FFTW_FORWARD')
    print 'Nthread:',Nthread
    %timeit input_array[:] = a; b_fftw = fft_plan()























np.random.seed(11)
X0 = np.random.random((1, 1024, 1024))

fast_model = FastMKSRegressionModel(Nbin=5, threads=8)
slow_model = MKSRegressionModel(Nbin=5)

fast_model.fit(X, y)
fast_model.resize_coeff(X0.shape[1:])
y_fast = fast_model.predict(X0.copy())

slow_model.fit(X, y)
slow_model.resize_coeff(X0.shape[1:])
y_slow = fast_model.predict(X0.copy())

print np.allclose(y_fast, y_slow)

%timeit slow_model.predict(X0.copy())

%timeit fast_model.predict(X0.copy())

from IPython.display import clear_output
import time
from pymks import FiPyCHModel
N = 41
np.random.seed(13)
phi0 = np.random.random((1, N, N))

model = FastMKSRegressionModel(Nbin=5)
model.fit(X, y)
model.resize_coeff((N, N))

fig = plt.figure()

fipymodel = FiPyCHModel()
fipy_phi1 = phi0.copy()
mks_phi1 = phi0.copy()
steps = 30
dt=1e-8
for i in range(steps):
    fipy_response = fipymodel.predict(fipy_phi1)
    mks_response = model.predict(mks_phi1)
    #Euler forward
    fipy_phi1 = fipy_response * dt + fipy_phi1
    mks_phi1 = mks_response * dt + mks_phi1
    mks_phi1 = np.minimum(mks_phi1, 1)
    mks_phi1 = np.maximum(mks_phi1, 0)
    #print phi0
    plt.subplot(1,2,1)
    plt.title('FiPy')
    a = plt.imshow(fipy_phi1.reshape((N, N)), vmin=0.0, vmax=1.0)
    plt.subplot(1,2,2)
    plt.title('MKS')
    b = plt.imshow(mks_phi1.reshape((N, N)), vmin=0.0, vmax=1.0)
    time.sleep(1)
    clear_output()
    display(fig)
    print i


np.random.seed(11)
X0 = np.random.random((1, 256, 256))
dt = 1e-8
model = FastMKSRegressionModel(Nbin=7, threads=8)
model.fit(X, y)
model.resize_coeff(X0.shape[1:])
data = []
X1 = X0.copy()
steps = 3000
for step in range(steps + 1):
    y1 = model.predict(X1)
    X1 = numexpr.evaluate('X1 + dt * y1')
    X1 = np.minimum(X1, 1)
    X1 = np.maximum(X1, 0)
    if step % (steps / 200) == 0:
        data.append(X1[0])
        print step,

from IPython.display import clear_output
import time

fig = plt.figure()
ax = plt.axes()
a = plt.imshow(X0[0].copy(), vmin=0.0, vmax=1.0)
plt.colorbar()
for d in data:
    a.set_array(d.copy())
    plt.show()
    time.sleep(1)
    clear_output()
    display(fig)

# JSAnimation import available at https://github.com/jakevdp/JSAnimation
from JSAnimation import IPython_display
from matplotlib import animation

# create a simple animation
fig = plt.figure()
ax = plt.axes()
im = plt.imshow(X0[0], vmin=0.0, vmax=1.0)
plt.colorbar()

def init():
    im.set_array(X0[0])
    return im

def animate(i):
    im.set_array(data[i])
    return im

anim = animation.FuncAnimation(fig, animate, init_func=init,
                               frames=len(data), interval=1, blit=True)

from IPython.display import HTML
plt.close(anim._fig)
HTML(IPython_display.anim_to_html(anim))