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

* Copyright (c) Wolf Vollprecht, Johan Mabille and Sylvain Corlay *

* Copyright (c) QuantStack *

* *

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

* *

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

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

#include <numeric>

#include "xtensor/core/xmath.hpp"

#include "xtensor/containers/xarray.hpp"

#include "xtensor/containers/xfixed.hpp"

#define FORCE_IMPORT_ARRAY

#include "xtensor-python/pyarray.hpp"

#include "xtensor-python/pytensor.hpp"

#include "xtensor-python/pyvectorize.hpp"

#include "xtensor/containers/xadapt.hpp"

#include "xtensor/views/xstrided_view.hpp"

namespace py = pybind11;

using complex_t = std::complex<double>;

// Examples

double example1(xt::pyarray<double>& m)

{

return m(0);

}

xt::pyarray<double> example2(xt::pyarray<double>& m)

{

return m + 2;

}

xt::xarray<int> example3_xarray(const xt::xarray<int>& m)

{

return xt::transpose(m) + 2;

}

xt::xarray<int, xt::layout_type::column_major> example3_xarray_colmajor(

const xt::xarray<int, xt::layout_type::column_major>& m)

{

return xt::transpose(m) + 2;

}

xt::xtensor<int, 3> example3_xtensor3(const xt::xtensor<int, 3>& m)

{

return xt::transpose(m) + 2;

}

xt::xtensor<int, 2> example3_xtensor2(const xt::xtensor<int, 2>& m)

{

return xt::transpose(m) + 2;

}

xt::xtensor<int, 2, xt::layout_type::column_major> example3_xtensor2_colmajor(

const xt::xtensor<int, 2, xt::layout_type::column_major>& m)

{

return xt::transpose(m) + 2;

}

xt::xtensor_fixed<int, xt::xshape<4, 3, 2>> example3_xfixed3(const xt::xtensor_fixed<int, xt::xshape<2, 3, 4>>& m)

{

return xt::transpose(m) + 2;

}

xt::xtensor_fixed<int, xt::xshape<3, 2>> example3_xfixed2(const xt::xtensor_fixed<int, xt::xshape<2, 3>>& m)

{

return xt::transpose(m) + 2;

}

xt::xtensor_fixed<int, xt::xshape<3, 2>, xt::layout_type::column_major> example3_xfixed2_colmajor(

const xt::xtensor_fixed<int, xt::xshape<2, 3>, xt::layout_type::column_major>& m)

{

return xt::transpose(m) + 2;

}

// Readme Examples

double readme_example1(xt::pyarray<double>& m)

{

auto sines = xt::sin(m);

return std::accumulate(sines.cbegin(), sines.cend(), 0.0);

}

double readme_example2(double i, double j)

{

return std::sin(i) - std::cos(j);

}

auto complex_overload(const xt::pyarray<std::complex<double>>& a)

{

return a;

}

auto no_complex_overload(const xt::pyarray<double>& a)

{

return a;

}

auto complex_overload_reg(const std::complex<double>& a)

{

return a;

}

auto no_complex_overload_reg(const double& a)

{

return a;

}

//

// Operator examples

//

xt::pyarray<double> array_addition(const xt::pyarray<double>& m, const xt::pyarray<double>& n)

{

return m + n;

}

xt::pyarray<double> array_subtraction(xt::pyarray<double>& m, xt::pyarray<double>& n)

{

return m - n;

}

xt::pyarray<double> array_multiplication(xt::pyarray<double>& m, xt::pyarray<double>& n)

{

return m * n;

}

xt::pyarray<double> array_division(xt::pyarray<double>& m, xt::pyarray<double>& n)

{

return m / n;

}

// Vectorize Examples

int add(int i, int j)

{

return i + j;

}

template <class T> std::string typestring() { return "Unknown"; }

template <> std::string typestring<uint8_t>() { return "uint8"; }

template <> std::string typestring<int8_t>() { return "int8"; }

template <> std::string typestring<uint16_t>() { return "uint16"; }

template <> std::string typestring<int16_t>() { return "int16"; }

template <> std::string typestring<uint32_t>() { return "uint32"; }

template <> std::string typestring<int32_t>() { return "int32"; }

template <> std::string typestring<uint64_t>() { return "uint64"; }

template <> std::string typestring<int64_t>() { return "int64"; }

template <class T>

inline std::string int_overload(xt::pyarray<T>& m)

{

return typestring<T>();

}

void dump_numpy_constant()

