Quickly drawing grids of rectangles, and updating their colors with VisPy

Here’s a demo run of the code:

Background

From the VisPy website:

VisPy is a Python library for interactive scientific visualization that is designed to be fast, scalable, and easy to use.

While looking for a near real time data visualization alternative to the venerable matplotlib, I came across this jaw dropping demo:

Absolutely insane, achieving that kind of performance in python is amazing to say the least. This demo in particular seems like it would be more likely to come from a pygame application at the least, but looks more like it would be a Unity project.

The VisPy project is massive, but luckily, there is a set of really good examples included in the repo. Reminds me of the Arduino standard library in this way. After through all of running these, I didn’t find exactly what I was looking for.

For how simple the finished product looks, the learning curve on the way there was surprisingly steep. Hopefully this post saves you some time.

Code + Explanation

I’d like to be able to draw a grid n * n rectangles, and control their side lengths and the amount of space between them. Being able to update the color of the rectangle is also a requirement.

The goal is to be able to visualize audio signals on a low resolution grid of LEDs.
This software is to be the basis for iterating on ideas and working out bugs without having to redesign hardware.

This example uses a single vispy.visuals.collections.PolygonCollection object to draw the rectangles rather than using a bunch of RectangleVisual objects. Moreover, I found that once there were hundreds of these RectangleVisual objects, I could not achieve the FPS range I was after.  Searching online, I found an explanation for this here:

Any solution that requires more than 10 visuals will perform terribly as it requires a separate GL program for each visual. Each program is drawn sequentially so it takes the GPU a long time to get through each one. It would be best to have one visual that draws all of these rectangles with the appropriate color (I’d never heard of a wafermap before)

So reader, if you’re working on something similar, I hope this snippet is a good starting point:

"""
vispy_grid_of_rectangles.py

Create a grid of rectangles of arbitrary size, and update their color using vispy.

You'll need to install the following packages:
    * mypy==0.720
    * numpy==1.18.2
    * vispy==0.6.4
    * PyQt5==5.14

4/20/2020 - Devon Bray - www.esologic.com/quickly-drawing-grids-of-rectangles-and-updating-their-colors-with-vispy
"""

import random
from typing import Optional, Union

import numpy as np
from vispy import app, gloo
from vispy.util.event import Event
from vispy.visuals.collections import PathCollection, PolygonCollection

# Both the `PolygonCollection` and `PathCollection` used in this example support 3D polygons,
# But we're only concerned with the 2D case here, so this global constant is set to avoid confusion.
Z_POSITION_DEFAULT = 0.0


def rectangle_polygon_vertices_2d(
    horizontal_side_length: float, vertical_side_length: float,
) -> np.ndarray:
    """
    Create a list of vertices that represent a rectangle centered around (0, 0, 0).
    The `_side_length` parameters should be floats greater than zero and less than one.
    These are 2D rectangles, not boxes, so changing the `z_position` will move the rectangle along
    the z axis.
    :param horizontal_side_length: how long the horizontal sizes should be
    :param vertical_side_length: how long the vertical sides should be
    :return: the vertices (x, y, z) as a numpy array.
    """
    return (
        np.array(
            [
                [1.0, 1.0, Z_POSITION_DEFAULT],  # top right
                [-1.0, 1.0, Z_POSITION_DEFAULT],  # top left
                [-1.0, -1.0, Z_POSITION_DEFAULT],  # bottom left
                [1.0, -1.0, Z_POSITION_DEFAULT],  # bottom right
            ]
        )
        * (horizontal_side_length, vertical_side_length, 1)
        / 2
    )


def offset_rectangle(position: int, side_length_float: float, buffer_length_float: float) -> float:
    """
    Given the position of a rectangle plus it's side/buffer lengths, create the float offset from
    the outer edge.
    :param position: The column/row number of the current rectangle
    :param side_length_float: the side length of the rectangle as a float
    :param buffer_length_float: the buffer (distance between rectangles) as a float
    :return: the offset distance for the current rectangle as a float
    """

    # Since the rectangle is currently centered around (0, 0), exactly HALF of it is already above
    # the origin.
    middle_offset = side_length_float / 2

    # These are the side lengths of the previous rectangles
    length_of_rectangles = position * side_length_float

    # These are the lengths of the previous buffers between rectangles
    border_lengths = (position + 1) * buffer_length_float

    return middle_offset + length_of_rectangles + border_lengths


