xref: /petsc/src/dm/tests/ex30.c (revision 58d68138c660dfb4e9f5b03334792cd4f2ffd7cc)
1 static char help[] = "Tests DMSLICED operations\n\n";
2 
3 #include <petscdmsliced.h>
4 
5 int main(int argc, char *argv[]) {
6   char         mat_type[256] = MATAIJ; /* default matrix type */
7   MPI_Comm     comm;
8   PetscMPIInt  rank, size;
9   DM           slice;
10   PetscInt     i, bs = 1, N = 5, n, m, rstart, ghosts[2], *d_nnz, *o_nnz, dfill[4] = {1, 0, 0, 1}, ofill[4] = {1, 1, 1, 1};
11   PetscReal    alpha = 1, K = 1, rho0 = 1, u0 = 0, sigma = 0.2;
12   PetscBool    useblock = PETSC_TRUE;
13   PetscScalar *xx;
14   Mat          A;
15   Vec          x, b, lf;
16 
17   PetscFunctionBeginUser;
18   PetscCall(PetscInitialize(&argc, &argv, 0, help));
19   comm = PETSC_COMM_WORLD;
20   PetscCallMPI(MPI_Comm_size(comm, &size));
21   PetscCallMPI(MPI_Comm_rank(comm, &rank));
22 
23   PetscOptionsBegin(comm, 0, "Options for DMSliced test", 0);
24   {
25     PetscCall(PetscOptionsInt("-n", "Global number of nodes", "", N, &N, NULL));
26     PetscCall(PetscOptionsInt("-bs", "Block size (1 or 2)", "", bs, &bs, NULL));
27     if (bs != 1) {
28       PetscCheck(bs == 2, PETSC_COMM_WORLD, PETSC_ERR_SUP, "Block size must be 1 or 2");
29       PetscCall(PetscOptionsReal("-alpha", "Inverse time step for wave operator", "", alpha, &alpha, NULL));
30       PetscCall(PetscOptionsReal("-K", "Bulk modulus of compressibility", "", K, &K, NULL));
31       PetscCall(PetscOptionsReal("-rho0", "Reference density", "", rho0, &rho0, NULL));
32       PetscCall(PetscOptionsReal("-u0", "Reference velocity", "", u0, &u0, NULL));
33       PetscCall(PetscOptionsReal("-sigma", "Width of Gaussian density perturbation", "", sigma, &sigma, NULL));
34       PetscCall(PetscOptionsBool("-block", "Use block matrix assembly", "", useblock, &useblock, NULL));
35     }
36     PetscCall(PetscOptionsString("-sliced_mat_type", "Matrix type to use (aij or baij)", "", mat_type, mat_type, sizeof(mat_type), NULL));
37   }
38   PetscOptionsEnd();
39 
40   /* Split ownership, set up periodic grid in 1D */
41   n = PETSC_DECIDE;
42   PetscCall(PetscSplitOwnership(comm, &n, &N));
43   rstart = 0;
44   PetscCallMPI(MPI_Scan(&n, &rstart, 1, MPIU_INT, MPI_SUM, comm));
45   rstart -= n;
46   ghosts[0] = (N + rstart - 1) % N;
47   ghosts[1] = (rstart + n) % N;
48 
49   PetscCall(PetscMalloc2(n, &d_nnz, n, &o_nnz));
50   for (i = 0; i < n; i++) {
51     if (size > 1 && (i == 0 || i == n - 1)) {
52       d_nnz[i] = 2;
53       o_nnz[i] = 1;
54     } else {
55       d_nnz[i] = 3;
56       o_nnz[i] = 0;
57     }
58   }
59   PetscCall(DMSlicedCreate(comm, bs, n, 2, ghosts, d_nnz, o_nnz, &slice)); /* Currently does not copy X_nnz so we can't free them until after DMSlicedGetMatrix */
60 
61   if (!useblock) PetscCall(DMSlicedSetBlockFills(slice, dfill, ofill)); /* Irrelevant for baij formats */
62   PetscCall(DMSetMatType(slice, mat_type));
63   PetscCall(DMCreateMatrix(slice, &A));
64   PetscCall(PetscFree2(d_nnz, o_nnz));
65   PetscCall(MatSetOption(A, MAT_NEW_NONZERO_ALLOCATION_ERR, PETSC_TRUE));
66 
67   PetscCall(DMCreateGlobalVector(slice, &x));
68   PetscCall(VecDuplicate(x, &b));
69 
70   PetscCall(VecGhostGetLocalForm(x, &lf));
71   PetscCall(VecGetSize(lf, &m));
72   PetscCheck(m == (n + 2) * bs, PETSC_COMM_SELF, PETSC_ERR_ARG_SIZ, "size of local form %" PetscInt_FMT ", expected %" PetscInt_FMT, m, (n + 2) * bs);
73   PetscCall(VecGetArray(lf, &xx));
74   for (i = 0; i < n; i++) {
75     PetscInt        row[2], col[9], im, ip;
76     PetscScalar     v[12];
77     const PetscReal xref = 2.0 * (rstart + i) / N - 1; /* [-1,1] */
78     const PetscReal h    = 1.0 / N;                    /* grid spacing */
79     im                   = (i == 0) ? n : i - 1;
80     ip                   = (i == n - 1) ? n + 1 : i + 1;
81     switch (bs) {
82     case 1: /* Laplacian with periodic boundaries */
83       col[0] = im;
84       col[1] = i;
85       col[2] = ip;
86       v[0]   = -h;
87       v[1]   = 2 * h;
88       v[2]   = -h;
89       PetscCall(MatSetValuesLocal(A, 1, &i, 3, col, v, INSERT_VALUES));
90       xx[i] = PetscSinReal(xref * PETSC_PI);
91       break;
92     case 2: /* Linear acoustic wave operator in variables [rho, u], central differences, periodic, timestep 1/alpha */
93       v[0]  = -0.5 * u0;
94       v[1]  = -0.5 * K;
95       v[2]  = alpha;
96       v[3]  = 0;
97       v[4]  = 0.5 * u0;
98       v[5]  = 0.5 * K;
99       v[6]  = -0.5 / rho0;
100       v[7]  = -0.5 * u0;
101       v[8]  = 0;
102       v[9]  = alpha;
103       v[10] = 0.5 / rho0;
104       v[11] = 0.5 * u0;
105       if (useblock) {
106         row[0] = i;
107         col[0] = im;
108         col[1] = i;
109         col[2] = ip;
110         PetscCall(MatSetValuesBlockedLocal(A, 1, row, 3, col, v, INSERT_VALUES));
111       } else {
112         row[0] = 2 * i;
113         row[1] = 2 * i + 1;
114         col[0] = 2 * im;
115         col[1] = 2 * im + 1;
116         col[2] = 2 * i;
117         col[3] = 2 * ip;
118         col[4] = 2 * ip + 1;
119         v[3]   = v[4];
120         v[4]   = v[5]; /* pack values in first row */
121         PetscCall(MatSetValuesLocal(A, 1, row, 5, col, v, INSERT_VALUES));
122         col[2] = 2 * i + 1;
123         v[8]   = v[9];
124         v[9]   = v[10];
125         v[10]  = v[11]; /* pack values in second row */
126         PetscCall(MatSetValuesLocal(A, 1, row + 1, 5, col, v + 6, INSERT_VALUES));
127       }
128       /* Set current state (gaussian density perturbation) */
129       xx[2 * i]     = 0.2 * PetscExpReal(-PetscSqr(xref) / (2 * PetscSqr(sigma)));
130       xx[2 * i + 1] = 0;
131       break;
132     default: SETERRQ(PETSC_COMM_SELF, PETSC_ERR_SUP, "not implemented for block size %" PetscInt_FMT, bs);
133     }
134   }
135   PetscCall(VecRestoreArray(lf, &xx));
136   PetscCall(VecGhostRestoreLocalForm(x, &lf));
137   PetscCall(MatAssemblyBegin(A, MAT_FINAL_ASSEMBLY));
138   PetscCall(MatAssemblyEnd(A, MAT_FINAL_ASSEMBLY));
139 
140   PetscCall(MatMult(A, x, b));
141   PetscCall(MatView(A, PETSC_VIEWER_STDOUT_WORLD));
142   PetscCall(VecView(x, PETSC_VIEWER_STDOUT_WORLD));
143   PetscCall(VecView(b, PETSC_VIEWER_STDOUT_WORLD));
144 
145   /* Update the ghosted values, view the result on rank 0. */
146   PetscCall(VecGhostUpdateBegin(b, INSERT_VALUES, SCATTER_FORWARD));
147   PetscCall(VecGhostUpdateEnd(b, INSERT_VALUES, SCATTER_FORWARD));
148   if (rank == 0) {
149     PetscCall(VecGhostGetLocalForm(b, &lf));
150     PetscCall(PetscViewerASCIIPrintf(PETSC_VIEWER_STDOUT_SELF, "Local form of b on rank 0, last two nodes are ghost nodes\n"));
151     PetscCall(VecView(lf, PETSC_VIEWER_STDOUT_SELF));
152     PetscCall(VecGhostRestoreLocalForm(b, &lf));
153   }
154 
155   PetscCall(DMDestroy(&slice));
156   PetscCall(VecDestroy(&x));
157   PetscCall(VecDestroy(&b));
158   PetscCall(MatDestroy(&A));
159   PetscCall(PetscFinalize());
160   return 0;
161 }
162 
163 /*TEST
164 
165    test:
166       nsize: 2
167       args: -bs 2 -block 0 -sliced_mat_type baij -alpha 10 -u0 0.1
168 
169    test:
170       suffix: 2
171       nsize: 2
172       args: -bs 2 -block 1 -sliced_mat_type aij -alpha 10 -u0 0.1
173 
174    test:
175       suffix: 3
176       nsize: 2
177       args: -bs 2 -block 0 -sliced_mat_type aij -alpha 10 -u0 0.1
178 
179 TEST*/
180