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" */

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

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

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

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

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

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

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

 82:   /* Set the residual function and Jacobian routines for least squares. */
 83:   TaoSetResidualRoutine(tao, res, EvaluateResidual, (void *)&user);
 84:   /* Jacobian matrix fixed as user.A for Linear least square problem. */
 85:   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:   TaoBRGNSetRegularizerObjectiveAndGradientRoutine(tao, EvaluateRegularizerObjectiveAndGradient, (void *)&user);
 89:   /* User defined regularizer Hessian setup, here is identiy shell matrix */
 90:   MatCreate(PETSC_COMM_SELF, &Hreg);
 91:   MatSetSizes(Hreg, PETSC_DECIDE, PETSC_DECIDE, user.N, user.N);
 92:   MatSetType(Hreg, MATSHELL);
 93:   MatSetUp(Hreg);
 94:   MatShellSetOperation(Hreg, MATOP_MULT, (void (*)(void))EvaluateRegularizerHessianProd);
 95:   TaoBRGNSetRegularizerHessianRoutine(tao, Hreg, EvaluateRegularizerHessian, (void *)&user);

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

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

102:   /* Perform the Solve */
103:   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:   PetscViewerBinaryOpen(PETSC_COMM_SELF, resultFile, FILE_MODE_WRITE, &fd);
107:   VecView(x, fd);
108:   PetscViewerDestroy(&fd);

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

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

119:   /* Free PETSc data structures */
120:   VecDestroy(&x);
121:   VecDestroy(&res);
122:   MatDestroy(&Hreg);
123:   /* Free user data structures */
124:   MatDestroy(&user.A);
125:   MatDestroy(&user.D);
126:   VecDestroy(&user.b);
127:   VecDestroy(&user.xGT);
128:   VecDestroy(&user.xlb);
129:   VecDestroy(&user.xub);
130:   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:   /* Compute Ax - b */
141:   MatMult(user->A, X, F);
142:   VecAXPY(F, -1, user->b);
143:   PetscLogFlops(2.0 * user->M * user->N);
144:   return 0;
145: }

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

154: /* ------------------------------------------------------------ */
155: PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao tao, Vec X, PetscReal *f_reg, Vec G_reg, void *ptr)
156: {
157:   /* compute regularizer objective = 0.5*x'*x */
158:   VecDot(X, X, f_reg);
159:   *f_reg *= 0.5;
160:   /* compute regularizer gradient = x */
161:   VecCopy(X, G_reg);
162:   return 0;
163: }

165: PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg, Vec in, Vec out)
166: {
167:   VecCopy(in, out);
168:   return 0;
169: }

171: /* ------------------------------------------------------------ */
172: PetscErrorCode EvaluateRegularizerHessian(Tao tao, Vec X, Mat Hreg, void *ptr)
173: {
174:   /* Hessian for regularizer objective = 0.5*x'*x is identity matrix, and is not changing*/
175:   return 0;
176: }

178: /* ------------------------------------------------------------ */
179: PetscErrorCode FormStartingPoint(Vec X, AppCtx *user)
180: {
181:   VecSet(X, 0.0);
182:   return 0;
183: }

185: /* ---------------------------------------------------------------------- */
186: PetscErrorCode InitializeUserData(AppCtx *user)
187: {
188:   PetscInt    k, n;                                  /* indices for row and columns of D. */
189:   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". */
190:   PetscInt    dictChoice = 1;                        /* choose from 0:identity, 1:gradient1D, 2:gradient2D, 3:DCT etc */
191:   PetscViewer fd;                                    /* used to load data from file */
192:   PetscReal   v;


195:   /*
196:   Matrix Vector read and write refer to:
197:   https://petsc.org/release/src/mat/tutorials/ex10.c
198:   https://petsc.org/release/src/mat/tutorials/ex12.c
199:  */
200:   /* Load the A matrix, b vector, and xGT vector from a binary file. */
201:   PetscViewerBinaryOpen(PETSC_COMM_WORLD, dataFile, FILE_MODE_READ, &fd);
202:   MatCreate(PETSC_COMM_WORLD, &user->A);
203:   MatSetType(user->A, MATSEQAIJ);
204:   MatLoad(user->A, fd);
205:   VecCreate(PETSC_COMM_WORLD, &user->b);
206:   VecLoad(user->b, fd);
207:   VecCreate(PETSC_COMM_WORLD, &user->xGT);
208:   VecLoad(user->xGT, fd);
209:   PetscViewerDestroy(&fd);
210:   VecDuplicate(user->xGT, &(user->xlb));
211:   VecSet(user->xlb, 0.0);
212:   VecDuplicate(user->xGT, &(user->xub));
213:   VecSet(user->xub, PETSC_INFINITY);

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

218:   /* 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.
219:   if (dictChoice == 0) {
220:     user->D = NULL;
221:     return 0;
222:   }
223:   */

225:   /* Specify D */
226:   /* (1) Specify D Size */
227:   switch (dictChoice) {
228:   case 0: /* 0:identity */
229:     user->K = user->N;
230:     break;
231:   case 1: /* 1:gradient1D */
232:     user->K = user->N - 1;
233:     break;
234:   }

236:   MatCreate(PETSC_COMM_SELF, &user->D);
237:   MatSetSizes(user->D, PETSC_DECIDE, PETSC_DECIDE, user->K, user->N);
238:   MatSetFromOptions(user->D);
239:   MatSetUp(user->D);

241:   /* (2) Specify D Content */
242:   switch (dictChoice) {
243:   case 0: /* 0:identity */
244:     for (k = 0; k < user->K; k++) {
245:       v = 1.0;
246:       MatSetValues(user->D, 1, &k, 1, &k, &v, INSERT_VALUES);
247:     }
248:     break;
249:   case 1: /* 1:gradient1D.  [-1, 1, 0,...; 0, -1, 1, 0, ...] */
250:     for (k = 0; k < user->K; k++) {
251:       v = 1.0;
252:       n = k + 1;
253:       MatSetValues(user->D, 1, &k, 1, &n, &v, INSERT_VALUES);
254:       v = -1.0;
255:       MatSetValues(user->D, 1, &k, 1, &k, &v, INSERT_VALUES);
256:     }
257:     break;
258:   }
259:   MatAssemblyBegin(user->D, MAT_FINAL_ASSEMBLY);
260:   MatAssemblyEnd(user->D, MAT_FINAL_ASSEMBLY);

262:   return 0;
263: }

265: /*TEST

267:    build:
268:       requires: !complex !single !__float128 !defined(PETSC_USE_64BIT_INDICES)

270:    test:
271:       localrunfiles: tomographyData_A_b_xGT
272:       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

274:    test:
275:       suffix: 2
276:       localrunfiles: tomographyData_A_b_xGT
277:       args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type l2prox -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6

279:    test:
280:       suffix: 3
281:       localrunfiles: tomographyData_A_b_xGT
282:       args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type user -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6

284: TEST*/