import numpy as np
from matplotlib import pyplot as plt
from matplotlib import colors

def simplify_contours(qcset, thresh=plt.rcParams['path.simplify_threshold']):
    for c in qcset.collections:
        paths = c.get_paths()
        for i, p in enumerate(paths):
            p.simplify_threshold = thresh / plt.gcf().dpi
            paths[i] = p.cleaned(simplify=True)

xmax = 3.2
vmax = 4
levels = np.linspace(-vmax, vmax, 22)
X, Y = np.mgrid[-xmax:xmax:400j, -xmax:xmax:400j]
V  = 1.0 / np.maximum(np.sqrt((X + 1)**2 + (Y + 1)**2), 1e-5)
V += 1.0 / np.maximum(np.sqrt((X - 1)**2 + (Y - 1)**2), 1e-5)
V -= 1.0 / np.maximum(np.sqrt((X + 1)**2 + (Y - 1)**2), 1e-5)
V -= 1.0 / np.maximum(np.sqrt((X - 1)**2 + (Y + 1)**2), 1e-5)
plt.figure(figsize=(6, 6)).add_axes([0, 0, 1, 1])
contf = plt.contourf(X, Y, V, levels=levels, cmap='RdBu_r', extend='both',
       norm=colors.SymLogNorm(linthresh=1.5, vmin=-vmax, vmax=vmax))
simplify_contours(contf, plt.rcParams['lines.linewidth'])
cont = plt.contour(X, Y, V, levels=contf.levels, colors='k', linestyles='solid')
plt.xticks([]), plt.yticks([])
plt.gca().set_aspect(aspect='equal')
plt.gca().axis('off')
for i in -1,1:
    for j in -1,1:
        plt.text(i, j, {-1:u'\u2212', 1:'+'}[i*j],
                 size=18, ha='center', va='center')
plt.savefig('QuadrupoleContour.svg')