Search code examples
pythonmatplotlibprobability-distribution

3D Probability Density Plots in Python

I am working on a data set consisting of 'pulse length' values and 5 or 6 'voltage values' corresponding to each of the pulse lengths. The first value is pulse length followed by voltages. Find the table below.

15 -56V -47V -53V -50V -50V

16 -49V -46V -52V -47V -50V

17 -50V -51V -47V -50V -49V

18 -50V -51V -48V -48V -45V

19 -49V -51V -45V -47V -52V

20 -45V -47V -50V -47V -54V

21 -46V -52V -52V -49V -54V

22 -53V -51V -53V -56V -52V

23 -52V -45V -51V -56V -53V

24 -51V -52V -54V -58V -52V

25 -56V -53V -57V -55V -53V

26 -53V -52V -55V -52V

27 -54V -49V -56V -54V

28 -52V -52V -57V -56V -53V

29 -63V -60V -54V -58V -61V

30 -59V -70V -61V

I want the X and Y axis to be pulse length and voltage and I want the Z axis to be its probability distribution. I have a 2D plot for the same using the set of 'voltage values' and its probability. enter image description here In the picture, the red graph corresponds to one pulse length and green graph corresponds to another pulse length. I tried doing the 3D plot the same way using an example of multivariate normal distribution example from stack overflow (Plot normal distribution in 3D). Since I have very little experience with 3D plots, I am unable to plot multiple surface plots on the same surface with different Y axis 'pulse length' values. The code that I tried is given below.

 import numpy as np
 import matplotlib
 import matplotlib.pyplot as plt
 from matplotlib.mlab import bivariate_normal
 from mpl_toolkits.mplot3d import Axes3D

 #Parameters to set
 mu_x = -48.8
 sigma_x = np.sqrt(6.5)

 mu_y = 0
 sigma_y = np.sqrt(16)

 #Create grid and multivariate normal
 x = range(-100,0)
 y = range(15,30)
 X, Y = np.meshgrid(x,y)
 Z = bivariate_normal(X,Y,sigma_x,sigma_y,mu_x,mu_y)


 #Make a 3D plot
 fig = plt.figure()
 ax = fig.gca(projection='3d')
 ax.plot_surface(X, Y, Z,cmap='Reds',linewidth=0, antialiased=True, 
 zorder = 0.5)

 ax.set_xlabel('Voltage')
 ax.set_ylabel('Pulse Length')
 ax.set_zlabel('Normal Distribution')
 plt.show()

I would be really thankful if someone could help me do the same for multiple pulse lengths. Thank you.

Solution

  • I don't know what kind of plot you want to achieve exactly but from what I understood, you want something like the below figure. I am putting only the relevant/modified code below. It's also not clear what variable is your pulse length. Since you have many pulse lengths, you can put the function to define mu_x, 'mu_y', Z in a for loop and plot several 3d surfaces.

    # Create grid and multivariate normal
x = np.linspace(-100, 0, 200) # Create a mesh of 200 x-points
y = np.linspace(-30, 30, 200) # Create a mesh of 200 y-points

X, Y = np.meshgrid(x,y)
Z = bivariate_normal(X,Y,sigma_x,sigma_y,mu_x,mu_y)
Z2 = bivariate_normal(X,Y,sigma_x,sigma_y,mu_x-20,mu_y+10)

fig = plt.figure(figsize=(10, 8))
ax = fig.gca(projection='3d')
ax.plot_surface(X, Y, Z,cmap='Reds',linewidth=0, antialiased=True, zorder = 0.5)
ax.plot_surface(X, Y, Z2,cmap='Blues',linewidth=0, alpha=0.5, antialiased=True, zorder = 0.5)

    Output

    enter image description here