Actual source code: matusfft.c

  1: #define PETSCMAT_DLL

  3: /*
  4:     Provides an implementation of the Unevenly Sampled FFT algorithm as a Mat.
  5:     Testing examples can be found in ~/src/mat/examples/tests
  6: */

  8: #include "private/matimpl.h"          /*I "petscmat.h" I*/
  9: #include "petscda.h"                  /*I "petscda.h"  I*/ /* Unlike equispaced FFT, USFFT requires geometric information encoded by a DA */
 10: #include "fftw3.h"

 12: typedef struct {
 13:   DA             inda;
 14:   DA             outda;
 15:   PetscInt       ndim;
 16:   PetscInt       *indim;
 17:   PetscInt       *outdim;
 18:   PetscInt       m,n,M,N;
 19:   PetscInt       dof;
 20:   fftw_plan      p_forward,p_backward;
 21:   unsigned       p_flag; /* planner flags, FFTW_ESTIMATE,FFTW_MEASURE, FFTW_PATIENT, FFTW_EXHAUSTIVE */
 22: } Mat_USFFT;


 27: PetscErrorCode MatApply_USFFT_Private(Mat_USFFT *usfft, fftw_plan *plan, int direction, Vec x,Vec y)
 28: {
 30:   PetscScalar    *x_array, *y_array;

 33:   /* NB: for now we use outdim for both x and y; this will change once a full USFFT is implemented */
 34:   VecGetArray(x,&x_array);
 35:   VecGetArray(y,&y_array);
 36:   if (!*plan){ /* create a plan then execute it*/
 37:     if(usfft->dof == 1) {
 38: #ifdef PETSC_DEBUG_USFFT
 39:       PetscPrintf(PETSC_COMM_WORLD, "direction = %d, usfft->ndim = %d\n", direction, usfft->ndim);
 40:       for(int ii = 0; ii < usfft->ndim; ++ii) {
 41:         PetscPrintf(PETSC_COMM_WORLD, "usfft->outdim[%d] = %d\n", ii, usfft->outdim[ii]);
 42:       }
 43: #endif 
 44: #if 0
 45:       *plan = fftw_plan_dft(usfft->ndim,usfft->outdim,(fftw_complex*)x_array,(fftw_complex*)y_array,direction,usfft->p_flag);
 46: #endif
 47:       switch (usfft->ndim){
 48:       case 1:
 49:         *plan = fftw_plan_dft_1d(usfft->outdim[0],(fftw_complex*)x_array,(fftw_complex*)y_array,direction,usfft->p_flag);
 50:         break;
 51:       case 2:
 52:         *plan = fftw_plan_dft_2d(usfft->outdim[0],usfft->outdim[1],(fftw_complex*)x_array,(fftw_complex*)y_array,direction,usfft->p_flag);
 53:         break;
 54:       case 3:
 55:         *plan = fftw_plan_dft_3d(usfft->outdim[0],usfft->outdim[1],usfft->outdim[2],(fftw_complex*)x_array,(fftw_complex*)y_array,direction,usfft->p_flag);
 56:         break;
 57:       default:
 58:         *plan = fftw_plan_dft(usfft->ndim,usfft->outdim,(fftw_complex*)x_array,(fftw_complex*)y_array,direction,usfft->p_flag);
 59:         break;
 60:       }
 61:       fftw_execute(*plan);
 62:     }/* if(dof == 1) */
 63:     else { /* if(dof > 1) */
 64:       *plan = fftw_plan_many_dft(/*rank*/usfft->ndim, /*n*/usfft->outdim, /*howmany*/usfft->dof,
 65:                                  (fftw_complex*)x_array, /*nembed*/usfft->outdim, /*stride*/usfft->dof, /*dist*/1,
 66:                                  (fftw_complex*)y_array, /*nembed*/usfft->outdim, /*stride*/usfft->dof, /*dist*/1,
 67:                                  /*sign*/direction, /*flags*/usfft->p_flag);
 68:       fftw_execute(*plan);
 69:     }/* if(dof > 1) */
 70:   }/* if(!*plan) */
 71:   else { /* if(*plan) */
 72:     /* use existing plan */
 73:     fftw_execute_dft(*plan,(fftw_complex*)x_array,(fftw_complex*)y_array);
 74:   }
 75:   VecRestoreArray(y,&y_array);
 76:   VecRestoreArray(x,&x_array);
 77:   return(0);
 78: }/* MatApply_FFTW_Private() */


 83: PetscErrorCode MatMult_SeqUSFFT(Mat A,Vec x,Vec y)
 84: {
 86:   Mat_USFFT       *usfft = (Mat_USFFT*)A->data;

 89:   /* NB: for now we use outdim for both x and y; this will change once a full USFFT is implemented */
 90:   MatApply_USFFT_Private(usfft, &usfft->p_forward, FFTW_FORWARD, x,y);
 91:   return(0);
 92: }

 96: PetscErrorCode MatMultTranspose_SeqUSFFT(Mat A,Vec x,Vec y)
 97: {
 99:   Mat_USFFT       *usfft = (Mat_USFFT*)A->data;
101:   /* NB: for now we use outdim for both x and y; this will change once a full USFFT is implemented */
102:   MatApply_USFFT_Private(usfft, &usfft->p_backward, FFTW_BACKWARD, x,y);
103:   return(0);
104: }

