/***************************************************************************

* Copyright (c) 2016, Johan Mabille and Sylvain Corlay *

* *

* Distributed under the terms of the BSD 3-Clause License. *

* *

* The full license is in the file LICENSE, distributed with this software. *

****************************************************************************/

#ifndef PY_ARRAY_HPP

#define PY_ARRAY_HPP

#include <algorithm>

#include <cstddef>

#include <vector>

#include "xtensor/xbuffer_adaptor.hpp"

#include "xtensor/xiterator.hpp"

#include "xtensor/xsemantic.hpp"

#include "pycontainer.hpp"

#include "pystrides_adaptor.hpp"

#include "xtensor_type_caster_base.hpp"

namespace xt

{

template <class T, layout_type L = layout_type::dynamic>

class pyarray;

}

namespace pybind11

{

namespace detail

{

template <class T, xt::layout_type L>

struct handle_type_name<xt::pyarray<T, L>>

{

static PYBIND11_DESCR name()

{

return _("numpy.ndarray[") + npy_format_descriptor<T>::name() + _("]");

}

};

template <typename T, xt::layout_type L>

struct pyobject_caster<xt::pyarray<T, L>>

{

using type = xt::pyarray<T, L>;

bool load(handle src, bool convert)

{

if (!convert)

{

if (!xt::detail::check_array<T>(src))

{

return false;

}

}

value = type::ensure(src);

return static_cast<bool>(value);

}

static handle cast(const handle& src, return_value_policy, handle)

{

return src.inc_ref();

}

PYBIND11_TYPE_CASTER(type, handle_type_name<type>::name());

};

// Type caster for casting ndarray to xexpression<pyarray>

template <typename T, xt::layout_type L>

struct type_caster<xt::xexpression<xt::pyarray<T, L>>> : pyobject_caster<xt::pyarray<T, L>>

{

using Type = xt::xexpression<xt::pyarray<T, L>>;

operator Type&()

{

return this->value;

}

operator const Type&()

{

return this->value;

}

};

// Type caster for casting xarray to ndarray

template <class T, xt::layout_type L>

struct type_caster<xt::xarray<T, L>> : xtensor_type_caster_base<xt::xarray<T, L>>

{

};

}

}

namespace xt

