CODEWITHSUNDEEP
pythonmatplotlibscientific-computing
skdys
skdys

Reputation: 221

RGB color codes in 3D quiver in matplotlib

I am trying to write a function to plot a vector field of unit vectors. The data is recorded in the variable snapshot which is a 4d array in which the field is saved in angular coordinates (theta, phi). After having converted the data in cartesian coordinates, I would like to plot them using the values of u,v,w as an RGB color code in order to clearly distinguish the zone where the vectors are aligned.

I have written this code but I am getting an error because the color array is not recognized. How can I fix it? I have seen other questions similar to mine but I am unable to understand how to apply the solution to my case.

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)

Upvotes: 0

Views: 782

Answers (1)

skdys
skdys

Reputation: 221

After some trying and using the suggestion of ImportanceOfBeingErnest I've solved the problem in the following way.

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)

Upvotes: 1

Related Questions