{

std::cout << "NPY_BOOL = " << NPY_BOOL << std::endl;

std::cout << "NPY_BYTE = " << NPY_BYTE << std::endl;

std::cout << "NPY_UBYTE = " << NPY_UBYTE << std::endl;

std::cout << "NPY_INT8 = " << NPY_INT8 << std::endl;

std::cout << "NPY_UINT8 = " << NPY_UINT8 << std::endl;

std::cout << "NPY_SHORT = " << NPY_SHORT << std::endl;

std::cout << "NPY_USHORT = " << NPY_USHORT << std::endl;

std::cout << "NPY_INT16 = " << NPY_INT16 << std::endl;

std::cout << "NPY_UINT16 = " << NPY_UINT16 << std::endl;

std::cout << "NPY_INT = " << NPY_INT << std::endl;

std::cout << "NPY_UINT = " << NPY_UINT << std::endl;

std::cout << "NPY_INT32 = " << NPY_INT32 << std::endl;

std::cout << "NPY_UINT32 = " << NPY_UINT32 << std::endl;

std::cout << "NPY_LONG = " << NPY_LONG << std::endl;

std::cout << "NPY_ULONG = " << NPY_ULONG << std::endl;

std::cout << "NPY_LONGLONG = " << NPY_LONGLONG << std::endl;

std::cout << "NPY_ULONGLONG = " << NPY_ULONGLONG << std::endl;

std::cout << "NPY_INT64 = " << NPY_INT64 << std::endl;

std::cout << "NPY_UINT64 = " << NPY_UINT64 << std::endl;

}

struct A

{

double a;

int b;

char c;

std::array<double, 3> x;

};

struct B

{

double a;

int b;

};

class C

{

public:

using array_type = xt::xarray<double, xt::layout_type::row_major>;

C() : m_array{0, 0, 0, 0} {}

array_type & array() { return m_array; }

private:

array_type m_array;

};

struct test_native_casters

{

using array_type = xt::xarray<double>;

array_type a = xt::ones<double>({50, 50});

const auto & get_array()

{

return a;

}

auto get_strided_view()

{

return xt::strided_view(a, {xt::range(0, 1), xt::range(0, 3, 2)});

}

auto get_array_adapter()

{

using shape_type = std::vector<size_t>;

shape_type shape = {2, 2};

shape_type stride = {3, 2};

return xt::adapt(a.data(), 4, xt::no_ownership(), shape, stride);

}

auto get_tensor_adapter()

{

using shape_type = std::array<size_t, 2>;

shape_type shape = {2, 2};

shape_type stride = {3, 2};

return xt::adapt(a.data(), 4, xt::no_ownership(), shape, stride);

}

auto get_owning_array_adapter()

{

size_t size = 100;

int * data = new int[size];

std::fill(data, data + size, 1);

using shape_type = std::vector<size_t>;

shape_type shape = {size};

return xt::adapt(std::move(data), size, xt::acquire_ownership(), shape);

}

};

xt::pyarray<A> dtype_to_python()

{

A a1{123, 321, 'a', {1, 2, 3}};

A a2{111, 222, 'x', {5, 5, 5}};

return xt::pyarray<A>({a1, a2});

}

xt::pyarray<B> dtype_from_python(xt::pyarray<B>& b)

{

if (b(0).a != 1 || b(0).b != 'p' || b(1).a != 123 || b(1).b != 'c')

{

throw std::runtime_error("FAIL");

}

b(0).a = 123.;

b(0).b = 'w';

return b;

}

void char_array(xt::pyarray<char[20]>& carr)

{

if (strcmp(carr(2), "python"))

{

throw std::runtime_error("TEST FAILED!");

}

std::fill(&carr(2)[0], &carr(2)[0] + 20, 0);

carr(2)[0] = 'c';

carr(2)[1] = '+';

carr(2)[2] = '+';

carr(2)[3] = '\0';

}

void row_major_tensor(xt::pytensor<double, 3, xt::layout_type::row_major>& arg)

{

if (!std::is_same<decltype(arg.begin()), double*>::value)

{

throw std::runtime_error("TEST FAILED");

}

}

void col_major_array(xt::pyarray<double, xt::layout_type::column_major>& arg)

{

if (!std::is_same<decltype(arg.template begin<xt::layout_type::column_major>()), double*>::value)

{

throw std::runtime_error("TEST FAILED");

}

}

xt::pytensor<int, 0> xscalar(const xt::pytensor<int, 1>& arg)

