TMorville TMorville - 4 months ago 99
Python Question

Python/Matplotlib: 2d random walk with kde joint density contour in a 3d plot

I'm struggling with creating a quite complex 3d figure in python, specifically using iPython notebook. I can partition the content of the graph into two sections:

The (x,y) plane: Here a two-dimensional random walk is bobbing around, let's call it G(). I would like to plot part of this trajectory on the (x,y) plane. Say, 10% of all the data points of G(). As G() bobs around, it visits some (x,y) pairs more frequently than others. I would like to estimate this density of G() using a kernel estimation approach and draw it as contour lines on the (x,y) plane.

The (z) plane: Here, I would like to draw a mesh or (transparent) surface plot of the information theoretical surprise of a bivariate normal. Surprise is simply -log(p(i)) or the negative (base 2) logarithm of outcome i. Given the bivariate normal, each (x,y) pair has some probability p(x,y) and the surprise of this is simply -log(p(x,y)).

Essentially these two graphs are independent. Assume the interval of the random walk G() is [xmin,xmax],[ymin,ymax] and of size N. The bivariate normal in the z-plane should be drawn from the same interval, such that for each (x,y) pair in the random walk, I can draw a (dashed) line from some subset of the random walk n < N to the bivariate normal. Assume that G(10) = (5,5) then I would like to draw a dashed line from (5,5) up the Z-axes, until it hits the bivariate normal.

So far, I've managed to plot G() in a 3-d space, and estimate the density f(X,Y) using scipy.stats.gaussian_kde. In another (2d) graph, I have the sort of contour lines I want. What I don't have, is the contour lines in the 3d-plot using the estimated KDE density. I also don't have the bivariate normal plot, or the projection of a few random points from the random walk, to the surface of the bivariate normal. I've added a hand drawn figure, which might ease intuition (ignore the label on the z-axis and the fact that there is no mesh.. difficult to draw!)

Any input, even just partial, such as how to draw the contour lines in the (x,y) plane of the 3d graph, or a mesh of a bivariate normal would be much appreciated.

Thanks!

import matplotlib as mpl
import matplotlib.pyplot as plt
import random
import numpy as np
import seaborn as sns
import scipy
from mpl_toolkits.mplot3d import Axes3D
%matplotlib inline

def randomwalk():

mpl.rcParams['legend.fontsize'] = 10

xyz = []
cur = [0, 0]

for _ in range(400):
axis = random.randrange(0, 2)
cur[axis] += random.choice([-1, 1])
xyz.append(cur[:])


x, y = zip(*xyz)
data = np.vstack([x,y])
kde = scipy.stats.gaussian_kde(data)
density = kde(data)

fig1 = plt.figure()

ax = fig1.gca(projection='3d')
ax.plot(x, y, label='Random walk')
sns.kdeplot(data[0,:], data[1,:], 0)
ax.scatter(x[-1], y[-1], c='b', marker='o') # End point
ax.legend()

fig2 = plt.figure()
sns.kdeplot(data[0,:], data[1,:])


Calling randomwalk() initialises and plots this:

enter image description here

enter image description here

Edit #1:

Made some progress, actually the only thing I need is to restrict the height of the dashed vertical lines to the bivariate. Any ideas?

import matplotlib as mpl
import matplotlib.pyplot as plt
import random
import numpy as np
import seaborn as sns
import scipy
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.mlab import bivariate_normal

%matplotlib inline

# Data for random walk

def randomwalk():

mpl.rcParams['legend.fontsize'] = 10

xyz = []
cur = [0, 0]

for _ in range(40):
axis = random.randrange(0, 2)
cur[axis] += random.choice([-1, 1])
xyz.append(cur[:])

# Get density

x, y = zip(*xyz)
data = np.vstack([x,y])
kde = scipy.stats.gaussian_kde(data)
density = kde(data)


# Data for bivariate gaussian

a = np.linspace(-7.5, 7.5, 20)
b = a
X,Y = np.meshgrid(a, b)
Z = bivariate_normal(X, Y)
surprise_Z = -np.log(Z)

# Get random points from walker and plot up z-axis to the gaussian

M = data[:,np.random.choice(20,5)].T

# Plot figure

fig = plt.figure(figsize=(10, 7))

ax = fig.gca(projection='3d')
ax.plot(x, y, 'grey', label='Random walk') # Walker
ax.scatter(x[-1], y[-1], c='k', marker='o') # End point

ax.legend()

surf = ax.plot_surface(X, Y, surprise_Z, rstride=1, cstride=1,
cmap = plt.cm.gist_heat_r, alpha=0.1, linewidth=0.1)

#fig.colorbar(surf, shrink=0.5, aspect=7, cmap=plt.cm.gray_r)

for i in range(5):
ax.plot([M[i,0], M[i,0]],[M[i,1], M[i,1]], [0,10],'k--',alpha=0.8, linewidth=0.5)

ax.set_zlim(0, 50)
ax.set_xlim(-10, 10)
ax.set_ylim(-10, 10)


enter image description here

Answer

Final code,

import matplotlib as mpl
import matplotlib.pyplot as plt
import random
import numpy as np
import seaborn as sns
import scipy
from mpl_toolkits.mplot3d import Axes3D
from matplotlib.mlab import bivariate_normal

%matplotlib inline

    # Data for random walk

def randomwalk():

    mpl.rcParams['legend.fontsize'] = 10

    xyz = []
    cur = [0, 0]

    for _ in range(50):
        axis = random.randrange(0, 2)
        cur[axis] += random.choice([-1, 1])
        xyz.append(cur[:])

    # Get density

    x, y = zip(*xyz)
    data = np.vstack([x,y])
    kde = scipy.stats.gaussian_kde(data)
    density = kde(data)


    # Data for bivariate gaussian 

    a = np.linspace(-7.5, 7.5, 100)
    b = a
    X,Y = np.meshgrid(a, b)
    Z = bivariate_normal(X, Y)
    surprise_Z = -np.log(Z)

    # Get random points from walker and plot up z-axis to the gaussian

    M = data[:,np.random.choice(50,10)].T

    # Plot figure

    fig = plt.figure(figsize=(10, 7))

    ax = fig.gca(projection='3d')
    ax.plot(x, y, 'grey', label='Random walk') # Walker

    ax.legend()

    surf = ax.plot_surface(X, Y, surprise_Z, rstride=1, cstride=1, 
        cmap = plt.cm.gist_heat_r, alpha=0.1, linewidth=0.1)

    #fig.colorbar(surf, shrink=0.5, aspect=7, cmap=plt.cm.gray_r)

    for i in range(10):
        x = [M[i,0], M[i,0]]
        y = [M[i,1], M[i,1]]
        z = [0,-np.log(bivariate_normal(M[i,0],M[i,1]))]
        ax.plot(x,y,z,'k--',alpha=0.8, linewidth=0.5)
        ax.scatter(x, y, z, c='k', marker='o')

enter image description here

Comments