1 #ifndef PYTHONIC_NUMPY_PARTIAL_SUM_HPP 2 #define PYTHONIC_NUMPY_PARTIAL_SUM_HPP 3 4 #include "pythonic/include/numpy/partial_sum.hpp" 5 6 #include "pythonic/types/ndarray.hpp" 7 #include "pythonic/builtins/ValueError.hpp" 8 9 PYTHONIC_NS_BEGIN 10 11 namespace numpy 12 { 13 14 /** 15 * The cast is perform to be numpy compliant 16 * 17 * a = numpy.array([1, 256]) 18 * In [10]: numpy.mod.accumulate(a, dtype=numpy.uint32) 19 * Out[10]: array([1, 1], dtype=uint32) 20 * In [11]: numpy.mod.accumulate(a, dtype=numpy.uint8) 21 * Out[11]: array([1, 0], dtype=uint8) 22 */ 23 namespace 24 { 25 26 template <class Op, size_t N, class A> 27 struct _partial_sum { 28 template <class E, class F> operator ()numpy::__anone97633850111::_partial_sum29 A operator()(E const &e, F &o) 30 { 31 auto it_begin = e.begin(); 32 A acc = _partial_sum<Op, N - 1, A>{}((*it_begin), o); 33 ++it_begin; 34 for (; it_begin < e.end(); ++it_begin) 35 acc = _partial_sum<Op, N - 1, A>{}(*it_begin, o, acc); 36 return acc; 37 } 38 template <class E, class F> operator ()numpy::__anone97633850111::_partial_sum39 A operator()(E const &e, F &o, A acc) 40 { 41 for (auto const &value : e) 42 acc = _partial_sum<Op, N - 1, A>{}(value, o, acc); 43 return acc; 44 } 45 }; 46 47 template <class Op, class A> 48 struct _partial_sum<Op, 1, A> { 49 template <class E, class F> operator ()numpy::__anone97633850111::_partial_sum50 A operator()(E const &e, F &o) 51 { 52 auto it_begin = e.begin(); 53 A acc = *it_begin; 54 *o = acc; 55 ++it_begin, ++o; 56 for (; it_begin < e.end(); ++it_begin, ++o) { 57 acc = Op{}(acc, (A)*it_begin); 58 *o = acc; 59 } 60 return acc; 61 } 62 template <class E, class F> operator ()numpy::__anone97633850111::_partial_sum63 A operator()(E e, F &o, A acc) 64 { 65 for (auto const &value : e) { 66 acc = Op{}(acc, (A)value); 67 *o = acc; 68 ++o; 69 } 70 return acc; 71 } 72 }; 73 } 74 75 template <class Op, class E, class dtype> 76 types::ndarray<typename dtype::type, types::pshape<long>> partial_sum(E const & expr,dtype d)77 partial_sum(E const &expr, dtype d) 78 { 79 const long count = expr.flat_size(); 80 types::ndarray<typename dtype::type, types::pshape<long>> the_partial_sum{ 81 types::make_tuple(count), builtins::None}; 82 auto begin_it = the_partial_sum.begin(); 83 _partial_sum<Op, E::value, typename dtype::type>{}(expr, begin_it); 84 return the_partial_sum; 85 } 86 87 template <class Op, class E, class dtype> partial_sum(E const & expr,long axis,dtype d)88 auto partial_sum(E const &expr, long axis, dtype d) -> 89 typename std::enable_if<E::value == 1, 90 decltype(partial_sum<Op, E, dtype>(expr))>::type 91 { 92 if (axis != 0) 93 throw types::ValueError("axis out of bounds"); 94 return partial_sum<Op, E, dtype>(expr); 95 } 96 97 template <class Op, class E, class dtype> 98 typename std::enable_if<E::value != 1, partial_sum_type<Op, E, dtype>>::type partial_sum(E const & expr,long axis,dtype d)99 partial_sum(E const &expr, long axis, dtype d) 100 { 101 if (axis < 0 || size_t(axis) >= E::value) 102 throw types::ValueError("axis out of bounds"); 103 104 auto shape = sutils::getshape(expr); 105 partial_sum_type<Op, E, dtype> the_partial_sum{shape, builtins::None}; 106 if (axis == 0) { 107 auto it_begin = the_partial_sum.begin(); 108 _partial_sum<Op, 1, partial_sum_type2<Op, E, dtype>>{}(expr, it_begin); 109 } else { 110 std::transform( 111 expr.begin(), expr.end(), the_partial_sum.begin(), 112 [axis, d]( 113 typename std::iterator_traits<typename E::iterator>::value_type 114 other) { return partial_sum<Op>(other, axis - 1, d); }); 115 } 116 return the_partial_sum; 117 } 118 } 119 PYTHONIC_NS_END 120 121 #endif 122