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: /* User-defined application context */
     32: typedef struct {
     33:   /* Working space. linear least square:  f(x) = A*x - b */
     34:   PetscReal A[M][N]; /* array of coefficients */
     35:   PetscReal b[M];    /* array of observations */
     36:   PetscReal xGT[M];  /* array of ground truth object, which can be used to compare the reconstruction result */
     37:   PetscReal D[K][N]; /* array of coefficients for 0.5*||Ax-b||^2 + lambda*||D*x||_1 */
     38:   PetscReal J[M][N]; /* dense jacobian matrix array. For linear least square, J = A. For nonlinear least square, it is different from A */
     39:   PetscInt  idm[M];  /* Matrix row, column indices for jacobian and dictionary */
     40:   PetscInt  idn[N];
     41:   PetscInt  idk[K];
     42: } AppCtx;

     44: /* User provided Routines */
     45: PetscErrorCode InitializeUserData(AppCtx *);
     46: PetscErrorCode FormStartingPoint(Vec);
     47: PetscErrorCode FormDictionaryMatrix(Mat, AppCtx *);
     48: PetscErrorCode EvaluateFunction(Tao, Vec, Vec, void *);
     49: PetscErrorCode EvaluateJacobian(Tao, Vec, Mat, Mat, void *);

     51: /*--------------------------------------------------------------------*/
     52: int main(int argc, char **argv)
     53: {
     54:   Vec       x, f; /* solution, function f(x) = A*x-b */
     55:   Mat       J, D; /* Jacobian matrix, Transform matrix */
     56:   Tao       tao;  /* Tao solver context */
     57:   PetscInt  i;    /* iteration information */
     58:   PetscReal hist[100], resid[100];
     59:   PetscInt  lits[100];
     60:   AppCtx    user; /* user-defined work context */

     62:   PetscFunctionBeginUser;
     63:   PetscCall(PetscInitialize(&argc, &argv, NULL, help));

     65:   /* Allocate solution and vector function vectors */
     66:   PetscCall(VecCreateSeq(PETSC_COMM_SELF, N, &x));
     67:   PetscCall(VecCreateSeq(PETSC_COMM_SELF, M, &f));

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

     73:   for (i = 0; i < M; i++) user.idm[i] = i;
     74:   for (i = 0; i < N; i++) user.idn[i] = i;
     75:   for (i = 0; i < K; i++) user.idk[i] = i;

     77:   /* Create TAO solver and set desired solution method */
     78:   PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
     79:   PetscCall(TaoSetType(tao, TAOBRGN));

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

     84:   /* Set initial guess */
     85:   PetscCall(FormStartingPoint(x));

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

     90:   /* Bind x to tao->solution. */
     91:   PetscCall(TaoSetSolution(tao, x));
     92:   /* Bind D to tao->data->D */
     93:   PetscCall(TaoBRGNSetDictionaryMatrix(tao, D));

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

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

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

    104:   /* Perform the Solve */
    105:   PetscCall(TaoSolve(tao));

    107:   /* XH: Debug: View the result, function and Jacobian.  */
    108:   PetscCall(PetscPrintf(PETSC_COMM_SELF, "-------- result x, residual f=A*x-b, and Jacobian=A. -------- \n"));
    109:   PetscCall(VecView(x, PETSC_VIEWER_STDOUT_SELF));
    110:   PetscCall(VecView(f, PETSC_VIEWER_STDOUT_SELF));
    111:   PetscCall(MatView(J, PETSC_VIEWER_STDOUT_SELF));
    112:   PetscCall(MatView(D, PETSC_VIEWER_STDOUT_SELF));

    114:   /* Free TAO data structures */
    115:   PetscCall(TaoDestroy(&tao));

    117:   /* Free PETSc data structures */
    118:   PetscCall(VecDestroy(&x));
    119:   PetscCall(VecDestroy(&f));
    120:   PetscCall(MatDestroy(&J));
    121:   PetscCall(MatDestroy(&D));

    123:   PetscCall(PetscFinalize());
    124:   return 0;
    125: }

    127: /*--------------------------------------------------------------------*/
    128: PetscErrorCode EvaluateFunction(Tao tao, Vec X, Vec F, void *ptr)
    129: {
    130:   AppCtx          *user = (AppCtx *)ptr;
    131:   PetscInt         m, n;
    132:   const PetscReal *x;
    133:   PetscReal       *b = user->b, *f;

    135:   PetscFunctionBegin;
    136:   PetscCall(VecGetArrayRead(X, &x));
    137:   PetscCall(VecGetArray(F, &f));

    139:   /* 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 */
    140:   for (m = 0; m < M; m++) {
    141:     f[m] = -b[m];
    142:     for (n = 0; n < N; n++) f[m] += user->A[m][n] * x[n];
    143:   }
    144:   PetscCall(VecRestoreArrayRead(X, &x));
    145:   PetscCall(VecRestoreArray(F, &f));
    146:   PetscCall(PetscLogFlops(2.0 * M * N));
    147:   PetscFunctionReturn(PETSC_SUCCESS);
    148: }

    150: /*------------------------------------------------------------*/
    151: /* J[m][n] = df[m]/dx[n] */
    152: PetscErrorCode EvaluateJacobian(Tao tao, Vec X, Mat J, Mat Jpre, void *ptr)
    153: {
    154:   AppCtx          *user = (AppCtx *)ptr;
    155:   PetscInt         m, n;
    156:   const PetscReal *x;

    158:   PetscFunctionBegin;
    159:   PetscCall(VecGetArrayRead(X, &x)); /* not used for linear least square, but keep for future nonlinear least square) */
    160:   /* 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?
    161:     For nonlinear least square, we require x to compute J, keep codes here for future nonlinear least square*/
    162:   for (m = 0; m < M; ++m) {
    163:     for (n = 0; n < N; ++n) user->J[m][n] = user->A[m][n];
    164:   }

    166:   PetscCall(MatSetValues(J, M, user->idm, N, user->idn, (PetscReal *)user->J, INSERT_VALUES));
    167:   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
    168:   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));

    170:   PetscCall(VecRestoreArrayRead(X, &x)); /* not used for linear least square, but keep for future nonlinear least square) */
    171:   PetscCall(PetscLogFlops(0));           /* 0 for linear least square, >0 for nonlinear least square */
    172:   PetscFunctionReturn(PETSC_SUCCESS);
    173: }

    175: /* ------------------------------------------------------------ */
    176: /* Currently fixed matrix, in future may be dynamic for D(x)? */
    177: PetscErrorCode FormDictionaryMatrix(Mat D, AppCtx *user)
    178: {
    179:   PetscFunctionBegin;
    180:   PetscCall(MatSetValues(D, K, user->idk, N, user->idn, (PetscReal *)user->D, INSERT_VALUES));
    181:   PetscCall(MatAssemblyBegin(D, MAT_FINAL_ASSEMBLY));
    182:   PetscCall(MatAssemblyEnd(D, MAT_FINAL_ASSEMBLY));

    184:   PetscCall(PetscLogFlops(0)); /* 0 for fixed dictionary matrix, >0 for varying dictionary matrix */
    185:   PetscFunctionReturn(PETSC_SUCCESS);
    186: }

    188: /* ------------------------------------------------------------ */
    189: PetscErrorCode FormStartingPoint(Vec X)
    190: {
    191:   PetscFunctionBegin;
    192:   PetscCall(VecSet(X, 0.0));
    193:   PetscFunctionReturn(PETSC_SUCCESS);
    194: }

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

    202:   PetscFunctionBegin;
    203:   /* b = A*x while x = [0;0;1;0;0] here*/
    204:   m      = 0;
    205:   b[m++] = 0.28;
    206:   b[m++] = 0.55;
    207:   b[m++] = 0.96;

    209:   /* MATLAB generated random matrix, uniformly distributed in [0,1] with 2 digits accuracy. rng(0); A = rand(M, N); A = round(A*100)/100;
    210:   A = [0.81  0.91  0.28  0.96  0.96
    211:        0.91  0.63  0.55  0.16  0.49
    212:        0.13  0.10  0.96  0.97  0.80]
    213:   */
    214:   m               = 0;
    215:   n               = 0;
    216:   user->A[m][n++] = 0.81;
    217:   user->A[m][n++] = 0.91;
    218:   user->A[m][n++] = 0.28;
    219:   user->A[m][n++] = 0.96;
    220:   user->A[m][n++] = 0.96;
    221:   ++m;
    222:   n               = 0;
    223:   user->A[m][n++] = 0.91;
    224:   user->A[m][n++] = 0.63;
    225:   user->A[m][n++] = 0.55;
    226:   user->A[m][n++] = 0.16;
    227:   user->A[m][n++] = 0.49;
    228:   ++m;
    229:   n               = 0;
    230:   user->A[m][n++] = 0.13;
    231:   user->A[m][n++] = 0.10;
    232:   user->A[m][n++] = 0.96;
    233:   user->A[m][n++] = 0.97;
    234:   user->A[m][n++] = 0.80;

    236:   /* initialize to 0 */
    237:   for (k = 0; k < K; k++) {
    238:     for (n = 0; n < N; n++) user->D[k][n] = 0.0;
    239:   }
    240:   /* Choice I: set D to identity matrix of size N*N for testing */
    241:   /* for (k=0; k<K; k++) user->D[k][k] = 1.0; */
    242:   /* Choice II: set D to Backward difference matrix of size (N-1)*N, with zero extended boundary assumption */
    243:   for (k = 0; k < K; k++) {
    244:     user->D[k][k]     = -1.0;
    245:     user->D[k][k + 1] = 1.0;
    246:   }
    247:   PetscFunctionReturn(PETSC_SUCCESS);
    248: }

    250: /*TEST

    252:    build:
    253:       requires: !complex !single !quad !defined(PETSC_USE_64BIT_INDICES)

    255:    test:
    256:       localrunfiles: cs1Data_A_b_xGT
    257:       args: -tao_monitor_short -tao_max_it 100 -tao_type pounders -tao_gatol 1.e-6

    259:    test:
    260:       suffix: 2
    261:       localrunfiles: cs1Data_A_b_xGT
    262:       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

    264:    test:
    265:       suffix: 3
    266:       localrunfiles: cs1Data_A_b_xGT
    267:       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

    269:    test:
    270:       suffix: 4
    271:       localrunfiles: cs1Data_A_b_xGT
    272:       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

    274:    test:
    275:       suffix: 5
    276:       localrunfiles: cs1Data_A_b_xGT
    277:       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

    279: TEST*/