prax1telis
prax1telis

Reputation: 317

Plot a 1D gaussian distribution on a plane in 3D plot python

I have the following code and along plots it generates. My aim is to plot on the second figure (right) a 1D Gaussian Distribution on the red plane shown.

The aim of this is to show that the overlap (which represents the conditional) is a Gaussian Distribution. I am not interested in the exact variance of the distribution to be correct but just show this visually.

Is there any straightforward way to do this in python?

Thanks, P

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

#Make a 3D plot
fig = plt.figure(figsize=plt.figaspect(0.5))

################ First Plot ##############
#Parameters to set
mu_x = 0
sigma_x = np.sqrt(5)

mu_y = 0
sigma_y = np.sqrt(5)

#Create grid and multivariate normal
x = np.linspace(-10,10,500)
y = np.linspace(-10,10,500)
X, Y = np.meshgrid(x,y)
Z = bivariate_normal(X,Y,sigma_x,sigma_y,mu_x,mu_y)

# Create plane
x_p = 2
y_p = np.linspace(-10,10,500)
z_p = np.linspace(0,0.02,500)
Y_p, Z_p = np.meshgrid(y_p, z_p)


# ax = fig.gca(projection='3d')
ax = fig.add_subplot(1,2,1, projection='3d')
ax.plot_surface(X, Y, Z, cmap='viridis',linewidth=0)
ax.plot_surface(x_p, Y_p, Z_p, color='r',linewidth=0, alpha=0.5)
plt.tight_layout()

################ Second Plot ##############
x_p = 2
y_p = np.linspace(-10,10,500)
z_p = np.linspace(0,0.02,500)
Y_p, Z_p = np.meshgrid(y_p, z_p)


# ax2 = fig.gca(projection='3d')
ax2 = fig.add_subplot(1,2,2,projection='3d')
ax2.plot_surface(x_p, Y_p, Z_p, color='r',linewidth=0, alpha=0.3)
plt.show()

Code output

Upvotes: 4

Views: 1606

Answers (1)

yacola
yacola

Reputation: 3013

You can try to get the closest coordinates of X within tolerance tol to the plan x_p = 2 with np.where for instance, then use the resulting index idx_x_p as a mask to select the corresponding Y and Z values. That leads you to the following code :


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

#Parameters to set for Gaussian distribution
mu_x = 0
sigma_x = np.sqrt(5)
mu_y = 0
sigma_y = np.sqrt(5)

#Create grid and multivariate normal
x = np.linspace(-10,10,500)
y = np.linspace(-10,10,500)
X, Y = np.meshgrid(x,y)
Z = bivariate_normal(X,Y,sigma_x,sigma_y,mu_x,mu_y)

# Create plane
x_p = 2
y_p = np.linspace(-10,10,500)
z_p = np.linspace(0,0.02,500)
Y_p, Z_p = np.meshgrid(y_p, z_p)

# Finding closest idx values of X mesh to x_p
tol = 1e-4
idx_x_p = (np.where(x < x_p+tol) and np.where(x > x_p-tol))[0][0]
# Select the corresponding values of X, Y, Z (carefully switch X and Y)
x_c, y_c, z_c = Y[idx_x_p], X[idx_x_p], Z[idx_x_p]

# Plot
fig = plt.figure(figsize=plt.figaspect(0.5))
ax = fig.add_subplot(1, 1, 1, projection='3d')
ax.plot_surface(X, Y, Z, cmap='viridis',linewidth=0,zorder=0)
ax.plot_surface(x_p, Y_p, Z_p, color='r',linewidth=0, alpha=0.5,zorder=5)
ax.plot(x_c,y_c,z_c,zorder=10)

plt.tight_layout()

which show the Gaussian-shaped overlap for different x_p values. Let's say for x_p in np.linspace(-10,10,20) :

gaussian_overlap_gif

Upvotes: 6

Related Questions