Date: 2022-03-31T04:49:06+00:00

Actual source code: tomography.c

      1: /* XH:
      2:     Todo: add cs1f.F90 and adjust makefile.
      3:     Todo: maybe provide code template to generate 1D/2D/3D gradient, DCT transform matrix for D etc.
      4: */
      5: /*
      6:    Include "petsctao.h" so that we can use TAO solvers.  Note that this
      7:    file automatically includes libraries such as:
      8:      petsc.h       - base PETSc routines   petscvec.h - vectors
      9:      petscsys.h    - system routines        petscmat.h - matrices
     10:      petscis.h     - index sets            petscksp.h - Krylov subspace methods
     11:      petscviewer.h - viewers               petscpc.h  - preconditioners

     13: */

     15: #include <petsctao.h>

     17: /*
     18: Description:   BRGN tomography reconstruction example .
     19:                0.5*||Ax-b||^2 + lambda*g(x)
     20: Reference:     None
     21: */

     23: static char help[] = "Finds the least-squares solution to the under constraint linear model Ax = b, with regularizer. \n\
     24:             A is a M*N real matrix (M<N), x is sparse. A good regularizer is an L1 regularizer. \n\
     25:             We find the sparse solution by solving 0.5*||Ax-b||^2 + lambda*||D*x||_1, where lambda (by default 1e-4) is a user specified weight.\n\
     26:             D is the K*N transform matrix so that D*x is sparse. By default D is identity matrix, so that D*x = x.\n";
     27: /*T
     28:    Concepts: TAO^Solving a system of nonlinear equations, nonlinear least squares
     29:    Routines: TaoCreate();
     30:    Routines: TaoSetType();
     31:    Routines: TaoSetSeparableObjectiveRoutine();
     32:    Routines: TaoSetJacobianRoutine();
     33:    Routines: TaoSetSolution();
     34:    Routines: TaoSetFromOptions();
     35:    Routines: TaoSetConvergenceHistory(); TaoGetConvergenceHistory();
     36:    Routines: TaoSolve();
     37:    Routines: TaoView(); TaoDestroy();
     38:    Processors: 1
     39: T*/

     41: /* User-defined application context */
     42: typedef struct {
     43:   /* Working space. linear least square:  res(x) = A*x - b */
     44:   PetscInt  M,N,K;            /* Problem dimension: A is M*N Matrix, D is K*N Matrix */
     45:   Mat       A,D;              /* Coefficients, Dictionary Transform of size M*N and K*N respectively. For linear least square, Jacobian Matrix J = A. For nonlinear least square, it is different from A */
     46:   Vec       b,xGT,xlb,xub;    /* observation b, ground truth xGT, the lower bound and upper bound of x*/
     47: } AppCtx;

     49: /* User provided Routines */
     50: PetscErrorCode InitializeUserData(AppCtx *);
     51: PetscErrorCode FormStartingPoint(Vec,AppCtx *);
     52: PetscErrorCode EvaluateResidual(Tao,Vec,Vec,void *);
     53: PetscErrorCode EvaluateJacobian(Tao,Vec,Mat,Mat,void *);
     54: PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao,Vec,PetscReal *,Vec,void*);
     55: PetscErrorCode EvaluateRegularizerHessian(Tao,Vec,Mat,void*);
     56: PetscErrorCode EvaluateRegularizerHessianProd(Mat,Vec,Vec);

     58: /*--------------------------------------------------------------------*/
     59: int main(int argc,char **argv)
     60: {
     61:   Vec            x,res;              /* solution, function res(x) = A*x-b */
     62:   Mat            Hreg;               /* regularizer Hessian matrix for user specified regularizer*/
     63:   Tao            tao;                /* Tao solver context */
     64:   PetscReal      hist[100],resid[100],v1,v2;
     65:   PetscInt       lits[100];
     66:   AppCtx         user;               /* user-defined work context */
     67:   PetscViewer    fd;   /* used to save result to file */
     68:   char           resultFile[] = "tomographyResult_x";  /* Debug: change from "tomographyResult_x" to "cs1Result_x" */

     70:   PetscInitialize(&argc,&argv,(char *)0,help);

     72:   /* Create TAO solver and set desired solution method */
     73:   TaoCreate(PETSC_COMM_SELF,&tao);
     74:   TaoSetType(tao,TAOBRGN);

     76:   /* User set application context: A, D matrice, and b vector. */
     77:   InitializeUserData(&user);

     79:   /* Allocate solution vector x,  and function vectors Ax-b, */
     80:   VecCreateSeq(PETSC_COMM_SELF,user.N,&x);
     81:   VecCreateSeq(PETSC_COMM_SELF,user.M,&res);

     83:   /* Set initial guess */
     84:   FormStartingPoint(x,&user);

     86:   /* Bind x to tao->solution. */
     87:   TaoSetSolution(tao,x);
     88:   /* Sets the upper and lower bounds of x */
     89:   TaoSetVariableBounds(tao,user.xlb,user.xub);

     91:   /* Bind user.D to tao->data->D */
     92:   TaoBRGNSetDictionaryMatrix(tao,user.D);

     94:   /* Set the residual function and Jacobian routines for least squares. */
     95:   TaoSetResidualRoutine(tao,res,EvaluateResidual,(void*)&user);
     96:   /* Jacobian matrix fixed as user.A for Linear least square problem. */
     97:   TaoSetJacobianResidualRoutine(tao,user.A,user.A,EvaluateJacobian,(void*)&user);

     99:   /* User set the regularizer objective, gradient, and hessian. Set it the same as using l2prox choice, for testing purpose.  */
    100:   TaoBRGNSetRegularizerObjectiveAndGradientRoutine(tao,EvaluateRegularizerObjectiveAndGradient,(void*)&user);
    101:   /* User defined regularizer Hessian setup, here is identiy shell matrix */
    102:   MatCreate(PETSC_COMM_SELF,&Hreg);
    103:   MatSetSizes(Hreg,PETSC_DECIDE,PETSC_DECIDE,user.N,user.N);
    104:   MatSetType(Hreg,MATSHELL);
    105:   MatSetUp(Hreg);
    106:   MatShellSetOperation(Hreg,MATOP_MULT,(void (*)(void))EvaluateRegularizerHessianProd);
    107:   TaoBRGNSetRegularizerHessianRoutine(tao,Hreg,EvaluateRegularizerHessian,(void*)&user);

    109:   /* Check for any TAO command line arguments */
    110:   TaoSetFromOptions(tao);

    112:   TaoSetConvergenceHistory(tao,hist,resid,0,lits,100,PETSC_TRUE);

    114:   /* Perform the Solve */
    115:   TaoSolve(tao);

    117:   /* Save x (reconstruction of object) vector to a binary file, which maybe read from Matlab and convert to a 2D image for comparison. */
    118:   PetscViewerBinaryOpen(PETSC_COMM_SELF,resultFile,FILE_MODE_WRITE,&fd);
    119:   VecView(x,fd);
    120:   PetscViewerDestroy(&fd);

    122:   /* compute the error */
    123:   VecAXPY(x,-1,user.xGT);
    124:   VecNorm(x,NORM_2,&v1);
    125:   VecNorm(user.xGT,NORM_2,&v2);
    126:   PetscPrintf(PETSC_COMM_SELF, "relative reconstruction error: ||x-xGT||/||xGT|| = %6.4e.\n", (double)(v1/v2));

    128:   /* Free TAO data structures */
    129:   TaoDestroy(&tao);

    131:    /* Free PETSc data structures */
    132:   VecDestroy(&x);
    133:   VecDestroy(&res);
    134:   MatDestroy(&Hreg);
    135:   /* Free user data structures */
    136:   MatDestroy(&user.A);
    137:   MatDestroy(&user.D);
    138:   VecDestroy(&user.b);
    139:   VecDestroy(&user.xGT);
    140:   VecDestroy(&user.xlb);
    141:   VecDestroy(&user.xub);
    142:   PetscFinalize();
    143:   return 0;
    144: }

    146: /*--------------------------------------------------------------------*/
    147: /* Evaluate residual function A(x)-b in least square problem ||A(x)-b||^2 */
    148: PetscErrorCode EvaluateResidual(Tao tao,Vec X,Vec F,void *ptr)
    149: {
    150:   AppCtx         *user = (AppCtx *)ptr;

    152:   /* Compute Ax - b */
    153:   MatMult(user->A,X,F);
    154:   VecAXPY(F,-1,user->b);
    155:   PetscLogFlops(2.0*user->M*user->N);
    156:   return 0;
    157: }

    159: /*------------------------------------------------------------*/
    160: PetscErrorCode EvaluateJacobian(Tao tao,Vec X,Mat J,Mat Jpre,void *ptr)
    161: {
    162:   /* Jacobian is not changing here, so use a empty dummy function here.  J[m][n] = df[m]/dx[n] = A[m][n] for linear least square */
    163:   return 0;
    164: }

    166: /* ------------------------------------------------------------ */
    167: PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao tao,Vec X,PetscReal *f_reg,Vec G_reg,void *ptr)
    168: {
    169:   /* compute regularizer objective = 0.5*x'*x */
    170:   VecDot(X,X,f_reg);
    171:   *f_reg *= 0.5;
    172:   /* compute regularizer gradient = x */
    173:   VecCopy(X,G_reg);
    174:   return 0;
    175: }

    177: PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg,Vec in,Vec out)
    178: {
    179:   VecCopy(in,out);
    180:   return 0;
    181: }

    183: /* ------------------------------------------------------------ */
    184: PetscErrorCode EvaluateRegularizerHessian(Tao tao,Vec X,Mat Hreg,void *ptr)
    185: {
    186:   /* Hessian for regularizer objective = 0.5*x'*x is identity matrix, and is not changing*/
    187:   return 0;
    188: }

    190: /* ------------------------------------------------------------ */
    191: PetscErrorCode FormStartingPoint(Vec X,AppCtx *user)
    192: {
    193:   VecSet(X,0.0);
    194:   return 0;
    195: }

    197: /* ---------------------------------------------------------------------- */
    198: PetscErrorCode InitializeUserData(AppCtx *user)
    199: {
    200:   PetscInt       k,n; /* indices for row and columns of D. */
    201:   char           dataFile[] = "tomographyData_A_b_xGT";   /* Matrix A and vectors b, xGT(ground truth) binary files generated by Matlab. Debug: change from "tomographyData_A_b_xGT" to "cs1Data_A_b_xGT". */
    202:   PetscInt       dictChoice = 1; /* choose from 0:identity, 1:gradient1D, 2:gradient2D, 3:DCT etc */
    203:   PetscViewer    fd;   /* used to load data from file */
    204:   PetscReal      v;

    207:   /*
    208:   Matrix Vector read and write refer to:
    209:   https://petsc.org/release/src/mat/tutorials/ex10.c
    210:   https://petsc.org/release/src/mat/tutorials/ex12.c
    211:  */
    212:   /* Load the A matrix, b vector, and xGT vector from a binary file. */
    213:   PetscViewerBinaryOpen(PETSC_COMM_WORLD,dataFile,FILE_MODE_READ,&fd);
    214:   MatCreate(PETSC_COMM_WORLD,&user->A);
    215:   MatSetType(user->A,MATSEQAIJ);
    216:   MatLoad(user->A,fd);
    217:   VecCreate(PETSC_COMM_WORLD,&user->b);
    218:   VecLoad(user->b,fd);
    219:   VecCreate(PETSC_COMM_WORLD,&user->xGT);
    220:   VecLoad(user->xGT,fd);
    221:   PetscViewerDestroy(&fd);
    222:   VecDuplicate(user->xGT,&(user->xlb));
    223:   VecSet(user->xlb,0.0);
    224:   VecDuplicate(user->xGT,&(user->xub));
    225:   VecSet(user->xub,PETSC_INFINITY);

    227:   /* Specify the size */
    228:   MatGetSize(user->A,&user->M,&user->N);

    230:   /* shortcut, when D is identity matrix, we may just specify it as NULL, and brgn will treat D*x as x without actually computing D*x.
    231:   if (dictChoice == 0) {
    232:     user->D = NULL;
    233:     return 0;
    234:   }
    235:   */

    237:   /* Speficy D */
    238:   /* (1) Specify D Size */
    239:   switch (dictChoice) {
    240:     case 0: /* 0:identity */
    241:       user->K = user->N;
    242:       break;
    243:     case 1: /* 1:gradient1D */
    244:       user->K = user->N-1;
    245:       break;
    246:   }

    248:   MatCreate(PETSC_COMM_SELF,&user->D);
    249:   MatSetSizes(user->D,PETSC_DECIDE,PETSC_DECIDE,user->K,user->N);
    250:   MatSetFromOptions(user->D);
    251:   MatSetUp(user->D);

    253:   /* (2) Specify D Content */
    254:   switch (dictChoice) {
    255:     case 0: /* 0:identity */
    256:       for (k=0; k<user->K; k++) {
    257:         v = 1.0;
    258:         MatSetValues(user->D,1,&k,1,&k,&v,INSERT_VALUES);
    259:       }
    260:       break;
    261:     case 1: /* 1:gradient1D.  [-1, 1, 0,...; 0, -1, 1, 0, ...] */
    262:       for (k=0; k<user->K; k++) {
    263:         v = 1.0;
    264:         n = k+1;
    265:         MatSetValues(user->D,1,&k,1,&n,&v,INSERT_VALUES);
    266:         v = -1.0;
    267:         MatSetValues(user->D,1,&k,1,&k,&v,INSERT_VALUES);
    268:       }
    269:       break;
    270:   }
    271:   MatAssemblyBegin(user->D,MAT_FINAL_ASSEMBLY);
    272:   MatAssemblyEnd(user->D,MAT_FINAL_ASSEMBLY);

    274:   return 0;
    275: }

    277: /*TEST

    279:    build:
    280:       requires: !complex !single !__float128 !defined(PETSC_USE_64BIT_INDICES)

    282:    test:
    283:       localrunfiles: tomographyData_A_b_xGT
    284:       args: -tao_max_it 1000 -tao_brgn_regularization_type l1dict -tao_brgn_regularizer_weight 1e-8 -tao_brgn_l1_smooth_epsilon 1e-6 -tao_gatol 1.e-8

    286:    test:
    287:       suffix: 2
    288:       localrunfiles: tomographyData_A_b_xGT
    289:       args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type l2prox -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6

    291:    test:
    292:       suffix: 3
    293:       localrunfiles: tomographyData_A_b_xGT
    294:       args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type user -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6

    296: TEST*/