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):


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

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 {
                    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;

fn interop_module_rs(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(update_cells, m)?)?;


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"),
    pos = (3, 3)
          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:
    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:
        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))):

            # 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)):

                    # 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))):

        cells = update_cells(cells)
        draw(surface, cells, cellsize, grid)

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

Screenshot of the simulation running:

Screenshot of the Conway's Game of Life simulation after 10 frames from the initial, random set-up

Any help will be highly appreciated.

  • 1
    \$\begingroup\$ You could unconditionally check "is num_alive == 3 ?", plus test for "was already alive and has 2 neighbors". More importantly, you repeatedly allocate and zero a new PyArray2 where simple double buffering would suffice. In order to use several threads you probably want to create a number_of_neighbors array, wait for that calculation to finish, then come back and make birth/death decisions. \$\endgroup\$
    – J_H
    Commented May 19, 2023 at 20:43

1 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.

  • \$\begingroup\$ Following your instructions I was able to get a new version of my function to work, however, it slightly negatively affected its performance (using map_collect()); when I try to use par_map_collect() I am shown the following error: error[E0277]: `UnsafeCell<PyObject>` cannot be shared between threads safely which points to let nxt = Zip::indexed(cur.readonly().as_array()).par_map_collect( > | < (r, c), _| { \$\endgroup\$
    – chubercik
    Commented May 31, 2023 at 22:04
  • \$\begingroup\$ @chubercik, I think the issue on both counts is probably what's in the closure. It seems like you are calling some function that invokes python functionality, which a) will hurt performance and b) isn't allowed in a parallel context. \$\endgroup\$ Commented Jun 1, 2023 at 4:43
  • \$\begingroup\$ thanks for the heads-up, I was able to locate the issue - the easy fix was to covert the cur variable before Zip and change the indexing syntax a little. It seems to work, though performance-wise it looks to be on par with my original solution: this might be the result of some sort of bottleneck on the Python script's side; I will investigate it further, however as my original question has been answered I am no longer looking for help this. Thanks again for your time :) \$\endgroup\$
    – chubercik
    Commented Jun 1, 2023 at 9:52
  • \$\begingroup\$ @chubercik, is it on par with parallelism or without? \$\endgroup\$ Commented Jun 1, 2023 at 15:12
  • \$\begingroup\$ with, but I noticed a mistake and in the end, the paralleled version actually outperforms the original as the size of the world matrix grows. \$\endgroup\$
    – chubercik
    Commented Jun 1, 2023 at 16:07

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.