108: PetscErrorCode MatDestroy_SeqUSFFT(Mat A)
109: {
110:   Mat_USFFT       *usfft = (Mat_USFFT*)A->data;

114:   fftw_destroy_plan(usfft->p_forward);
115:   fftw_destroy_plan(usfft->p_backward);
116:   PetscFree(usfft->indim);
117:   PetscFree(usfft->outdim);
118:   PetscFree(usfft);
119:   PetscObjectChangeTypeName((PetscObject)A,0);
120:   return(0);
121: }/* MatDestroy_SeqUSFFT() */


126: /*@
127:       MatCreateSeqUSFFT - Creates a matrix object that provides sequential USFFT
128:   via the external package FFTW

130:    Collective on MPI_Comm

132:    Input Parameter:
133: +   da - geometry of the domain encoded by a DA

135:    Output Parameter:
136: .   A  - the matrix

138:   Options Database Keys:
139: + -mat_usfft_plannerflags - set the FFTW planner flags

141:    Level: intermediate
142:    
143: @*/
144: PetscErrorCode  MatCreateSeqUSFFT(DA inda, DA outda, Mat* A)
145: {
147:   Mat_USFFT      *usfft;
148:   PetscInt       m,n,M,N,i;
149:   const char     *p_flags[]={"FFTW_ESTIMATE","FFTW_MEASURE","FFTW_PATIENT","FFTW_EXHAUSTIVE"};
150:   PetscTruth     flg;
151:   PetscInt       p_flag;
152:   PetscInt       dof, ndim;
153:   PetscInt       dim[3];
154:   MPI_Comm       comm;
155:   PetscInt       size;

158:   PetscObjectGetComm((PetscObject)inda, &comm);
159:   MPI_Comm_size(comm, &size);
160:   if (size > 1) SETERRQ(PETSC_ERR_USER, "Parallel DA (in) not yet supported by USFFT");
161:   PetscObjectGetComm((PetscObject)outda, &comm);
162:   MPI_Comm_size(comm, &size);
163:   if (size > 1) SETERRQ(PETSC_ERR_USER, "Parallel DA (out) not yet supported by USFFT");
164:   MatCreate(comm,A);
165:   PetscNewLog(*A,Mat_USFFT,&usfft);
166:   (*A)->data = (void*)usfft;
167:   usfft->inda = inda;
168:   usfft->outda = outda;
169:   /* inda */
170:   DAGetInfo(usfft->inda, &ndim, dim+0, dim+1, dim+2, PETSC_NULL, PETSC_NULL, PETSC_NULL, &dof, PETSC_NULL, PETSC_NULL, PETSC_NULL);
171:   if (ndim <= 0) SETERRQ1(PETSC_ERR_USER,"ndim %d must be > 0",ndim);
172:   if (dof <= 0) SETERRQ1(PETSC_ERR_USER,"dof %d must be > 0",dof);
173:   usfft->ndim = ndim;
174:   usfft->dof = dof;
175:   /* Store input dimensions */
176:   /* NB: we reverse the DA dimensions, since the DA ordering (natural on x-y-z, with x varying the fastest) 
177:      is the order opposite of that assumed by FFTW: z varying the fastest */
178:   PetscMalloc((usfft->ndim+1)*sizeof(PetscInt),&usfft->indim);
179:   for(i = usfft->ndim; i > 0; --i) {
180:     usfft->indim[usfft->ndim-i] = dim[i-1];
181:   }
182:   /* outda */
183:   DAGetInfo(usfft->outda, &ndim, dim+0, dim+1, dim+2, PETSC_NULL, PETSC_NULL, PETSC_NULL, &dof, PETSC_NULL, PETSC_NULL, PETSC_NULL);
184:   if (ndim != usfft->ndim) SETERRQ2(PETSC_ERR_USER,"in and out DA dimensions must match: %d != %d",usfft->ndim, ndim);
185:   if (dof != usfft->dof) SETERRQ2(PETSC_ERR_USER,"in and out DA dof must match: %d != %d",usfft->dof, dof);
186:   /* Store output dimensions */
187:   /* NB: we reverse the DA dimensions, since the DA ordering (natural on x-y-z, with x varying the fastest) 
188:      is the order opposite of that assumed by FFTW: z varying the fastest */
189:   PetscMalloc((usfft->ndim+1)*sizeof(PetscInt),&usfft->outdim);
190:   for(i = usfft->ndim; i > 0; --i) {
191:     usfft->outdim[usfft->ndim-i] = dim[i-1];
192:   }
193:   /* mat sizes */
194:   m = 1; n = 1;
195:   for (i=0; i<usfft->ndim; i++){
196:     if (usfft->indim[i] <= 0) SETERRQ2(PETSC_ERR_USER,"indim[%d]=%d must be > 0",i,usfft->indim[i]);
197:     if (usfft->outdim[i] <= 0) SETERRQ2(PETSC_ERR_USER,"outdim[%d]=%d must be > 0",i,usfft->outdim[i]);
198:     n *= usfft->indim[i];
199:     m *= usfft->outdim[i];
200:   }
201:   N = n*usfft->dof;
202:   M = m*usfft->dof;
203:   MatSetSizes(*A,M,N,M,N);  /* "in size" is the number of columns, "out size" is the number of rows" */
204:   PetscObjectChangeTypeName((PetscObject)*A,MATSEQUSFFT);
205:   usfft->m = m; usfft->n = n; usfft->M = M; usfft->N = N;
206:   /* FFTW */
207:   usfft->p_forward  = 0;
208:   usfft->p_backward = 0;
209:   usfft->p_flag     = FFTW_ESTIMATE;
210:   /* set Mat ops */
211:   (*A)->ops->mult          = MatMult_SeqUSFFT;
212:   (*A)->ops->multtranspose = MatMultTranspose_SeqUSFFT;
213:   (*A)->assembled          = PETSC_TRUE;
214:   (*A)->ops->destroy       = MatDestroy_SeqUSFFT;
215:   /* get runtime options */
216:   PetscOptionsBegin(((PetscObject)(*A))->comm,((PetscObject)(*A))->prefix,"USFFT Options","Mat");
217:   PetscOptionsEList("-mat_usfft_fftw_plannerflags","Planner Flags","None",p_flags,4,p_flags[0],&p_flag,&flg);
218:   if (flg) {usfft->p_flag = (unsigned)p_flag;}
219:   PetscOptionsEnd();
220:   return(0);
221: }/* MatCreateSeqUSFFT() */