%matplotlib inline

import numpy as np
import matplotlib.pyplot as plt

"""
Explanation: matplotlib
Matplotlib is the core plotting package in scientific python.  There are others to explore as well (which we'll chat about on slack).
Note: the latest version of matplotlib (2.0) introduced a number of style changes.  This is the version we use here.
Also, there are different interfaces for interacting with matplotlib, an interactive, function-driven (state machine) commandset and an object-oriented version.  Usually for interactive work, we use the state interface.
We want matplotlib to work inline in the notebook.
End of explanation
"""


x = np.linspace(0,2.0*np.pi, num=100)
y = np.cos(x)

plt.plot(x,y)
plt.xlabel(r"$x$")
plt.ylabel(r"$\cos(x)$", fontsize="x-large")
plt.xlim(0, 2.0*np.pi)

"""
Explanation: Matplotlib concepts
Matplotlib was designed with the following goals (from mpl docs):

Plots should look great -- publication quality (e.g. antialiased)
Postscript output for inclusion with TeX documents
Embeddable in a graphical user interface for application development
Code should be easy to understand it and extend
Making plots should be easy

Matplotlib is mostly for 2-d data, but there are some basic 3-d (surface) interfaces.
Volumetric data requires a different approach
Gallery
Matplotlib has a great gallery on their webpage -- find something there close to what you are trying to do and use it as a starting point:
http://matplotlib.org/gallery.html
Importing
There are several different interfaces for matplotlib (see http://matplotlib.org/faq/usage_faq.html)
Basic ideas:

matplotlib is the entire package
matplotlib.pyplot is a module within matplotlib that provides easy access to the core plotting routines
pylab combines pyplot and numpy into a single namespace to give a MatLab like interface.  You should avoid this&mdash;it might be removed in the future.

There are a number of modules that extend its behavior, e.g. basemap for plotting on a sphere, mplot3d for 3-d surfaces
Anatomy of a figure
Figures are the highest level obect and can inlcude multiple axes

(figure from: http://matplotlib.org/faq/usage_faq.html#parts-of-a-figure )
Backends
Interactive backends: pygtk, wxpython, tkinter, ...
Hardcopy backends: PNG, PDF, PS, SVG, ...
Basic plotting
plot() is the most basic command.  Here we also see that we can use LaTeX notation for the axes
End of explanation
"""


plt.plot(x, np.sin(x), "ro", label="sine")
plt.plot(x, np.cos(x), "bx", label="cosine")
plt.xlim(0.0, 2.0*np.pi)
plt.legend(frameon=False, loc=5)

"""
Explanation: Note that when we use the plot() command like this, matplotlib automatically creates a figure and an axis for us and it draws the plot on this for us.  This is the state machine interface.  
<div style="background-color:yellow; padding: 10px"><h3><span class="fa fa-flash"></span> Quick Exercise:</h3></div>

We can plot 2 lines on a plot simply by calling plot twice.  Make a plot with both sin(x) and cos(x) drawn
<hr>

we can use symbols instead of lines pretty easily too&mdash;and label them
End of explanation
"""


plt.plot(x, np.sin(x), "r--", linewidth=3.0)
plt.plot(x, np.cos(x), "b-")
plt.xlim(0.0, 2.0*np.pi)

"""
Explanation: most functions take a number of optional named argumets too
End of explanation
"""


line = plt.plot(x, np.sin(x))
plt.setp(line)

"""
Explanation: there is a command setp() that can also set the properties.  We can get the list of settable properties as
End of explanation
"""


plt.subplot(211)

x = np.linspace(0,5,100)
plt.plot(x, x**3 - 4*x)
plt.xlabel("x")
plt.ylabel(r"$x^3 - 4x$", fontsize="large")

plt.subplot(212)

plt.plot(x, np.exp(-x**2))
plt.xlabel("x")
plt.ylabel("Gaussian")

# log scale
ax = plt.gca()
ax.set_yscale("log")

# get the figure and set its size
f = plt.gcf()
f.set_size_inches(6,8)

# tight_layout() makes sure things don't overlap
plt.tight_layout()
plt.savefig("test.png")

"""
Explanation: Multiple axes
there are a wide range of methods for putting multiple axes on a grid.  We'll look at the simplest method.  All plotting commands apply to the current set of axes
End of explanation
"""


f = plt.figure()
ax = f.add_subplot(111)

x = np.linspace(0, 2*np.pi, 100)
y = np.sin(x)
ax.plot(x, y)
f

"""
Explanation: Object oriented interface
In the object oriented interface, we create a figure object, add an axis, and then interact through those objects directly.
End of explanation
"""


ax.set_xlabel("x")
ax.set_ylabel("y")
f

