xref: /petsc/src/mat/tests/ex123.c (revision ccb4e88a40f0b86eaeca07ff64c64e4de2fae686)
1 static char help[] = "Test MatSetPreallocationCOO and MatSetValuesCOO\n\n";
2 
3 #include <petscmat.h>
4 #define MyMatView(a,b) PetscPrintf(PetscObjectComm((PetscObject)(a)),"LINE %d\n",__LINE__),MatView(a,b);
5 #define MyVecView(a,b) PetscPrintf(PetscObjectComm((PetscObject)(a)),"LINE %d\n",__LINE__),VecView(a,b);
6 int main(int argc,char **args)
7 {
8   Mat            A,At,AAt;
9   Vec            x,y,z;
10   PetscLayout    rmap,cmap;
11   PetscInt       n1 = 11, n2 = 9;
12   PetscInt       i1[] = {   7,  6,  2,  0,  4,  1,  1,  0,  2,  2,  1 };
13   PetscInt       j1[] = {   1,  4,  3,  5,  3,  3,  4,  5,  0,  3,  1 };
14   PetscInt       i2[] = {   7,  6,  2,  0,  4,  1,  1,  2, 1 };
15   PetscInt       j2[] = {   1,  4,  3,  5,  3,  3,  4,  0, 1 };
16   PetscScalar    v1[] = { -1., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10.};
17   PetscScalar    v2[] = {  1.,-1.,-2.,-3.,-4.,-5.,-6.,-7.,-8.,-9.,-10.};
18   PetscInt       N = 6, m = 8, rstart, cstart, i;
19   PetscMPIInt    size;
20   PetscBool      loc = PETSC_FALSE;
21   PetscBool      locdiag = PETSC_TRUE, ismpiaij;
22   PetscErrorCode ierr;
23 
24   ierr = PetscInitialize(&argc,&args,(char*)0,help);if (ierr) return ierr;
25   ierr = PetscOptionsGetBool(NULL,NULL,"-loc",&loc,NULL);CHKERRQ(ierr);
26   ierr = PetscOptionsGetBool(NULL,NULL,"-locdiag",&locdiag,NULL);CHKERRQ(ierr);
27 
28   ierr = MatCreate(PETSC_COMM_WORLD,&A);CHKERRQ(ierr);
29   if (loc) {
30     if (locdiag) {
31       ierr = MatSetSizes(A,m,N,PETSC_DECIDE,PETSC_DECIDE);CHKERRQ(ierr);
32     } else {
33       ierr = MatSetSizes(A,m,m+N,PETSC_DECIDE,PETSC_DECIDE);CHKERRQ(ierr);
34     }
35   } else {
36     ierr = MatSetSizes(A,m,PETSC_DECIDE,PETSC_DECIDE,N);CHKERRQ(ierr);
37   }
38   ierr = MatSetFromOptions(A);CHKERRQ(ierr);
39   ierr = MatGetLayouts(A,&rmap,&cmap);CHKERRQ(ierr);
40   ierr = PetscLayoutSetUp(rmap);CHKERRQ(ierr);
41   ierr = PetscLayoutSetUp(cmap);CHKERRQ(ierr);
42   ierr = MatCreateVecs(A,&x,&y);CHKERRQ(ierr);
43   ierr = MatCreateVecs(A,NULL,&z);CHKERRQ(ierr);
44   ierr = VecSet(x,1.);CHKERRQ(ierr);
45   ierr = VecSet(z,2.);CHKERRQ(ierr);
46   ierr = PetscLayoutGetRange(rmap,&rstart,NULL);CHKERRQ(ierr);
47   ierr = PetscLayoutGetRange(cmap,&cstart,NULL);CHKERRQ(ierr);
48   for (i = 0; i < n1; i++) i1[i] += rstart;
49   for (i = 0; i < n2; i++) i2[i] += rstart;
50   if (loc) {
51     if (locdiag) {
52       for (i = 0; i < n1; i++) j1[i] += cstart;
53       for (i = 0; i < n2; i++) j2[i] += cstart;
54     } else {
55       for (i = 0; i < n1; i++) j1[i] += cstart + m;
56       for (i = 0; i < n2; i++) j2[i] += cstart + m;
57     }
58   }
59 
60   /* test with repeated entries */
61   ierr = MatSetPreallocationCOO(A,n1,i1,j1);CHKERRQ(ierr);
62   ierr = MatSetValuesCOO(A,v1,ADD_VALUES);CHKERRQ(ierr);
63   ierr = MyMatView(A,NULL);CHKERRQ(ierr);
64   ierr = MatMult(A,x,y);CHKERRQ(ierr);
65   ierr = MyVecView(y,NULL);CHKERRQ(ierr);
66   ierr = MatSetValuesCOO(A,v2,ADD_VALUES);CHKERRQ(ierr);
67   ierr = MyMatView(A,NULL);CHKERRQ(ierr);
68   ierr = MatMultAdd(A,x,y,y);CHKERRQ(ierr);
69   ierr = MyVecView(y,NULL);CHKERRQ(ierr);
70   ierr = MatTranspose(A,MAT_INITIAL_MATRIX,&At);CHKERRQ(ierr);
71   ierr = MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr);
72   ierr = MyMatView(AAt,NULL);CHKERRQ(ierr);
73   ierr = MatDestroy(&AAt);CHKERRQ(ierr);
74   ierr = MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr);
75   ierr = MyMatView(AAt,NULL);CHKERRQ(ierr);
76   ierr = MatDestroy(&AAt);CHKERRQ(ierr);
77   ierr = MatDestroy(&At);CHKERRQ(ierr);
78   /* INSERT_VALUES will overwrite matrix entries but
79      still perform the sum of the repeated entries */
80   ierr = MatSetValuesCOO(A,v2,INSERT_VALUES);CHKERRQ(ierr);
81   ierr = MyMatView(A,NULL);CHKERRQ(ierr);
82 
83   /* test with unique entries */
84   ierr = MatSetPreallocationCOO(A,n2,i2,j2);CHKERRQ(ierr);
85   ierr = MatSetValuesCOO(A,v1,ADD_VALUES);CHKERRQ(ierr);
86   ierr = MyMatView(A,NULL);CHKERRQ(ierr);
87   ierr = MatMult(A,x,y);CHKERRQ(ierr);
88   ierr = MyVecView(y,NULL);CHKERRQ(ierr);
89   ierr = MatSetValuesCOO(A,v2,ADD_VALUES);CHKERRQ(ierr);
90   ierr = MyMatView(A,NULL);CHKERRQ(ierr);
91   ierr = MatMultAdd(A,x,y,z);CHKERRQ(ierr);
92   ierr = MyVecView(z,NULL);CHKERRQ(ierr);
93   ierr = MatSetPreallocationCOO(A,n2,i2,j2);CHKERRQ(ierr);
94   ierr = MatSetValuesCOO(A,v1,INSERT_VALUES);CHKERRQ(ierr);
95   ierr = MyMatView(A,NULL);CHKERRQ(ierr);
96   ierr = MatMult(A,x,y);CHKERRQ(ierr);
97   ierr = MyVecView(y,NULL);CHKERRQ(ierr);
98   ierr = MatSetValuesCOO(A,v2,INSERT_VALUES);CHKERRQ(ierr);
99   ierr = MyMatView(A,NULL);CHKERRQ(ierr);
100   ierr = MatMultAdd(A,x,y,z);CHKERRQ(ierr);
101   ierr = MyVecView(z,NULL);CHKERRQ(ierr);
102   ierr = MatTranspose(A,MAT_INITIAL_MATRIX,&At);CHKERRQ(ierr);
103   ierr = MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr);
104   ierr = MyMatView(AAt,NULL);CHKERRQ(ierr);
105   ierr = MatDestroy(&AAt);CHKERRQ(ierr);
106   ierr = MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr);
107   ierr = MyMatView(AAt,NULL);CHKERRQ(ierr);
108   ierr = MatDestroy(&AAt);CHKERRQ(ierr);
109   ierr = MatDestroy(&At);CHKERRQ(ierr);
110 
111   /* test providing diagonal first, the offdiagonal */
112   ierr = MPI_Comm_size(PetscObjectComm((PetscObject)A),&size);CHKERRMPI(ierr);
113   ierr = PetscObjectBaseTypeCompare((PetscObject)A,MATMPIAIJ,&ismpiaij);CHKERRQ(ierr);
114   if (ismpiaij && size > 1) {
115     Mat               lA,lB;
116     const PetscInt    *garray,*iA,*jA,*iB,*jB;
117     const PetscScalar *vA,*vB;
118     PetscScalar       *coo_v;
119     PetscInt          *coo_i,*coo_j;
120     PetscInt          i,j,nA,nB,nnz;
121     PetscBool         flg;
122 
123     ierr = MatMPIAIJGetSeqAIJ(A,&lA,&lB,&garray);CHKERRQ(ierr);
124     ierr = MatSeqAIJGetArrayRead(lA,&vA);CHKERRQ(ierr);
125     ierr = MatSeqAIJGetArrayRead(lB,&vB);CHKERRQ(ierr);
126     ierr = MatGetRowIJ(lA,0,PETSC_FALSE,PETSC_FALSE,&nA,&iA,&jA,&flg);CHKERRQ(ierr);
127     ierr = MatGetRowIJ(lB,0,PETSC_FALSE,PETSC_FALSE,&nB,&iB,&jB,&flg);CHKERRQ(ierr);
128     nnz  = iA[nA] + iB[nB];
129     ierr = PetscMalloc3(nnz,&coo_i,nnz,&coo_j,nnz,&coo_v);CHKERRQ(ierr);
130     nnz  = 0;
131     for (i=0;i<nA;i++) {
132       for (j=iA[i];j<iA[i+1];j++,nnz++) {
133         coo_i[nnz] = i+rstart;
134         coo_j[nnz] = jA[j]+cstart;
135         coo_v[nnz] = vA[j];
136       }
137     }
138     for (i=0;i<nB;i++) {
139       for (j=iB[i];j<iB[i+1];j++,nnz++) {
140         coo_i[nnz] = i+rstart;
141         coo_j[nnz] = garray[jB[j]];
142         coo_v[nnz] = vB[j];
143       }
144     }
145     ierr = MatRestoreRowIJ(lA,0,PETSC_FALSE,PETSC_FALSE,&nA,&iA,&jA,&flg);CHKERRQ(ierr);
146     ierr = MatRestoreRowIJ(lB,0,PETSC_FALSE,PETSC_FALSE,&nB,&iB,&jB,&flg);CHKERRQ(ierr);
147     ierr = MatSeqAIJRestoreArrayRead(lA,&vA);CHKERRQ(ierr);
148     ierr = MatSeqAIJRestoreArrayRead(lB,&vB);CHKERRQ(ierr);
149 
150     ierr = MatSetPreallocationCOO(A,nnz,coo_i,coo_j);CHKERRQ(ierr);
151     ierr = MatSetValuesCOO(A,coo_v,ADD_VALUES);CHKERRQ(ierr);
152     ierr = MyMatView(A,NULL);CHKERRQ(ierr);
153     ierr = MatMult(A,x,y);CHKERRQ(ierr);
154     ierr = MyVecView(y,NULL);CHKERRQ(ierr);
155     ierr = MatSetValuesCOO(A,coo_v,INSERT_VALUES);CHKERRQ(ierr);
156     ierr = MyMatView(A,NULL);CHKERRQ(ierr);
157     ierr = MatMult(A,x,y);CHKERRQ(ierr);
158     ierr = MyVecView(y,NULL);CHKERRQ(ierr);
159     ierr = MatTranspose(A,MAT_INITIAL_MATRIX,&At);CHKERRQ(ierr);
160     ierr = MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr);
161     ierr = MyMatView(AAt,NULL);CHKERRQ(ierr);
162     ierr = MatDestroy(&AAt);CHKERRQ(ierr);
163     ierr = MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr);
164     ierr = MyMatView(AAt,NULL);CHKERRQ(ierr);
165     ierr = MatDestroy(&AAt);CHKERRQ(ierr);
166     ierr = MatDestroy(&At);CHKERRQ(ierr);
167 
168     ierr = PetscFree3(coo_i,coo_j,coo_v);CHKERRQ(ierr);
169   }
170   ierr = VecDestroy(&z);CHKERRQ(ierr);
171   ierr = VecDestroy(&x);CHKERRQ(ierr);
172   ierr = VecDestroy(&y);CHKERRQ(ierr);
173   ierr = MatDestroy(&A);CHKERRQ(ierr);
174   ierr = PetscFinalize();
175   return ierr;
176 }
177 
178 /*TEST
179 
180    test:
181      suffix: 1
182      filter: grep -v type
183      diff_args: -j
184      args: -mat_type {{seqaij mpiaij}}
185 
186    test:
187      requires: cuda
188      suffix: 1_cuda
189      filter: grep -v type
190      diff_args: -j
191      args: -mat_type {{seqaijcusparse mpiaijcusparse}}
192      output_file: output/ex123_1.out
193 
194    test:
195      suffix: 2
196      nsize: 7
197      filter: grep -v type
198      diff_args: -j
199      args: -mat_type mpiaij
200 
201    test:
202      requires: cuda
203      suffix: 2_cuda
204      nsize: 7
205      filter: grep -v type
206      diff_args: -j
207      args: -mat_type mpiaijcusparse
208      output_file: output/ex123_2.out
209 
210    test:
211      suffix: 3
212      nsize: 3
213      filter: grep -v type
214      diff_args: -j
215      args: -mat_type mpiaij -loc
216 
217    test:
218      requires: cuda
219      suffix: 3_cuda
220      nsize: 3
221      filter: grep -v type
222      diff_args: -j
223      args: -mat_type mpiaijcusparse -loc
224      output_file: output/ex123_3.out
225 
226    test:
227      suffix: 4
228      nsize: 4
229      filter: grep -v type
230      diff_args: -j
231      args: -mat_type mpiaij -loc -locdiag 0
232 
233    test:
234      requires: cuda
235      suffix: 4_cuda
236      nsize: 4
237      filter: grep -v type
238      diff_args: -j
239      args: -mat_type mpiaijcusparse -loc -locdiag 0
240      output_file: output/ex123_4.out
241 
242 TEST*/
243