PyReadonlyArray2 到 Vec<T>

Question

我想将 Python 中的 2D Numpy 数组转换为 Rust 向量。 具体来说，数组的第二维是已知的。 也就是说，数组的形状为 [VARIABLE, 1024]（类型为 double） 我想知道如何将其转换为 Rust 向量。 具体来说，Vec<[f64; 1024]>

我知道我应该使用 rust-numpy 来做到这一点 第一步应该是将其转换为 PyReadonlyArray2 但我不知道接下来该怎么办？

Answer 1

ndarray

不支持常量维度，因此您不会找到转换为数组的支持（

nalgebra

确实支持它们，但看起来不支持转换为

Vec

）。因此，第一步是转换为

Vec<f64>

，然后将其转换为

Vec<[f64; 1024]>

（这可以在没有不安全代码的情况下使用像

bytemuck

这样的板条箱来完成）。

如果我们想要高效地完成第一步，那么第一步也可能会非常复杂。如果我们可以直接分配并使用

memcpy()

复制内容就好了；然而，只有当数组是连续的（即不是不连续的视图）时，我们才能这样做，并且它是C布局，即行优先布局，而不是Fortran布局（列优先）。

总而言之：

const ARRAY_SIZE: usize = 1024;

fn numpy_to_vec(arr: &PyReadonlyArray2<f64>) -> PyResult<Vec<[f64; ARRAY_SIZE]>> {
    if arr.shape()[1] != ARRAY_SIZE {
        return Err(PyTypeError::new_err(format!(
            "expected column size to be {ARRAY_SIZE}",
        )));
    }

    let data = if arr.is_c_contiguous() {
        arr.as_slice().unwrap().to_vec()
    } else {
        let mut data = Vec::with_capacity(arr.len());
        let arr = arr.as_array();
        for row in arr.axis_iter(Axis(0)) {
            match row.as_slice() {
                Some(row) => data.extend_from_slice(row),
                None => data.extend(row),
            }
        }
        data
    };

    Ok(bytemuck::allocation::cast_vec(data))
}