/***************************************************************************
* 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
#include
#include "pybind11/numpy.h"
#include "xtensor/xexpression.hpp"
#include "xtensor/xsemantic.hpp"
#include "xtensor/xiterator.hpp"
namespace xt
{
using pybind_array = pybind11::array;
using buffer_info = pybind11::buffer_info;
/***********************
* pyarray declaration *
***********************/
template
class pyarray;
template
struct array_inner_types
{
using temporary_type = pyarray;
};
template
class pyarray_backstrides
{
public:
using array_type = A;
using value_type = typename array_type::size_type;
using size_type = typename array_type::size_type;
pyarray_backstrides(const A& a);
value_type operator[](size_type i) const;
private:
const pybind_array* p_a;
};
/**
* @class pyarray
* @brief Wrapper on the Python buffer protocol.
*/
template
class pyarray : public pybind_array,
public xarray_semantic
{
public:
using self_type = pyarray;
using base_type = pybind_array;
using semantic_base = xarray_semantic;
using value_type = T;
using reference = T&;
using const_reference = const T&;
using pointer = T*;
using const_pointer = const T*;
using size_type = std::size_t;
using difference_type = std::ptrdiff_t;
using stepper = xstepper;
using const_stepper = xstepper;
using iterator = xiterator;
using const_iterator = xiterator;
using storage_iterator = T*;
using const_storage_iterator = const T*;
using shape_type = xshape;
using strides_type = xstrides;
using backstrides_type = pyarray_backstrides;
using closure_type = const self_type&;
PYBIND11_OBJECT_CVT(pyarray, pybind_array, is_non_null, m_ptr = ensure_(m_ptr));
pyarray();
explicit pyarray(const buffer_info& info);
pyarray(const xshape& shape,
const xstrides& strides,
const T* ptr = nullptr,
handle base = handle());
explicit pyarray(const xshape& shape,
const T* ptr = nullptr,
handle base = handle());
explicit pyarray(size_type count,
const T* ptr = nullptr,
handle base = handle());
size_type dimension() const;
const shape_type& shape() const;
const strides_type& strides() const;
backstrides_type backstrides() const;
void reshape(const shape_type& shape);
void reshape(const shape_type& shape, layout l);
void reshape(const shape_type& shape, const strides_type& strides);
template
reference operator()(Args... args);
template
const_reference operator()(Args... args) const;
template
pointer data(Args... args);
template
const_pointer data(Args... args) const;
bool broadcast_shape(shape_type& shape) const;
bool is_trivial_broadcast(const strides_type& strides) const;
iterator begin();
iterator end();
const_iterator begin() const;
const_iterator end() const;
const_iterator cbegin() const;
const_iterator cend() const;
iterator xbegin(const shape_type& shape);
iterator xend(const shape_type& shape);
const_iterator xbegin(const shape_type& shape) const;
const_iterator xend(const shape_type& shape) const;
const_iterator cxbegin(const shape_type& shape) const;
const_iterator cxend(const shape_type& shape) const;
stepper stepper_begin(const shape_type& shape);
stepper stepper_end(const shape_type& shape);
const_stepper stepper_begin(const shape_type& shape) const;
const_stepper stepper_end(const shape_type& shape) const;
storage_iterator storage_begin();
storage_iterator storage_end();
const_storage_iterator storage_begin() const;
const_storage_iterator storage_end() const;
template
pyarray(const xexpression& e);
template
pyarray& operator=(const xexpression& e);
private:
template
auto index_at(Args... args) const -> size_type;
static constexpr auto itemsize() -> size_type;
static bool is_non_null(PyObject* ptr);
static PyObject *ensure_(PyObject* ptr);
mutable shape_type m_shape;
mutable strides_type m_strides;
};
/**************************************
* pyarray_backstrides implementation *
**************************************/
template
inline pyarray_backstrides
::pyarray_backstrides(const A& a)
: p_a(&a)
{
}
template
inline auto pyarray_backstrides::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] / sizeof(typename A::value_type);
return res;
}
/**************************
* pyarray implementation *
**************************/
template
inline pyarray::pyarray()
: pybind_array()
{
}
template
inline pyarray::pyarray(const buffer_info& info)
: pybind_array(info)
{
}
template
inline pyarray::pyarray(const xshape& shape,
const xstrides& strides,
const T *ptr,
handle base)
: pybind_array(shape, strides, ptr, base)
{
}
template
inline pyarray::pyarray(const xshape& shape,
const T* ptr,
handle base)
: pybind_array(shape, ptr, base)
{
}
template
inline pyarray::pyarray(size_type count,
const T* ptr,
handle base)
: pybind_array(count, ptr, base)
{
}
template
inline auto pyarray::dimension() const -> size_type
{
return pybind_array::ndim();
}
template
inline auto pyarray::shape() const -> const shape_type&
{
// Until we have the CRTP on shape types, we copy the shape.
m_shape.resize(dimension());
std::copy(pybind_array::shape(), pybind_array::shape() + dimension(), m_shape.begin());
return m_shape;
}
template
inline auto pyarray::strides() const -> const strides_type&
{
m_strides.resize(dimension());
std::transform(pybind_array::strides(), pybind_array::strides() + dimension(), m_strides.begin(),
[](size_type str) { return str / sizeof(value_type); });
return m_strides;
}
template
inline auto pyarray::backstrides() const -> backstrides_type
{
backstrides_type tmp(*this);
return tmp;
}
template
void pyarray::reshape(const shape_type& shape)
{
if (!m_ptr || shape.size() != dimension() || !std::equal(shape.begin(), shape.end(), pybind_array::shape()))
{
reshape(shape, layout::row_major);
}
}
template
void pyarray::reshape(const shape_type& shape, layout l)
{
strides_type strides(shape.size());
size_type data_size = sizeof(value_type);
if (l == layout::row_major)
{
for (size_type i = strides.size(); i != 0; --i)
{
strides[i - 1] = data_size;
data_size = strides[i - 1] * shape[i - 1];
if (shape[i - 1] == 1)
{
strides[i - 1] = 0;
}
}
}
else
{
for (size_type i = 0; i < strides.size(); ++i)
{
strides[i] = data_size;
data_size = strides[i] * shape[i];
if (shape[i] == 1)
{
strides[i] = 0;
}
}
}
reshape(shape, strides);
}
template
void pyarray::reshape(const shape_type& shape, const strides_type& strides)
{
self_type tmp(shape, strides);
*this = std::move(tmp);
}
template
template
inline auto pyarray::operator()(Args... args) -> reference
{
if (sizeof...(args) != dimension())
{
pybind_array::fail_dim_check(sizeof...(args), "index dimension mismatch");
}
// not using pybind_array::offset_at() / index_at() here so as to avoid another dimension check.
return *(static_cast(pybind_array::mutable_data()) + pybind_array::get_byte_offset(args...) / itemsize());
}
template
template
inline auto pyarray::operator()(Args... args) const -> const_reference
{
if (sizeof...(args) != dimension())
{
pybind_array::fail_dim_check(sizeof...(args), "index dimension mismatch");
}
// not using pybind_array::offset_at() / index_at() here so as to avoid another dimension check.
return *(static_cast(pybind_array::data()) + pybind_array::get_byte_offset(args...) / itemsize());
}
template
template
inline auto pyarray::data(Args... args) -> pointer
{
return static_cast(pybind_array::mutable_data(args...));
}
template
template
inline auto pyarray::data(Args... args) const -> const_pointer
{
return static_cast(pybind_array::data(args...));
}
template
bool pyarray::broadcast_shape(shape_type& shape) const
{
return xt::broadcast_shape(this->shape(), shape);
}
template
bool pyarray::is_trivial_broadcast(const strides_type& strides) const
{
return strides.size() == dimension() &&
std::equal(strides.begin(), strides.end(), this->strides().begin());
}
template
inline auto pyarray::begin() -> iterator
{
return xbegin(shape());
}
template
inline auto pyarray::end() -> iterator
{
return xend(shape());
}
template
inline auto pyarray::begin() const -> const_iterator
{
return xbegin(shape());
}
template
inline auto pyarray::end() const -> const_iterator
{
return xend(shape());
}
template
inline auto pyarray::cbegin() const -> const_iterator
{
return begin();
}
template
inline auto pyarray::cend() const -> const_iterator
{
return end();
}
template
inline auto pyarray::xbegin(const shape_type& shape) -> iterator
{
return iterator(stepper_begin(shape), shape);
}
template
inline auto pyarray::xend(const shape_type& shape) -> iterator
{
return iterator(stepper_end(shape), shape);
}
template
inline auto pyarray::xbegin(const shape_type& shape) const -> const_iterator
{
return const_iterator(stepper_begin(shape), shape);
}
template
inline auto pyarray::xend(const shape_type& shape) const -> const_iterator
{
return const_iterator(stepper_end(shape), shape);
}
template
inline auto pyarray::cxbegin(const shape_type& shape) const -> const_iterator
{
return xbegin(shape);
}
template
inline auto pyarray::cxend(const shape_type& shape) const -> const_iterator
{
return xend(shape);
}
template
inline auto pyarray::stepper_begin(const shape_type& shape) -> stepper
{
size_type offset = shape.size() - dimension();
return stepper(this, storage_begin(), offset);
}
template
inline auto pyarray::stepper_end(const shape_type& shape) -> stepper
{
size_type offset = shape.size() - dimension();
return stepper(this, storage_end(), offset);
}
template
inline auto pyarray::stepper_begin(const shape_type& shape) const -> const_stepper
{
size_type offset = shape.size() - dimension();
return const_stepper(this, storage_begin(), offset);
}
template
inline auto pyarray::stepper_end(const shape_type& shape) const -> const_stepper
{
size_type offset = shape.size() - dimension();
return const_stepper(this, storage_end(), offset);
}
template
inline auto pyarray::storage_begin() -> storage_iterator
{
return reinterpret_cast(PyArray_GET_(m_ptr, data));
}
template
inline auto pyarray::storage_end() -> storage_iterator
{
return storage_begin() + pybind_array::size();
}
template
inline auto pyarray::storage_begin() const -> const_storage_iterator
{
return reinterpret_cast(PyArray_GET_(m_ptr, data));
}
template
inline auto pyarray::storage_end() const -> const_storage_iterator
{
return storage_begin() + pybind_array::size();
}
template
template
inline pyarray::pyarray(const xexpression& e)
: pybind_array()
{
semantic_base::assign(e);
}
template
template
inline auto pyarray::operator=(const xexpression& e) -> self_type&
{
return semantic_base::operator=(e);
}
// Private methods
template
template
inline auto pyarray::index_at(Args... args) const -> size_type
{
return pybind_array::offset_at(args...) / itemsize();
}
template
constexpr auto pyarray::itemsize() -> size_type
{
return sizeof(value_type);
}
template
inline bool pyarray::is_non_null(PyObject* ptr)
{
return ptr != nullptr;
}
template
inline PyObject* pyarray::ensure_(PyObject* ptr)
{
if (ptr == nullptr)
{
return nullptr;
}
auto& api = pybind11::detail::npy_api::get();
PyObject *result = api.PyArray_FromAny_(ptr, pybind11::dtype::of().release().ptr(), 0, 0,
pybind11::detail::npy_api::NPY_ENSURE_ARRAY_ | ExtraFlags, nullptr);
if (!result)
{
PyErr_Clear();
}
Py_DECREF(ptr);
return result;
}
}
#endif