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 PetscFunctionBeginUser; 58 PetscCall(PetscInitialize(&argc,&argv,(char *)0,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 identiy 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,(void (*)(void))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 */ 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 PetscLogFlops(2.0*user->M*user->N); 145 PetscFunctionReturn(0); 146 } 147 148 /*------------------------------------------------------------*/ 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(0); 154 } 155 156 /* ------------------------------------------------------------ */ 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(0); 166 } 167 168 PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg,Vec in,Vec out) 169 { 170 PetscFunctionBegin; 171 PetscCall(VecCopy(in,out)); 172 PetscFunctionReturn(0); 173 } 174 175 /* ------------------------------------------------------------ */ 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(0); 181 } 182 183 /* ------------------------------------------------------------ */ 184 PetscErrorCode FormStartingPoint(Vec X,AppCtx *user) 185 { 186 PetscFunctionBegin; 187 PetscCall(VecSet(X,0.0)); 188 PetscFunctionReturn(0); 189 } 190 191 /* ---------------------------------------------------------------------- */ 192 PetscErrorCode InitializeUserData(AppCtx *user) 193 { 194 PetscInt k,n; /* indices for row and columns of D. */ 195 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". */ 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 200 PetscFunctionBegin; 201 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 /* Load the A matrix, b vector, and xGT vector from a binary file. */ 208 PetscCall(PetscViewerBinaryOpen(PETSC_COMM_WORLD,dataFile,FILE_MODE_READ,&fd)); 209 PetscCall(MatCreate(PETSC_COMM_WORLD,&user->A)); 210 PetscCall(MatSetType(user->A,MATSEQAIJ)); 211 PetscCall(MatLoad(user->A,fd)); 212 PetscCall(VecCreate(PETSC_COMM_WORLD,&user->b)); 213 PetscCall(VecLoad(user->b,fd)); 214 PetscCall(VecCreate(PETSC_COMM_WORLD,&user->xGT)); 215 PetscCall(VecLoad(user->xGT,fd)); 216 PetscCall(PetscViewerDestroy(&fd)); 217 PetscCall(VecDuplicate(user->xGT,&(user->xlb))); 218 PetscCall(VecSet(user->xlb,0.0)); 219 PetscCall(VecDuplicate(user->xGT,&(user->xub))); 220 PetscCall(VecSet(user->xub,PETSC_INFINITY)); 221 222 /* Specify the size */ 223 PetscCall(MatGetSize(user->A,&user->M,&user->N)); 224 225 /* 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. 226 if (dictChoice == 0) { 227 user->D = NULL; 228 PetscFunctionReturn(0); 229 } 230 */ 231 232 /* Speficy D */ 233 /* (1) Specify D Size */ 234 switch (dictChoice) { 235 case 0: /* 0:identity */ 236 user->K = user->N; 237 break; 238 case 1: /* 1:gradient1D */ 239 user->K = user->N-1; 240 break; 241 } 242 243 PetscCall(MatCreate(PETSC_COMM_SELF,&user->D)); 244 PetscCall(MatSetSizes(user->D,PETSC_DECIDE,PETSC_DECIDE,user->K,user->N)); 245 PetscCall(MatSetFromOptions(user->D)); 246 PetscCall(MatSetUp(user->D)); 247 248 /* (2) Specify D Content */ 249 switch (dictChoice) { 250 case 0: /* 0:identity */ 251 for (k=0; k<user->K; k++) { 252 v = 1.0; 253 PetscCall(MatSetValues(user->D,1,&k,1,&k,&v,INSERT_VALUES)); 254 } 255 break; 256 case 1: /* 1:gradient1D. [-1, 1, 0,...; 0, -1, 1, 0, ...] */ 257 for (k=0; k<user->K; k++) { 258 v = 1.0; 259 n = k+1; 260 PetscCall(MatSetValues(user->D,1,&k,1,&n,&v,INSERT_VALUES)); 261 v = -1.0; 262 PetscCall(MatSetValues(user->D,1,&k,1,&k,&v,INSERT_VALUES)); 263 } 264 break; 265 } 266 PetscCall(MatAssemblyBegin(user->D,MAT_FINAL_ASSEMBLY)); 267 PetscCall(MatAssemblyEnd(user->D,MAT_FINAL_ASSEMBLY)); 268 269 PetscFunctionReturn(0); 270 } 271 272 /*TEST 273 274 build: 275 requires: !complex !single !__float128 !defined(PETSC_USE_64BIT_INDICES) 276 277 test: 278 localrunfiles: tomographyData_A_b_xGT 279 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 280 281 test: 282 suffix: 2 283 localrunfiles: tomographyData_A_b_xGT 284 args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type l2prox -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6 285 286 test: 287 suffix: 3 288 localrunfiles: tomographyData_A_b_xGT 289 args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type user -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6 290 291 TEST*/ 292