Actual source code: ex13.c
1: static char help[] = "Tests setup-time weighted ASM scaling callbacks.\n";
3: #include <petscksp.h>
5: typedef struct {
6: PetscInt calls;
7: PetscReal scale, normalizer;
8: } ScalingCtx;
10: static PetscErrorCode ComputeScaling(PC pc, PetscInt local, Vec scaling, PetscCtx ctx)
11: {
12: ScalingCtx *user = (ScalingCtx *)ctx;
13: IS *is;
14: const PetscInt *indices;
15: PetscInt n;
16: PetscScalar *values;
18: PetscFunctionBeginUser;
19: PetscCall(PCASMGetLocalSubdomains(pc, NULL, &is, NULL));
20: PetscCall(ISGetLocalSize(is[local], &n));
21: PetscCall(ISGetIndices(is[local], &indices));
22: PetscCall(VecGetArray(scaling, &values));
23: for (PetscInt i = 0; i < n; ++i) values[i] = user->scale * (local + 1) * (indices[i] + 1) / user->normalizer;
24: PetscCall(VecRestoreArray(scaling, &values));
25: PetscCall(ISRestoreIndices(is[local], &indices));
26: ++user->calls;
27: PetscFunctionReturn(PETSC_SUCCESS);
28: }
30: static PetscErrorCode CheckMatApply(PC, Vec, PetscReal, PetscInt, PetscScalar);
32: static PetscErrorCode CheckApply(PC pc, Vec x, Vec y, PetscReal scale, PetscInt nrhs)
33: {
34: const PetscScalar *values;
35: PetscInt start, end;
37: PetscFunctionBeginUser;
38: PetscCall(PCApply(pc, x, y));
39: PetscCall(VecGetOwnershipRange(y, &start, &end));
40: PetscCall(VecGetArrayRead(y, &values));
41: for (PetscInt i = start; i < end; ++i) PetscCheck(PetscAbsScalar(values[i - start] - scale * (i + 1)) < 100 * PETSC_MACHINE_EPSILON, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Incorrect weighted correction");
42: PetscCall(VecRestoreArrayRead(y, &values));
43: if (nrhs) PetscCall(CheckMatApply(pc, x, scale, nrhs, 0.0));
44: PetscFunctionReturn(PETSC_SUCCESS);
45: }
47: static PetscErrorCode CheckMatApply(PC pc, Vec x, PetscReal scale, PetscInt nrhs, PetscScalar coupling)
48: {
49: Mat X, Y;
50: PetscScalar *values;
51: const PetscScalar *rhs, *result;
52: PetscInt m, n, start, ldx, ldy;
54: PetscFunctionBeginUser;
55: PetscCall(VecGetLocalSize(x, &m));
56: PetscCall(VecGetSize(x, &n));
57: PetscCall(VecGetOwnershipRange(x, &start, NULL));
58: PetscCall(MatCreateDense(PETSC_COMM_WORLD, m, PETSC_DECIDE, n, nrhs, NULL, &X));
59: PetscCall(MatDuplicate(X, MAT_DO_NOT_COPY_VALUES, &Y));
60: PetscCall(MatDenseGetLDA(X, &ldx));
61: PetscCall(MatDenseGetLDA(Y, &ldy));
62: PetscCall(MatDenseGetArray(X, &values));
63: for (PetscInt j = 0; j < nrhs; ++j)
64: for (PetscInt i = 0; i < m; ++i) values[i + j * ldx] = (start + i + 1) * (j + 1) - 2 * j;
65: PetscCall(MatDenseRestoreArray(X, &values));
66: PetscCall(MatDenseGetArrayRead(X, &rhs));
67: for (PetscInt transpose = 0; transpose < 2; ++transpose) {
68: if (transpose) PetscCall(PCMatApplyTranspose(pc, X, Y));
69: else PetscCall(PCMatApply(pc, X, Y));
70: PetscCall(MatDenseGetArrayRead(Y, &result));
71: for (PetscInt j = 0; j < nrhs; ++j) {
72: for (PetscInt i = 0; i < m; ++i) {
73: PetscScalar expected = (start + i + 1) * rhs[i + j * ldx];
75: // Blocks [1 coupling; 0 1] stay within each rank's four rows; weights are proportional to the global row + 1.
76: if (transpose && i % 2) expected -= coupling * (start + i) * rhs[i - 1 + j * ldx];
77: if (!transpose && !(i % 2)) expected -= coupling * (start + i + 1) * rhs[i + 1 + j * ldx];
78: expected *= scale;
79: PetscCheck(PetscAbsScalar(result[i + j * ldy] - expected) <= 100 * PETSC_MACHINE_EPSILON * PetscMax(1.0, PetscAbsScalar(expected)), PETSC_COMM_SELF, PETSC_ERR_PLIB, "Incorrect weighted matrix correction");
80: }
81: }
82: PetscCall(MatDenseRestoreArrayRead(Y, &result));
83: }
84: PetscCall(MatDenseRestoreArrayRead(X, &rhs));
85: PetscCall(MatDestroy(&X));
86: PetscCall(MatDestroy(&Y));
87: PetscFunctionReturn(PETSC_SUCCESS);
88: }
90: int main(int argc, char **argv)
91: {
92: PC pc;
93: Mat A;
94: Vec x, y, *weights;
95: IS *is;
96: PetscInt nlocal, n, start, end, nrhs = 0;
97: PetscMPIInt rank, size;
98: ScalingCtx ctx = {0, 2.0, 0.0};
100: PetscFunctionBeginUser;
101: PetscCall(PetscInitialize(&argc, &argv, NULL, help));
102: PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank));
103: PetscCallMPI(MPI_Comm_size(PETSC_COMM_WORLD, &size));
104: PetscCall(PetscOptionsGetInt(NULL, NULL, "-nrhs", &nrhs, NULL));
105: PetscCheck(nrhs >= 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_OUTOFRANGE, "Number of right-hand sides must be nonnegative");
106: nlocal = nrhs ? 1 : rank + 1;
107: n = 4 * size;
108: ctx.normalizer = nrhs ? size : size * (size + 1.0) * (size + 2.0) / 6.0;
109: PetscCall(MatCreateAIJ(PETSC_COMM_WORLD, 4, 4, n, n, nrhs ? 2 : 1, NULL, 0, NULL, &A));
110: PetscCall(MatGetOwnershipRange(A, &start, &end));
111: for (PetscInt i = start; i < end; ++i) PetscCall(MatSetValue(A, i, i, 1.0, INSERT_VALUES));
112: if (nrhs)
113: for (PetscInt i = start; i < end; i += 2) PetscCall(MatSetValue(A, i, i + 1, 0.0, INSERT_VALUES));
114: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
115: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
116: PetscCall(MatCreateVecs(A, &x, &y));
117: PetscCall(VecSet(x, 1.0));
118: PetscCall(PCCreate(PETSC_COMM_WORLD, &pc));
119: PetscCall(PCSetType(pc, PCASM));
120: PetscCall(PCASMSetType(pc, PC_ASM_WEIGHTED));
121: PetscCall(PCASMWeightedSetComputeScaling(pc, ComputeScaling, &ctx));
122: for (PetscInt reset = 0; reset < 2; ++reset) {
123: PetscCall(PCSetOperators(pc, A, A));
124: PetscCall(PetscMalloc1(nlocal, &is));
125: for (PetscInt i = 0; i < nlocal; ++i) PetscCall(ISCreateStride(PETSC_COMM_SELF, n, n - 1, -1, &is[i]));
126: PetscCall(PCASMSetLocalSubdomains(pc, nlocal, is, NULL));
127: for (PetscInt i = 0; i < nlocal; ++i) PetscCall(ISDestroy(&is[i]));
128: PetscCall(PetscFree(is));
129: ctx.calls = 0;
130: ctx.scale = 2.0;
131: PetscCall(CheckApply(pc, x, y, 2.0, nrhs)); // implicit setup, including after reset
132: PetscCall(CheckApply(pc, x, y, 2.0, nrhs));
133: PetscCheck(ctx.calls == nlocal, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Callback repeated without setup");
134: ctx.scale = 6.0;
135: PetscCall(MatScale(A, 2.0));
136: PetscCall(CheckApply(pc, x, y, 3.0, nrhs)); // matrix change recomputes the weights
137: PetscCheck(ctx.calls == 2 * nlocal, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Callback did not run on rebuild");
138: PetscCall(PCASMWeightedGetScaling(pc, NULL, &weights));
139: for (PetscInt i = 0; i < nlocal; ++i) PetscCall(VecSet(weights[i], 0.0));
140: PetscCall(PCASMWeightedSetScaling(pc, nlocal, weights));
141: PetscCall(CheckApply(pc, x, y, 0.0, nrhs)); // explicit weights apply until the next rebuild
142: PetscCall(MatScale(A, 0.5));
143: PetscCall(CheckApply(pc, x, y, 6.0, nrhs));
144: PetscCheck(ctx.calls == 3 * nlocal, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Callback did not replace explicit weights");
145: if (!reset) PetscCall(PCReset(pc));
146: }
147: PetscCall(PCASMWeightedSetComputeScaling(pc, NULL, NULL));
148: PetscCall(MatScale(A, 2.0));
149: PetscCall(CheckApply(pc, x, y, 3.0, nrhs)); // disabling preserves the last weights
150: PetscCheck(ctx.calls == 3 * nlocal, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Disabled callback was called");
151: if (nrhs) {
152: PetscScalar coupling = 0.25;
154: #if PetscDefined(USE_COMPLEX)
155: coupling += 0.125 * PETSC_i;
156: #endif
157: for (PetscInt i = start; i < end; i += 2) PetscCall(MatSetValue(A, i, i + 1, 2.0 * coupling, INSERT_VALUES));
158: PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
159: PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
160: PetscCall(CheckMatApply(pc, x, 3.0, nrhs, coupling));
161: PetscCheck(ctx.calls == 3 * nlocal, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Disabled callback was called after operator change");
162: }
163: PetscCall(PCDestroy(&pc));
164: PetscCall(VecDestroy(&x));
165: PetscCall(VecDestroy(&y));
166: PetscCall(MatDestroy(&A));
167: PetscCall(PetscFinalize());
168: return 0;
169: }
171: /*TEST
173: test:
174: nsize: {{1 2}}
175: output_file: output/empty.out
177: test:
178: suffix: mat
179: nsize: {{1 2}}
180: args: -nrhs {{1 3}}
181: output_file: output/empty.out
183: TEST*/