How to wrap function with NDArray input and output with PyO3?

Question

I want to wrap a function that takes an one-dimensional NDArray (rust-numpy) and an usize as parameters, and returns a one-dimensional array using PyO3 to call the code from python. Unfortunately, I can't find a good example of how to deal with the arrays in PyO3. This is the code that I have so far:

use numpy::ndarray::prelude::*;
use pyo3::prelude::*;
use pyo3::wrap_pyfunction;

#[pyfunction]
fn sma(prices: &Array1<f32>, period: usize) -> PyResult<Array1<f32>> {
    let length = prices.len() - period +1;
    let mut result = Array1::<f32>::zeros(length);

    for i in 0..length {
        let slice = prices.slice(s![i..i+period]);
        result[i] = slice.sum()/(period as f32);
    }

    Ok(result)
}

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

I added the decorators, and the module wrapping function. Now I am getting this error:

error[E0277]: the trait bound `Result<ArrayBase<OwnedRepr<f32>, Dim<[usize; 1]>>, PyErr>: IntoPyCallbackOutput<_>` is not satisfied.

This is my cargo.toml:

[package]
name = "panther"
version = "0.1.0"
edition = "2021"

[lib]
name = "panther"
crate-type = ["cdylib"]

[dependencies]
pyo3 = { version = "0.12.3", features = ["extension-module"] }
numpy = "0.15"

Can someone show me how to actually modify the SMA function for use with PyO3? Thank you!

score 1 · Accepted Answer · answered Feb 18 '22 at 07:54

So I ended up figuring it out:

#[pyfunction]
fn sma(a: Vec<f32>, period: usize) -> PyResult<Vec<f32>> {
    let prices = Array::from_vec(a);
    let length = prices.len() - period +1;
    let mut result = Array1::<f32>::zeros(length);

    for i in 0..length {
        let slice = prices.slice(s![i..i+period]);
        result[i] = slice.sum()/(period as f32);
    }

    Ok(Array::to_vec(&result))
}

The issue is that PyO3 needs NDArray to be defined in order for it to wrap it for python. Since I wasn't doing that, I'm using a Vec, the python equivalent of a list of floats. I am then converting this to a NDArray by using Array::from_vec(). This allows us to get the values from the Python list and convert it to an NDArray to process with NumPy functions in Rust. The list is then converted back to a normal list with the Array::to_vec() function to send back to Python as a normal list.

How to wrap function with NDArray input and output with PyO3?

1 Answers1