from __future__ import division
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import rcParams

rcParams['font.family'] = 'serif'
rcParams['font.size'] = 16


def fun(x, y):
    return np.abs(y + x) + 3*np.abs(y - x)


X, Y = np.mgrid[-3:3:301j, -3:3:301j]
Z = fun(X, Y)
plt.figure(figsize=(10, 10))
levels = range(1,7)
CS = plt.contour(X, Y, Z, levels=levels, colors='k')
plt.clabel(CS)
plt.arrow(-2, -2, 0.5, 0, fc='r', ec='r', length_includes_head=True)
plt.arrow(-2, -2, 0, 0.5, fc='r', ec='r', length_includes_head=True)
plt.axis('image')
plt.xlabel(r'$x$', size=18)
plt.ylabel(r'$y$', size=18)
plt.title(r'$f(x,y) = |x + y| + 3|y -x|$')
plt.savefig('nonsmooth coordinate descent.svg')
plt.show()