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
12
13 */
14
15 #include <petsctao.h>
16
17 /*
18 Description: BRGN tomography reconstruction example .
19 0.5*||Ax-b||^2 + lambda*g(x)
20 Reference: None
21 */
22
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
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;
35
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);
44
45 /*--------------------------------------------------------------------*/
main(int argc,char ** argv)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" */
56
57 PetscFunctionBeginUser;
58 PetscCall(PetscInitialize(&argc, &argv, NULL, help));
59
60 /* Create TAO solver and set desired solution method */
61 PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
62 PetscCall(TaoSetType(tao, TAOBRGN));
63
64 /* User set application context: A, D matrice, and b vector. */
65 PetscCall(InitializeUserData(&user));
66
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));
70
71 /* Set initial guess */
72 PetscCall(FormStartingPoint(x, &user));
73
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));
78
79 /* Bind user.D to tao->data->D */
80 PetscCall(TaoBRGNSetDictionaryMatrix(tao, user.D));
81
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));
86
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, (PetscErrorCodeFn *)EvaluateRegularizerHessianProd));
95 PetscCall(TaoBRGNSetRegularizerHessianRoutine(tao, Hreg, EvaluateRegularizerHessian, (void *)&user));
96
97 /* Check for any TAO command line arguments */
98 PetscCall(TaoSetFromOptions(tao));
99
100 PetscCall(TaoSetConvergenceHistory(tao, hist, resid, 0, lits, 100, PETSC_TRUE));
101
102 /* Perform the Solve */
103 PetscCall(TaoSolve(tao));
104
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));
109
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)));
115
116 /* Free TAO data structures */
117 PetscCall(TaoDestroy(&tao));
118
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 }
133
134 /*--------------------------------------------------------------------*/
135 /* Evaluate residual function A(x)-b in least square problem ||A(x)-b||^2 */
EvaluateResidual(Tao tao,Vec X,Vec F,void * ptr)136 PetscErrorCode EvaluateResidual(Tao tao, Vec X, Vec F, void *ptr)
137 {
138 AppCtx *user = (AppCtx *)ptr;
139
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 }
147
148 /*------------------------------------------------------------*/
EvaluateJacobian(Tao tao,Vec X,Mat J,Mat Jpre,void * ptr)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 }
155
156 /* ------------------------------------------------------------ */
EvaluateRegularizerObjectiveAndGradient(Tao tao,Vec X,PetscReal * f_reg,Vec G_reg,void * ptr)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 }
167
EvaluateRegularizerHessianProd(Mat Hreg,Vec in,Vec out)168 PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg, Vec in, Vec out)
169 {
170 PetscFunctionBegin;
171 PetscCall(VecCopy(in, out));
172 PetscFunctionReturn(PETSC_SUCCESS);
173 }
174
175 /* ------------------------------------------------------------ */
EvaluateRegularizerHessian(Tao tao,Vec X,Mat Hreg,void * ptr)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 }
182
183 /* ------------------------------------------------------------ */
FormStartingPoint(Vec X,AppCtx * user)184 PetscErrorCode FormStartingPoint(Vec X, AppCtx *user)
185 {
186 PetscFunctionBegin;
187 PetscCall(VecSet(X, 0.0));
188 PetscFunctionReturn(PETSC_SUCCESS);
189 }
190
191 /* ---------------------------------------------------------------------- */
InitializeUserData(AppCtx * user)192 PetscErrorCode InitializeUserData(AppCtx *user)
193 {
194 PetscInt k, n; /* indices for row and columns of D. */
195 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". */
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;
199 PetscBool flg;
200
201 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 PetscCall(PetscOptionsGetString(NULL, NULL, "-path", path, sizeof(path), &flg));
208 PetscCheck(flg, PETSC_COMM_WORLD, PETSC_ERR_USER, "Must specify -path ${DATAFILESPATH}/tao/tomography");
209 /* Load the A matrix, b vector, and xGT vector from a binary file. */
210 PetscCall(PetscSNPrintf(dataFile, sizeof(dataFile), "%s/tomographyData_A_b_xGT", path));
211 PetscCall(PetscViewerBinaryOpen(PETSC_COMM_WORLD, dataFile, FILE_MODE_READ, &fd));
212 PetscCall(MatCreate(PETSC_COMM_WORLD, &user->A));
213 PetscCall(MatSetType(user->A, MATSEQAIJ));
214 PetscCall(MatLoad(user->A, fd));
215 PetscCall(VecCreate(PETSC_COMM_WORLD, &user->b));
216 PetscCall(VecLoad(user->b, fd));
217 PetscCall(VecCreate(PETSC_COMM_WORLD, &user->xGT));
218 PetscCall(VecLoad(user->xGT, fd));
219 PetscCall(PetscViewerDestroy(&fd));
220 PetscCall(VecDuplicate(user->xGT, &user->xlb));
221 PetscCall(VecSet(user->xlb, 0.0));
222 PetscCall(VecDuplicate(user->xGT, &user->xub));
223 PetscCall(VecSet(user->xub, PETSC_INFINITY));
224
225 /* Specify the size */
226 PetscCall(MatGetSize(user->A, &user->M, &user->N));
227
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 */
234
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 }
245
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));
250
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 }
273
274 /*TEST
275
276 build:
277 requires: !complex !single !__float128 !defined(PETSC_USE_64BIT_INDICES)
278
279 testset:
280 requires: datafilespath
281 args: -path ${DATAFILESPATH}/tao/tomography
282
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
285
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
289
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
293
294 TEST*/
295