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
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 /*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*/
40
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;
48
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);
57
58 /*--------------------------------------------------------------------*/
main(int argc,char ** argv)59 int main(int argc,char **argv)
60 {
61 PetscErrorCode ierr; /* used to check for functions returning nonzeros */
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" */
70
71 ierr = PetscInitialize(&argc,&argv,(char *)0,help);if (ierr) return ierr;
72
73 /* Create TAO solver and set desired solution method */
74 ierr = TaoCreate(PETSC_COMM_SELF,&tao);CHKERRQ(ierr);
75 ierr = TaoSetType(tao,TAOBRGN);CHKERRQ(ierr);
76
77 /* User set application context: A, D matrice, and b vector. */
78 ierr = InitializeUserData(&user);CHKERRQ(ierr);
79
80 /* Allocate solution vector x, and function vectors Ax-b, */
81 ierr = VecCreateSeq(PETSC_COMM_SELF,user.N,&x);CHKERRQ(ierr);
82 ierr = VecCreateSeq(PETSC_COMM_SELF,user.M,&res);CHKERRQ(ierr);
83
84 /* Set initial guess */
85 ierr = FormStartingPoint(x,&user);CHKERRQ(ierr);
86
87 /* Bind x to tao->solution. */
88 ierr = TaoSetInitialVector(tao,x);CHKERRQ(ierr);
89 /* Sets the upper and lower bounds of x */
90 ierr = TaoSetVariableBounds(tao,user.xlb,user.xub);CHKERRQ(ierr);
91
92 /* Bind user.D to tao->data->D */
93 ierr = TaoBRGNSetDictionaryMatrix(tao,user.D);CHKERRQ(ierr);
94
95 /* Set the residual function and Jacobian routines for least squares. */
96 ierr = TaoSetResidualRoutine(tao,res,EvaluateResidual,(void*)&user);CHKERRQ(ierr);
97 /* Jacobian matrix fixed as user.A for Linear least sqaure problem. */
98 ierr = TaoSetJacobianResidualRoutine(tao,user.A,user.A,EvaluateJacobian,(void*)&user);CHKERRQ(ierr);
99
100 /* User set the regularizer objective, gradient, and hessian. Set it the same as using l2prox choice, for testing purpose. */
101 ierr = TaoBRGNSetRegularizerObjectiveAndGradientRoutine(tao,EvaluateRegularizerObjectiveAndGradient,(void*)&user);CHKERRQ(ierr);
102 /* User defined regularizer Hessian setup, here is identiy shell matrix */
103 ierr = MatCreate(PETSC_COMM_SELF,&Hreg);CHKERRQ(ierr);
104 ierr = MatSetSizes(Hreg,PETSC_DECIDE,PETSC_DECIDE,user.N,user.N);CHKERRQ(ierr);
105 ierr = MatSetType(Hreg,MATSHELL);CHKERRQ(ierr);
106 ierr = MatSetUp(Hreg);CHKERRQ(ierr);
107 ierr = MatShellSetOperation(Hreg,MATOP_MULT,(void (*)(void))EvaluateRegularizerHessianProd);CHKERRQ(ierr);
108 ierr = TaoBRGNSetRegularizerHessianRoutine(tao,Hreg,EvaluateRegularizerHessian,(void*)&user);CHKERRQ(ierr);
109
110 /* Check for any TAO command line arguments */
111 ierr = TaoSetFromOptions(tao);CHKERRQ(ierr);
112
113 ierr = TaoSetConvergenceHistory(tao,hist,resid,0,lits,100,PETSC_TRUE);CHKERRQ(ierr);
114
115 /* Perform the Solve */
116 ierr = TaoSolve(tao);CHKERRQ(ierr);
117
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 ierr = PetscViewerBinaryOpen(PETSC_COMM_SELF,resultFile,FILE_MODE_WRITE,&fd);CHKERRQ(ierr);
120 ierr = VecView(x,fd);CHKERRQ(ierr);
121 ierr = PetscViewerDestroy(&fd);CHKERRQ(ierr);
122
123 /* compute the error */
124 ierr = VecAXPY(x,-1,user.xGT);CHKERRQ(ierr);
125 ierr = VecNorm(x,NORM_2,&v1);CHKERRQ(ierr);
126 ierr = VecNorm(user.xGT,NORM_2,&v2);CHKERRQ(ierr);
127 ierr = PetscPrintf(PETSC_COMM_SELF, "relative reconstruction error: ||x-xGT||/||xGT|| = %6.4e.\n", (double)(v1/v2));CHKERRQ(ierr);
128
129 /* Free TAO data structures */
130 ierr = TaoDestroy(&tao);CHKERRQ(ierr);
131
132 /* Free PETSc data structures */
133 ierr = VecDestroy(&x);CHKERRQ(ierr);
134 ierr = VecDestroy(&res);CHKERRQ(ierr);
135 ierr = MatDestroy(&Hreg);CHKERRQ(ierr);
136 /* Free user data structures */
137 ierr = MatDestroy(&user.A);CHKERRQ(ierr);
138 ierr = MatDestroy(&user.D);CHKERRQ(ierr);
139 ierr = VecDestroy(&user.b);CHKERRQ(ierr);
140 ierr = VecDestroy(&user.xGT);CHKERRQ(ierr);
141 ierr = VecDestroy(&user.xlb);CHKERRQ(ierr);
142 ierr = VecDestroy(&user.xub);CHKERRQ(ierr);
143 ierr = PetscFinalize();
144 return ierr;
145 }
146
147 /*--------------------------------------------------------------------*/
148 /* Evaluate residual function A(x)-b in least square problem ||A(x)-b||^2 */
EvaluateResidual(Tao tao,Vec X,Vec F,void * ptr)149 PetscErrorCode EvaluateResidual(Tao tao,Vec X,Vec F,void *ptr)
150 {
151 AppCtx *user = (AppCtx *)ptr;
152 PetscErrorCode ierr;
153
154 PetscFunctionBegin;
155 /* Compute Ax - b */
156 ierr = MatMult(user->A,X,F);CHKERRQ(ierr);
157 ierr = VecAXPY(F,-1,user->b);CHKERRQ(ierr);
158 PetscLogFlops(2.0*user->M*user->N);
159 PetscFunctionReturn(0);
160 }
161
162 /*------------------------------------------------------------*/
EvaluateJacobian(Tao tao,Vec X,Mat J,Mat Jpre,void * ptr)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 */
166 PetscFunctionBegin;
167 PetscFunctionReturn(0);
168 }
169
170 /* ------------------------------------------------------------ */
EvaluateRegularizerObjectiveAndGradient(Tao tao,Vec X,PetscReal * f_reg,Vec G_reg,void * ptr)171 PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao tao,Vec X,PetscReal *f_reg,Vec G_reg,void *ptr)
172 {
173 PetscErrorCode ierr;
174
175 PetscFunctionBegin;
176 /* compute regularizer objective = 0.5*x'*x */
177 ierr = VecDot(X,X,f_reg);CHKERRQ(ierr);
178 *f_reg *= 0.5;
179 /* compute regularizer gradient = x */
180 ierr = VecCopy(X,G_reg);CHKERRQ(ierr);
181 PetscFunctionReturn(0);
182 }
183
EvaluateRegularizerHessianProd(Mat Hreg,Vec in,Vec out)184 PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg,Vec in,Vec out)
185 {
186 PetscErrorCode ierr;
187 PetscFunctionBegin;
188 ierr = VecCopy(in,out);CHKERRQ(ierr);
189 PetscFunctionReturn(0);
190 }
191
192 /* ------------------------------------------------------------ */
EvaluateRegularizerHessian(Tao tao,Vec X,Mat Hreg,void * ptr)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*/
196 PetscFunctionBegin;
197 PetscFunctionReturn(0);
198 }
199
200 /* ------------------------------------------------------------ */
FormStartingPoint(Vec X,AppCtx * user)201 PetscErrorCode FormStartingPoint(Vec X,AppCtx *user)
202 {
203 PetscErrorCode ierr;
204 PetscFunctionBegin;
205 ierr = VecSet(X,0.0);CHKERRQ(ierr);
206 PetscFunctionReturn(0);
207 }
208
209 /* ---------------------------------------------------------------------- */
InitializeUserData(AppCtx * user)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 */
216 PetscErrorCode ierr;
217 PetscReal v;
218
219 PetscFunctionBegin;
220
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 ierr = PetscViewerBinaryOpen(PETSC_COMM_WORLD,dataFile,FILE_MODE_READ,&fd);CHKERRQ(ierr);
228 ierr = MatCreate(PETSC_COMM_WORLD,&user->A);CHKERRQ(ierr);
229 ierr = MatSetType(user->A,MATSEQAIJ);CHKERRQ(ierr);
230 ierr = MatLoad(user->A,fd);CHKERRQ(ierr);
231 ierr = VecCreate(PETSC_COMM_WORLD,&user->b);CHKERRQ(ierr);
232 ierr = VecLoad(user->b,fd);CHKERRQ(ierr);
233 ierr = VecCreate(PETSC_COMM_WORLD,&user->xGT);CHKERRQ(ierr);
234 ierr = VecLoad(user->xGT,fd);CHKERRQ(ierr);
235 ierr = PetscViewerDestroy(&fd);CHKERRQ(ierr);
236 ierr = VecDuplicate(user->xGT,&(user->xlb));CHKERRQ(ierr);
237 ierr = VecSet(user->xlb,0.0);CHKERRQ(ierr);
238 ierr = VecDuplicate(user->xGT,&(user->xub));CHKERRQ(ierr);
239 ierr = VecSet(user->xub,PETSC_INFINITY);CHKERRQ(ierr);
240
241 /* Specify the size */
242 ierr = MatGetSize(user->A,&user->M,&user->N);CHKERRQ(ierr);
243
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 PetscFunctionReturn(0);
248 }
249 */
250
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 }
261
262 ierr = MatCreate(PETSC_COMM_SELF,&user->D);CHKERRQ(ierr);
263 ierr = MatSetSizes(user->D,PETSC_DECIDE,PETSC_DECIDE,user->K,user->N);CHKERRQ(ierr);
264 ierr = MatSetFromOptions(user->D);CHKERRQ(ierr);
265 ierr = MatSetUp(user->D);CHKERRQ(ierr);
266
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 ierr = MatSetValues(user->D,1,&k,1,&k,&v,INSERT_VALUES);CHKERRQ(ierr);
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 ierr = MatSetValues(user->D,1,&k,1,&n,&v,INSERT_VALUES);CHKERRQ(ierr);
280 v = -1.0;
281 ierr = MatSetValues(user->D,1,&k,1,&k,&v,INSERT_VALUES);CHKERRQ(ierr);
282 }
283 break;
284 }
285 ierr = MatAssemblyBegin(user->D,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
286 ierr = MatAssemblyEnd(user->D,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
287
288 PetscFunctionReturn(0);
289 }
290
291 /*TEST
292
293 build:
294 requires: !complex !single !__float128 !define(PETSC_USE_64BIT_INDICES)
295
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
299
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
304
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
309
310 TEST*/
311