xref: /petsc/src/tao/matrix/adamat.c (revision 40badf4fbc550ac1f60bd080eaff6de6d55b946d)
1 #include <petscmat.h>              /*I  "mat.h"  I*/
2 
3 PETSC_INTERN PetscErrorCode MatCreateADA(Mat,Vec, Vec, Mat*);
4 
5 typedef struct {
6   Mat      A;
7   Vec      D1;
8   Vec      D2;
9   Vec      W;
10   Vec      W2;
11   Vec      ADADiag;
12   PetscInt GotDiag;
13 } _p_TaoMatADACtx;
14 typedef  _p_TaoMatADACtx* TaoMatADACtx;
15 
16 static PetscErrorCode MatMult_ADA(Mat mat,Vec a,Vec y)
17 {
18   TaoMatADACtx   ctx;
19   PetscReal      one = 1.0;
20 
21   PetscFunctionBegin;
22   CHKERRQ(MatShellGetContext(mat,&ctx));
23   CHKERRQ(MatMult(ctx->A,a,ctx->W));
24   if (ctx->D1) {
25     CHKERRQ(VecPointwiseMult(ctx->W,ctx->D1,ctx->W));
26   }
27   CHKERRQ(MatMultTranspose(ctx->A,ctx->W,y));
28   if (ctx->D2) {
29     CHKERRQ(VecPointwiseMult(ctx->W2, ctx->D2, a));
30     CHKERRQ(VecAXPY(y, one, ctx->W2));
31   }
32   PetscFunctionReturn(0);
33 }
34 
35 static PetscErrorCode MatMultTranspose_ADA(Mat mat,Vec a,Vec y)
36 {
37   PetscFunctionBegin;
38   CHKERRQ(MatMult_ADA(mat,a,y));
39   PetscFunctionReturn(0);
40 }
41 
42 static PetscErrorCode MatDiagonalSet_ADA(Mat M,Vec D, InsertMode mode)
43 {
44   TaoMatADACtx   ctx;
45   PetscReal      zero=0.0,one = 1.0;
46 
47   PetscFunctionBegin;
48   PetscCheck(mode != INSERT_VALUES,PetscObjectComm((PetscObject)M),PETSC_ERR_SUP,"Cannot insert diagonal entries of this matrix type, can only add");
49   CHKERRQ(MatShellGetContext(M,&ctx));
50   if (!ctx->D2) {
51     CHKERRQ(VecDuplicate(D,&ctx->D2));
52     CHKERRQ(VecSet(ctx->D2, zero));
53   }
54   CHKERRQ(VecAXPY(ctx->D2, one, D));
55   PetscFunctionReturn(0);
56 }
57 
58 static PetscErrorCode MatDestroy_ADA(Mat mat)
59 {
60   TaoMatADACtx   ctx;
61 
62   PetscFunctionBegin;
63   CHKERRQ(MatShellGetContext(mat,&ctx));
64   CHKERRQ(VecDestroy(&ctx->W));
65   CHKERRQ(VecDestroy(&ctx->W2));
66   CHKERRQ(VecDestroy(&ctx->ADADiag));
67   CHKERRQ(MatDestroy(&ctx->A));
68   CHKERRQ(VecDestroy(&ctx->D1));
69   CHKERRQ(VecDestroy(&ctx->D2));
70   CHKERRQ(PetscFree(ctx));
71   PetscFunctionReturn(0);
72 }
73 
74 static PetscErrorCode MatView_ADA(Mat mat,PetscViewer viewer)
75 {
76   PetscFunctionBegin;
77   PetscFunctionReturn(0);
78 }
79 
80 static PetscErrorCode MatShift_ADA(Mat Y, PetscReal a)
81 {
82   TaoMatADACtx   ctx;
83 
84   PetscFunctionBegin;
85   CHKERRQ(MatShellGetContext(Y,&ctx));
86   CHKERRQ(VecShift(ctx->D2,a));
87   PetscFunctionReturn(0);
88 }
89 
90 static PetscErrorCode MatDuplicate_ADA(Mat mat,MatDuplicateOption op,Mat *M)
91 {
92   TaoMatADACtx      ctx;
93   Mat               A2;
94   Vec               D1b=NULL,D2b;
95 
96   PetscFunctionBegin;
97   CHKERRQ(MatShellGetContext(mat,&ctx));
98   CHKERRQ(MatDuplicate(ctx->A,op,&A2));
99   if (ctx->D1) {
100     CHKERRQ(VecDuplicate(ctx->D1,&D1b));
101     CHKERRQ(VecCopy(ctx->D1,D1b));
102   }
103   CHKERRQ(VecDuplicate(ctx->D2,&D2b));
104   CHKERRQ(VecCopy(ctx->D2,D2b));
105   CHKERRQ(MatCreateADA(A2,D1b,D2b,M));
106   if (ctx->D1) {
107     CHKERRQ(PetscObjectDereference((PetscObject)D1b));
108   }
109   CHKERRQ(PetscObjectDereference((PetscObject)D2b));
110   CHKERRQ(PetscObjectDereference((PetscObject)A2));
111   PetscFunctionReturn(0);
112 }
113 
114 static PetscErrorCode MatEqual_ADA(Mat A,Mat B,PetscBool *flg)
115 {
116   TaoMatADACtx   ctx1,ctx2;
117 
118   PetscFunctionBegin;
119   CHKERRQ(MatShellGetContext(A,&ctx1));
120   CHKERRQ(MatShellGetContext(B,&ctx2));
121   CHKERRQ(VecEqual(ctx1->D2,ctx2->D2,flg));
122   if (*flg==PETSC_TRUE) {
123     CHKERRQ(VecEqual(ctx1->D1,ctx2->D1,flg));
124   }
125   if (*flg==PETSC_TRUE) {
126     CHKERRQ(MatEqual(ctx1->A,ctx2->A,flg));
127   }
128   PetscFunctionReturn(0);
129 }
130 
131 static PetscErrorCode MatScale_ADA(Mat mat, PetscReal a)
132 {
133   TaoMatADACtx   ctx;
134 
135   PetscFunctionBegin;
136   CHKERRQ(MatShellGetContext(mat,&ctx));
137   CHKERRQ(VecScale(ctx->D1,a));
138   if (ctx->D2) {
139     CHKERRQ(VecScale(ctx->D2,a));
140   }
141   PetscFunctionReturn(0);
142 }
143 
144 static PetscErrorCode MatTranspose_ADA(Mat mat,MatReuse reuse,Mat *B)
145 {
146   TaoMatADACtx   ctx;
147 
148   PetscFunctionBegin;
149   CHKERRQ(MatShellGetContext(mat,&ctx));
150   if (reuse == MAT_INITIAL_MATRIX) {
151     CHKERRQ(MatDuplicate(mat,MAT_COPY_VALUES,B));
152   } else if (reuse == MAT_REUSE_MATRIX) {
153     CHKERRQ(MatCopy(mat,*B,SAME_NONZERO_PATTERN));
154   } else SETERRQ(PetscObjectComm((PetscObject)mat),PETSC_ERR_SUP,"Does not support inplace transpose");
155   PetscFunctionReturn(0);
156 }
157 
158 static PetscErrorCode MatADAComputeDiagonal(Mat mat)
159 {
160   PetscInt       i,m,n,low,high;
161   PetscScalar    *dtemp,*dptr;
162   TaoMatADACtx   ctx;
163 
164   PetscFunctionBegin;
165   CHKERRQ(MatShellGetContext(mat,&ctx));
166   CHKERRQ(MatGetOwnershipRange(mat, &low, &high));
167   CHKERRQ(MatGetSize(mat,&m,&n));
168 
169   CHKERRQ(PetscMalloc1(n,&dtemp));
170   for (i=0; i<n; i++) {
171     CHKERRQ(MatGetColumnVector(ctx->A, ctx->W, i));
172     CHKERRQ(VecPointwiseMult(ctx->W,ctx->W,ctx->W));
173     CHKERRQ(VecDotBegin(ctx->D1, ctx->W,dtemp+i));
174   }
175   for (i=0; i<n; i++) {
176     CHKERRQ(VecDotEnd(ctx->D1, ctx->W,dtemp+i));
177   }
178 
179   CHKERRQ(VecGetArray(ctx->ADADiag,&dptr));
180   for (i=low; i<high; i++) {
181     dptr[i-low]= dtemp[i];
182   }
183   CHKERRQ(VecRestoreArray(ctx->ADADiag,&dptr));
184   CHKERRQ(PetscFree(dtemp));
185   PetscFunctionReturn(0);
186 }
187 
188 static PetscErrorCode MatGetDiagonal_ADA(Mat mat,Vec v)
189 {
190   PetscReal       one=1.0;
191   TaoMatADACtx    ctx;
192 
193   PetscFunctionBegin;
194   CHKERRQ(MatShellGetContext(mat,&ctx));
195   CHKERRQ(MatADAComputeDiagonal(mat));
196   CHKERRQ(VecCopy(ctx->ADADiag,v));
197   if (ctx->D2) {
198     CHKERRQ(VecAXPY(v, one, ctx->D2));
199   }
200   PetscFunctionReturn(0);
201 }
202 
203 static PetscErrorCode MatCreateSubMatrix_ADA(Mat mat,IS isrow,IS iscol,MatReuse cll, Mat *newmat)
204 {
205   PetscInt          low,high;
206   IS                ISrow;
207   Vec               D1,D2;
208   Mat               Atemp;
209   TaoMatADACtx      ctx;
210   PetscBool         isequal;
211 
212   PetscFunctionBegin;
213   CHKERRQ(ISEqual(isrow,iscol,&isequal));
214   PetscCheck(isequal,PETSC_COMM_SELF,PETSC_ERR_SUP,"Only for identical column and row indices");
215   CHKERRQ(MatShellGetContext(mat,&ctx));
216 
217   CHKERRQ(MatGetOwnershipRange(ctx->A,&low,&high));
218   CHKERRQ(ISCreateStride(PetscObjectComm((PetscObject)mat),high-low,low,1,&ISrow));
219   CHKERRQ(MatCreateSubMatrix(ctx->A,ISrow,iscol,cll,&Atemp));
220   CHKERRQ(ISDestroy(&ISrow));
221 
222   if (ctx->D1) {
223     CHKERRQ(VecDuplicate(ctx->D1,&D1));
224     CHKERRQ(VecCopy(ctx->D1,D1));
225   } else {
226     D1 = NULL;
227   }
228 
229   if (ctx->D2) {
230     Vec D2sub;
231 
232     CHKERRQ(VecGetSubVector(ctx->D2,isrow,&D2sub));
233     CHKERRQ(VecDuplicate(D2sub,&D2));
234     CHKERRQ(VecCopy(D2sub,D2));
235     CHKERRQ(VecRestoreSubVector(ctx->D2,isrow,&D2sub));
236   } else {
237     D2 = NULL;
238   }
239 
240   CHKERRQ(MatCreateADA(Atemp,D1,D2,newmat));
241   CHKERRQ(MatShellGetContext(*newmat,&ctx));
242   CHKERRQ(PetscObjectDereference((PetscObject)Atemp));
243   if (ctx->D1) {
244     CHKERRQ(PetscObjectDereference((PetscObject)D1));
245   }
246   if (ctx->D2) {
247     CHKERRQ(PetscObjectDereference((PetscObject)D2));
248   }
249   PetscFunctionReturn(0);
250 }
251 
252 static PetscErrorCode MatCreateSubMatrices_ADA(Mat A,PetscInt n, IS *irow,IS *icol,MatReuse scall,Mat **B)
253 {
254   PetscInt       i;
255 
256   PetscFunctionBegin;
257   if (scall == MAT_INITIAL_MATRIX) {
258     CHKERRQ(PetscCalloc1(n+1,B));
259   }
260   for (i=0; i<n; i++) {
261     CHKERRQ(MatCreateSubMatrix_ADA(A,irow[i],icol[i],scall,&(*B)[i]));
262   }
263   PetscFunctionReturn(0);
264 }
265 
266 static PetscErrorCode MatGetColumnVector_ADA(Mat mat,Vec Y, PetscInt col)
267 {
268   PetscInt       low,high;
269   PetscScalar    zero=0.0,one=1.0;
270 
271   PetscFunctionBegin;
272   CHKERRQ(VecSet(Y, zero));
273   CHKERRQ(VecGetOwnershipRange(Y,&low,&high));
274   if (col>=low && col<high) {
275     CHKERRQ(VecSetValue(Y,col,one,INSERT_VALUES));
276   }
277   CHKERRQ(VecAssemblyBegin(Y));
278   CHKERRQ(VecAssemblyEnd(Y));
279   CHKERRQ(MatMult_ADA(mat,Y,Y));
280   PetscFunctionReturn(0);
281 }
282 
283 PETSC_INTERN PetscErrorCode MatConvert_ADA(Mat mat,MatType newtype,Mat *NewMat)
284 {
285   PetscMPIInt    size;
286   PetscBool      sametype, issame, isdense, isseqdense;
287   TaoMatADACtx   ctx;
288 
289   PetscFunctionBegin;
290   CHKERRQ(MatShellGetContext(mat,&ctx));
291   CHKERRMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat),&size));
292 
293   CHKERRQ(PetscObjectTypeCompare((PetscObject)mat,newtype,&sametype));
294   CHKERRQ(PetscObjectTypeCompare((PetscObject)mat,MATSAME,&issame));
295   CHKERRQ(PetscObjectTypeCompare((PetscObject)mat,MATMPIDENSE,&isdense));
296   CHKERRQ(PetscObjectTypeCompare((PetscObject)mat,MATSEQDENSE,&isseqdense));
297 
298   if (sametype || issame) {
299     CHKERRQ(MatDuplicate(mat,MAT_COPY_VALUES,NewMat));
300   } else if (isdense) {
301     PetscInt          i,j,low,high,m,n,M,N;
302     const PetscScalar *dptr;
303     Vec               X;
304 
305     CHKERRQ(VecDuplicate(ctx->D2,&X));
306     CHKERRQ(MatGetSize(mat,&M,&N));
307     CHKERRQ(MatGetLocalSize(mat,&m,&n));
308     CHKERRQ(MatCreateDense(PetscObjectComm((PetscObject)mat),m,m,N,N,NULL,NewMat));
309     CHKERRQ(MatGetOwnershipRange(*NewMat,&low,&high));
310     for (i=0;i<M;i++) {
311       CHKERRQ(MatGetColumnVector_ADA(mat,X,i));
312       CHKERRQ(VecGetArrayRead(X,&dptr));
313       for (j=0; j<high-low; j++) {
314         CHKERRQ(MatSetValue(*NewMat,low+j,i,dptr[j],INSERT_VALUES));
315       }
316       CHKERRQ(VecRestoreArrayRead(X,&dptr));
317     }
318     CHKERRQ(MatAssemblyBegin(*NewMat,MAT_FINAL_ASSEMBLY));
319     CHKERRQ(MatAssemblyEnd(*NewMat,MAT_FINAL_ASSEMBLY));
320     CHKERRQ(VecDestroy(&X));
321   } else if (isseqdense && size==1) {
322     PetscInt          i,j,low,high,m,n,M,N;
323     const PetscScalar *dptr;
324     Vec               X;
325 
326     CHKERRQ(VecDuplicate(ctx->D2,&X));
327     CHKERRQ(MatGetSize(mat,&M,&N));
328     CHKERRQ(MatGetLocalSize(mat,&m,&n));
329     CHKERRQ(MatCreateSeqDense(PetscObjectComm((PetscObject)mat),N,N,NULL,NewMat));
330     CHKERRQ(MatGetOwnershipRange(*NewMat,&low,&high));
331     for (i=0;i<M;i++) {
332       CHKERRQ(MatGetColumnVector_ADA(mat,X,i));
333       CHKERRQ(VecGetArrayRead(X,&dptr));
334       for (j=0; j<high-low; j++) {
335         CHKERRQ(MatSetValue(*NewMat,low+j,i,dptr[j],INSERT_VALUES));
336       }
337       CHKERRQ(VecRestoreArrayRead(X,&dptr));
338     }
339     CHKERRQ(MatAssemblyBegin(*NewMat,MAT_FINAL_ASSEMBLY));
340     CHKERRQ(MatAssemblyEnd(*NewMat,MAT_FINAL_ASSEMBLY));
341     CHKERRQ(VecDestroy(&X));
342   } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_PLIB,"No support to convert objects to that type");
343   PetscFunctionReturn(0);
344 }
345 
346 static PetscErrorCode MatNorm_ADA(Mat mat,NormType type,PetscReal *norm)
347 {
348   TaoMatADACtx   ctx;
349 
350   PetscFunctionBegin;
351   CHKERRQ(MatShellGetContext(mat,&ctx));
352   if (type == NORM_FROBENIUS) {
353     *norm = 1.0;
354   } else if (type == NORM_1 || type == NORM_INFINITY) {
355     *norm = 1.0;
356   } else SETERRQ(PETSC_COMM_SELF,PETSC_ERR_SUP,"No two norm");
357   PetscFunctionReturn(0);
358 }
359 
360 /*@C
361    MatCreateADA - Creates a matrix M=A^T D1 A + D2 where D1, D2 are diagonal
362 
363    Collective on matrix
364 
365    Input Parameters:
366 +  mat - matrix of arbitrary type
367 .  d1 - A vector defining a diagonal matrix
368 -  d2 - A vector defining a diagonal matrix
369 
370    Output Parameters:
371 .  J - New matrix whose operations are defined in terms of mat, D1, and D2.
372 
373    Notes:
374    The user provides the input data and is responsible for destroying
375    this data after matrix J has been destroyed.
376 
377    Level: developer
378 
379 .seealso: MatCreate()
380 @*/
381 PetscErrorCode MatCreateADA(Mat mat,Vec d1, Vec d2, Mat *J)
382 {
383   MPI_Comm       comm = PetscObjectComm((PetscObject)mat);
384   TaoMatADACtx   ctx;
385   PetscInt       nloc,n;
386 
387   PetscFunctionBegin;
388   CHKERRQ(PetscNew(&ctx));
389   ctx->A=mat;
390   ctx->D1=d1;
391   ctx->D2=d2;
392   if (d1) {
393     CHKERRQ(VecDuplicate(d1,&ctx->W));
394     CHKERRQ(PetscObjectReference((PetscObject)d1));
395   } else {
396     ctx->W = NULL;
397   }
398   if (d2) {
399     CHKERRQ(VecDuplicate(d2,&ctx->W2));
400     CHKERRQ(VecDuplicate(d2,&ctx->ADADiag));
401     CHKERRQ(PetscObjectReference((PetscObject)d2));
402   } else {
403     ctx->W2      = NULL;
404     ctx->ADADiag = NULL;
405   }
406 
407   ctx->GotDiag = 0;
408   CHKERRQ(PetscObjectReference((PetscObject)mat));
409 
410   CHKERRQ(VecGetLocalSize(d2,&nloc));
411   CHKERRQ(VecGetSize(d2,&n));
412 
413   CHKERRQ(MatCreateShell(comm,nloc,nloc,n,n,ctx,J));
414   CHKERRQ(MatShellSetManageScalingShifts(*J));
415   CHKERRQ(MatShellSetOperation(*J,MATOP_MULT,(void(*)(void))MatMult_ADA));
416   CHKERRQ(MatShellSetOperation(*J,MATOP_DESTROY,(void(*)(void))MatDestroy_ADA));
417   CHKERRQ(MatShellSetOperation(*J,MATOP_VIEW,(void(*)(void))MatView_ADA));
418   CHKERRQ(MatShellSetOperation(*J,MATOP_MULT_TRANSPOSE,(void(*)(void))MatMultTranspose_ADA));
419   CHKERRQ(MatShellSetOperation(*J,MATOP_DIAGONAL_SET,(void(*)(void))MatDiagonalSet_ADA));
420   CHKERRQ(MatShellSetOperation(*J,MATOP_SHIFT,(void(*)(void))MatShift_ADA));
421   CHKERRQ(MatShellSetOperation(*J,MATOP_EQUAL,(void(*)(void))MatEqual_ADA));
422   CHKERRQ(MatShellSetOperation(*J,MATOP_SCALE,(void(*)(void))MatScale_ADA));
423   CHKERRQ(MatShellSetOperation(*J,MATOP_TRANSPOSE,(void(*)(void))MatTranspose_ADA));
424   CHKERRQ(MatShellSetOperation(*J,MATOP_GET_DIAGONAL,(void(*)(void))MatGetDiagonal_ADA));
425   CHKERRQ(MatShellSetOperation(*J,MATOP_CREATE_SUBMATRICES,(void(*)(void))MatCreateSubMatrices_ADA));
426   CHKERRQ(MatShellSetOperation(*J,MATOP_NORM,(void(*)(void))MatNorm_ADA));
427   CHKERRQ(MatShellSetOperation(*J,MATOP_DUPLICATE,(void(*)(void))MatDuplicate_ADA));
428   CHKERRQ(MatShellSetOperation(*J,MATOP_CREATE_SUBMATRIX,(void(*)(void))MatCreateSubMatrix_ADA));
429 
430   CHKERRQ(PetscLogObjectParent((PetscObject)(*J),(PetscObject)ctx->W));
431   CHKERRQ(PetscLogObjectParent((PetscObject)mat,(PetscObject)(*J)));
432 
433   CHKERRQ(MatSetOption(*J,MAT_SYMMETRIC,PETSC_TRUE));
434   PetscFunctionReturn(0);
435 }
436