numpy/
convert.rs

1//! Defines conversion traits between Rust types and NumPy data types.
2
3use std::{mem, os::raw::c_int, ptr};
4
5use ndarray::{ArrayBase, Data, Dim, Dimension, IntoDimension, Ix1, OwnedRepr};
6use pyo3::{Bound, Python};
7
8use crate::array::{PyArray, PyArrayMethods};
9use crate::dtype::Element;
10use crate::error::MAX_DIMENSIONALITY_ERR;
11use crate::npyffi::{self, npy_intp};
12use crate::slice_container::PySliceContainer;
13
14/// Conversion trait from owning Rust types into [`PyArray`].
15///
16/// This trait takes ownership of `self`, which means it holds a pointer into the Rust heap.
17///
18/// In addition, some destructive methods like `resize` cannot be used with NumPy arrays constructed using this trait.
19///
20/// # Example
21///
22/// ```
23/// use numpy::{PyArray, IntoPyArray, PyArrayMethods};
24/// use pyo3::Python;
25///
26/// Python::with_gil(|py| {
27///     let py_array = vec![1, 2, 3].into_pyarray(py);
28///
29///     assert_eq!(py_array.readonly().as_slice().unwrap(), &[1, 2, 3]);
30///
31///     // Array cannot be resized when its data is owned by Rust.
32///     unsafe {
33///         assert!(py_array.resize(100).is_err());
34///     }
35/// });
36/// ```
37pub trait IntoPyArray: Sized {
38    /// The element type of resulting array.
39    type Item: Element;
40    /// The dimension type of the resulting array.
41    type Dim: Dimension;
42
43    /// Consumes `self` and moves its data into a NumPy array.
44    fn into_pyarray<'py>(self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>>;
45}
46
47impl<T: Element> IntoPyArray for Box<[T]> {
48    type Item = T;
49    type Dim = Ix1;
50
51    fn into_pyarray<'py>(self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
52        let container = PySliceContainer::from(self);
53        let dims = Dim([container.len]);
54        let strides = [mem::size_of::<T>() as npy_intp];
55        // The data pointer is derived only after dissolving `Box` into `PySliceContainer`
56        // to avoid unsound aliasing of Box<[T]> which is currently noalias,
57        // c.f. https://github.com/rust-lang/unsafe-code-guidelines/issues/326
58        let data_ptr = container.ptr as *mut T;
59        unsafe { PyArray::from_raw_parts(py, dims, strides.as_ptr(), data_ptr, container) }
60    }
61}
62
63impl<T: Element> IntoPyArray for Vec<T> {
64    type Item = T;
65    type Dim = Ix1;
66
67    fn into_pyarray<'py>(mut self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
68        let dims = Dim([self.len()]);
69        let strides = [mem::size_of::<T>() as npy_intp];
70        let data_ptr = self.as_mut_ptr();
71        unsafe {
72            PyArray::from_raw_parts(
73                py,
74                dims,
75                strides.as_ptr(),
76                data_ptr,
77                PySliceContainer::from(self),
78            )
79        }
80    }
81}
82
83impl<A, D> IntoPyArray for ArrayBase<OwnedRepr<A>, D>
84where
85    A: Element,
86    D: Dimension,
87{
88    type Item = A;
89    type Dim = D;
90
91    fn into_pyarray<'py>(self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
92        PyArray::from_owned_array(py, self)
93    }
94}
95
96/// Conversion trait from borrowing Rust types to [`PyArray`].
97///
98/// This trait takes `&self` by reference, which means it allocates in Python heap and then copies the elements there.
99///
100/// # Examples
101///
102/// ```
103/// use numpy::{PyArray, ToPyArray, PyArrayMethods};
104/// use pyo3::Python;
105///
106/// Python::with_gil(|py| {
107///     let py_array = vec![1, 2, 3].to_pyarray(py);
108///
109///     assert_eq!(py_array.readonly().as_slice().unwrap(), &[1, 2, 3]);
110/// });
111/// ```
112///
113/// Due to copying the elments, this method converts non-contiguous arrays to C-order contiguous arrays.
114///
115/// ```
116/// use numpy::prelude::*;
117/// use numpy::{PyArray, ToPyArray};
118/// use ndarray::{arr3, s};
119/// use pyo3::Python;
120///
121/// Python::with_gil(|py| {
122///     let array = arr3(&[[[1, 2, 3], [4, 5, 6]], [[7, 8, 9], [10, 11, 12]]]);
123///     let py_array = array.slice(s![.., 0..1, ..]).to_pyarray(py);
124///
125///     assert_eq!(py_array.readonly().as_array(), arr3(&[[[1, 2, 3]], [[7, 8, 9]]]));
126///     assert!(py_array.is_c_contiguous());
127/// });
128/// ```
129pub trait ToPyArray {
130    /// The element type of resulting array.
131    type Item: Element;
132    /// The dimension type of the resulting array.
133    type Dim: Dimension;
134
135    /// Copies the content pointed to by `&self` into a newly allocated NumPy array.
136    fn to_pyarray<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>>;
137}
138
139impl<T: Element> ToPyArray for [T] {
140    type Item = T;
141    type Dim = Ix1;
142
143    fn to_pyarray<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
144        PyArray::from_slice(py, self)
145    }
146}
147
148impl<S, D, A> ToPyArray for ArrayBase<S, D>
149where
150    S: Data<Elem = A>,
151    D: Dimension,
152    A: Element,
153{
154    type Item = A;
155    type Dim = D;
156
157    fn to_pyarray<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
158        let len = self.len();
159        match self.order() {
160            Some(flag) if A::IS_COPY => {
161                // if the array is contiguous, copy it by `copy_nonoverlapping`.
162                let strides = self.npy_strides();
163                unsafe {
164                    let array = PyArray::new_uninit(py, self.raw_dim(), strides.as_ptr(), flag);
165                    ptr::copy_nonoverlapping(self.as_ptr(), array.data(), len);
166                    array
167                }
168            }
169            _ => {
170                // if the array is not contiguous, copy all elements by `ArrayBase::iter`.
171                let dim = self.raw_dim();
172                unsafe {
173                    let array = PyArray::<A, _>::new(py, dim, false);
174                    let mut data_ptr = array.data();
175                    for item in self.iter() {
176                        data_ptr.write(item.clone_ref(py));
177                        data_ptr = data_ptr.add(1);
178                    }
179                    array
180                }
181            }
182        }
183    }
184}
185
186#[cfg(feature = "nalgebra")]
187impl<N, R, C, S> ToPyArray for nalgebra::Matrix<N, R, C, S>
188where
189    N: nalgebra::Scalar + Element,
190    R: nalgebra::Dim,
191    C: nalgebra::Dim,
192    S: nalgebra::Storage<N, R, C>,
193{
194    type Item = N;
195    type Dim = crate::Ix2;
196
197    /// Note that the NumPy array always has Fortran memory layout
198    /// matching the [memory layout][memory-layout] used by [`nalgebra`].
199    ///
200    /// [memory-layout]: https://nalgebra.org/docs/faq/#what-is-the-memory-layout-of-matrices
201    fn to_pyarray<'py>(&self, py: Python<'py>) -> Bound<'py, PyArray<Self::Item, Self::Dim>> {
202        unsafe {
203            let array = PyArray::<N, _>::new(py, (self.nrows(), self.ncols()), true);
204            let mut data_ptr = array.data();
205            if self.data.is_contiguous() {
206                ptr::copy_nonoverlapping(self.data.ptr(), data_ptr, self.len());
207            } else {
208                for item in self.iter() {
209                    data_ptr.write(item.clone_ref(py));
210                    data_ptr = data_ptr.add(1);
211                }
212            }
213            array
214        }
215    }
216}
217
218pub(crate) trait ArrayExt {
219    fn npy_strides(&self) -> [npyffi::npy_intp; 32];
220    fn order(&self) -> Option<c_int>;
221}
222
223impl<A, S, D> ArrayExt for ArrayBase<S, D>
224where
225    S: Data<Elem = A>,
226    D: Dimension,
227{
228    fn npy_strides(&self) -> [npyffi::npy_intp; 32] {
229        let strides = self.strides();
230        let itemsize = mem::size_of::<A>() as isize;
231
232        assert!(strides.len() <= 32, "{}", MAX_DIMENSIONALITY_ERR);
233
234        let mut new_strides = [0; 32];
235
236        for i in 0..strides.len() {
237            new_strides[i] = (strides[i] * itemsize) as npyffi::npy_intp;
238        }
239
240        new_strides
241    }
242
243    fn order(&self) -> Option<c_int> {
244        if self.is_standard_layout() {
245            Some(npyffi::NPY_ORDER::NPY_CORDER as _)
246        } else if self.ndim() > 1 && self.raw_view().reversed_axes().is_standard_layout() {
247            Some(npyffi::NPY_ORDER::NPY_FORTRANORDER as _)
248        } else {
249            None
250        }
251    }
252}
253
254/// Utility trait to specify the dimensions of an array.
255pub trait ToNpyDims: Dimension + Sealed {
256    #[doc(hidden)]
257    fn ndim_cint(&self) -> c_int {
258        self.ndim() as c_int
259    }
260    #[doc(hidden)]
261    fn as_dims_ptr(&mut self) -> *mut npyffi::npy_intp {
262        self.slice_mut().as_ptr() as *mut npyffi::npy_intp
263    }
264    #[doc(hidden)]
265    fn to_npy_dims(&mut self) -> npyffi::PyArray_Dims {
266        npyffi::PyArray_Dims {
267            ptr: self.as_dims_ptr(),
268            len: self.ndim_cint(),
269        }
270    }
271}
272
273mod sealed {
274    pub trait Sealed {}
275}
276
277use sealed::Sealed;
278
279impl<D> ToNpyDims for D where D: Dimension {}
280
281/// Trait implemented by types that can be used to index an array.
282///
283/// This is equivalent to [`ndarray::NdIndex`] but accounts for
284/// NumPy strides being in units of bytes instead of elements.
285///
286/// All types which implement [`IntoDimension`] implement this trait as well.
287/// This includes at least
288/// - [tuple](https://doc.rust-lang.org/stable/std/primitive.tuple.html)
289/// - [array](https://doc.rust-lang.org/stable/std/primitive.array.html)
290/// - [slice](https://doc.rust-lang.org/stable/std/primitive.slice.html)
291pub trait NpyIndex: IntoDimension + Sealed {
292    #[doc(hidden)]
293    fn get_checked<T>(self, dims: &[usize], strides: &[isize]) -> Option<isize>;
294    #[doc(hidden)]
295    fn get_unchecked<T>(self, strides: &[isize]) -> isize;
296}
297
298impl<D: IntoDimension> Sealed for D {}
299
300impl<D: IntoDimension> NpyIndex for D {
301    fn get_checked<T>(self, dims: &[usize], strides: &[isize]) -> Option<isize> {
302        let indices = self.into_dimension();
303        let indices = indices.slice();
304
305        if indices.len() != dims.len() {
306            return None;
307        }
308        if indices.iter().zip(dims).any(|(i, d)| i >= d) {
309            return None;
310        }
311
312        Some(get_unchecked_impl::<T>(indices, strides))
313    }
314
315    fn get_unchecked<T>(self, strides: &[isize]) -> isize {
316        let indices = self.into_dimension();
317        let indices = indices.slice();
318        get_unchecked_impl::<T>(indices, strides)
319    }
320}
321
322fn get_unchecked_impl<T>(indices: &[usize], strides: &[isize]) -> isize {
323    let size = mem::size_of::<T>() as isize;
324
325    indices
326        .iter()
327        .zip(strides)
328        .map(|(&i, stride)| stride * i as isize / size)
329        .sum()
330}
numpy/convert.rs

numpy/
convert.rs
