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";

     28: /* User-defined application context */
     29: typedef struct {
     30:   /* Working space. linear least square:  res(x) = A*x - b */
     31:   PetscInt M, N, K;          /* Problem dimension: A is M*N Matrix, D is K*N Matrix */
     32:   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 */
     33:   Vec      b, xGT, xlb, xub; /* observation b, ground truth xGT, the lower bound and upper bound of x*/
     34: } AppCtx;

     36: /* User provided Routines */
     37: PetscErrorCode InitializeUserData(AppCtx *);
     38: PetscErrorCode FormStartingPoint(Vec, AppCtx *);
     39: PetscErrorCode EvaluateResidual(Tao, Vec, Vec, void *);
     40: PetscErrorCode EvaluateJacobian(Tao, Vec, Mat, Mat, void *);
     41: PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao, Vec, PetscReal *, Vec, void *);
     42: PetscErrorCode EvaluateRegularizerHessian(Tao, Vec, Mat, void *);
     43: PetscErrorCode EvaluateRegularizerHessianProd(Mat, Vec, Vec);

     45: /*--------------------------------------------------------------------*/
     46: int main(int argc, char **argv)
     47: {
     48:   Vec         x, res; /* solution, function res(x) = A*x-b */
     49:   Mat         Hreg;   /* regularizer Hessian matrix for user specified regularizer*/
     50:   Tao         tao;    /* Tao solver context */
     51:   PetscReal   hist[100], resid[100], v1, v2;
     52:   PetscInt    lits[100];
     53:   AppCtx      user;                                /* user-defined work context */
     54:   PetscViewer fd;                                  /* used to save result to file */
     55:   char        resultFile[] = "tomographyResult_x"; /* Debug: change from "tomographyResult_x" to "cs1Result_x" */

     57:   PetscFunctionBeginUser;
     58:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));

     60:   /* Create TAO solver and set desired solution method */
     61:   PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
     62:   PetscCall(TaoSetType(tao, TAOBRGN));

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

     67:   /* Allocate solution vector x,  and function vectors Ax-b, */
     68:   PetscCall(VecCreateSeq(PETSC_COMM_SELF, user.N, &x));
     69:   PetscCall(VecCreateSeq(PETSC_COMM_SELF, user.M, &res));

     71:   /* Set initial guess */
     72:   PetscCall(FormStartingPoint(x, &user));

     74:   /* Bind x to tao->solution. */
     75:   PetscCall(TaoSetSolution(tao, x));
     76:   /* Sets the upper and lower bounds of x */
     77:   PetscCall(TaoSetVariableBounds(tao, user.xlb, user.xub));

     79:   /* Bind user.D to tao->data->D */
     80:   PetscCall(TaoBRGNSetDictionaryMatrix(tao, user.D));

     82:   /* Set the residual function and Jacobian routines for least squares. */
     83:   PetscCall(TaoSetResidualRoutine(tao, res, EvaluateResidual, (void *)&user));
     84:   /* Jacobian matrix fixed as user.A for Linear least square problem. */
     85:   PetscCall(TaoSetJacobianResidualRoutine(tao, user.A, user.A, EvaluateJacobian, (void *)&user));

     87:   /* User set the regularizer objective, gradient, and hessian. Set it the same as using l2prox choice, for testing purpose.  */
     88:   PetscCall(TaoBRGNSetRegularizerObjectiveAndGradientRoutine(tao, EvaluateRegularizerObjectiveAndGradient, (void *)&user));
     89:   /* User defined regularizer Hessian setup, here is identity shell matrix */
     90:   PetscCall(MatCreate(PETSC_COMM_SELF, &Hreg));
     91:   PetscCall(MatSetSizes(Hreg, PETSC_DECIDE, PETSC_DECIDE, user.N, user.N));
     92:   PetscCall(MatSetType(Hreg, MATSHELL));
     93:   PetscCall(MatSetUp(Hreg));
     94:   PetscCall(MatShellSetOperation(Hreg, MATOP_MULT, (PetscErrorCodeFn *)EvaluateRegularizerHessianProd));
     95:   PetscCall(TaoBRGNSetRegularizerHessianRoutine(tao, Hreg, EvaluateRegularizerHessian, (void *)&user));

     97:   /* Check for any TAO command line arguments */
     98:   PetscCall(TaoSetFromOptions(tao));

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

    102:   /* Perform the Solve */
    103:   PetscCall(TaoSolve(tao));

    105:   /* Save x (reconstruction of object) vector to a binary file, which maybe read from MATLAB and convert to a 2D image for comparison. */
    106:   PetscCall(PetscViewerBinaryOpen(PETSC_COMM_SELF, resultFile, FILE_MODE_WRITE, &fd));
    107:   PetscCall(VecView(x, fd));
    108:   PetscCall(PetscViewerDestroy(&fd));

    110:   /* compute the error */
    111:   PetscCall(VecAXPY(x, -1, user.xGT));
    112:   PetscCall(VecNorm(x, NORM_2, &v1));
    113:   PetscCall(VecNorm(user.xGT, NORM_2, &v2));
    114:   PetscCall(PetscPrintf(PETSC_COMM_SELF, "relative reconstruction error: ||x-xGT||/||xGT|| = %6.4e.\n", (double)(v1 / v2)));

    116:   /* Free TAO data structures */
    117:   PetscCall(TaoDestroy(&tao));

    119:   /* Free PETSc data structures */
    120:   PetscCall(VecDestroy(&x));
    121:   PetscCall(VecDestroy(&res));
    122:   PetscCall(MatDestroy(&Hreg));
    123:   /* Free user data structures */
    124:   PetscCall(MatDestroy(&user.A));
    125:   PetscCall(MatDestroy(&user.D));
    126:   PetscCall(VecDestroy(&user.b));
    127:   PetscCall(VecDestroy(&user.xGT));
    128:   PetscCall(VecDestroy(&user.xlb));
    129:   PetscCall(VecDestroy(&user.xub));
    130:   PetscCall(PetscFinalize());
    131:   return 0;
    132: }

    134: /*--------------------------------------------------------------------*/
    135: /* Evaluate residual function A(x)-b in least square problem ||A(x)-b||^2 */
    136: PetscErrorCode EvaluateResidual(Tao tao, Vec X, Vec F, void *ptr)
    137: {
    138:   AppCtx *user = (AppCtx *)ptr;

    140:   PetscFunctionBegin;
    141:   /* Compute Ax - b */
    142:   PetscCall(MatMult(user->A, X, F));
    143:   PetscCall(VecAXPY(F, -1, user->b));
    144:   PetscCall(PetscLogFlops(2.0 * user->M * user->N));
    145:   PetscFunctionReturn(PETSC_SUCCESS);
    146: }

    148: /*------------------------------------------------------------*/
    149: PetscErrorCode EvaluateJacobian(Tao tao, Vec X, Mat J, Mat Jpre, void *ptr)
    150: {
    151:   /* 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 */
    152:   PetscFunctionBegin;
    153:   PetscFunctionReturn(PETSC_SUCCESS);
    154: }

    156: /* ------------------------------------------------------------ */
    157: PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao tao, Vec X, PetscReal *f_reg, Vec G_reg, void *ptr)
    158: {
    159:   PetscFunctionBegin;
    160:   /* compute regularizer objective = 0.5*x'*x */
    161:   PetscCall(VecDot(X, X, f_reg));
    162:   *f_reg *= 0.5;
    163:   /* compute regularizer gradient = x */
    164:   PetscCall(VecCopy(X, G_reg));
    165:   PetscFunctionReturn(PETSC_SUCCESS);
    166: }

    168: PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg, Vec in, Vec out)
    169: {
    170:   PetscFunctionBegin;
    171:   PetscCall(VecCopy(in, out));
    172:   PetscFunctionReturn(PETSC_SUCCESS);
    173: }

    175: /* ------------------------------------------------------------ */
    176: PetscErrorCode EvaluateRegularizerHessian(Tao tao, Vec X, Mat Hreg, void *ptr)
    177: {
    178:   /* Hessian for regularizer objective = 0.5*x'*x is identity matrix, and is not changing*/
    179:   PetscFunctionBegin;
    180:   PetscFunctionReturn(PETSC_SUCCESS);
    181: }

    183: /* ------------------------------------------------------------ */
    184: PetscErrorCode FormStartingPoint(Vec X, AppCtx *user)
    185: {
    186:   PetscFunctionBegin;
    187:   PetscCall(VecSet(X, 0.0));
    188:   PetscFunctionReturn(PETSC_SUCCESS);
    189: }

    191: /* ---------------------------------------------------------------------- */
    192: PetscErrorCode InitializeUserData(AppCtx *user)
    193: {
    194:   PetscInt    k, n;                                                   /* indices for row and columns of D. */
    195:   char        dataFile[PETSC_MAX_PATH_LEN], path[PETSC_MAX_PATH_LEN]; /* 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". */
    196:   PetscInt    dictChoice = 1;                                         /* choose from 0:identity, 1:gradient1D, 2:gradient2D, 3:DCT etc */
    197:   PetscViewer fd;                                                     /* used to load data from file */
    198:   PetscReal   v;
    199:   PetscBool   flg;

    201:   PetscFunctionBegin;
    202:   /*
    203:   Matrix Vector read and write refer to:
    204:   https://petsc.org/release/src/mat/tutorials/ex10.c
    205:   https://petsc.org/release/src/mat/tutorials/ex12.c
    206:  */
    207:   PetscCall(PetscOptionsGetString(NULL, NULL, "-path", path, sizeof(path), &flg));
    208:   PetscCheck(flg, PETSC_COMM_WORLD, PETSC_ERR_USER, "Must specify -path ${DATAFILESPATH}/tao/tomography");
    209:   /* Load the A matrix, b vector, and xGT vector from a binary file. */
    210:   PetscCall(PetscSNPrintf(dataFile, sizeof(dataFile), "%s/tomographyData_A_b_xGT", path));
    211:   PetscCall(PetscViewerBinaryOpen(PETSC_COMM_WORLD, dataFile, FILE_MODE_READ, &fd));
    212:   PetscCall(MatCreate(PETSC_COMM_WORLD, &user->A));
    213:   PetscCall(MatSetType(user->A, MATSEQAIJ));
    214:   PetscCall(MatLoad(user->A, fd));
    215:   PetscCall(VecCreate(PETSC_COMM_WORLD, &user->b));
    216:   PetscCall(VecLoad(user->b, fd));
    217:   PetscCall(VecCreate(PETSC_COMM_WORLD, &user->xGT));
    218:   PetscCall(VecLoad(user->xGT, fd));
    219:   PetscCall(PetscViewerDestroy(&fd));
    220:   PetscCall(VecDuplicate(user->xGT, &user->xlb));
    221:   PetscCall(VecSet(user->xlb, 0.0));
    222:   PetscCall(VecDuplicate(user->xGT, &user->xub));
    223:   PetscCall(VecSet(user->xub, PETSC_INFINITY));

    225:   /* Specify the size */
    226:   PetscCall(MatGetSize(user->A, &user->M, &user->N));

    228:   /* 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.
    229:   if (dictChoice == 0) {
    230:     user->D = NULL;
    231:     PetscFunctionReturn(PETSC_SUCCESS);
    232:   }
    233:   */

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

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

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

    274: /*TEST

    276:    build:
    277:       requires: !complex !single !__float128 !defined(PETSC_USE_64BIT_INDICES)

    279:    testset:
    280:       requires: datafilespath
    281:       args: -path ${DATAFILESPATH}/tao/tomography

    283:       test:
    284:          args: -tao_max_it 10 -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:          args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type l2prox -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6 -tao_brgn_subsolver_tao_monitor

    290:       test:
    291:          suffix: 3
    292:          args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type user -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6 -tao_brgn_subsolver_tao_monitor

    294: TEST*/