Цветовые коды RGB в 3D колчане в матплотлибе - PullRequest
Новая
       53

Цветовые коды RGB в 3D колчане в матплотлибе

0 голосов
/ 10 октября 2018

Я пытаюсь написать функцию для построения векторного поля единичных векторов.Данные записываются в снимок переменной, который представляет собой массив 4d, в котором поле сохраняется в угловых координатах (тета, фи).После преобразования данных в декартовы координаты я хотел бы построить их, используя значения 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)

1 Ответ

0 голосов
/ 10 октября 2018

После некоторых попыток и использования предложения 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)
...