Plotting a 3-dimensional superball shape

I'm trying to plot a 3D superball in python matplotlib, where a superball is defined as a general mathematical shape that can be used to describe rounded cubes using a shape parameter p, where for p = 1 the shape is equal to that of a sphere.

This paper claims that the superball is defined by using modified spherical coordinates with:

x = r*cos(u)**1/p * sin(v)**1/p
y = r*cos(u)**1/p * sin(v)**1/p
z = r*cos(v)**1/p

with u = phi and v = theta.

I managed to get the code running, at least for p = 1 which generates a sphere - exactly as it should do:

import matplotlib.pyplot as plt
import numpy as np


fig = plt.figure()
ax = fig.add_subplot(projection='3d')

r, p = 1, 1

# Make data
u = np.linspace(0, 2 * np.pi, 100)
v = np.linspace(0, np.pi, 100)
u, v = np.meshgrid(u, v)

x = r * np.cos(u)**(1/p) * np.sin(v)**(1/p)
y = r * np.sin(u)**(1/p) * np.sin(v)**(1/p)
z = r * np.cos(v)**(1/p)

# Plot the surface
ax.plot_surface(x, y, z)

plt.show()

This is a 3D plot of the code above for p = 1.

However, as I put in any other value for p, e.g. 2, it's giving me only a partial shape, while it should actually give me a full superball.

This is a 3D plot of the code above for p = 2.

I believe the fix is more of mathematical nature, but how can this be fixed?

When plotting a regular sphere, we transform positive and negative coordinates differently:

• Positives: x**0.5
• Negatives: -1 * abs(x)**0.5

For the superball variants, apply the same logic using np.sign and np.abs:

power = lambda base, exp: np.sign(base) * np.abs(base)**exp

x = r * power(np.cos(u), 1/p) * power(np.sin(v), 1/p)
y = r * power(np.sin(u), 1/p) * power(np.sin(v), 1/p)
z = r * power(np.cos(v), 1/p)

Full example for p = 4:

import matplotlib.pyplot as plt
import numpy as np

fig, ax = plt.subplots(subplot_kw={'projection': '3d'})

r, p = 1, 4

# Make the data
u = np.linspace(0, 2 * np.pi)
v = np.linspace(0, np.pi)
u, v = np.meshgrid(u, v)

# Transform the coordinates
#   Positives: base**exp
#   Negatives: -abs(base)**exp
power = lambda base, exp: np.sign(base) * np.abs(base)**exp
x = r * power(np.cos(u), 1/p) * power(np.sin(v), 1/p)
y = r * power(np.sin(u), 1/p) * power(np.sin(v), 1/p)
z = r * power(np.cos(v), 1/p)

# Plot the surface
ax.plot_surface(x, y, z)
plt.show()