2: static char help[] = "Tests MatShift(), MatScale(), and MatDiagonalScale() for SHELL and NEST matrices\n\n";

  4: #include <petscmat.h>

  6: typedef struct _n_User *User;
  7: struct _n_User {
  8:   Mat B;
  9: };

 13: static PetscErrorCode MatMult_User(Mat A,Vec X,Vec Y)
 14: {
 15:   User           user;

 19:   MatShellGetContext(A,&user);
 20:   MatMult(user->B,X,Y);
 21:   return(0);
 22: }

 26: static PetscErrorCode MatMultTranspose_User(Mat A,Vec X,Vec Y)
 27: {
 28:   User           user;

 32:   MatShellGetContext(A,&user);
 33:   MatMultTranspose(user->B,X,Y);
 34:   return(0);
 35: }

 39: static PetscErrorCode MatGetDiagonal_User(Mat A,Vec X)
 40: {
 41:   User           user;

 45:   MatShellGetContext(A,&user);
 46:   MatGetDiagonal(user->B,X);
 47:   return(0);
 48: }

 52: static PetscErrorCode TestMatrix(Mat A,Vec X,Vec Y,Vec Z)
 53: {
 55:   Vec            W1,W2;
 56:   Mat            E;
 57:   const char     *mattypename;
 58:   PetscViewer    viewer = PETSC_VIEWER_STDOUT_WORLD;

 61:   VecDuplicate(X,&W1);
 62:   VecDuplicate(X,&W2);
 63:   MatScale(A,31);
 64:   MatShift(A,37);
 65:   MatDiagonalScale(A,X,Y);
 66:   MatScale(A,41);
 67:   MatDiagonalScale(A,Y,Z);
 68:   MatComputeExplicitOperator(A,&E);

 70:   PetscObjectGetType((PetscObject)A,&mattypename);
 71:   PetscViewerASCIIPrintf(viewer,"Matrix of type: %s\n",mattypename);
 72:   MatView(E,viewer);
 73:   MatMult(A,Z,W1);
 74:   MatMultTranspose(A,W1,W2);
 75:   VecView(W2,viewer);
 76:   MatGetDiagonal(A,W2);
 77:   VecView(W2,viewer);
 78:   MatDestroy(&E);
 79:   VecDestroy(&W1);
 80:   VecDestroy(&W2);
 81:   return(0);
 82: }

 86: int main(int argc,char **args)
 87: {
 88:   const PetscScalar xvals[] = {11,13},yvals[] = {17,19},zvals[] = {23,29};
 89:   const PetscInt    inds[]  = {0,1};
 90:   PetscScalar       avals[] = {2,3,5,7};
 91:   Mat               A,S,D[4],N;
 92:   Vec               X,Y,Z;
 93:   User              user;
 94:   PetscInt          i;
 95:   PetscErrorCode    ierr;

 97:   PetscInitialize(&argc,&args,(char*)0,help);
 98:   MatCreateSeqAIJ(PETSC_COMM_WORLD,2,2,2,NULL,&A);
 99:   MatSetUp(A);
100:   MatSetValues(A,2,inds,2,inds,avals,INSERT_VALUES);
101:   VecCreateSeq(PETSC_COMM_WORLD,2,&X);
102:   VecDuplicate(X,&Y);
103:   VecDuplicate(X,&Z);
104:   VecSetValues(X,2,inds,xvals,INSERT_VALUES);
105:   VecSetValues(Y,2,inds,yvals,INSERT_VALUES);
106:   VecSetValues(Z,2,inds,zvals,INSERT_VALUES);
107:   MatAssemblyBegin(A,MAT_FINAL_ASSEMBLY);
108:   MatAssemblyEnd(A,MAT_FINAL_ASSEMBLY);
109:   VecAssemblyBegin(X);
110:   VecAssemblyBegin(Y);
111:   VecAssemblyBegin(Z);
112:   VecAssemblyEnd(X);
113:   VecAssemblyEnd(Y);
114:   VecAssemblyEnd(Z);

116:   PetscNew(&user);
117:   user->B = A;

119:   MatCreateShell(PETSC_COMM_WORLD,2,2,2,2,user,&S);
120:   MatSetUp(S);
121:   MatShellSetOperation(S,MATOP_MULT,(void (*)(void))MatMult_User);
122:   MatShellSetOperation(S,MATOP_MULT_TRANSPOSE,(void (*)(void))MatMultTranspose_User);
123:   MatShellSetOperation(S,MATOP_GET_DIAGONAL,(void (*)(void))MatGetDiagonal_User);

125:   for (i=0; i<4; i++) {
126:     MatCreateSeqDense(PETSC_COMM_WORLD,1,1,&avals[i],&D[i]);
127:   }
128:   MatCreateNest(PETSC_COMM_WORLD,2,NULL,2,NULL,D,&N);
129:   MatSetUp(N);

131:   TestMatrix(S,X,Y,Z);
132:   TestMatrix(A,X,Y,Z);
133:   TestMatrix(N,X,Y,Z);

135:   for (i=0; i<4; i++) {MatDestroy(&D[i]);}
136:   MatDestroy(&A);
137:   MatDestroy(&S);
138:   MatDestroy(&N);
139:   VecDestroy(&X);
140:   VecDestroy(&Y);
141:   VecDestroy(&Z);
142:   PetscFree(user);
143:   PetscFinalize();
144:   return 0;
145: }