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