If you have an ellipsoid specified by an arbitrary covariance matrix cov
and offset bias
, you perform a simpler version of @minillinim's answer by vectorizing the operations.
Starting with
u = np.linspace(0, 2 * np.pi, 100)
v = np.linspace(0, np.pi, 100)
Make a unit sphere
x = np.outer(np.cos(u), np.sin(v))
y = np.outer(np.sin(u), np.sin(v))
z = np.outer(np.ones_like(u), np.cos(v))
sphere = np.stack((x, y, z), axis=-1)[..., None]
Compute the standard deviation matrix, which has the same rotation as the covariance, but scales by the square root of the eigenvalues:
e, v = np.linalg.eig(cov)
s = v @ np.diag(np.sqrt(e)) @ v.T
Now transform the sphere:
ellipsoid = (s @ sphere).squeeze(-1) + bias
You can plot the result pretty much as before:
ax.plot_surface(*ellipsoid.transpose(2, 0, 1), rstride=4, cstride=4, color='b', alpha=0.75)
For reference, u
, v
have shape (100,)
, which makes x
, y
, z
into (100, 100)
arrays. sphere
is (100, 100, 3, 1)
, which makes it a 100x100 array of 3x1 vectors as far as the broadcasting @
operator is concerned. s @ sphere
has the same size, so squeezing out the last unit axis makes it suitable for broadcasting for addition with bias
. Finally, ellipsoid.transpose(2, 0, 1)
has shape (3, 100, 100)
, which can be star-expanded as three separate arrays of x-, y- and z-values into the call to plot_surface
.