{

/**************************

* pybackstrides_iterator *

**************************/

template <class B>

class pybackstrides_iterator

{

public:

using self_type = pybackstrides_iterator<B>;

using value_type = typename B::value_type;

using pointer = const value_type*;

using reference = value_type;

using difference_type = std::ptrdiff_t;

using iterator_category = std::random_access_iterator_tag;

inline pybackstrides_iterator(const B* b, std::size_t offset)

: p_b(b), m_offset(offset)

{

}

inline reference operator*() const

{

return p_b->operator[](m_offset);

}

inline pointer operator->() const

{

// Returning the address of a temporary

value_type res = p_b->operator[](m_offset);

return &res;

}

inline reference operator[](difference_type n) const

{

return p_b->operator[](m_offset + n);

}

inline self_type& operator++()

{

++m_offset;

return *this;

}

inline self_type& operator--()

{

--m_offset;

return *this;

}

inline self_type operator++(int)

{

self_type tmp(*this);

++m_offset;

return tmp;

}

inline self_type operator--(int)

{

self_type tmp(*this);

--m_offset;

return tmp;

}

inline self_type& operator+=(difference_type n)

{

m_offset += n;

return *this;

}

inline self_type& operator-=(difference_type n)

{

m_offset -= n;

return *this;

}

inline self_type operator+(difference_type n) const

{

return self_type(p_b, m_offset + n);

}

inline self_type operator-(difference_type n) const

{

return self_type(p_b, m_offset - n);

}

inline self_type operator-(const self_type& rhs) const

{

self_type tmp(*this);

tmp -= (m_offset - rhs.m_offset);

return tmp;

}

inline std::size_t offset() const

{

return m_offset;

}

private:

const B* p_b;

std::size_t m_offset;

};

template <class B>

inline bool operator==(const pybackstrides_iterator<B>& lhs,

const pybackstrides_iterator<B>& rhs)

{

return lhs.offset() == rhs.offset();

}

template <class B>

inline bool operator!=(const pybackstrides_iterator<B>& lhs,

const pybackstrides_iterator<B>& rhs)

{

return !(lhs == rhs);

}

template <class B>

inline bool operator<(const pybackstrides_iterator<B>& lhs,

const pybackstrides_iterator<B>& rhs)

{

return lhs.offset() < rhs.offset();

}

template <class B>

inline bool operator<=(const pybackstrides_iterator<B>& lhs,

const pybackstrides_iterator<B>& rhs)

{

return (lhs < rhs) || (lhs == rhs);

}

template <class B>

inline bool operator>(const pybackstrides_iterator<B>& lhs,

const pybackstrides_iterator<B>& rhs)

{

return !(lhs <= rhs);

}

template <class B>

inline bool operator>=(const pybackstrides_iterator<B>& lhs,

const pybackstrides_iterator<B>& rhs)

{

return !(lhs < rhs);

}

template <class A>

class pyarray_backstrides

{

public:

using self_type = pyarray_backstrides<A>;

using array_type = A;

using value_type = typename array_type::size_type;

using const_reference = value_type;

using const_pointer = const value_type*;

using size_type = typename array_type::size_type;

using difference_type = typename array_type::difference_type;

using const_iterator = pybackstrides_iterator<self_type>;

pyarray_backstrides() = default;

pyarray_backstrides(const array_type& a);

bool empty() const;

size_type size() const;

value_type operator[](size_type i) const;

const_reference front() const;

const_reference back() const;

const_iterator begin() const;

const_iterator end() const;

const_iterator cbegin() const;

const_iterator cend() const;

private:

const array_type* p_a;

};

template <class T, layout_type L>

struct xiterable_inner_types<pyarray<T, L>>

: xcontainer_iterable_types<pyarray<T, L>>

{

};

template <class T, layout_type L>

struct xcontainer_inner_types<pyarray<T, L>>

{

using storage_type = xbuffer_adaptor<T*>;

using shape_type = std::vector<typename storage_type::size_type>;

using strides_type = std::vector<typename storage_type::difference_type>;

using backstrides_type = pyarray_backstrides<pyarray<T, L>>;

using inner_shape_type = xbuffer_adaptor<std::size_t*>;

using inner_strides_type = pystrides_adaptor<sizeof(T)>;

using inner_backstrides_type = backstrides_type;

using temporary_type = pyarray<T, L>;

static constexpr layout_type layout = L;

};

/**

* @class pyarray

* @brief Multidimensional container providing the xtensor container semantics to a numpy array.

*

* pyarray is similar to the xarray container in that it has a dynamic dimensionality. Reshapes of

* a pyarray container are reflected in the underlying numpy array.

*

* @tparam T The type of the element stored in the pyarray.

* @sa pytensor

*/

template <class T, layout_type L>

class pyarray : public pycontainer<pyarray<T, L>>,

public xcontainer_semantic<pyarray<T, L>>

{

public:

using self_type = pyarray<T, L>;

using semantic_base = xcontainer_semantic<self_type>;

using base_type = pycontainer<self_type>;

using storage_type = typename base_type::storage_type;

using value_type = typename base_type::value_type;

using reference = typename base_type::reference;

using const_reference = typename base_type::const_reference;

using pointer = typename base_type::pointer;

using size_type = typename base_type::size_type;

using shape_type = typename base_type::shape_type;

using strides_type = typename base_type::strides_type;

using backstrides_type = typename base_type::backstrides_type;

using inner_shape_type = typename base_type::inner_shape_type;

using inner_strides_type = typename base_type::inner_strides_type;

using inner_backstrides_type = typename base_type::inner_backstrides_type;

pyarray();

pyarray(const value_type& t);

pyarray(nested_initializer_list_t<T, 1> t);

pyarray(nested_initializer_list_t<T, 2> t);

pyarray(nested_initializer_list_t<T, 3> t);

pyarray(nested_initializer_list_t<T, 4> t);

pyarray(nested_initializer_list_t<T, 5> t);

pyarray(pybind11::handle h, pybind11::object::borrowed_t);

pyarray(pybind11::handle h, pybind11::object::stolen_t);

pyarray(const pybind11::object& o);

explicit pyarray(const shape_type& shape, layout_type l = layout_type::row_major);

explicit pyarray(const shape_type& shape, const_reference value, layout_type l = layout_type::row_major);

explicit pyarray(const shape_type& shape, const strides_type& strides, const_reference value);

explicit pyarray(const shape_type& shape, const strides_type& strides);

template <class S = shape_type>

static pyarray from_shape(S&& s);

pyarray(const self_type& rhs);

self_type& operator=(const self_type& rhs);

pyarray(self_type&&) = default;

self_type& operator=(self_type&& e) = default;

template <class E>

pyarray(const xexpression<E>& e);

template <class E>

self_type& operator=(const xexpression<E>& e);

using base_type::begin;

using base_type::end;

static self_type ensure(pybind11::handle h);

static bool check_(pybind11::handle h);

private:

inner_shape_type m_shape;

inner_strides_type m_strides;

mutable inner_backstrides_type m_backstrides;

storage_type m_storage;

void init_array(const shape_type& shape, const strides_type& strides);

void init_from_python();

const inner_shape_type& shape_impl() const noexcept;

const inner_strides_type& strides_impl() const noexcept;

const inner_backstrides_type& backstrides_impl() const noexcept;

storage_type& storage_impl() noexcept;

const storage_type& storage_impl() const noexcept;

friend class xcontainer<pyarray<T, L>>;

friend class pycontainer<pyarray<T, L>>;

};

/**************************************

* pyarray_backstrides implementation *

**************************************/

template <class A>

inline pyarray_backstrides<A>::pyarray_backstrides(const array_type& a)

: p_a(&a)

{

}

template <class A>

inline bool pyarray_backstrides<A>::empty() const

{

return p_a->dimension() == 0;

}

template <class A>

inline auto pyarray_backstrides<A>::size() const -> size_type

{

return p_a->dimension();

}

template <class A>

inline auto pyarray_backstrides<A>::operator[](size_type i) const -> value_type

{

value_type sh = p_a->shape()[i];

value_type res = sh == 1 ? 0 : (sh - 1) * p_a->strides()[i];

return res;

}

template <class A>

inline auto pyarray_backstrides<A>::front() const -> const_reference

{

value_type sh = p_a->shape()[0];

value_type res = sh == 1 ? 0 : (sh - 1) * p_a->strides()[0];

return res;

}

template <class A>

inline auto pyarray_backstrides<A>::back() const -> const_reference

{

auto index = p_a->size() - 1;

value_type sh = p_a->shape()[index];

value_type res = sh == 1 ? 0 : (sh - 1) * p_a->strides()[index];

return res;

}

template <class A>

inline auto pyarray_backstrides<A>::begin() const -> const_iterator

{

return cbegin();

}

template <class A>

inline auto pyarray_backstrides<A>::end() const -> const_iterator

{

return cend();

}

template <class A>

inline auto pyarray_backstrides<A>::cbegin() const -> const_iterator

{

return const_iterator(this, 0);

}

template <class A>

inline auto pyarray_backstrides<A>::cend() const -> const_iterator

{

return const_iterator(this, size());

}

/**************************

* pyarray implementation *

**************************/

/**

* @name Constructors

*/

//@{

template <class T, layout_type L>

inline pyarray<T, L>::pyarray()

: base_type()

{

// TODO: avoid allocation

shape_type shape = xtl::make_sequence<shape_type>(0, size_type(1));

strides_type strides = xtl::make_sequence<strides_type>(0, size_type(0));

init_array(shape, strides);

detail::default_initialize(m_storage);

}

/**

* Allocates a pyarray with nested initializer lists.

*/

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(const value_type& t)

: base_type()

{

base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);

nested_copy(m_storage.begin(), t);

}

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(nested_initializer_list_t<T, 1> t)

