Parallelization of the tick function in a Game of Life simulation using Rust and Python (+ NumPy)

Question

I've been writing a small Rust module recently with the sole purpose of speeding up my Python program, which is a Conway's Game of Life simulation. The function written in Rust is called from Python using the maturin package and PyO3 crate (on Rust's side).

I've come up with a working prototype, but there are still a few things I'd like to polish out before I'm happy with the project:

1. the function, albeit working, looks a little verbose and not how Rust ought to be written - I am new to the language, so clunky code is to be expected; a complete rewrite of the function might be beneficial in achieving the desired effect,
2. the function uses a PyArray2 object to represent the Game of Life Universe (due to the fact that on Python's side this is stored as a NumPy's ndarray) and stores individual cells as u8 type, which I think is an overkill - we don't need 8 bits per cell when one bit is plenty enough, but I haven't got a clue how to properly store them in such a manner,
3. most importantly: parallelization - I'd like my function to fully utilize the CPU cores by running multiple threads in the update function (for example, by splitting the updated universe in half and dispatching work to two threads). I have attempted to achieve this in a few different ways but haven't come up with anything of value, hence why I'm asking for help).

Attached below is the source code of both my Rust and Python files (I only care about improving the update_cells() function on Rust's side, but I thought showing the entirety of the project's code might ease the work required to test the function):

lib.rs:

use numpy::PyArray2;
use pyo3::prelude::*;

#[pyfunction]
fn update_cells<'a>(py: Python<'a>, cur: &'a PyArray2<u8>) -> PyResult<&'a PyArray2<u8>> {
    let width = cur.shape()[0];
    let height = cur.shape()[1];
    let nxt = PyArray2::<u8>::zeros(py, [width, height], false);

    for r in 0..width {
        for c in 0..height {
            let mut num_alive: u8 = 0;
            for i in -1..2 {
                for j in -1..2 {
                    if i == 0 && j == 0 {
                        continue;
                    }
                    let row = (r as i32 + i) as usize;
                    let col = (c as i32 + j) as usize;
                    if row < width && col < height {
                        num_alive += cur.get_owned([row, col]).unwrap();
                    }
                }
            }
            if (cur.get_owned([r, c]).unwrap() == 1 && 2 <= num_alive && num_alive <= 3)
                || (cur.get_owned([r, c]).unwrap() == 0 && num_alive == 3)
            {
                unsafe {
                    *nxt.get_mut([r, c]).unwrap() = 1;
                }
            }
        }
    }
    Ok(nxt)
}

#[pymodule]
fn interop_module_rs(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(update_cells, m)?)?;
    Ok(())
}

main.py:

import os
from interop_module_rs import update_cells

os.environ["PYGAME_HIDE_SUPPORT_PROMPT"] = '1'

import numpy as np  # noqa: E402
import pygame as pg  # noqa: E402
from numpy.typing import ArrayLike  # noqa: E402

PATH = os.path.dirname(os.path.abspath(__file__))
MAX_FPS = 60  # set to 0 to disable


def draw(surface: pg.Surface,
         cur: ArrayLike,
         sz: int,
         grid: pg.Surface) -> None:
    cur = np.flip(np.rot90(cur*255), 0)
    cur = np.repeat(np.repeat(cur, sz, axis=0), sz, axis=1)
    surf = pg.surfarray.make_surface(cur)
    surface.blit(surf, (0, 0))
    surface.blit(grid, (0, 0))


def init_gun(dimx: int,
             dimy: int) -> ArrayLike:
    cells = np.zeros((dimy, dimx)).astype(np.uint8)
    pattern = np.loadtxt(
        os.path.join(PATH, "init_gun.csv"),
        dtype=np.uint8,
        delimiter=",")
    pos = (3, 3)
    cells[pos[0]:pos[0]+pattern.shape[0],
          pos[1]:pos[1]+pattern.shape[1]] = pattern
    return cells


def random_init(dimx: int,
                dimy: int) -> ArrayLike:
    return np.random.randint(0, 2, (dimy, dimx)).astype(np.uint8)


def create_grid(surface: pg.Surface,
                cellsize: int) -> pg.Surface:
    grid = pg.Surface(surface.get_size())
    for x in range(0, grid.get_width(), cellsize):
        pg.draw.line(grid, (32, 32, 32), (x, 0), (x, grid.get_height()))
    for y in range(0, grid.get_height(), cellsize):
        pg.draw.line(grid, (32, 32, 32), (0, y), (grid.get_width(), y))
    grid.set_colorkey((0, 0, 0))
    return grid


def main(dimx: int,
         dimy: int,
         cellsize: int) -> None:
    pg.init()
    surface = pg.display.set_mode((dimx*cellsize, dimy*cellsize))
    pg.display.set_caption("Game of Life")

    # cells = init_gun(dimx, dimy)
    cells = random_init(dimx, dimy)
    grid = create_grid(surface, cellsize)

    clock = pg.time.Clock()
    while True:
        clock.tick(MAX_FPS)
        pg.display.set_caption(f"Game of Life - {int(clock.get_fps())} FPS")

        for event in pg.event.get():
            # quit game
            if (event.type == pg.QUIT or
                (event.type == pg.KEYDOWN and
                 event.key in (pg.K_ESCAPE, pg.K_q))):
                pg.quit()
                return

            # pause game
            if (event.type == pg.KEYDOWN and
               event.key in (pg.K_SPACE, pg.K_p)):
                while True:
                    pg.display.set_caption("Game of Life - PAUSED")
                    event = pg.event.wait()
                    if (event.type == pg.KEYDOWN and
                       event.key in (pg.K_SPACE, pg.K_p)):
                        break

                    # check if user wants to quit during pause
                    if (event.type == pg.QUIT or
                        (event.type == pg.KEYDOWN and
                         event.key in (pg.K_ESCAPE, pg.K_q))):
                        pg.quit()
                        return

        cells = update_cells(cells)
        draw(surface, cells, cellsize, grid)
        pg.display.update()


if __name__ == "__main__":
    main(64, 64, 9)

Screenshot of the simulation running:

Any help will be highly appreciated.

Answer 1

The numpy PyArray types are somewhat limited. Fortunately, they have an as_array method that will give you an Array from the ndarray crate which is must more powerful. It is unsafe for PyArray but you can use .readonly().as_array() safely.

Once you've got an array you can used numpy::ndarray::Zip, something like this:

Zip::indexed(cur.readonly().as_array()).map_collect(|(r, c), _| {
  // calculate the correct value for row r and column c and return it
})

This will create a new array of the same shape as cur but with values computed by calling the closure.

You can also call par_map_collect instead of map_collect and it will do it using parallelism.

Then you can convert that array into a numpy array using numpy::PyArray2::from_owned_array.

                let row = (r as i32 + i) as usize;
                let col = (c as i32 + j) as usize;

What happens if r = 0 and i = -1 or c = 0 and j = 1? In dev mode this should panic. If you really want to use usizes that wrap around you should use std::num::Wrapping or the wrapping_add function.