2
\$\begingroup\$

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()
\$\endgroup\$
1
  • 1
    \$\begingroup\$ This is cool. It looks like the growth of a slime mold. \$\endgroup\$ Commented May 5, 2022 at 19:30

1 Answer 1

Reset to default
2
\$\begingroup\$

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()
\$\endgroup\$
1
  • \$\begingroup\$ 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). \$\endgroup\$ Commented May 6, 2022 at 10:08

You must log in to answer this question.

Start asking to get answers

Find the answer to your question by asking.

Ask question

Explore related questions

See similar questions with these tags.