"""
Explanation: Notice that with the state machine interface, each cell created a new figure and worked on that.  Here, our f is a figure object, and we can refer to that figure object across multiple cells to build our figure.
End of explanation
"""


def g(x, y):
    return np.exp(-((x-0.5)**2)/0.1**2 - ((y-0.5)**2)/0.2**2)

N = 100

x = np.linspace(0.0,1.0,N)
y = x.copy()

xv, yv = np.meshgrid(x, y)

"""
Explanation: Visualizing 2-d array data
2-d datasets consist of (x, y) pairs and a value associated with that point.  Here we create a 2-d Gaussian, using the meshgrid() function to define a rectangular set of points.
End of explanation
"""


plt.imshow(g(xv,yv), origin="lower")
plt.colorbar()

"""
Explanation: A "heatmap" style plot assigns colors to the data values.  A lot of work has gone into the latest matplotlib to define a colormap that works good for colorblindness and black-white printing.
End of explanation
"""


plt.contour(g(xv,yv))
ax = plt.axis("scaled")   # this adjusts the size of image to make x and y lengths equal


"""
Explanation: Sometimes we want to show just contour lines&mdash;like on a topographic map.  The contour() function does this for us.
End of explanation
"""


def y_experiment(a1, a2, sigma, x):
    """ return the experimental data in a linear + random fashion a1
        is the intercept, a2 is the slope, and sigma is the error """

    N = len(x)

    # randn gives samples from the "standard normal" distribution
    r = np.random.randn(N)
    y = a1 + a2*x + sigma*r
    return y

N = 40
x = np.linspace(0.0, 100.0, N)
sigma = 25.0*np.ones(N)
y = y_experiment(10.0, 3.0, sigma, x)

plt.errorbar(x, y, yerr=sigma, fmt="o")

"""
Explanation: <div style="background-color:yellow; padding: 10px"><h3><span class="fa fa-flash"></span> Quick Exercise:</h3></div>

Contour plots can label the contours, using the plt.clabel() function.
Try adding labels to this contour plot.
<hr>

Error bars
For experiments, we often have errors associated with the $y$ values.  Here we create some data and add some noise to it, then plot it with errors.
End of explanation
"""


xx = np.linspace(0, 2.0*np.pi, 1000)
plt.plot(xx, np.sin(xx), color="r")
plt.text(np.pi/2, np.sin(np.pi/2), r"maximum")
ax = plt.gca()
ax.spines['right'].set_visible(False)
ax.spines['top'].set_visible(False)
ax.xaxis.set_ticks_position('bottom')                                           
ax.yaxis.set_ticks_position('left') 

#example from http://matplotlib.org/examples/pylab_examples/annotation_demo.html
fig = plt.figure()
ax = fig.add_subplot(111, projection='polar')
r = np.arange(0, 1, 0.001)
theta = 2*2*np.pi*r
line, = ax.plot(theta, r, color='#ee8d18', lw=3)

ind = 800
thisr, thistheta = r[ind], theta[ind]
ax.plot([thistheta], [thisr], 'o')
ax.annotate('a polar annotation',
            xy=(thistheta, thisr),  # theta, radius
            xytext=(0.05, 0.05),    # fraction, fraction
            textcoords='figure fraction',
            arrowprops=dict(facecolor='black', shrink=0.05),
            horizontalalignment='left',
            verticalalignment='bottom',
            )


"""
Explanation: Annotations
adding text and annotations is easy
End of explanation
"""


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

# parametric curves
N = 100
theta = np.linspace(-4*np.pi, 4*np.pi, N)
z = np.linspace(-2, 2, N)
r = z**2 + 1

x = r*np.sin(theta)
y = r*np.cos(theta)

ax.plot(x,y,z)

fig = plt.figure()
ax = fig.gca(projection="3d")
X = np.arange(-5,5, 0.25)
Y = np.arange(-5,5, 0.25)
X, Y = np.meshgrid(X, Y)
R = np.sqrt(X**2 + Y**2)
Z = np.sin(R)

surf = ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap="coolwarm")

# we can use setp to investigate and set options here too
plt.setp(surf)
plt.setp(surf,lw=0)
plt.setp(surf, facecolor="red")


# and the view (note: most interactive backends will allow you to rotate this freely)
ax = plt.gca()
ax.azim = 90
ax.elev = 40

"""
Explanation: Surface plots
matplotlib can't deal with true 3-d data (i.e., x,y,z + a value), but it can plot 2-d surfaces and lines in 3-d.
End of explanation
"""