: base_type()

{

base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);

nested_copy(m_storage.begin(), t);

}

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(nested_initializer_list_t<T, 2> t)

: base_type()

{

base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);

nested_copy(m_storage.begin(), t);

}

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(nested_initializer_list_t<T, 3> t)

: base_type()

{

base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);

nested_copy(m_storage.begin(), t);

}

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(nested_initializer_list_t<T, 4> t)

: base_type()

{

base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);

nested_copy(m_storage.begin(), t);

}

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(nested_initializer_list_t<T, 5> t)

: base_type()

{

base_type::resize(xt::shape<shape_type>(t), layout_type::row_major);

nested_copy(m_storage.begin(), t);

}

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(pybind11::handle h, pybind11::object::borrowed_t b)

: base_type(h, b)

{

init_from_python();

}

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(pybind11::handle h, pybind11::object::stolen_t s)

: base_type(h, s)

{

init_from_python();

}

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(const pybind11::object& o)

: base_type(o)

{

init_from_python();

}

/**

* Allocates an uninitialized pyarray with the specified shape and

* layout.

* @param shape the shape of the pyarray

* @param l the layout of the pyarray

*/

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(const shape_type& shape, layout_type l)

: base_type()

{

strides_type strides(shape.size());

compute_strides(shape, l, strides);

init_array(shape, strides);

}

