%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d.axes3d import Axes3D

def plot(t, r, w, dx, dy, ph):
    rx = np.arange(-r, r, w)
    ry = np.arange(-r, r, w)
    c = 1 if ph is None else 2
    fig = plt.figure(figsize=(c * 9 - 4, 5))
    fig.suptitle(t)

    if ph is not None:
        ax = fig.add_subplot(1, 2, 1, projection='3d')

        x, y, z = np.meshgrid(rx, ry, 0)
        ax.quiver(x, y, z, dx(x, y), dy(x, y), 0, length=w, color='red', normalize=True)

        x, y = np.meshgrid(rx, ry)
        ax.plot_surface(x, y, ph(x, y), cmap=plt.cm.bone)

    ax = fig.add_subplot(1, c, c)

    if ph is not None:
        r2 = np.arange(-r, r, w / 4)
        x, y = np.meshgrid(r2, r2)

        cs1 = ax.contourf(x, y, ph(x, y), 10, cmap=plt.cm.bone)
        #cs2 = ax.contour(cs1, colors='b')

        cbar = plt.colorbar(cs1)
        #cbar.add_lines(cs2)

    x, y = np.meshgrid(rx, ry)
    ax.quiver(x, y, dx(x, y), dy(x, y), color='red')

dx1, dy1, ph1 = lambda x, y: x, lambda x, y: -y, lambda x, y: (x**2-y**2)/2
plot(r"$z \mapsto x\,dx-y\,dy = d(x^2-y^2)/2$", 5, 0.5, dx1, dy1, ph1)

plot(r"$\bar z \mapsto x\,dx+y\,dy = d(x^2+y^2)/2$", 5, 0.5,
     lambda x, y: x, lambda x, y: y, lambda x, y: (x**2 + y**2) / 2)

dx2, dy2, ph2 = lambda x, y: x**2 - y**2, lambda x, y: -2*x*y, lambda x, y: x**3/3 - x*y**2
plot(r"$z^2 \mapsto (x^2-y^2)dx-2xy\,dy = d(x^3/3-xy^2)$", 5, 0.5, dx2, dy2, ph2)

plot(r"$\bar z^2 \mapsto (x^2-y^2)dx+2xy\,dy$", 5, 0.5,
     lambda x, y: x**2 - y**2, lambda x, y: 2*x*y, None)

dx3, dy3, ph3 = lambda x, y: x**3 - 3*x*y**2, lambda x, y: y**3 - 3*x**2*y, lambda x, y: (x**4 - 6*x**2*y**2 + y**4) / 4
plot(r"$z^3 \mapsto (x^3-3xy^2)dx+(y^3-3x^2y)dy = d(x^4-6x^2y^2+y^4)/4$", 5, 0.5, dx3, dy3, ph3)

makef = lambda a, b, c: lambda f3, f2, f1: lambda x, y: a * f3(x, y) + b * f2(x, y) + c * f1(x, y)
f = makef(1, 5, 10)
plot("$z^3+5z^2+10$", 5, 0.5, f(dx3, dx2, dx1), f(dy3, dy2, dy1), f(ph3, ph2, ph1))

plot("$e^z \mapsto e^x\cos\,y\,dx-e^x\sin\,y\,dy = d(e^x\cos\,y)$", 5, 0.5,
     lambda x, y: np.exp(x) * np.cos(y), lambda x, y: -np.exp(x) * np.sin(y), lambda x, y: np.exp(x) * np.cos(y))

plot("$\sin\,z \mapsto \sin\,x\,\cosh\,y\,dx-\cos\,x\,\sinh\,y\,dy = d(-\cos\,x\,\cosh\,y)$", 5, 0.5,
     lambda x, y: np.sin(x) * np.cosh(y), lambda x, y: -np.cos(x) * np.sinh(y), lambda x, y: -np.cos(x) * np.cosh(y))