Actual source code: tomography.c
petsc-3.14.6 2021-03-30
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 - sysem 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: TaoSetInitialVector();
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: {
62: Vec x,res; /* solution, function res(x) = A*x-b */
63: Mat Hreg; /* regularizer Hessian matrix for user specified regularizer*/
64: Tao tao; /* Tao solver context */
65: PetscReal hist[100],resid[100],v1,v2;
66: PetscInt lits[100];
67: AppCtx user; /* user-defined work context */
68: PetscViewer fd; /* used to save result to file */
69: char resultFile[] = "tomographyResult_x"; /* Debug: change from "tomographyResult_x" to "cs1Result_x" */
71: PetscInitialize(&argc,&argv,(char *)0,help);if (ierr) return ierr;
73: /* Create TAO solver and set desired solution method */
74: TaoCreate(PETSC_COMM_SELF,&tao);
75: TaoSetType(tao,TAOBRGN);
77: /* User set application context: A, D matrice, and b vector. */
78: InitializeUserData(&user);
80: /* Allocate solution vector x, and function vectors Ax-b, */
81: VecCreateSeq(PETSC_COMM_SELF,user.N,&x);
82: VecCreateSeq(PETSC_COMM_SELF,user.M,&res);
84: /* Set initial guess */
85: FormStartingPoint(x,&user);
87: /* Bind x to tao->solution. */
88: TaoSetInitialVector(tao,x);
89: /* Sets the upper and lower bounds of x */
90: TaoSetVariableBounds(tao,user.xlb,user.xub);
92: /* Bind user.D to tao->data->D */
93: TaoBRGNSetDictionaryMatrix(tao,user.D);
95: /* Set the residual function and Jacobian routines for least squares. */
96: TaoSetResidualRoutine(tao,res,EvaluateResidual,(void*)&user);
97: /* Jacobian matrix fixed as user.A for Linear least sqaure problem. */
98: TaoSetJacobianResidualRoutine(tao,user.A,user.A,EvaluateJacobian,(void*)&user);
100: /* User set the regularizer objective, gradient, and hessian. Set it the same as using l2prox choice, for testing purpose. */
101: TaoBRGNSetRegularizerObjectiveAndGradientRoutine(tao,EvaluateRegularizerObjectiveAndGradient,(void*)&user);
102: /* User defined regularizer Hessian setup, here is identiy shell matrix */
103: MatCreate(PETSC_COMM_SELF,&Hreg);
104: MatSetSizes(Hreg,PETSC_DECIDE,PETSC_DECIDE,user.N,user.N);
105: MatSetType(Hreg,MATSHELL);
106: MatSetUp(Hreg);
107: MatShellSetOperation(Hreg,MATOP_MULT,(void (*)(void))EvaluateRegularizerHessianProd);
108: TaoBRGNSetRegularizerHessianRoutine(tao,Hreg,EvaluateRegularizerHessian,(void*)&user);
110: /* Check for any TAO command line arguments */
111: TaoSetFromOptions(tao);
113: TaoSetConvergenceHistory(tao,hist,resid,0,lits,100,PETSC_TRUE);
115: /* Perform the Solve */
116: TaoSolve(tao);
118: /* Save x (reconstruction of object) vector to a binary file, which maybe read from Matlab and convert to a 2D image for comparison. */
119: PetscViewerBinaryOpen(PETSC_COMM_SELF,resultFile,FILE_MODE_WRITE,&fd);
120: VecView(x,fd);
121: PetscViewerDestroy(&fd);
123: /* compute the error */
124: VecAXPY(x,-1,user.xGT);
125: VecNorm(x,NORM_2,&v1);
126: VecNorm(user.xGT,NORM_2,&v2);
127: PetscPrintf(PETSC_COMM_SELF, "relative reconstruction error: ||x-xGT||/||xGT|| = %6.4e.\n", (double)(v1/v2));
129: /* Free TAO data structures */
130: TaoDestroy(&tao);
132: /* Free PETSc data structures */
133: VecDestroy(&x);
134: VecDestroy(&res);
135: MatDestroy(&Hreg);
136: /* Free user data structures */
137: MatDestroy(&user.A);
138: MatDestroy(&user.D);
139: VecDestroy(&user.b);
140: VecDestroy(&user.xGT);
141: VecDestroy(&user.xlb);
142: VecDestroy(&user.xub);
143: PetscFinalize();
144: return ierr;
145: }
147: /*--------------------------------------------------------------------*/
148: /* Evaluate residual function A(x)-b in least square problem ||A(x)-b||^2 */
149: PetscErrorCode EvaluateResidual(Tao tao,Vec X,Vec F,void *ptr)
150: {
151: AppCtx *user = (AppCtx *)ptr;
155: /* Compute Ax - b */
156: MatMult(user->A,X,F);
157: VecAXPY(F,-1,user->b);
158: PetscLogFlops(2.0*user->M*user->N);
159: return(0);
160: }
162: /*------------------------------------------------------------*/
163: PetscErrorCode EvaluateJacobian(Tao tao,Vec X,Mat J,Mat Jpre,void *ptr)
164: {
165: /* 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 */
167: return(0);
168: }
170: /* ------------------------------------------------------------ */
171: PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao tao,Vec X,PetscReal *f_reg,Vec G_reg,void *ptr)
172: {
176: /* compute regularizer objective = 0.5*x'*x */
177: VecDot(X,X,f_reg);
178: *f_reg *= 0.5;
179: /* compute regularizer gradient = x */
180: VecCopy(X,G_reg);
181: return(0);
182: }
184: PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg,Vec in,Vec out)
185: {
188: VecCopy(in,out);
189: return(0);
190: }
192: /* ------------------------------------------------------------ */
193: PetscErrorCode EvaluateRegularizerHessian(Tao tao,Vec X,Mat Hreg,void *ptr)
194: {
195: /* Hessian for regularizer objective = 0.5*x'*x is identity matrix, and is not changing*/
197: return(0);
198: }
200: /* ------------------------------------------------------------ */
201: PetscErrorCode FormStartingPoint(Vec X,AppCtx *user)
202: {
205: VecSet(X,0.0);
206: return(0);
207: }
209: /* ---------------------------------------------------------------------- */
210: PetscErrorCode InitializeUserData(AppCtx *user)
211: {
212: PetscInt k,n; /* indices for row and columns of D. */
213: 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". */
214: PetscInt dictChoice = 1; /* choose from 0:identity, 1:gradient1D, 2:gradient2D, 3:DCT etc */
215: PetscViewer fd; /* used to load data from file */
217: PetscReal v;
221: /*
222: Matrix Vector read and write refer to:
223: https://www.mcs.anl.gov/petsc/petsc-current/src/mat/tutorials/ex10.c
224: https://www.mcs.anl.gov/petsc/petsc-current/src/mat/tutorials/ex12.c
225: */
226: /* Load the A matrix, b vector, and xGT vector from a binary file. */
227: PetscViewerBinaryOpen(PETSC_COMM_WORLD,dataFile,FILE_MODE_READ,&fd);
228: MatCreate(PETSC_COMM_WORLD,&user->A);
229: MatSetType(user->A,MATSEQAIJ);
230: MatLoad(user->A,fd);
231: VecCreate(PETSC_COMM_WORLD,&user->b);
232: VecLoad(user->b,fd);
233: VecCreate(PETSC_COMM_WORLD,&user->xGT);
234: VecLoad(user->xGT,fd);
235: PetscViewerDestroy(&fd);
236: VecDuplicate(user->xGT,&(user->xlb));
237: VecSet(user->xlb,0.0);
238: VecDuplicate(user->xGT,&(user->xub));
239: VecSet(user->xub,PETSC_INFINITY);
241: /* Specify the size */
242: MatGetSize(user->A,&user->M,&user->N);
244: /* 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.
245: if (dictChoice == 0) {
246: user->D = NULL;
247: return(0);
248: }
249: */
251: /* Speficy D */
252: /* (1) Specify D Size */
253: switch (dictChoice) {
254: case 0: /* 0:identity */
255: user->K = user->N;
256: break;
257: case 1: /* 1:gradient1D */
258: user->K = user->N-1;
259: break;
260: }
262: MatCreate(PETSC_COMM_SELF,&user->D);
263: MatSetSizes(user->D,PETSC_DECIDE,PETSC_DECIDE,user->K,user->N);
264: MatSetFromOptions(user->D);
265: MatSetUp(user->D);
267: /* (2) Specify D Content */
268: switch (dictChoice) {
269: case 0: /* 0:identity */
270: for (k=0; k<user->K; k++) {
271: v = 1.0;
272: MatSetValues(user->D,1,&k,1,&k,&v,INSERT_VALUES);
273: }
274: break;
275: case 1: /* 1:gradient1D. [-1, 1, 0,...; 0, -1, 1, 0, ...] */
276: for (k=0; k<user->K; k++) {
277: v = 1.0;
278: n = k+1;
279: MatSetValues(user->D,1,&k,1,&n,&v,INSERT_VALUES);
280: v = -1.0;
281: MatSetValues(user->D,1,&k,1,&k,&v,INSERT_VALUES);
282: }
283: break;
284: }
285: MatAssemblyBegin(user->D,MAT_FINAL_ASSEMBLY);
286: MatAssemblyEnd(user->D,MAT_FINAL_ASSEMBLY);
288: return(0);
289: }
291: /*TEST
293: build:
294: requires: !complex !single !__float128 !define(PETSC_USE_64BIT_INDICES)
296: test:
297: localrunfiles: tomographyData_A_b_xGT
298: 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
300: test:
301: suffix: 2
302: localrunfiles: tomographyData_A_b_xGT
303: args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type l2prox -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6
305: test:
306: suffix: 3
307: localrunfiles: tomographyData_A_b_xGT
308: args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type user -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6
310: TEST*/