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: PetscFunctionBeginUser;
58: PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
60: /* Create TAO solver and set desired solution method */
61: PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
62: PetscCall(TaoSetType(tao, TAOBRGN));
64: /* User set application context: A, D matrice, and b vector. */
65: PetscCall(InitializeUserData(&user));
67: /* Allocate solution vector x, and function vectors Ax-b, */
68: PetscCall(VecCreateSeq(PETSC_COMM_SELF, user.N, &x));
69: PetscCall(VecCreateSeq(PETSC_COMM_SELF, user.M, &res));
71: /* Set initial guess */
72: PetscCall(FormStartingPoint(x, &user));
74: /* Bind x to tao->solution. */
75: PetscCall(TaoSetSolution(tao, x));
76: /* Sets the upper and lower bounds of x */
77: PetscCall(TaoSetVariableBounds(tao, user.xlb, user.xub));
79: /* Bind user.D to tao->data->D */
80: PetscCall(TaoBRGNSetDictionaryMatrix(tao, user.D));
82: /* Set the residual function and Jacobian routines for least squares. */
83: PetscCall(TaoSetResidualRoutine(tao, res, EvaluateResidual, (void *)&user));
84: /* Jacobian matrix fixed as user.A for Linear least square problem. */
85: PetscCall(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: PetscCall(TaoBRGNSetRegularizerObjectiveAndGradientRoutine(tao, EvaluateRegularizerObjectiveAndGradient, (void *)&user));
89: /* User defined regularizer Hessian setup, here is identity shell matrix */
90: PetscCall(MatCreate(PETSC_COMM_SELF, &Hreg));
91: PetscCall(MatSetSizes(Hreg, PETSC_DECIDE, PETSC_DECIDE, user.N, user.N));
92: PetscCall(MatSetType(Hreg, MATSHELL));
93: PetscCall(MatSetUp(Hreg));
94: PetscCall(MatShellSetOperation(Hreg, MATOP_MULT, (void (*)(void))EvaluateRegularizerHessianProd));
95: PetscCall(TaoBRGNSetRegularizerHessianRoutine(tao, Hreg, EvaluateRegularizerHessian, (void *)&user));
97: /* Check for any TAO command line arguments */
98: PetscCall(TaoSetFromOptions(tao));
100: PetscCall(TaoSetConvergenceHistory(tao, hist, resid, 0, lits, 100, PETSC_TRUE));
102: /* Perform the Solve */
103: PetscCall(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: PetscCall(PetscViewerBinaryOpen(PETSC_COMM_SELF, resultFile, FILE_MODE_WRITE, &fd));
107: PetscCall(VecView(x, fd));
108: PetscCall(PetscViewerDestroy(&fd));
110: /* compute the error */
111: PetscCall(VecAXPY(x, -1, user.xGT));
112: PetscCall(VecNorm(x, NORM_2, &v1));
113: PetscCall(VecNorm(user.xGT, NORM_2, &v2));
114: PetscCall(PetscPrintf(PETSC_COMM_SELF, "relative reconstruction error: ||x-xGT||/||xGT|| = %6.4e.\n", (double)(v1 / v2)));
116: /* Free TAO data structures */
117: PetscCall(TaoDestroy(&tao));
119: /* Free PETSc data structures */
120: PetscCall(VecDestroy(&x));
121: PetscCall(VecDestroy(&res));
122: PetscCall(MatDestroy(&Hreg));
123: /* Free user data structures */
124: PetscCall(MatDestroy(&user.A));
125: PetscCall(MatDestroy(&user.D));
126: PetscCall(VecDestroy(&user.b));
127: PetscCall(VecDestroy(&user.xGT));
128: PetscCall(VecDestroy(&user.xlb));
129: PetscCall(VecDestroy(&user.xub));
130: PetscCall(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: PetscFunctionBegin;
141: /* Compute Ax - b */
142: PetscCall(MatMult(user->A, X, F));
143: PetscCall(VecAXPY(F, -1, user->b));
144: PetscCall(PetscLogFlops(2.0 * user->M * user->N));
145: PetscFunctionReturn(PETSC_SUCCESS);
146: }
148: /*------------------------------------------------------------*/
149: PetscErrorCode EvaluateJacobian(Tao tao, Vec X, Mat J, Mat Jpre, void *ptr)
150: {
151: /* 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 */
152: PetscFunctionBegin;
153: PetscFunctionReturn(PETSC_SUCCESS);
154: }
156: /* ------------------------------------------------------------ */
157: PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao tao, Vec X, PetscReal *f_reg, Vec G_reg, void *ptr)
158: {
159: PetscFunctionBegin;
160: /* compute regularizer objective = 0.5*x'*x */
161: PetscCall(VecDot(X, X, f_reg));
162: *f_reg *= 0.5;
163: /* compute regularizer gradient = x */
164: PetscCall(VecCopy(X, G_reg));
165: PetscFunctionReturn(PETSC_SUCCESS);
166: }
168: PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg, Vec in, Vec out)
169: {
170: PetscFunctionBegin;
171: PetscCall(VecCopy(in, out));
172: PetscFunctionReturn(PETSC_SUCCESS);
173: }
175: /* ------------------------------------------------------------ */
176: PetscErrorCode EvaluateRegularizerHessian(Tao tao, Vec X, Mat Hreg, void *ptr)
177: {
178: /* Hessian for regularizer objective = 0.5*x'*x is identity matrix, and is not changing*/
179: PetscFunctionBegin;
180: PetscFunctionReturn(PETSC_SUCCESS);
181: }
183: /* ------------------------------------------------------------ */
184: PetscErrorCode FormStartingPoint(Vec X, AppCtx *user)
185: {
186: PetscFunctionBegin;
187: PetscCall(VecSet(X, 0.0));
188: PetscFunctionReturn(PETSC_SUCCESS);
189: }
191: /* ---------------------------------------------------------------------- */
192: PetscErrorCode InitializeUserData(AppCtx *user)
193: {
194: PetscInt k, n; /* indices for row and columns of D. */
195: 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". */
196: PetscInt dictChoice = 1; /* choose from 0:identity, 1:gradient1D, 2:gradient2D, 3:DCT etc */
197: PetscViewer fd; /* used to load data from file */
198: PetscReal v;
200: PetscFunctionBegin;
202: /*
203: Matrix Vector read and write refer to:
204: https://petsc.org/release/src/mat/tutorials/ex10.c
205: https://petsc.org/release/src/mat/tutorials/ex12.c
206: */
207: /* Load the A matrix, b vector, and xGT vector from a binary file. */
208: PetscCall(PetscViewerBinaryOpen(PETSC_COMM_WORLD, dataFile, FILE_MODE_READ, &fd));
209: PetscCall(MatCreate(PETSC_COMM_WORLD, &user->A));
210: PetscCall(MatSetType(user->A, MATSEQAIJ));
211: PetscCall(MatLoad(user->A, fd));
212: PetscCall(VecCreate(PETSC_COMM_WORLD, &user->b));
213: PetscCall(VecLoad(user->b, fd));
214: PetscCall(VecCreate(PETSC_COMM_WORLD, &user->xGT));
215: PetscCall(VecLoad(user->xGT, fd));
216: PetscCall(PetscViewerDestroy(&fd));
217: PetscCall(VecDuplicate(user->xGT, &(user->xlb)));
218: PetscCall(VecSet(user->xlb, 0.0));
219: PetscCall(VecDuplicate(user->xGT, &(user->xub)));
220: PetscCall(VecSet(user->xub, PETSC_INFINITY));
222: /* Specify the size */
223: PetscCall(MatGetSize(user->A, &user->M, &user->N));
225: /* 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.
226: if (dictChoice == 0) {
227: user->D = NULL;
228: PetscFunctionReturn(PETSC_SUCCESS);
229: }
230: */
232: /* Specify D */
233: /* (1) Specify D Size */
234: switch (dictChoice) {
235: case 0: /* 0:identity */
236: user->K = user->N;
237: break;
238: case 1: /* 1:gradient1D */
239: user->K = user->N - 1;
240: break;
241: }
243: PetscCall(MatCreate(PETSC_COMM_SELF, &user->D));
244: PetscCall(MatSetSizes(user->D, PETSC_DECIDE, PETSC_DECIDE, user->K, user->N));
245: PetscCall(MatSetFromOptions(user->D));
246: PetscCall(MatSetUp(user->D));
248: /* (2) Specify D Content */
249: switch (dictChoice) {
250: case 0: /* 0:identity */
251: for (k = 0; k < user->K; k++) {
252: v = 1.0;
253: PetscCall(MatSetValues(user->D, 1, &k, 1, &k, &v, INSERT_VALUES));
254: }
255: break;
256: case 1: /* 1:gradient1D. [-1, 1, 0,...; 0, -1, 1, 0, ...] */
257: for (k = 0; k < user->K; k++) {
258: v = 1.0;
259: n = k + 1;
260: PetscCall(MatSetValues(user->D, 1, &k, 1, &n, &v, INSERT_VALUES));
261: v = -1.0;
262: PetscCall(MatSetValues(user->D, 1, &k, 1, &k, &v, INSERT_VALUES));
263: }
264: break;
265: }
266: PetscCall(MatAssemblyBegin(user->D, MAT_FINAL_ASSEMBLY));
267: PetscCall(MatAssemblyEnd(user->D, MAT_FINAL_ASSEMBLY));
269: PetscFunctionReturn(PETSC_SUCCESS);
270: }
272: /*TEST
274: build:
275: requires: !complex !single !__float128 !defined(PETSC_USE_64BIT_INDICES)
277: test:
278: localrunfiles: tomographyData_A_b_xGT
279: 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
281: test:
282: suffix: 2
283: localrunfiles: tomographyData_A_b_xGT
284: args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type l2prox -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6
286: test:
287: suffix: 3
288: localrunfiles: tomographyData_A_b_xGT
289: args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type user -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6
291: TEST*/