{

return xt::sum(arg);

}

template <class T>

using ndarray = xt::pyarray<T, xt::layout_type::row_major>;

void test_rm(ndarray<int>const& x)

{

ndarray<int> y = x;

ndarray<int> z = xt::zeros<int>({10});

}

PYBIND11_MODULE(xtensor_python_test, m)

{

xt::import_numpy();

m.doc() = "Test module for xtensor python bindings";

m.def("example1", example1);

m.def("example2", example2);

m.def("example3_xarray", example3_xarray);

m.def("example3_xarray_colmajor", example3_xarray_colmajor);

m.def("example3_xtensor3", example3_xtensor3);

m.def("example3_xtensor2", example3_xtensor2);

m.def("example3_xtensor2_colmajor", example3_xtensor2_colmajor);

m.def("example3_xfixed3", example3_xfixed3);

m.def("example3_xfixed2", example3_xfixed2);

m.def("example3_xfixed2_colmajor", example3_xfixed2_colmajor);

m.def("complex_overload", no_complex_overload);

m.def("complex_overload", complex_overload);

m.def("complex_overload_reg", no_complex_overload_reg);

m.def("complex_overload_reg", complex_overload_reg);

m.def("readme_example1", readme_example1);

m.def("readme_example2", xt::pyvectorize(readme_example2));

m.def("array_addition", array_addition);

m.def("array_subtraction", array_subtraction);

m.def("array_multiplication", array_multiplication);

m.def("array_division", array_division);

m.def("vectorize_example1", xt::pyvectorize(add));

m.def("rect_to_polar", xt::pyvectorize([](complex_t x) { return std::abs(x); }));

m.def("compare_shapes", [](const xt::pyarray<double>& a, const xt::pyarray<double>& b) {

return a.shape() == b.shape();

});

m.def("test_rm", test_rm);

m.def("int_overload", int_overload<uint8_t>);

m.def("int_overload", int_overload<int8_t>);

m.def("int_overload", int_overload<uint16_t>);

m.def("int_overload", int_overload<int16_t>);

m.def("int_overload", int_overload<uint32_t>);

m.def("int_overload", int_overload<int32_t>);

m.def("int_overload", int_overload<uint64_t>);

m.def("int_overload", int_overload<int64_t>);

m.def("dump_numpy_constant", dump_numpy_constant);

// Register additional dtypes

PYBIND11_NUMPY_DTYPE(A, a, b, c, x);

PYBIND11_NUMPY_DTYPE(B, a, b);

m.def("dtype_to_python", dtype_to_python);

m.def("dtype_from_python", dtype_from_python);

m.def("char_array", char_array);

m.def("col_major_array", col_major_array);

m.def("row_major_tensor", row_major_tensor);

m.def("xscalar", xscalar);

py::class_<C>(m, "C")

.def(py::init<>())

.def_property_readonly(

"copy",

[](C & self) { return self.array(); }

)

.def_property_readonly(

"ref",

[](C & self) -> C::array_type & { return self.array(); }

)

;

m.def("simple_array", [](xt::pyarray<int>) { return 1; } );

m.def("simple_tensor", [](xt::pytensor<int, 1>) { return 2; } );

m.def("diff_shape_overload", [](xt::pytensor<int, 1> a) { return 1; });

m.def("diff_shape_overload", [](xt::pytensor<int, 2> a) { return 2; });

py::class_<test_native_casters>(m, "test_native_casters")

.def(py::init<>())

.def("get_array", &test_native_casters::get_array, py::return_value_policy::reference_internal) // memory managed by the class instance

.def("get_strided_view", &test_native_casters::get_strided_view, py::keep_alive<0, 1>()) // keep_alive<0, 1>() => do not free "self" before the returned view

.def("get_array_adapter", &test_native_casters::get_array_adapter, py::keep_alive<0, 1>()) // keep_alive<0, 1>() => do not free "self" before the returned adapter

.def("get_tensor_adapter", &test_native_casters::get_tensor_adapter, py::keep_alive<0, 1>()) // keep_alive<0, 1>() => do not free "self" before the returned adapter

.def("get_owning_array_adapter", &test_native_casters::get_owning_array_adapter) // auto memory management as the adapter owns its memory

.def("view_keep_alive_member_function", [](test_native_casters & self, xt::pyarray<double> & a) // keep_alive<0, 2>() => do not free second parameter before the returned view

{return xt::reshape_view(a, {a.size(), });},

py::keep_alive<0, 2>());

}