Question

我正在尝试编写一个函数来绘制单位矢量的矢量场。数据记录在变量快照中，该快照是一个4d数组，其中的字段保存在角坐标（θ，phi）中。在将数据转换为笛卡尔坐标后，我想使用u，v，w的值作为RGB颜色代码来绘制它们，以便清楚地区分向量对齐的区域。

我已经编写了这段代码，但是由于无法识别颜色阵列而出现错误。我该如何解决？我看到了其他与我类似的问题，但是我不明白如何将解决方案应用于我的案子。

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

def sph2xyz(theta, phi):
    """
    Convert spherical coordinates to unit vector
    :param theta: theta angle
    :param phi: phi angle
    :return: (x, y, z) coordinates
    """
    x = np.sin(theta) * np.cos(phi)
    y = np.sin(theta) * np.sin(phi)
    z = np.cos(theta)
    return np.array([x, y, z])

def plot_state(snapshot):
    """
    Plot system state
    """
    nx = snapshot.shape[0]
    ny = snapshot.shape[1]
    nz = snapshot.shape[2]

    x, y, z = np.meshgrid(np.arange(0, nx),
                          np.arange(0, ny),
                          np.arange(0, nz))

    u = np.zeros(shape=(nx, ny, nz))
    v = np.zeros(shape=(nx, ny, nz))
    w = np.zeros(shape=(nx, ny, nz))

    for i, j, k in np.ndindex(nx, ny, nz):
        u[i, j, k], v[i, j, k], w[i, j, k] = sph2xyz(snapshot[i, j, k, 0], snapshot[i, j, k, 1])

    c = np.zeros(shape=(nx*3, ny*3, nz*3, 3))
    c[:, :, :, 0] = np.tile(u, (3, 3, 3))
    c[:, :, :, 1] = np.tile(v, (3, 3, 3))
    c[:, :, :, 2] = np.tile(w, (3, 3, 3))
    c = np.abs(c)

    fig = plt.figure()
    ax: Axes3D = fig.gca(projection='3d')
    ax.quiver(x, y, z, u, v, w, pivot='middle' , color=c)

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    plt.show()



s = np.zeros(shape=(5,5,5,2))
plot_state(s)

Answer 1

尝试并使用ImportanceOfBeingErnest的建议后，我已通过以下方式解决了该问题。

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

def sph2xyz(theta, phi):
    """
    Convert spherical coordinates to unit vector
    :param theta: theta angle
    :param phi: phi angle
    :return: (x, y, z) coordinates
    """
    x = np.sin(theta) * np.cos(phi)
    y = np.sin(theta) * np.sin(phi)
    z = np.cos(theta)
    return np.array([x, y, z])

def plot_state(snapshot):
    """
    Plot system state
    """
    nx = snapshot.shape[0]
    ny = snapshot.shape[1]
    nz = snapshot.shape[2]

    x, y, z = np.meshgrid(np.arange(0, nx),
                          np.arange(0, ny),
                          np.arange(0, nz))

    u = np.zeros(shape=(nx, ny, nz))
    v = np.zeros(shape=(nx, ny, nz))
    w = np.zeros(shape=(nx, ny, nz))

    for i, j, k in np.ndindex(nx, ny, nz):
        u[i, j, k], v[i, j, k], w[i, j, k] = sph2xyz(snapshot[i, j, k, 0], snapshot[i, j, k, 1])

    c = np.zeros(shape=(nx, ny, nz, 4))
    c[:, :, :, 0] = u
    c[:, :, :, 1] = v
    c[:, :, :, 2] = w
    c[:, :, :, 3] = np.ones(shape=(nx,ny,nz))
    c = np.abs(c)

    c2 = np.zeros(shape=(nx*ny*nz, 4))
    l = 0
    for i,j,k in np.ndindex((nx,ny,nz)):
        c2[l]=c[i,j,k]
        l+=1

    c3 = np.concatenate((c2, np.repeat(c2,2, axis=0)), axis=0)

    fig = plt.figure()
    ax: Axes3D = fig.gca(projection='3d')
    ax.quiver(x, y, z, u, v, w, pivot='middle' , color=c3)

    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')

    plt.show()


s = np.random.uniform(size=(5,5,5,2))
plot_state(s)

matplotlib中的3D颤抖中的RGB颜色代码

1 个答案: