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