0

I have the code below. If you run that program a scatter plot is created. The plot is very messy. I want every point to represent one car. I only want to see the point of the current position of each car in the animation. I am not very good at matplotlib, but how would I do that?

from matplotlib import pyplot as plt
import random
import math
from matplotlib import animation
import numpy as np

roadLength    = 100
numFrames     = 200

numCars = 60
velocities = []
positions     = []
theta         = []
r             = []
color         = []
probability = 0.5
vmax = 1
flowr = []

fig = plt.figure()
ax = fig.add_subplot(111, projection='polar')
ax.axis('off')


for i in range(numCars):
    positions.append(i * 2)
    theta.append(0)
    r.append(1)
    color.append(i)
    velocities.append(0)


def pos(k):
    rand = random.uniform(0,1)
    if velocities[k] < vmax:
        velocities[k] +=  1
    if k < numCars-1 and (positions[k+1] - positions[k]) <= velocities[k]:
        velocities[k] = positions[k+1] - positions[k] - 1
    if k == numCars-1 and (roadLength - (positions[k] - positions[0])) <= velocities[k]:
        velocities[k] = roadLength - (positions[k] - positions[0]) - 1

    if rand < probability and velocities[k] > 0:
        velocities[k] = velocities[k] - 1
    positions[k] = positions[k] + velocities[k]
    return positions[k]




def animate(frameNr):
    for i in range(numCars):
        theta[i] = pos(i) * 2 * math.pi / roadLength
        flowr.append(velocities[i])
    ax.scatter(theta, r, c=color)
    return ax.scatter(theta, r, c=color),


def simulate():
    anim = animation.FuncAnimation(fig, animate,
                               frames=numFrames, interval=10, blit=True, repeat=False)

    plt.show()

simulate()

Thanks in advance!

Edit:

I fixed so that it became a point plot I am now storing velocity, position and lane in an object for a later problem. My question is: How do I make the points have different colors?

from matplotlib import pyplot as plt
import random
import math
from matplotlib import animation
import numpy as np
from operator import attrgetter

roadLength    = 100
numFrames     = 200
nlanes = 1
numCars = 20
posss =[]
theta         = []
r             = []
color         = []
probability = 0.5
vmax = 1
flowr = []
cars=[]


class Car:
    def __init__(self, position, velocity, lane):
        self.position = position
        self.velocity = velocity
        self.lane = lane






def pos(car,k):
    rand = random.uniform(0,1)
    if car[k].velocity < vmax:
        car[k].velocity +=  1
    if k < numCars-1 and (car[k+1].position - car[k].position) <= car[k].velocity:
        car[k].velocity = car[k+1].position - car[k].position - 1
    if k == numCars-1 and (roadLength - (car[k].position - car[0].position)) <= car[k].velocity:
        car[k].velocity = roadLength - (car[k].position - car[0].position) - 1

    if rand < probability and car[k].velocity > 0:
        car[k].velocity = car[k].velocity - 1
    car[k].position = car[k].position + car[k].velocity
    return car[k].position

for i in range(numCars):
    cars.append(Car(i, 0, 1))
    theta.append(0)
    r.append(1)
    color.append(i)
    posss.append(i)

fig = plt.figure()
ax = fig.add_subplot(111)
point, = ax.plot(posss, r, 'o')
ax.set_xlim(-10, 1.2*numFrames)
ax.set_ylim(-2, nlanes + 3)

def animate(frameNr):
    sort_cars = sorted(cars, key=attrgetter("position"))
    for i in range(numCars):
        pos(sort_cars,i)
    for k in range(numCars):
        theta[k]=cars[k].position


    point.set_data(theta, r)
    return point,


def simulate():
    anim = animation.FuncAnimation(fig, animate,
                               frames=numFrames, interval=10, blit=True, repeat=False)

    plt.show()

simulate()

1 Answers1

0

plot() returns a Line2D object, in which all points have the same color.

If you want your points to have different colors, you should use scatter(). The logic is exactly the same, except you use set_offsets(<2D array>) instead of set_data().

See this question for more details

Diziet Asahi
  • 38,379
  • 7
  • 60
  • 75