class GridOfRectangles(app.Canvas):
    """
    Draw a grid of rectangles using `PolygonCollection` and `PathCollection` objects.
    Every 1 second, change each rectangles to a random color.
    """

    def __init__(
        self: "GridOfRectangles",
        rows: int,
        columns: int,
        horizontal_side_length_pixels: int,
        vertical_side_length_pixels: int,
        buffer_length_pixels: int,
        **kwargs: Optional[Union[bool, str]]
    ) -> None:
        """
        Constructor, note that `self.polygons` and `self.edges_of_polygons` are created before the
        parent class is initialized.
        :param rows: the desired number of rows in the grid
        :param columns: the desired number of columns in the grid
        :param horizontal_side_length_pixels: the side horizontal length (width) of each of the
        rectangles in pixels
        :param vertical_side_length_pixels: the vertical side length (height) of each rectangle in
        pixels
        :param buffer_length_pixels: the number of pixels between each rectangle.
        :param kwargs: keyword arguments to be passed to the parent canvas constructor.
        :return: None
        """

        self.polygons = PolygonCollection("agg", color="shared")
        self.edges_of_polygons = PathCollection("agg", color="shared")

        app.Canvas.__init__(
            self,
            size=(
                (columns * horizontal_side_length_pixels) + ((columns + 1) * buffer_length_pixels),
                (rows * vertical_side_length_pixels) + ((rows + 1) * buffer_length_pixels),
            ),
            **kwargs
        )

        self._timer = app.Timer(interval=1, connect=self.on_timer, start=True)

        horizontal_side = (horizontal_side_length_pixels / self.physical_size[0]) * 2
        horizontal_buffer = (buffer_length_pixels / self.physical_size[0]) * 2

        vertical_side_length = (vertical_side_length_pixels / self.physical_size[1]) * 2
        vertical_buffer_length = (buffer_length_pixels / self.physical_size[1]) * 2

        self.num_rectangles = columns * rows

        for num_column in range(columns):
            for num_row in range(rows):

                xyz_offset = (
                    -1 + (offset_rectangle(num_column, horizontal_side, horizontal_buffer)),  # X
                    1
                    - (
                        offset_rectangle(num_row, vertical_side_length, vertical_buffer_length)
                    ),  # Y
                    Z_POSITION_DEFAULT,  # Z
                )

                polygon_points = (
                    rectangle_polygon_vertices_2d(horizontal_side, vertical_side_length)
                    + xyz_offset
                )

                # The default color for each of this will be black.
                # The color of the polygons will be changed with `on_timer`, but the edges will
                # stay as is.
                self.polygons.append(polygon_points, color=(0, 0, 0, 1))
                self.edges_of_polygons.append(polygon_points, closed=True, color=(0, 0, 0, 1))

        self.edges_of_polygons["linewidth"] = 1
        self.edges_of_polygons["viewport"] = 0, 0, self.physical_size[0], self.physical_size[1]

    def on_draw(self: "GridOfRectangles", event: Event) -> None:
        """
        Called each time a frame is to be written, as often as possible
        :param event: The vispy event context
        :return: None
        """
        gloo.clear("white")
        self.polygons.draw()
        self.edges_of_polygons.draw()
        self.update()

    def on_timer(self: "GridOfRectangles", event: Event) -> None:
        """
        This function sets the color of each of the polygons to a random color.
        Called by the timer, every 1 second.
        :param event: The vispy event context
        :return: None
        """
        self.polygons["color"] = np.array(
            list(
                (random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1), 1)
                for _ in range(self.num_rectangles)
            )
        )
        self.update()


if __name__ == "__main__":
    canvas = GridOfRectangles(
        rows=500,
        columns=500,
        horizontal_side_length_pixels=3,
        vertical_side_length_pixels=3,
        buffer_length_pixels=1,
        show=True,
        keys="interactive",
    )
    gloo.set_viewport(0, 0, canvas.size[0], canvas.size[1])
    gloo.set_state("translucent", depth_test=False)
    canvas.measure_fps()
    app.run()

With a grid of 100 x 100, I’m able update the visual at ~30 FPS, which is way more that I would ever need for my application. If you wanted to extend this to more objects, it seems like the bet approach is to use markers.

Thanks for reading.

2 Comments

Leave a Reply

Your email address will not be published. Required fields are marked *

This site uses Akismet to reduce spam. Learn how your comment data is processed.