PyReadonlyArray2 to Vec<T>

Question

I would like to convert a 2D Numpy array in Python to a Rust vector. Specifically, the second dimension of the array is known. That is, the array is in shape [VARIABLE, 1024] (with type double) I would like to know how I could convert it to a Rust vector. Specifically, Vec<[f64; 1024]>

I know that I should use rust-numpy to do this And the first step should be converting it into PyReadonlyArray2 But I'm not sure what should I do next?

Answer 1

ndarray does not support constant dimensions, so you won't find support for converting into arrays (nalgebra does support them, but it looks like it doesn't have support for converting to Vec). So the first step will be to convert to Vec<f64>, and then transmute that into Vec<[f64; 1024]> (this can be done without unsafe code with a crate like bytemuck).

The first step can also be quite involved if we want to do it efficiently. It will be good if we can just allocate directly and copy the contents using memcpy(); however, we can only do that if the array is contiguous (i.e. not a view that is not contiguous), and it is in C layout, i.e. row-major layout, and not Fortran layout (column-major).

All in all:

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

Edit: Turns out this can be done with ndarray methods without much trouble, and this is also faster in case of non-contiguous layout:

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 mut data = arr
        .as_array()
        .as_standard_layout()
        .into_owned()
        .into_raw_vec();

    // Guard against case in which capacity%ARRAY_SIZE!=0, in which `bytemuck::allocation::cast_vec()` will fail.
    data.shrink_to_fit();

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