%PDF- %PDF-
| Direktori : /usr/lib/python3/dist-packages/pythran/pythonic/numpy/ |
| Current File : //usr/lib/python3/dist-packages/pythran/pythonic/numpy/trace.hpp |
#ifndef PYTHONIC_NUMPY_TRACE_HPP
#define PYTHONIC_NUMPY_TRACE_HPP
#include "pythonic/include/numpy/trace.hpp"
#include "pythonic/utils/functor.hpp"
PYTHONIC_NS_BEGIN
namespace numpy
{
template <class T>
typename T::dtype trace(T const &expr, int offset)
{
static_assert(T::value == 2, "Not Implemented : Trace for dimension != 2");
typename T::dtype res = 0;
long y_offset = std::max(-offset, 0);
long x_offset = std::max(0, offset);
long size = std::min(expr.flat_size() - y_offset,
expr.fast(0).flat_size() - x_offset);
if (offset < 0)
for (long i = 0; i < size; i++)
res += expr.fast(i + offset).fast(i);
else if (offset > 0)
for (long i = 0; i < size; i++)
res += expr.fast(i).fast(i + offset);
else
for (long i = 0; i < size; i++)
res += expr.fast(i).fast(i);
return res;
}
}
PYTHONIC_NS_END
#endif