Animación con PyGame + Exportarla como GIF

Con el objetivo de representar gráficamente el proceso de criba de Eratóstenes para obtener numeros primos, utilizamos PyGame, Imageio y NumPy para generar la siguiente animación.

Código completo al final del artículo.

Generar la animación

Para comenzar, recordemos la función a la que habíamos llegado en el artículo anterior.

def get_prime_numbers(max_number):
    numbers = [True, True] + [True] * (max_number-1)
    last_prime_number = 2
    i = last_prime_number
    
    while last_prime_number**2 <= max_number:
        i += last_prime_number
        while i <= max_number:
            numbers[i] = False
            i += last_prime_number
        j = last_prime_number + 1
        while j < max_number:
            if numbers[j]:
                last_prime_number = j
                break
            j += 1
        i = last_prime_number
    
    return [i + 2 for i, not_crossed in enumerate(numbers[2:]) if not_crossed]

La idea es, a medida que la función selecciona y descarta números, representarlo en la pantalla utilizando PyGame. Para evitar que el código anterior bloquee el bucle principal de PyGame, haremos las modificaciones a continuación.

def get_prime_numbers(max_number):
    numbers = [True, True] + [True] * (max_number-1)
    last_prime_number = 2
    i = last_prime_number
    
    while last_prime_number**2 <= max_number:
        yield ACTION_SELECTED, i
        i += last_prime_number
        while i <= max_number:
            numbers[i] = False
            yield ACTION_CROSSED_OUT, i
            i += last_prime_number
        j = last_prime_number + 1
        while j < max_number:
            if numbers[j]:
                last_prime_number = j
                break
            j += 1
        i = last_prime_number
    
    for i, not_crossed in enumerate(numbers[2:]):
        if not_crossed:
            yield ACTION_PRIME_NUMBER, i + 2

Previamente definiendo las siguientes constantes.

ACTION_SELECTED = 0
ACTION_CROSSED_OUT = 1
ACTION_PRIME_NUMBER = 2

Convertimos la función en un generador para dar lugar a la ejecución del bucle principal a medida que se realiza el proceso de criba. Las constantes anteriores indicarán qué tipo recuadro se debe dibujar (verde para números primos o rojo para números descartados) y cuánto tiempo debe esperarse. Volveremos a esto más adelante.

Seguimos definiendo otras tres constantes que representan la cantidad de números en pantalla, las columnas y las filas, respectivamente.

NUMBERS = 120
COLS = 10
ROWS = NUMBERS / COLS

Para evitar el uso de objetos globales, implementamos una clase principal, Animation, y una estructura básica de PyGame.

import pygame

# ...

class Animation(object):

    def start(self):
        pygame.init()

        self.screen = pygame.display.set_mode((515, 500))
        pygame.display.set_caption("Criba de Eratóstenes")
        run = True
        # Pintar pantalla de blanco.
        self.screen.fill((255, 255, 255))
        # Fuente para dibujar los números.
        self.font = pygame.font.Font(None, 20)

        while run:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    run = False
                    break
            
            pygame.display.flip()


if __name__ == "__main__":
    Animation().start()

Hasta el momento tenemos simplemente una pantalla en blanco. Continuaremos por colocar la cantidad de numeros definidos (120). Antes de hacerlo, definimos la siguiente función que retorna la fila y columna correspondiente para un determinado número dentro del rango (1-120).

def get_number_position(number):
    col = number % COLS
    if col == 0:
        col = 10
    row = (number - col) / COLS
    return row, col

Luego, dentro de la clase Animation, incluimos la función draw_rect, que será la encargada de dibujar un número con su respectivo recuadro de color en la posición determinada por get_number_position.

    def draw_rect(self, number, color):
        row, col = get_number_position(number)
        self.screen.fill(color, (15 + col*40, 15 + row*40, 30, 30))
        text = self.font.render(str(number), 1, (240, 240, 240))
        w, h = text.get_size()
        self.screen.blit(text,
            (15 + col*40 + (30 - w)/2, 15 + row*40 + (30 - h)/2))

Por último, antes de while run:, procedemos a dibujar los números en pantalla (exceptuando el 1).

        numbers = range(1, NUMBERS + 1)
        for number in numbers:
            if number == 1:
                continue
            self.draw_rect(number, (190, 190, 190))

Ahora bien, luego del código anterior, obtenemos la lista de números primos y creamos un par de variables para el control del tiempo en la animación.

        prime_numbers = get_prime_numbers(NUMBERS)
        prev_ticks = 0
        wait = 100

Finalmente, antes de pygame.display.flip(), ubicamos la siguiente lógica.

            ticks = pygame.time.get_ticks()
            dif = ticks - prev_ticks
            if dif > wait:
                try:
                    action, number = next(prime_numbers)
                except StopIteration:
                    run = False
                else:
                    if action in (ACTION_SELECTED, ACTION_PRIME_NUMBER):
                        self.draw_rect(number, (0, 190, 70))
                    elif action == ACTION_CROSSED_OUT:
                        self.draw_rect(number, (255, 130, 130))
                        
                    if action in (ACTION_CROSSED_OUT, ACTION_PRIME_NUMBER):
                        wait = 75
                    else:
                        wait = 750
                    
                    prev_ticks = ticks

