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