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