Plotting 3-D Solid of Revolution in Python Using Matplotlib - python-3.x

Question

Given the region bounded by the curves y=x^2, y=(x-2)^2 and the axis.
I want to plot the 3-D solid rotated about the x-axis.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
# Define the function to rotate
def f(x):
return x**2
def g(x):
return (x-2)**2
# Define the range of x values to plot
x = np.linspace(0, 1, 100)
x2=np.linspace(1, 2, 100)
# Define the range of angles to rotate over
theta = np.linspace(0, 2*np.pi, 100)
# Create a meshgrid of x and theta values
X, Theta = np.meshgrid(x, theta)
X2, Theta = np.meshgrid(x2, theta)
# Calculate the corresponding cylindrical coordinates
R = X
Y = R*np.sin(Theta)
Z = R*np.cos(Theta)*f(X)
R2 = X2
Y2 = R2*np.sin(Theta)
Z2 = R2*np.cos(Theta)*g(X2)
# Create the 3D plot
fig = plt.figure(figsize = (11,8))
ax = fig.add_subplot(111, projection='3d')
ax.plot_surface(X, Y, Z)
ax.plot_surface(X2, Y2, Z2)
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
plt.show()
Output:
As you can see, it works fine for the first curve y = x^2 (blue) but it's not rendering correctly for y=(x-2)^2 (orange). Why is it doing that?
The code and output attached above.

Answer 1

I used a trick to make the plotting process easier.
Instead of rotating around the x-axis, it is much easier rotating around the z-axis using spherical coordinates. matplotlib has intuitive example of utilizing spherical coordinates to draw a ball. Hence, we can swap the axis (e.g. treat the x-axis in the 2D plot as the z-axis in the 3D plot), compute the required spherical coordinates from the given two functions, and then convert back to Cartesian for plotting.
Since we swap the coordinates, eventually we have to rotate the plot and manually assign the axis label.
import matplotlib.pyplot as plt
import numpy as np
from typing import Tuple, Callable
fig = plt.figure(figsize=(10, 10))
ax = fig.add_subplot(projection='3d')
# Define the function to rotate
def f(x):
return x**2
def g(x):
return (x-2)**2
def get_p(phi: np.ndarray, f: Callable, x0: float = 0) -> np.ndarray:
"""Get the distance p
Let the origin be O and a line starting from O with its angle relative to
x-axis being phi intersect with the curve y = f(x) at point Q, the distance
p is the length of the line segment OQ.
:param phi: the angle relative to x-axis
:type phi: np.ndarray
:param f: the curve to be rotated around its x-axis
:type f: Callable
:param x0: starting estimate of x-coord of intersection Q. Use this to
control which intersection is desired. default to 0
:type x0: optional, float
:return: an array of distance, corresponding to each given phi
:rtype: np.ndarray
"""
ks = np.tan(phi)
x = []
for k in ks:
func = lambda x : f(x) - k * x
# we only look for one root
x.append(scipy.optimize.fsolve(func, x0)[0])
x = np.array(x)
y = x * ks
return np.sqrt(x**2 + y**2)
def get_xyz(
theta: np.ndarray, phi: np.ndarray, p: np.ndarray,
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
"""Produce the Cartesian coordinates from the given spherical coordinates.
For reference, see: https://mathinsight.org/spherical_coordinates#:~:text=In%20summary%2C%20the%20formulas%20for,%CE%B8z%3D%CF%81cos%CF%95.
:param theta: in the 3D coordinate, given its origin O, a point P and its
projection Q on the XY plane, theta is the angle between line segment
OQ and positive x-axis.
:type theta: np.ndarray
:param phi: using the same setup as described above, phi is the angle
between line segment OP and positive z-axis
:type phi: np.ndarray
:param p: using the same setup as described above, p is the length of line
segment OP.
:type p: np.ndarray
:return: the Cartesian coordinates converted from the spherical coordinates
in the form of (x, y, z)
:rtype: Tuple[np.ndarray, np.ndarray, np.ndarray]
"""
return (
np.outer(np.cos(theta), np.sin(phi) * p),
np.outer(np.sin(theta), np.sin(phi) * p),
np.outer(np.ones(np.size(theta)), np.cos(phi) * p),
)
# Make data
theta = np.linspace(0, 2 * np.pi, 100)
phi_intercept = np.pi / 4 # the angle relative to x-axis when the two curves meet
# Plot y = x^2 half
phi2 = np.linspace(0, phi_intercept, 50)
p2 = get_p(phi2, f, x0=1)
ax.plot_surface(*get_xyz(theta, phi2, p2))
# Plot y = (x - 2)^2 half
phi1 = np.linspace(0, phi_intercept, 50)
p1 = get_p(phi1, g, x0=1)
ax.plot_surface(*get_xyz(theta, phi1, p1))
# Set plot properties
ax.set_box_aspect([1,1,1])
# x axis in the 2D plot becomes z here
ax.set_zlim(0, 2)
ax.set_zlabel('X')
# y axis in the 2D plot is still y here
ax.set_ylim(-1, 1)
ax.set_ylabel('Y')
# the new z axis after rotation becomes x here
ax.set_xlim(-1, 1)
ax.set_xlabel('Z')
# rotate the plot
ax.view_init(10, 0, -90)
plt.savefig('demo.png', dpi=100)