Plot density function on sphere surface using plotly (python)

Question

I'm interested in plotting a real-valued function f(x,y,z)=a, where (x,y,z) is a 3D point on the sphere and a is a real number. I calculate the Cartesian coordinates of the points of the sphere as follows, but I have no clue on how to visualize the value of f on each of those points.

import plotly.graph_objects as go
import numpy as np

fig = go.Figure(layout=go.Layout(title=go.layout.Title(text=title), hovermode=False))

# Create mesh grid for spherical coordinates
phi, theta = np.mgrid[0.0:np.pi:100j, 0.0:2.0 * np.pi:100j]

# Get Cartesian mesh grid
x = np.sin(phi) * np.cos(theta)
y = np.sin(phi) * np.sin(theta)
z = np.cos(phi)

# Plot sphere surface
self.fig.add_surface(x=x, y=y, z=z, opacity=0.35)

fig.show()

I would imagine/expect/like a visualization like this

Additionally, I also have the gradient of f calculated in closed-form (i.e., for each (x,y,z) I calculate the 3D-dimensional gradient of f). Is there a way of plotting this vector field, similarly to what is shown in the figure above?

Answer 1

Here's an answer that's far from perfect, but hopefully that's enough for you to build on.

For the sphere itself, I don't know of any "shortcut" to do something like that in plotly, so my approach is simply to manually create a sphere mesh. Generating the vertices is simple, for example like you did - the slightly more tricky part is figuring out the vertex indices for the triangles (which depends on the vertex generation scheme). There are various algorithms to do that smoothly (i.e. generating a sphere with no "tip"), I hacked something crude just for the demonstration. Then we can use the Mesh3d object to display the sphere along with the intensities and your choice of colormap:

N = 100  # Sphere resolution (both rings and segments, can be separated to different constants)
theta, z = np.meshgrid(np.linspace(-np.pi, np.pi, N), np.linspace(-1, 1, N))
r = np.sqrt(1 - z ** 2)
x = r * np.cos(theta)
y = r * np.sin(theta)
x = x.ravel()
y = y.ravel()
z = z.ravel()

# Triangle indices
indices = np.arange(N * (N - 1) - 1)
i1 = np.concatenate([indices, (indices // N + 1) * N + (indices + 1) % N])
i2 = np.concatenate([indices + N, indices // N * N + (indices + 1) % N])
i3 = np.concatenate([(indices // N + 1) * N + (indices + 1) % N, indices])

# Point intensity function
def f(x, y, z):
    return (np.cos(x * 2) + np.sin(y ** 2) + np.sin(z) + 3) / 6


fig = go.Figure(data=[
    go.Mesh3d(
        x=x,
        y=y,
        z=z,
        colorbar_title='f(x, y, z)',
        colorscale=[[0, 'gold'],
                    [0.5, 'mediumturquoise'],
                    [1, 'magenta']],
        intensity = f(x, y, z),
        i = i1,
        j = i2,
        k = i3,
        name='y',
        showscale=True
    )
])

fig.show()

This yields the following interactive plot:

To add the vector field you can use the Cone plot; this requires some tinkering because when I simply draw the cones at the same x, y, z position as the sphere, some of the cones are partially or fully occluded by the sphere. So I generate another sphere, with a slightly larger radius, and place the cones there. I also played with some lighting parameters to make it black like in your example. The full code looks like this:

N = 100  # Sphere resolution (both rings and segments, can be separated to different constants)
theta, z = np.meshgrid(np.linspace(-np.pi, np.pi, N), np.linspace(-1, 1, N))
r = np.sqrt(1 - z ** 2)
x = r * np.cos(theta)
y = r * np.sin(theta)
x = x.ravel()
y = y.ravel()
z = z.ravel()

# Triangle indices
indices = np.arange(N * (N - 1) - 1)
i1 = np.concatenate([indices, (indices // N + 1) * N + (indices + 1) % N])
i2 = np.concatenate([indices + N, indices // N * N + (indices + 1) % N])
i3 = np.concatenate([(indices // N + 1) * N + (indices + 1) % N, indices])

# Point intensity function
def f(x, y, z):
    return (np.cos(x * 2) + np.sin(y ** 2) + np.sin(z) + 3) / 6


# Vector field function
def grad_f(x, y, z):
    return np.stack([np.cos(3 * y + 5 * x),
                     np.sin(z * y),
                     np.cos(4 * x - 3 * y + z * 7)], axis=1)

# Second sphere for placing cones
N2 = 50  # Smaller resolution (again rings and segments combined)
R2 = 1.05  # Slightly larger radius
theta2, z2 = np.meshgrid(np.linspace(-np.pi, np.pi, N2), np.linspace(-R2, R2, N2))
r2 = np.sqrt(R2 ** 2 - z2 ** 2)
x2 = r2 * np.cos(theta2)
y2 = r2 * np.sin(theta2)
x2 = x2.ravel()
y2 = y2.ravel()
z2 = z2.ravel()

uvw = grad_f(x2, y2, z2)

fig = go.Figure(data=[
    go.Mesh3d(
        x=x,
        y=y,
        z=z,
        colorbar_title='f(x, y, z)',
        colorscale=[[0, 'gold'],
                    [0.5, 'mediumturquoise'],
                    [1, 'magenta']],
        intensity = f(x, y, z),
        i = i1,
        j = i2,
        k = i3,
        name='y',
        showscale=True
    ),
    go.Cone(
        x=x2, y=y2, z=z2, u=uvw[:, 0], v=uvw[:, 1], w=uvw[:, 2], sizemode='absolute', sizeref=2, anchor='tail',
        lighting_ambient=0, lighting_diffuse=0, opacity=.2
    )
])

fig.show()

And yields this plot:

Hope this helps. There are a lot of tweaks to the display, and certainly better ways to construct a sphere mesh (e.g. see this article), so there should be a lot of freedom there (albeit at the cost of some work).

Good luck!