1 #include "petscsys.h" 2 static char help[] = "Tests DMSwarm with DMShell\n\n"; 3 4 #include <petscsf.h> 5 #include <petscdm.h> 6 #include <petscdmshell.h> 7 #include <petscdmda.h> 8 #include <petscdmswarm.h> 9 #include <petsc/private/dmimpl.h> 10 11 PetscErrorCode _DMLocatePoints_DMDARegular_IS(DM dm, Vec pos, IS *iscell) 12 { 13 PetscInt p, n, bs, npoints, si, sj, milocal, mjlocal, mx, my; 14 DM dmregular; 15 PetscInt *cellidx; 16 const PetscScalar *coor; 17 PetscReal dx, dy; 18 PetscMPIInt rank; 19 20 PetscFunctionBegin; 21 PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank)); 22 PetscCall(VecGetLocalSize(pos, &n)); 23 PetscCall(VecGetBlockSize(pos, &bs)); 24 npoints = n / bs; 25 26 PetscCall(PetscMalloc1(npoints, &cellidx)); 27 PetscCall(DMGetApplicationContext(dm, &dmregular)); 28 PetscCall(DMDAGetCorners(dmregular, &si, &sj, NULL, &milocal, &mjlocal, NULL)); 29 PetscCall(DMDAGetInfo(dmregular, NULL, &mx, &my, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL, NULL)); 30 31 dx = 2.0 / ((PetscReal)mx); 32 dy = 2.0 / ((PetscReal)my); 33 34 PetscCall(VecGetArrayRead(pos, &coor)); 35 for (p = 0; p < npoints; p++) { 36 PetscReal coorx, coory; 37 PetscInt mi, mj; 38 39 coorx = PetscRealPart(coor[2 * p]); 40 coory = PetscRealPart(coor[2 * p + 1]); 41 42 mi = (PetscInt)((coorx - (-1.0)) / dx); 43 mj = (PetscInt)((coory - (-1.0)) / dy); 44 45 cellidx[p] = DMLOCATEPOINT_POINT_NOT_FOUND; 46 47 if ((mj >= sj) && (mj < sj + mjlocal)) { 48 if ((mi >= si) && (mi < si + milocal)) cellidx[p] = (mi - si) + (mj - sj) * milocal; 49 } 50 if (coorx < -1.0) cellidx[p] = DMLOCATEPOINT_POINT_NOT_FOUND; 51 if (coorx > 1.0) cellidx[p] = DMLOCATEPOINT_POINT_NOT_FOUND; 52 if (coory < -1.0) cellidx[p] = DMLOCATEPOINT_POINT_NOT_FOUND; 53 if (coory > 1.0) cellidx[p] = DMLOCATEPOINT_POINT_NOT_FOUND; 54 } 55 PetscCall(VecRestoreArrayRead(pos, &coor)); 56 PetscCall(ISCreateGeneral(PETSC_COMM_SELF, npoints, cellidx, PETSC_OWN_POINTER, iscell)); 57 PetscFunctionReturn(PETSC_SUCCESS); 58 } 59 60 PetscErrorCode DMLocatePoints_DMDARegular(DM dm, Vec pos, DMPointLocationType ltype, PetscSF cellSF) 61 { 62 IS iscell; 63 PetscSFNode *cells; 64 PetscInt p, bs, npoints, nfound; 65 const PetscInt *boxCells; 66 67 PetscFunctionBegin; 68 PetscCall(_DMLocatePoints_DMDARegular_IS(dm, pos, &iscell)); 69 PetscCall(VecGetLocalSize(pos, &npoints)); 70 PetscCall(VecGetBlockSize(pos, &bs)); 71 npoints = npoints / bs; 72 73 PetscCall(PetscMalloc1(npoints, &cells)); 74 PetscCall(ISGetIndices(iscell, &boxCells)); 75 76 for (p = 0; p < npoints; p++) { 77 cells[p].rank = 0; 78 cells[p].index = DMLOCATEPOINT_POINT_NOT_FOUND; 79 cells[p].index = boxCells[p]; 80 } 81 PetscCall(ISRestoreIndices(iscell, &boxCells)); 82 PetscCall(ISDestroy(&iscell)); 83 nfound = npoints; 84 PetscCall(PetscSFSetGraph(cellSF, npoints, nfound, NULL, PETSC_OWN_POINTER, cells, PETSC_OWN_POINTER)); 85 PetscCall(ISDestroy(&iscell)); 86 PetscFunctionReturn(PETSC_SUCCESS); 87 } 88 89 PetscErrorCode DMGetNeighbors_DMDARegular(DM dm, PetscInt *nneighbors, const PetscMPIInt **neighbors) 90 { 91 DM dmregular; 92 93 PetscFunctionBegin; 94 PetscCall(DMGetApplicationContext(dm, &dmregular)); 95 PetscCall(DMGetNeighbors(dmregular, nneighbors, neighbors)); 96 PetscFunctionReturn(PETSC_SUCCESS); 97 } 98 99 PetscErrorCode SwarmViewGP(DM dms, const char prefix[]) 100 { 101 PetscReal *array; 102 PetscInt *iarray; 103 PetscInt npoints, p, bs; 104 FILE *fp; 105 char name[PETSC_MAX_PATH_LEN]; 106 PetscMPIInt rank; 107 108 PetscFunctionBegin; 109 PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank)); 110 PetscCall(PetscSNPrintf(name, PETSC_MAX_PATH_LEN - 1, "%s-rank%d.gp", prefix, rank)); 111 fp = fopen(name, "w"); 112 PetscCheck(fp, PETSC_COMM_SELF, PETSC_ERR_FILE_OPEN, "Cannot open file %s", name); 113 PetscCall(DMSwarmGetLocalSize(dms, &npoints)); 114 PetscCall(DMSwarmGetField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array)); 115 PetscCall(DMSwarmGetField(dms, "itag", NULL, NULL, (void **)&iarray)); 116 for (p = 0; p < npoints; p++) fprintf(fp, "%+1.4e %+1.4e %1.4e\n", array[2 * p], array[2 * p + 1], (double)iarray[p]); 117 PetscCall(DMSwarmRestoreField(dms, "itag", NULL, NULL, (void **)&iarray)); 118 PetscCall(DMSwarmRestoreField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array)); 119 fclose(fp); 120 PetscFunctionReturn(PETSC_SUCCESS); 121 } 122 123 /* 124 Create a DMShell and attach a regularly spaced DMDA for point location 125 Override methods for point location 126 */ 127 PetscErrorCode ex3_1(void) 128 { 129 DM dms, dmcell, dmregular; 130 PetscMPIInt rank; 131 PetscInt p, bs, nlocal, overlap, mx, tk; 132 PetscReal dx; 133 PetscReal *array, dt; 134 PetscInt *iarray; 135 PetscRandom rand; 136 137 PetscFunctionBegin; 138 PetscCallMPI(MPI_Comm_rank(PETSC_COMM_WORLD, &rank)); 139 140 /* Create a regularly spaced DMDA */ 141 mx = 40; 142 overlap = 0; 143 PetscCall(DMDACreate2d(PETSC_COMM_WORLD, DM_BOUNDARY_NONE, DM_BOUNDARY_NONE, DMDA_STENCIL_BOX, mx, mx, PETSC_DECIDE, PETSC_DECIDE, 1, overlap, NULL, NULL, &dmregular)); 144 PetscCall(DMSetFromOptions(dmregular)); 145 PetscCall(DMSetUp(dmregular)); 146 147 dx = 2.0 / ((PetscReal)mx); 148 PetscCall(DMDASetUniformCoordinates(dmregular, -1.0 + 0.5 * dx, 1.0 - 0.5 * dx, -1.0 + 0.5 * dx, 1.0 - 0.5 * dx, -1.0, 1.0)); 149 150 /* Create a DMShell for point location purposes */ 151 PetscCall(DMShellCreate(PETSC_COMM_WORLD, &dmcell)); 152 PetscCall(PetscObjectSetName((PetscObject)dmcell, "celldm")); 153 PetscCall(DMSetApplicationContext(dmcell, dmregular)); 154 dmcell->ops->locatepoints = DMLocatePoints_DMDARegular; 155 dmcell->ops->getneighbors = DMGetNeighbors_DMDARegular; 156 157 /* Create the swarm */ 158 PetscCall(DMCreate(PETSC_COMM_WORLD, &dms)); 159 PetscCall(DMSetType(dms, DMSWARM)); 160 PetscCall(DMSetDimension(dms, 2)); 161 PetscCall(PetscObjectSetName((PetscObject)dms, "Particles")); 162 163 PetscCall(DMSwarmSetType(dms, DMSWARM_PIC)); 164 PetscCall(DMSwarmSetCellDM(dms, dmcell)); 165 166 PetscCall(DMSwarmRegisterPetscDatatypeField(dms, "itag", 1, PETSC_INT)); 167 PetscCall(DMSwarmFinalizeFieldRegister(dms)); 168 { 169 PetscInt si, sj, milocal, mjlocal; 170 const PetscScalar *LA_coors; 171 Vec coors; 172 PetscInt cnt; 173 174 PetscCall(DMDAGetCorners(dmregular, &si, &sj, NULL, &milocal, &mjlocal, NULL)); 175 PetscCall(DMGetCoordinates(dmregular, &coors)); 176 PetscCall(VecGetArrayRead(coors, &LA_coors)); 177 PetscCall(DMSwarmSetLocalSizes(dms, milocal * mjlocal, 4)); 178 PetscCall(DMSwarmGetLocalSize(dms, &nlocal)); 179 PetscCall(DMSwarmGetField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array)); 180 cnt = 0; 181 PetscCall(PetscRandomCreate(PETSC_COMM_SELF, &rand)); 182 PetscCall(PetscRandomSetInterval(rand, -1.0, 1.0)); 183 for (p = 0; p < nlocal; p++) { 184 PetscReal px, py, rx, ry, r2; 185 186 PetscCall(PetscRandomGetValueReal(rand, &rx)); 187 PetscCall(PetscRandomGetValueReal(rand, &ry)); 188 189 px = PetscRealPart(LA_coors[2 * p + 0]) + 0.1 * rx * dx; 190 py = PetscRealPart(LA_coors[2 * p + 1]) + 0.1 * ry * dx; 191 192 r2 = px * px + py * py; 193 if (r2 < 0.75 * 0.75) { 194 array[bs * cnt + 0] = px; 195 array[bs * cnt + 1] = py; 196 cnt++; 197 } 198 } 199 PetscCall(PetscRandomDestroy(&rand)); 200 PetscCall(DMSwarmRestoreField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array)); 201 PetscCall(VecRestoreArrayRead(coors, &LA_coors)); 202 PetscCall(DMSwarmSetLocalSizes(dms, cnt, 4)); 203 204 PetscCall(DMSwarmGetLocalSize(dms, &nlocal)); 205 PetscCall(DMSwarmGetField(dms, "itag", &bs, NULL, (void **)&iarray)); 206 for (p = 0; p < nlocal; p++) iarray[p] = (PetscInt)rank; 207 PetscCall(DMSwarmRestoreField(dms, "itag", &bs, NULL, (void **)&iarray)); 208 } 209 210 PetscCall(DMView(dms, PETSC_VIEWER_STDOUT_WORLD)); 211 PetscCall(SwarmViewGP(dms, "step0")); 212 213 dt = 0.1; 214 for (tk = 1; tk < 20; tk++) { 215 char prefix[PETSC_MAX_PATH_LEN]; 216 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "Step %" PetscInt_FMT " \n", tk)); 217 /* push points */ 218 PetscCall(DMSwarmGetLocalSize(dms, &nlocal)); 219 PetscCall(DMSwarmGetField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array)); 220 for (p = 0; p < nlocal; p++) { 221 PetscReal cx, cy, vx, vy; 222 223 cx = array[2 * p]; 224 cy = array[2 * p + 1]; 225 vx = cy; 226 vy = -cx; 227 228 array[2 * p] += dt * vx; 229 array[2 * p + 1] += dt * vy; 230 } 231 PetscCall(DMSwarmRestoreField(dms, DMSwarmPICField_coor, &bs, NULL, (void **)&array)); 232 233 /* migrate points */ 234 PetscCall(DMSwarmMigrate(dms, PETSC_TRUE)); 235 /* view points */ 236 PetscCall(PetscSNPrintf(prefix, PETSC_MAX_PATH_LEN - 1, "step%" PetscInt_FMT, tk)); 237 /* should use the regular SwarmView() api, not one for a particular type */ 238 PetscCall(SwarmViewGP(dms, prefix)); 239 } 240 PetscCall(DMDestroy(&dmregular)); 241 PetscCall(DMDestroy(&dmcell)); 242 PetscCall(DMDestroy(&dms)); 243 PetscFunctionReturn(PETSC_SUCCESS); 244 } 245 246 int main(int argc, char **argv) 247 { 248 PetscFunctionBeginUser; 249 PetscCall(PetscInitialize(&argc, &argv, NULL, help)); 250 PetscCall(ex3_1()); 251 PetscCall(PetscFinalize()); 252 return 0; 253 } 254 255 /*TEST 256 257 build: 258 requires: double !complex 259 260 test: 261 filter: grep -v atomic 262 filter_output: grep -v atomic 263 TEST*/ 264