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