CODEWITHSUNDEEP
pythonmatplotlib
AN_
AN_

Reputation: 143

Making a 3D, moving scatterplot with matplotlib

I want to make a 3D plot in matplotlib from an already existing 2D plot.
In the 2D plot, there are dots, which are randomly moving inside a circle. Now I have made the axes 3D and as far as I can see, the dots are in a 3D position as well. But they are not moving, and I get an error message: 'ValueError: too many values to unpack (expected 2)'.
I didn't find where I have made the error (I guess somewhere the program still wants only the x,y coordinates).
Could you help me? 

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

def get_initial_coordinates():
    x_coord =[random.uniform(3, 7) for i in range(n_particles)]
    y_coord = [random.uniform(3, 7) for i in range(n_particles)]
    z_coord = [random.uniform(3, 7) for i in range(n_particles)]
    return x_coord, y_coord, z_coord

def get_initial_velocities():
    x_vel = [3 * (np.random.random() - 0.5) * box_width for i in range(n_particles)]
    y_vel = [3 * (np.random.random() - 0.5) * box_width for i in range(n_particles)]
    z_vel = [3 * (np.random.random() - 0.5) * box_width for i in range(n_particles)]
    return x_vel, y_vel, z_vel

def take_step(x_coord, y_coord, z_coord, x_vel, y_vel, z_vel):
    for i in range(n_particles):
     x_coord[i] += x_vel[i]*dt
     y_coord[i] += y_vel[i]*dt
     z_coord[i] += z_vel[i]*dt

    return x_coord, y_coord,z_coord, x_vel, y_vel,z_vel

n_particles = 40
box_width = 10
n_steps = 5000
dt = 0.001

x_coord, y_coord, z_coord = get_initial_coordinates()
x_vel, y_vel, z_vel = get_initial_velocities()
for i in range(n_steps):
    x_coord,y_coord,z_coord,x_vel,y_vel,z_vel= take_step(x_coord,y_coord,z_coord,x_vel,y_vel,z_vel)
#---------------------
fig = pyplot.figure()
ax= Axes3D(fig)
d, = ax.plot([x_coord[i] for i in range(n_particles)],
             [y_coord[i] for i in range(n_particles)],
             [z_coord[i] for i in range(n_particles)], 'ro')
x_coord, y_coord, z_coord = get_initial_coordinates()
x_vel, y_vel, z_vel = get_initial_velocities()
circle = plt.Circle((5, 5), 3, color='blue', fill=False)
ax.add_artist(circle)

def animate(i):
    take_step(x_coord, y_coord, z_coord, x_vel, y_vel, z_vel)
    d.set_data([x_coord[i]for i in range(n_particles)],
               [y_coord[i]for i in range(n_particles)],
               [z_coord[i]for i in range(n_particles)])
    return d,

anim = animation.FuncAnimation(fig, animate, frames=200, interval=20)
pyplot.show()

Upvotes: 1

Views: 209

Answers (1)

p479h
p479h

Reputation: 199

set_data() is hinherited from Line2D, so it can only take 2 lists, one for x one for y. You have to use .set_3d_properties() for the z component.

def animate(i):
    take_step(x_coord, y_coord, z_coord, x_vel, y_vel, z_vel)
    d.set_data([x_coord[i]for i in range(n_particles)],
               [y_coord[i]for i in range(n_particles)],
               )
    d.set_3d_properties([z_coord[i]for i in range(n_particles)])
    return d,

Upvotes: 1

Related Questions