import numpy as np
import matplotlib.pyplot as plt

def arbitrary_function(x, y):
    if x < 0:
        return -1.0
    else:
        return np.sin(x) * np.exp(y)

xs = [-1, 1, 2, 3, 4.]
ys = [1, 1, 2, 3, 4.]
arbitrary_function(xs, ys)

vectorized_function = np.vectorize(arbitrary_function)
vectorized_function(xs, ys)

xs = np.linspace(-1, 1, num=5)
ys = np.linspace(-1, 1, num=5)

print xs
print ys

xx, yy = np.meshgrid(xs, ys)

print xx
print yy

plt.imshow(vectorized_function(xx, yy),
           interpolation='nearest',
           extent=[-1, 1, -1, 1],
           origin='lower',
          )

fine_xs = np.linspace(-1, 1, num=200)
fine_ys = np.linspace(-1, 1, num=200)
fine_xx, fine_yy = np.meshgrid(fine_xs, fine_ys)

plt.imshow(vectorized_function(fine_xx, fine_yy),
           interpolation='nearest',
           extent=[-1, 1, -1, 1],
           origin='lower',
          )

def arbitrary_function_2(x, y, a):
    if x < 0:
       return -1.0
    else:
        return np.sin(a*x) * np.exp(y)

vectorized_function_2 = np.vectorize(arbitrary_function_2)

plt.imshow(vectorized_function_2(fine_xx, fine_yy, 3.0),
           extent=[-1, 1, -1, 1],
           origin='lower',
          )

plt.imshow(vectorized_function_2(fine_xx, fine_yy, 6.0),
           extent=[-1, 1, -1, 1],
           origin='lower',
          )