Actual source code: matusfft.c

  1: /*
  2:     Provides an implementation of the Unevenly Sampled FFT algorithm as a Mat.
  3:     Testing examples can be found in ~/src/mat/tests FIX: should these be moved to dm/da/tests?
  4: */

  6: #include <petsc/private/matimpl.h>
  7: #include <petscdmda.h>
  8: #include <fftw3.h>

 10: typedef struct {
 11:   PetscInt  dim;
 12:   Vec       sampleCoords;
 13:   PetscInt  dof;
 14:   DM        freqDA;     /* frequency DMDA */
 15:   PetscInt *freqSizes;  /* sizes of the frequency DMDA, one per each dim */
 16:   DM        resampleDa; /* the Battle-Lemarie interpolant DMDA */
 17:   Vec       resample;   /* Vec of samples, one per dof per sample point */
 18:   fftw_plan p_forward, p_backward;
 19:   unsigned  p_flag; /* planner flags, FFTW_ESTIMATE, FFTW_MEASURE, FFTW_PATIENT, FFTW_EXHAUSTIVE */
 20: } Mat_USFFT;

 22: #if 0
 23: static PetscErrorCode MatApply_USFFT_Private(Mat A, fftw_plan *plan, int direction, Vec x, Vec y)
 24: {
 25:   #if 0
 26:   PetscScalar    *r_array, *y_array;
 27:   Mat_USFFT* = (Mat_USFFT*)A->data;
 28:   #endif

 30:   PetscFunctionBegin;
 31:   #if 0
 32:   /* resample x to usfft->resample */
 33:   PetscCall(MatResample_USFFT_Private(A, x));

 35:   /* NB: for now we use outdim for both x and y; this will change once a full USFFT is implemented */
 36:   PetscCall(VecGetArray(usfft->resample,&r_array));
 37:   PetscCall(VecGetArray(y,&y_array));
 38:   if (!*plan) { /* create a plan then execute it*/
 39:     if (usfft->dof == 1) {
 40:     #if defined(PETSC_DEBUG_USFFT)
 41:       PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "direction = %d, usfft->ndim = %d\n", direction, usfft->ndim));
 42:       for (int ii = 0; ii < usfft->ndim; ++ii) {
 43:         PetscCall(PetscPrintf(PetscObjectComm((PetscObject)A), "usfft->outdim[%d] = %d\n", ii, usfft->outdim[ii]));
 44:       }
 45:     #endif

 47:       switch (usfft->dim) {
 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:   PetscCall(VecRestoreArray(y,&y_array));
 76:   PetscCall(VecRestoreArray(x,&x_array));
 77:   #endif
 78:   PetscFunctionReturn(PETSC_SUCCESS);
 79: } /* MatApply_USFFT_Private() */

 81: PetscErrorCode MatUSFFT_ProjectOnBattleLemarie_Private(Vec x,double *r)
 82: /* Project onto the Battle-Lemarie function centered around r */
 83: {
 84:   PetscScalar    *x_array, *y_array;

 86:   PetscFunctionBegin;
 87:   PetscFunctionReturn(PETSC_SUCCESS);
 88: } /* MatUSFFT_ProjectOnBattleLemarie_Private() */

 90: PetscErrorCode MatInterpolate_USFFT_Private(Vec x,Vec y)
 91: {
 92:   PetscScalar    *x_array, *y_array;

 94:   PetscFunctionBegin;
 95:   PetscFunctionReturn(PETSC_SUCCESS);
 96: } /* MatInterpolate_USFFT_Private() */

 98: PetscErrorCode MatMult_SeqUSFFT(Mat A,Vec x,Vec y)
 99: {
100:   Mat_USFFT      *usfft = (Mat_USFFT*)A->data;

102:   PetscFunctionBegin;
103:   /* NB: for now we use outdim for both x and y; this will change once a full USFFT is implemented */
104:   PetscCall(MatApply_USFFT_Private(A, &usfft->p_forward, FFTW_FORWARD, x,y));
105:   PetscFunctionReturn(PETSC_SUCCESS);
106: }

108: PetscErrorCode MatMultTranspose_SeqUSFFT(Mat A,Vec x,Vec y)
109: {
110:   Mat_USFFT      *usfft = (Mat_USFFT*)A->data;

112:   PetscFunctionBegin;
113:   /* NB: for now we use outdim for both x and y; this will change once a full USFFT is implemented */
114:   PetscCall(MatApply_USFFT_Private(usfft, &usfft->p_backward, FFTW_BACKWARD, x,y));
115:   PetscFunctionReturn(PETSC_SUCCESS);
116: }

118: PetscErrorCode MatDestroy_SeqUSFFT(Mat A)
119: {
120:   Mat_USFFT      *usfft = (Mat_USFFT*)A->data;

122:   PetscFunctionBegin;
123:   fftw_destroy_plan(usfft->p_forward);
124:   fftw_destroy_plan(usfft->p_backward);
125:   PetscCall(PetscFree(usfft->indim));
126:   PetscCall(PetscFree(usfft->outdim));
127:   PetscCall(PetscFree(usfft));
128:   PetscCall(PetscObjectChangeTypeName((PetscObject)A,0));
129:   PetscFunctionReturn(PETSC_SUCCESS);
130: }

132: /*@C
133:       MatCreateSeqUSFFT - Creates a matrix object that provides sequential USFFT
134:   via the external package FFTW

136:    Collective

138:    Input Parameter:
139: .  da - geometry of the domain encoded by a `DMDA`

141:    Output Parameter:
142: .  A  - the matrix

144:   Options Database Key:
145: . -mat_usfft_plannerflags - set the FFTW planner flags

147:    Level: intermediate

149: .seealso: `Mat`, `Vec`, `DMDA`, `DM`
150: @*/
151: PetscErrorCode  MatCreateSeqUSFFT(Vec sampleCoords, DMDA freqDA, Mat *A)
152: {
153:   Mat_USFFT      *usfft;
154:   PetscInt       m,n,M,N,i;
155:   const char     *p_flags[]={"FFTW_ESTIMATE","FFTW_MEASURE","FFTW_PATIENT","FFTW_EXHAUSTIVE"};
156:   PetscBool      flg;
157:   PetscInt       p_flag;
158:   PetscInt       dof, dim, freqSizes[3];
159:   MPI_Comm       comm;
160:   PetscInt       size;

162:   PetscFunctionBegin;
163:   PetscCall(PetscObjectGetComm((PetscObject)inda, &comm));
164:   PetscCallMPI(MPI_Comm_size(comm, &size));
165:   PetscCheck(size <= 1,comm,PETSC_ERR_USER, "Parallel DMDA (in) not yet supported by USFFT");
166:   PetscCall(PetscObjectGetComm((PetscObject)outda, &comm));
167:   PetscCallMPI(MPI_Comm_size(comm, &size));
168:   PetscCheck(size <= 1,comm,PETSC_ERR_USER, "Parallel DMDA (out) not yet supported by USFFT");
169:   PetscCall(MatCreate(comm,A));
170:   PetscCall(PetscNew(&usfft));
171:   (*A)->data   = (void*)usfft;
172:   usfft->inda  = inda;
173:   usfft->outda = outda;
174:   /* inda */
175:   PetscCall(DMDAGetInfo(usfft->inda, &ndim, dim+0, dim+1, dim+2, NULL, NULL, NULL, &dof, NULL, NULL, NULL));
176:   PetscCheck(ndim > 0,PETSC_COMM_SELF,PETSC_ERR_USER,"ndim %d must be > 0",ndim);
177:   PetscCheck(dof > 0,PETSC_COMM_SELF,PETSC_ERR_USER,"dof %d must be > 0",dof);
178:   usfft->ndim   = ndim;
179:   usfft->dof    = dof;
180:   usfft->freqDA = freqDA;
181:   /* NB: we reverse the freq and resample DMDA sizes, since the DMDA ordering (natural on x-y-z, with x varying the fastest)
182:      is the order opposite of that assumed by FFTW: z varying the fastest */
183:   PetscCall(PetscMalloc1(usfft->ndim+1,&usfft->indim));
184:   for (i = usfft->ndim; i > 0; --i) usfft->indim[usfft->ndim-i] = dim[i-1];

186:   /* outda */
187:   PetscCall(DMDAGetInfo(usfft->outda, &ndim, dim+0, dim+1, dim+2, NULL, NULL, NULL, &dof, NULL, NULL, NULL));
188:   PetscCheck(ndim == usfft->ndim,PETSC_COMM_SELF,PETSC_ERR_USER,"in and out DMDA dimensions must match: %d != %d",usfft->ndim, ndim);
189:   PetscCheck(dof == usfft->dof,PETSC_COMM_SELF,PETSC_ERR_USER,"in and out DMDA dof must match: %d != %d",usfft->dof, dof);
190:   /* Store output dimensions */
191:   /* NB: we reverse the DMDA dimensions, since the DMDA ordering (natural on x-y-z, with x varying the fastest)
192:      is the order opposite of that assumed by FFTW: z varying the fastest */
193:   PetscCall(PetscMalloc1(usfft->ndim+1,&usfft->outdim));
194:   for (i = usfft->ndim; i > 0; --i) usfft->outdim[usfft->ndim-i] = dim[i-1];

196:   /* TODO: Use the new form of DMDACreate() */
197:   #if 0
198:   PetscCall(DMDACreate(comm,usfft->dim, DMDA_NONPERIODIC, DMDA_STENCIL_STAR, usfft->freqSizes[0], usfft->freqSizes[1], usfft->freqSizes[2],
199:                        PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, dof, 0, NULL, NULL, NULL,  0, &usfft->resampleDA));
200:   #endif
201:   PetscCall(DMDAGetVec(usfft->resampleDA, usfft->resample));

203:   /* CONTINUE: Need to build the connectivity "Sieve" attaching sample points to the resample points they are close to */

205:   /* CONTINUE: recalculate matrix sizes based on the connectivity "Sieve" */
206:   /* mat sizes */
207:   m = 1; n = 1;
208:   for (i=0; i<usfft->ndim; i++) {
209:     PetscCheck(usfft->indim[i] > 0,PETSC_COMM_SELF,PETSC_ERR_USER,"indim[%d]=%d must be > 0",i,usfft->indim[i]);
210:     PetscCheck(usfft->outdim[i] > 0,PETSC_COMM_SELF,PETSC_ERR_USER,"outdim[%d]=%d must be > 0",i,usfft->outdim[i]);
211:     n *= usfft->indim[i];
212:     m *= usfft->outdim[i];
213:   }
214:   N        = n*usfft->dof;
215:   M        = m*usfft->dof;
216:   PetscCall(MatSetSizes(*A,M,N,M,N)); /* "in size" is the number of columns, "out size" is the number of rows" */
217:   PetscCall(PetscObjectChangeTypeName((PetscObject)*A,MATSEQUSFFT));
218:   usfft->m = m; usfft->n = n; usfft->M = M; usfft->N = N;
219:   /* FFTW */
220:   usfft->p_forward  = 0;
221:   usfft->p_backward = 0;
222:   usfft->p_flag     = FFTW_ESTIMATE;
223:   /* set Mat ops */
224:   (*A)->ops->mult          = MatMult_SeqUSFFT;
225:   (*A)->ops->multtranspose = MatMultTranspose_SeqUSFFT;
226:   (*A)->assembled          = PETSC_TRUE;
227:   (*A)->ops->destroy       = MatDestroy_SeqUSFFT;
228:   /* get runtime options */
229:   PetscOptionsBegin(((PetscObject)*A)->comm,((PetscObject)*A)->prefix,"USFFT Options","Mat");
230:   PetscCall(PetscOptionsEList("-mat_usfft_fftw_plannerflags","Planner Flags","None",p_flags,4,p_flags[0],&p_flag,&flg));
231:   if (flg) usfft->p_flag = (unsigned)p_flag;
232:   PetscOptionsEnd();
233:   PetscFunctionReturn(PETSC_SUCCESS);
234: } /* MatCreateSeqUSFFT() */

236: #endif