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 identity 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 PetscCall(PetscLogFlops(2.0 * user->M * user->N)); 145 PetscFunctionReturn(PETSC_SUCCESS); 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(PETSC_SUCCESS); 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(PETSC_SUCCESS); 166 } 167 168 PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg, Vec in, Vec out) 169 { 170 PetscFunctionBegin; 171 PetscCall(VecCopy(in, out)); 172 PetscFunctionReturn(PETSC_SUCCESS); 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(PETSC_SUCCESS); 181 } 182 183 /* ------------------------------------------------------------ */ 184 PetscErrorCode FormStartingPoint(Vec X, AppCtx *user) 185 { 186 PetscFunctionBegin; 187 PetscCall(VecSet(X, 0.0)); 188 PetscFunctionReturn(PETSC_SUCCESS); 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 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(PETSC_SUCCESS); 228 } 229 */ 230 231 /* Specify 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 PetscFunctionReturn(PETSC_SUCCESS); 268 } 269 270 /*TEST 271 272 build: 273 requires: !complex !single !__float128 !defined(PETSC_USE_64BIT_INDICES) 274 275 test: 276 localrunfiles: tomographyData_A_b_xGT 277 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 278 279 test: 280 suffix: 2 281 localrunfiles: tomographyData_A_b_xGT 282 args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type l2prox -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6 283 284 test: 285 suffix: 3 286 localrunfiles: tomographyData_A_b_xGT 287 args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type user -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6 288 289 TEST*/ 290