[Blitz-devel] n-dimensional convolution

[Daniel O. <da...@da...>]

Hello,

I just implemented n-dimensional convolution on blitz-arrays and  
thought it might be of use to someone or integrated into the library.  
For convenience it has two parameters specifying the range (full/ 
valid/same as in Matlab) and the dimensions over which to convolve.  
It also has a convenience cross-corellation function which performs  
the appropriate reverses before convolution. I put this code under GPL.

this is the convolution code:

/*
*  convn.h
*
*  Created by Daniel Oberhoff on 8/27/06.
*
*/

#include <blitz/array.h>
#include <blitz/bzdebug.h>

#include <iostream>
using namespace std;

#if defined(BZ_DEBUG)
     #include <blitz/vecexpr.h>
#endif

//this function shifts d2 by offset and then constrains the domains
//to where they overlap
template<int N> inline void mask_constrain_nd(
           blitz::RectDomain<N>& d1,
           blitz::RectDomain<N>& d2,
           blitz::TinyVector<int,N> offset
                                               )
{

#if defined(BZ_DEBUG)
//    blitz::TinyVector<int,N> max_offset = d1.ubound();
//    max_offset -= d1.lbound();
//    blitz::TinyVector<bool,N> condition = offset <= max_offset;
//    BZPRECONDITION(all(condition));
#endif

     //this could be done using tinyvec-et, but they
     //slow down the compiler too much...
     for (int i = 0; i < N; ++i)
     {
         int offset_i = offset[i];
         int msize_i = d2.ubound(i) - d2.lbound(i);
         if (offset_i > 0)
         {
             d1.lbound(i) += offset_i;
             int ubound1_i = d1.ubound(i);
             int lbound1_i = d1.lbound(i);
             int shrink_i  = ubound1_i - (lbound1_i + msize_i);
             if (shrink_i > 0) d1.ubound(i) -= shrink_i;
             d2.ubound(i) = d2.lbound(i) + (d1.ubound(i) - d1.lbound 
(i));
         }
         else if (offset_i < 0)
         {
             int ubound1_i = d1.ubound(i);
             int lbound1_i = d1.lbound(i);
             int shrink_i  = ubound1_i - (lbound1_i + msize_i +  
offset_i);
             if (shrink_i > 0) d1.ubound(i) -= shrink_i;
             d2.lbound(i) = d2.ubound(i) - (d1.ubound(i) - d1.lbound 
(i));
         }
         else //offset_i == 0
         {
             int ssize_i = d1.ubound(i) - d1.lbound(i);
             if (ssize_i > msize_i)
                 d1.ubound(i) -= ssize_i - msize_i;
             else if (ssize_i < msize_i)
                 d2.ubound(i) -= msize_i - ssize_i;
         }
     }
}

//3 possible ways to do convolution:
//Full: the full convolution, implicitly zero-padding m1 with m2.shape 
() on each side,
//        result is bigger than m1 by m2.shape()-1 where it is odd  
and m2.shape() where it is even
//Valid: the 'inner' convolution, no padding,
//          result is smaller than m1 by m2.shape()-1 where it is odd  
and m2.shape() where it is even
//Same: retains the shape, implicitly zero-pads m1 by m2.shape()/2
enum conv_shape {Full, Same, Valid};

