xref: /petsc/src/tao/leastsquares/tutorials/cs1.c (revision 58d68138c660dfb4e9f5b03334792cd4f2ffd7cc)
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 /* 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;
43 
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 *);
50 
51 /*--------------------------------------------------------------------*/
52 int main(int argc, char **argv) {
53   Vec       x, f; /* solution, function f(x) = A*x-b */
54   Mat       J, D; /* Jacobian matrix, Transform matrix */
55   Tao       tao;  /* Tao solver context */
56   PetscInt  i;    /* iteration information */
57   PetscReal hist[100], resid[100];
58   PetscInt  lits[100];
59   AppCtx    user; /* user-defined work context */
60 
61   PetscFunctionBeginUser;
62   PetscCall(PetscInitialize(&argc, &argv, (char *)0, help));
63 
64   /* Allocate solution and vector function vectors */
65   PetscCall(VecCreateSeq(PETSC_COMM_SELF, N, &x));
66   PetscCall(VecCreateSeq(PETSC_COMM_SELF, M, &f));
67 
68   /* Allocate Jacobian and Dictionary matrix. */
69   PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, M, N, NULL, &J));
70   PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, K, N, NULL, &D)); /* XH: TODO: dense -> sparse/dense/shell etc, do it on fly  */
71 
72   for (i = 0; i < M; i++) user.idm[i] = i;
73   for (i = 0; i < N; i++) user.idn[i] = i;
74   for (i = 0; i < K; i++) user.idk[i] = i;
75 
76   /* Create TAO solver and set desired solution method */
77   PetscCall(TaoCreate(PETSC_COMM_SELF, &tao));
78   PetscCall(TaoSetType(tao, TAOBRGN));
79 
80   /* User set application context: A, D matrice, and b vector. */
81   PetscCall(InitializeUserData(&user));
82 
83   /* Set initial guess */
84   PetscCall(FormStartingPoint(x));
85 
86   /* Fill the content of matrix D from user application Context */
87   PetscCall(FormDictionaryMatrix(D, &user));
88 
89   /* Bind x to tao->solution. */
90   PetscCall(TaoSetSolution(tao, x));
91   /* Bind D to tao->data->D */
92   PetscCall(TaoBRGNSetDictionaryMatrix(tao, D));
93 
94   /* Set the function and Jacobian routines. */
95   PetscCall(TaoSetResidualRoutine(tao, f, EvaluateFunction, (void *)&user));
96   PetscCall(TaoSetJacobianResidualRoutine(tao, J, J, EvaluateJacobian, (void *)&user));
97 
98   /* Check for any TAO command line arguments */
99   PetscCall(TaoSetFromOptions(tao));
100 
101   PetscCall(TaoSetConvergenceHistory(tao, hist, resid, 0, lits, 100, PETSC_TRUE));
102 
103   /* Perform the Solve */
104   PetscCall(TaoSolve(tao));
105 
106   /* XH: Debug: View the result, function and Jacobian.  */
107   PetscCall(PetscPrintf(PETSC_COMM_SELF, "-------- result x, residual f=A*x-b, and Jacobian=A. -------- \n"));
108   PetscCall(VecView(x, PETSC_VIEWER_STDOUT_SELF));
109   PetscCall(VecView(f, PETSC_VIEWER_STDOUT_SELF));
110   PetscCall(MatView(J, PETSC_VIEWER_STDOUT_SELF));
111   PetscCall(MatView(D, PETSC_VIEWER_STDOUT_SELF));
112 
113   /* Free TAO data structures */
114   PetscCall(TaoDestroy(&tao));
115 
116   /* Free PETSc data structures */
117   PetscCall(VecDestroy(&x));
118   PetscCall(VecDestroy(&f));
119   PetscCall(MatDestroy(&J));
120   PetscCall(MatDestroy(&D));
121 
122   PetscCall(PetscFinalize());
123   return 0;
124 }
125 
126 /*--------------------------------------------------------------------*/
127 PetscErrorCode EvaluateFunction(Tao tao, Vec X, Vec F, void *ptr) {
128   AppCtx          *user = (AppCtx *)ptr;
129   PetscInt         m, n;
130   const PetscReal *x;
131   PetscReal       *b = user->b, *f;
132 
133   PetscFunctionBegin;
134   PetscCall(VecGetArrayRead(X, &x));
135   PetscCall(VecGetArray(F, &f));
136 
137   /* 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 */
138   for (m = 0; m < M; m++) {
139     f[m] = -b[m];
140     for (n = 0; n < N; n++) { f[m] += user->A[m][n] * x[n]; }
141   }
142   PetscCall(VecRestoreArrayRead(X, &x));
143   PetscCall(VecRestoreArray(F, &f));
144   PetscLogFlops(2.0 * M * N);
145   PetscFunctionReturn(0);
146 }
147 
148 /*------------------------------------------------------------*/
149 /* J[m][n] = df[m]/dx[n] */
150 PetscErrorCode EvaluateJacobian(Tao tao, Vec X, Mat J, Mat Jpre, void *ptr) {
151   AppCtx          *user = (AppCtx *)ptr;
152   PetscInt         m, n;
153   const PetscReal *x;
154 
155   PetscFunctionBegin;
156   PetscCall(VecGetArrayRead(X, &x)); /* not used for linear least square, but keep for future nonlinear least square) */
157   /* 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?
158     For nonlinear least square, we require x to compute J, keep codes here for future nonlinear least square*/
159   for (m = 0; m < M; ++m) {
160     for (n = 0; n < N; ++n) { user->J[m][n] = user->A[m][n]; }
161   }
162 
163   PetscCall(MatSetValues(J, M, user->idm, N, user->idn, (PetscReal *)user->J, INSERT_VALUES));
164   PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY));
165   PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY));
166 
167   PetscCall(VecRestoreArrayRead(X, &x)); /* not used for linear least square, but keep for future nonlinear least square) */
168   PetscLogFlops(0);                      /* 0 for linear least square, >0 for nonlinear least square */
169   PetscFunctionReturn(0);
170 }
171 
172 /* ------------------------------------------------------------ */
173 /* Currently fixed matrix, in future may be dynamic for D(x)? */
174 PetscErrorCode FormDictionaryMatrix(Mat D, AppCtx *user) {
175   PetscFunctionBegin;
176   PetscCall(MatSetValues(D, K, user->idk, N, user->idn, (PetscReal *)user->D, INSERT_VALUES));
177   PetscCall(MatAssemblyBegin(D, MAT_FINAL_ASSEMBLY));
178   PetscCall(MatAssemblyEnd(D, MAT_FINAL_ASSEMBLY));
179 
180   PetscLogFlops(0); /* 0 for fixed dictionary matrix, >0 for varying dictionary matrix */
181   PetscFunctionReturn(0);
182 }
183 
184 /* ------------------------------------------------------------ */
185 PetscErrorCode FormStartingPoint(Vec X) {
186   PetscFunctionBegin;
187   PetscCall(VecSet(X, 0.0));
188   PetscFunctionReturn(0);
189 }
190 
191 /* ---------------------------------------------------------------------- */
192 PetscErrorCode InitializeUserData(AppCtx *user) {
193   PetscReal *b = user->b; /* **A=user->A, but we don't kown the dimension of A in this way, how to fix? */
194   PetscInt   m, n, k;     /* loop index for M,N,K dimension. */
195 
196   PetscFunctionBegin;
197   /* b = A*x while x = [0;0;1;0;0] here*/
198   m      = 0;
199   b[m++] = 0.28;
200   b[m++] = 0.55;
201   b[m++] = 0.96;
202 
203   /* matlab generated random matrix, uniformly distributed in [0,1] with 2 digits accuracy. rng(0); A = rand(M, N); A = round(A*100)/100;
204   A = [0.81  0.91  0.28  0.96  0.96
205        0.91  0.63  0.55  0.16  0.49
206        0.13  0.10  0.96  0.97  0.80]
207   */
208   m               = 0;
209   n               = 0;
210   user->A[m][n++] = 0.81;
211   user->A[m][n++] = 0.91;
212   user->A[m][n++] = 0.28;
213   user->A[m][n++] = 0.96;
214   user->A[m][n++] = 0.96;
215   ++m;
216   n               = 0;
217   user->A[m][n++] = 0.91;
218   user->A[m][n++] = 0.63;
219   user->A[m][n++] = 0.55;
220   user->A[m][n++] = 0.16;
221   user->A[m][n++] = 0.49;
222   ++m;
223   n               = 0;
224   user->A[m][n++] = 0.13;
225   user->A[m][n++] = 0.10;
226   user->A[m][n++] = 0.96;
227   user->A[m][n++] = 0.97;
228   user->A[m][n++] = 0.80;
229 
230   /* initialize to 0 */
231   for (k = 0; k < K; k++) {
232     for (n = 0; n < N; n++) { user->D[k][n] = 0.0; }
233   }
234   /* Choice I: set D to identity matrix of size N*N for testing */
235   /* for (k=0; k<K; k++) user->D[k][k] = 1.0; */
236   /* Choice II: set D to Backward difference matrix of size (N-1)*N, with zero extended boundary assumption */
237   for (k = 0; k < K; k++) {
238     user->D[k][k]     = -1.0;
239     user->D[k][k + 1] = 1.0;
240   }
241 
242   PetscFunctionReturn(0);
243 }
244 
245 /*TEST
246 
247    build:
248       requires: !complex !single !quad !defined(PETSC_USE_64BIT_INDICES)
249 
250    test:
251       localrunfiles: cs1Data_A_b_xGT
252       args: -tao_smonitor -tao_max_it 100 -tao_type pounders -tao_gatol 1.e-6
253 
254    test:
255       suffix: 2
256       localrunfiles: cs1Data_A_b_xGT
257       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
258 
259    test:
260       suffix: 3
261       localrunfiles: cs1Data_A_b_xGT
262       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
263 
264    test:
265       suffix: 4
266       localrunfiles: cs1Data_A_b_xGT
267       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
268 
269    test:
270       suffix: 5
271       localrunfiles: cs1Data_A_b_xGT
272       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
273 
274 TEST*/
275