In [132]:
from matplotlib import pylab
import nengo
import random
import numpy as np
import gzip as gz
import cPickle
from cPickle import load
try:
    import Image
except ImportError:
    from PIL import Image
from scipy.sparse.linalg import svds
import scipy
from scipy import ndimage
import matplotlib.pyplot as plt
import matplotlib.animation as animation

#%matplotlib inline #Makes visualizations appar inline (Commented out because animation popup as new window)

In [133]:
def load_img(path, dims):
    """Load the image at path and return an array representing the raster.
    Flattens image. Shifts pixel activations such that 0 represents gray,
    normalizes the output array.
    Keyword arguments:
    path -- str, path of the image to be loaded.
    dims -- (w, h), where w,h are ints indicating dimensions of the image (in
        px)."""

    img = Image.open(path).resize(dims).getdata()
    img.convert('L')
    img = subtract(array(img).flatten(), 127.5)
    return img/norm(img)


def load_data(filename):
    """Uncompress, unpickle and return a .pkl.gz file.
    Keyword arguments:
    filename -- str, a valid file path"""

    return load(gz.open(filename))

def load_mini_mnist(option=None):
    """Load and return the first \%10 of the images in the mnist dataset.
    Does not return labels. Pass in 'train', 'valid' or 'test' if you want to
    load a specific subset of the dataset.
    Keyword arguments:
    option -- str (default=None)."""

    mini_mnist = load(gz.open('./mini_mnist.pkl.gz', 'rb'))
    if option == 'train':
        return mini_mnist[0]
    elif option == 'valid':
        return mini_mnist[1]
    elif option == 'test':
        return mini_mnist[2]
    else:
        return mini_mnist

In [134]:
def rotate_img(img, degrees):
    '''
    img is the dim**2 by 1 vector representing the pixel values.
    Rotates image the degrees passed in counterclockwise
    Returns the Reshaped image (to original shape which is the one dimensional vector)
    dim is a global variable
    '''
    original = img.shape
    
    newImg = scipy.ndimage.interpolation.rotate(np.reshape(img, (dim,dim), 'F'),degrees,reshape=False)
    newImg = np.reshape(newImg, original, 'F')
    return newImg

def move_img(img, (x,y)):
    '''Shifts image by amount of coordinates passed in
    Reshapes image to original
    '''
    original = img.shape
    
    newImg = scipy.ndimage.interpolation.shift(np.reshape(img, (dim,dim), 'F'),(y,x))
    
    newImg = np.reshape(newImg, original, 'F')
    return newImg
    
def add_two_images(img1,img2):
    return np.add(img1,img2)

In [135]:
def add_dif_shape_matrices(big,small,x,y):
    '''Helper function for resize'''
    #http://stackoverflow.com/questions/9886303/adding-different-sized-shaped-displaced-numpy-matrices
    b1 = big
    b2 = small

    pos_v, pos_h = x, y  # offset
    v_range1 = slice(max(0, pos_v), max(min(pos_v + b2.shape[0], b1.shape[0]), 0))
    h_range1 = slice(max(0, pos_h), max(min(pos_h + b2.shape[1], b1.shape[1]), 0))

    v_range2 = slice(max(0, -pos_v), min(-pos_v + b1.shape[0], b2.shape[0]))
    h_range2 = slice(max(0, -pos_h), min(-pos_h + b1.shape[1], b2.shape[1]))

    b1[v_range1, h_range1] += b2[v_range2, h_range2]
    
    return b1

def resize_img(img, scale):
    '''Resizes image to the scale passed in
    Reshapes image to original
    '''
    original = img.shape
    
    newImg = scipy.ndimage.interpolation.zoom(np.reshape(img, (dim,dim), 'F'),scale)
    #new img is not the right size, add it to an empy matrix with correct size
    newImg = add_dif_shape_matrices(np.zeros((dim,dim)), newImg,0,0)
    
    newImg = np.reshape(newImg, original, 'F')
    return newImg

In [136]:
conn_synapse = 0.1 #post synaptic time constant to use for filtering (pstc) - what does changing this do?
probe_synapse = 0.01 #pstc
multiplier = 2 #not used
n_neurons = 5000
direct = False #Direct - function computed explicitly instead of in neurons 
stop_time = 3.0
run_time = 3.0 #in seconds

In [137]:
dim = 28 #size of the image
mnist = load_mini_mnist()
train = mnist[0] #collection of training images
img = mnist[1][0] #image to be used for testing
compress_size = 400 #?
basis, S, V = svds(train.T, k=compress_size) #Used for encoding and decoding information 
    #a set of vectors representing what a hand drawn number should look like?

