xref: /petsc/src/tao/constrained/tutorials/tomographyADMM.c (revision c4762a1b19cd2af06abeed90e8f9d34fb975dd94)
1*c4762a1bSJed Brown #include <petsctao.h>
2*c4762a1bSJed Brown /*
3*c4762a1bSJed Brown Description:   ADMM tomography reconstruction example .
4*c4762a1bSJed Brown                0.5*||Ax-b||^2 + lambda*g(x)
5*c4762a1bSJed Brown Reference:     BRGN Tomography Example
6*c4762a1bSJed Brown */
7*c4762a1bSJed Brown 
8*c4762a1bSJed Brown static char help[] = "Finds the ADMM solution to the under constraint linear model Ax = b, with regularizer. \n\
9*c4762a1bSJed Brown                       A is a M*N real matrix (M<N), x is sparse. A good regularizer is an L1 regularizer. \n\
10*c4762a1bSJed Brown                       We first split the operator into 0.5*||Ax-b||^2, f(x), and lambda*||x||_1, g(z), where lambda is user specified weight. \n\
11*c4762a1bSJed Brown                       g(z) could be either ||z||_1, or ||z||_2^2. Default closed form solution for NORM1 would be soft-threshold, which is \n\
12*c4762a1bSJed Brown                       natively supported in admm.c with -tao_admm_regularizer_type soft-threshold. Or user can use regular TAO solver for  \n\
13*c4762a1bSJed Brown                       either NORM1 or NORM2 or TAOSHELL, with -reg {1,2,3} \n\
14*c4762a1bSJed Brown                       Then, we augment both f and g, and solve it via ADMM. \n\
15*c4762a1bSJed Brown                       D is the M*N transform matrix so that D*x is sparse. \n";
16*c4762a1bSJed Brown 
17*c4762a1bSJed Brown typedef struct {
18*c4762a1bSJed Brown   PetscInt  M,N,K,reg;
19*c4762a1bSJed Brown   PetscReal lambda,eps,mumin;
20*c4762a1bSJed Brown   Mat       A,ATA,H,Hx,D,Hz,DTD,HF;
21*c4762a1bSJed Brown   Vec       c,xlb,xub,x,b,workM,workN,workN2,workN3,xGT;    /* observation b, ground truth xGT, the lower bound and upper bound of x*/
22*c4762a1bSJed Brown } AppCtx;
23*c4762a1bSJed Brown 
24*c4762a1bSJed Brown /*------------------------------------------------------------*/
25*c4762a1bSJed Brown 
26*c4762a1bSJed Brown PetscErrorCode NullJacobian(Tao tao,Vec X,Mat J,Mat Jpre,void *ptr)
27*c4762a1bSJed Brown {
28*c4762a1bSJed Brown   PetscFunctionBegin;
29*c4762a1bSJed Brown   PetscFunctionReturn(0);
30*c4762a1bSJed Brown }
31*c4762a1bSJed Brown 
32*c4762a1bSJed Brown /*------------------------------------------------------------*/
33*c4762a1bSJed Brown 
34*c4762a1bSJed Brown static PetscErrorCode TaoShellSolve_SoftThreshold(Tao tao)
35*c4762a1bSJed Brown {
36*c4762a1bSJed Brown   PetscErrorCode ierr;
37*c4762a1bSJed Brown   PetscReal      lambda, mu;
38*c4762a1bSJed Brown   AppCtx         *user;
39*c4762a1bSJed Brown   Vec            out,work,y,x;
40*c4762a1bSJed Brown   Tao            admm_tao,misfit;
41*c4762a1bSJed Brown 
42*c4762a1bSJed Brown   PetscFunctionBegin;
43*c4762a1bSJed Brown   user = NULL;
44*c4762a1bSJed Brown   mu   = 0;
45*c4762a1bSJed Brown   ierr = TaoGetADMMParentTao(tao,&admm_tao);CHKERRQ(ierr);
46*c4762a1bSJed Brown   ierr = TaoADMMGetMisfitSubsolver(admm_tao, &misfit);CHKERRQ(ierr);
47*c4762a1bSJed Brown   ierr = TaoADMMGetSpectralPenalty(admm_tao,&mu);CHKERRQ(ierr);
48*c4762a1bSJed Brown   ierr = TaoShellGetContext(tao, (void**) &user);CHKERRQ(ierr);
49*c4762a1bSJed Brown 
50*c4762a1bSJed Brown   lambda = user->lambda;
51*c4762a1bSJed Brown   work   = user->workN;
52*c4762a1bSJed Brown   ierr   = TaoGetSolutionVector(tao, &out);CHKERRQ(ierr);
53*c4762a1bSJed Brown   ierr   = TaoGetSolutionVector(misfit, &x);CHKERRQ(ierr);
54*c4762a1bSJed Brown   ierr   = TaoADMMGetDualVector(admm_tao, &y);CHKERRQ(ierr);
55*c4762a1bSJed Brown 
56*c4762a1bSJed Brown   /* Dx + y/mu */
57*c4762a1bSJed Brown   ierr = MatMult(user->D,x,work);CHKERRQ(ierr);
58*c4762a1bSJed Brown   ierr = VecAXPY(work,1/mu,y);CHKERRQ(ierr);
59*c4762a1bSJed Brown 
60*c4762a1bSJed Brown   /* soft thresholding */
61*c4762a1bSJed Brown   ierr = TaoSoftThreshold(work, -lambda/mu, lambda/mu, out);CHKERRQ(ierr);
62*c4762a1bSJed Brown   PetscFunctionReturn(0);
63*c4762a1bSJed Brown }
64*c4762a1bSJed Brown 
65*c4762a1bSJed Brown /*------------------------------------------------------------*/
66*c4762a1bSJed Brown 
67*c4762a1bSJed Brown PetscErrorCode MisfitObjectiveAndGradient(Tao tao,Vec X,PetscReal *f,Vec g,void *ptr)
68*c4762a1bSJed Brown {
69*c4762a1bSJed Brown   AppCtx         *user = (AppCtx*)ptr;
70*c4762a1bSJed Brown   PetscErrorCode ierr;
71*c4762a1bSJed Brown 
72*c4762a1bSJed Brown   PetscFunctionBegin;
73*c4762a1bSJed Brown   /* Objective  0.5*||Ax-b||_2^2 */
74*c4762a1bSJed Brown   ierr = MatMult(user->A,X,user->workM);CHKERRQ(ierr);
75*c4762a1bSJed Brown   ierr = VecAXPY(user->workM,-1,user->b);CHKERRQ(ierr);
76*c4762a1bSJed Brown   ierr = VecDot(user->workM,user->workM,f);CHKERRQ(ierr);
77*c4762a1bSJed Brown   *f  *= 0.5;
78*c4762a1bSJed Brown   /* Gradient. ATAx-ATb */
79*c4762a1bSJed Brown   ierr = MatMult(user->ATA,X,user->workN);CHKERRQ(ierr);
80*c4762a1bSJed Brown   ierr = MatMultTranspose(user->A,user->b,user->workN2);CHKERRQ(ierr);
81*c4762a1bSJed Brown   ierr = VecWAXPY(g,-1.,user->workN2,user->workN);CHKERRQ(ierr);
82*c4762a1bSJed Brown   PetscFunctionReturn(0);
83*c4762a1bSJed Brown }
84*c4762a1bSJed Brown 
85*c4762a1bSJed Brown /*------------------------------------------------------------*/
86*c4762a1bSJed Brown 
87*c4762a1bSJed Brown PetscErrorCode RegularizerObjectiveAndGradient1(Tao tao,Vec X,PetscReal *f_reg,Vec G_reg,void *ptr)
88*c4762a1bSJed Brown {
89*c4762a1bSJed Brown   AppCtx         *user = (AppCtx*)ptr;
90*c4762a1bSJed Brown   PetscErrorCode ierr;
91*c4762a1bSJed Brown 
92*c4762a1bSJed Brown   PetscFunctionBegin;
93*c4762a1bSJed Brown   /* compute regularizer objective
94*c4762a1bSJed Brown    * f = f + lambda*sum(sqrt(y.^2+epsilon^2) - epsilon), where y = D*x */
95*c4762a1bSJed Brown   ierr    = VecCopy(X,user->workN2);CHKERRQ(ierr);
96*c4762a1bSJed Brown   ierr    = VecPow(user->workN2,2.);CHKERRQ(ierr);
97*c4762a1bSJed Brown   ierr    = VecShift(user->workN2,user->eps*user->eps);CHKERRQ(ierr);
98*c4762a1bSJed Brown   ierr    = VecSqrtAbs(user->workN2);CHKERRQ(ierr);
99*c4762a1bSJed Brown   ierr    = VecCopy(user->workN2, user->workN3);CHKERRQ(ierr);
100*c4762a1bSJed Brown   ierr    = VecShift(user->workN2,-user->eps);CHKERRQ(ierr);
101*c4762a1bSJed Brown   ierr    = VecSum(user->workN2,f_reg);CHKERRQ(ierr);
102*c4762a1bSJed Brown   *f_reg *= user->lambda;
103*c4762a1bSJed Brown   /* compute regularizer gradient = lambda*x */
104*c4762a1bSJed Brown   ierr = VecPointwiseDivide(G_reg,X,user->workN3);CHKERRQ(ierr);
105*c4762a1bSJed Brown   ierr = VecScale(G_reg,user->lambda);CHKERRQ(ierr);
106*c4762a1bSJed Brown   PetscFunctionReturn(0);
107*c4762a1bSJed Brown }
108*c4762a1bSJed Brown 
109*c4762a1bSJed Brown /*------------------------------------------------------------*/
110*c4762a1bSJed Brown 
111*c4762a1bSJed Brown PetscErrorCode RegularizerObjectiveAndGradient2(Tao tao,Vec X,PetscReal *f_reg,Vec G_reg,void *ptr)
112*c4762a1bSJed Brown {
113*c4762a1bSJed Brown   AppCtx         *user = (AppCtx*)ptr;
114*c4762a1bSJed Brown   PetscErrorCode ierr;
115*c4762a1bSJed Brown   PetscReal      temp;
116*c4762a1bSJed Brown 
117*c4762a1bSJed Brown   PetscFunctionBegin;
118*c4762a1bSJed Brown   /* compute regularizer objective = lambda*|z|_2^2 */
119*c4762a1bSJed Brown   ierr   = VecDot(X,X,&temp);CHKERRQ(ierr);
120*c4762a1bSJed Brown   *f_reg = 0.5*user->lambda*temp;
121*c4762a1bSJed Brown   /* compute regularizer gradient = lambda*z */
122*c4762a1bSJed Brown   ierr = VecCopy(X,G_reg);
123*c4762a1bSJed Brown   ierr = VecScale(G_reg,user->lambda);CHKERRQ(ierr);
124*c4762a1bSJed Brown   PetscFunctionReturn(0);
125*c4762a1bSJed Brown }
126*c4762a1bSJed Brown 
127*c4762a1bSJed Brown /*------------------------------------------------------------*/
128*c4762a1bSJed Brown 
129*c4762a1bSJed Brown static PetscErrorCode HessianMisfit(Tao tao, Vec x, Mat H, Mat Hpre, void *ptr)
130*c4762a1bSJed Brown {
131*c4762a1bSJed Brown   PetscFunctionBegin;
132*c4762a1bSJed Brown   PetscFunctionReturn(0);
133*c4762a1bSJed Brown }
134*c4762a1bSJed Brown 
135*c4762a1bSJed Brown /*------------------------------------------------------------*/
136*c4762a1bSJed Brown 
137*c4762a1bSJed Brown static PetscErrorCode HessianReg(Tao tao, Vec x, Mat H, Mat Hpre, void *ptr)
138*c4762a1bSJed Brown {
139*c4762a1bSJed Brown   AppCtx         *user = (AppCtx*)ptr;
140*c4762a1bSJed Brown   PetscErrorCode ierr;
141*c4762a1bSJed Brown 
142*c4762a1bSJed Brown   PetscFunctionBegin;
143*c4762a1bSJed Brown   ierr = MatMult(user->D,x,user->workN);CHKERRQ(ierr);
144*c4762a1bSJed Brown   ierr = VecPow(user->workN2,2.);CHKERRQ(ierr);
145*c4762a1bSJed Brown   ierr = VecShift(user->workN2,user->eps*user->eps);CHKERRQ(ierr);
146*c4762a1bSJed Brown   ierr = VecSqrtAbs(user->workN2);CHKERRQ(ierr);
147*c4762a1bSJed Brown   ierr = VecShift(user->workN2,-user->eps);CHKERRQ(ierr);
148*c4762a1bSJed Brown   ierr = VecReciprocal(user->workN2);CHKERRQ(ierr);
149*c4762a1bSJed Brown   ierr = VecScale(user->workN2,user->eps*user->eps);CHKERRQ(ierr);
150*c4762a1bSJed Brown   ierr = MatDiagonalSet(H,user->workN2,INSERT_VALUES);CHKERRQ(ierr);
151*c4762a1bSJed Brown   PetscFunctionReturn(0);
152*c4762a1bSJed Brown }
153*c4762a1bSJed Brown 
154*c4762a1bSJed Brown /*------------------------------------------------------------*/
155*c4762a1bSJed Brown 
156*c4762a1bSJed Brown PetscErrorCode FullObjGrad(Tao tao,Vec X,PetscReal *f,Vec g,void *ptr)
157*c4762a1bSJed Brown {
158*c4762a1bSJed Brown   AppCtx         *user = (AppCtx*)ptr;
159*c4762a1bSJed Brown   PetscErrorCode ierr;
160*c4762a1bSJed Brown   PetscReal      f_reg;
161*c4762a1bSJed Brown 
162*c4762a1bSJed Brown   PetscFunctionBegin;
163*c4762a1bSJed Brown   /* Objective  0.5*||Ax-b||_2^2 + lambda*||x||_2^2*/
164*c4762a1bSJed Brown   ierr = MatMult(user->A,X,user->workM);CHKERRQ(ierr);
165*c4762a1bSJed Brown   ierr = VecAXPY(user->workM,-1,user->b);CHKERRQ(ierr);
166*c4762a1bSJed Brown   ierr = VecDot(user->workM,user->workM,f);CHKERRQ(ierr);
167*c4762a1bSJed Brown   ierr = VecNorm(X,NORM_2,&f_reg);CHKERRQ(ierr);
168*c4762a1bSJed Brown   *f  *= 0.5;
169*c4762a1bSJed Brown   *f  += user->lambda*f_reg*f_reg;
170*c4762a1bSJed Brown   /* Gradient. ATAx-ATb + 2*lambda*x */
171*c4762a1bSJed Brown   ierr = MatMult(user->ATA,X,user->workN);CHKERRQ(ierr);
172*c4762a1bSJed Brown   ierr = MatMultTranspose(user->A,user->b,user->workN2);CHKERRQ(ierr);
173*c4762a1bSJed Brown   ierr = VecWAXPY(g,-1.,user->workN2,user->workN);CHKERRQ(ierr);
174*c4762a1bSJed Brown   ierr = VecAXPY(g,2*user->lambda,X);CHKERRQ(ierr);
175*c4762a1bSJed Brown   PetscFunctionReturn(0);
176*c4762a1bSJed Brown }
177*c4762a1bSJed Brown /*------------------------------------------------------------*/
178*c4762a1bSJed Brown 
179*c4762a1bSJed Brown static PetscErrorCode HessianFull(Tao tao, Vec x, Mat H, Mat Hpre, void *ptr)
180*c4762a1bSJed Brown {
181*c4762a1bSJed Brown   PetscFunctionBegin;
182*c4762a1bSJed Brown   PetscFunctionReturn(0);
183*c4762a1bSJed Brown }
184*c4762a1bSJed Brown /*------------------------------------------------------------*/
185*c4762a1bSJed Brown 
186*c4762a1bSJed Brown 
187*c4762a1bSJed Brown PetscErrorCode InitializeUserData(AppCtx *user)
188*c4762a1bSJed Brown {
189*c4762a1bSJed Brown   char           dataFile[] = "tomographyData_A_b_xGT";   /* Matrix A and vectors b, xGT(ground truth) binary files generated by Matlab. Debug: change from "tomographyData_A_b_xGT" to "cs1Data_A_b_xGT". */
190*c4762a1bSJed Brown   PetscViewer    fd;   /* used to load data from file */
191*c4762a1bSJed Brown   PetscErrorCode ierr;
192*c4762a1bSJed Brown   PetscInt       k,n;
193*c4762a1bSJed Brown   PetscScalar    v;
194*c4762a1bSJed Brown   PetscFunctionBegin;
195*c4762a1bSJed Brown 
196*c4762a1bSJed Brown   /* Load the A matrix, b vector, and xGT vector from a binary file. */
197*c4762a1bSJed Brown   ierr = PetscViewerBinaryOpen(PETSC_COMM_WORLD,dataFile,FILE_MODE_READ,&fd);CHKERRQ(ierr);
198*c4762a1bSJed Brown   ierr = MatCreate(PETSC_COMM_WORLD,&user->A);CHKERRQ(ierr);
199*c4762a1bSJed Brown   ierr = MatSetType(user->A,MATAIJ);CHKERRQ(ierr);
200*c4762a1bSJed Brown   ierr = MatLoad(user->A,fd);CHKERRQ(ierr);
201*c4762a1bSJed Brown   ierr = VecCreate(PETSC_COMM_WORLD,&user->b);CHKERRQ(ierr);
202*c4762a1bSJed Brown   ierr = VecLoad(user->b,fd);CHKERRQ(ierr);
203*c4762a1bSJed Brown   ierr = VecCreate(PETSC_COMM_WORLD,&user->xGT);CHKERRQ(ierr);
204*c4762a1bSJed Brown   ierr = VecLoad(user->xGT,fd);CHKERRQ(ierr);
205*c4762a1bSJed Brown   ierr = PetscViewerDestroy(&fd);CHKERRQ(ierr);
206*c4762a1bSJed Brown 
207*c4762a1bSJed Brown   ierr = MatGetSize(user->A,&user->M,&user->N);CHKERRQ(ierr);
208*c4762a1bSJed Brown 
209*c4762a1bSJed Brown   ierr = MatCreate(PETSC_COMM_WORLD,&user->D);CHKERRQ(ierr);
210*c4762a1bSJed Brown   ierr = MatSetSizes(user->D,PETSC_DECIDE,PETSC_DECIDE,user->N,user->N);CHKERRQ(ierr);
211*c4762a1bSJed Brown   ierr = MatSetFromOptions(user->D);CHKERRQ(ierr);
212*c4762a1bSJed Brown   ierr = MatSetUp(user->D);CHKERRQ(ierr);
213*c4762a1bSJed Brown   for (k=0; k<user->N; k++) {
214*c4762a1bSJed Brown     v = 1.0;
215*c4762a1bSJed Brown     n = k+1;
216*c4762a1bSJed Brown     if (k< user->N -1) {
217*c4762a1bSJed Brown       ierr = MatSetValues(user->D,1,&k,1,&n,&v,INSERT_VALUES);CHKERRQ(ierr);
218*c4762a1bSJed Brown     }
219*c4762a1bSJed Brown     v    = -1.0;
220*c4762a1bSJed Brown     ierr = MatSetValues(user->D,1,&k,1,&k,&v,INSERT_VALUES);CHKERRQ(ierr);
221*c4762a1bSJed Brown   }
222*c4762a1bSJed Brown   ierr = MatAssemblyBegin(user->D,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
223*c4762a1bSJed Brown   ierr = MatAssemblyEnd(user->D,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
224*c4762a1bSJed Brown 
225*c4762a1bSJed Brown   ierr = MatTransposeMatMult(user->D,user->D,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&user->DTD);CHKERRQ(ierr);
226*c4762a1bSJed Brown 
227*c4762a1bSJed Brown   ierr = MatCreate(PETSC_COMM_WORLD,&user->Hz);CHKERRQ(ierr);
228*c4762a1bSJed Brown   ierr = MatSetSizes(user->Hz,PETSC_DECIDE,PETSC_DECIDE,user->N,user->N);CHKERRQ(ierr);
229*c4762a1bSJed Brown   ierr = MatSetFromOptions(user->Hz);CHKERRQ(ierr);
230*c4762a1bSJed Brown   ierr = MatSetUp(user->Hz);CHKERRQ(ierr);
231*c4762a1bSJed Brown   ierr = MatAssemblyBegin(user->Hz,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
232*c4762a1bSJed Brown   ierr = MatAssemblyEnd(user->Hz,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
233*c4762a1bSJed Brown 
234*c4762a1bSJed Brown   ierr = VecCreate(PETSC_COMM_WORLD,&(user->x));CHKERRQ(ierr);
235*c4762a1bSJed Brown   ierr = VecCreate(PETSC_COMM_WORLD,&(user->workM));CHKERRQ(ierr);
236*c4762a1bSJed Brown   ierr = VecCreate(PETSC_COMM_WORLD,&(user->workN));CHKERRQ(ierr);
237*c4762a1bSJed Brown   ierr = VecCreate(PETSC_COMM_WORLD,&(user->workN2));CHKERRQ(ierr);
238*c4762a1bSJed Brown   ierr = VecSetSizes(user->x,PETSC_DECIDE,user->N);CHKERRQ(ierr);
239*c4762a1bSJed Brown   ierr = VecSetSizes(user->workM,PETSC_DECIDE,user->M);CHKERRQ(ierr);
240*c4762a1bSJed Brown   ierr = VecSetSizes(user->workN,PETSC_DECIDE,user->N);CHKERRQ(ierr);
241*c4762a1bSJed Brown   ierr = VecSetSizes(user->workN2,PETSC_DECIDE,user->N);CHKERRQ(ierr);
242*c4762a1bSJed Brown   ierr = VecSetFromOptions(user->x);CHKERRQ(ierr);
243*c4762a1bSJed Brown   ierr = VecSetFromOptions(user->workM);CHKERRQ(ierr);
244*c4762a1bSJed Brown   ierr = VecSetFromOptions(user->workN);CHKERRQ(ierr);
245*c4762a1bSJed Brown   ierr = VecSetFromOptions(user->workN2);CHKERRQ(ierr);
246*c4762a1bSJed Brown 
247*c4762a1bSJed Brown   ierr = VecDuplicate(user->workN,&(user->workN3));CHKERRQ(ierr);
248*c4762a1bSJed Brown   ierr = VecDuplicate(user->x,&(user->xlb));CHKERRQ(ierr);
249*c4762a1bSJed Brown   ierr = VecDuplicate(user->x,&(user->xub));CHKERRQ(ierr);
250*c4762a1bSJed Brown   ierr = VecDuplicate(user->x,&(user->c));CHKERRQ(ierr);
251*c4762a1bSJed Brown   ierr = VecSet(user->xlb,0.0);CHKERRQ(ierr);
252*c4762a1bSJed Brown   ierr = VecSet(user->c,0.0);CHKERRQ(ierr);
253*c4762a1bSJed Brown   ierr = VecSet(user->xub,PETSC_INFINITY);CHKERRQ(ierr);
254*c4762a1bSJed Brown 
255*c4762a1bSJed Brown   ierr = MatTransposeMatMult(user->A,user->A, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &(user->ATA));CHKERRQ(ierr);
256*c4762a1bSJed Brown   ierr = MatTransposeMatMult(user->A,user->A, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &(user->Hx));CHKERRQ(ierr);
257*c4762a1bSJed Brown   ierr = MatTransposeMatMult(user->A,user->A, MAT_INITIAL_MATRIX, PETSC_DEFAULT, &(user->HF));CHKERRQ(ierr);
258*c4762a1bSJed Brown 
259*c4762a1bSJed Brown   ierr = MatAssemblyBegin(user->ATA,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
260*c4762a1bSJed Brown   ierr = MatAssemblyEnd(user->ATA,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
261*c4762a1bSJed Brown   ierr = MatAssemblyBegin(user->Hx,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
262*c4762a1bSJed Brown   ierr = MatAssemblyEnd(user->Hx,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
263*c4762a1bSJed Brown   ierr = MatAssemblyBegin(user->HF,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
264*c4762a1bSJed Brown   ierr = MatAssemblyEnd(user->HF,MAT_FINAL_ASSEMBLY);CHKERRQ(ierr);
265*c4762a1bSJed Brown 
266*c4762a1bSJed Brown   user->lambda = 1.e-8;
267*c4762a1bSJed Brown   user->eps    = 1.e-3;
268*c4762a1bSJed Brown   user->reg    = 2;
269*c4762a1bSJed Brown   user->mumin  = 5.e-6;
270*c4762a1bSJed Brown 
271*c4762a1bSJed Brown   ierr = PetscOptionsBegin(PETSC_COMM_WORLD, NULL, "Configure separable objection example", "tomographyADMM.c");CHKERRQ(ierr);
272*c4762a1bSJed Brown   ierr = PetscOptionsInt("-reg","Regularization scheme for z solver (1,2)", "tomographyADMM.c", user->reg, &(user->reg), NULL);CHKERRQ(ierr);
273*c4762a1bSJed Brown   ierr = PetscOptionsReal("-lambda", "The regularization multiplier. 1 default", "tomographyADMM.c", user->lambda, &(user->lambda), NULL);CHKERRQ(ierr);
274*c4762a1bSJed Brown   ierr = PetscOptionsReal("-eps", "L1 norm epsilon padding", "tomographyADMM.c", user->eps, &(user->eps), NULL);CHKERRQ(ierr);
275*c4762a1bSJed Brown   ierr = PetscOptionsReal("-mumin", "Minimum value for ADMM spectral penalty", "tomographyADMM.c", user->mumin, &(user->mumin), NULL);CHKERRQ(ierr);
276*c4762a1bSJed Brown   ierr = PetscOptionsEnd();CHKERRQ(ierr);
277*c4762a1bSJed Brown   PetscFunctionReturn(0);
278*c4762a1bSJed Brown }
279*c4762a1bSJed Brown 
280*c4762a1bSJed Brown /*------------------------------------------------------------*/
281*c4762a1bSJed Brown 
282*c4762a1bSJed Brown PetscErrorCode DestroyContext(AppCtx *user)
283*c4762a1bSJed Brown {
284*c4762a1bSJed Brown   PetscErrorCode ierr;
285*c4762a1bSJed Brown 
286*c4762a1bSJed Brown   PetscFunctionBegin;
287*c4762a1bSJed Brown   ierr = MatDestroy(&user->A);CHKERRQ(ierr);
288*c4762a1bSJed Brown   ierr = MatDestroy(&user->ATA);CHKERRQ(ierr);
289*c4762a1bSJed Brown   ierr = MatDestroy(&user->Hx);CHKERRQ(ierr);
290*c4762a1bSJed Brown   ierr = MatDestroy(&user->Hz);CHKERRQ(ierr);
291*c4762a1bSJed Brown   ierr = MatDestroy(&user->HF);CHKERRQ(ierr);
292*c4762a1bSJed Brown   ierr = MatDestroy(&user->D);CHKERRQ(ierr);
293*c4762a1bSJed Brown   ierr = MatDestroy(&user->DTD);CHKERRQ(ierr);
294*c4762a1bSJed Brown   ierr = VecDestroy(&user->xGT);CHKERRQ(ierr);
295*c4762a1bSJed Brown   ierr = VecDestroy(&user->xlb);CHKERRQ(ierr);
296*c4762a1bSJed Brown   ierr = VecDestroy(&user->xub);CHKERRQ(ierr);
297*c4762a1bSJed Brown   ierr = VecDestroy(&user->b);CHKERRQ(ierr);
298*c4762a1bSJed Brown   ierr = VecDestroy(&user->x);CHKERRQ(ierr);
299*c4762a1bSJed Brown   ierr = VecDestroy(&user->c);CHKERRQ(ierr);
300*c4762a1bSJed Brown   ierr = VecDestroy(&user->workN3);CHKERRQ(ierr);
301*c4762a1bSJed Brown   ierr = VecDestroy(&user->workN2);CHKERRQ(ierr);
302*c4762a1bSJed Brown   ierr = VecDestroy(&user->workN);CHKERRQ(ierr);
303*c4762a1bSJed Brown   ierr = VecDestroy(&user->workM);CHKERRQ(ierr);
304*c4762a1bSJed Brown   PetscFunctionReturn(0);
305*c4762a1bSJed Brown }
306*c4762a1bSJed Brown 
307*c4762a1bSJed Brown /*------------------------------------------------------------*/
308*c4762a1bSJed Brown 
309*c4762a1bSJed Brown int main(int argc,char **argv)
310*c4762a1bSJed Brown {
311*c4762a1bSJed Brown   PetscErrorCode ierr;
312*c4762a1bSJed Brown   Tao            tao,misfit,reg;
313*c4762a1bSJed Brown   PetscReal      v1,v2;
314*c4762a1bSJed Brown   AppCtx*        user;
315*c4762a1bSJed Brown   PetscViewer    fd;
316*c4762a1bSJed Brown   char           resultFile[] = "tomographyResult_x";
317*c4762a1bSJed Brown 
318*c4762a1bSJed Brown   ierr = PetscInitialize(&argc,&argv,(char*)0,help);if (ierr) return ierr;
319*c4762a1bSJed Brown   ierr = PetscNew(&user);CHKERRQ(ierr);
320*c4762a1bSJed Brown   ierr = InitializeUserData(user);CHKERRQ(ierr);
321*c4762a1bSJed Brown 
322*c4762a1bSJed Brown   ierr = TaoCreate(PETSC_COMM_WORLD, &tao);CHKERRQ(ierr);
323*c4762a1bSJed Brown   ierr = TaoSetType(tao, TAOADMM);CHKERRQ(ierr);
324*c4762a1bSJed Brown   ierr = TaoSetInitialVector(tao, user->x);CHKERRQ(ierr);
325*c4762a1bSJed Brown   /* f(x) + g(x) for parent tao */
326*c4762a1bSJed Brown   ierr = TaoADMMSetSpectralPenalty(tao,1.);CHKERRQ(ierr);
327*c4762a1bSJed Brown   ierr = TaoSetObjectiveAndGradientRoutine(tao, FullObjGrad, (void*)user);CHKERRQ(ierr);
328*c4762a1bSJed Brown   ierr = MatShift(user->HF,user->lambda);CHKERRQ(ierr);
329*c4762a1bSJed Brown   ierr = TaoSetHessianRoutine(tao, user->HF, user->HF, HessianFull, (void*)user);CHKERRQ(ierr);
330*c4762a1bSJed Brown 
331*c4762a1bSJed Brown   /* f(x) for misfit tao */
332*c4762a1bSJed Brown   ierr = TaoADMMSetMisfitObjectiveAndGradientRoutine(tao, MisfitObjectiveAndGradient, (void*)user);CHKERRQ(ierr);
333*c4762a1bSJed Brown   ierr = TaoADMMSetMisfitHessianRoutine(tao, user->Hx, user->Hx, HessianMisfit, (void*)user);CHKERRQ(ierr);
334*c4762a1bSJed Brown   ierr = TaoADMMSetMisfitHessianChangeStatus(tao,PETSC_FALSE);CHKERRQ(ierr);
335*c4762a1bSJed Brown   ierr = TaoADMMSetMisfitConstraintJacobian(tao,user->D,user->D,NullJacobian,(void*)user);CHKERRQ(ierr);
336*c4762a1bSJed Brown 
337*c4762a1bSJed Brown   /* g(x) for regularizer tao */
338*c4762a1bSJed Brown   if (user->reg == 1) {
339*c4762a1bSJed Brown     ierr = TaoADMMSetRegularizerObjectiveAndGradientRoutine(tao, RegularizerObjectiveAndGradient1, (void*)user);CHKERRQ(ierr);
340*c4762a1bSJed Brown     ierr = TaoADMMSetRegularizerHessianRoutine(tao, user->Hz, user->Hz, HessianReg, (void*)user);CHKERRQ(ierr);
341*c4762a1bSJed Brown     ierr = TaoADMMSetRegHessianChangeStatus(tao,PETSC_TRUE);CHKERRQ(ierr);
342*c4762a1bSJed Brown   } else if (user->reg == 2) {
343*c4762a1bSJed Brown     ierr = TaoADMMSetRegularizerObjectiveAndGradientRoutine(tao, RegularizerObjectiveAndGradient2, (void*)user);CHKERRQ(ierr);
344*c4762a1bSJed Brown     ierr = MatShift(user->Hz,1);CHKERRQ(ierr);
345*c4762a1bSJed Brown     ierr = MatScale(user->Hz,user->lambda);CHKERRQ(ierr);
346*c4762a1bSJed Brown     ierr = TaoADMMSetRegularizerHessianRoutine(tao, user->Hz, user->Hz, HessianMisfit, (void*)user);CHKERRQ(ierr);
347*c4762a1bSJed Brown     ierr = TaoADMMSetRegHessianChangeStatus(tao,PETSC_TRUE);CHKERRQ(ierr);
348*c4762a1bSJed Brown   } else if (user->reg != 3) SETERRQ(PETSC_COMM_WORLD, 1, "Incorrect Reg type"); /* TaoShell case */
349*c4762a1bSJed Brown 
350*c4762a1bSJed Brown   /* Set type for the misfit solver */
351*c4762a1bSJed Brown   ierr = TaoADMMGetMisfitSubsolver(tao, &misfit);CHKERRQ(ierr);
352*c4762a1bSJed Brown   ierr = TaoADMMGetRegularizationSubsolver(tao, &reg);CHKERRQ(ierr);
353*c4762a1bSJed Brown   ierr = TaoSetType(misfit,TAONLS);CHKERRQ(ierr);
354*c4762a1bSJed Brown   if (user->reg == 3) {
355*c4762a1bSJed Brown     ierr = TaoSetType(reg,TAOSHELL);CHKERRQ(ierr);
356*c4762a1bSJed Brown     ierr = TaoShellSetContext(reg, (void*) user);CHKERRQ(ierr);
357*c4762a1bSJed Brown     ierr = TaoShellSetSolve(reg, TaoShellSolve_SoftThreshold);CHKERRQ(ierr);
358*c4762a1bSJed Brown   } else {
359*c4762a1bSJed Brown     ierr = TaoSetType(reg,TAONLS);CHKERRQ(ierr);
360*c4762a1bSJed Brown   }
361*c4762a1bSJed Brown   ierr = TaoSetVariableBounds(misfit,user->xlb,user->xub);CHKERRQ(ierr);
362*c4762a1bSJed Brown 
363*c4762a1bSJed Brown   /* Soft Thresholding solves the ADMM problem with the L1 regularizer lambda*||z||_1 and the x-z=0 constraint */
364*c4762a1bSJed Brown   ierr = TaoADMMSetRegularizerCoefficient(tao, user->lambda);CHKERRQ(ierr);
365*c4762a1bSJed Brown   ierr = TaoADMMSetRegularizerConstraintJacobian(tao,NULL,NULL,NullJacobian,(void*)user);CHKERRQ(ierr);
366*c4762a1bSJed Brown   ierr = TaoADMMSetMinimumSpectralPenalty(tao,user->mumin);CHKERRQ(ierr);
367*c4762a1bSJed Brown 
368*c4762a1bSJed Brown   ierr = TaoADMMSetConstraintVectorRHS(tao,user->c);CHKERRQ(ierr);
369*c4762a1bSJed Brown   ierr = TaoSetFromOptions(tao);CHKERRQ(ierr);
370*c4762a1bSJed Brown   ierr = TaoSolve(tao);CHKERRQ(ierr);
371*c4762a1bSJed Brown 
372*c4762a1bSJed Brown   /* Save x (reconstruction of object) vector to a binary file, which maybe read from Matlab and convert to a 2D image for comparison. */
373*c4762a1bSJed Brown   ierr = PetscViewerBinaryOpen(PETSC_COMM_WORLD,resultFile,FILE_MODE_WRITE,&fd);CHKERRQ(ierr);
374*c4762a1bSJed Brown   ierr = VecView(user->x,fd);CHKERRQ(ierr);
375*c4762a1bSJed Brown   ierr = PetscViewerDestroy(&fd);CHKERRQ(ierr);
376*c4762a1bSJed Brown 
377*c4762a1bSJed Brown   /* compute the error */
378*c4762a1bSJed Brown   ierr = VecAXPY(user->x,-1,user->xGT);CHKERRQ(ierr);
379*c4762a1bSJed Brown   ierr = VecNorm(user->x,NORM_2,&v1);CHKERRQ(ierr);
380*c4762a1bSJed Brown   ierr = VecNorm(user->xGT,NORM_2,&v2);CHKERRQ(ierr);
381*c4762a1bSJed Brown   ierr = PetscPrintf(PETSC_COMM_WORLD, "relative reconstruction error: ||x-xGT||/||xGT|| = %6.4e.\n", (double)(v1/v2));CHKERRQ(ierr);
382*c4762a1bSJed Brown 
383*c4762a1bSJed Brown   /* Free TAO data structures */
384*c4762a1bSJed Brown   ierr = TaoDestroy(&tao);CHKERRQ(ierr);
385*c4762a1bSJed Brown   ierr = DestroyContext(user);CHKERRQ(ierr);
386*c4762a1bSJed Brown   ierr = PetscFree(user);CHKERRQ(ierr);
387*c4762a1bSJed Brown   ierr = PetscFinalize();
388*c4762a1bSJed Brown   return ierr;
389*c4762a1bSJed Brown }
390*c4762a1bSJed Brown 
391*c4762a1bSJed Brown /*TEST
392*c4762a1bSJed Brown 
393*c4762a1bSJed Brown    build:
394*c4762a1bSJed Brown       requires: !complex !single !__float128 !define(PETSC_USE_64BIT_INDICES)
395*c4762a1bSJed Brown 
396*c4762a1bSJed Brown    test:
397*c4762a1bSJed Brown       suffix: 1
398*c4762a1bSJed Brown       localrunfiles: tomographyData_A_b_xGT
399*c4762a1bSJed Brown       args:  -lambda 1.e-8 -tao_monitor -tao_type nls -tao_nls_pc_type icc
400*c4762a1bSJed Brown 
401*c4762a1bSJed Brown    test:
402*c4762a1bSJed Brown       suffix: 2
403*c4762a1bSJed Brown       localrunfiles: tomographyData_A_b_xGT
404*c4762a1bSJed Brown       args:  -reg 2 -lambda 1.e-8 -tao_admm_dual_update update_basic -tao_admm_regularizer_type regularizer_user -tao_max_it 20 -tao_monitor -tao_admm_tolerance_update_factor 1.e-8  -misfit_tao_nls_pc_type icc -misfit_tao_monitor -reg_tao_monitor
405*c4762a1bSJed Brown 
406*c4762a1bSJed Brown    test:
407*c4762a1bSJed Brown       suffix: 3
408*c4762a1bSJed Brown       localrunfiles: tomographyData_A_b_xGT
409*c4762a1bSJed Brown       args:  -lambda 1.e-8 -tao_admm_dual_update update_basic -tao_admm_regularizer_type regularizer_soft_thresh -tao_max_it 20 -tao_monitor -tao_admm_tolerance_update_factor 1.e-8 -misfit_tao_nls_pc_type icc -misfit_tao_monitor
410*c4762a1bSJed Brown 
411*c4762a1bSJed Brown    test:
412*c4762a1bSJed Brown       suffix: 4
413*c4762a1bSJed Brown       localrunfiles: tomographyData_A_b_xGT
414*c4762a1bSJed Brown       args:  -lambda 1.e-8 -tao_admm_dual_update update_adaptive -tao_admm_regularizer_type regularizer_soft_thresh -tao_max_it 20 -tao_monitor -misfit_tao_monitor -misfit_tao_nls_pc_type icc
415*c4762a1bSJed Brown 
416*c4762a1bSJed Brown    test:
417*c4762a1bSJed Brown       suffix: 5
418*c4762a1bSJed Brown       localrunfiles: tomographyData_A_b_xGT
419*c4762a1bSJed Brown       args:  -reg 2 -lambda 1.e-8 -tao_admm_dual_update update_adaptive -tao_admm_regularizer_type regularizer_user -tao_max_it 20 -tao_monitor -tao_admm_tolerance_update_factor 1.e-8 -misfit_tao_monitor -reg_tao_monitor -misfit_tao_nls_pc_type icc
420*c4762a1bSJed Brown 
421*c4762a1bSJed Brown    test:
422*c4762a1bSJed Brown       suffix: 6
423*c4762a1bSJed Brown       localrunfiles: tomographyData_A_b_xGT
424*c4762a1bSJed Brown       args:  -reg 3 -lambda 1.e-8 -tao_admm_dual_update update_adaptive -tao_admm_regularizer_type regularizer_user -tao_max_it 20 -tao_monitor -tao_admm_tolerance_update_factor 1.e-8 -misfit_tao_monitor -reg_tao_monitor -misfit_tao_nls_pc_type icc
425*c4762a1bSJed Brown 
426*c4762a1bSJed Brown TEST*/
427