'Coloring by 2D location in 3D scatter plot

I have a 3D scatter plot of first three principal components plotted by matplotlib.pyplot. For each point plotted, I also have a separate 2x1 numpy array which depicts its location in a 2D environment. Is it possible to use the 2x1 arrays to introduce colouring by location in my 3D scatter plot?

This is an approximation of my code:

import numpy as np
import matplotlib.pyplot as plt

def plotThreeComponents():

    lv_1 = np.array([0,1,2,3,7])
    lv_2 = np.array([3,2,1,0,7])
    lv_3 = np.array([0,1,0,5,7])

    location = np.array([[5,7,0,4,7],
                         [1,1,2,4,7]])

    plt.figure()

    ax = plt.axes(projection='3d')
    ax.scatter3D(lv_1, lv_2, lv_3, c=location, s = 4)

    ax.set_xlabel('PC1')
    ax.set_ylabel('PC2')
    ax.set_zlabel('PC3')

    plt.show()

    pass

plotThreeComponents()

I also tried to add this to the function:

colorList = []
for i in range(location.shape[1]):
    colorList.append(location[:,i])

To generate a list where each element is an array of shape (2, ) and then changed

ax.scatter3D(lv_1, lv_2, lv_3, c=colorList, s = 4)

However, I am still getting an Error message stating that "'c' argument has 10 elements, which is inconsistent with 'x' and 'y' with size 5.", even if len(colorList) = 5

Column vectors of 'location' array depict coordinates in 2D space. Location and lv_1, lv_2, lv_3 are linked indirectly.

what should be 'c' in 'ax.scatter3D' so that I get coloring by 2D location?

Thank you for any suggestions.



Sources

This article follows the attribution requirements of Stack Overflow and is licensed under CC BY-SA 3.0.

Source: Stack Overflow

Solution Source