Pygame and Numpy Animations

Solution 1:

[...] shape (y,x,4,k) where y is the height, x is the width, 4 is the number of channels (red, green, blue, alpha), and the k is the number of frames [...]

Yes it is possible, with a single line of code. The following line convets a (y, x, 4, k) numpy array (data) to a list of pygme.Surface objects (surf_list):

surf_list = [pygame.image.frombuffer(d[:,:,[2, 1, 0, 3]].flatten(), (data.shape[1::-1]), 'RGBA') for d in data.transpose(3, 0, 1, 2)]

respectively

surf_list = []
for d in data.transpose(3, 0, 1, 2):
    bytes = d[:,:,[2, 1, 0, 3]].flatten()
    size = data.shape[1::-1]
    format = 'RGBA'
    surface = pygame.image.frombuffer(bytes, size, format)
    surf_list.append(surface)

Explanation:

Use numpy.traspose to bring the 3rd (frame) axis moving axis 2 to the front (see Iterating over arbitrary dimension of numpy.array) and iterate through the frames:

for d in data.transpose(3, 0, 1, 2):

Create a pygame.Surface from each frame by pygame.image.frombuffer():

surface = pygame.image.frombuffer(bytes, size, format)

pygame.image.frombuffer() has 3 arguments, bytes, size, format. bytes is the 1 dimensional byte array of pixel data. numpy.ndarray.flatten return a copy of the array collapsed into one dimension. Most likely the order of the color channels is BGRA rather than RGBA. Hence you have to swap the red and blue color channel (d[:,:,[2, 1, 0, 3]]). You can skip this if the order of the color channels is RGBA:

bytes = d[:,:,[2, 1, 0, 3]].flatten() # for BGRA

bytes = d.flatten() # for RGBA

size is a tuple with 2 elements (x, y) and specifies the size of the image. The size can be get form numpy.ndarray.shape:

size = data.shape[1::-1]

or

size = (data.shape[1], data.shape[0])

format specifies the image format and has to be 'RGBA' ('BGRA' doesn't exist):

format = 'RGBA'

See the minimal example, which crates a (y, x, 4, k) numpy array (data ) and converts it the a list of pygme.Surface objects (surf_list):

import pygame
import numpy as np
pygame.init()
window = pygame.display.set_mode((400, 400))
clock = pygame.time.Clock()

radius = 100
frames = 20
data = np.zeros(shape = (radius*2, radius*2, 4, frames), dtype = "uint8")
for x in range(data.shape[0]): 
    for y in range(data.shape[1]):
        px, py = x - data.shape[0]/2, y - data.shape[1]/2
        for i in range(frames):
            maxY2 = (radius*radius - px*px) * pow(abs(i-frames/2) / frames, 2)
            if px*px + py* py < radius*radius:
                if py * py < maxY2:
                    data[y, x, (0, 3), i] = 255, 255
                    if (px*px + py*py)*4 > radius*radius:
                        data[y, x, (1, 2), i] = 255, 255
                else:
                    data[y, x, (2, 3), i] = 255, 255 
                   
surf_list = [pygame.image.frombuffer(d[:,:,[2, 1, 0, 3]].flatten(), (data.shape[1::-1]), 'RGBA') for d in data.transpose(3, 0, 1, 2)]

count = 0
run = False
while not run:
    clock.tick(20)
    for event in pygame.event.get():
        if event.type == pygame.QUIT:
            run = True

    window.fill(0)
    window.blit(surf_list[count], surf_list[0].get_rect(center = window.get_rect().center))
    pygame.display.flip()
    count = (count + 1) %  len(surf_list)

pygame.quit()