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: petscsys.h - system routines
9: petscis.h - index sets
10: petscvec.h - vectors
11: petscmat.h - matrices
12: petscpc.h - preconditioners
13: petscksp.h - Krylov subspace methods
14: */
16: #include <petsctao.h>
18: /*
19: Description: BRGN tomography reconstruction example .
20: 0.5*||Ax-b||^2 + lambda*g(x)
21: Reference: None
22: */
24: static char help[] = "Finds the least-squares solution to the under constraint linear model Ax = b, with regularizer. \n\
25: A is a M*N real matrix (M<N), x is sparse. A good regularizer is an L1 regularizer. \n\
26: 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\
27: 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";
29: /* User-defined application context */
30: typedef struct {
31: /* Working space. linear least square: res(x) = A*x - b */
32: PetscInt M, N, K; /* Problem dimension: A is M*N Matrix, D is K*N Matrix */
33: 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 */
34: Vec b, xGT, xlb, xub; /* observation b, ground truth xGT, the lower bound and upper bound of x*/
35: } AppCtx;
37: /* User provided Routines */
38: PetscErrorCode InitializeUserData(AppCtx *);
39: PetscErrorCode FormStartingPoint(Vec, AppCtx *);
40: PetscErrorCode EvaluateResidual(Tao, Vec, Vec, void *);
41: PetscErrorCode EvaluateJacobian(Tao, Vec, Mat, Mat, void *);
42: PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao, Vec, PetscReal *, Vec, void *);
43: PetscErrorCode EvaluateRegularizerHessian(Tao, Vec, Mat, void *);
44: PetscErrorCode EvaluateRegularizerHessianProd(Mat, Vec, Vec);
46: /*--------------------------------------------------------------------*/
47: int main(int argc, char **argv)
48: {
49: Vec x, res; /* solution, function res(x) = A*x-b */
50: Mat Hreg; /* regularizer Hessian matrix for user specified regularizer*/
51: Tao tao; /* Tao solver context */
52: PetscReal hist[100], resid[100], v1, v2;
53: PetscInt lits[100];
54: AppCtx user; /* user-defined work context */
55: PetscViewer fd; /* used to save result to file */
56: char resultFile[] = "tomographyResult_x"; /* Debug: change from "tomographyResult_x" to "cs1Result_x" */
58: PetscFunctionBeginUser;
59: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
61: /* Create TAO solver and set desired solution method */
62: PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
63: PetscCall(TaoSetType(tao, TAOBRGN));
65: /* User set application context: A, D matrice, and b vector. */
66: PetscCall(InitializeUserData(&user));
68: /* Allocate solution vector x, and function vectors Ax-b, */
69: PetscCall(VecCreateSeq(PETSC_COMM_SELF, user.N, &x));
70: PetscCall(VecCreateSeq(PETSC_COMM_SELF, user.M, &res));
72: /* Set initial guess */
73: PetscCall(FormStartingPoint(x, &user));
75: /* Bind x to tao->solution. */
76: PetscCall(TaoSetSolution(tao, x));
77: /* Sets the upper and lower bounds of x */
78: PetscCall(TaoSetVariableBounds(tao, user.xlb, user.xub));
80: /* Bind user.D to tao->data->D */
81: PetscCall(TaoBRGNSetDictionaryMatrix(tao, user.D));
83: /* Set the residual function and Jacobian routines for least squares. */
84: PetscCall(TaoSetResidualRoutine(tao, res, EvaluateResidual, (void *)&user));
85: /* Jacobian matrix fixed as user.A for Linear least square problem. */
86: PetscCall(TaoSetJacobianResidualRoutine(tao, user.A, user.A, EvaluateJacobian, (void *)&user));
88: /* User set the regularizer objective, gradient, and hessian. Set it the same as using l2prox choice, for testing purpose. */
89: PetscCall(TaoBRGNSetRegularizerObjectiveAndGradientRoutine(tao, EvaluateRegularizerObjectiveAndGradient, (void *)&user));
90: /* User defined regularizer Hessian setup, here is identity shell matrix */
91: PetscCall(MatCreate(PETSC_COMM_SELF, &Hreg));
92: PetscCall(MatSetSizes(Hreg, PETSC_DECIDE, PETSC_DECIDE, user.N, user.N));
93: PetscCall(MatSetType(Hreg, MATSHELL));
94: PetscCall(MatSetUp(Hreg));
95: PetscCall(MatShellSetOperation(Hreg, MATOP_MULT, (PetscErrorCodeFn *)EvaluateRegularizerHessianProd));
96: PetscCall(TaoBRGNSetRegularizerHessianRoutine(tao, Hreg, EvaluateRegularizerHessian, (void *)&user));
98: /* Check for any TAO command line arguments */
99: PetscCall(TaoSetFromOptions(tao));
101: PetscCall(TaoSetConvergenceHistory(tao, hist, resid, NULL, lits, 100, PETSC_TRUE));
103: /* Perform the Solve */
104: PetscCall(TaoSolve(tao));
106: /* Save x (reconstruction of object) vector to a binary file, which maybe read from MATLAB and convert to a 2D image for comparison. */
107: PetscCall(PetscViewerBinaryOpen(PETSC_COMM_SELF, resultFile, FILE_MODE_WRITE, &fd));
108: PetscCall(VecView(x, fd));
109: PetscCall(PetscViewerDestroy(&fd));
111: /* compute the error */
112: PetscCall(VecAXPY(x, -1, user.xGT));
113: PetscCall(VecNorm(x, NORM_2, &v1));
114: PetscCall(VecNorm(user.xGT, NORM_2, &v2));
115: PetscCall(PetscPrintf(PETSC_COMM_SELF, "relative reconstruction error: ||x-xGT||/||xGT|| = %6.4e.\n", (double)(v1 / v2)));
117: /* Free TAO data structures */
118: PetscCall(TaoDestroy(&tao));
120: /* Free PETSc data structures */
121: PetscCall(VecDestroy(&x));
122: PetscCall(VecDestroy(&res));
123: PetscCall(MatDestroy(&Hreg));
124: /* Free user data structures */
125: PetscCall(MatDestroy(&user.A));
126: PetscCall(MatDestroy(&user.D));
127: PetscCall(VecDestroy(&user.b));
128: PetscCall(VecDestroy(&user.xGT));
129: PetscCall(VecDestroy(&user.xlb));
130: PetscCall(VecDestroy(&user.xub));
131: PetscCall(PetscFinalize());
132: return 0;
133: }
135: /*--------------------------------------------------------------------*/
136: /* Evaluate residual function A(x)-b in least square problem ||A(x)-b||^2 */
137: PetscErrorCode EvaluateResidual(Tao tao, Vec X, Vec F, void *ptr)
138: {
139: AppCtx *user = (AppCtx *)ptr;
141: PetscFunctionBegin;
142: /* Compute Ax - b */
143: PetscCall(MatMult(user->A, X, F));
144: PetscCall(VecAXPY(F, -1, user->b));
145: PetscCall(PetscLogFlops(2.0 * user->M * user->N));
146: PetscFunctionReturn(PETSC_SUCCESS);
147: }
149: /*------------------------------------------------------------*/
150: PetscErrorCode EvaluateJacobian(Tao tao, Vec X, Mat J, Mat Jpre, void *ptr)
151: {
152: /* 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 */
153: PetscFunctionBegin;
154: PetscFunctionReturn(PETSC_SUCCESS);
155: }
157: /* ------------------------------------------------------------ */
158: PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao tao, Vec X, PetscReal *f_reg, Vec G_reg, void *ptr)
159: {
160: PetscFunctionBegin;
161: /* compute regularizer objective = 0.5*x'*x */
162: PetscCall(VecDot(X, X, f_reg));
163: *f_reg *= 0.5;
164: /* compute regularizer gradient = x */
165: PetscCall(VecCopy(X, G_reg));
166: PetscFunctionReturn(PETSC_SUCCESS);
167: }
169: PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg, Vec in, Vec out)
170: {
171: PetscFunctionBegin;
172: PetscCall(VecCopy(in, out));
173: PetscFunctionReturn(PETSC_SUCCESS);
174: }
176: /* ------------------------------------------------------------ */
177: PetscErrorCode EvaluateRegularizerHessian(Tao tao, Vec X, Mat Hreg, void *ptr)
178: {
179: /* Hessian for regularizer objective = 0.5*x'*x is identity matrix, and is not changing*/
180: PetscFunctionBegin;
181: PetscFunctionReturn(PETSC_SUCCESS);
182: }
184: /* ------------------------------------------------------------ */
185: PetscErrorCode FormStartingPoint(Vec X, AppCtx *user)
186: {
187: PetscFunctionBegin;
188: PetscCall(VecSet(X, 0.0));
189: PetscFunctionReturn(PETSC_SUCCESS);
190: }
192: /* ---------------------------------------------------------------------- */
193: PetscErrorCode InitializeUserData(AppCtx *user)
194: {
195: PetscInt k, n; /* indices for row and columns of D. */
196: char dataFile[PETSC_MAX_PATH_LEN], path[PETSC_MAX_PATH_LEN]; /* 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". */
197: PetscInt dictChoice = 1; /* choose from 0:identity, 1:gradient1D, 2:gradient2D, 3:DCT etc */
198: PetscViewer fd; /* used to load data from file */
199: PetscReal v;
200: PetscBool flg;
202: PetscFunctionBegin;
203: /*
204: Matrix Vector read and write refer to:
205: https://petsc.org/release/src/mat/tutorials/ex10.c
206: https://petsc.org/release/src/mat/tutorials/ex12.c
207: */
208: PetscCall(PetscOptionsGetString(NULL, NULL, "-path", path, sizeof(path), &flg));
209: PetscCheck(flg, PETSC_COMM_WORLD, PETSC_ERR_USER, "Must specify -path ${DATAFILESPATH}/tao/tomography");
210: /* Load the A matrix, b vector, and xGT vector from a binary file. */
211: PetscCall(PetscSNPrintf(dataFile, sizeof(dataFile), "%s/tomographyData_A_b_xGT", path));
212: PetscCall(PetscViewerBinaryOpen(PETSC_COMM_WORLD, dataFile, FILE_MODE_READ, &fd));
213: PetscCall(MatCreate(PETSC_COMM_WORLD, &user->A));
214: PetscCall(MatSetType(user->A, MATSEQAIJ));
215: PetscCall(MatLoad(user->A, fd));
216: PetscCall(VecCreate(PETSC_COMM_WORLD, &user->b));
217: PetscCall(VecLoad(user->b, fd));
218: PetscCall(VecCreate(PETSC_COMM_WORLD, &user->xGT));
219: PetscCall(VecLoad(user->xGT, fd));
220: PetscCall(PetscViewerDestroy(&fd));
221: PetscCall(VecDuplicate(user->xGT, &user->xlb));
222: PetscCall(VecDuplicate(user->xGT, &user->xub));
223: PetscCall(VecSet(user->xub, PETSC_INFINITY));
225: /* Specify the size */
226: PetscCall(MatGetSize(user->A, &user->M, &user->N));
228: /* 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.
229: if (dictChoice == 0) {
230: user->D = NULL;
231: PetscFunctionReturn(PETSC_SUCCESS);
232: }
233: */
235: /* Specify D */
236: /* (1) Specify D Size */
237: switch (dictChoice) {
238: case 0: /* 0:identity */
239: user->K = user->N;
240: break;
241: case 1: /* 1:gradient1D */
242: user->K = user->N - 1;
243: break;
244: }
246: PetscCall(MatCreate(PETSC_COMM_SELF, &user->D));
247: PetscCall(MatSetSizes(user->D, PETSC_DECIDE, PETSC_DECIDE, user->K, user->N));
248: PetscCall(MatSetFromOptions(user->D));
249: PetscCall(MatSetUp(user->D));
251: /* (2) Specify D Content */
252: switch (dictChoice) {
253: case 0: /* 0:identity */
254: for (k = 0; k < user->K; k++) {
255: v = 1.0;
256: PetscCall(MatSetValues(user->D, 1, &k, 1, &k, &v, INSERT_VALUES));
257: }
258: break;
259: case 1: /* 1:gradient1D. [-1, 1, 0,...; 0, -1, 1, 0, ...] */
260: for (k = 0; k < user->K; k++) {
261: v = 1.0;
262: n = k + 1;
263: PetscCall(MatSetValues(user->D, 1, &k, 1, &n, &v, INSERT_VALUES));
264: v = -1.0;
265: PetscCall(MatSetValues(user->D, 1, &k, 1, &k, &v, INSERT_VALUES));
266: }
267: break;
268: }
269: PetscCall(MatAssemblyBegin(user->D, MAT_FINAL_ASSEMBLY));
270: PetscCall(MatAssemblyEnd(user->D, MAT_FINAL_ASSEMBLY));
271: PetscFunctionReturn(PETSC_SUCCESS);
272: }
274: /*TEST
276: build:
277: requires: !complex !single !__float128 !defined(PETSC_USE_64BIT_INDICES)
279: testset:
280: requires: datafilespath
281: args: -path ${DATAFILESPATH}/tao/tomography
283: test:
284: args: -tao_max_it 10 -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: args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type l2prox -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6 -tao_brgn_subsolver_tao_monitor
290: test:
291: suffix: 3
292: args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type user -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6 -tao_brgn_subsolver_tao_monitor
294: TEST*/