xref: /petsc/src/tao/leastsquares/tutorials/cs1.c (revision d756bedd70a89ca052be956bccd75c5761cb2ab4)
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
9 
10 */
11 
12 #include <petsctao.h>
13 
14 /*
15 Description:   Compressive sensing test example 1.
16                0.5*||Ax-b||^2 + lambda*||D*x||_1
17                Xiang Huang: Nov 19, 2018
18 
19 Reference:     None
20 */
21 
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";
26 
27 #define M 3
28 #define N 5
29 #define K 4
30 
31 typedef enum {
32   TEST_L1DICT,
33   TEST_LM,
34   TEST_NONE
35 } TestType;
36 
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;
51 
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 *);
58 
SetTaoOptionsFromUserOptions(Tao tao,AppCtx * ctx)59 static PetscErrorCode SetTaoOptionsFromUserOptions(Tao tao, AppCtx *ctx)
60 {
61   PetscBool isbrgn;
62 
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 }
82 
TestOutType(Tao tao,AppCtx * ctx)83 static PetscErrorCode TestOutType(Tao tao, AppCtx *ctx)
84 {
85   PetscBool isbrgn;
86 
87   PetscFunctionBeginUser;
88   PetscCall(PetscObjectTypeCompare((PetscObject)tao, TAOBRGN, &isbrgn));
89   if (isbrgn) {
90     TaoBRGNRegularizationType type;
91 
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 }
107 
ProcessOptions(MPI_Comm comm,AppCtx * ctx)108 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *ctx)
109 {
110   const char *testTypes[3] = {"l1dict", "lm", "none"};
111   PetscInt    run;
112 
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 }
124 
125 /*--------------------------------------------------------------------*/
main(int argc,char ** argv)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 */
135 
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));
142 
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  */
146 
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;
150 
151   /* Create TAO solver and set desired solution method */
152   PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
153   PetscCall(TaoSetType(tao, TAOBRGN));
154 
155   /* User set application context: A, D matrice, and b vector. */
156   PetscCall(InitializeUserData(&user));
157 
158   /* Set initial guess */
159   PetscCall(FormStartingPoint(x));
160 
161   /* Fill the content of matrix D from user application Context */
162   PetscCall(FormDictionaryMatrix(D, &user));
163 
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));
170 
171   /* Set the function and Jacobian routines. */
172   PetscCall(TaoSetResidualRoutine(tao, f, EvaluateFunction, (void *)&user));
173   PetscCall(TaoSetJacobianResidualRoutine(tao, J, J, EvaluateJacobian, (void *)&user));
174 
175   /* Check for any TAO command line arguments */
176   PetscCall(TaoSetFromOptions(tao));
177 
178   PetscCall(TaoSetConvergenceHistory(tao, hist, resid, 0, lits, 100, PETSC_TRUE));
179 
180   /* Perform the Solve */
181   PetscCall(TaoSolve(tao));
182 
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));
192 
193   /* Free TAO data structures */
194   PetscCall(TaoDestroy(&tao));
195 
196   /* Free PETSc data structures */
197   PetscCall(VecDestroy(&x));
198   PetscCall(VecDestroy(&f));
199   PetscCall(MatDestroy(&J));
200   PetscCall(MatDestroy(&D));
201 
202   PetscCall(PetscFinalize());
203   return 0;
204 }
205 
206 /*--------------------------------------------------------------------*/
EvaluateFunction(Tao tao,Vec X,Vec F,void * ptr)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;
213 
214   PetscFunctionBegin;
215   PetscCall(VecGetArrayRead(X, &x));
216   PetscCall(VecGetArray(F, &f));
217 
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 }
228 
229 /*------------------------------------------------------------*/
230 /* J[m][n] = df[m]/dx[n] */
EvaluateJacobian(Tao tao,Vec X,Mat J,Mat Jpre,void * ptr)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;
236 
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   }
244 
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));
248 
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 }
253 
254 /* ------------------------------------------------------------ */
255 /* Currently fixed matrix, in future may be dynamic for D(x)? */
FormDictionaryMatrix(Mat D,AppCtx * user)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));
262 
263   PetscCall(PetscLogFlops(0)); /* 0 for fixed dictionary matrix, >0 for varying dictionary matrix */
264   PetscFunctionReturn(PETSC_SUCCESS);
265 }
266 
267 /* ------------------------------------------------------------ */
FormStartingPoint(Vec X)268 PetscErrorCode FormStartingPoint(Vec X)
269 {
270   PetscFunctionBegin;
271   PetscCall(VecSet(X, 0.0));
272   PetscFunctionReturn(PETSC_SUCCESS);
273 }
274 
275 /* ---------------------------------------------------------------------- */
InitializeUserData(AppCtx * user)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. */
280 
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;
287 
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;
314 
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 }
328 
329 /*TEST
330 
331    build:
332       requires: !complex !single !quad !defined(PETSC_USE_64BIT_INDICES) !__float128
333 
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
337 
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 -tao_brgn_subsolver_tao_monitor
342 
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 -tao_brgn_subsolver_tao_monitor
347 
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 -tao_brgn_subsolver_tao_monitor
352 
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 -tao_brgn_subsolver_tao_monitor
357 
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
362 
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
367 
368 TEST*/
369