Actual source code: ex8.c
2: static char help[] = "Shows how to add a new MatOperation to AIJ MatType\n\n";
4: #include <petscmat.h>
5: #include <petscblaslapack.h>
7: static PetscErrorCode MatScaleUserImpl_SeqAIJ(Mat inA, PetscScalar alpha)
8: {
9: PetscFunctionBegin;
10: PetscCall(MatScale(inA, alpha));
11: PetscFunctionReturn(PETSC_SUCCESS);
12: }
14: extern PetscErrorCode MatScaleUserImpl(Mat, PetscScalar);
16: static PetscErrorCode MatScaleUserImpl_MPIAIJ(Mat A, PetscScalar aa)
17: {
18: Mat AA, AB;
20: PetscFunctionBegin;
21: PetscCall(MatMPIAIJGetSeqAIJ(A, &AA, &AB, NULL));
22: PetscCall(MatScaleUserImpl(AA, aa));
23: PetscCall(MatScaleUserImpl(AB, aa));
24: PetscFunctionReturn(PETSC_SUCCESS);
25: }
27: /* This routine registers MatScaleUserImpl_SeqAIJ() and
28: MatScaleUserImpl_MPIAIJ() as methods providing MatScaleUserImpl()
29: functionality for SeqAIJ and MPIAIJ matrix-types */
30: PetscErrorCode RegisterMatScaleUserImpl(Mat mat)
31: {
32: PetscMPIInt size;
34: PetscFunctionBegin;
35: PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size));
36: if (size == 1) { /* SeqAIJ Matrix */
37: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ));
38: } else { /* MPIAIJ Matrix */
39: Mat AA, AB;
40: PetscCall(MatMPIAIJGetSeqAIJ(mat, &AA, &AB, NULL));
41: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", MatScaleUserImpl_MPIAIJ));
42: PetscCall(PetscObjectComposeFunction((PetscObject)AA, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ));
43: PetscCall(PetscObjectComposeFunction((PetscObject)AB, "MatScaleUserImpl_C", MatScaleUserImpl_SeqAIJ));
44: }
45: PetscFunctionReturn(PETSC_SUCCESS);
46: }
48: /* This routine deregisters MatScaleUserImpl_SeqAIJ() and
49: MatScaleUserImpl_MPIAIJ() as methods providing MatScaleUserImpl()
50: functionality for SeqAIJ and MPIAIJ matrix-types */
51: PetscErrorCode DeRegisterMatScaleUserImpl(Mat mat)
52: {
53: PetscMPIInt size;
55: PetscFunctionBegin;
56: PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)mat), &size));
57: if (size == 1) { /* SeqAIJ Matrix */
58: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", NULL));
59: } else { /* MPIAIJ Matrix */
60: Mat AA, AB;
61: PetscCall(MatMPIAIJGetSeqAIJ(mat, &AA, &AB, NULL));
62: PetscCall(PetscObjectComposeFunction((PetscObject)mat, "MatScaleUserImpl_C", NULL));
63: PetscCall(PetscObjectComposeFunction((PetscObject)AA, "MatScaleUserImpl_C", NULL));
64: PetscCall(PetscObjectComposeFunction((PetscObject)AB, "MatScaleUserImpl_C", NULL));
65: }
66: PetscFunctionReturn(PETSC_SUCCESS);
67: }
69: /* this routines queries the already registered MatScaleUserImp_XXX
70: implementations for the given matrix, and calls the correct
71: routine. i.e if MatType is SeqAIJ, MatScaleUserImpl_SeqAIJ() gets
72: called, and if MatType is MPIAIJ, MatScaleUserImpl_MPIAIJ() gets
73: called */
74: PetscErrorCode MatScaleUserImpl(Mat mat, PetscScalar a)
75: {
76: PetscErrorCode (*f)(Mat, PetscScalar);
78: PetscFunctionBegin;
79: PetscCall(PetscObjectQueryFunction((PetscObject)mat, "MatScaleUserImpl_C", &f));
80: if (f) PetscCall((*f)(mat, a));
81: PetscFunctionReturn(PETSC_SUCCESS);
82: }
84: /* Main user code that uses MatScaleUserImpl() */
86: int main(int argc, char **args)
87: {
88: Mat mat;
89: PetscInt i, j, m = 2, n, Ii, J;
90: PetscScalar v, none = -1.0;
91: PetscMPIInt rank, size;
93: PetscFunctionBeginUser;
94: PetscCall(PetscInitialize(&argc, &args, (char *)0, help));
95: PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
96: PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
97: n = 2 * size;
99: /* create the matrix */
100: PetscCall(MatCreate(PETSC_COMM_WORLD, &mat));
101: PetscCall(MatSetSizes(mat, PETSC_DECIDE, PETSC_DECIDE, m * n, m * n));
102: PetscCall(MatSetType(mat, MATAIJ));
103: PetscCall(MatSetUp(mat));
105: /* register user defined MatScaleUser() operation for both SeqAIJ
106: and MPIAIJ types */
107: PetscCall(RegisterMatScaleUserImpl(mat));
109: /* assemble the matrix */
110: for (i = 0; i < m; i++) {
111: for (j = 2 * rank; j < 2 * rank + 2; j++) {
112: v = -1.0;
113: Ii = j + n * i;
114: if (i > 0) {
115: J = Ii - n;
116: PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES));
117: }
118: if (i < m - 1) {
119: J = Ii + n;
120: PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES));
121: }
122: if (j > 0) {
123: J = Ii - 1;
124: PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES));
125: }
126: if (j < n - 1) {
127: J = Ii + 1;
128: PetscCall(MatSetValues(mat, 1, &Ii, 1, &J, &v, INSERT_VALUES));
129: }
130: v = 4.0;
131: PetscCall(MatSetValues(mat, 1, &Ii, 1, &Ii, &v, INSERT_VALUES));
132: }
133: }
134: PetscCall(MatAssemblyBegin(mat, MAT_FINAL_ASSEMBLY));
135: PetscCall(MatAssemblyEnd(mat, MAT_FINAL_ASSEMBLY));
137: /* check the matrix before and after scaling by -1.0 */
138: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Matrix _before_ MatScaleUserImpl() operation\n"));
139: PetscCall(MatView(mat, PETSC_VIEWER_STDOUT_WORLD));
140: PetscCall(MatScaleUserImpl(mat, none));
141: PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Matrix _after_ MatScaleUserImpl() operation\n"));
142: PetscCall(MatView(mat, PETSC_VIEWER_STDOUT_WORLD));
144: /* deregister user defined MatScaleUser() operation for both SeqAIJ
145: and MPIAIJ types */
146: PetscCall(DeRegisterMatScaleUserImpl(mat));
147: PetscCall(MatDestroy(&mat));
148: PetscCall(PetscFinalize());
149: return 0;
150: }
152: /*TEST
154: test:
156: test:
157: suffix: 2
158: nsize: 2
160: TEST*/