# plotting2d.py
'''
Examples for plotting 2d graphs, including images.
'''

import numpy as np
import matplotlib.pyplot as plt
# or:
import pylab as plt


#
# Basic line plotting.
#

plt.plot([1, 2, 3, 4])
plt.ylabel('some numbers')
plt.show()

# Activate interactive mode (plot automatically drawn).
plt.ion()
# Deactivate:
plt.ioff()

# For non-interactive need:
plt.show() # Show the plot.
plt.draw() # (Re)draw the plot.
# plt.draw is used in interactive mode to update a figure that has been
# altered, but not automatically re-drawn.

# Change plotting style.
plt.plot([1, 2, 3, 4], [1, 4, 9, 16], 'ro')

# Plot multiple data sets with plotting style string.
t = np.arange(0, 5, 0.2)
plt.plot(t, t, 'r--', t, t**2, 'bs', t, t**3, 'g^')

# The plot function has many arguments.
plt.plot(t, t**2, alpha=0.8, color='r', linestyle='--', linewidth=2, label='first')
plt.plot(t, np.sin(t), linestyle='', markeredgecolor='black', markeredgewidth=2,
         markerfacecolor='r', markersize=15, marker='o', label='second')

# Add some information.
plt.grid()
plt.legend()
plt.xlabel(r'$x$', fontsize=25)
plt.ylabel(r'$y^2 + \int_{0}^{\infty}$', fontsize=25)
plt.title('Hello Matplotlib!')

# Logarithmic scale.
t = np.linspace(1, 10, 10)
plt.semilogy(t, np.exp(t))
t = np.linspace(1, 1000, 1000)
plt.loglog(t, t**-3)
plt.xscale('linear')

# Plot with error bars.
x = np.linspace(0, 2*np.pi, 30)
y = np.sin(x)
err = np.random.random(len(x))/5
plt.errorbar(x, y, yerr=err)
plt.xlim(xmin=0, xmax=2*np.pi)

# Add annotations.
plt.annotate('high', xy=(np.pi/2, 1), xycoords='data', xytext=(np.pi/2+0.5, 1.2),
             textcoords='data', arrowprops=dict(arrowstyle="->"))
plt.annotate(r'$-\frac{1}{\pi}\left[\int_{-\infty}^{\infty} e^{-x^2}\ {\rm d}x\right]^2$',
             xy=(3*np.pi/2, -1), xycoords='data', xytext=(4, -0.5),
             textcoords='data', arrowprops=dict(arrowstyle="fancy"))


#
# Simple 2d color map plots.
#

# Plot on a uniform mesh.
x = np.linspace(0, 2*np.pi, 100)
y = np.linspace(0, 4*np.pi, 200)
xx, yy = np.meshgrid(x, y, indexing='ij')
zz = np.sin(xx) + np.cos(yy)
plt.imshow(zz.T, cmap='gist_heat', interpolation='nearest', vmin=-0.5, vmax=2.5,
           origin='lower', extent=[x.min(), x.max(), y.min(), y.max()],
           aspect='auto')
plt.xlabel('x')
plt.ylabel(r'$y$')
plt.colorbar()

# Plot on a non-uniform mesh.
# Still need a regular grid for this.
x = np.linspace(-1, 1, 40)
y = np.linspace(-1, 1, 40)
xx, yy = np.meshgrid(x, y, indexing='ij')
# Add some deformation.
phi = np.arctan2(xx, yy)
rr = np.sqrt(xx**2+yy**2)
xx = xx + np.cos(phi)*(rr**4-rr**2)/5
yy = yy + np.sin(phi)*(rr**4-rr**2)/5
zz = np.sin(2*np.pi*xx) + np.cos(2*np.pi*yy)
pcol = plt.pcolormesh(xx, yy, zz, cmap='rainbow', linewidth=0, rasterized=True)
pcol.set_edgecolor('face')
plt.xlim(xmin=xx.min())

# Plot contours of the map.
plt.contourf(xx, yy, zz)


#
# Simple 3d plots.
#

from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
#ax = fig.gca(projection='3d')
ax = Axes3D(fig)

# line plot
theta = np.linspace(-4 * np.pi, 4 * np.pi, 100)
z = np.linspace(-2, 2, 100)
r = z**2 + 1
x = r * np.sin(theta)
y = r * np.cos(theta)
ax.plot(x, y, z, label='parametric curve')
ax.legend()

# surface plot
x = np.linspace(0, 2*np.pi, 100)
y = np.linspace(0, 4*np.pi, 200)
xx, yy = np.meshgrid(x, y, indexing='ij')
zz = np.sin(xx) + np.cos(yy)
ax.plot_surface(xx, yy, zz, rstride=8, cstride=8, alpha=0.3)


#
# Plotting vector fields.
#

# Plot streamlines.
x = np.linspace(-3, 3, 100)
y = np.linspace(-3, 3, 100)
xx, yy = np.meshgrid(x, y, indexing='ij')
uu = -1 - xx**2 + yy
vv = 1 + xx - yy**2
speed = np.sqrt(uu**2 + vv**2)
# NB: We need the 1d arrays as first input parameters.
strm = plt.streamplot(x, y, uu, vv, color=speed, linewidth=2, cmap=plt.cm.autumn)
plt.colorbar(strm.lines)

plt.streamplot(x, y, uu, vv, density=[0.5, 1])

plt.streamplot(x, y, uu, vv, density=0.6, color='k', linewidth=speed)

