#!/usr/bin/python2 import numpy as np from mpl_toolkits.mplot3d import Axes3D from mpl_toolkits.mplot3d.art3d import Poly3DCollection from matplotlib import cm import matplotlib.pyplot as plt fig = plt.figure() ax = fig.add_subplot(111, projection='3d') vmin, vmax = -np.pi, np.pi plot_args = {'cmap':'YlGnBu_r', 'linewidth': 0.4, 'antialiased': True} cmap = cm.get_cmap(plot_args['cmap']) def modpi(x, ref): return (x + np.pi/2 - ref) % (np.pi) - np.pi/2 + ref xa = np.linspace(-1, 1, 31) xc = 0.5 * (xa[:-1] + xa[1:]) ya = np.linspace(-1, 1, 31) yc = 0.5 * (ya[:-1] + ya[1:]) y, x = np.meshgrid(ya, xa) z = np.arctan(y / x) ym, xm = np.meshgrid(yc, xc) zc = np.arctan(ym / xm) colors = cmap((zc.flatten() - vmin) / (vmax - vmin)) verts = [] for ix in range(len(xa)-1): for iy in range(len(ya)-1): zm = zc[ix, iy] polygon = [ (xa[ix], ya[iy], modpi(z[ix,iy], zm)), (xa[ix], ya[iy+1], modpi(z[ix,iy+1], zm)), (xa[ix+1], ya[iy+1], modpi(z[ix+1,iy+1], zm)), (xa[ix+1], ya[iy], modpi(z[ix+1,iy], zm))] for i, p in enumerate(polygon): if (p[0], p[1]) == (0., 0.): z1 = polygon[(i-1)%len(polygon)][2] z2 = polygon[(i+1)%len(polygon)][2] new_points = [(p[0], p[1], z1), (p[0], p[1], z2)] polygon = polygon[:i] + new_points + polygon[i+1:] verts.append(polygon) ax.add_collection3d(Poly3DCollection(verts, facecolors=colors, **plot_args)) #ax.plot_surface(x, y, z, **plot_args) ax.view_init(azim=-70, elev=60) ax.set_xlim(-1, 1) ax.set_ylim(-1, 1) ax.set_zlim(vmin, vmax) plt.xticks([-1, -0.5, 0, 0.5, 1], [r"$-1$", r"$-1/2$", r"$0$", r"$1/2$", r"$1$"]) plt.yticks([-1, -0.5, 0, 0.5, 1], [r"$-1$", r"$-1/2$", r"$0$", r"$1/2$", r"$1$"]) ax.set_zticks([-np.pi, 0, np.pi]) ax.set_zticklabels([r"$-\pi$", r"$0$", r"$\pi$"]) ax.w_xaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) ax.w_yaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) ax.w_zaxis.set_pane_color((1.0, 1.0, 1.0, 0.0)) ax.set_xlabel(r"$x$", fontsize=18) ax.set_ylabel(r"$y$", fontsize=18) ax.set_zlabel(r"$\arctan(y/x)$", fontsize=18) plt.savefig("atan.svg", bbox_inches="tight", transparent=True) plt.show()