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