I've been writing a small Rust module recently with the sole purpose of speeding up my Python program, which is a Convay'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.