1 static char help[] = "Test DMPlexGetLETKFLocalizationMatrix.\n\n";
2
3 #include <petscdmplex.h>
4 #include <petscdmda.h>
5
main(int argc,char ** argv)6 int main(int argc, char **argv)
7 {
8 DM dm;
9 Mat H, Q = NULL;
10 PetscInt nvertexobs, ndof = 1, n_state_global;
11 PetscInt dim = 1, n, vStart, vEnd;
12 PetscInt faces[3] = {1, 1, 1};
13 PetscReal lower[3] = {0.0, 0.0, 0.0};
14 PetscReal upper[3] = {1.0, 1.0, 1.0};
15 Vec Vecxyz[3] = {NULL, NULL, NULL};
16 PetscBool isda, isplex, print = PETSC_FALSE;
17 char type[256] = DMPLEX;
18
19 PetscFunctionBeginUser;
20 PetscCall(PetscInitialize(&argc, &argv, NULL, help));
21
22 /* Get dimension and from options. We need the data here and Plex does not have access functions */
23 PetscOptionsBegin(PETSC_COMM_WORLD, "", "DMField Tutorial Options", "DM");
24 PetscCall(PetscOptionsFList("-dm_type", "DM implementation on which to define field", "ex20.c", DMList, type, type, 256, NULL));
25 PetscCall(PetscStrncmp(type, DMPLEX, 256, &isplex));
26 PetscCall(PetscStrncmp(type, DMDA, 256, &isda));
27 PetscCall(PetscOptionsGetBool(NULL, NULL, "-ex20_print", &print, NULL));
28 PetscCall(PetscOptionsGetInt(NULL, NULL, "-dm_plex_dim", &dim, NULL));
29 PetscCheck(dim <= 3 && dim >= 1, PETSC_COMM_WORLD, PETSC_ERR_ARG_INCOMP, "dm_plex_dim = %" PetscInt_FMT, dim);
30 n = 3;
31 PetscCall(PetscOptionsGetIntArray(NULL, NULL, "-dm_plex_box_faces", faces, &n, NULL));
32 PetscCheck(n == 0 || n == dim, PETSC_COMM_WORLD, PETSC_ERR_ARG_INCOMP, "dm_plex_box_faces dimension %" PetscInt_FMT " does not match requested dimension %" PetscInt_FMT, n, dim);
33 n = 3;
34 PetscCall(PetscOptionsGetRealArray(NULL, NULL, "-dm_plex_box_lower", lower, &n, NULL));
35 PetscCheck(n == 0 || n == dim, PETSC_COMM_WORLD, PETSC_ERR_ARG_INCOMP, "dm_plex_box_lower dimension %" PetscInt_FMT " does not match requested dimension %" PetscInt_FMT, n, dim);
36 n = 3;
37 PetscCall(PetscOptionsGetRealArray(NULL, NULL, "-dm_plex_box_upper", upper, &n, NULL));
38 PetscCheck(n == 0 || n == dim, PETSC_COMM_WORLD, PETSC_ERR_ARG_INCOMP, "dm_plex_box_upper dimension %" PetscInt_FMT " does not match requested dimension %" PetscInt_FMT, n, dim);
39 PetscOptionsEnd();
40
41 if (isplex) {
42 PetscCall(DMCreate(PETSC_COMM_WORLD, &dm));
43 PetscCall(DMSetType(dm, DMPLEX));
44 PetscCall(DMSetFromOptions(dm));
45 {
46 PetscSection section;
47 PetscInt pStart, pEnd;
48
49 PetscCall(DMPlexGetChart(dm, &pStart, &pEnd));
50 PetscCall(DMPlexGetDepthStratum(dm, 0, &vStart, &vEnd));
51 PetscCall(PetscSectionCreate(PETSC_COMM_WORLD, §ion));
52 PetscCall(PetscSectionSetNumFields(section, 1));
53 PetscCall(PetscSectionSetChart(section, pStart, pEnd));
54 for (PetscInt v = vStart; v < vEnd; ++v) PetscCall(PetscSectionSetDof(section, v, 1));
55 PetscCall(PetscSectionSetUp(section));
56 PetscCall(DMSetLocalSection(dm, section));
57 PetscCall(PetscSectionDestroy(§ion));
58
59 for (PetscInt d = 0; d < dim; d++) {
60 Vec loc_vec;
61 Vec coordinates;
62 PetscSection coordSection, s;
63 const PetscScalar *coordArray;
64 PetscScalar *xArray;
65
66 PetscCall(DMCreateGlobalVector(dm, &Vecxyz[d]));
67 PetscCall(PetscObjectSetName((PetscObject)Vecxyz[d], d == 0 ? "x_coordinate" : (d == 1 ? "y_coordinate" : "z_coordinate")));
68 PetscCall(DMGetLocalVector(dm, &loc_vec));
69
70 PetscCall(DMGetLocalSection(dm, &s));
71 PetscCall(DMGetCoordinatesLocal(dm, &coordinates));
72 PetscCall(DMGetCoordinateSection(dm, &coordSection));
73 PetscCall(VecGetArrayRead(coordinates, &coordArray));
74 PetscCall(VecGetArray(loc_vec, &xArray));
75
76 for (PetscInt v = vStart; v < vEnd; v++) {
77 PetscInt cOff, sOff;
78
79 PetscCall(PetscSectionGetOffset(coordSection, v, &cOff));
80 PetscCall(PetscSectionGetOffset(s, v, &sOff));
81 xArray[sOff] = coordArray[cOff + d];
82 }
83 PetscCall(VecRestoreArrayRead(coordinates, &coordArray));
84 PetscCall(VecRestoreArray(loc_vec, &xArray));
85
86 PetscCall(DMLocalToGlobal(dm, loc_vec, INSERT_VALUES, Vecxyz[d]));
87 PetscCall(DMRestoreLocalVector(dm, &loc_vec));
88 PetscCall(VecGetSize(Vecxyz[d], &n_state_global));
89 }
90 }
91
92 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Created DMPlex in %" PetscInt_FMT "D with faces (%" PetscInt_FMT ", %" PetscInt_FMT ", %" PetscInt_FMT "), global vector size %" PetscInt_FMT "\n", dim, faces[0], faces[1], faces[2], n_state_global));
93 } else if (isda) {
94 switch (dim) {
95 case 1:
96 PetscCall(DMDACreate1d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, faces[0], ndof, 1, NULL, &dm));
97 break;
98 case 2:
99 PetscCall(DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, DM_BOUNDARY_NONE, DMDA_STENCIL_BOX, faces[0], faces[1] + 1, PETSC_DETERMINE, PETSC_DETERMINE, ndof, 1, NULL, NULL, &dm));
100 break;
101 default:
102 PetscCall(DMDACreate3d(PETSC_COMM_WORLD, DM_BOUNDARY_PERIODIC, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_BOX, faces[0], faces[1] + 1, faces[2] + 1, PETSC_DETERMINE, PETSC_DETERMINE, PETSC_DETERMINE, ndof, 1, NULL, NULL, NULL, &dm));
103 break;
104 }
105 PetscCall(DMSetUp(dm));
106 PetscCall(DMDASetUniformCoordinates(dm, lower[0], upper[0], lower[1], upper[1], lower[2], upper[2]));
107 {
108 Vec coord;
109 PetscCall(DMGetCoordinates(dm, &coord));
110 for (PetscInt d = 0; d < dim; d++) {
111 PetscCall(DMCreateGlobalVector(dm, &Vecxyz[d]));
112 PetscCall(VecSetFromOptions(Vecxyz[d]));
113 PetscCall(PetscObjectSetName((PetscObject)Vecxyz[d], d == 0 ? "x_coordinate" : (d == 1 ? "y_coordinate" : "z_coordinate")));
114 PetscCall(VecStrideGather(coord, d, Vecxyz[d], INSERT_VALUES));
115 PetscCall(VecGetSize(Vecxyz[d], &n_state_global));
116 }
117 }
118 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Created DMDA of type %s in %" PetscInt_FMT "D with faces (%" PetscInt_FMT ", %" PetscInt_FMT ", %" PetscInt_FMT "), global vector size %" PetscInt_FMT "\n", type, dim, faces[0], faces[1], faces[2], n_state_global));
119 } else SETERRQ(PETSC_COMM_WORLD, PETSC_ERR_SUP, "This test does not run for DM type %s", type);
120 PetscCall(DMViewFromOptions(dm, NULL, "-ex20_dm_view")); // PetscSleep(10);
121
122 /* Set number of local observations to use: 3^dim */
123 nvertexobs = 1;
124 for (PetscInt d = 0; d < dim && d < 2; d++) nvertexobs *= 3;
125 PetscCheck(nvertexobs > 0, PETSC_COMM_WORLD, PETSC_ERR_ARG_INCOMP, "nvertexobs %" PetscInt_FMT " must be > 0 locally for now", nvertexobs);
126
127 /* Count observations (every other vertex in each dimension) */
128 PetscInt nobs_local = 0;
129 PetscBool *isObs;
130 PetscInt nloc;
131
132 PetscCall(VecGetLocalSize(Vecxyz[0], &nloc));
133 PetscCall(PetscMalloc1(nloc, &isObs));
134 {
135 const PetscScalar *coords[3];
136 PetscReal gridSpacing[3];
137 for (PetscInt d = 0; d < dim; d++) PetscCall(VecGetArrayRead(Vecxyz[d], &coords[d]));
138 for (PetscInt d = 0; d < dim; d++) gridSpacing[d] = (upper[d] - lower[d]) / faces[d];
139
140 for (PetscInt v = 0; v < nloc; v++) {
141 PetscReal c[3] = {0.0, 0.0, 0.0};
142
143 isObs[v] = PETSC_TRUE;
144 for (PetscInt d = 0; d < dim; d++) c[d] = PetscRealPart(coords[d][v]);
145 /* Check if this vertex is at an observation location (every other grid point) */
146 for (PetscInt d = 0; d < dim; d++) {
147 PetscReal relCoord = c[d] - lower[d];
148 PetscInt gridIdx = (PetscInt)PetscFloorReal(relCoord / gridSpacing[d] + 0.5);
149 PetscCheck(PetscAbsReal(relCoord - gridIdx * gridSpacing[d]) < PETSC_SMALL, PETSC_COMM_WORLD, PETSC_ERR_ARG_INCOMP, "Error vertex v %" PetscInt_FMT " (dim %" PetscInt_FMT "): %g not on grid (h= %g, distance to grid %g)", v, d, (double)c[d], (double)gridSpacing[d], (double)PetscAbsReal(relCoord - gridIdx * gridSpacing[d]));
150 if (gridIdx % 2 != 0) {
151 isObs[v] = PETSC_FALSE;
152 break;
153 }
154 }
155 if (isObs[v]) nobs_local++;
156 }
157 for (PetscInt d = 0; d < dim; d++) PetscCall(VecRestoreArrayRead(Vecxyz[d], &coords[d]));
158 }
159
160 /* Create H matrix n_obs X n_state */
161 PetscCall(MatCreate(PETSC_COMM_WORLD, &H));
162 PetscCall(MatSetSizes(H, nobs_local, PETSC_DECIDE, PETSC_DECIDE, n_state_global)); //
163 PetscCall(MatSetBlockSizes(H, 1, ndof));
164 PetscCall(MatSetType(H, MATAIJ));
165 PetscCall(MatSeqAIJSetPreallocation(H, 1, NULL));
166 PetscCall(MatMPIAIJSetPreallocation(H, 1, NULL, 1, NULL)); // assumes boolean observation operator, could use interpolation
167 PetscCall(PetscObjectSetName((PetscObject)H, "H_observation_operator"));
168 PetscCall(MatSetFromOptions(H));
169
170 /* Fill H matrix */
171 PetscInt globalRowIdx, globalColIdx, obsIdx = 0;
172 PetscCall(VecGetOwnershipRange(Vecxyz[0], &globalColIdx, NULL));
173 PetscCall(MatGetOwnershipRange(H, &globalRowIdx, NULL));
174 for (PetscInt v = 0; v < nloc; v++) {
175 if (isObs[v]) {
176 PetscInt grow = globalRowIdx + obsIdx++, gcol = globalColIdx + v;
177 PetscCall(MatSetValue(H, grow, gcol, 1.0, INSERT_VALUES));
178 }
179 }
180 PetscCall(PetscFree(isObs));
181 PetscCall(MatAssemblyBegin(H, MAT_FINAL_ASSEMBLY));
182 PetscCall(MatAssemblyEnd(H, MAT_FINAL_ASSEMBLY));
183
184 /* View H */
185 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Observation Operator H:\n"));
186 if (print) PetscCall(MatView(H, PETSC_VIEWER_STDOUT_WORLD));
187
188 /* Perturb interior vertex coordinates */
189 {
190 PetscScalar *coords[3];
191 PetscInt nloc;
192 unsigned long seed = 123456789;
193
194 PetscCall(VecGetLocalSize(Vecxyz[0], &nloc));
195 for (PetscInt d = 0; d < dim; d++) PetscCall(VecGetArray(Vecxyz[d], &coords[d]));
196
197 for (PetscInt v = 0; v < nloc; v++) {
198 for (PetscInt d = 0; d < dim; d++) {
199 PetscReal noise, gridSpacing = (upper[d] - lower[d]) / faces[d];
200
201 seed = (1103515245 * seed + 12345) % 2147483648;
202 noise = (PetscReal)seed / 2147483648.0;
203 coords[d][v] += (noise - 0.5) * 0.001 * gridSpacing;
204 }
205 }
206 for (PetscInt d = 0; d < dim; d++) PetscCall(VecRestoreArray(Vecxyz[d], &coords[d]));
207 }
208
209 /* Call the function */
210 PetscCall(DMPlexGetLETKFLocalizationMatrix(nvertexobs, nobs_local, ndof, Vecxyz, H, &Q));
211 PetscCall(PetscObjectSetName((PetscObject)Q, "Q_localization"));
212
213 // View Q
214 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Localization Matrix Q:\n"));
215 if (print) PetscCall(MatView(Q, PETSC_VIEWER_STDOUT_WORLD));
216
217 /* Cleanup */
218 for (PetscInt d = 0; d < dim; d++) PetscCall(VecDestroy(&Vecxyz[d]));
219 PetscCall(MatDestroy(&H));
220 PetscCall(MatDestroy(&Q));
221 PetscCall(DMDestroy(&dm));
222 PetscCall(PetscFinalize());
223 return 0;
224 }
225
226 /*TEST
227
228 test:
229 requires: kokkos_kernels
230 suffix: 1
231 diff_args: -j
232 args: -dm_plex_dim 1 -dm_plex_box_faces 16 -dm_plex_simplex 0 -dm_plex_box_bd periodic -dm_plex_box_upper 5 -ex20_print -ex20_dm_view -ex20_dm_view -mat_type aijkokkos -dm_vec_type kokkos
233
234 test:
235 requires: kokkos_kernels
236 suffix: 2
237 diff_args: -j
238 args: -dm_plex_dim 2 -dm_plex_box_faces 7,7 -dm_plex_simplex 0 -dm_plex_box_bd periodic,none -dm_plex_box_upper 5,5 -ex20_print -ex20_dm_view -ex20_dm_view -mat_type aijkokkos -dm_vec_type kokkos
239
240 test:
241 requires: kokkos_kernels
242 suffix: da2
243 diff_args: -j
244 args: -dm_type da -dm_plex_dim 2 -dm_plex_box_faces 7,7 -dm_plex_box_upper 5,5 -ex20_print -ex20_dm_view -ex20_dm_view -mat_type aijkokkos -vec_type kokkos
245
246 test:
247 requires: kokkos_kernels
248 suffix: 3
249 diff_args: -j
250 args: -dm_plex_dim 3 -dm_plex_box_faces 5,5,5 -dm_plex_simplex 0 -dm_plex_box_bd periodic,none,none -dm_plex_box_upper 5,5,5 -ex20_print -ex20_dm_view -mat_type aijkokkos -dm_vec_type kokkos
251
252 TEST*/
253