repository·main·Indexed 23 days ago
https://github.com/pyo3/rust-numpyPyO3-based Rust bindings for the NumPy C-API, providing high-performance bridges between Rust and the Python NumPy ecosystem. It enables the creation of Python extension modules that operate on NumPy arrays using types like PyArrayDyn and PyReadonlyArrayDyn, and supports array-like arguments via PyArrayLike. The library integrates with the ndarray crate for computational operations and supports parallel execution and BLAS via optional feature flags.
To use rust-numpy, ensure your environment meets the following requirements:
numpy must be installed in your Python environment (e.g., via pip install numpy). Recommended version is >= 1.16.0.PyO3 and ndarray.The rust-parallel example demonstrates how to build a rust-numpy extension that leverages optional ndarray feature flags. Specifically, it shows how to enable:
rayon feature in ndarray to enable parallel processing across arrays.BLAS (Basic Linear Algebra Subprograms) via ndarray features for high-performance linear algebra operations.To implement similar functionality in your own extension, you must configure your Cargo.toml to enable these specific feature flags for the ndarray dependency.
The project uses nox to automate environment management. Running nox inside the directory will:
tests/test_exp.py.maturin develop from within a virtualenv. This installs the extension into the current environment, allowing you to import it directly in Python.maturin developrust-linalg package is an example extension that demonstrates how to use rust-numpy in conjunction with the ndarray-linalg crate. This extension is designed to link against a system-provided OpenBLAS implementation for linear algebra operations.To create a Python extension module in Rust that operates on NumPy arrays, you can use PyReadonlyArrayDyn for immutable access and PyArrayDyn for mutable or owned arrays.
Key patterns:
PyReadonlyArrayDyn<'py, T> to receive arrays from Python. Convert them to ndarray views using .as_array() to perform computations.&Bound<'py, PyArrayDyn<T>> and access the underlying mutable array via unsafe { x.as_array_mut() }..into_pyarray(py) on an ndarray type to convert it back into a Python-managed NumPy array.[lib]
name = "rust_ext"
crate-type = ["cdylib"]
[dependencies]
pyo3 = { version = "0.29" }
numpy = "0.29"#[pyo3::pymodule]
mod rust_ext {
use numpy::ndarray::{ArrayD, ArrayViewD, ArrayViewMutD};
use numpy::{IntoPyArray, PyArrayDyn, PyReadonlyArrayDyn, PyArrayMethods};
use pyo3::{pyfunction, PyResult, Python, Bound};
// example using immutable borrows producing a new array
fn axpy(a: f64, x: ArrayViewD<'_, f64>, y: ArrayViewD<'_, f64>) -> ArrayD<f64> {
a * &x + &y
}
// example using a mutable borrow to modify an array in-place
fn mult(a: f64, mut x: ArrayViewMutD<'_, f64>) {
x *= a;
}
// wrapper of `axpy`
#[pyfunction(name = "axpy")]
fn axpy_py<'py>(
py: Python<'py>,
a: f64,
x: PyReadonlyArrayDyn<'py, f64>,
y: PyReadonlyArrayDyn<'py, f64>,
) -> Bound<'py, PyArrayDyn<f64>> {
let x = x.as_array();
let y = y.as_array();
let z = axpy(a, x, y);
z.into_pyarray(py)
}
// wrapper of `mult`
#[pyfunction(name = "mult")]
fn mult_py<'py>(a: f64, x: &Bound<'py, PyArrayDyn<f64>>) {
let x = unsafe { x.as_array_mut() };
mult(a, x);
}
}When using PyArrayLike, you can control whether NumPy should attempt to cast the input data to the requested type T using the C type parameter.
TypeMustMatch (Default)Use this when the element type must match the specified type T exactly. If the input cannot be interpreted as type T without casting (e.g., passing a list of floats to a function expecting i32), the function call will fail.
AllowTypeChangeUse this when you want NumPy to cast the input to type T using numpy.asarray. This is useful for flexibility, but be aware that it may result in precision loss (e.g., casting float to int).
Use the IntoPyArray trait to move ownership of Rust data (like Vec<T>, Box<[T]>, or ndarray::ArrayBase) into a NumPy array. This is highly efficient because it avoids copying data; instead, it passes a pointer to the Rust-allocated memory to Python.
Important Limitations:
resize method. Attempting to do so will result in an error.use numpy::{PyArray, IntoPyArray, PyArrayMethods};
use pyo3::Python;
Python::attach(|py| {
let py_array = vec![1, 2, 3].into_pyarray(py);
assert_eq!(py_array.readonly().as_slice().unwrap(), &[1, 2, 3]);
// Array cannot be resized when its data is owned by Rust.
unsafe {
assert!(py_array.resize(100).is_err());
}
});The numpy::datetime module provides Rust wrappers for NumPy's datetime64 and timedelta64 types. These types are designed for scientific applications, meaning they use flexible units to support a wide range of scales (up to $2^{64}$ years) or high precision (down to $10^{-18}$ seconds), but they ignore calendars and time zones.
You can use Datetime<U> and Timedelta<U> where U is a type implementing the Unit trait, specifying the temporal resolution (e.g., Days, Seconds, Nanoseconds).
use numpy::{datetime::{units, Datetime, Timedelta}, PyArray1, PyArrayMethods};
use pyo3::{Python, types::PyAnyMethods, ffi::c_str};
// Example: Casting a Python NumPy array to a Rust PyArray1 of Datetime<Days>
let array = py
.eval(
c_str!("np.array([np.datetime64('2017-04-21')])"),
None,
Some(&locals),
)?
.cast_into::<PyArray1<Datetime<units::Days>>>()?;The Element trait defines types that can be stored in a NumPy array. It is an unsafe trait because the implementer must guarantee that the type is safe for NumPy to manage (e.g., handling memory and reference counting correctly).
Key requirements for Element:
IS_COPY: A constant indicating if the type is trivially copyable. This must be false for object types or records containing object-type fields.get_dtype(py): Returns the associated PyArrayDescr.clone_ref(py): Creates a clone of the value while the GIL is held.Note on Python Objects: To store Python objects in a NumPy array, implement Element for Py<PyAny>. For these types, IS_COPY is false.
If the nalgebra feature is enabled, rust-numpy provides seamless integration between NumPy arrays and nalgebra matrices:
ToPyArray for nalgebra::Matrix to convert nalgebra matrices into NumPy arrays.PyReadonlyArray::try_as_matrix to treat a NumPy array as a read-only nalgebra matrix slice.PyReadwriteArray::try_as_matrix_mut to treat a NumPy array as a mutable nalgebra matrix slice.When calling mathematical functions like inner(), dot(), or einsum(), the result can be either a NumPy array or a Python scalar. The ArrayOrScalar<'a, 'py, T> trait is used to define the expected return type in a generic way.
It is implemented for:
Bound<'py, PyArray<T, D>>: When you expect the result to be a NumPy array.T: When you expect the result to be a scalar of type T (where T is a NumPy Element and implements FromPyObject).