# Plot the vectors.
plt.quiver(xx[::4, ::4].T, yy[::4, ::4].T, uu[::4, ::4], vv[::4, ::4],
           pivot='mid', color='b')



#
# Figure container.
#

fig = plt.figure(num=0, figsize=(8, 6))
print(fig.number)
fig.clf()

# Get current figure.
fig2 = plt.figure(num=1, figsize=(10, 8))
ff = plt.gcf()
ff.number

# All plots without figure specification are plotted in the active figure.
x = np.linspace(0, 2*np.pi, 100)
y = np.sin(x)
z = np.cos(x)
plt.plot(x, y)

# Clear the current figure.
plt.clf()
fig.clf()

# Configure the figure beforehand.
plt.rc("figure.subplot", left=0.12)
plt.rc("figure.subplot", right=.88)
plt.rc("figure.subplot", bottom=.18)
plt.rc("figure.subplot", top=.95)
# All figures will the be created with this configuration.


#
# Axes container.
#

fig = plt.figure(num=0, figsize=(8, 6))
ax1 = fig.add_subplot(2, 1, 1) # two rows, one column, first plot
ax2 = fig.add_subplot(2, 1, 2) # two rows, one column, second plot
ax1.plot(x, y)
ax2.plot(x, z)
ax2.set_xlabel(r'$x$')
plt.ylabel('z') # Changes the active axes.
plt.sca(ax1)
plt.ylabel('y') # Changes the active axes.
ax2.cla()
plt.cla()

# Access the figure, in which the axes lives in.
ax2.figure
# Access the axes in the figure.
for ax in fig.axes:
    ax.grid()

# Alternatively, but less control:
fig = plt.figure(num=0, figsize=(8, 6))
plt.subplot(211)
plt.plot(x, y)
plt.subplot(212)
plt.plot(x, z)

# Create axes at arbitrary location.
fig = plt.figure(num=0, figsize=(8, 6))
ax1 = fig.add_subplot(2, 1, 1) # two rows, one column, first plot
ax2 = fig.add_subplot(2, 1, 2) # two rows, one column, second plot
ax3 = fig.add_axes([0.35, 0.1, 0.9, 0.7])
plt.plot(x, y)

# Create axes using subplot.
fig = plt.figure(figsize=(8, 6))
ax = fig.add_subplot(111)


#
# Axis container.
#

fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
ax.plot(x, y)
axis = ax.xaxis

# Access ticks information.
axis.get_ticklocs()
axis.get_ticklabels()[0].get_text()
for tick_label in axis.get_ticklabels():
    print(tick_label.get_text())

# Change ticks information.
labels = [item.get_text() for item in ax.get_xticklabels()]
i = 1
for l in range(len(labels)):
    labels[l] = r'${0}$'.format(i)
    i *= 2
ax.set_xticklabels(labels)

# Change ticks style.
for tick in ax.xaxis.get_major_ticks():
    tick.label.set_fontname('serif')
for label in ax.xaxis.get_ticklabels():
    label.set_position((0, -0.03))


#
# Ticks container.
#

import matplotlib.ticker as ticker
formatter = ticker.FormatStrFormatter('$%1.2f')
ax.xaxis.set_major_formatter(formatter)
for tick in ax.xaxis.get_major_ticks():
    tick.label1On = False
    tick.label2On = True
    tick.label2.set_color('green')
plt.draw()
plt.show()


#
# More on plotting objects.
#

# Lines.
fig = plt.figure()
ax = fig.add_subplot(1, 1, 1)
lines = ax.plot(x, y, x, z)
ll = lines[0]
ll.set_color('red')
ll.set_ydata(np.sqrt(x))


#
# Legends.
#

x = np.linspace(0, 2*np.pi, 100)
y = np.sin(x)
z = np.cos(x)

plt.plot(x, y, label='first')
plt.plot(x, z, label='second')

leg = plt.legend(loc=0, shadow=False, fancybox=False, numpoints=1)
#leg = plt.gca().get_legend()
# Change the font size of the legend.
ltext = leg.get_texts() # all the text.Text instance in the legend
for i in range(len(ltext)):
    legLine = ltext[i]
    legLine.set_fontsize(25)
    legLine.set_family('serif')
frame = leg.get_frame()
frame.set_facecolor('1')
leg.draw_frame(False)


#
# 2d data.
#

# Image data.
plt.figure()
x = np.linspace(0, 2*np.pi, 100)
y = np.linspace(0, 4*np.pi, 200)
xx, yy = np.meshgrid(x, y, indexing='ij')
zz = np.sin(xx) + np.cos(yy)
im = plt.imshow(zz.T)
im.set_data(zz**2)

# Color bar.
cb_label = r'$\langle\eta J^2\rangle_{xy}$'
cb = plt.colorbar(format='%.2e')
cb.set_label(cb_label, fontsize=25)
cbytick_obj = plt.getp(cb.ax.axes, 'yticklabels')
plt.setp(cbytick_obj, fontsize=15, family='serif')


#
# Two plots next to each other.
#

x = np.linspace(0, 2*np.pi, 100)
y = np.sin(x)
z = np.cos(x)
fig, ax = plt.subplots(1, 2, sharey=True, figsize=(8, 6))
ax[0].plot(x, y)
ax[1].plot(x, z)
plt.ylim(ymax=2)
fig.subplots_adjust(wspace=0)


#
# Save into a file.
#
fig.savefig('my_plot.eps')
fig.savefig('my_plot.png', dpi=120, format='png')