In [1]:
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
from nengo.utils.compat import is_iterable
import collections

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

In [2]:
#The code in this cell is for reading the images from the MNIST database and not part of the brain model.
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 [3]:
#NOT USED

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

In [4]:
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 [5]:
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 [6]:
#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 [7]:
def stim_func(t):
    '''returns the image for first 0.1s'''
    if t < 0.01:
        return img
    else:
        return [0 for _ in range(len(img))]

In [8]:
def connection_func(x):
    '''takes the output from the first ensemble and rotates it 1 degree'''
    return rotate_img(x,1)

In [9]:
#Courtesy of Terry and Mike
original = np.random.random((28, 28))

flat = original.flatten()


def get_transform(width, height, angle, offset_x, offset_y):
    transform = np.zeros((width*height, width*height))
    for i in range(width):
        for j in range(height):
            ii = (i-offset_x)*np.cos(angle) - (j-offset_y)*np.sin(angle)
            jj = (i-offset_x)*np.sin(angle) + (j-offset_y)*np.cos(angle)
            
            iii = int(ii+offset_x)
            jjj = int(jj+offset_y)
            
            x2 = (ii+offset_x) - iii
            y2 = (jj+offset_y) - jjj
            x1 = 1-x2
            y1 = 1-y2
            
            amount = np.array([x1*y1, x1*y2, x2*y2, x2*y1])
            amount /= np.sum(amount)
            
            
            if 0 <=iii <width and 0<=jjj<height:        
                transform[i+width*j, iii+width*jjj] += amount[0]
                if jjj+1 < height:
                    transform[i+width*j, iii+width*(jjj+1)] += amount[1]
                    if iii+1 < width:
                        transform[i+width*j, iii+1+width*(jjj+1)] += amount[2]
                if iii+1 < width:
                    transform[i+width*j, iii+1+width*jjj] += amount[3]
    return transform

t = get_transform(dim, dim, np.pi/180, 14, 14)

print(t.shape)


(784L, 784L)

In [10]:
test_in = img
print(test_in.shape)
test_out = np.dot(t, test_in)

#for _ in range(10):
#    test_out = np.dot(t, test_out)
#    pylab.imshow(np.reshape(test_out, (dim,dim), 'F'), cmap='Greys_r')
#    plt.show()


(784L,)

In [16]:
#Function to control inhibition of ensembles
def inhibit_rotate(t):
    #A value of zero gives no inhibition
    '''if t < 0.2:
        return 0
    if t <0.3:
        return (t*10 - 2)*(dim**2)
    else:
        return dim**2'''
    return 0 if t < 0.3 else dim**2
    
def inhibit_identity(t):
    '''if t < 0.2:
        return dim**2
    if t <0.3:
        return (3 - t*10 )*(dim**2)
    else:
        return 0'''
    return dim**2 if t < 0.25 else 0

In [17]:
#A network is primarily used for grouping together related objects and connections for visualization purposes
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)
    
    '''An array of ensembles. This acts, in some ways, like a single high-dimensional ensemble,
    but actually consists of many sub-ensembles, each one representing a separate dimension. 
    This tends to be much faster to create and can be more accurate than having one huge 
    high-dimensional ensemble. However, since the neurons represent different dimensions separately,
    we cannot compute nonlinear interactions between those dimensions.'''
    ensArr = nengo.networks.EnsembleArray(100, dim**2, ens_dimensions=1,neuron_type=neuron_type)
    #incresing num neurons has smaller effect on run time
    
    #Connect each pixel of the input to its own ensemble
    nengo.Connection(ipt,ensArr.input)

    #Add ensemble to collect rotated output. Separate ensemble used to allow inhibition
    ensRotated = nengo.networks.EnsembleArray(100, dim**2, ens_dimensions=1,neuron_type=neuron_type)
    
    #Add ensemble to collect untransformed output, for when no transformation is occurring
    ensIdentity = nengo.networks.EnsembleArray(100, dim**2, ens_dimensions=1,neuron_type=neuron_type)
    
    #Connect output of ens array to second ensemble and transform using rotation matrix
    conn = nengo.Connection(ensArr.output,ensRotated.input, transform = t)
    
    #Connect rotated ensemble information back to the input of the enseble array
    nengo.Connection(ensRotated.output, ensArr.input)
    
    
    #Connect output of ens array to third ensemble and no transformation
    conn = nengo.Connection(ensArr.output,ensIdentity.input)
    
    #Connect rotated ensemble information back to the input of the enseble array
    nengo.Connection(ensIdentity.output, ensArr.input)
    
    
    
    #https://pythonhosted.org/nengo/examples/learn_communication_channel.html

    #Node to control inhibition of rotation ensemble
    inhib_rot = nengo.Node(inhibit_rotate)
    
    #nengo.Connection(inhib, ensRotated.neurons, transform=[[-1]] * ensRotated.n_neurons)
    #nengo.Connection(inhib, ensRotated.neurons, transform=-10*np.ones((n_neurons,1)))
    
    #Must apply inhibition to each ensemble in the array
    for e in ensRotated.ea_ensembles:
        nengo.Connection(inhib_rot, e.neurons, transform=[[-10]] * ensRotated.n_neurons)
        
    #Node to control inhibition of identity ensemble
    inhib_id = nengo.Node(inhibit_identity)
    
    #Must apply inhibition to each ensemble in the array
    for e in ensIdentity.ea_ensembles:
        nengo.Connection(inhib_id, e.neurons, transform=[[-10]] * ensIdentity.n_neurons)
    
    #Gathering output of ensArr
    probe = nengo.Probe(ensArr.output,# attr='decoded_output',#sample_every=0.001,
                       synapse=probe_synapse)

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

In [19]:
sim.run(run_time)


Simulation finished in 0:08:46.                                                 

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

In [ ]:
'''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)

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

In [20]:
'''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 [ ]:
# save the output
#cPickle.dump(sim.data[probe], open( "Buffer_rotations_in_connection_ensemble_array_direct.p", "wb" ) )
#cPickle.dump(sim.data[probe], open( "Buffer_rotations_in_connection_ensemble_array_LIF_100_stim0.01.p", "wb" ) )