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