import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation

xdim = 400
ydim =400
freq = 1. # spatial frequency of the grating
theta = np.pi/3. #orientation of the grating

#rectangular mask

from matplotlib import path

nx, ny = xdim/2, ydim/2
xoff,yoff = 20,10

poly_verts = [(xoff,ny-yoff), (xdim-xoff,ny-yoff), (xdim-xoff,ny+yoff),(xoff,ny+yoff)]

x, y = np.meshgrid(np.arange(xdim), np.arange(ydim))
x, y = x.flatten(), y.flatten()
points = np.vstack((x,y)).T

p = path.Path(poly_verts)
rectImg = p.contains_points(points)
rectImg = rectImg.reshape((ydim,xdim))

#print grid

#circular mask

imageSize = xdim;
r = (imageSize-1)/4; cx = cy = imageSize/2; 
[X, Y] = np.meshgrid(np.arange(imageSize), np.arange(imageSize))
diskImg = (X-cx)**2 + (Y-cy)**2 <= r**2



#compare grid = diskImg with grid = rectImg

grid = diskImg
#grid = rectImg

fig = plt.figure()

def grating(x, y):
    temp = np.cos((2.*(np.pi)*freq)*(np.cos(theta)*x + np.sin(theta)*y))
    return temp

x = np.linspace(0, 2 * np.pi, xdim)
y = np.linspace(0, 2 * np.pi, ydim).reshape(-1, 1)

im = plt.imshow(grating(x, y), cmap=plt.get_cmap('Greys'))

def updatefig(*args):
    global x,y
    x += np.pi / 15.
#    im.set_array(grating(x,y)*diskImg)
    im.set_array(grating(x,y)*grid)
    return im

ani = animation.FuncAnimation(fig, updatefig, interval=10, blit=False) # blit needs to be set to False
plt.show()