Date: 2022-08-02T03:14:15+00:00

Actual source code: matusfft.c

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

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

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

     23: 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: #if 0
     31:   /* resample x to usfft->resample */
     32:   MatResample_USFFT_Private(A, x);

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

     46:       switch (usfft->dim) {
     47:       case 1:
     48:         *plan = fftw_plan_dft_1d(usfft->outdim[0],(fftw_complex*)x_array,(fftw_complex*)y_array,direction,usfft->p_flag);
     49:         break;
     50:       case 2:
     51:         *plan = fftw_plan_dft_2d(usfft->outdim[0],usfft->outdim[1],(fftw_complex*)x_array,(fftw_complex*)y_array,direction,usfft->p_flag);
     52:         break;
     53:       case 3:
     54:         *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);
     55:         break;
     56:       default:
     57:         *plan = fftw_plan_dft(usfft->ndim,usfft->outdim,(fftw_complex*)x_array,(fftw_complex*)y_array,direction,usfft->p_flag);
     58:         break;
     59:       }
     60:       fftw_execute(*plan);
     61:     } /* if (dof == 1) */
     62:     else { /* if (dof > 1) */
     63:       *plan = fftw_plan_many_dft(/*rank*/usfft->ndim, /*n*/usfft->outdim, /*howmany*/usfft->dof,
     64:                                  (fftw_complex*)x_array, /*nembed*/usfft->outdim, /*stride*/usfft->dof, /*dist*/1,
     65:                                  (fftw_complex*)y_array, /*nembed*/usfft->outdim, /*stride*/usfft->dof, /*dist*/1,
     66:                                  /*sign*/direction, /*flags*/usfft->p_flag);
     67:       fftw_execute(*plan);
     68:     } /* if (dof > 1) */
     69:   } /* if (!*plan) */
     70:   else {  /* if (*plan) */
     71:     /* use existing plan */
     72:     fftw_execute_dft(*plan,(fftw_complex*)x_array,(fftw_complex*)y_array);
     73:   }
     74:   VecRestoreArray(y,&y_array);
     75:   VecRestoreArray(x,&x_array);
     76: #endif
     77:   return 0;
     78: } /* MatApply_USFFT_Private() */

     80: #if 0
     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:   return 0;
     87: } /* MatUSFFT_ProjectOnBattleLemarie_Private() */

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

     93:   return 0;
     94: } /* MatInterpolate_USFFT_Private() */

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

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

    105: PetscErrorCode MatMultTranspose_SeqUSFFT(Mat A,Vec x,Vec y)
    106: {
    107:   Mat_USFFT      *usfft = (Mat_USFFT*)A->data;

    109:   /* NB: for now we use outdim for both x and y; this will change once a full USFFT is implemented */
    110:   MatApply_USFFT_Private(usfft, &usfft->p_backward, FFTW_BACKWARD, x,y);
    111:   return 0;
    112: }

    114: PetscErrorCode MatDestroy_SeqUSFFT(Mat A)
    115: {
    116:   Mat_USFFT      *usfft = (Mat_USFFT*)A->data;

    118:   fftw_destroy_plan(usfft->p_forward);
    119:   fftw_destroy_plan(usfft->p_backward);
    120:   PetscFree(usfft->indim);
    121:   PetscFree(usfft->outdim);
    122:   PetscFree(usfft);
    123:   PetscObjectChangeTypeName((PetscObject)A,0);
    124:   return 0;
    125: } /* MatDestroy_SeqUSFFT() */

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

    131:    Collective

    133:    Input Parameter:
    134: .   da - geometry of the domain encoded by a DMDA

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

    139:   Options Database Keys:
    140: . -mat_usfft_plannerflags - set the FFTW planner flags

    142:    Level: intermediate

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

    157:   PetscObjectGetComm((PetscObject)inda, &comm);
    158:   MPI_Comm_size(comm, &size);
    160:   PetscObjectGetComm((PetscObject)outda, &comm);
    161:   MPI_Comm_size(comm, &size);
    163:   MatCreate(comm,A);
    164:   PetscNewLog(*A,&usfft);
    165:   (*A)->data   = (void*)usfft;
    166:   usfft->inda  = inda;
    167:   usfft->outda = outda;
    168:   /* inda */
    169:   DMDAGetInfo(usfft->inda, &ndim, dim+0, dim+1, dim+2, NULL, NULL, NULL, &dof, NULL, NULL, NULL);
    172:   usfft->ndim   = ndim;
    173:   usfft->dof    = dof;
    174:   usfft->freqDA = freqDA;
    175:   /* NB: we reverse the freq and resample DMDA sizes, since the DMDA ordering (natural on x-y-z, with x varying the fastest)
    176:      is the order opposite of that assumed by FFTW: z varying the fastest */
    177:   PetscMalloc1(usfft->ndim+1,&usfft->indim);
    178:   for (i = usfft->ndim; i > 0; --i) usfft->indim[usfft->ndim-i] = dim[i-1];

    180:   /* outda */
    181:   DMDAGetInfo(usfft->outda, &ndim, dim+0, dim+1, dim+2, NULL, NULL, NULL, &dof, NULL, NULL, NULL);
    184:   /* Store output dimensions */
    185:   /* NB: we reverse the DMDA dimensions, since the DMDA ordering (natural on x-y-z, with x varying the fastest)
    186:      is the order opposite of that assumed by FFTW: z varying the fastest */
    187:   PetscMalloc1(usfft->ndim+1,&usfft->outdim);
    188:   for (i = usfft->ndim; i > 0; --i) usfft->outdim[usfft->ndim-i] = dim[i-1];

    190:   /* TODO: Use the new form of DMDACreate() */
    191: #if 0
    192:   DMDACreate(comm,usfft->dim, DMDA_NONPERIODIC, DMDA_STENCIL_STAR, usfft->freqSizes[0], usfft->freqSizes[1], usfft->freqSizes[2],
    193:                     PETSC_DECIDE, PETSC_DECIDE, PETSC_DECIDE, dof, 0, NULL, NULL, NULL,  0, &(usfft->resampleDA));
    194: #endif
    195:   DMDAGetVec(usfft->resampleDA, usfft->resample);

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

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

    230: #endif