Merge pull request #119 from kngwyu/slice_to_pyarray

Fix to_pyarray for not-contiguous array

Author: kngwyu

src/array.rs

@@ -296,7 +296,7 @@ impl<T, D> PyArray<T, D> {
     }
 
     /// Returns the pointer to the first element of the inner array.
-    unsafe fn data(&self) -> *mut T {
+    pub(crate) unsafe fn data(&self) -> *mut T {
         let ptr = self.as_array_ptr();
         (*ptr).data as *mut T
     }
@@ -355,7 +355,7 @@ impl<T: TypeNum, D: Dimension> PyArray<T, D> {
     pub(crate) unsafe fn new_<'py, ID>(
         py: Python<'py>,
         dims: ID,
-        strides: *mut npy_intp,
+        strides: *const npy_intp,
         flag: c_int,
     ) -> &'py Self
     where
@@ -367,7 +367,7 @@ impl<T: TypeNum, D: Dimension> PyArray<T, D> {
             dims.ndim_cint(),
             dims.as_dims_ptr(),
             T::typenum_default(),
-            strides,                // strides
+            strides as *mut _,      // strides
             ptr::null_mut(),        // data
             0,                      // itemsize
             flag,                   // flag

src/convert.rs

@@ -58,7 +58,7 @@ where
     type Item = A;
     type Dim = D;
     fn into_pyarray<'py>(self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
-        let strides = npy_strides(&self);
+        let strides = self.npy_strides();
         let dim = self.raw_dim();
         let boxed = self.into_raw_vec().into_boxed_slice();
         unsafe { PyArray::from_boxed_slice(py, dim, strides.as_ptr(), boxed) }
@@ -71,12 +71,27 @@ where
 /// elements there**.
 /// # Example
 /// ```
-/// # fn main() {
 /// use numpy::{PyArray, ToPyArray};
 /// let gil = pyo3::Python::acquire_gil();
 /// let py_array = vec![1, 2, 3].to_pyarray(gil.python());
 /// assert_eq!(py_array.as_slice().unwrap(), &[1, 2, 3]);
-/// # }
+/// ```
+///
+/// This method converts a not-contiguous array to C-order contiguous array.
+/// # Example
+/// ```
+/// use numpy::{PyArray, ToPyArray};
+/// use ndarray::{arr3, s};
+/// let gil = pyo3::Python::acquire_gil();
+/// let py = gil.python();
+/// let a = arr3(&[[[ 1,  2,  3], [ 4,  5,  6]],
+///                [[ 7,  8,  9], [10, 11, 12]]]);
+/// let slice = a.slice(s![.., 0..1, ..]);
+/// let sliced = arr3(&[[[ 1,  2,  3]],
+///                     [[ 7,  8,  9]]]);
+/// let py_slice = slice.to_pyarray(py);
+/// assert_eq!(py_slice.as_array(), sliced);
+/// pyo3::py_run!(py, py_slice, "assert py_slice.flags['C_CONTIGUOUS']");
 /// ```
 pub trait ToPyArray {
     type Item: TypeNum;
@@ -102,26 +117,107 @@ where
     type Dim = D;
     fn to_pyarray<'py>(&self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
         let len = self.len();
-        let mut strides = npy_strides(self);
-        unsafe {
-            let array = PyArray::new_(py, self.raw_dim(), strides.as_mut_ptr() as *mut npy_intp, 0);
-            array.copy_ptr(self.as_ptr(), len);
-            array
+        if let Some(order) = self.order() {
+            // if the array is contiguous, copy it by `copy_ptr`.
+            let strides = self.npy_strides();
+            unsafe {
+                let array = PyArray::new_(py, self.raw_dim(), strides.as_ptr(), order.to_flag());
+                array.copy_ptr(self.as_ptr(), len);
+                array
+            }
+        } else {
+            // if the array is not contiguous, copy all elements by `ArrayBase::iter`.
+            let dim = self.raw_dim();
+            let strides = NpyStrides::from_dim(&dim, mem::size_of::<A>());
+            unsafe {
+                let array = PyArray::<A, _>::new_(py, dim, strides.as_ptr(), 0);
+                let data_ptr = array.data();
+                for (i, item) in self.iter().enumerate() {
+                    data_ptr.offset(i as isize).write(*item);
+                }
+                array
+            }
         }
     }
 }
 
-fn npy_strides<S, D, A>(array: &ArrayBase<S, D>) -> Vec<npyffi::npy_intp>
+enum Order {
+    Standard,
+    Fortran,
+}
+
+impl Order {
+    fn to_flag(&self) -> c_int {
+        match self {
+            Order::Standard => 0,
+            Order::Fortran => 1,
+        }
+    }
+}
+
+trait ArrayExt {
+    fn npy_strides(&self) -> NpyStrides;
+    fn order(&self) -> Option<Order>;
+}
+
+impl<A, S, D> ArrayExt for ArrayBase<S, D>
 where
     S: Data<Elem = A>,
     D: Dimension,
-    A: TypeNum,
 {
-    array
-        .strides()
-        .into_iter()
-        .map(|n| n * mem::size_of::<A>() as npyffi::npy_intp)
-        .collect()
+    fn npy_strides(&self) -> NpyStrides {
+        NpyStrides::new(
+            self.strides().into_iter().map(|&x| x as npyffi::npy_intp),
+            mem::size_of::<A>(),
+        )
+    }
+
+    fn order(&self) -> Option<Order> {
+        if self.is_standard_layout() {
+            Some(Order::Standard)
+        } else if self.ndim() > 1 && self.raw_view().reversed_axes().is_standard_layout() {
+            Some(Order::Fortran)
+        } else {
+            None
+        }
+    }
+}
+
+/// Numpy strides with short array optimization
+enum NpyStrides {
+    Short([npyffi::npy_intp; 8]),
+    Long(Vec<npyffi::npy_intp>),
+}
+
+impl NpyStrides {
+    fn as_ptr(&self) -> *const npy_intp {
+        match self {
+            NpyStrides::Short(inner) => inner.as_ptr(),
+            NpyStrides::Long(inner) => inner.as_ptr(),
+        }
+    }
+    fn from_dim<D: Dimension>(dim: &D, type_size: usize) -> Self {
+        Self::new(
+            dim.default_strides()
+                .slice()
+                .into_iter()
+                .map(|&x| x as npyffi::npy_intp),
+            type_size,
+        )
+    }
+    fn new(strides: impl ExactSizeIterator<Item = npyffi::npy_intp>, type_size: usize) -> Self {
+        let len = strides.len();
+        let type_size = type_size as npyffi::npy_intp;
+        if len <= 8 {
+            let mut res = [0; 8];
+            for (i, s) in strides.enumerate() {
+                res[i] = s * type_size;
+            }
+            NpyStrides::Short(res)
+        } else {
+            NpyStrides::Long(strides.map(|n| n as npyffi::npy_intp * type_size).collect())
+        }
+    }
 }
 
 /// Utility trait to specify the dimention of array

src/error.rs

@@ -1,8 +1,5 @@
 //! Defines error types.
-
-use crate::array::PyArray;
-use crate::convert::ToNpyDims;
-use crate::types::{NpyDataType, TypeNum};
+use crate::types::NpyDataType;
 use pyo3::{exceptions as exc, PyErr, PyResult, Python};
 use std::error;
 use std::fmt;

src/types.rs

@@ -68,7 +68,7 @@ impl NpyDataType {
     }
 }
 
-pub trait TypeNum {
+pub trait TypeNum: std::fmt::Debug + Copy {
     fn is_same_type(other: i32) -> bool;
     fn npy_data_type() -> NpyDataType;
     fn typenum_default() -> i32;

tests/to_py.rs

@@ -108,3 +108,36 @@ fn into_pyarray_cant_resize() {
     let arr = a.into_pyarray(gil.python());
     assert!(arr.resize(100).is_err())
 }
+
+#[test]
+fn forder_to_pyarray() {
+    let gil = pyo3::Python::acquire_gil();
+    let py = gil.python();
+    let matrix = Array2::from_shape_vec([4, 2], vec![0, 1, 2, 3, 4, 5, 6, 7]).unwrap();
+    let fortran_matrix = matrix.reversed_axes();
+    let fmat_py = fortran_matrix.to_pyarray(py);
+    assert_eq!(fmat_py.as_array(), array![[0, 2, 4, 6], [1, 3, 5, 7]],);
+    pyo3::py_run!(py, fmat_py, "assert fmat_py.flags['F_CONTIGUOUS']")
+}
+
+#[test]
+fn slice_to_pyarray() {
+    let gil = pyo3::Python::acquire_gil();
+    let py = gil.python();
+    let matrix = Array2::from_shape_vec([4, 2], vec![0, 1, 2, 3, 4, 5, 6, 7]).unwrap();
+    let slice = matrix.slice(s![1..4; -1, ..]);
+    let slice_py = slice.to_pyarray(py);
+    assert_eq!(slice_py.as_array(), array![[6, 7], [4, 5], [2, 3]],);
+    pyo3::py_run!(py, slice_py, "assert slice_py.flags['C_CONTIGUOUS']")
+}
+
+#[test]
+fn forder_into_pyarray() {
+    let gil = pyo3::Python::acquire_gil();
+    let py = gil.python();
+    let matrix = Array2::from_shape_vec([4, 2], vec![0, 1, 2, 3, 4, 5, 6, 7]).unwrap();
+    let fortran_matrix = matrix.reversed_axes();
+    let fmat_py = fortran_matrix.into_pyarray(py);
+    assert_eq!(fmat_py.as_array(), array![[0, 2, 4, 6], [1, 3, 5, 7]],);
+    pyo3::py_run!(py, fmat_py, "assert fmat_py.flags['F_CONTIGUOUS']")
+}