Date: 2025-03-28T21:10:44+00:00

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:   Note:
    150:   This does not currently exist. There is some code in place but apparently unfinished and commented out with #ifdef 0

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

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

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

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

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

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

    239: #endif