1 static char help[] = "Test DMPlexGetLETKFLocalizationMatrix.\n\n"; 2 3 #include <petscdmplex.h> 4 #include <petscdmda.h> 5 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