//the n-dimensional convolition on 2 arrays
//flags gives the option not to convolve over all dimensions
//where no convolition should be done (flags[i]==false) m2 must be  
the same shape as m1!
template<typename T, int N> inline blitz::Array<T,N> conv_n(
     const blitz::Array<T, N>& m1,
     const blitz::Array<T, N>& m2,
     conv_shape shape = Full,
     blitz::TinyVector<bool,N> flags = true
                                                             )
{
     typedef blitz::Array<T, N> T_matrix;
     typedef blitz::RectDomain<N> Domain;
     typedef typename T_matrix::T_index T_index;
     typedef typename T_matrix::iterator T_iter;

     T_index    vi;
     T_index    lbound = m1.lbound();
     T_index ubound = m1.ubound();
     T_index mshift = 0;

     T_index    rshape = m1.shape();

     //determine the domain of the convolution
     for (int i = 0; i < N; i++)
     {
         if (flags[i])
         {
             mshift[i] = m2.extent(i) / 2;
             switch(shape)
             {
                 case Same:
                     ubound[i] = m1.ubound(i);
                     break;
                 case Full:
                     lbound[i] = m1.lbound(i) - mshift[i];
                     ubound[i] = m1.ubound(i) + mshift[i];
                     break;
                 case Valid:
                     lbound[i] = m1.lbound(i) + mshift[i];
                     ubound[i] = m1.ubound(i) - mshift[i];
                     break;
             }
             rshape[i] = ubound[i] - lbound[i] + 1;
         }
     }

     T_matrix result(rshape);
     T_index shift0 = -mshift;
     shift0 += lbound;
     T_index baseshift = m2.lbound();
     baseshift -= m1.lbound();
     for (T_iter iout = result.begin(); iout != result.end(); ++iout)
     {
         Domain d1 = m1.domain();
         Domain d2 = m2.domain();
         T_index shift = shift0;
         T_index pos = iout.position();

         for (int i = 0; i < N; ++i)
         {
             if (!flags[i])
             {
                 d1.lbound(i) = d1.ubound(i) = pos[i];
                 d2.lbound(i) = d2.ubound(i) = pos[i] + baseshift[i];
             } else shift[i] += pos[i];
         }

         mask_constrain_nd(d1,d2,shift);

         cout << "shift: " << shift << endl;
         cout << "1: [ " << d1.lbound() << ", " << d1.ubound() <<  
" ]" << endl;
         cout << "2: [ " << d2.lbound() << ", " << d2.ubound() <<  
" ]" << endl;

         *iout = blitz::sum(m1(d1)*m2(d2));
     }

     return result;
}

//convenience function for cross-corellation, takes the same  
arguments as conv_n
//but reverses m2 where on the dimensions where flags[i] is set  
before invoking conv_n
template<typename T, int N> inline blitz::Array<T,N>  
cross_correllation_n(
     const blitz::Array<T, N>& m1,
     const blitz::Array<T, N>& m2,
     conv_shape shape = Full,
     blitz::TinyVector<bool,N> flags = true
                                                             )
{
     const blitz::Array<T, N> mm2 = m2;

     for (int i = 0; i < N; ++i)
     {
         if (flags[i])
         {
             mm2.reference(mm2.reverse(i));
         }
     }
     return conv_n(m1, mm2, shape, flags);
}



here is a little test app:



#define BZ_DEBUG

#include "../convn.h"
#include <iostream>

using namespace blitz;
using namespace std;

int main(void)
{
     typedef Array<int,2> Matrix;

     Matrix A(5,5), B(5,5);

     A =    0, 0, 1, 0, 0,
         0, 0, 1, 0, 0,
         1, 1, 1, 1, 1,
         0, 0, 1, 0, 0,
         0, 0, 1, 0, 0;

     B =    0, 0, 0, 0, 0,
         0, 1, 1, 1, 0,
         0, 1, 1, 1, 0,
         0, 1, 1, 1, 0,
         0, 0, 0, 0, 0;

//    B =    1, 1, 1,
//         1, 1, 1,
//         1, 1, 1;
// breaks the 1d conv as the non-convolved shape has to be the same

     Matrix CS = conv_n(A,B,Same);
     Matrix CF = conv_n(A,B,Full);
     Matrix CV = conv_n(A,B,Valid);

     TinyVector<bool,2> flags1 = true;
     flags1[1] = false;

     Matrix C1S = conv_n(A,B,Same,flags1);
     Matrix C1F = conv_n(A,B,Full,flags1);
     Matrix C1V = conv_n(A,B,Valid,flags1);

     cout << "A = " << A << endl;
     cout << "B = " << B << endl;
     cout << "CS = " << CS << endl;
     cout << "CF = " << CF << endl;
     cout << "CV = " << CV << endl;
     cout << "C1S = " << C1S << endl;
     cout << "C1F = " << C1F << endl;
     cout << "C1V = " << C1V << endl;
     return 0;
}

Cheers,

Daniel Oberhoff

p.s: I have some suggestions for blitz: iterations over domains. A  
method RectDomain::intersect(other) from which this code would benefit.