Trouble with rendering speed of the animation of a simulation?

Question

The following code is a conversion from some old java I wrote to python.

It shows the beginnings of a simulation of an ant colony.

I am finding the animation speed very slow - and I'm wondering if I am doing something wrong - or if it's nothing to do with the animation, and everything to do with not exploiting the vector methods of numpy (which, in fairness, I struggle to understand).

I am aware that there are 5k points being mapped. But I've seen matplotlib demos handling many more than that.

I am relatively new to python, matplotlib, and numpy - and do not really understand how to profile. I particularly do not 'get' numpy - nor do I really understand broadcasting (though I read about it in 'absolute_beginners'(.

import random
from math import sqrt, cos, sin, radians
import numpy as np
import matplotlib.pyplot as plt
from matplotlib import use as mpl_use
from matplotlib.animation import FuncAnimation
from matplotlib.colors import LinearSegmentedColormap
mpl_use("QtAgg") # PyCharm seems to need this for animations.
history = 100 # trail size per mover.
class Mover:
 offs = [i for i in range(history)] # used to reference the mover history.
 def __init__(self, colour, window_limits):
 # colour = R G B tuple of floats between 0 (black) and 1 (full)
 r, g, b = colour
 self.velocity = 0.75
 self.ccmap = LinearSegmentedColormap.from_list("", [[r, g, b, 1.0 - i / history] for i in range(history)])
 self.mv = 0., 1. # starting unit vector.
 self.mx = [home_x for i in range(history)]
 self.my = [home_y for i in range(history)]
 self.plt = plt.scatter(self.mx, self.my, c=Mover.offs, cmap=self.ccmap)
 self.normalising = False
 self.wrapping = True
 self.w_limits = window_limits
 self.steer(radians(random.uniform(-180, 180)))
 def start(self):
 x, y = self.mx[0], self.my[0] # copy most recent state.
 self.mx.insert(0, x)
 self.my.insert(0, y)
 if len(self.mx) > history: # remove the oldest state if too long..
 self.mx.pop()
 self.my.pop()
 def normalise(self):
 x, y = self.mv
 mag = sqrt(x*x + y*y)
 if mag != 1.0:
 if mag != 0:
 x /= mag
 y /= mag
 else:
 x = 0.0
 y = 1.0
 self.mv = x, y
 def steer(self, theta):
 # theta is in radians
 x, y = self.mv
 u = x*cos(theta) - y*sin(theta) # x1 = x0cos(θ) – y0sin(θ)
 v = x*sin(theta) + y*cos(theta) # y1 = x0sin(θ) + y0cos(θ)
 self.mv = u, v
 if self.normalising:
 self.normalise()
 def transit(self):
 u, v = self.mv
 self.mx[0] += self.velocity * u
 self.my[0] += self.velocity * v
 if self.wrapping:
 self.wrap()
 def wrap(self):
 x, y = self.mx[0], self.my[0]
 if x > self.w_limits:
 x -= self.w_limits
 if y > self.w_limits:
 y -= self.w_limits
 if x < 0:
 x += self.w_limits
 if y < 0:
 y += self.w_limits
 self.mx[0], self.my[0] = x, y
 def update_plot(self):
 self.plt.set_offsets(np.c_[self.mx, self.my])
window_limits = 100
home_x, home_y = 50.0, 50.0
wl = window_limits
fig, ax = plt.subplots()
ax.set_xlim(0, wl)
ax.set_ylim(0, wl)
mover_count = 50
movers = []
colours = plt.cm.turbo
color_normal = colours.N/mover_count
for m in range(mover_count):
 col = colours.colors[int(m*color_normal)]
 mover = Mover(col, wl)
 movers.append(mover)
def init():
 ax.set_xlim(0, window_limits)
 ax.set_ylim(0, window_limits)
 return [o.plt for o in movers]
def animate(frame: int):
 th_var = 30
 for mov in movers:
 mov.start()
 mov.steer(radians(random.uniform(-th_var, th_var)))
 mov.transit()
 mov.update_plot()
 return [mv.plt for mv in movers]
ani = FuncAnimation(fig, animate, init_func=init, blit=True)
plt.show()

• Reinderien: This is cool. It looks like the growth of a slime mold.

Answer 1

Don't import from math when you have Numpy.

I didn't find mpl_use to be necessary.

history, since it's a constant, should be HISTORY.

Don't [i for i in range(history)]; range(history) alone has the same effect since you don't need to mutate the result.

Add PEP484 type hints.

Expanding your colour thus:

r, g, b = colour

is not needed, since you can just use *colour in your list.

You never normalise so I've omitted this from my example code.

Your wrapping code can be greatly simplified by using a single call to np.mod.

From fig, ax = plt.subplots() onward, none of that code should exist in the global namespace and should be moved to functions.

The biggest factor causing apparent slowness is that you've not overridden the default interval=, so the animation appears choppy. A first pass that addresses some of the above:

import random
from typing import Iterator
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.collections import PathCollection
from matplotlib.colors import LinearSegmentedColormap, ListedColormap
HISTORY = 100 # trail size per mover.
HOME_X, HOME_Y = 50, 50
WINDOW_LIMITS = 100
class Mover:
 VELOCITY = 0.75
 WRAPPING = True
 def __init__(self, colour: list[float], window_limits: int) -> None:
 # colour = R G B tuple of floats between 0 (black) and 1 (full)
 cmap_array = np.empty((HISTORY, 4))
 cmap_array[:, :3] = colour
 cmap_array[:, 3] = np.linspace(start=1, stop=0, num=HISTORY)
 ccmap = LinearSegmentedColormap.from_list(name="", colors=cmap_array)
 offsets = range(HISTORY)
 self.mv = 0., 1. # starting unit vector.
 self.mxy = [
 [HOME_X] * HISTORY,
 [HOME_Y] * HISTORY,
 ]
 self.plt: PathCollection = plt.scatter(*self.mxy, c=offsets, cmap=ccmap)
 self.w_limits = window_limits
 self.steer(np.radians(random.uniform(-180, 180)))
 def start(self) -> None:
 x, y = self.mxy # copy most recent state.
 x.insert(0, x[0])
 y.insert(0, y[0])
 if len(x) > HISTORY: # remove the oldest state if too long..
 x.pop()
 y.pop()
 def steer(self, theta: float) -> None:
 # theta is in radians
 x, y = self.mv
 cosa, sina = np.cos(theta), np.sin(theta)
 u = x*cosa - y*sina # x1 = x0cos(θ) – y0sin(θ)
 v = x*sina + y*cosa # y1 = x0sin(θ) + y0cos(θ)
 self.mv = u, v
 def transit(self) -> None:
 u, v = self.mv
 x, y = self.mxy
 x[0] += self.VELOCITY * u
 y[0] += self.VELOCITY * v
 if self.WRAPPING:
 self.wrap()
 def wrap(self) -> None:
 x, y = self.mxy
 if x[0] > self.w_limits:
 x[0] -= self.w_limits
 if y[0] > self.w_limits:
 y[0] -= self.w_limits
 if x[0] < 0:
 x[0] += self.w_limits
 if y[0] < 0:
 y[0] += self.w_limits
 def update_plot(self) -> None:
 self.plt.set_offsets(np.array(self.mxy).T)
def make_movers(colours: ListedColormap, mover_count: int = 50) -> Iterator[Mover]:
 colour_indices = np.linspace(start=0, stop=colours.N-1, num=mover_count, dtype=int)
 for c in colour_indices:
 col = colours.colors[c]
 yield Mover(col, WINDOW_LIMITS)
def main() -> None:
 def init() -> list[plt.Artist]:
 ax.set_xlim(0, WINDOW_LIMITS)
 ax.set_ylim(0, WINDOW_LIMITS)
 return artists
 def animate(frame: int) -> list[plt.Artist]:
 theta = np.radians(30)
 for mov in movers:
 mov.start()
 mov.steer(random.uniform(-theta, theta))
 mov.transit()
 mov.update_plot()
 return artists
 fig, ax = plt.subplots()
 movers = tuple(make_movers(plt.cm.turbo))
 artists = [mv.plt for mv in movers]
 ani = FuncAnimation(fig=fig, func=animate, init_func=init, blit=True, interval=10)
 plt.show()
if __name__ == '__main__':
 main()

A vectorised implementation is possible, but doesn't make much of a difference; the slowest part is matplotlib itself:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.collections import PathCollection
from matplotlib.colors import LinearSegmentedColormap, ListedColormap
from numpy.random import default_rng
HISTORY = 100 # trail size per mover.
HOME_X, HOME_Y = 50, 50
WINDOW_LIMITS = 100
VELOCITY = 0.75
rand = default_rng()
def start_movers(pos: np.ndarray) -> None:
 # pos is n_movers * HISTORY * 2, newest first
 # shift the data one row down, treating it as a queue
 pos[:, 1:, :] = pos[:, :-1, :]
def steer_movers(vel: np.ndarray, theta: float) -> None:
 # vel is n_movers * 2
 x, y = vel.T
 theta = rand.uniform(low=-theta, high=theta, size=vel.shape[0])
 cosa, sina = np.cos(theta), np.sin(theta)
 u = x*cosa - y*sina
 v = x*sina + y*cosa
 vel[:, 0] = u
 vel[:, 1] = v
def transit_movers(pos: np.ndarray, vel: np.ndarray) -> None:
 pos[:, 0, :] += VELOCITY * vel
def wrap_movers(pos: np.ndarray) -> None:
 np.mod(pos, WINDOW_LIMITS, out=pos)
class Mover:
 def __init__(self, cmap_array: np.ndarray, pos: np.ndarray) -> None:
 self.pos = pos
 ccmap = LinearSegmentedColormap.from_list(name="", colors=cmap_array)
 self.plt: PathCollection = plt.scatter(*self.pos.T, c=range(HISTORY), cmap=ccmap)
 def update_plot(self) -> None:
 self.plt.set_offsets(self.pos)
def make_movers(
 colours: ListedColormap,
 mover_count: int = 50,
) -> tuple[
 np.ndarray, # positions
 np.ndarray, # velocities
 list[Mover],
]:
 pos = np.empty((mover_count, HISTORY, 2))
 pos[..., 0] = HOME_X
 pos[..., 1] = HOME_Y
 vel = np.empty((mover_count, 2))
 theta = rand.uniform(low=-np.pi, high=np.pi, size=mover_count)
 vel[:, 0] = np.cos(theta)
 vel[:, 1] = np.sin(theta)
 colour_indices = np.linspace(start=0, stop=colours.N-1, num=mover_count, dtype=int)
 cmap_array = np.empty((mover_count, HISTORY, 4))
 cmap_array[:, :, :3] = np.array(colours.colors)[colour_indices, np.newaxis, :]
 cmap_array[:, :, 3] = np.linspace(start=1, stop=0, num=HISTORY)
 movers = []
 for i_mover, cmap in enumerate(cmap_array):
 movers.append(Mover(cmap_array=cmap, pos=pos[i_mover, ...]))
 return pos, vel, movers
def main() -> None:
 def init() -> list[plt.Artist]:
 ax.set_xlim(0, WINDOW_LIMITS)
 ax.set_ylim(0, WINDOW_LIMITS)
 return artists
 def animate(frame: int) -> list[plt.Artist]:
 start_movers(pos)
 steer_movers(vel, theta=np.radians(30))
 transit_movers(pos, vel)
 wrap_movers(pos)
 for mov in movers:
 mov.update_plot()
 return artists
 fig, ax = plt.subplots()
 pos, vel, movers = make_movers(plt.cm.turbo)
 artists = [mv.plt for mv in movers]
 ani = FuncAnimation(fig=fig, func=animate, init_func=init, blit=True, interval=10)
 plt.show()
if __name__ == '__main__':
 main()

A different style of animation is possible where you don't cull old points based on your queueing code and instead disable blit and draw over your old data, but this quickly hits scalability limits once enough data build up:

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
from matplotlib.collections import PathCollection
from matplotlib.colors import LinearSegmentedColormap, ListedColormap
from numpy.random import default_rng
HOME_X, HOME_Y = 50, 50
WINDOW_LIMITS = 100
VELOCITY = 0.75
rand = default_rng()
def steer_movers(vel: np.ndarray, theta: float) -> None:
 # vel is n_movers * 2
 x, y = vel.T
 theta = rand.uniform(low=-theta, high=theta, size=vel.shape[0])
 cosa, sina = np.cos(theta), np.sin(theta)
 u = x*cosa - y*sina
 v = x*sina + y*cosa
 vel[:, 0] = u
 vel[:, 1] = v
def transit_movers(pos: np.ndarray, vel: np.ndarray) -> None:
 pos += VELOCITY * vel
def wrap_movers(pos: np.ndarray) -> None:
 np.mod(pos, WINDOW_LIMITS, out=pos)
def make_movers(
 colours: ListedColormap,
 mover_count: int = 50,
) -> tuple[
 np.ndarray, # positions
 np.ndarray, # velocities
 LinearSegmentedColormap, # colour map, non-history-based
]:
 pos = np.empty((mover_count, 2))
 pos[:, 0] = HOME_X
 pos[:, 1] = HOME_Y
 vel = np.empty((mover_count, 2))
 theta = rand.uniform(low=-np.pi, high=np.pi, size=mover_count)
 vel[:, 0] = np.cos(theta)
 vel[:, 1] = np.sin(theta)
 colour_indices = np.linspace(start=0, stop=colours.N-1, num=mover_count, dtype=int)
 cmap_array = np.ones((mover_count, 4))
 cmap_array[:, :3] = np.array(colours.colors)[colour_indices, :]
 cmap = LinearSegmentedColormap.from_list(name="", colors=cmap_array)
 return pos, vel, cmap
def main() -> None:
 def init() -> list[plt.Artist]:
 ax.set_xlim(0, WINDOW_LIMITS)
 ax.set_ylim(0, WINDOW_LIMITS)
 return []
 def animate(frame: int) -> list[plt.Artist]:
 steer_movers(vel, theta=np.radians(30))
 transit_movers(pos, vel)
 wrap_movers(pos)
 scatter: PathCollection = plt.scatter(
 *pos.T,
 c=range(pos.shape[0]),
 cmap=cmap,
 )
 return [scatter]
 fig, ax = plt.subplots()
 pos, vel, cmap = make_movers(plt.cm.turbo)
 ani = FuncAnimation(fig=fig, func=animate, init_func=init, blit=False, interval=10)
 plt.show()
if __name__ == '__main__':
 main()

• Konchog: I think that I have enough to work with for the time being. However, I would like to find a python animation/renderer that can handle much larger numbers of ants. While I've had enough of Java, I see dramatically faster results via Graphics2D than in MatPlotLib. (edit/saw the blit off - not what I am looking for).