python - 与原点距离的颜色点

Question

这是使用 PyVista 在 Python 中创建的 Hopf 圆环：

import numpy as np
import pyvista as pv

A = 0.44
n = 3
def Gamma(t):
    alpha = np.pi/2 - (np.pi/2-A)*np.cos(n*t)
    beta = t + A*np.sin(2*n*t)
    return np.array([
      np.sin(alpha) * np.cos(beta),
      np.sin(alpha) * np.sin(beta),
      np.cos(alpha)
    ])

def HopfInverse(p, phi):
    return np.array([
      (1+p[2])*np.cos(phi),
      p[0]*np.sin(phi) - p[1]*np.cos(phi), 
      p[0]*np.cos(phi) + p[1]*np.sin(phi),
      (1+p[2])*np.sin(phi)
    ]) / np.sqrt(2*(1+p[2]))

def Stereo(q):
    return 2*q[0:3] / (1-q[3])

def F(t, phi):
    return Stereo(HopfInverse(Gamma(t), phi))

angle = np.linspace(0, 2*np.pi, 300)
angle2 = np.linspace(0, np.pi, 150)
theta, phi = np.meshgrid(angle, angle2)
x, y, z = F(theta, phi)

# Display the mesh
grid = pv.StructuredGrid(x, y, z)
grid.plot(smooth_shading=True)

我想在这个表面上添加一个调色板。圆环以原点 (0,0,0) 为中心。我想根据到原点的距离获得一种颜色。

使用Matplotlib，我会：

from mpl_toolkits.mplot3d import Axes3D
import matplotlib.colors as mcolors
from matplotlib import cm
import matplotlib.pyplot as plt
import numpy as np

A = 0.44
n = 3

......

colorfunction = (X**2+Y**2+Z**2)
norm = mcolors.Normalize(colorfunction.min(),colorfunction.max())

# Display the mesh
fig = plt.figure()
ax = fig.gca(projection = '3d')
ax.plot_surface(z, x, y, rstride = 1, cstride = 1, facecolors=cm.jet(norm(colorfunction)))
plt.show()

---

编辑

我有一个解决方案，但我不控制颜色：

grid = pv.StructuredGrid(x, y, z)
grid['Data'] = grid.points
grid.plot(smooth_shading=True, scalars="Data")

Answer 1

作为旁注，至少对我来说，自己计算点的大小并将其设置为标量更清楚（而不是依赖矢量数据的大小作为颜色映射的标量，即使这是支持且有效的）。

您缺少的只是colourmap 的选择。与 matplotlib 一样，默认值是viridis. 相反，它似乎是你想要jet的（尽管我建议不要这样做；在大多数情况下，对于数据可视化来说，感知统一的颜色图更可取）：

import numpy as np
import pyvista as pv

A = 0.44
n = 3
def Gamma(t):
    alpha = np.pi/2 - (np.pi/2-A)*np.cos(n*t)
    beta = t + A*np.sin(2*n*t)
    return np.array([
      np.sin(alpha) * np.cos(beta),
      np.sin(alpha) * np.sin(beta),
      np.cos(alpha)
    ])

def HopfInverse(p, phi):
    return np.array([
      (1+p[2])*np.cos(phi),
      p[0]*np.sin(phi) - p[1]*np.cos(phi),
      p[0]*np.cos(phi) + p[1]*np.sin(phi),
      (1+p[2])*np.sin(phi)
    ]) / np.sqrt(2*(1+p[2]))

def Stereo(q):
    return 2*q[0:3] / (1-q[3])

def F(t, phi):
    return Stereo(HopfInverse(Gamma(t), phi))

angle = np.linspace(0, 2 * np.pi, 300)
theta, phi = np.meshgrid(angle, angle)
x, y, z = F(theta, phi)
grid = pv.StructuredGrid(x, y, z)

# convert to PolyData and clean to remove the seam
cleaned_poly = grid.extract_geometry().clean(tolerance=1e-6)

# add distance from origin as scalars
cleaned_poly.point_data['distance'] = np.linalg.norm(cleaned_poly.points, axis=1)
# this also makes these the default scalars

cleaned_poly.plot(smooth_shading=True, cmap='jet')  # but don't use jet if possible