Search code examples
pythonmatplotlibvisualization

How to fix alignment of projection from (x,y,z) coordinates onto xy-plane in matplotlib 3d plot?

I was trying to make a 3D visualization of the joint probability mass function with the following code:

import math
import numpy as np
import matplotlib.pyplot as plt

def f(x, y):
    if(1 <= x + y and x + y <= 4):
        return (math.comb(3, x) * math.comb(2, y) * math.comb(3, 4 - x - y)) / math.comb(8, 4)
    else:
        return 0.0

x_domain = np.array([0, 1, 2, 3])
y_domain = np.array([0, 1, 2])

X, Y = np.meshgrid(x_domain, y_domain)
X = np.ravel(X)
Y = np.ravel(Y)

Z = np.zeros_like(X, dtype=float)
for i in range(len(X)):
    Z[i] = f(X[i], Y[i])

fig = plt.figure(figsize=(8, 4))
ax = fig.add_subplot(1, 1, 1, projection="3d")

ax.scatter(X, Y, Z) # plots the induvidual points
for i in range(len(X)):  # draws lines from xy plane up (x,y,z)
    ax.plot([X[i], X[i]], [Y[i], Y[i]], [0, Z[i]], color="r")

ax.set_xticks(x_domain)  
ax.set_yticks(y_domain)  
plt.show()

Which gave the following result:

3D scatter plot

As you can see, the stems of lines from the xy-plane do not align with the integer coordinates. I have searched through the matplotlib documentation and cant find the cause for this (unless I have overlooked something). Does anyone know how to fix the alignment issue?

EDIT It seems stack overflow could not render the mathjax equation, so I removed it.

Solution

  • Actually, they DO align with the integer grid, as you will see by rotating your plot a bit.

    It's just that the z axis doesn't naturally start from 0, which is where the end of your whatever-they-are are.

    Add the line

    ax.set_zlim( 0.0, 0.25 )

    just before plt.show() and you will be good to go.

    enter image description here