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: PetscInitialize(&argc,&argv,(char *)0,help);
59: /* Create TAO solver and set desired solution method */
60: TaoCreate(PETSC_COMM_SELF,&tao);
61: TaoSetType(tao,TAOBRGN);
63: /* User set application context: A, D matrice, and b vector. */
64: InitializeUserData(&user);
66: /* Allocate solution vector x, and function vectors Ax-b, */
67: VecCreateSeq(PETSC_COMM_SELF,user.N,&x);
68: VecCreateSeq(PETSC_COMM_SELF,user.M,&res);
70: /* Set initial guess */
71: FormStartingPoint(x,&user);
73: /* Bind x to tao->solution. */
74: TaoSetSolution(tao,x);
75: /* Sets the upper and lower bounds of x */
76: TaoSetVariableBounds(tao,user.xlb,user.xub);
78: /* Bind user.D to tao->data->D */
79: TaoBRGNSetDictionaryMatrix(tao,user.D);
81: /* Set the residual function and Jacobian routines for least squares. */
82: TaoSetResidualRoutine(tao,res,EvaluateResidual,(void*)&user);
83: /* Jacobian matrix fixed as user.A for Linear least square problem. */
84: TaoSetJacobianResidualRoutine(tao,user.A,user.A,EvaluateJacobian,(void*)&user);
86: /* User set the regularizer objective, gradient, and hessian. Set it the same as using l2prox choice, for testing purpose. */
87: TaoBRGNSetRegularizerObjectiveAndGradientRoutine(tao,EvaluateRegularizerObjectiveAndGradient,(void*)&user);
88: /* User defined regularizer Hessian setup, here is identiy shell matrix */
89: MatCreate(PETSC_COMM_SELF,&Hreg);
90: MatSetSizes(Hreg,PETSC_DECIDE,PETSC_DECIDE,user.N,user.N);
91: MatSetType(Hreg,MATSHELL);
92: MatSetUp(Hreg);
93: MatShellSetOperation(Hreg,MATOP_MULT,(void (*)(void))EvaluateRegularizerHessianProd);
94: TaoBRGNSetRegularizerHessianRoutine(tao,Hreg,EvaluateRegularizerHessian,(void*)&user);
96: /* Check for any TAO command line arguments */
97: TaoSetFromOptions(tao);
99: TaoSetConvergenceHistory(tao,hist,resid,0,lits,100,PETSC_TRUE);
101: /* Perform the Solve */
102: TaoSolve(tao);
104: /* Save x (reconstruction of object) vector to a binary file, which maybe read from Matlab and convert to a 2D image for comparison. */
105: PetscViewerBinaryOpen(PETSC_COMM_SELF,resultFile,FILE_MODE_WRITE,&fd);
106: VecView(x,fd);
107: PetscViewerDestroy(&fd);
109: /* compute the error */
110: VecAXPY(x,-1,user.xGT);
111: VecNorm(x,NORM_2,&v1);
112: VecNorm(user.xGT,NORM_2,&v2);
113: PetscPrintf(PETSC_COMM_SELF, "relative reconstruction error: ||x-xGT||/||xGT|| = %6.4e.\n", (double)(v1/v2));
115: /* Free TAO data structures */
116: TaoDestroy(&tao);
118: /* Free PETSc data structures */
119: VecDestroy(&x);
120: VecDestroy(&res);
121: MatDestroy(&Hreg);
122: /* Free user data structures */
123: MatDestroy(&user.A);
124: MatDestroy(&user.D);
125: VecDestroy(&user.b);
126: VecDestroy(&user.xGT);
127: VecDestroy(&user.xlb);
128: VecDestroy(&user.xub);
129: PetscFinalize();
130: return 0;
131: }
133: /*--------------------------------------------------------------------*/
134: /* Evaluate residual function A(x)-b in least square problem ||A(x)-b||^2 */
135: PetscErrorCode EvaluateResidual(Tao tao,Vec X,Vec F,void *ptr)
136: {
137: AppCtx *user = (AppCtx *)ptr;
139: /* Compute Ax - b */
140: MatMult(user->A,X,F);
141: VecAXPY(F,-1,user->b);
142: PetscLogFlops(2.0*user->M*user->N);
143: return 0;
144: }
146: /*------------------------------------------------------------*/
147: PetscErrorCode EvaluateJacobian(Tao tao,Vec X,Mat J,Mat Jpre,void *ptr)
148: {
149: /* 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 */
150: return 0;
151: }
153: /* ------------------------------------------------------------ */
154: PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao tao,Vec X,PetscReal *f_reg,Vec G_reg,void *ptr)
155: {
156: /* compute regularizer objective = 0.5*x'*x */
157: VecDot(X,X,f_reg);
158: *f_reg *= 0.5;
159: /* compute regularizer gradient = x */
160: VecCopy(X,G_reg);
161: return 0;
162: }
164: PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg,Vec in,Vec out)
165: {
166: VecCopy(in,out);
167: return 0;
168: }
170: /* ------------------------------------------------------------ */
171: PetscErrorCode EvaluateRegularizerHessian(Tao tao,Vec X,Mat Hreg,void *ptr)
172: {
173: /* Hessian for regularizer objective = 0.5*x'*x is identity matrix, and is not changing*/
174: return 0;
175: }
177: /* ------------------------------------------------------------ */
178: PetscErrorCode FormStartingPoint(Vec X,AppCtx *user)
179: {
180: VecSet(X,0.0);
181: return 0;
182: }
184: /* ---------------------------------------------------------------------- */
185: PetscErrorCode InitializeUserData(AppCtx *user)
186: {
187: PetscInt k,n; /* indices for row and columns of D. */
188: 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". */
189: PetscInt dictChoice = 1; /* choose from 0:identity, 1:gradient1D, 2:gradient2D, 3:DCT etc */
190: PetscViewer fd; /* used to load data from file */
191: PetscReal v;
194: /*
195: Matrix Vector read and write refer to:
196: https://petsc.org/release/src/mat/tutorials/ex10.c
197: https://petsc.org/release/src/mat/tutorials/ex12.c
198: */
199: /* Load the A matrix, b vector, and xGT vector from a binary file. */
200: PetscViewerBinaryOpen(PETSC_COMM_WORLD,dataFile,FILE_MODE_READ,&fd);
201: MatCreate(PETSC_COMM_WORLD,&user->A);
202: MatSetType(user->A,MATSEQAIJ);
203: MatLoad(user->A,fd);
204: VecCreate(PETSC_COMM_WORLD,&user->b);
205: VecLoad(user->b,fd);
206: VecCreate(PETSC_COMM_WORLD,&user->xGT);
207: VecLoad(user->xGT,fd);
208: PetscViewerDestroy(&fd);
209: VecDuplicate(user->xGT,&(user->xlb));
210: VecSet(user->xlb,0.0);
211: VecDuplicate(user->xGT,&(user->xub));
212: VecSet(user->xub,PETSC_INFINITY);
214: /* Specify the size */
215: MatGetSize(user->A,&user->M,&user->N);
217: /* 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.
218: if (dictChoice == 0) {
219: user->D = NULL;
220: return 0;
221: }
222: */
224: /* Speficy D */
225: /* (1) Specify D Size */
226: switch (dictChoice) {
227: case 0: /* 0:identity */
228: user->K = user->N;
229: break;
230: case 1: /* 1:gradient1D */
231: user->K = user->N-1;
232: break;
233: }
235: MatCreate(PETSC_COMM_SELF,&user->D);
236: MatSetSizes(user->D,PETSC_DECIDE,PETSC_DECIDE,user->K,user->N);
237: MatSetFromOptions(user->D);
238: MatSetUp(user->D);
240: /* (2) Specify D Content */
241: switch (dictChoice) {
242: case 0: /* 0:identity */
243: for (k=0; k<user->K; k++) {
244: v = 1.0;
245: MatSetValues(user->D,1,&k,1,&k,&v,INSERT_VALUES);
246: }
247: break;
248: case 1: /* 1:gradient1D. [-1, 1, 0,...; 0, -1, 1, 0, ...] */
249: for (k=0; k<user->K; k++) {
250: v = 1.0;
251: n = k+1;
252: MatSetValues(user->D,1,&k,1,&n,&v,INSERT_VALUES);
253: v = -1.0;
254: MatSetValues(user->D,1,&k,1,&k,&v,INSERT_VALUES);
255: }
256: break;
257: }
258: MatAssemblyBegin(user->D,MAT_FINAL_ASSEMBLY);
259: MatAssemblyEnd(user->D,MAT_FINAL_ASSEMBLY);
261: return 0;
262: }
264: /*TEST
266: build:
267: requires: !complex !single !__float128 !defined(PETSC_USE_64BIT_INDICES)
269: test:
270: localrunfiles: tomographyData_A_b_xGT
271: 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
273: test:
274: suffix: 2
275: localrunfiles: tomographyData_A_b_xGT
276: args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type l2prox -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6
278: test:
279: suffix: 3
280: localrunfiles: tomographyData_A_b_xGT
281: args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type user -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6
283: TEST*/