/**

* Allocates a pyarray with the specified shape and layout. Elements

* are initialized to the specified value.

* @param shape the shape of the pyarray

* @param value the value of the elements

* @param l the layout of the pyarray

*/

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(const shape_type& shape, const_reference value, layout_type l)

: base_type()

{

strides_type strides(shape.size());

compute_strides(shape, l, strides);

init_array(shape, strides);

std::fill(m_storage.begin(), m_storage.end(), value);

}

/**

* Allocates an uninitialized pyarray with the specified shape and strides.

* Elements are initialized to the specified value.

* @param shape the shape of the pyarray

* @param strides the strides of the pyarray

* @param value the value of the elements

*/

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(const shape_type& shape, const strides_type& strides, const_reference value)

: base_type()

{

init_array(shape, strides);

std::fill(m_storage.begin(), m_storage.end(), value);

}

/**

* Allocates an uninitialized pyarray with the specified shape and strides.

* @param shape the shape of the pyarray

* @param strides the strides of the pyarray

*/

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(const shape_type& shape, const strides_type& strides)

: base_type()

{

init_array(shape, strides);

}

/**

* Allocates and returns an pyarray with the specified shape.

* @param shape the shape of the pyarray

*/

template <class T, layout_type L>

template <class S>

inline pyarray<T, L> pyarray<T, L>::from_shape(S&& shape)

{

auto shp = xtl::forward_sequence<shape_type>(shape);

return self_type(shp);

}

//@}

/**

* @name Copy semantic

*/

//@{

/**

* The copy constructor.

*/

template <class T, layout_type L>

inline pyarray<T, L>::pyarray(const self_type& rhs)

: base_type(), semantic_base(rhs)

{

auto tmp = pybind11::reinterpret_steal<pybind11::object>(

PyArray_NewLikeArray(rhs.python_array(), NPY_KEEPORDER, nullptr, 1));

if (!tmp)

{

throw std::runtime_error("NumPy: unable to create ndarray");

}

this->m_ptr = tmp.release().ptr();

init_from_python();

std::copy(rhs.storage().cbegin(), rhs.storage().cend(), this->storage().begin());

}

/**

* The assignment operator.

*/

template <class T, layout_type L>

inline auto pyarray<T, L>::operator=(const self_type& rhs) -> self_type&

{

self_type tmp(rhs);

*this = std::move(tmp);

return *this;

}

//@}

/**

* @name Extended copy semantic

*/

//@{

/**

* The extended copy constructor.

*/

template <class T, layout_type L>

template <class E>

inline pyarray<T, L>::pyarray(const xexpression<E>& e)

: base_type()

