如何使用 PyO3 包装带有 NDArray 输入和输出的函数？

Question

我想包装一个以一维 NDArray (rust-numpy) 和 usize 作为参数的函数，returns 一个使用 PyO3 调用代码的一维数组 python.不幸的是，我找不到一个很好的例子来说明如何处理 PyO3 中的数组。这是我到目前为止的代码：

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(())
}

我添加了装饰器和模块包装功能。现在我收到这个错误：

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

这是我的 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"

有人可以告诉我如何实际修改 SMA 函数以用于 PyO3 吗？谢谢！

Answer 1

所以我最终弄明白了：

#[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))
}

问题是 PyO3 需要定义 NDArray 才能为 python 包装它。因为我没有这样做，所以我使用了 Vec，python 相当于一个浮点数列表。然后我使用 Array::from_vec() 将其转换为 NDArray。这允许我们从 Python 列表中获取值并将其转换为 NDArray 以使用 Rust 中的 NumPy 函数进行处理。然后使用 Array::to_vec() 函数将该列表转换回普通列表，作为普通列表发送回 Python。

如何使用 PyO3 包装带有 NDArray 输入和输出的函数？

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

python

numpy

rust

pyo3

rust-ndarray