0

I'm trying to make an animated 3-D scatter plot with the ability to plot a dynamic number of classes as different colors. This is one of the attempts. I've included the whole code in case it is helpful, and marked the trouble spot with a row of stars:

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

x_arr,y_arr,depth_arr,time_arr,ml_arr,cluster_arr = np.loadtxt(data, unpack=5, usecols=(0, 1, 2, 5, 6))
class Point:
        def __init__(self,x,y,depth,time,cluster):
                self.x=x
                self.y=y
                self.depth=depth
                self.time=time
                self.cluster=cluster
points = []
for i in range(0,len(x_arr)):
        points.append(Point(x_arr[i],y_arr[i],depth_arr[i],time_arr[i],cluster_arr[i]))
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.set_xlim(min(x_arr), max(x_arr))
ax.set_ylim(min(y_arr), max(y_arr))
ax.set_zlim(min(depth_arr), max(depth_arr))
colors_1 = plt.cm.jet(np.linspace(0,max(cluster_arr),max(cluster_arr)+1))
colors = colors_1.reshape(-1,4)
def plot_points(time):
        x = []
        y = []
        z = []
        clust = []
        points_cp = list(np.copy(points))
        for i in range(0,(int(max(cluster_arr))+1)):
                for event in points_cp:
                        if event.cluster == i:
                                if event.time < time:
                                        points_cp.remove(event)
                                elif event.time <= time + 86400:
                                        x.append(event.x)
                                        y.append(event.y)
                                        z.append(event.depth)
                                        clust.append(event.cluster)
                                        points_cp.remove(event)
# **************************************************************
        color_ind = 0
        first_ind = 0
        last_ind = 0
        for i in range(0,len(x)):
                if clust[i] != color_ind:
                        last_ind = i
                        for i in range(0,len(x)):
                                ax.scatter(x[first_ind:last_ind],y[first_ind:last_ind],z[first_ind:last_ind],c=colors[int(color_ind)])
                                color_ind = clust[i]
                                first_ind = i

time = np.linspace(min(time_arr),max(time_arr),100)

ani = animation.FuncAnimation(fig,plot_points,time)
plt.show()

This gives me a plot with the correct colors, but once a point is plotted, it remains throughout the entire animation.

I have also tried set_x, set_color, etc., but this doesn't work with a loop (it is updated with each iteration, so that only the last class is actually plotted), and I need to use a for loop to accommodate a variable number of classes. I've tried using a colormap with a fixed extent, but have been unsuccessful, as colormapping doesn't work with the plot function, and I haven't been able to get the rest of the code to work with a scatter function.

Thanks in advance for your help, and my apologies if the code is a little wonky. I'm pretty new to this.

wrf
  • 1

0 Answers0