In [138]:
#Need same number of vectors in basis as number of neurons (randomly sample from basis)
expanded_basis = np.array([random.choice(basis.T) for _ in range(n_neurons)])

In [139]:
def stim_func(t):
    '''returns the image for first 0.1s'''
    if t < 0.1:
        return img
    else:
        return [0 for _ in range(len(img))]

In [140]:
def stim_func_rot(t):
    if t < 0.1:
        return 0
    elif t<0.3:
        return 10
    else:
        return 0

In [157]:
def stim_func_translate(t):
    if t < 0.3:
        return (0,0)
    elif t<0.6:
        return (0,0.2)
    else:
        return (0,0)

In [142]:
def connection_func(x):
    '''takes the output from the first ensemble and rotates it by degrees specified by last value'''
    return rotate_img(x[:-1],x[-1])
    #return rotate_img(x,10) - (x)

In [143]:
def node_func(t,x):
    newImg = rotate_img(x[:-3],x[-3])
    newImg = move_img(newImg,(x[-2],x[-1]))
    return newImg

In [158]:
with nengo.Network() as net:
    
    if direct:
        neuron_type = nengo.Direct() #function computed explicitly, instead of in neurons
    else:
        neuron_type = nengo.LIF() #spiking version of the leaky integrate-and-fire neuron model

    #Input stimulus - provide data to the ensemble
    ipt = nengo.Node(stim_func)
    
    ipt2 = nengo.Node(stim_func_rot)
    ipt3 = nengo.Node(stim_func_translate)
    
    ensArr = nengo.networks.EnsembleArray(100, dim**2+3, ens_dimensions=1) #incresing num neurons has small effect on run time
    
    #Node that brings all information from individual ensembles together and computes the rotation
    node = nengo.Node(node_func,size_in = dim**2+3, size_out =dim**2)
    
    for i in range(ensArr.n_ensembles-3):
        #Connect the input (img) to each ensemble individually 
        nengo.Connection(ipt[i],ensArr.ea_ensembles[i])
        #Connect each ensemble to the same node
        nengo.Connection(ensArr.ea_ensembles[i],node[i])
        #Connect the node output to each ensemble individually
        nengo.Connection(node[i],ensArr.ea_ensembles[i])
    
    #Connect input for rate of rotation to its own ensemble
    nengo.Connection(ipt2,ensArr.ea_ensembles[-3])
    nengo.Connection(ensArr.ea_ensembles[-3],node[-3])
    
    nengo.Connection(ipt3[0],ensArr.ea_ensembles[-2])
    nengo.Connection(ipt3[1],ensArr.ea_ensembles[-1])
    nengo.Connection(ensArr.ea_ensembles[-2],node[-2])
    nengo.Connection(ensArr.ea_ensembles[-1],node[-1])
    
      
    ens = nengo.Ensemble(n_neurons, dimensions = dim**2)
    nengo.Connection(node, ens)
    
    probe = nengo.Probe(ens, attr='decoded_output',#sample_every=0.001,
                       synapse=probe_synapse)

In [159]:
sim = nengo.Simulator(net)

In [160]:
sim.run(run_time)


Simulation finished in 0:03:47.                                                 

In [147]:
pylab.imshow(np.reshape(img, (dim,dim), 'F'), cmap='Greys_r')


Out[147]:
<matplotlib.image.AxesImage at 0x79ee7f28>

In [148]:
'''Image at stop time'''
pylab.imshow(np.reshape([0. if x < 0.00001 else x for x in sim.data[probe][int(stop_time*1000)-1]], 
                             (dim, dim), 'F'), cmap=plt.get_cmap('Greys_r'),animated=True)


Out[148]:
<matplotlib.image.AxesImage at 0x67503940>

In [149]:
'''Image at start time'''
pylab.imshow(np.reshape([0. if x < 0.00001 else x for x in sim.data[probe][1]], 
                             (dim, dim), 'F'), cmap=plt.get_cmap('Greys_r'),animated=True)


Out[149]:
<matplotlib.image.AxesImage at 0x3bbd208>

In [161]:
'''Animation for Probe output'''
fig = plt.figure()

def updatefig(i):
    im = pylab.imshow(np.reshape([0. if x < 0.00001 else x for x in sim.data[probe][i]],
                                 (dim, dim), 'F'), cmap=plt.get_cmap('Greys_r'),animated=True)
    
    return im,

ani = animation.FuncAnimation(fig, updatefig, interval=1, blit=True)
plt.show()

In [162]:
# save the output
#cPickle.dump(sim.data[probe], open( "Buffer_manipulations_ensemble_array_scalar_LIF.p", "wb" ) )