%PDF- %PDF-
Direktori : /proc/985914/root/lib/python3/dist-packages/pythran/pythonic/include/numpy/ |
Current File : //proc/985914/root/lib/python3/dist-packages/pythran/pythonic/include/numpy/arange.hpp |
#ifndef PYTHONIC_INCLUDE_NUMPY_ARANGE_HPP #define PYTHONIC_INCLUDE_NUMPY_ARANGE_HPP #include "pythonic/include/utils/functor.hpp" #include "pythonic/include/operator_/pos.hpp" #include "pythonic/include/types/ndarray.hpp" PYTHONIC_NS_BEGIN namespace numpy { namespace details { #ifdef USE_XSIMD template <class T> struct arange_simd_iterator { using vector_type = xsimd::simd_type<T>; vector_type curr_; vector_type step_; long index_; arange_simd_iterator(T start, T step, long n) : curr_{}, step_{static_cast<T>(vector_type::size * step)}, index_{static_cast<long>(n / vector_type::size)} { T from[vector_type::size]; for (size_t i = 0; i < vector_type::size; ++i) from[i] = start + i * step; curr_ = xsimd::load_unaligned(from); } vector_type operator*() const { return curr_; } arange_simd_iterator &operator++() { curr_ += step_; ++index_; return *this; } arange_simd_iterator &operator+=(long n) { curr_ += n * step_; index_ += n; return *this; } arange_simd_iterator operator+(long n) const { arange_simd_iterator other{*this}; return other += n; } arange_simd_iterator &operator--() { curr_ -= step_; --index_; return *this; } long operator-(arange_simd_iterator const &other) const { return index_ - other.index_; } bool operator!=(arange_simd_iterator const &other) const { return index_ != other.index_; } bool operator==(arange_simd_iterator const &other) const { return index_ == other.index_; } bool operator<(arange_simd_iterator const &other) const { return index_ < other.index_; } arange_simd_iterator & operator=(arange_simd_iterator const &other) = default; }; #endif template <class T> struct arange_index { T start, step; long size; using iterator = types::nditerator<arange_index>; using const_iterator = types::const_nditerator<arange_index>; using dtype = T; using value_type = dtype; using shape_t = types::pshape<long>; #ifdef USE_XSIMD using simd_iterator = arange_simd_iterator<T>; using simd_iterator_nobroadcast = simd_iterator; template <class vectorizer> simd_iterator vbegin(vectorizer) const { return {start, step, 0}; } template <class vectorizer> simd_iterator vend(vectorizer) const { return {static_cast<T>(start + size * step), step, size}; } #endif static constexpr size_t value = 1; static constexpr bool is_strided = false; static constexpr bool is_vectorizable = types::is_vectorizable<T>::value; T fast(long i) const { return start + i * step; } dtype load(long i) const { return fast(i); } template <size_t I> long shape() const { return size; } types::ndarray<dtype, shape_t> operator[](types::slice s) const { auto ns = s.normalize(size); arange_index r{start + s.lower * step, step * ns.step, ns.size()}; return { types::numpy_expr<pythonic::operator_::functor::pos, arange_index>{ r}}; } types::ndarray<dtype, shape_t> operator()(types::slice s) const { return operator[](s); } types::ndarray<dtype, shape_t> operator[](types::contiguous_slice s) const { auto ns = s.normalize(size); arange_index r{start + s.lower * step, step, ns.size()}; return { types::numpy_expr<pythonic::operator_::functor::pos, arange_index>{ r}}; } types::ndarray<dtype, shape_t> operator()(types::contiguous_slice s) const { return operator[](s); } template <class... S> auto operator()(S const &... s) const -> typename std::enable_if< (sizeof...(S) > 1), decltype(std::declval<types::ndarray<dtype, shape_t>>()(s...))>::type { return types::ndarray<dtype, shape_t>{ types::numpy_expr<pythonic::operator_::functor::pos, arange_index>{ *this}}(s...); } const_iterator begin() const { return {*this, 0}; } const_iterator end() const { return {*this, size}; } iterator begin() { return {*this, 0}; } iterator end() { return {*this, size}; } }; } template <class T, class U, class S = long, class dtype = types::dtype_t<typename __combined<T, U, S>::type>> types::numpy_expr<pythonic::operator_::functor::pos, details::arange_index<typename dtype::type>> arange(T begin, U end, S step = S(1), dtype d = dtype()); template <class T> types::numpy_expr<pythonic::operator_::functor::pos, details::arange_index<typename types::dtype_t<T>::type>> arange(T end); DEFINE_FUNCTOR(pythonic::numpy, arange); } PYTHONIC_NS_END #endif