Calculating the gradient over a thrust::device_vector

I am currently transferring code from local C++ to CUDA, while working with thrust::device_vectors. Now there is a function to calculate the gradient where I also need to have access to non only the current element, but also the surrounding elements. In the original code I wrote the following:

void calc_grad(Complex *grad, const Complex *data, const int size)
{
    for (int i = 1; i < size - 1; i++) {
        grad[i] = 0.5 * (data[i + 1] - data[i - 1]);
    }
    grad[0] = data[1] - data[0];
    grad[size - 1] = data[size - 1] - data[size - 2];
}

Is it possible to create a functor for thrust out of that, such that I can call it in thrust::transform()? Up til now I only can access one element at the time, without getting the surrounding elements. Or is it not possible anyway, due to depending on elements before and afterwards, which are changed?

The formula for the code came from the matlab function gradient, shown here: http://se.mathworks.com/help/matlab/ref/gradient.html?requestedDomain=www.mathworks.com

Solution

A single Thrust::transform() can do most part of your work. All you need to do is to shift the data a little bit so that grad[i], data[i-1] and data[i+1] are aligned.

thrust::transform(data.begin()+2,
                  data.end(),
                  data.begin(),
                  grad.begin()+1,
                  (_1 - _2) * 0.5);

Then you can deal with the boundary cases.

Edit

And you can also include the boundary cases in one transform. With the following form of transform, your functor Grad should be able to know the index of the data he is dealing with, by the first functor parameter. Based on the index, he can then choose 2 out of the 3 elements from the second functor parameter, which is a tuple, to do the correct calculation.

Everything here is not tested. I'm not sure if data.begin()-1 works. You may also be careful on the Complex type.

thrust::transform(
  thrust::make_counting_iterator(int(0)),
  thrust::make_counting_iterator(int(0)) + size,
  thrust::make_zip_iterator(
      thrust::make_tuple(
          data.begin()-1,
          data.begin(),
          data.begin()+1)),
  grad.begin(),
  Grad(size)
);

The functor is something like this.

struct Grad {
  int size_;
  Grad(int s) :
      size_(s) {
  }
  template<typename T, typename Tuple>
  __host__ __device__
  inline T operator()(const int& idx, const Tuple& d) {
    if (idx == 0) {
      return d.get<2>() - d.get<1>();
    } else if (idx == size_ - 1) {
      return d.get<1>() - d.get<0>();
    } else {
      return 0.5 * (d.get<2>() - d.get<0>());
    }
  }
};