We Keep Coding

Matplotlib: Coloring scatter plot by density relative to another data set

I'm new to Python and having some trouble with matplotlib. I currently have data that is contained in two numpy arrays, call them x and y, that I am plotting on a scatter plot with coordinates for each point (x, y) (i.e I have points x[0], y[0] and x1, y1 and so on on my plot). I have been using the following code segment to color the points in my scatter plot based on the spatial density of nearby points (found this on another stackoverflow post):
http://prntscr.com/abqowk
import numpy as np
import matplotlib.pyplot as plt
from scipy.stats import gaussian_kde
x = np.random.normal(size=1000)
y = x*3 + np.random.normal(size=1000)
xy = np.vstack([x,y])
z = gaussian_kde(xy)(xy)
idx = z.argsort()
fig,ax = plt.subplots()
ax.scatter(x,y,c=z,s=50,edgecolor='')
plt.show()
Output:
I've been using it without being sure exactly how it works (namely the point density calculation - if someone could explain how exactly that works, would also be much appreciated).
However, now I'd like to color code by the ratio of the spatial density of points in x,y to that of the spatial density of points in another set of numpy arrays, call them x2, y2. That is, I would like to make a plot such that I can identify how the density of points in x,y compares to the points in x2,y2 on the same scatter plot. Could someone please explain how I could go about doing this?
Thanks in advance for your help!

I've been trying to do the same thing based on that same earlier post, and I think I just figured it out! The trick is to use matplotlib.colors.Normalize() to define a scale and then weight it according to some data set (xnorm,ynorm):
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.colors as mplc
import matplotlib.cm as cm
from scipy.stats import gaussian_kde
def kdeplot(x,y,xnorm,ynorm):
xy = np.vstack([x,y])
z = gaussian_kde(xy)(xy)
wt = 1.0*len(x)/(len(xnorm)*1.0)
norm = mplc.Normalize(vmin=0, vmax=8/wt)
cmap = cm.gnuplot
idx = z.argsort()
x, y, z = x[idx], y[idx], z[idx]
args = (x,y)
kwargs = {'c':z,'s':10,'edgecolor':'','cmap':cmap,'norm':norm}
return args, kwargs
# (x1,y1) is some data set whose density map coloring you
# want to scale to (xnorm,ynorm)
args,kwargs = kdeplot(x1,y1,xnorm,ynorm)
plt.scatter(*args,**kwargs)
I used trial and error to optimize my normalization for my particular data and choice of colormap. Here's what my data looks like scaled to itself; here's my data scaled to some comparison data (which is on the bottom of that image).
I'm not sure this method is entirely general, but it works in my case: I know that my data and the comparison data are in similar regions of parameter space, and they both have gaussian scatter, so I can use a naive linear scaling determined by the number of data points and it results in something that gives the right idea visually.

Basic scatter plot with reference data on diagonal (identity line)

I have two arrays x,y obtained from a machine learning calculations and I wish to make a scatter plot with the reference data x on the diagonal in a way to visualize better the predicted values y against the true ones x. Please can you suggest me how to do it in python or gnuplot?

import numpy as np
import matplotlib.pyplot as plt
N = 50
x = np.random.rand(N)
y = np.random.rand(N)
colors = np.random.rand(N)
plt.scatter(x, y, c=colors)
plt.plot( [0,1],[0,1] )
plt.savefig('a.png')
This will produce:
Check this page for more information.

a simple example:
import matplotlib.pyplot as plt
import numpy as np
x=np.linspace(0,100,101)
y=np.random.normal(x) # add some noise
plt.plot(x,y,'r.') # x vs y
plt.plot(x,x,'k-') # identity line
plt.xlim(0,100)
plt.ylim(0,100)
plt.show()

In matplotlib, you can also draw an "infinite" line in order to avoid having to define the exact coordinates. For example, if you have an axes ax, you can do:
pt = (0, 0)
ax.axline(pt, slope=1, color='black')
where pt is an intersection point. Note if pt isn't included in the limits of the plot, the limits will be modified to include it.

Waterfall plot python?

Is there a python module that will do a waterfall plot like MATLAB does? I googled 'numpy waterfall', 'scipy waterfall', and 'matplotlib waterfall', but did not find anything.

You can do a waterfall in matplotlib using the PolyCollection class. See this specific example to have more details on how to do a waterfall using this class.
Also, you might find this blog post useful, since the author shows that you might obtain some 'visual bug' in some specific situation (depending on the view angle chosen).
Below is an example of a waterfall made with matplotlib (image from the blog post):
(source: austringer.net)