{

// TODO: prevent intermediary shape allocation

shape_type shape = xtl::forward_sequence<shape_type>(e.derived_cast().shape());

strides_type strides = xtl::make_sequence<strides_type>(shape.size(), size_type(0));

compute_strides(shape, layout_type::row_major, strides);

init_array(shape, strides);

semantic_base::assign(e);

}

/**

* The extended assignment operator.

*/

template <class T, layout_type L>

template <class E>

inline auto pyarray<T, L>::operator=(const xexpression<E>& e) -> self_type&

{

return semantic_base::operator=(e);

}

//@}

template <class T, layout_type L>

inline auto pyarray<T, L>::ensure(pybind11::handle h) -> self_type

{

return base_type::ensure(h);

}

template <class T, layout_type L>

inline bool pyarray<T, L>::check_(pybind11::handle h)

{

return base_type::check_(h);

}

template <class T, layout_type L>

inline void pyarray<T, L>::init_array(const shape_type& shape, const strides_type& strides)

{

strides_type adapted_strides(strides);

std::transform(strides.begin(), strides.end(), adapted_strides.begin(),

[](auto v) { return sizeof(T) * v; });

int flags = NPY_ARRAY_ALIGNED;

if (!std::is_const<T>::value)

{

flags |= NPY_ARRAY_WRITEABLE;

}

auto dtype = pybind11::detail::npy_format_descriptor<T>::dtype();

npy_intp* shape_data = reinterpret_cast<npy_intp*>(const_cast<size_type*>(shape.data()));

npy_intp* strides_data = reinterpret_cast<npy_intp*>(adapted_strides.data());

auto tmp = pybind11::reinterpret_steal<pybind11::object>(

PyArray_NewFromDescr(&PyArray_Type, (PyArray_Descr*) dtype.release().ptr(), static_cast<int>(shape.size()), shape_data, strides_data,

nullptr, flags, nullptr));

if (!tmp)

{

throw std::runtime_error("NumPy: unable to create ndarray");

}

this->m_ptr = tmp.release().ptr();

init_from_python();

}

template <class T, layout_type L>

inline void pyarray<T, L>::init_from_python()

{

m_shape = inner_shape_type(reinterpret_cast<size_type*>(PyArray_SHAPE(this->python_array())),

static_cast<size_type>(PyArray_NDIM(this->python_array())));

m_strides = inner_strides_type(reinterpret_cast<size_type*>(PyArray_STRIDES(this->python_array())),

static_cast<size_type>(PyArray_NDIM(this->python_array())));

if (L != layout_type::dynamic && !do_strides_match(m_shape, m_strides, L))

{

throw std::runtime_error("NumPy: passing container with bad strides for layout (is it a view?).");

}

m_backstrides = backstrides_type(*this);

m_storage = storage_type(reinterpret_cast<pointer>(PyArray_DATA(this->python_array())),

this->get_min_stride() * static_cast<size_type>(PyArray_SIZE(this->python_array())));

}

template <class T, layout_type L>

inline auto pyarray<T, L>::shape_impl() const noexcept -> const inner_shape_type&

{

return m_shape;

}

template <class T, layout_type L>

inline auto pyarray<T, L>::strides_impl() const noexcept -> const inner_strides_type&

{

return m_strides;

}

template <class T, layout_type L>

inline auto pyarray<T, L>::backstrides_impl() const noexcept -> const inner_backstrides_type&

{

// m_backstrides wraps the numpy array backstrides, which is a raw pointer.

// The address of the raw pointer stored in the wrapper would be invalidated when the pyarray is copied.

// Hence, we build a new backstrides object (cheap wrapper around the underlying pointer) upon access.

m_backstrides = backstrides_type(*this);

return m_backstrides;

}

template <class T, layout_type L>

inline auto pyarray<T, L>::storage_impl() noexcept -> storage_type&

{

return m_storage;

}

template <class T, layout_type L>

inline auto pyarray<T, L>::storage_impl() const noexcept -> const storage_type&

{

return m_storage;

}

}

#endif