def to_lonlat(x,y,z):
    SMALL = 1.e-100
    rho = np.sqrt((x + SMALL)**2 + (y + SMALL)**2)
    R = np.sqrt(rho**2 + (z + SMALL)**2)
    
    theta = np.degrees(np.arctan2(rho, z + SMALL))
    phi = np.degrees(np.arctan2(y + SMALL, x + SMALL))
    
    # latitude is 90 - the spherical theta
    return (phi, 90-theta)


from mpl_toolkits.basemap import Basemap

# other projections are allowed, e.g. "ortho", moll"
map = Basemap(projection='moll', lat_0 = 45, lon_0 = 45,
              resolution = 'l', area_thresh = 1000.)

map.drawmapboundary()

map.drawmeridians(np.arange(0, 360, 15), color="0.5", latmax=90)
map.drawparallels(np.arange(-90, 90, 15), color="0.5", latmax=90) #, labels=[1,0,0,1])

# unit vectors (+x, +y, +z)
points = [(1,0,0), (0,1,0), (0,0,1)]
labels = ["+x", "+y", "+z"]

for i in range(len(points)):
    p = points[i]
    print(p)
    lon, lat = to_lonlat(p[0], p[1], p[2])
    xp, yp = map(lon, lat)
    s = plt.text(xp, yp, labels[i], color="b", zorder=10)

# draw a great circle arc between two points
lats = [0, 0]
lons = [0, 90]

map.drawgreatcircle(lons[0], lats[0], lons[1], lats[1], linewidth=2, color="r")



"""
Explanation: Plotting on a sphere
the map funcationality expects stuff in longitude and latitude, so if you want to plot x,y,z on the surface of a sphere using the idea of spherical coordinates, remember that the spherical angle from z (theta) is co-latitude
note: you need the python-basemap package installed for this to work
This also illustrates getting access to a matplotlib toolkit
End of explanation
"""


map = Basemap(projection='ortho', lat_0 = 45, lon_0 = 45,
              resolution = 'l', area_thresh = 1000.)

map.drawcoastlines()
map.drawmapboundary()

"""
Explanation: also, if you really are interested in earth...
End of explanation
"""


N = 10000
r = np.random.randn(N)
plt.hist(r, normed=True, bins=20)

x = np.linspace(-5,5,200)
sigma = 1.0
plt.plot(x, np.exp(-x**2/(2*sigma**2))/(sigma*np.sqrt(2.0*np.pi)),
         c="r", lw=2)
plt.xlabel("x")


"""
Explanation: Histograms
here we generate a bunch of gaussian-normalized random numbers and make a histogram.  The probability distribution should match
$$y(x) = \frac{1}{\sigma \sqrt{2\pi}} e^{-x^2/(2\sigma^2)}$$
End of explanation
"""


data = np.loadtxt("test1.exact.128.out")
print(data.shape)

plt.plot(data[:,1], data[:,2]/np.max(data[:,2]), label=r"$\rho$")
plt.plot(data[:,1], data[:,3]/np.max(data[:,3]), label=r"$u$")
plt.plot(data[:,1], data[:,4]/np.max(data[:,4]), label=r"$p$")
plt.plot(data[:,1], data[:,5]/np.max(data[:,5]), label=r"$T$")
plt.ylim(0,1.1)
plt.legend(frameon=False, loc="best", fontsize=12)

"""
Explanation: Plotting data from a file
numpy.loadtxt() provides an easy way to read columns of data from an ASCII file
End of explanation
"""


from ipywidgets import interact

plt.figure()
x = np.linspace(0,1,100)
def plotsin(f):
    plt.plot(x, np.sin(2*np.pi*x*f))
    plt.show()

interact(plotsin, f=(1,10,0.1))

# interactive histogram
def hist(N, sigma):
    r = sigma*np.random.randn(N)
    plt.hist(r, normed=True, bins=20)

    x = np.linspace(-5,5,200)
    
    plt.plot(x, np.exp(-x**2/(2*sigma**2))/(sigma*np.sqrt(2.0*np.pi)),
             c="r", lw=2)
    plt.xlabel("x")
    plt.show()
    
interact(hist, N=(100,10000,10), sigma=(0.5,5,0.1))

"""
Explanation: Interactivity
matplotlib has it's own set of widgets that you can use, but recently, Jupyter / Ipython gained the interact() function
(see http://ipywidgets.readthedocs.io/en/latest/examples/Using%20Interact.html )
Note: something changed in mpl 2.0 that we now need a plt.show() here.  See:
https://github.com/jupyter-widgets/ipywidgets/issues/1179
End of explanation
"""


plt.xkcd()

"""
Explanation: Final fun
if you want to make things look hand-drawn in the style of xkcd, rerun these examples after doing
plt.xkcd()
End of explanation
"""