Have a look at mplot3d:
# copied from
# http://matplotlib.sourceforge.net/mpl_examples/mplot3d/wire3d_demo.py
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
import numpy as np
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
X, Y, Z = axes3d.get_test_data(0.05)
ax.plot_wireframe(X, Y, Z, rstride=10, cstride=10)
plt.show()
I don't know how to get results as nice as Matlab does.
If you want more, you may also have a look at MayaVi: http://mayavi.sourceforge.net/

The Wikipedia type of Waterfall chart one can obtain also like this:
import numpy as np
import pandas as pd
def waterfall(series):
df = pd.DataFrame({'pos':np.maximum(series,0),'neg':np.minimum(series,0)})
blank = series.cumsum().shift(1).fillna(0)
df.plot(kind='bar', stacked=True, bottom=blank, color=['r','b'])
step = blank.reset_index(drop=True).repeat(3).shift(-1)
step[1::3] = np.nan
plt.plot(step.index, step.values,'k')
test = pd.Series(-1 + 2 * np.random.rand(10), index=list('abcdefghij'))
waterfall(test)

I have generated a function that replicates the matlab waterfall behaviour in matplotlib. That is:
It generates the 3D shape as many independent and parallel 2D curves
Its color comes from a colormap in the z values
I started from two examples in matplotlib documentation: multicolor lines and multiple lines in 3d plot. From these examples, I only saw possible to draw lines whose color varies following a given colormap according to its z value following the example, which is reshaping the input array to draw the line by segments of 2 points and setting the color of the segment to the z mean value between these 2 points.
Thus, given the input matrixes n,m matrixes X,Y and Z, the function loops over the smallest dimension between n,m to plot each of the waterfall plot independent lines as a line collection of the 2 points segments as explained above.
def waterfall_plot(fig,ax,X,Y,Z,**kwargs):
'''
Make a waterfall plot
Input:
fig,ax : matplotlib figure and axes to populate
Z : n,m numpy array. Must be a 2d array even if only one line should be plotted
X,Y : n,m array
kwargs : kwargs are directly passed to the LineCollection object
'''
# Set normalization to the same values for all plots
norm = plt.Normalize(Z.min().min(), Z.max().max())
# Check sizes to loop always over the smallest dimension
n,m = Z.shape
if n>m:
X=X.T; Y=Y.T; Z=Z.T
m,n = n,m
for j in range(n):
# reshape the X,Z into pairs
points = np.array([X[j,:], Z[j,:]]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
# The values used by the colormap are the input to the array parameter
lc = LineCollection(segments, cmap='plasma', norm=norm, array=(Z[j,1:]+Z[j,:-1])/2, **kwargs)
line = ax.add_collection3d(lc,zs=(Y[j,1:]+Y[j,:-1])/2, zdir='y') # add line to axes
fig.colorbar(lc) # add colorbar, as the normalization is the same for all
# it doesent matter which of the lc objects we use
ax.auto_scale_xyz(X,Y,Z) # set axis limits
Therefore, plots looking like matlab waterfall can be easily generated with the same input matrixes as a matplotlib surface plot:
import numpy as np; import matplotlib.pyplot as plt
from matplotlib.collections import LineCollection
from mpl_toolkits.mplot3d import Axes3D
# Generate data
x = np.linspace(-2,2, 500)
y = np.linspace(-2,2, 60)
X,Y = np.meshgrid(x,y)
Z = np.sin(X**2+Y**2)-.2*X
# Generate waterfall plot
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
waterfall_plot(fig,ax,X,Y,Z,linewidth=1.5,alpha=0.5)
ax.set_xlabel('X'); ax.set_ylabel('Y'); ax.set_zlabel('Z')
fig.tight_layout()
The function assumes that when generating the meshgrid, the x array is the longest, and by default the lines have fixed y, and its the x coordinate what varies. However, if the size of the y array is longer, the matrixes are transposed, generating the lines with fixed x. Thus, generating the meshgrid with the sizes inverted (len(x)=60 and len(y)=500) yields:
To see what are the possibilities of the **kwargs argument, refer to the LineCollection class documantation and to its set_ methods.

Create a stack of polar plots using Matplotlib/Python

I need to generate a stack of 2D polar plots (a 3D cylindrical plot) so that I can view a distorted cylinder. I want to use matplotlib since I already have it installed and want to distribute my code to others who only have matplotlib. For example, say I have a bunch of 2-D arrays. Is there any way I can do this without having to download an external package? Here's my code.
#!usr/bin/env python
import matplotlib.pyplot as plt
import numpy as np
x = np.arange(-180.0,190.0,10)
theta = (np.pi/180.0 )*x # in radians
A0 = 55.0
offset = 60.0
R = [116.225,115.105,114.697,115.008,115.908,117.184,118.61,119.998,121.224,122.216,\
122.93,123.323,123.343,122.948,122.134,120.963,119.575,118.165,116.941,116.074,115.66\
,115.706,116.154,116.913,117.894,119.029,120.261,121.518,122.684,123.594,124.059,\
123.917,123.096,121.661,119.821,117.894,116.225]
fig = plt.figure()
ax = fig.add_axes([0.1,0.1,0.8,0.8],polar=True) # Polar plot
ax.plot(theta,R,lw=2.5)
ax.set_rmax(1.5*(A0)+offset)
plt.show()
I have 10 more similar 2D polar plots and I want to stack them up nicely. If there's any better way to visualize a distorted cylinder in 3D, I'm totally open to suggestions. Any help would be appreciated. Thanks!

If you want to stack polar charts using matplotlib, one approach is to use the Axes3D module. You'll notice that I used polar coordinates first and then converted them back to Cartesian when I was ready to plot them.
from numpy import *
from mpl_toolkits.mplot3d import Axes3D
import matplotlib.pyplot as plt
n = 1000
fig = plt.figure()
ax = fig.gca(projection='3d')
for k in linspace(0, 5, 5):
THETA = linspace(0, 2*pi, n)
R = ones(THETA.shape)*cos(THETA*k)
# Convert to Cartesian coordinates
X = R*cos(THETA)
Y = R*sin(THETA)
ax.plot(X, Y, k-2)
plt.show()
If you play with the last argument of ax.plot, it controls the height of each slice. For example, if you want to project all of your data down to a single axis you would use ax.plot(X, Y, 0). For a more exotic example, you can map the height of the data onto a function, say a saddle ax.plot(X, Y, -X**2+Y**2 ). By playing with the colors as well, you could in theory represent multiple 4 dimensional datasets (though I'm not sure how clear this would be). Examples below:

how do you radially 'sweep out' a 1D array to plot 3d figure in python? (to represent a wavefunction)

effectively I have a large 1D array of heights. As a small example consider:
u=array([0,1,2,1,0,2,4,6,4,2,1])
and a 1D array, the same size as u, of radial values which the heights correspond to, e.g.:
r=array([0,1,2,3,4,5,6,7,8,9,10])
Obviously plotting these with:
pylab.plot(r,u)
gives a nice 2D plot.
How can one sweep this out around 360 degrees, to give a 3D contour/surface plot?
If you can imagine it should look like a series of concentric, circular ridges, like for the wavefunction of an atom.
any help would be much appreciated!

You're better off with something more 3D oriented than matplotlib, in this case...
Here's a quick example using mayavi:
from enthought.mayavi import mlab
import numpy as np
# Generate some random data along a straight line in the x-direction
num = 100
x = np.arange(num)
y, z = np.ones(num), np.ones(num)
s = np.cumsum(np.random.random(num) - 0.5)
# Plot using mayavi's mlab api
fig = mlab.figure()
# First we need to make a line source from our data
line = mlab.pipeline.line_source(x,y,z,s)
# Then we apply the "tube" filter to it, and vary the radius by "s"
tube = mlab.pipeline.tube(line, tube_sides=20, tube_radius=1.0)
tube.filter.vary_radius = 'vary_radius_by_scalar'
# Now we display the tube as a surface
mlab.pipeline.surface(tube)
# And finally visualize the result
mlab.show()

#!/usr/bin/python
from mpl_toolkits.mplot3d import Axes3D
import matplotlib
import numpy as np
from scipy.interpolate import interp1d
from matplotlib import cm
from matplotlib import pyplot as plt
step = 0.04
maxval = 1.0
fig = plt.figure()
ax = Axes3D(fig)
u=np.array([0,1,2,1,0,2,4,6,4,2,1])
r=np.array([0,1,2,3,4,5,6,7,8,9,10])
f=interp1d(r,u)
# walk along the circle
p = np.linspace(0,2*np.pi,50)
R,P = np.meshgrid(r,p)
# transform them to cartesian system
X,Y = R*np.cos(P),R*np.sin(P)
Z=f(R)
ax.plot_surface(X, Y, Z, rstride=1, cstride=1, cmap=cm.jet)
ax.set_xticks([])
plt.show()