use std::{mem, os::raw::c_int, ptr};
use ndarray::{ArrayBase, Data, Dim, Dimension, IntoDimension, Ix1, OwnedRepr};
use pyo3::Python;
use crate::array::PyArray;
use crate::dtype::Element;
use crate::error::MAX_DIMENSIONALITY_ERR;
use crate::npyffi::{self, npy_intp};
use crate::sealed::Sealed;
use crate::slice_container::PySliceContainer;
pub trait IntoPyArray {
type Item: Element;
type Dim: Dimension;
fn into_pyarray<'py>(self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim>;
}
impl<T: Element> IntoPyArray for Box<[T]> {
type Item = T;
type Dim = Ix1;
fn into_pyarray<'py>(self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
let container = PySliceContainer::from(self);
let dims = Dim([container.len]);
let strides = [mem::size_of::<T>() as npy_intp];
let data_ptr = container.ptr as *mut T;
unsafe { PyArray::from_raw_parts(py, dims, strides.as_ptr(), data_ptr, container) }
}
}
impl<T: Element> IntoPyArray for Vec<T> {
type Item = T;
type Dim = Ix1;
fn into_pyarray<'py>(mut self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
let dims = Dim([self.len()]);
let strides = [mem::size_of::<T>() as npy_intp];
let data_ptr = self.as_mut_ptr();
unsafe {
PyArray::from_raw_parts(
py,
dims,
strides.as_ptr(),
data_ptr,
PySliceContainer::from(self),
)
}
}
}
impl<A, D> IntoPyArray for ArrayBase<OwnedRepr<A>, D>
where
A: Element,
D: Dimension,
{
type Item = A;
type Dim = D;
fn into_pyarray<'py>(self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
PyArray::from_owned_array(py, self)
}
}
pub trait ToPyArray {
type Item: Element;
type Dim: Dimension;
fn to_pyarray<'py>(&self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim>;
}
impl<T: Element> ToPyArray for [T] {
type Item = T;
type Dim = Ix1;
fn to_pyarray<'py>(&self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
PyArray::from_slice(py, self)
}
}
impl<S, D, A> ToPyArray for ArrayBase<S, D>
where
S: Data<Elem = A>,
D: Dimension,
A: Element,
{
type Item = A;
type Dim = D;
fn to_pyarray<'py>(&self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
let len = self.len();
match self.order() {
Some(flag) if A::IS_COPY => {
let strides = self.npy_strides();
unsafe {
let array = PyArray::new_uninit(py, self.raw_dim(), strides.as_ptr(), flag);
ptr::copy_nonoverlapping(self.as_ptr(), array.data(), len);
array
}
}
_ => {
let dim = self.raw_dim();
unsafe {
let array = PyArray::<A, _>::new(py, dim, false);
let mut data_ptr = array.data();
for item in self.iter() {
data_ptr.write(item.clone());
data_ptr = data_ptr.add(1);
}
array
}
}
}
}
}
#[cfg(feature = "nalgebra")]
impl<N, R, C, S> ToPyArray for nalgebra::Matrix<N, R, C, S>
where
N: nalgebra::Scalar + Element,
R: nalgebra::Dim,
C: nalgebra::Dim,
S: nalgebra::Storage<N, R, C>,
{
type Item = N;
type Dim = crate::Ix2;
fn to_pyarray<'py>(&self, py: Python<'py>) -> &'py PyArray<Self::Item, Self::Dim> {
unsafe {
let array = PyArray::<N, _>::new(py, (self.nrows(), self.ncols()), true);
let mut data_ptr = array.data();
if self.data.is_contiguous() {
ptr::copy_nonoverlapping(self.data.ptr(), data_ptr, self.len());
} else {
for item in self.iter() {
data_ptr.write(item.clone());
data_ptr = data_ptr.add(1);
}
}
array
}
}
}
pub(crate) trait ArrayExt {
fn npy_strides(&self) -> [npyffi::npy_intp; 32];
fn order(&self) -> Option<c_int>;
}
impl<A, S, D> ArrayExt for ArrayBase<S, D>
where
S: Data<Elem = A>,
D: Dimension,
{
fn npy_strides(&self) -> [npyffi::npy_intp; 32] {
let strides = self.strides();
let itemsize = mem::size_of::<A>() as isize;
assert!(strides.len() <= 32, "{}", MAX_DIMENSIONALITY_ERR);
let mut new_strides = [0; 32];
for i in 0..strides.len() {
new_strides[i] = (strides[i] * itemsize) as npyffi::npy_intp;
}
new_strides
}
fn order(&self) -> Option<c_int> {
if self.is_standard_layout() {
Some(npyffi::NPY_ORDER::NPY_CORDER as _)
} else if self.ndim() > 1 && self.raw_view().reversed_axes().is_standard_layout() {
Some(npyffi::NPY_ORDER::NPY_FORTRANORDER as _)
} else {
None
}
}
}
pub trait ToNpyDims: Dimension + Sealed {
#[doc(hidden)]
fn ndim_cint(&self) -> c_int {
self.ndim() as c_int
}
#[doc(hidden)]
fn as_dims_ptr(&mut self) -> *mut npyffi::npy_intp {
self.slice_mut().as_ptr() as *mut npyffi::npy_intp
}
#[doc(hidden)]
fn to_npy_dims(&mut self) -> npyffi::PyArray_Dims {
npyffi::PyArray_Dims {
ptr: self.as_dims_ptr(),
len: self.ndim_cint(),
}
}
}
impl<D> ToNpyDims for D where D: Dimension {}
pub trait NpyIndex: IntoDimension + Sealed {
#[doc(hidden)]
fn get_checked<T>(self, dims: &[usize], strides: &[isize]) -> Option<isize>;
#[doc(hidden)]
fn get_unchecked<T>(self, strides: &[isize]) -> isize;
}
impl<D: IntoDimension> Sealed for D {}
impl<D: IntoDimension> NpyIndex for D {
fn get_checked<T>(self, dims: &[usize], strides: &[isize]) -> Option<isize> {
let indices = self.into_dimension();
let indices = indices.slice();
if indices.len() != dims.len() {
return None;
}
if indices.iter().zip(dims).any(|(i, d)| i >= d) {
return None;
}
Some(get_unchecked_impl::<T>(indices, strides))
}
fn get_unchecked<T>(self, strides: &[isize]) -> isize {
let indices = self.into_dimension();
let indices = indices.slice();
get_unchecked_impl::<T>(indices, strides)
}
}
fn get_unchecked_impl<T>(indices: &[usize], strides: &[isize]) -> isize {
let size = mem::size_of::<T>() as isize;
indices
.iter()
.zip(strides)
.map(|(&i, stride)| stride * i as isize / size)
.sum()
}