Skip to content

Instantly share code, notes, and snippets.

@alexlib

alexlib/rustimport_jupyter_openpiv_rust.ipynb

Last active April 9, 2026 22:29
Show Gist options

  • Select an option

    Learn more about clone URLs
  • Save alexlib/ab0db711f9133b9709c573f1f96aca17 to your computer and use it in GitHub Desktop.

Select an option

Learn more about clone URLs
Save alexlib/ab0db711f9133b9709c573f1f96aca17 to your computer and use it in GitHub Desktop.
Download ZIP
rustimport_jupyter_openpiv_rust.ipynb
Display the source blob
Display the rendered blob
Raw
rustimport_jupyter_openpiv_rust.ipynb
{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "view-in-github",
"colab_type": "text"
},
"source": [
"<a href=\"https://colab.research.google.com/gist/alexlib/ab0db711f9133b9709c573f1f96aca17/rustimport_jupyter_openpiv_rust.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "R60MRpVGEykF"
},
"source": [
"# Python Extensions in Rust with Jupyter Notebooks"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "3XOBCO5XEykX",
"outputId": "e7ab0d4b-8ec8-4d18-8c51-f0a683775fbb"
},
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"Requirement already satisfied: rustimport_jupyter in /usr/local/lib/python3.12/dist-packages (0.1.3)\n",
"Requirement already satisfied: polars==0.20.2 in /usr/local/lib/python3.12/dist-packages (0.20.2)\n",
"Requirement already satisfied: ipython in /usr/local/lib/python3.12/dist-packages (from rustimport_jupyter) (7.34.0)\n",
"Requirement already satisfied: rustimport in /usr/local/lib/python3.12/dist-packages (from rustimport_jupyter) (1.7.1)\n",
"Requirement already satisfied: setuptools>=18.5 in /usr/local/lib/python3.12/dist-packages (from ipython->rustimport_jupyter) (75.2.0)\n",
"Requirement already satisfied: jedi>=0.16 in /usr/local/lib/python3.12/dist-packages (from ipython->rustimport_jupyter) (0.19.2)\n",
"Requirement already satisfied: decorator in /usr/local/lib/python3.12/dist-packages (from ipython->rustimport_jupyter) (4.4.2)\n",
"Requirement already satisfied: pickleshare in /usr/local/lib/python3.12/dist-packages (from ipython->rustimport_jupyter) (0.7.5)\n",
"Requirement already satisfied: traitlets>=4.2 in /usr/local/lib/python3.12/dist-packages (from ipython->rustimport_jupyter) (5.7.1)\n",
"Requirement already satisfied: prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0 in /usr/local/lib/python3.12/dist-packages (from ipython->rustimport_jupyter) (3.0.52)\n",
"Requirement already satisfied: pygments in /usr/local/lib/python3.12/dist-packages (from ipython->rustimport_jupyter) (2.20.0)\n",
"Requirement already satisfied: backcall in /usr/local/lib/python3.12/dist-packages (from ipython->rustimport_jupyter) (0.2.0)\n",
"Requirement already satisfied: matplotlib-inline in /usr/local/lib/python3.12/dist-packages (from ipython->rustimport_jupyter) (0.2.1)\n",
"Requirement already satisfied: pexpect>4.3 in /usr/local/lib/python3.12/dist-packages (from ipython->rustimport_jupyter) (4.9.0)\n",
"Requirement already satisfied: toml>=0.10.2 in /usr/local/lib/python3.12/dist-packages (from rustimport->rustimport_jupyter) (0.10.2)\n",
"Requirement already satisfied: parso<0.9.0,>=0.8.4 in /usr/local/lib/python3.12/dist-packages (from jedi>=0.16->ipython->rustimport_jupyter) (0.8.6)\n",
"Requirement already satisfied: ptyprocess>=0.5 in /usr/local/lib/python3.12/dist-packages (from pexpect>4.3->ipython->rustimport_jupyter) (0.7.0)\n",
"Requirement already satisfied: wcwidth in /usr/local/lib/python3.12/dist-packages (from prompt-toolkit!=3.0.0,!=3.0.1,<3.1.0,>=2.0.0->ipython->rustimport_jupyter) (0.6.0)\n",
"\u001b[1minfo:\u001b[0m downloading installer\n",
"\u001b[1m\u001b[33mwarn:\u001b[0m It looks like you have an existing rustup settings file at:\n",
"\u001b[1m\u001b[33mwarn:\u001b[0m /root/.rustup/settings.toml\n",
"\u001b[1m\u001b[33mwarn:\u001b[0m Rustup will install the default toolchain as specified in the settings file,\n",
"\u001b[1m\u001b[33mwarn:\u001b[0m instead of the one inferred from the default host triple.\n",
"\u001b[1minfo:\u001b[0m profile set to default\n",
"\u001b[1minfo:\u001b[0m default host triple is x86_64-unknown-linux-gnu\n",
"\u001b[1m\u001b[33mwarn:\u001b[0m Updating existing toolchain, profile choice will be ignored\n",
"\u001b[1minfo:\u001b[0m syncing channel updates for stable-x86_64-unknown-linux-gnu\n",
"\u001b[1minfo:\u001b[0m default toolchain set to stable-x86_64-unknown-linux-gnu\n",
"\n",
" \u001b[1mstable-x86_64-unknown-linux-gnu unchanged\u001b[0m - rustc 1.94.1 (e408947bf 2026-03-25)\n",
"\n",
"\u001b[0m\u001b[1m\n",
"Rust is installed now. Great!\n",
"\u001b[0m\n",
"To get started you may need to restart your current shell.\n",
"This would reload your \u001b[0m\u001b[1mPATH\u001b[0m environment variable to include\n",
"Cargo's bin directory ($HOME/.cargo/bin).\n",
"\n",
"To configure your current shell, you need to source\n",
"the corresponding \u001b[0m\u001b[1menv\u001b[0m file under $HOME/.cargo.\n",
"\n",
"This is usually done by running one of the following (note the leading DOT):\n",
". \"$HOME/.cargo/env\" # For sh/bash/zsh/ash/dash/pdksh\n",
"source \"$HOME/.cargo/env.fish\" # For fish\n",
"source \"~/.cargo/env.nu\" # For nushell\n",
"source \"$HOME/.cargo/env.tcsh\" # For tcsh\n",
". \"$HOME/.cargo/env.ps1\" # For pwsh\n",
"source \"$HOME/.cargo/env.xsh\" # For xonsh\n",
"The rustimport_jupyter extension is already loaded. To reload it, use:\n",
" %reload_ext rustimport_jupyter\n"
]
}
],
"source": [
"# Install dependencies for google colab\n",
"import os\n",
"import sys\n",
"\n",
"IN_COLAB = \"google.colab\" in sys.modules\n",
"if IN_COLAB:\n",
" !pip install -U rustimport_jupyter polars==0.20.2\n",
" !curl --proto '=https' --tlsv1.2 -sSf https://sh.rustup.rs | sh -s -- -y\n",
" os.environ[\"PATH\"] += \":/root/.cargo/bin\"\n",
" import rustimport_jupyter\n",
" %load_ext rustimport_jupyter"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "BTfxKu3LEykm"
},
"source": [
"The Rust programming language has gotten more prominent for writing compiled Python extensions. Currently, there is a bunch of boilerplate for wrapping writing up a Rust function and making it callable from Python. I enjoy exploring and prototyping code in Jupyter Notebooks, so I developed [rustimport_jupyter](https://github.com/thomasjpfan/rustimport_jupyter) to compile Rust code in Jupyter and have the compiled code available in Python! In this blog post, I will showcase a simple function, NumPy function, and Polar expression plugins. This blog post is runnable as a [notebook on Google Colab](https://colab.research.google.com/github/thomasjpfan/thomasjpfan.github.io/blob/main/content/notebooks/20231227_rustimport_jupyter.ipynb)."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "34mXlPJpEykp"
},
"source": [
"## Simple Rust Functions\n",
"\n",
"`rustimport_jupyter` builds on top of [rustimport](https://github.com/mityax/rustimport) to compile Python extensions written in Rust from Jupyter notebooks. After installing the [rustimport_jupyter](https://pypi.org/project/rustimport-jupyter/) package from PyPI, we load the magic from within a Jupyter notebook:"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "T7lMEuBhEyky",
"outputId": "413b3d95-e690-4615-a6b9-788fc64874e7"
},
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"The rustimport_jupyter extension is already loaded. To reload it, use:\n",
" %reload_ext rustimport_jupyter\n"
]
}
],
"source": [
"%load_ext rustimport_jupyter"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "vLACR89BEykz"
},
"source": [
"Next, we define a `double` function in Rust and prefixing the cell with the `%%rustimport` marker:"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {
"id": "5Ah20mv9Eyk0"
},
"outputs": [],
"source": [
"%%rustimport\n",
"use pyo3::prelude::*;\n",
"\n",
"#[pyfunction]\n",
"fn double(x: i32) -> i32 {\n",
" 2 * x\n",
"}"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ajvatJ8qEyk6"
},
"source": [
"The `%%rustimport` marker compiles the Rust code and imports the `double` function into the Jupyter notebook environment. This means, we can directly call it from Python!"
]
},
{
"cell_type": "code",
"execution_count": 18,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "rIbHJq9GEylE",
"outputId": "0da2a973-e1aa-494c-d3a7-5a83ec6d377d"
},
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"68"
]
},
"metadata": {},
"execution_count": 18
}
],
"source": [
"double(34)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "YR7FqKYCEylN"
},
"source": [
"By default, `%%rustimport` is compiles without Rust optimizations. We can enable these optimizations by adding the `--release` flag:"
]
},
{
"cell_type": "code",
"execution_count": 19,
"metadata": {
"id": "OiWw8AYdEylP"
},
"outputs": [],
"source": [
"%%rustimport --release\n",
"use pyo3::prelude::*;\n",
"\n",
"#[pyfunction]\n",
"fn triple(x: i32) -> i32 {\n",
" 3 * x\n",
"}"
]
},
{
"cell_type": "code",
"execution_count": 20,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "jRrbItq3EylU",
"outputId": "c45095fe-6177-4076-cd5a-5bc59f761ff8"
},
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"21"
]
},
"metadata": {},
"execution_count": 20
}
],
"source": [
"triple(7)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "orSltr0bEylV"
},
"source": [
"## NumPy in Rust\n",
"\n",
"Rust's ecosystem contains many [third party libraries](https://crates.io) that is useful for writing our custom functions. [rustimport](https://github.com/mityax/rustimport#customizing-an-extension) defines a custom `//:` comment syntax that we can use to pull in write our extensions. In this next example, we use [PyO3/rust-numpy](https://github.com/PyO3/rust-numpy) to define a NumPy function that computes `a*x-y` in Rust:"
]
},
{
"cell_type": "code",
"execution_count": 21,
"metadata": {
"id": "Q63ArSViEylZ"
},
"outputs": [],
"source": [
"%%rustimport --release\n",
"//: [dependencies]\n",
"//: pyo3 = { version = \"0.21\", features = [\"extension-module\"] }\n",
"//: numpy = \"0.21\"\n",
"\n",
"use pyo3::prelude::*;\n",
"use numpy::ndarray::{ArrayD, ArrayViewD};\n",
"use numpy::{IntoPyArray, PyArrayDyn, PyReadonlyArrayDyn};\n",
"\n",
"fn axsy(a: f64, x: ArrayViewD<'_, f64>, y: ArrayViewD<'_, f64>) -> ArrayD<f64> {\n",
" a * &x - &y\n",
"}\n",
"\n",
"#[pyfunction]\n",
"#[pyo3(name = \"axsy\")]\n",
"fn axsy_py<'py>(\n",
" py: Python<'py>,\n",
" a: f64,\n",
" x: PyReadonlyArrayDyn<'py, f64>,\n",
" y: PyReadonlyArrayDyn<'py, f64>,\n",
") -> Bound<'py, PyArrayDyn<f64>> {\n",
" let x = x.as_array();\n",
" let y = y.as_array();\n",
" let z = axsy(a, x, y);\n",
" z.into_pyarray_bound(py)\n",
"}\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "IGdRGEdPEylg"
},
"source": [
"The `pyo3(name = \"axsy\")` Rust macro exports the compiled function as `axsy` in Python. We can now use `axsy` directly in Jupyter:"
]
},
{
"cell_type": "code",
"execution_count": 22,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "EuRMn4MMEylh",
"outputId": "e735ebf7-dfb7-4acd-d7a9-f5ba659f86e8"
},
"outputs": [
{
"output_type": "execute_result",
"data": {
"text/plain": [
"array([ 0.3, -8.2, 5.6])"
]
},
"metadata": {},
"execution_count": 22
}
],
"source": [
"import numpy as np\n",
"\n",
"a = 2.4\n",
"x = np.array([1.0, -3.0, 4.0], dtype=np.float64)\n",
"y = np.array([2.1, 1.0, 4.0], dtype=np.float64)\n",
"\n",
"axsy(a, x, y)"
]
},
{
"cell_type": "code",
"source": [
"%%rustimport --release\n",
"//: [dependencies]\n",
"//: pyo3 = { version = \"0.21\", features = [\"extension-module\"] }\n",
"//: numpy = \"0.21\"\n",
"//: ndarray = \"0.15\"\n",
"\n",
"use pyo3::prelude::*;\n",
"use numpy::{PyArray2, PyArrayMethods, ToPyArray, PyUntypedArrayMethods};\n",
"use ndarray::{Array2, s};\n",
"\n",
"#[pyfunction]\n",
"pub fn run_piv(py: Python, img1: &Bound<'_, PyArray2<u8>>, img2: &Bound<'_, PyArray2<u8>>, win_size: usize, search_size: usize) -> PyResult<Py<PyArray2<isize>>> {\n",
" let view1 = img1.readonly();\n",
" let view2 = img2.readonly();\n",
" let view1 = view1.as_array();\n",
" let view2 = view2.as_array();\n",
"\n",
" let (h, w) = (view1.nrows(), view1.ncols());\n",
" let mut vectors = Array2::zeros((h / win_size, w / win_size));\n",
"\n",
" for i in (0..h - win_size).step_by(win_size) {\n",
" for j in (0..w - win_size).step_by(win_size) {\n",
" let window1 = view1.slice(s![i..i+win_size, j..j+win_size]);\n",
" let mut max_corr = -1.0;\n",
" let mut best_shift = 0isize;\n",
"\n",
" let half_search = (search_size / 2) as isize;\n",
" for dy in -half_search..half_search {\n",
" for dx in -half_search..half_search {\n",
" let y_start = i as isize + dy;\n",
" let x_start = j as isize + dx;\n",
" if y_start < 0 || x_start < 0 || (y_start as usize + win_size) > h || (x_start as usize + win_size) > w {\n",
" continue;\n",
" }\n",
" let window2 = view2.slice(s![y_start as usize..y_start as usize + win_size, x_start as usize..x_start as usize + win_size]);\n",
" let mut corr = 0.0;\n",
" for r in 0..win_size {\n",
" for c in 0..win_size {\n",
" corr += (window1[[r, c]] as f32) * (window2[[r, c]] as f32);\n",
" }\n",
" }\n",
" if corr > max_corr {\n",
" max_corr = corr;\n",
" best_shift = dx;\n",
" }\n",
" }\n",
" }\n",
" vectors[[i / win_size, j / win_size]] = best_shift;\n",
" }\n",
" }\n",
" Ok(vectors.to_pyarray_bound(py).unbind())\n",
"}"
],
"metadata": {
"id": "-hO1y0JtJpg9"
},
"execution_count": 23,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "e6e1f623",
"outputId": "d492e6a7-fb0b-4d7b-d5a0-a78eaf6b5f2d"
},
"source": [
"import numpy as np\n",
"import time\n",
"\n",
"def generate_synthetic_data(size=512, shift=3):\n",
" img1 = np.random.randint(0, 256, (size, size), dtype=np.uint8)\n",
" img2 = np.roll(img1, shift, axis=0)\n",
" img2 = np.roll(img2, shift, axis=1)\n",
" return img1, img2\n",
"\n",
"def basic_piv_python(img1, img2, win_size=32, search_size=10):\n",
" h, w = img1.shape\n",
" vectors = np.zeros((h // win_size, w // win_size), dtype=np.int64)\n",
" for i in range(0, h - win_size, win_size):\n",
" for j in range(0, w - win_size, win_size):\n",
" window1 = img1[i:i+win_size, j:j+win_size].astype(np.float32)\n",
" max_corr = -1\n",
" best_shift = 0\n",
" for dy in range(-search_size // 2, search_size // 2):\n",
" for dx in range(-search_size // 2, search_size // 2):\n",
" y_start, x_start = i + dy, j + dx\n",
" if y_start < 0 or x_start < 0 or y_start + win_size > h or x_start + win_size > w:\n",
" continue\n",
" window2 = img2[y_start:y_start+win_size, x_start:x_start+win_size].astype(np.float32)\n",
" corr = np.sum(window1 * window2)\n",
" if corr > max_corr:\n",
" max_corr = corr\n",
" best_shift = dx\n",
" vectors[i//win_size, j//win_size] = best_shift\n",
" return vectors\n",
"\n",
"# Initialize data\n",
"img1, img2 = generate_synthetic_data(size=256)\n",
"win, search = 16, 10\n",
"\n",
"print(\"--- Benchmark Results ---\")\n",
"# Test Pure Python\n",
"start = time.time()\n",
"res_py = basic_piv_python(img1, img2, win, search)\n",
"py_time = time.time() - start\n",
"print(f\"Pure Python Time: {py_time:.4f}s\")\n",
"\n",
"# Test Rust-backed Python\n",
"start = time.time()\n",
"# run_piv was imported by the previous %%rustimport cell\n",
"res_rs = run_piv(img1, img2, win, search)\n",
"rs_time = time.time() - start\n",
"print(f\"Rust-backed Python Time: {rs_time:.4f}s\")\n",
"\n",
"print(f\"Speedup: {py_time / rs_time:.1f}x\")"
],
"execution_count": 24,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"--- Benchmark Results ---\n",
"Pure Python Time: 0.1507s\n",
"Rust-backed Python Time: 0.0097s\n",
"Speedup: 15.5x\n"
]
}
]
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "4867cfe6",
"outputId": "4309cd02-50fc-4c11-8926-f90c5dd682d1"
},
"source": [
"def vectorized_piv_numpy(img1, img2, win_size=32, search_size=10):\n",
" # Pre-convert to float32 once to avoid overhead in the loop\n",
" img1_f = img1.astype(np.float32)\n",
" img2_f = img2.astype(np.float32)\n",
"\n",
" h, w = img1.shape\n",
" vectors = np.zeros((h // win_size, w // win_size), dtype=np.int64)\n",
" half_search = search_size // 2\n",
"\n",
" for i in range(0, h - win_size, win_size):\n",
" for j in range(0, w - win_size, win_size):\n",
" window1 = img1_f[i:i+win_size, j:j+win_size]\n",
" max_corr = -1\n",
" best_shift = 0\n",
"\n",
" for dy in range(-half_search, half_search):\n",
" for dx in range(-half_search, half_search):\n",
" y_s, x_s = i + dy, j + dx\n",
" if y_s < 0 or x_s < 0 or y_s + win_size > h or x_s + win_size > w:\n",
" continue\n",
"\n",
" window2 = img2_f[y_s:y_s+win_size, x_s:x_s+win_size]\n",
" # Optimized NumPy operation\n",
" corr = np.vdot(window1, window2)\n",
"\n",
" if corr > max_corr:\n",
" max_corr = corr\n",
" best_shift = dx\n",
" vectors[i//win_size, j//win_size] = best_shift\n",
" return vectors\n",
"\n",
"print(\"--- Optimized Three-Way Benchmark ---\")\n",
"# 1. Pure Python\n",
"start = time.time()\n",
"res_py = basic_piv_python(img1, img2, win, search)\n",
"py_time = time.time() - start\n",
"print(f\"Basic Python Time: {py_time:.4f}s\")\n",
"\n",
"# 2. NumPy Vectorized (Optimized)\n",
"start = time.time()\n",
"res_np = vectorized_piv_numpy(img1, img2, win, search)\n",
"np_time = time.time() - start\n",
"print(f\"NumPy Optimized Time: {np_time:.4f}s\")\n",
"\n",
"# 3. Rust\n",
"start = time.time()\n",
"res_rs = run_piv(img1, img2, win, search)\n",
"rs_time = time.time() - start\n",
"print(f\"Rust-backed Time: {rs_time:.4f}s\")\n",
"\n",
"print(\"\\nNew Speedup relative to Basic Python:\")\n",
"print(f\"NumPy (Optimized): {py_time / np_time:.1f}x\")\n",
"print(f\"Rust: {py_time / rs_time:.1f}x\")\n",
"\n",
"# Consistency check between the two FFT implementations\n",
"matches = np.sum(res_py == res_rs)\n",
"print(f\"Consistency between Python and Rust versions: {matches}/{res_py.size} matches\")\n",
"\n",
"# Consistency check between the two FFT implementations\n",
"matches = np.sum(res_np == res_rs)\n",
"print(f\"Consistency between Numpy and Rust versions: {matches}/{res_py.size} matches\")"
],
"execution_count": 25,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"--- Optimized Three-Way Benchmark ---\n",
"Basic Python Time: 0.1500s\n",
"NumPy Optimized Time: 0.0572s\n",
"Rust-backed Time: 0.0099s\n",
"\n",
"New Speedup relative to Basic Python:\n",
"NumPy (Optimized): 2.6x\n",
"Rust: 15.2x\n",
"Consistency between Python and Rust versions: 256/256 matches\n",
"Consistency between Numpy and Rust versions: 256/256 matches\n"
]
}
]
},
{
"cell_type": "code",
"source": [
"def fft_piv_numpy(img1, img2, win_size):\n",
" \"\"\"Perform PIV using FFT-based cross-correlation in NumPy.\"\"\"\n",
" img1 = img1.astype(np.float32)\n",
" img2 = img2.astype(np.float32)\n",
" h, w = img1.shape\n",
" vectors = np.zeros((h // win_size, w // win_size), dtype=np.int64)\n",
"\n",
" for i in range(0, h - win_size, win_size):\n",
" for j in range(0, w - win_size, win_size):\n",
" sub1 = img1[i:i+win_size, j:j+win_size]\n",
" sub2 = img2[i:i+win_size, j:j+win_size]\n",
"\n",
" # Standard FFT-based cross-correlation\n",
" f1 = np.fft.fft2(sub1)\n",
" f2 = np.fft.fft2(np.flip(sub2))\n",
" corr = np.real(np.fft.ifft2(f1 * f2))\n",
"\n",
" # Find the displacement\n",
" max_idx = np.unravel_index(np.argmax(corr), corr.shape)\n",
" vectors[i//win_size, j//win_size] = max_idx[1]\n",
" return vectors"
],
"metadata": {
"id": "N21BNq_0Vm4t"
},
"execution_count": 26,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"id": "xC9A6SRhL5jM"
},
"source": [
"%%rustimport --release\n",
"//: [dependencies]\n",
"//: pyo3 = { version = \"0.21\", features = [\"extension-module\"] }\n",
"//: numpy = \"0.21\"\n",
"//: ndarray = \"0.15\"\n",
"//: rustfft = \"6.1\"\n",
"\n",
"use pyo3::prelude::*;\n",
"use numpy::{PyArray2, PyArrayMethods, ToPyArray};\n",
"use ndarray::{Array2, s};\n",
"use rustfft::{FftPlanner, num_complex::Complex};\n",
"\n",
"#[pyfunction]\n",
"pub fn run_piv_fft_rust(py: Python, img1: &Bound<'_, PyArray2<u8>>, img2: &Bound<'_, PyArray2<u8>>, win_size: usize) -> PyResult<Py<PyArray2<isize>>> {\n",
" let view1 = img1.readonly();\n",
" let view2 = img2.readonly();\n",
" let view1 = view1.as_array();\n",
" let view2 = view2.as_array();\n",
"\n",
" let (h, w) = (view1.nrows(), view1.ncols());\n",
" let mut vectors = Array2::zeros((h / win_size, w / win_size));\n",
"\n",
" let mut planner = FftPlanner::new();\n",
" let fft = planner.plan_fft_forward(win_size * win_size);\n",
" let ifft = planner.plan_fft_inverse(win_size * win_size);\n",
"\n",
" for i in (0..h - win_size).step_by(win_size) {\n",
" for j in (0..w - win_size).step_by(win_size) {\n",
" let sub1 = view1.slice(s![i..i+win_size, j..j+win_size]);\n",
" let sub2 = view2.slice(s![i..i+win_size, j..j+win_size]);\n",
"\n",
" let mut buf1: Vec<Complex<f32>> = sub1.iter().map(|&x| Complex::new(x as f32, 0.0)).collect();\n",
" let mut buf2: Vec<Complex<f32>> = sub2.iter().map(|&x| Complex::new(x as f32, 0.0)).collect();\n",
" buf2.reverse();\n",
"\n",
" fft.process(&mut buf1);\n",
" fft.process(&mut buf2);\n",
"\n",
" for k in 0..buf1.len() {\n",
" buf1[k] *= buf2[k];\n",
" }\n",
"\n",
" ifft.process(&mut buf1);\n",
"\n",
" let mut max_val = -1.0;\n",
" let mut max_pos = 0;\n",
" for (idx, val) in buf1.iter().enumerate() {\n",
" if val.re > max_val {\n",
" max_val = val.re;\n",
" max_pos = idx % win_size;\n",
" }\n",
" }\n",
" vectors[[i / win_size, j / win_size]] = max_pos as isize;\n",
" }\n",
" }\n",
" Ok(vectors.to_pyarray_bound(py).unbind())\n",
"}"
],
"execution_count": 27,
"outputs": []
},
{
"cell_type": "code",
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "xC9A6SRhL5jL",
"outputId": "ab906cc1-aabd-43cb-a9d0-1f5bc392b911"
},
"source": [
"import time\n",
"import numpy as np\n",
"\n",
"print(\"--- FFT Performance: Rust vs Python ---\")\n",
"\n",
"# Re-initialize data for consistency\n",
"img1_bench, img2_bench = generate_synthetic_data(size=256)\n",
"win_bench = 16\n",
"\n",
"start = time.time()\n",
"res_fft_py = fft_piv_numpy(img1_bench, img2_bench, win_bench)\n",
"py_fft_time = time.time() - start\n",
"print(f\"FFT (Python/NumPy) Time: {py_fft_time:.4f}s\")\n",
"\n",
"start = time.time()\n",
"# This will now be defined once cell xC9A6SRhL5jM compiles\n",
"res_fft_rs = run_piv_fft_rust(img1_bench, img2_bench, win_bench)\n",
"rs_fft_time = time.time() - start\n",
"print(f\"FFT (Rust) Time: {rs_fft_time:.4f}s\")\n",
"\n",
"print(f\"\\nRust FFT Speedup: {py_fft_time / rs_fft_time:.1f}x\")\n",
"\n",
"matches = np.sum(res_fft_py == res_fft_rs)\n",
"print(f\"Consistency between FFT versions: {matches}/{res_fft_py.size} matches\")"
],
"execution_count": 28,
"outputs": [
{
"output_type": "stream",
"name": "stdout",
"text": [
"--- FFT Performance: Rust vs Python ---\n",
"FFT (Python/NumPy) Time: 0.0541s\n",
"FFT (Rust) Time: 0.0010s\n",
"\n",
"Rust FFT Speedup: 53.6x\n",
"Consistency between FFT versions: 256/256 matches\n"
]
}
]
},
{
"cell_type": "code",
"source": [],
"metadata": {
"id": "KxaEI2_AXQ9_"
},
"execution_count": null,
"outputs": []
}
],
"metadata": {
"kernelspec": {
"display_name": "rustimport-jupyter",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.7"
},
"colab": {
"provenance": [],
"include_colab_link": true
}
},
"nbformat": 4,
"nbformat_minor": 0
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment