Actual source code: cs1.c

  1: /* XH: todo add cs1f.F90 and asjust makefile */
  2: /*
  3:    Include "petsctao.h" so that we can use TAO solvers.  Note that this
  4:    file automatically includes libraries such as:
  5:      petsc.h       - base PETSc routines   petscvec.h - vectors
  6:      petscsys.h    - system routines        petscmat.h - matrices
  7:      petscis.h     - index sets            petscksp.h - Krylov subspace methods
  8:      petscviewer.h - viewers               petscpc.h  - preconditioners

 10: */

 12: #include <petsctao.h>

 14: /*
 15: Description:   Compressive sensing test example 1.
 16:                0.5*||Ax-b||^2 + lambda*||D*x||_1
 17:                Xiang Huang: Nov 19, 2018

 19: Reference:     None
 20: */

 22: static char help[] = "Finds the least-squares solution to the under constraint linear model Ax = b, with L1-norm regularizer. \n\
 23:             A is a M*N real matrix (M<N), x is sparse. \n\
 24:             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\
 25:             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: #define M 3
 28: #define N 5
 29: #define K 4

 31: typedef enum {
 32:   TEST_L1DICT,
 33:   TEST_LM,
 34:   TEST_NONE
 35: } TestType;

 37: /* User-defined application context */
 38: typedef struct {
 39:   /* Working space. linear least square:  f(x) = A*x - b */
 40:   PetscReal A[M][N]; /* array of coefficients */
 41:   PetscReal b[M];    /* array of observations */
 42:   PetscReal xGT[M];  /* array of ground truth object, which can be used to compare the reconstruction result */
 43:   PetscReal D[K][N]; /* array of coefficients for 0.5*||Ax-b||^2 + lambda*||D*x||_1 */
 44:   PetscReal J[M][N]; /* dense jacobian matrix array. For linear least square, J = A. For nonlinear least square, it is different from A */
 45:   PetscInt  idm[M];  /* Matrix row, column indices for jacobian and dictionary */
 46:   PetscInt  idn[N];
 47:   PetscInt  idk[K];
 48:   TestType  tType;
 49:   PetscBool view_sol;
 50: } AppCtx;

 52: /* User provided Routines */
 53: PetscErrorCode InitializeUserData(AppCtx *);
 54: PetscErrorCode FormStartingPoint(Vec);
 55: PetscErrorCode FormDictionaryMatrix(Mat, AppCtx *);
 56: PetscErrorCode EvaluateFunction(Tao, Vec, Vec, void *);
 57: PetscErrorCode EvaluateJacobian(Tao, Vec, Mat, Mat, void *);

 59: static PetscErrorCode SetTaoOptionsFromUserOptions(Tao tao, AppCtx *ctx)
 60: {
 61:   PetscBool isbrgn;

 63:   PetscFunctionBeginUser;
 64:   PetscCall(PetscObjectTypeCompare((PetscObject)tao, TAOBRGN, &isbrgn));
 65:   if (isbrgn) {
 66:     switch (ctx->tType) {
 67:     case TEST_LM:
 68:       PetscCall(TaoBRGNSetRegularizationType(tao, TAOBRGN_REGULARIZATION_LM));
 69:       break;
 70:     case TEST_L1DICT:
 71:       PetscCall(TaoBRGNSetRegularizationType(tao, TAOBRGN_REGULARIZATION_L1DICT));
 72:       PetscCall(TaoBRGNSetRegularizerWeight(tao, 0.0001));
 73:       PetscCall(TaoBRGNSetL1SmoothEpsilon(tao, 1.e-6));
 74:       break;
 75:     case TEST_NONE:
 76:     default:
 77:       break;
 78:     }
 79:   }
 80:   PetscFunctionReturn(PETSC_SUCCESS);
 81: }

 83: static PetscErrorCode TestOutType(Tao tao, AppCtx *ctx)
 84: {
 85:   PetscBool isbrgn;

 87:   PetscFunctionBeginUser;
 88:   PetscCall(PetscObjectTypeCompare((PetscObject)tao, TAOBRGN, &isbrgn));
 89:   if (isbrgn) {
 90:     TaoBRGNRegularizationType type;

 92:     PetscCall(TaoBRGNGetRegularizationType(tao, &type));
 93:     switch (ctx->tType) {
 94:     case TEST_LM:
 95:       PetscCheck(type == TAOBRGN_REGULARIZATION_LM, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_NOTSAMETYPE, "BRGN Regularization type is not LM!");
 96:       break;
 97:     case TEST_L1DICT:
 98:       PetscCheck(type == TAOBRGN_REGULARIZATION_L1DICT, PetscObjectComm((PetscObject)tao), PETSC_ERR_ARG_NOTSAMETYPE, "BRGN Regularization type is not L1DICT!");
 99:       break;
100:     case TEST_NONE:
101:     default:
102:       break;
103:     }
104:   }
105:   PetscFunctionReturn(PETSC_SUCCESS);
106: }

108: static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *ctx)
109: {
110:   const char *testTypes[3] = {"l1dict", "lm", "none"};
111:   PetscInt    run;

113:   PetscFunctionBeginUser;
114:   ctx->tType    = TEST_NONE;
115:   ctx->view_sol = PETSC_TRUE;
116:   PetscOptionsBegin(comm, "", "Least squares coverage", "");
117:   PetscCall(PetscOptionsBool("-view_sol", "Penalize deviation from both goals", "cs1.c", ctx->view_sol, &ctx->view_sol, NULL));
118:   run = ctx->tType;
119:   PetscCall(PetscOptionsEList("-test_type", "The coverage test type", "cs1.c", testTypes, 3, testTypes[ctx->tType], &run, NULL));
120:   ctx->tType = (TestType)run;
121:   PetscOptionsEnd();
122:   PetscFunctionReturn(PETSC_SUCCESS);
123: }

125: /*--------------------------------------------------------------------*/
126: int main(int argc, char **argv)
127: {
128:   Vec       x, f; /* solution, function f(x) = A*x-b */
129:   Mat       J, D; /* Jacobian matrix, Transform matrix */
130:   Tao       tao;  /* Tao solver context */
131:   PetscInt  i;    /* iteration information */
132:   PetscReal hist[100], resid[100];
133:   PetscInt  lits[100];
134:   AppCtx    user; /* user-defined work context */

136:   PetscFunctionBeginUser;
137:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));
138:   PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user));
139:   /* Allocate solution and vector function vectors */
140:   PetscCall(VecCreateSeq(PETSC_COMM_SELF, N, &x));
141:   PetscCall(VecCreateSeq(PETSC_COMM_SELF, M, &f));

143:   /* Allocate Jacobian and Dictionary matrix. */
144:   PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, M, N, NULL, &J));
145:   PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, K, N, NULL, &D)); /* XH: TODO: dense -> sparse/dense/shell etc, do it on fly  */

147:   for (i = 0; i < M; i++) user.idm[i] = i;
148:   for (i = 0; i < N; i++) user.idn[i] = i;
149:   for (i = 0; i < K; i++) user.idk[i] = i;

151:   /* Create TAO solver and set desired solution method */
152:   PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
153:   PetscCall(TaoSetType(tao, TAOBRGN));

155:   /* User set application context: A, D matrice, and b vector. */
156:   PetscCall(InitializeUserData(&user));

158:   /* Set initial guess */
159:   PetscCall(FormStartingPoint(x));

161:   /* Fill the content of matrix D from user application Context */
162:   PetscCall(FormDictionaryMatrix(D, &user));

164:   /* If needed, set options via function for testing purpose */
165:   PetscCall(SetTaoOptionsFromUserOptions(tao, &user));
166:   /* Bind x to tao->solution. */
167:   PetscCall(TaoSetSolution(tao, x));
168:   /* Bind D to tao->data->D */
169:   PetscCall(TaoBRGNSetDictionaryMatrix(tao, D));

171:   /* Set the function and Jacobian routines. */
172:   PetscCall(TaoSetResidualRoutine(tao, f, EvaluateFunction, (void *)&user));
173:   PetscCall(TaoSetJacobianResidualRoutine(tao, J, J, EvaluateJacobian, (void *)&user));

175:   /* Check for any TAO command line arguments */
176:   PetscCall(TaoSetFromOptions(tao));

178:   PetscCall(TaoSetConvergenceHistory(tao, hist, resid, 0, lits, 100, PETSC_TRUE));

180:   /* Perform the Solve */
181:   PetscCall(TaoSolve(tao));

183:   /* XH: Debug: View the result, function and Jacobian.  */
184:   if (user.view_sol) {
185:     PetscCall(PetscPrintf(PETSC_COMM_SELF, "-------- result x, residual f=A*x-b, and Jacobian=A. -------- \n"));
186:     PetscCall(VecView(x, PETSC_VIEWER_STDOUT_SELF));
187:     PetscCall(VecView(f, PETSC_VIEWER_STDOUT_SELF));
188:     PetscCall(MatView(J, PETSC_VIEWER_STDOUT_SELF));
189:     PetscCall(MatView(D, PETSC_VIEWER_STDOUT_SELF));
190:   }
191:   PetscCall(TestOutType(tao, &user));

193:   /* Free TAO data structures */
194:   PetscCall(TaoDestroy(&tao));

196:   /* Free PETSc data structures */
197:   PetscCall(VecDestroy(&x));
198:   PetscCall(VecDestroy(&f));
199:   PetscCall(MatDestroy(&J));
200:   PetscCall(MatDestroy(&D));

202:   PetscCall(PetscFinalize());
203:   return 0;
204: }

206: /*--------------------------------------------------------------------*/
207: PetscErrorCode EvaluateFunction(Tao tao, Vec X, Vec F, void *ptr)
208: {
209:   AppCtx          *user = (AppCtx *)ptr;
210:   PetscInt         m, n;
211:   const PetscReal *x;
212:   PetscReal       *b = user->b, *f;

214:   PetscFunctionBegin;
215:   PetscCall(VecGetArrayRead(X, &x));
216:   PetscCall(VecGetArray(F, &f));

218:   /* Even for linear least square, we do not direct use matrix operation f = A*x - b now, just for future modification and compatibility for nonlinear least square */
219:   for (m = 0; m < M; m++) {
220:     f[m] = -b[m];
221:     for (n = 0; n < N; n++) f[m] += user->A[m][n] * x[n];
222:   }
223:   PetscCall(VecRestoreArrayRead(X, &x));
224:   PetscCall(VecRestoreArray(F, &f));
225:   PetscCall(PetscLogFlops(2.0 * M * N));
226:   PetscFunctionReturn(PETSC_SUCCESS);
227: }

229: /*------------------------------------------------------------*/
230: /* J[m][n] = df[m]/dx[n] */
231: PetscErrorCode EvaluateJacobian(Tao tao, Vec X, Mat J, Mat Jpre, void *ptr)
232: {
233:   AppCtx          *user = (AppCtx *)ptr;
234:   PetscInt         m, n;
235:   const PetscReal *x;

237:   PetscFunctionBegin;
238:   PetscCall(VecGetArrayRead(X, &x)); /* not used for linear least square, but keep for future nonlinear least square) */
239:   /* XH: TODO:  For linear least square, we can just set J=A fixed once, instead of keep update it! Maybe just create a function getFixedJacobian?
240:     For nonlinear least square, we require x to compute J, keep codes here for future nonlinear least square*/
241:   for (m = 0; m < M; ++m) {
242:     for (n = 0; n < N; ++n) user->J[m][n] = user->A[m][n];
243:   }

245:   PetscCall(MatSetValues(J, M, user->idm, N, user->idn, (PetscReal *)user->J, INSERT_VALUES));
246:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
247:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));

249:   PetscCall(VecRestoreArrayRead(X, &x)); /* not used for linear least square, but keep for future nonlinear least square) */
250:   PetscCall(PetscLogFlops(0));           /* 0 for linear least square, >0 for nonlinear least square */
251:   PetscFunctionReturn(PETSC_SUCCESS);
252: }

254: /* ------------------------------------------------------------ */
255: /* Currently fixed matrix, in future may be dynamic for D(x)? */
256: PetscErrorCode FormDictionaryMatrix(Mat D, AppCtx *user)
257: {
258:   PetscFunctionBegin;
259:   PetscCall(MatSetValues(D, K, user->idk, N, user->idn, (PetscReal *)user->D, INSERT_VALUES));
260:   PetscCall(MatAssemblyBegin(D, MAT_FINAL_ASSEMBLY));
261:   PetscCall(MatAssemblyEnd(D, MAT_FINAL_ASSEMBLY));

263:   PetscCall(PetscLogFlops(0)); /* 0 for fixed dictionary matrix, >0 for varying dictionary matrix */
264:   PetscFunctionReturn(PETSC_SUCCESS);
265: }

267: /* ------------------------------------------------------------ */
268: PetscErrorCode FormStartingPoint(Vec X)
269: {
270:   PetscFunctionBegin;
271:   PetscCall(VecSet(X, 0.0));
272:   PetscFunctionReturn(PETSC_SUCCESS);
273: }

275: /* ---------------------------------------------------------------------- */
276: PetscErrorCode InitializeUserData(AppCtx *user)
277: {
278:   PetscReal *b = user->b; /* **A=user->A, but we don't know the dimension of A in this way, how to fix? */
279:   PetscInt   m, n, k;     /* loop index for M,N,K dimension. */

281:   PetscFunctionBegin;
282:   /* b = A*x while x = [0;0;1;0;0] here*/
283:   m      = 0;
284:   b[m++] = 0.28;
285:   b[m++] = 0.55;
286:   b[m++] = 0.96;

288:   /* MATLAB generated random matrix, uniformly distributed in [0,1] with 2 digits accuracy. rng(0); A = rand(M, N); A = round(A*100)/100;
289:   A = [0.81  0.91  0.28  0.96  0.96
290:        0.91  0.63  0.55  0.16  0.49
291:        0.13  0.10  0.96  0.97  0.80]
292:   */
293:   m               = 0;
294:   n               = 0;
295:   user->A[m][n++] = 0.81;
296:   user->A[m][n++] = 0.91;
297:   user->A[m][n++] = 0.28;
298:   user->A[m][n++] = 0.96;
299:   user->A[m][n++] = 0.96;
300:   ++m;
301:   n               = 0;
302:   user->A[m][n++] = 0.91;
303:   user->A[m][n++] = 0.63;
304:   user->A[m][n++] = 0.55;
305:   user->A[m][n++] = 0.16;
306:   user->A[m][n++] = 0.49;
307:   ++m;
308:   n               = 0;
309:   user->A[m][n++] = 0.13;
310:   user->A[m][n++] = 0.10;
311:   user->A[m][n++] = 0.96;
312:   user->A[m][n++] = 0.97;
313:   user->A[m][n++] = 0.80;

315:   /* initialize to 0 */
316:   for (k = 0; k < K; k++) {
317:     for (n = 0; n < N; n++) user->D[k][n] = 0.0;
318:   }
319:   /* Choice I: set D to identity matrix of size N*N for testing */
320:   /* for (k=0; k<K; k++) user->D[k][k] = 1.0; */
321:   /* Choice II: set D to Backward difference matrix of size (N-1)*N, with zero extended boundary assumption */
322:   for (k = 0; k < K; k++) {
323:     user->D[k][k]     = -1.0;
324:     user->D[k][k + 1] = 1.0;
325:   }
326:   PetscFunctionReturn(PETSC_SUCCESS);
327: }

329: /*TEST

331:    build:
332:       requires: !complex !single !quad !defined(PETSC_USE_64BIT_INDICES) !__float128

334:    test:
335:       localrunfiles: cs1Data_A_b_xGT
336:       args: -tao_monitor_short -tao_max_it 100 -tao_type pounders -tao_gatol 1.e-6

338:    test:
339:       suffix: 2
340:       localrunfiles: cs1Data_A_b_xGT
341:       args: -tao_monitor -tao_max_it 100 -tao_type brgn -tao_brgn_regularization_type l2prox -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6 -tao_brgn_subsolver_tao_bnk_ksp_converged_reason

343:    test:
344:       suffix: 3
345:       localrunfiles: cs1Data_A_b_xGT
346:       args: -tao_monitor -tao_max_it 100 -tao_type brgn -tao_brgn_regularization_type l1dict -tao_brgn_regularizer_weight 1e-8 -tao_brgn_l1_smooth_epsilon 1e-6 -tao_gatol 1.e-6

348:    test:
349:       suffix: 4
350:       localrunfiles: cs1Data_A_b_xGT
351:       args: -tao_monitor -tao_max_it 100 -tao_type brgn -tao_brgn_regularization_type l2pure -tao_brgn_regularizer_weight 1e-8 -tao_gatol 1.e-6

353:    test:
354:       suffix: 5
355:       localrunfiles: cs1Data_A_b_xGT
356:       args: -tao_monitor -tao_max_it 100 -tao_type brgn -tao_brgn_regularization_type lm -tao_gatol 1.e-6 -tao_brgn_subsolver_tao_type bnls

358:    test:
359:       suffix: view_lm
360:       localrunfiles: cs1Data_A_b_xGT
361:       args: -tao_type brgn -test_type lm -tao_gatol 1.e-6 -view_sol 0 -tao_view

363:    test:
364:       suffix: view_l1dict
365:       localrunfiles: cs1Data_A_b_xGT
366:       args: -tao_type brgn -test_type l1dict -tao_gatol 1.e-6 -view_sol 0 -tao_view

368: TEST*/