Rust中能否通过pyo3将nalgebra矩阵类型暴露给Python调用

这个需求完全可以实现，你当前代码无法编译的核心原因是：nalgebra的矩阵、向量类型没有实现pyo3的FromPyObject和IntoPy trait，pyo3无法自动完成Python对象和Rust nalgebra类型之间的互转。

最常用的兼容方案是将矩阵类型和Python numpy数组做互转，适配Python生态的常规使用习惯，实现步骤如下：

1. 补充Cargo.toml依赖 #

需要开启pyo3的numpy特性，引入numpy绑定库：

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

[dependencies]
pyo3 = { version = "0.20", features = ["extension-module", "numpy"] }
nalgebra = "0.32"
numpy = "0.20" # 版本需和所用pyo3版本匹配

2. 改写lib.rs实现类型互转 #

调整函数参数为numpy数组类型，手动完成numpy和nalgebra类型的转换：

extern crate nalgebra as na;
use pyo3::prelude::*;
use numpy::{PyArray1, PyArray2, IntoPyArray};
use na::{Vector3, Matrix3};

/// Formats the sum of two numbers as string.
#[pyfunction]
fn sum_as_string(a: usize, b: usize) -> PyResult<String> {
    Ok((a + b).to_string())
}


#[pyfunction]
fn matrix_math<'py>(py: Python<'py>, v3: &'py PyArray1<f64>, m3x3: &'py PyArray2<f64>) -> PyResult<&'py PyArray1<f64>>{
    //  numpy数组转nalgebra类型
    let v3_slice = unsafe { v3.as_slice()? };
    let rust_v3 = Vector3::from_row_slice(v3_slice);
    
    let m3_slice = unsafe { m3x3.as_slice()? };
    let rust_m3 = Matrix3::from_row_slice(m3_slice);
    
    // 执行矩阵计算
    let mxv = rust_m3 * rust_v3;
    
    // 计算结果转回numpy数组返回
    Ok(mxv.into_pyarray(py))
}

#[pymodule]
fn libpytest(_py: Python, m: &PyModule) -> PyResult<()> {
    m.add_function(wrap_pyfunction!(sum_as_string, m)?)?;
    m.add_function(wrap_pyfunction!(matrix_math, m)?)?;

    Ok(())
}

3. Python侧调用示例 #

直接传入numpy数组即可调用：

import numpy as np
import libpytest

# 注意dtype要和Rust侧的f64对应
vec = np.array([1.0, 2.0, 3.0], dtype=np.float64)
mat = np.array([[1.0, 0.0, 0.0], [0.0, 1.0, 0.0], [0.0, 0.0, 1.0]], dtype=np.float64)

result = libpytest.matrix_math(vec, mat)
print(result)
# 输出: [1. 2. 3.]

其他说明 #

• 如果需要使用动态大小的矩阵，只需读取numpy数组的shape属性，构造nalgebra的动态矩阵即可
• 如果需要在Python侧暴露自定义矩阵操作方法，可以给nalgebra类型套一层#[pyclass]自定义包装类，在包装类上实现需要暴露的方法即可

内容的提问来源于stack exchange，提问作者sbeskur