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 Vec x, res; /* solution, function res(x) = A*x-b */ 48 Mat Hreg; /* regularizer Hessian matrix for user specified regularizer*/ 49 Tao tao; /* Tao solver context */ 50 PetscReal hist[100], resid[100], v1, v2; 51 PetscInt lits[100]; 52 AppCtx user; /* user-defined work context */ 53 PetscViewer fd; /* used to save result to file */ 54 char resultFile[] = "tomographyResult_x"; /* Debug: change from "tomographyResult_x" to "cs1Result_x" */ 55 56 PetscFunctionBeginUser; 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 AppCtx *user = (AppCtx *)ptr; 137 138 PetscFunctionBegin; 139 /* Compute Ax - b */ 140 PetscCall(MatMult(user->A, X, F)); 141 PetscCall(VecAXPY(F, -1, user->b)); 142 PetscLogFlops(2.0 * user->M * user->N); 143 PetscFunctionReturn(0); 144 } 145 146 /*------------------------------------------------------------*/ 147 PetscErrorCode EvaluateJacobian(Tao tao, Vec X, Mat J, Mat Jpre, void *ptr) { 148 /* 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 */ 149 PetscFunctionBegin; 150 PetscFunctionReturn(0); 151 } 152 153 /* ------------------------------------------------------------ */ 154 PetscErrorCode EvaluateRegularizerObjectiveAndGradient(Tao tao, Vec X, PetscReal *f_reg, Vec G_reg, void *ptr) { 155 PetscFunctionBegin; 156 /* compute regularizer objective = 0.5*x'*x */ 157 PetscCall(VecDot(X, X, f_reg)); 158 *f_reg *= 0.5; 159 /* compute regularizer gradient = x */ 160 PetscCall(VecCopy(X, G_reg)); 161 PetscFunctionReturn(0); 162 } 163 164 PetscErrorCode EvaluateRegularizerHessianProd(Mat Hreg, Vec in, Vec out) { 165 PetscFunctionBegin; 166 PetscCall(VecCopy(in, out)); 167 PetscFunctionReturn(0); 168 } 169 170 /* ------------------------------------------------------------ */ 171 PetscErrorCode EvaluateRegularizerHessian(Tao tao, Vec X, Mat Hreg, void *ptr) { 172 /* Hessian for regularizer objective = 0.5*x'*x is identity matrix, and is not changing*/ 173 PetscFunctionBegin; 174 PetscFunctionReturn(0); 175 } 176 177 /* ------------------------------------------------------------ */ 178 PetscErrorCode FormStartingPoint(Vec X, AppCtx *user) { 179 PetscFunctionBegin; 180 PetscCall(VecSet(X, 0.0)); 181 PetscFunctionReturn(0); 182 } 183 184 /* ---------------------------------------------------------------------- */ 185 PetscErrorCode InitializeUserData(AppCtx *user) { 186 PetscInt k, n; /* indices for row and columns of D. */ 187 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". */ 188 PetscInt dictChoice = 1; /* choose from 0:identity, 1:gradient1D, 2:gradient2D, 3:DCT etc */ 189 PetscViewer fd; /* used to load data from file */ 190 PetscReal v; 191 192 PetscFunctionBegin; 193 194 /* 195 Matrix Vector read and write refer to: 196 https://petsc.org/release/src/mat/tutorials/ex10.c 197 https://petsc.org/release/src/mat/tutorials/ex12.c 198 */ 199 /* Load the A matrix, b vector, and xGT vector from a binary file. */ 200 PetscCall(PetscViewerBinaryOpen(PETSC_COMM_WORLD, dataFile, FILE_MODE_READ, &fd)); 201 PetscCall(MatCreate(PETSC_COMM_WORLD, &user->A)); 202 PetscCall(MatSetType(user->A, MATSEQAIJ)); 203 PetscCall(MatLoad(user->A, fd)); 204 PetscCall(VecCreate(PETSC_COMM_WORLD, &user->b)); 205 PetscCall(VecLoad(user->b, fd)); 206 PetscCall(VecCreate(PETSC_COMM_WORLD, &user->xGT)); 207 PetscCall(VecLoad(user->xGT, fd)); 208 PetscCall(PetscViewerDestroy(&fd)); 209 PetscCall(VecDuplicate(user->xGT, &(user->xlb))); 210 PetscCall(VecSet(user->xlb, 0.0)); 211 PetscCall(VecDuplicate(user->xGT, &(user->xub))); 212 PetscCall(VecSet(user->xub, PETSC_INFINITY)); 213 214 /* Specify the size */ 215 PetscCall(MatGetSize(user->A, &user->M, &user->N)); 216 217 /* 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. 218 if (dictChoice == 0) { 219 user->D = NULL; 220 PetscFunctionReturn(0); 221 } 222 */ 223 224 /* Speficy D */ 225 /* (1) Specify D Size */ 226 switch (dictChoice) { 227 case 0: /* 0:identity */ user->K = user->N; break; 228 case 1: /* 1:gradient1D */ user->K = user->N - 1; break; 229 } 230 231 PetscCall(MatCreate(PETSC_COMM_SELF, &user->D)); 232 PetscCall(MatSetSizes(user->D, PETSC_DECIDE, PETSC_DECIDE, user->K, user->N)); 233 PetscCall(MatSetFromOptions(user->D)); 234 PetscCall(MatSetUp(user->D)); 235 236 /* (2) Specify D Content */ 237 switch (dictChoice) { 238 case 0: /* 0:identity */ 239 for (k = 0; k < user->K; k++) { 240 v = 1.0; 241 PetscCall(MatSetValues(user->D, 1, &k, 1, &k, &v, INSERT_VALUES)); 242 } 243 break; 244 case 1: /* 1:gradient1D. [-1, 1, 0,...; 0, -1, 1, 0, ...] */ 245 for (k = 0; k < user->K; k++) { 246 v = 1.0; 247 n = k + 1; 248 PetscCall(MatSetValues(user->D, 1, &k, 1, &n, &v, INSERT_VALUES)); 249 v = -1.0; 250 PetscCall(MatSetValues(user->D, 1, &k, 1, &k, &v, INSERT_VALUES)); 251 } 252 break; 253 } 254 PetscCall(MatAssemblyBegin(user->D, MAT_FINAL_ASSEMBLY)); 255 PetscCall(MatAssemblyEnd(user->D, MAT_FINAL_ASSEMBLY)); 256 257 PetscFunctionReturn(0); 258 } 259 260 /*TEST 261 262 build: 263 requires: !complex !single !__float128 !defined(PETSC_USE_64BIT_INDICES) 264 265 test: 266 localrunfiles: tomographyData_A_b_xGT 267 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 268 269 test: 270 suffix: 2 271 localrunfiles: tomographyData_A_b_xGT 272 args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type l2prox -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6 273 274 test: 275 suffix: 3 276 localrunfiles: tomographyData_A_b_xGT 277 args: -tao_monitor -tao_max_it 1000 -tao_brgn_regularization_type user -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6 278 279 TEST*/ 280