Hacemos uso de la función get_ticks para controlar la duración de cada proceso de selección en pantalla. A medida que obtenemos instrucciones del generador prime_numbers (vía next(prime_numbers)), dibujamos en pantalla los recuadros con su correspondiente color (self.draw_rect). Finalmente, cuando la criba ha finalizado (se lanza StopIteration), terminamos el bucle principal (run = False).

Exportar como imagen GIF

Los módulos y funciones necesarias:

import imageio
from numpy import fliplr, rot90

Luego de wait = 100, creamos una lista para cada fragmento de la animación y otra para su duración.

        frames = []
        duration = []

Las cuales procedemos a completar luego de prev_ticks = ticks.

                    # Obtener vector de NumPy de la imagen actual.
                    frames.append(pygame.surfarray.array3d(self.screen))
                    # Conversión a segundos para Imageio.
                    duration.append(dif / 1000)

Por último, generamos la imagen GIF al final de la función start.

        print("Generando archivo GIF...")
        # Esperar 2 segundos al comenzar.
        duration[0] = 2
        # Esperar 5 segundos al finalizar.
        duration[-1] = 5
        # Rotar e invertir el vector para satisfacer la lectura
        # de Imageio.
        imageio.mimwrite("sieve.gif",
            (fliplr(rot90(f, 3)) for f in frames), duration=duration)
        print("Listo.")

Código completo

#!/usr/bin/env python
# -*- coding: utf-8 -*-

import imageio
import pygame
from numpy import fliplr, rot90

NUMBERS = 120
COLS = 10
ROWS = NUMBERS / COLS

ACTION_SELECTED = 0
ACTION_CROSSED_OUT = 1
ACTION_PRIME_NUMBER = 2


def get_prime_numbers(max_number):
    numbers = [True, True] + [True] * (max_number-1)
    last_prime_number = 2
    i = last_prime_number
    
    while last_prime_number**2 <= max_number:
        yield ACTION_SELECTED, i
        i += last_prime_number
        while i <= max_number:
            numbers[i] = False
            yield ACTION_CROSSED_OUT, i
            i += last_prime_number
        j = last_prime_number + 1
        while j < max_number:
            if numbers[j]:
                last_prime_number = j
                break
            j += 1
        i = last_prime_number
    
    for i, not_crossed in enumerate(numbers[2:]):
        if not_crossed:
            yield ACTION_PRIME_NUMBER, i + 2


def get_number_position(number):
    col = number % COLS
    if col == 0:
        col = 10
    row = (number - col) / COLS
    return row, col


class Animation(object):

    def draw_rect(self, number, color):
        row, col = get_number_position(number)
        self.screen.fill(color, (15 + col*40, 15 + row*40, 30, 30))
        text = self.font.render(str(number), 1, (240, 240, 240))
        w, h = text.get_size()
        self.screen.blit(text,
            (15 + col*40 + (30 - w)/2, 15 + row*40 + (30 - h)/2))

    def start(self):
        pygame.init()

        self.screen = pygame.display.set_mode((515, 500))
        pygame.display.set_caption("Criba de Eratóstenes")
        run = True
        self.screen.fill((255, 255, 255))
        self.font = pygame.font.Font(None, 20)

        numbers = range(1, NUMBERS + 1)
        for number in numbers:
            if number == 1:
                continue
            self.draw_rect(number, (190, 190, 190))
        
        prime_numbers = get_prime_numbers(NUMBERS)
        prev_ticks = 0
        wait = 100
        frames = []
        duration = []

        while run:
            for event in pygame.event.get():
                if event.type == pygame.QUIT:
                    run = False
                    break
            
            ticks = pygame.time.get_ticks()
            dif = ticks - prev_ticks
            if dif > wait:
                try:
                    action, number = next(prime_numbers)
                except StopIteration:
                    run = False
                else:
                    if action in (ACTION_SELECTED, ACTION_PRIME_NUMBER):
                        self.draw_rect(number, (0, 190, 70))
                    elif action == ACTION_CROSSED_OUT:
                        self.draw_rect(number, (255, 130, 130))
                        
                    if action in (ACTION_CROSSED_OUT, ACTION_PRIME_NUMBER):
                        wait = 75
                    else:
                        wait = 750
                    
                    prev_ticks = ticks
                    frames.append(pygame.surfarray.array3d(self.screen))
                    duration.append(dif / 1000)
            
            pygame.display.flip()
        
        print("Generando archivo GIF...")
        duration[0] = 2
        duration[-1] = 5
        imageio.mimwrite("sieve.gif",
            (fliplr(rot90(f, 3)) for f in frames), duration=duration)
        print("Listo.")


if __name__ == "__main__":
    Animation().start()