I have a set of (xyz) points and I am plotting them in a scatter plot. This set of points evolves in time, therefore, an iterator is used to update each step. At the moment, I am able to extract the maximum elevation (in the Z direction) of my plot, in order to define the max/min values and normalize the color map. However, I do not know how to link the value of the elevation of my plot to the color map (defined previously) once the step is updated.
Right now, the code works without updating the colors. I was trying to include the variable colormap as an iterator, but I couldn't find the commands to perform this task.
#Python moduels
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import cm
import matplotlib.animation
import pandas as pd
###############################################################################
#Generate data
###############################################################################
np.random.seed(10) #To obtain the same result
n_radii = 8
n_angles = 120
# Make radii and angles spaces (radius r=0 omitted to eliminate duplication).
radii = np.linspace(0.125, 1.0, n_radii)
angles = np.linspace(0, 2*np.pi, n_angles, endpoint=False)
# Repeat all angles for each radius.
angles = np.repeat(angles[..., np.newaxis], n_radii, axis=1)
# Convert polar (radii, angles) coords to cartesian (x, y) coords.
# (0, 0) is manually added at this stage, so there will be no duplicate
# points in the (x, y) plane.
x = np.append(0, (radii*np.cos(angles)).flatten())
y = np.append(0, (radii*np.sin(angles)).flatten())
rows = len(x)
dim = 3
samples = 10
data = np.zeros((rows,dim,samples)) # (20 xyz) points stored in 10 sets
for i in range(samples):
data[:,0,i] = x
data[:,1,i] = y
# Compute z to make the pringle surface.
data[:,2,i] = np.sin(-i*x*y)
index = np.arange(0,samples,1)
Frame = np.array([np.ones(rows)*index[i] for i in range(samples)]).flatten()
a = np.zeros((len(Frame),3))
for i in range(samples):
a[i*rows:(i+1)*rows,0] = data[:,0,i]
a[i*rows:(i+1)*rows,1] = data[:,1,i]
a[i*rows:(i+1)*rows,2] = data[:,2,i]
df = pd.DataFrame({"Frame": Frame ,"x" : a[:,0], "y" : a[:,1], "z" : a[:,2]})
##############################################################################
#SCATTER PLOT
################################################################################
#Function to update my scatter plot
def update_graph(num):
data = df[df['Frame'] == index[num]]
graph._offsets3d = (data.x, data.y, data.z)
title.set_text(r'How to update the color map for each plot based on the value of, Frame={}'.format(round(index[num],3)))
#Limits of Elevation (U3) in data
print('U3_min',min(a[:,2]))
print('U3_max',max(a[:,2]))
#Define figure
fig = plt.figure(figsize=(10, 8), dpi=200)
title = plt.suptitle(r'How to update the color map for each plot based on the value of, Frame={}'.format(round(index[0],3), fontsize=16))
# cmap will generate a tuple of RGBA values for a given number in the range 0.0 to 1.0
# (also 0 to 255 - not used in this example).
# To map our z values cleanly to this range, we create a Normalize object.
cmap = matplotlib.cm.get_cmap('coolwarm')
normalize = matplotlib.colors.Normalize(vmin=min(data[:,2,-1]), vmax=max(data[:,2,-1]))
colors = [cmap(normalize(value)) for value in data[:,2,-1]]
ax1 = fig.add_subplot(111, projection='3d')
ax1.view_init(45,30)
ax1.set_xlabel('X (mm)')
ax1.set_ylabel('Y (mm)')
ax1.set_zlabel('Z Elevation (mm)')
data=df[df['Frame']==1]
graph = ax1.scatter(data.x, data.y, data.z, color=colors)
## Optionally add a colorbar
cax, _ = matplotlib.colorbar.make_axes(ax1, shrink=0.5)
cbar = matplotlib.colorbar.ColorbarBase(cax, cmap=cmap, norm=normalize)
#Animation
ani = matplotlib.animation.FuncAnimation(fig, update_graph, samples,
interval=40, blit=False)
plt.show()
In the attached file, the colormap is defined for the last frame of the problem, where U3max=1 and U3min=-1. For example, during the first step (frame=0), there is no elevation, therefore, in the 'coolwarm' scale, the plot should be white, but it takes the values of blue and red as well.
Any suggestion is welcome.
Thanks in advance