static char help[] = "Test MatSetPreallocationCOO and MatSetValuesCOO\n\n"; #include #define MyMatView(a,b) (PetscPrintf(PetscObjectComm((PetscObject)(a)),"LINE %d\n",__LINE__),MatView(a,b)); #define MyVecView(a,b) (PetscPrintf(PetscObjectComm((PetscObject)(a)),"LINE %d\n",__LINE__),VecView(a,b)); int main(int argc,char **args) { Mat A,At,AAt; Vec x,y,z; ISLocalToGlobalMapping rl2g,cl2g; IS is; PetscLayout rmap,cmap; PetscInt *it,*jt; PetscInt n1 = 11, n2 = 9; PetscInt i1[] = { 7, 6, 2, 0, 4, 1, 1, 0, 2, 2, 1 , -1, -1}; PetscInt j1[] = { 1, 4, 3, 5, 3, 3, 4, 5, 0, 3, 1 , -1, -1}; PetscInt i2[] = { 7, 6, 2, 0, 4, 1, 1, 2, 1, -1, -1}; PetscInt j2[] = { 1, 4, 3, 5, 3, 3, 4, 0, 1, -1, -1}; PetscScalar v1[] = { -1., 1., 2., 3., 4., 5., 6., 7., 8., 9., 10., PETSC_MAX_REAL, PETSC_MAX_REAL}; PetscScalar v2[] = { 1.,-1.,-2.,-3.,-4.,-5.,-6.,-7.,-8.,-9.,-10., PETSC_MAX_REAL, PETSC_MAX_REAL}; PetscInt N = 6, m = 8, M, rstart, cstart, i; PetscMPIInt size; PetscBool loc = PETSC_FALSE; PetscBool locdiag = PETSC_TRUE; PetscBool localapi = PETSC_FALSE; PetscBool neg = PETSC_FALSE; PetscBool ismatis, ismpiaij; PetscErrorCode ierr; ierr = PetscInitialize(&argc,&args,(char*)0,help);if (ierr) return ierr; ierr = PetscOptionsGetBool(NULL,NULL,"-neg",&neg,NULL);CHKERRQ(ierr); ierr = PetscOptionsGetBool(NULL,NULL,"-loc",&loc,NULL);CHKERRQ(ierr); ierr = PetscOptionsGetBool(NULL,NULL,"-locdiag",&locdiag,NULL);CHKERRQ(ierr); ierr = PetscOptionsGetBool(NULL,NULL,"-localapi",&localapi,NULL);CHKERRQ(ierr); ierr = MatCreate(PETSC_COMM_WORLD,&A);CHKERRQ(ierr); if (loc) { if (locdiag) { ierr = MatSetSizes(A,m,N,PETSC_DECIDE,PETSC_DECIDE);CHKERRQ(ierr); } else { ierr = MatSetSizes(A,m,m+N,PETSC_DECIDE,PETSC_DECIDE);CHKERRQ(ierr); } } else { ierr = MatSetSizes(A,m,PETSC_DECIDE,PETSC_DECIDE,N);CHKERRQ(ierr); } ierr = MatSetFromOptions(A);CHKERRQ(ierr); ierr = MatGetLayouts(A,&rmap,&cmap);CHKERRQ(ierr); ierr = PetscLayoutSetUp(rmap);CHKERRQ(ierr); ierr = PetscLayoutSetUp(cmap);CHKERRQ(ierr); ierr = PetscLayoutGetRange(rmap,&rstart,NULL);CHKERRQ(ierr); ierr = PetscLayoutGetRange(cmap,&cstart,NULL);CHKERRQ(ierr); ierr = PetscLayoutGetSize(rmap,&M);CHKERRQ(ierr); ierr = PetscLayoutGetSize(cmap,&N);CHKERRQ(ierr); ierr = PetscObjectTypeCompare((PetscObject)A,MATIS,&ismatis);CHKERRQ(ierr); /* create fake l2g maps to test the local API */ ierr = ISCreateStride(PETSC_COMM_WORLD,M-rstart,rstart,1,&is);CHKERRQ(ierr); ierr = ISLocalToGlobalMappingCreateIS(is,&rl2g);CHKERRQ(ierr); ierr = ISDestroy(&is);CHKERRQ(ierr); ierr = ISCreateStride(PETSC_COMM_WORLD,N,0,1,&is);CHKERRQ(ierr); ierr = ISLocalToGlobalMappingCreateIS(is,&cl2g);CHKERRQ(ierr); ierr = ISDestroy(&is);CHKERRQ(ierr); ierr = MatSetLocalToGlobalMapping(A,rl2g,cl2g);CHKERRQ(ierr); ierr = ISLocalToGlobalMappingDestroy(&rl2g);CHKERRQ(ierr); ierr = ISLocalToGlobalMappingDestroy(&cl2g);CHKERRQ(ierr); ierr = MatCreateVecs(A,&x,&y);CHKERRQ(ierr); ierr = MatCreateVecs(A,NULL,&z);CHKERRQ(ierr); ierr = VecSet(x,1.);CHKERRQ(ierr); ierr = VecSet(z,2.);CHKERRQ(ierr); if (!localapi) for (i = 0; i < n1; i++) i1[i] += rstart; if (!localapi) for (i = 0; i < n2; i++) i2[i] += rstart; if (loc) { if (locdiag) { for (i = 0; i < n1; i++) j1[i] += cstart; for (i = 0; i < n2; i++) j2[i] += cstart; } else { for (i = 0; i < n1; i++) j1[i] += cstart + m; for (i = 0; i < n2; i++) j2[i] += cstart + m; } } if (neg) { n1 += 2; n2 += 2; } /* MatSetPreallocationCOOLocal maps the indices! */ ierr = PetscMalloc2(PetscMax(n1,n2),&it,PetscMax(n1,n2),&jt);CHKERRQ(ierr); /* test with repeated entries */ if (!localapi) { ierr = MatSetPreallocationCOO(A,n1,i1,j1);CHKERRQ(ierr); } else { ierr = PetscArraycpy(it,i1,n1);CHKERRQ(ierr); ierr = PetscArraycpy(jt,j1,n1);CHKERRQ(ierr); ierr = MatSetPreallocationCOOLocal(A,n1,it,jt);CHKERRQ(ierr); } ierr = MatSetValuesCOO(A,v1,ADD_VALUES);CHKERRQ(ierr); ierr = MatMult(A,x,y);CHKERRQ(ierr); ierr = MyMatView(A,NULL);CHKERRQ(ierr); ierr = MyVecView(y,NULL);CHKERRQ(ierr); ierr = MatSetValuesCOO(A,v2,ADD_VALUES);CHKERRQ(ierr); ierr = MatMultAdd(A,x,y,y);CHKERRQ(ierr); ierr = MyMatView(A,NULL);CHKERRQ(ierr); ierr = MyVecView(y,NULL);CHKERRQ(ierr); ierr = MatTranspose(A,MAT_INITIAL_MATRIX,&At);CHKERRQ(ierr); if (!ismatis) { ierr = MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); ierr = MatDestroy(&AAt);CHKERRQ(ierr); ierr = MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); ierr = MatDestroy(&AAt);CHKERRQ(ierr); } ierr = MatDestroy(&At);CHKERRQ(ierr); /* INSERT_VALUES will overwrite matrix entries but still perform the sum of the repeated entries */ ierr = MatSetValuesCOO(A,v2,INSERT_VALUES);CHKERRQ(ierr); ierr = MyMatView(A,NULL);CHKERRQ(ierr); /* test with unique entries */ if (!localapi) { ierr = MatSetPreallocationCOO(A,n2,i2,j2);CHKERRQ(ierr); } else { ierr = PetscArraycpy(it,i2,n2);CHKERRQ(ierr); ierr = PetscArraycpy(jt,j2,n2);CHKERRQ(ierr); ierr = MatSetPreallocationCOOLocal(A,n2,it,jt);CHKERRQ(ierr); } ierr = MatSetValuesCOO(A,v1,ADD_VALUES);CHKERRQ(ierr); ierr = MatMult(A,x,y);CHKERRQ(ierr); ierr = MyMatView(A,NULL);CHKERRQ(ierr); ierr = MyVecView(y,NULL);CHKERRQ(ierr); ierr = MatSetValuesCOO(A,v2,ADD_VALUES);CHKERRQ(ierr); ierr = MatMultAdd(A,x,y,z);CHKERRQ(ierr); ierr = MyMatView(A,NULL);CHKERRQ(ierr); ierr = MyVecView(z,NULL);CHKERRQ(ierr); if (!localapi) { ierr = MatSetPreallocationCOO(A,n2,i2,j2);CHKERRQ(ierr); } else { ierr = PetscArraycpy(it,i2,n2);CHKERRQ(ierr); ierr = PetscArraycpy(jt,j2,n2);CHKERRQ(ierr); ierr = MatSetPreallocationCOOLocal(A,n2,it,jt);CHKERRQ(ierr); } ierr = MatSetValuesCOO(A,v1,INSERT_VALUES);CHKERRQ(ierr); ierr = MatMult(A,x,y);CHKERRQ(ierr); ierr = MyMatView(A,NULL);CHKERRQ(ierr); ierr = MyVecView(y,NULL);CHKERRQ(ierr); ierr = MatSetValuesCOO(A,v2,INSERT_VALUES);CHKERRQ(ierr); ierr = MatMultAdd(A,x,y,z);CHKERRQ(ierr); ierr = MyMatView(A,NULL);CHKERRQ(ierr); ierr = MyVecView(z,NULL);CHKERRQ(ierr); ierr = MatTranspose(A,MAT_INITIAL_MATRIX,&At);CHKERRQ(ierr); if (!ismatis) { ierr = MatMatMult(A,At,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); ierr = MatDestroy(&AAt);CHKERRQ(ierr); ierr = MatMatMult(At,A,MAT_INITIAL_MATRIX,PETSC_DEFAULT,&AAt);CHKERRQ(ierr); ierr = MyMatView(AAt,NULL);CHKERRQ(ierr); ierr = MatDestroy(&AAt);CHKERRQ(ierr); } ierr = MatDestroy(&At);CHKERRQ(ierr); /* test providing diagonal first, then offdiagonal */ ierr = MPI_Comm_size(PetscObjectComm((PetscObject)A),&size);CHKERRMPI(ierr); ierr = PetscObjectBaseTypeCompare((PetscObject)A,MATMPIAIJ,&ismpiaij);CHKERRQ(ierr); if (ismpiaij && size > 1) { Mat lA,lB; const PetscInt *garray,*iA,*jA,*iB,*jB; const PetscScalar *vA,*vB; PetscScalar *coo_v; PetscInt *coo_i,*coo_j; PetscInt i,j,nA,nB,nnz; PetscBool flg; ierr = MatMPIAIJGetSeqAIJ(A,&lA,&lB,&garray);CHKERRQ(ierr); ierr = MatSeqAIJGetArrayRead(lA,&vA);CHKERRQ(ierr); ierr = MatSeqAIJGetArrayRead(lB,&vB);CHKERRQ(ierr); ierr = MatGetRowIJ(lA,0,PETSC_FALSE,PETSC_FALSE,&nA,&iA,&jA,&flg);CHKERRQ(ierr); ierr = MatGetRowIJ(lB,0,PETSC_FALSE,PETSC_FALSE,&nB,&iB,&jB,&flg);CHKERRQ(ierr); nnz = iA[nA] + iB[nB]; ierr = PetscMalloc3(nnz,&coo_i,nnz,&coo_j,nnz,&coo_v);CHKERRQ(ierr); nnz = 0; for (i=0;i