1 static char help[] = "Vlasov-Poisson example of central orbits\n"; 2 3 /* 4 To visualize the orbit, we can used 5 6 -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain -1 -ts_monitor_sp_swarm_phase 0 -draw_size 500,500 7 8 and we probably want it to run fast and not check convergence 9 10 -convest_num_refine 0 -ts_dt 0.01 -ts_max_steps 100 -ts_max_time 100 -output_step 10 11 */ 12 13 #include <petscts.h> 14 #include <petscdmplex.h> 15 #include <petscdmswarm.h> 16 #include <petsc/private/dmpleximpl.h> /* For norm and dot */ 17 #include <petscfe.h> 18 #include <petscds.h> 19 #include <petsc/private/petscfeimpl.h> /* For interpolation */ 20 #include <petscksp.h> 21 #include <petscsnes.h> 22 23 PETSC_EXTERN PetscErrorCode circleSingleX(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *); 24 PETSC_EXTERN PetscErrorCode circleSingleV(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *); 25 PETSC_EXTERN PetscErrorCode circleMultipleX(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *); 26 PETSC_EXTERN PetscErrorCode circleMultipleV(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *); 27 28 const char *EMTypes[] = {"primal", "mixed", "coulomb", "none", "EMType", "EM_", NULL}; 29 typedef enum { 30 EM_PRIMAL, 31 EM_MIXED, 32 EM_COULOMB, 33 EM_NONE 34 } EMType; 35 36 typedef struct { 37 PetscBool error; /* Flag for printing the error */ 38 PetscInt ostep; /* print the energy at each ostep time steps */ 39 PetscReal timeScale; /* Nondimensionalizing time scale */ 40 PetscReal sigma; /* Linear charge per box length */ 41 EMType em; /* Type of electrostatic model */ 42 SNES snes; 43 } AppCtx; 44 45 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options) 46 { 47 PetscFunctionBeginUser; 48 options->error = PETSC_FALSE; 49 options->ostep = 100; 50 options->timeScale = 1.0e-6; 51 options->sigma = 1.; 52 options->em = EM_COULOMB; 53 54 PetscOptionsBegin(comm, "", "Central Orbit Options", "DMSWARM"); 55 PetscCall(PetscOptionsBool("-error", "Flag to print the error", "ex6.c", options->error, &options->error, NULL)); 56 PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", "ex6.c", options->ostep, &options->ostep, NULL)); 57 PetscCall(PetscOptionsReal("-sigma", "Linear charge per box length", "ex6.c", options->sigma, &options->sigma, NULL)); 58 PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", "ex6.c", options->timeScale, &options->timeScale, NULL)); 59 PetscCall(PetscOptionsEnum("-em_type", "Type of electrostatic solver", "ex6.c", EMTypes, (PetscEnum)options->em, (PetscEnum *)&options->em, NULL)); 60 PetscOptionsEnd(); 61 PetscFunctionReturn(PETSC_SUCCESS); 62 } 63 64 static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm) 65 { 66 PetscFunctionBeginUser; 67 PetscCall(DMCreate(comm, dm)); 68 PetscCall(DMSetType(*dm, DMPLEX)); 69 PetscCall(DMSetFromOptions(*dm)); 70 PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view")); 71 PetscFunctionReturn(PETSC_SUCCESS); 72 } 73 74 static void laplacian_f1(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f1[]) 75 { 76 PetscInt d; 77 for (d = 0; d < dim; ++d) f1[d] = u_x[d]; 78 } 79 80 static void laplacian_g3(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g3[]) 81 { 82 PetscInt d; 83 for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0; 84 } 85 86 /* 87 / I grad\ /q\ = /0\ 88 \-div 0 / \u/ \f/ 89 */ 90 static void f0_q(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f0[]) 91 { 92 for (PetscInt c = 0; c < dim; ++c) f0[c] += u[uOff[0] + c]; 93 } 94 95 static void f1_q(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar f1[]) 96 { 97 for (PetscInt d = 0; d < dim; ++d) f1[d * dim + d] += u[uOff[1]]; 98 } 99 100 /* <t, q> */ 101 static void g0_qq(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g0[]) 102 { 103 PetscInt c; 104 for (c = 0; c < dim; ++c) g0[c * dim + c] += 1.0; 105 } 106 107 static void g2_qu(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g2[]) 108 { 109 for (PetscInt d = 0; d < dim; ++d) g2[d * dim + d] += 1.0; 110 } 111 112 static void g1_uq(PetscInt dim, PetscInt Nf, PetscInt NfAux, const PetscInt uOff[], const PetscInt uOff_x[], const PetscScalar u[], const PetscScalar u_t[], const PetscScalar u_x[], const PetscInt aOff[], const PetscInt aOff_x[], const PetscScalar a[], const PetscScalar a_t[], const PetscScalar a_x[], PetscReal t, PetscReal u_tShift, const PetscReal x[], PetscInt numConstants, const PetscScalar constants[], PetscScalar g1[]) 113 { 114 for (PetscInt d = 0; d < dim; ++d) g1[d * dim + d] += 1.0; 115 } 116 117 static PetscErrorCode CreateFEM(DM dm, AppCtx *user) 118 { 119 PetscFE feu, feq; 120 PetscDS ds; 121 PetscBool simplex; 122 PetscInt dim; 123 124 PetscFunctionBeginUser; 125 PetscCall(DMGetDimension(dm, &dim)); 126 PetscCall(DMPlexIsSimplex(dm, &simplex)); 127 if (user->em == EM_MIXED) { 128 DMLabel label; 129 130 PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, simplex, "field_", PETSC_DETERMINE, &feq)); 131 PetscCall(PetscObjectSetName((PetscObject)feq, "field")); 132 PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, "potential_", PETSC_DETERMINE, &feu)); 133 PetscCall(PetscObjectSetName((PetscObject)feu, "potential")); 134 PetscCall(PetscFECopyQuadrature(feq, feu)); 135 PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feq)); 136 PetscCall(DMSetField(dm, 1, NULL, (PetscObject)feu)); 137 PetscCall(DMCreateDS(dm)); 138 PetscCall(PetscFEDestroy(&feu)); 139 PetscCall(PetscFEDestroy(&feq)); 140 141 PetscCall(DMGetLabel(dm, "marker", &label)); 142 PetscCall(DMGetDS(dm, &ds)); 143 PetscCall(PetscDSSetResidual(ds, 0, f0_q, f1_q)); 144 PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_qq, NULL, NULL, NULL)); 145 PetscCall(PetscDSSetJacobian(ds, 0, 1, NULL, NULL, g2_qu, NULL)); 146 PetscCall(PetscDSSetJacobian(ds, 1, 0, NULL, g1_uq, NULL, NULL)); 147 } else if (user->em == EM_PRIMAL) { 148 PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &feu)); 149 PetscCall(PetscObjectSetName((PetscObject)feu, "potential")); 150 PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feu)); 151 PetscCall(DMCreateDS(dm)); 152 PetscCall(PetscFEDestroy(&feu)); 153 PetscCall(DMGetDS(dm, &ds)); 154 PetscCall(PetscDSSetResidual(ds, 0, NULL, laplacian_f1)); 155 PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3)); 156 } 157 PetscFunctionReturn(PETSC_SUCCESS); 158 } 159 160 static PetscErrorCode CreatePoisson(DM dm, AppCtx *user) 161 { 162 SNES snes; 163 Mat J; 164 MatNullSpace nullSpace; 165 166 PetscFunctionBeginUser; 167 PetscCall(CreateFEM(dm, user)); 168 PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes)); 169 PetscCall(SNESSetOptionsPrefix(snes, "em_")); 170 PetscCall(SNESSetDM(snes, dm)); 171 PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, user)); 172 PetscCall(SNESSetFromOptions(snes)); 173 174 PetscCall(DMCreateMatrix(dm, &J)); 175 PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace)); 176 PetscCall(MatSetNullSpace(J, nullSpace)); 177 PetscCall(MatNullSpaceDestroy(&nullSpace)); 178 PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL)); 179 PetscCall(MatDestroy(&J)); 180 user->snes = snes; 181 PetscFunctionReturn(PETSC_SUCCESS); 182 } 183 184 static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw) 185 { 186 PetscReal v0[1] = {1.}; 187 PetscInt dim; 188 189 PetscFunctionBeginUser; 190 PetscCall(DMGetDimension(dm, &dim)); 191 PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw)); 192 PetscCall(DMSetType(*sw, DMSWARM)); 193 PetscCall(DMSetDimension(*sw, dim)); 194 PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC)); 195 PetscCall(DMSwarmSetCellDM(*sw, dm)); 196 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR)); 197 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL)); 198 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT)); 199 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "initCoordinates", dim, PETSC_REAL)); 200 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "initVelocity", dim, PETSC_REAL)); 201 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL)); 202 PetscCall(DMSwarmFinalizeFieldRegister(*sw)); 203 PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw)); 204 PetscCall(DMSwarmInitializeCoordinates(*sw)); 205 PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0)); 206 PetscCall(DMSetFromOptions(*sw)); 207 PetscCall(DMSetApplicationContext(*sw, user)); 208 PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles")); 209 PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view")); 210 { 211 Vec gc, gc0, gv, gv0; 212 213 PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc)); 214 PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initCoordinates", &gc0)); 215 PetscCall(VecCopy(gc, gc0)); 216 PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, DMSwarmPICField_coor, &gc)); 217 PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initCoordinates", &gc0)); 218 PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "velocity", &gv)); 219 PetscCall(DMSwarmCreateGlobalVectorFromField(*sw, "initVelocity", &gv0)); 220 PetscCall(VecCopy(gv, gv0)); 221 PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "velocity", &gv)); 222 PetscCall(DMSwarmDestroyGlobalVectorFromField(*sw, "initVelocity", &gv0)); 223 } 224 { 225 const char *fieldnames[2] = {DMSwarmPICField_coor, "velocity"}; 226 227 PetscCall(DMSwarmVectorDefineFields(*sw, 2, fieldnames)); 228 } 229 PetscFunctionReturn(PETSC_SUCCESS); 230 } 231 232 static PetscErrorCode ComputeFieldAtParticles_Coulomb(SNES snes, DM sw, PetscReal E[]) 233 { 234 PetscReal *coords; 235 PetscInt dim, d, Np, p, q; 236 PetscMPIInt size; 237 238 PetscFunctionBegin; 239 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)snes), &size)); 240 PetscCheck(size == 1, PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "Coulomb code only works in serial"); 241 PetscCall(DMGetDimension(sw, &dim)); 242 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 243 244 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 245 for (p = 0; p < Np; ++p) { 246 PetscReal *pcoord = &coords[p * dim]; 247 PetscReal *pE = &E[p * dim]; 248 /* Calculate field at particle p due to particle q */ 249 for (q = 0; q < Np; ++q) { 250 PetscReal *qcoord = &coords[q * dim]; 251 PetscReal rpq[3], r; 252 253 if (p == q) continue; 254 for (d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d]; 255 r = DMPlex_NormD_Internal(dim, rpq); 256 for (d = 0; d < dim; ++d) pE[d] += rpq[d] / PetscPowRealInt(r, 3); 257 } 258 } 259 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 260 PetscFunctionReturn(PETSC_SUCCESS); 261 } 262 263 static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, PetscReal E[]) 264 { 265 DM dm; 266 PetscDS ds; 267 PetscFE fe; 268 Mat M_p; 269 Vec phi, locPhi, rho, f; 270 PetscReal *coords, chargeTol = 1e-13; 271 PetscInt dim, d, cStart, cEnd, c, Np; 272 const char **oldFields; 273 PetscInt Nf; 274 const char **tmp; 275 276 PetscFunctionBegin; 277 PetscCall(DMGetDimension(sw, &dim)); 278 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 279 PetscCall(SNESGetDM(snes, &dm)); 280 281 PetscCall(DMSwarmVectorGetField(sw, &Nf, &tmp)); 282 PetscCall(PetscMalloc1(Nf, &oldFields)); 283 for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscStrallocpy(tmp[f], (char **)&oldFields[f])); 284 PetscCall(DMSwarmVectorDefineField(sw, "w_q")); 285 PetscCall(DMCreateMassMatrix(sw, dm, &M_p)); 286 PetscCall(DMSwarmVectorDefineFields(sw, Nf, oldFields)); 287 for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscFree(oldFields[f])); 288 PetscCall(PetscFree(oldFields)); 289 290 /* Create the charges rho */ 291 PetscCall(DMGetGlobalVector(dm, &rho)); 292 PetscCall(PetscObjectSetName((PetscObject)rho, "rho")); 293 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f)); 294 PetscCall(PetscObjectSetName((PetscObject)f, "particle weight")); 295 PetscCall(MatMultTranspose(M_p, f, rho)); 296 PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view")); 297 PetscCall(VecViewFromOptions(f, NULL, "-weights_view")); 298 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f)); 299 PetscCall(MatDestroy(&M_p)); 300 { 301 PetscScalar sum; 302 PetscInt n; 303 PetscReal phi_0 = 1.; /* (sigma*sigma*sigma)*(timeScale*timeScale)/(m_e*q_e*epsi_0)*/ 304 305 /* Remove constant from rho */ 306 PetscCall(VecGetSize(rho, &n)); 307 PetscCall(VecSum(rho, &sum)); 308 PetscCall(VecShift(rho, -sum / n)); 309 PetscCall(VecSum(rho, &sum)); 310 PetscCheck(PetscAbsScalar(sum) < chargeTol, PetscObjectComm((PetscObject)sw), PETSC_ERR_PLIB, "Charge should have no DC component: %g", (double)PetscRealPart(sum)); 311 /* Nondimensionalize rho */ 312 PetscCall(VecScale(rho, phi_0)); 313 } 314 PetscCall(VecViewFromOptions(rho, NULL, "-poisson_rho_view")); 315 316 PetscCall(DMGetGlobalVector(dm, &phi)); 317 PetscCall(PetscObjectSetName((PetscObject)phi, "potential")); 318 PetscCall(VecSet(phi, 0.0)); 319 PetscCall(SNESSolve(snes, rho, phi)); 320 PetscCall(DMRestoreGlobalVector(dm, &rho)); 321 PetscCall(VecViewFromOptions(phi, NULL, "-phi_view")); 322 323 PetscCall(DMGetLocalVector(dm, &locPhi)); 324 PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi)); 325 PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi)); 326 PetscCall(DMRestoreGlobalVector(dm, &phi)); 327 328 PetscCall(DMGetDS(dm, &ds)); 329 PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe)); 330 PetscCall(DMSwarmSortGetAccess(sw)); 331 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd)); 332 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 333 for (c = cStart; c < cEnd; ++c) { 334 PetscTabulation tab; 335 PetscScalar *clPhi = NULL; 336 PetscReal *pcoord, *refcoord; 337 PetscReal v[3], J[9], invJ[9], detJ; 338 PetscInt *points; 339 PetscInt Ncp, cp; 340 341 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points)); 342 PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord)); 343 PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord)); 344 for (cp = 0; cp < Ncp; ++cp) 345 for (d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d]; 346 PetscCall(DMPlexCoordinatesToReference(dm, c, Ncp, pcoord, refcoord)); 347 PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab)); 348 PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v, J, invJ, &detJ)); 349 PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi)); 350 for (cp = 0; cp < Ncp; ++cp) { 351 const PetscReal *basisDer = tab->T[1]; 352 const PetscInt p = points[cp]; 353 354 for (d = 0; d < dim; ++d) E[p * dim + d] = 0.; 355 PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, invJ, NULL, cp, &E[p * dim])); 356 for (d = 0; d < dim; ++d) E[p * dim + d] *= -1.0; 357 } 358 PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi)); 359 PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord)); 360 PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord)); 361 PetscCall(PetscTabulationDestroy(&tab)); 362 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points)); 363 } 364 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 365 PetscCall(DMSwarmSortRestoreAccess(sw)); 366 PetscCall(DMRestoreLocalVector(dm, &locPhi)); 367 PetscFunctionReturn(PETSC_SUCCESS); 368 } 369 370 static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, PetscReal E[]) 371 { 372 DM dm, potential_dm; 373 IS potential_IS; 374 PetscDS ds; 375 PetscFE fe; 376 PetscFEGeom feGeometry; 377 Mat M_p; 378 Vec phi, locPhi, rho, f, temp_rho; 379 PetscQuadrature q; 380 PetscReal *coords, chargeTol = 1e-13; 381 PetscInt dim, d, cStart, cEnd, c, Np, pot_field = 1; 382 const char **oldFields; 383 PetscInt Nf; 384 const char **tmp; 385 386 PetscFunctionBegin; 387 PetscCall(DMGetDimension(sw, &dim)); 388 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 389 390 /* Create the charges rho */ 391 PetscCall(SNESGetDM(snes, &dm)); 392 PetscCall(DMGetGlobalVector(dm, &rho)); 393 PetscCall(PetscObjectSetName((PetscObject)rho, "rho")); 394 PetscCall(DMCreateSubDM(dm, 1, &pot_field, &potential_IS, &potential_dm)); 395 396 PetscCall(DMSwarmVectorGetField(sw, &Nf, &tmp)); 397 PetscCall(PetscMalloc1(Nf, &oldFields)); 398 for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscStrallocpy(tmp[f], (char **)&oldFields[f])); 399 PetscCall(DMSwarmVectorDefineField(sw, "w_q")); 400 PetscCall(DMCreateMassMatrix(sw, potential_dm, &M_p)); 401 PetscCall(DMSwarmVectorDefineFields(sw, Nf, oldFields)); 402 for (PetscInt f = 0; f < Nf; ++f) PetscCall(PetscFree(oldFields[f])); 403 PetscCall(PetscFree(oldFields)); 404 405 PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view")); 406 PetscCall(DMGetGlobalVector(potential_dm, &temp_rho)); 407 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f)); 408 PetscCall(PetscObjectSetName((PetscObject)f, "particle weight")); 409 PetscCall(VecViewFromOptions(f, NULL, "-weights_view")); 410 PetscCall(MatMultTranspose(M_p, f, temp_rho)); 411 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f)); 412 PetscCall(MatDestroy(&M_p)); 413 PetscCall(PetscObjectSetName((PetscObject)rho, "rho")); 414 PetscCall(VecViewFromOptions(rho, NULL, "-poisson_rho_view")); 415 PetscCall(VecISCopy(rho, potential_IS, SCATTER_FORWARD, temp_rho)); 416 PetscCall(DMRestoreGlobalVector(potential_dm, &temp_rho)); 417 PetscCall(DMDestroy(&potential_dm)); 418 PetscCall(ISDestroy(&potential_IS)); 419 { 420 PetscScalar sum; 421 PetscInt n; 422 PetscReal phi_0 = 1.; /*(sigma*sigma*sigma)*(timeScale*timeScale)/(m_e*q_e*epsi_0);*/ 423 424 /*Remove constant from rho*/ 425 PetscCall(VecGetSize(rho, &n)); 426 PetscCall(VecSum(rho, &sum)); 427 PetscCall(VecShift(rho, -sum / n)); 428 PetscCall(VecSum(rho, &sum)); 429 PetscCheck(PetscAbsScalar(sum) < chargeTol, PetscObjectComm((PetscObject)sw), PETSC_ERR_PLIB, "Charge should have no DC component: %g", (double)PetscRealPart(sum)); 430 /* Nondimensionalize rho */ 431 PetscCall(VecScale(rho, phi_0)); 432 } 433 PetscCall(DMGetGlobalVector(dm, &phi)); 434 PetscCall(PetscObjectSetName((PetscObject)phi, "potential")); 435 PetscCall(VecSet(phi, 0.0)); 436 PetscCall(SNESSolve(snes, NULL, phi)); 437 PetscCall(DMRestoreGlobalVector(dm, &rho)); 438 PetscCall(VecViewFromOptions(phi, NULL, "-phi_view")); 439 440 PetscCall(DMGetLocalVector(dm, &locPhi)); 441 PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi)); 442 PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi)); 443 PetscCall(DMRestoreGlobalVector(dm, &phi)); 444 445 PetscCall(DMGetDS(dm, &ds)); 446 PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe)); 447 PetscCall(DMSwarmSortGetAccess(sw)); 448 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd)); 449 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 450 for (c = cStart; c < cEnd; ++c) { 451 PetscTabulation tab; 452 PetscScalar *clPhi = NULL; 453 PetscReal *pcoord, *refcoord; 454 PetscReal v[3], J[9], invJ[9], detJ; 455 PetscInt *points; 456 PetscInt Ncp, cp; 457 458 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points)); 459 PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord)); 460 PetscCall(DMGetWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord)); 461 for (cp = 0; cp < Ncp; ++cp) 462 for (d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d]; 463 PetscCall(DMPlexCoordinatesToReference(dm, c, Ncp, pcoord, refcoord)); 464 PetscCall(PetscFECreateTabulation(fe, 1, Ncp, refcoord, 1, &tab)); 465 PetscCall(DMPlexComputeCellGeometryFEM(dm, c, NULL, v, J, invJ, &detJ)); 466 PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi)); 467 for (cp = 0; cp < Ncp; ++cp) { 468 const PetscInt p = points[cp]; 469 470 for (d = 0; d < dim; ++d) E[p * dim + d] = 0.; 471 PetscCall(PetscFEGetQuadrature(fe, &q)); 472 PetscCall(PetscFECreateCellGeometry(fe, q, &feGeometry)); 473 PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, &feGeometry, cp, &E[p * dim])); 474 PetscCall(PetscFEDestroyCellGeometry(fe, &feGeometry)); 475 } 476 PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi)); 477 PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &pcoord)); 478 PetscCall(DMRestoreWorkArray(dm, Ncp * dim, MPIU_REAL, &refcoord)); 479 PetscCall(PetscTabulationDestroy(&tab)); 480 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points)); 481 } 482 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 483 PetscCall(DMSwarmSortRestoreAccess(sw)); 484 PetscCall(DMRestoreLocalVector(dm, &locPhi)); 485 PetscFunctionReturn(PETSC_SUCCESS); 486 } 487 488 static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw, PetscReal E[]) 489 { 490 AppCtx *ctx; 491 PetscInt dim, Np; 492 493 PetscFunctionBegin; 494 PetscValidHeaderSpecific(snes, SNES_CLASSID, 1); 495 PetscValidHeaderSpecific(sw, DM_CLASSID, 2); 496 PetscAssertPointer(E, 3); 497 PetscCall(DMGetDimension(sw, &dim)); 498 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 499 PetscCall(DMGetApplicationContext(sw, &ctx)); 500 PetscCall(PetscArrayzero(E, Np * dim)); 501 502 switch (ctx->em) { 503 case EM_PRIMAL: 504 PetscCall(ComputeFieldAtParticles_Primal(snes, sw, E)); 505 break; 506 case EM_COULOMB: 507 PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E)); 508 break; 509 case EM_MIXED: 510 PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, E)); 511 break; 512 case EM_NONE: 513 break; 514 default: 515 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[ctx->em]); 516 } 517 PetscFunctionReturn(PETSC_SUCCESS); 518 } 519 520 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx) 521 { 522 DM sw; 523 SNES snes = ((AppCtx *)ctx)->snes; 524 const PetscReal *coords, *vel; 525 const PetscScalar *u; 526 PetscScalar *g; 527 PetscReal *E; 528 PetscInt dim, d, Np, p; 529 530 PetscFunctionBeginUser; 531 PetscCall(TSGetDM(ts, &sw)); 532 PetscCall(DMGetDimension(sw, &dim)); 533 PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords)); 534 PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel)); 535 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E)); 536 PetscCall(VecGetLocalSize(U, &Np)); 537 PetscCall(VecGetArrayRead(U, &u)); 538 PetscCall(VecGetArray(G, &g)); 539 540 PetscCall(ComputeFieldAtParticles(snes, sw, E)); 541 542 Np /= 2 * dim; 543 for (p = 0; p < Np; ++p) { 544 const PetscReal x0 = coords[p * dim + 0]; 545 const PetscReal vy0 = vel[p * dim + 1]; 546 const PetscReal omega = vy0 / x0; 547 548 for (d = 0; d < dim; ++d) { 549 g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d]; 550 g[(p * 2 + 1) * dim + d] = E[p * dim + d] - PetscSqr(omega) * u[(p * 2 + 0) * dim + d]; 551 } 552 } 553 PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords)); 554 PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel)); 555 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E)); 556 PetscCall(VecRestoreArrayRead(U, &u)); 557 PetscCall(VecRestoreArray(G, &g)); 558 PetscFunctionReturn(PETSC_SUCCESS); 559 } 560 561 /* J_{ij} = dF_i/dx_j 562 J_p = ( 0 1) 563 (-w^2 0) 564 TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator. 565 Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb 566 */ 567 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx) 568 { 569 DM sw; 570 const PetscReal *coords, *vel; 571 PetscInt dim, d, Np, p, rStart; 572 573 PetscFunctionBeginUser; 574 PetscCall(TSGetDM(ts, &sw)); 575 PetscCall(DMGetDimension(sw, &dim)); 576 PetscCall(VecGetLocalSize(U, &Np)); 577 PetscCall(MatGetOwnershipRange(J, &rStart, NULL)); 578 PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords)); 579 PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel)); 580 Np /= 2 * dim; 581 for (p = 0; p < Np; ++p) { 582 const PetscReal x0 = coords[p * dim + 0]; 583 const PetscReal vy0 = vel[p * dim + 1]; 584 const PetscReal omega = vy0 / x0; 585 PetscScalar vals[4] = {0., 1., -PetscSqr(omega), 0.}; 586 587 for (d = 0; d < dim; ++d) { 588 const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart}; 589 PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES)); 590 } 591 } 592 PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords)); 593 PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel)); 594 PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY)); 595 PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY)); 596 PetscFunctionReturn(PETSC_SUCCESS); 597 } 598 599 static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx) 600 { 601 DM sw; 602 const PetscScalar *v; 603 PetscScalar *xres; 604 PetscInt Np, p, dim, d; 605 606 PetscFunctionBeginUser; 607 PetscCall(TSGetDM(ts, &sw)); 608 PetscCall(DMGetDimension(sw, &dim)); 609 PetscCall(VecGetLocalSize(Xres, &Np)); 610 Np /= dim; 611 PetscCall(VecGetArrayRead(V, &v)); 612 PetscCall(VecGetArray(Xres, &xres)); 613 for (p = 0; p < Np; ++p) { 614 for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d]; 615 } 616 PetscCall(VecRestoreArrayRead(V, &v)); 617 PetscCall(VecRestoreArray(Xres, &xres)); 618 PetscFunctionReturn(PETSC_SUCCESS); 619 } 620 621 static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx) 622 { 623 DM sw; 624 SNES snes = ((AppCtx *)ctx)->snes; 625 const PetscScalar *x; 626 const PetscReal *coords, *vel; 627 PetscScalar *vres; 628 PetscReal *E; 629 PetscInt Np, p, dim, d; 630 631 PetscFunctionBeginUser; 632 PetscCall(TSGetDM(ts, &sw)); 633 PetscCall(DMGetDimension(sw, &dim)); 634 PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords)); 635 PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel)); 636 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E)); 637 PetscCall(VecGetLocalSize(Vres, &Np)); 638 PetscCall(VecGetArrayRead(X, &x)); 639 PetscCall(VecGetArray(Vres, &vres)); 640 PetscCheck(dim == 2, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension must be 2"); 641 642 PetscCall(ComputeFieldAtParticles(snes, sw, E)); 643 644 Np /= dim; 645 for (p = 0; p < Np; ++p) { 646 const PetscReal x0 = coords[p * dim + 0]; 647 const PetscReal vy0 = vel[p * dim + 1]; 648 const PetscReal omega = vy0 / x0; 649 650 for (d = 0; d < dim; ++d) vres[p * dim + d] = E[p * dim + d] - PetscSqr(omega) * x[p * dim + d]; 651 } 652 PetscCall(VecRestoreArrayRead(X, &x)); 653 PetscCall(VecRestoreArray(Vres, &vres)); 654 PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords)); 655 PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel)); 656 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E)); 657 PetscFunctionReturn(PETSC_SUCCESS); 658 } 659 660 static PetscErrorCode CreateSolution(TS ts) 661 { 662 DM sw; 663 Vec u; 664 PetscInt dim, Np; 665 666 PetscFunctionBegin; 667 PetscCall(TSGetDM(ts, &sw)); 668 PetscCall(DMGetDimension(sw, &dim)); 669 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 670 PetscCall(VecCreate(PETSC_COMM_WORLD, &u)); 671 PetscCall(VecSetBlockSize(u, dim)); 672 PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE)); 673 PetscCall(VecSetUp(u)); 674 PetscCall(TSSetSolution(ts, u)); 675 PetscCall(VecDestroy(&u)); 676 PetscFunctionReturn(PETSC_SUCCESS); 677 } 678 679 static PetscErrorCode SetProblem(TS ts) 680 { 681 AppCtx *user; 682 DM sw; 683 684 PetscFunctionBegin; 685 PetscCall(TSGetDM(ts, &sw)); 686 PetscCall(DMGetApplicationContext(sw, (void **)&user)); 687 // Define unified system for (X, V) 688 { 689 Mat J; 690 PetscInt dim, Np; 691 692 PetscCall(DMGetDimension(sw, &dim)); 693 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 694 PetscCall(MatCreate(PETSC_COMM_WORLD, &J)); 695 PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE)); 696 PetscCall(MatSetBlockSize(J, 2 * dim)); 697 PetscCall(MatSetFromOptions(J)); 698 PetscCall(MatSetUp(J)); 699 PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user)); 700 PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user)); 701 PetscCall(MatDestroy(&J)); 702 } 703 /* Define split system for X and V */ 704 { 705 Vec u; 706 IS isx, isv, istmp; 707 const PetscInt *idx; 708 PetscInt dim, Np, rstart; 709 710 PetscCall(TSGetSolution(ts, &u)); 711 PetscCall(DMGetDimension(sw, &dim)); 712 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 713 PetscCall(VecGetOwnershipRange(u, &rstart, NULL)); 714 PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp)); 715 PetscCall(ISGetIndices(istmp, &idx)); 716 PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx)); 717 PetscCall(ISRestoreIndices(istmp, &idx)); 718 PetscCall(ISDestroy(&istmp)); 719 PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp)); 720 PetscCall(ISGetIndices(istmp, &idx)); 721 PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv)); 722 PetscCall(ISRestoreIndices(istmp, &idx)); 723 PetscCall(ISDestroy(&istmp)); 724 PetscCall(TSRHSSplitSetIS(ts, "position", isx)); 725 PetscCall(TSRHSSplitSetIS(ts, "momentum", isv)); 726 PetscCall(ISDestroy(&isx)); 727 PetscCall(ISDestroy(&isv)); 728 PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user)); 729 PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user)); 730 } 731 PetscFunctionReturn(PETSC_SUCCESS); 732 } 733 734 PetscErrorCode circleSingleX(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx) 735 { 736 x[0] = p + 1.; 737 x[1] = 0.; 738 return PETSC_SUCCESS; 739 } 740 741 PetscErrorCode circleSingleV(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar v[], void *ctx) 742 { 743 v[0] = 0.; 744 v[1] = PetscSqrtReal(1000. / (p + 1.)); 745 return PETSC_SUCCESS; 746 } 747 748 /* Put 5 particles into each circle */ 749 PetscErrorCode circleMultipleX(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx) 750 { 751 const PetscInt n = 5; 752 const PetscReal r0 = (p / n) + 1.; 753 const PetscReal th0 = (2. * PETSC_PI * (p % n)) / n; 754 755 x[0] = r0 * PetscCosReal(th0); 756 x[1] = r0 * PetscSinReal(th0); 757 return PETSC_SUCCESS; 758 } 759 760 /* Put 5 particles into each circle */ 761 PetscErrorCode circleMultipleV(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar v[], void *ctx) 762 { 763 const PetscInt n = 5; 764 const PetscReal r0 = (p / n) + 1.; 765 const PetscReal th0 = (2. * PETSC_PI * (p % n)) / n; 766 const PetscReal omega = PetscSqrtReal(1000. / r0) / r0; 767 768 v[0] = -r0 * omega * PetscSinReal(th0); 769 v[1] = r0 * omega * PetscCosReal(th0); 770 return PETSC_SUCCESS; 771 } 772 773 /* 774 InitializeSolveAndSwarm - Set the solution values to the swarm coordinates and velocities, and also possibly set the initial values. 775 776 Input Parameters: 777 + ts - The TS 778 - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem 779 780 Output Parameter: 781 . u - The initialized solution vector 782 783 Level: advanced 784 785 .seealso: InitializeSolve() 786 */ 787 static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial) 788 { 789 DM sw; 790 Vec u, gc, gv, gc0, gv0; 791 IS isx, isv; 792 AppCtx *user; 793 794 PetscFunctionBeginUser; 795 PetscCall(TSGetDM(ts, &sw)); 796 PetscCall(DMGetApplicationContext(sw, &user)); 797 if (useInitial) { 798 PetscReal v0[1] = {1.}; 799 800 PetscCall(DMSwarmInitializeCoordinates(sw)); 801 PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0)); 802 PetscCall(DMSwarmMigrate(sw, PETSC_TRUE)); 803 PetscCall(TSReset(ts)); 804 PetscCall(CreateSolution(ts)); 805 PetscCall(SetProblem(ts)); 806 } 807 PetscCall(TSGetSolution(ts, &u)); 808 PetscCall(TSRHSSplitGetIS(ts, "position", &isx)); 809 PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv)); 810 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc)); 811 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initCoordinates", &gc0)); 812 if (useInitial) PetscCall(VecCopy(gc, gc0)); 813 PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc)); 814 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc)); 815 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initCoordinates", &gc0)); 816 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv)); 817 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "initVelocity", &gv0)); 818 if (useInitial) PetscCall(VecCopy(gv, gv0)); 819 PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv)); 820 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv)); 821 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "initVelocity", &gv0)); 822 PetscFunctionReturn(PETSC_SUCCESS); 823 } 824 825 static PetscErrorCode InitializeSolve(TS ts, Vec u) 826 { 827 PetscFunctionBegin; 828 PetscCall(TSSetSolution(ts, u)); 829 PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE)); 830 PetscFunctionReturn(PETSC_SUCCESS); 831 } 832 833 static PetscErrorCode ComputeError(TS ts, Vec U, Vec E) 834 { 835 MPI_Comm comm; 836 DM sw; 837 AppCtx *user; 838 const PetscScalar *u; 839 const PetscReal *coords, *vel; 840 PetscScalar *e; 841 PetscReal t; 842 PetscInt dim, Np, p; 843 844 PetscFunctionBeginUser; 845 PetscCall(PetscObjectGetComm((PetscObject)ts, &comm)); 846 PetscCall(TSGetDM(ts, &sw)); 847 PetscCall(DMGetApplicationContext(sw, &user)); 848 PetscCall(DMGetDimension(sw, &dim)); 849 PetscCall(TSGetSolveTime(ts, &t)); 850 PetscCall(VecGetArray(E, &e)); 851 PetscCall(VecGetArrayRead(U, &u)); 852 PetscCall(VecGetLocalSize(U, &Np)); 853 PetscCall(DMSwarmGetField(sw, "initCoordinates", NULL, NULL, (void **)&coords)); 854 PetscCall(DMSwarmGetField(sw, "initVelocity", NULL, NULL, (void **)&vel)); 855 Np /= 2 * dim; 856 for (p = 0; p < Np; ++p) { 857 /* TODO generalize initial conditions and project into plane instead of assuming x-y */ 858 const PetscReal r0 = DMPlex_NormD_Internal(dim, &coords[p * dim]); 859 const PetscReal th0 = PetscAtan2Real(coords[p * dim + 1], coords[p * dim + 0]); 860 const PetscReal v0 = DMPlex_NormD_Internal(dim, &vel[p * dim]); 861 const PetscReal omega = v0 / r0; 862 const PetscReal ct = PetscCosReal(omega * t + th0); 863 const PetscReal st = PetscSinReal(omega * t + th0); 864 const PetscScalar *x = &u[(p * 2 + 0) * dim]; 865 const PetscScalar *v = &u[(p * 2 + 1) * dim]; 866 const PetscReal xe[3] = {r0 * ct, r0 * st, 0.0}; 867 const PetscReal ve[3] = {-v0 * st, v0 * ct, 0.0}; 868 PetscInt d; 869 870 for (d = 0; d < dim; ++d) { 871 e[(p * 2 + 0) * dim + d] = x[d] - xe[d]; 872 e[(p * 2 + 1) * dim + d] = v[d] - ve[d]; 873 } 874 if (user->error) { 875 const PetscReal en = 0.5 * DMPlex_DotRealD_Internal(dim, v, v); 876 const PetscReal exen = 0.5 * PetscSqr(v0); 877 PetscCall(PetscPrintf(comm, "t %.4g: p%" PetscInt_FMT " error [%.2f %.2f] sol [(%.6lf %.6lf) (%.6lf %.6lf)] exact [(%.6lf %.6lf) (%.6lf %.6lf)] energy/exact energy %g / %g (%.10lf%%)\n", (double)t, p, (double)DMPlex_NormD_Internal(dim, &e[(p * 2 + 0) * dim]), (double)DMPlex_NormD_Internal(dim, &e[(p * 2 + 1) * dim]), (double)x[0], (double)x[1], (double)v[0], (double)v[1], (double)xe[0], (double)xe[1], (double)ve[0], (double)ve[1], (double)en, (double)exen, (double)(PetscAbsReal(exen - en) * 100. / exen))); 878 } 879 } 880 PetscCall(DMSwarmRestoreField(sw, "initCoordinates", NULL, NULL, (void **)&coords)); 881 PetscCall(DMSwarmRestoreField(sw, "initVelocity", NULL, NULL, (void **)&vel)); 882 PetscCall(VecRestoreArrayRead(U, &u)); 883 PetscCall(VecRestoreArray(E, &e)); 884 PetscFunctionReturn(PETSC_SUCCESS); 885 } 886 887 static PetscErrorCode EnergyMonitor(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx) 888 { 889 const PetscInt ostep = ((AppCtx *)ctx)->ostep; 890 const EMType em = ((AppCtx *)ctx)->em; 891 DM sw; 892 const PetscScalar *u; 893 PetscReal *coords, *E; 894 PetscReal enKin = 0., enEM = 0.; 895 PetscInt dim, d, Np, p, q; 896 897 PetscFunctionBeginUser; 898 if (step % ostep == 0) { 899 PetscCall(TSGetDM(ts, &sw)); 900 PetscCall(DMGetDimension(sw, &dim)); 901 PetscCall(VecGetArrayRead(U, &u)); 902 PetscCall(VecGetLocalSize(U, &Np)); 903 Np /= 2 * dim; 904 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 905 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E)); 906 if (!step) PetscCall(PetscPrintf(PetscObjectComm((PetscObject)ts), "Time Step Part Energy\n")); 907 for (p = 0; p < Np; ++p) { 908 const PetscReal v2 = DMPlex_DotRealD_Internal(dim, &u[(p * 2 + 1) * dim], &u[(p * 2 + 1) * dim]); 909 PetscReal *pcoord = &coords[p * dim]; 910 911 PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " %5" PetscInt_FMT " %10.4lf\n", (double)t, step, p, (double)(0.5 * v2))); 912 enKin += 0.5 * v2; 913 if (em == EM_NONE) { 914 continue; 915 } else if (em == EM_COULOMB) { 916 for (q = p + 1; q < Np; ++q) { 917 PetscReal *qcoord = &coords[q * dim]; 918 PetscReal rpq[3], r; 919 for (d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d]; 920 r = DMPlex_NormD_Internal(dim, rpq); 921 enEM += 1. / r; 922 } 923 } else if (em == EM_PRIMAL || em == EM_MIXED) { 924 for (d = 0; d < dim; ++d) enEM += E[p * dim + d]; 925 } 926 } 927 PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " KE\t %10.4lf\n", (double)t, step, (double)enKin)); 928 PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " PE\t %1.10g\n", (double)t, step, (double)enEM)); 929 PetscCall(PetscSynchronizedPrintf(PetscObjectComm((PetscObject)ts), "%.6lf %4" PetscInt_FMT " E\t %10.4lf\n", (double)t, step, (double)(enKin + enEM))); 930 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 931 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E)); 932 PetscCall(PetscSynchronizedFlush(PetscObjectComm((PetscObject)ts), NULL)); 933 PetscCall(VecRestoreArrayRead(U, &u)); 934 } 935 PetscFunctionReturn(PETSC_SUCCESS); 936 } 937 938 static PetscErrorCode SetUpMigrateParticles(TS ts, PetscInt n, PetscReal t, Vec x, PetscBool *flg, void *ctx) 939 { 940 DM sw; 941 942 PetscFunctionBeginUser; 943 *flg = PETSC_TRUE; 944 PetscCall(TSGetDM(ts, &sw)); 945 PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre")); 946 { 947 Vec u, gc, gv; 948 IS isx, isv; 949 950 PetscCall(TSGetSolution(ts, &u)); 951 PetscCall(TSRHSSplitGetIS(ts, "position", &isx)); 952 PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv)); 953 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc)); 954 PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc)); 955 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc)); 956 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv)); 957 PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv)); 958 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv)); 959 } 960 PetscFunctionReturn(PETSC_SUCCESS); 961 } 962 963 static PetscErrorCode MigrateParticles(TS ts, PetscInt nv, Vec vecsin[], Vec vecsout[], void *ctx) 964 { 965 DM sw; 966 967 PetscFunctionBeginUser; 968 PetscCall(TSGetDM(ts, &sw)); 969 PetscCall(DMSwarmMigrate(sw, PETSC_TRUE)); 970 PetscCall(CreateSolution(ts)); 971 PetscCall(SetProblem(ts)); 972 PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE)); 973 PetscFunctionReturn(PETSC_SUCCESS); 974 } 975 976 int main(int argc, char **argv) 977 { 978 DM dm, sw; 979 TS ts; 980 Vec u; 981 AppCtx user; 982 983 PetscCall(PetscInitialize(&argc, &argv, NULL, help)); 984 PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user)); 985 PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm)); 986 PetscCall(CreatePoisson(dm, &user)); 987 PetscCall(CreateSwarm(dm, &user, &sw)); 988 PetscCall(DMSetApplicationContext(sw, &user)); 989 990 PetscCall(TSCreate(PETSC_COMM_WORLD, &ts)); 991 PetscCall(TSSetProblemType(ts, TS_NONLINEAR)); 992 PetscCall(TSSetDM(ts, sw)); 993 PetscCall(TSSetMaxTime(ts, 0.1)); 994 PetscCall(TSSetTimeStep(ts, 0.00001)); 995 PetscCall(TSSetMaxSteps(ts, 100)); 996 PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP)); 997 PetscCall(TSMonitorSet(ts, EnergyMonitor, &user, NULL)); 998 PetscCall(TSSetFromOptions(ts)); 999 PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve)); 1000 PetscCall(TSSetComputeExactError(ts, ComputeError)); 1001 PetscCall(TSSetResize(ts, PETSC_FALSE, SetUpMigrateParticles, MigrateParticles, NULL)); 1002 1003 PetscCall(CreateSolution(ts)); 1004 PetscCall(TSGetSolution(ts, &u)); 1005 PetscCall(TSComputeInitialCondition(ts, u)); 1006 PetscCall(TSSolve(ts, NULL)); 1007 1008 PetscCall(SNESDestroy(&user.snes)); 1009 PetscCall(TSDestroy(&ts)); 1010 PetscCall(DMDestroy(&sw)); 1011 PetscCall(DMDestroy(&dm)); 1012 PetscCall(PetscFinalize()); 1013 return 0; 1014 } 1015 1016 /*TEST 1017 1018 build: 1019 requires: double !complex 1020 1021 testset: 1022 requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT) 1023 args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 1,1 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 \ 1024 -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \ 1025 -ts_type basicsymplectic\ 1026 -dm_view -output_step 50 -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10 1027 test: 1028 suffix: none_bsi_2d_1 1029 args: -ts_basicsymplectic_type 1 -em_type none -error 1030 test: 1031 suffix: none_bsi_2d_2 1032 args: -ts_basicsymplectic_type 2 -em_type none -error 1033 test: 1034 suffix: none_bsi_2d_3 1035 args: -ts_basicsymplectic_type 3 -em_type none -error 1036 test: 1037 suffix: none_bsi_2d_4 1038 args: -ts_basicsymplectic_type 4 -em_type none -error 1039 test: 1040 suffix: coulomb_bsi_2d_1 1041 args: -ts_basicsymplectic_type 1 1042 test: 1043 suffix: coulomb_bsi_2d_2 1044 args: -ts_basicsymplectic_type 2 1045 test: 1046 suffix: coulomb_bsi_2d_3 1047 args: -ts_basicsymplectic_type 3 1048 test: 1049 suffix: coulomb_bsi_2d_4 1050 args: -ts_basicsymplectic_type 4 1051 1052 testset: 1053 requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT) 1054 args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 1,1 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 \ 1055 -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \ 1056 -ts_type basicsymplectic\ 1057 -em_type primal -em_pc_type svd\ 1058 -dm_view -output_step 50 -error -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10\ 1059 -petscspace_degree 2 -petscfe_default_quadrature_order 3 -sigma 1.0e-8 -timeScale 2.0e-14 1060 test: 1061 suffix: poisson_bsi_2d_1 1062 args: -ts_basicsymplectic_type 1 1063 test: 1064 suffix: poisson_bsi_2d_2 1065 args: -ts_basicsymplectic_type 2 1066 test: 1067 suffix: poisson_bsi_2d_3 1068 args: -ts_basicsymplectic_type 3 1069 test: 1070 suffix: poisson_bsi_2d_4 1071 args: -ts_basicsymplectic_type 4 1072 1073 testset: 1074 requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT) 1075 args: -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \ 1076 -ts_type theta -ts_theta_theta 0.5 -ts_convergence_estimate -convest_num_refine 2 \ 1077 -mat_type baij -ksp_error_if_not_converged -em_pc_type svd\ 1078 -dm_view -output_step 50 -error -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10\ 1079 -pc_type svd -sigma 1.0e-8 -timeScale 2.0e-14 1080 test: 1081 suffix: im_2d_0 1082 args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 1,1 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 1083 1084 testset: 1085 requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT) 1086 args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 10,10 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 -petscpartitioner_type simple \ 1087 -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV -dm_swarm_num_species 1\ 1088 -ts_type basicsymplectic -ts_convergence_estimate -convest_num_refine 2 \ 1089 -em_snes_type ksponly -em_pc_type svd -em_type primal -petscspace_degree 1\ 1090 -dm_view -output_step 50\ 1091 -pc_type svd -sigma 1.0e-8 -timeScale 2.0e-14 -ts_dt 0.01 -ts_max_time 1.0 -ts_max_steps 10 1092 test: 1093 suffix: bsi_2d_mesh_1 1094 args: -ts_basicsymplectic_type 4 1095 test: 1096 suffix: bsi_2d_mesh_1_par_2 1097 nsize: 2 1098 args: -ts_basicsymplectic_type 4 1099 test: 1100 suffix: bsi_2d_mesh_1_par_3 1101 nsize: 3 1102 args: -ts_basicsymplectic_type 4 1103 test: 1104 suffix: bsi_2d_mesh_1_par_4 1105 nsize: 4 1106 args: -ts_basicsymplectic_type 4 -dm_swarm_num_particles 0,0,2,0 1107 1108 testset: 1109 requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT) 1110 args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 10,10 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 \ 1111 -dm_swarm_num_particles 10 -dm_swarm_coordinate_function circleMultipleX -dm_swarm_velocity_function circleMultipleV \ 1112 -ts_convergence_estimate -convest_num_refine 2 \ 1113 -em_pc_type lu\ 1114 -dm_view -output_step 50 -error\ 1115 -pc_type svd -sigma 1.0e-8 -timeScale 2.0e-14 -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10 1116 test: 1117 suffix: bsi_2d_multiple_1 1118 args: -ts_type basicsymplectic -ts_basicsymplectic_type 1 1119 test: 1120 suffix: bsi_2d_multiple_2 1121 args: -ts_type basicsymplectic -ts_basicsymplectic_type 2 1122 test: 1123 suffix: bsi_2d_multiple_3 1124 args: -ts_type basicsymplectic -ts_basicsymplectic_type 3 -ts_dt 0.001 1125 test: 1126 suffix: im_2d_multiple_0 1127 args: -ts_type theta -ts_theta_theta 0.5 \ 1128 -mat_type baij -ksp_error_if_not_converged -em_pc_type lu 1129 1130 testset: 1131 requires: defined(PETSC_HAVE_EXECUTABLE_EXPORT) 1132 args: -dm_plex_dim 2 -dm_plex_simplex 0 -dm_plex_box_faces 2,2 -dm_plex_box_lower -5,-5 -dm_plex_box_upper 5,5 \ 1133 -dm_swarm_num_particles 2 -dm_swarm_coordinate_function circleSingleX -dm_swarm_velocity_function circleSingleV \ 1134 -em_pc_type fieldsplit -ksp_rtol 1e-10 -em_ksp_type preonly -em_type mixed -em_ksp_error_if_not_converged\ 1135 -dm_view -output_step 50 -error -dm_refine 0\ 1136 -pc_type svd -sigma 1.0e-8 -timeScale 2.0e-14 -ts_dt 0.01 -ts_max_time 10.0 -ts_max_steps 10 1137 test: 1138 suffix: bsi_4_rt_1 1139 args: -ts_type basicsymplectic -ts_basicsymplectic_type 4\ 1140 -pc_fieldsplit_detect_saddle_point\ 1141 -pc_type fieldsplit\ 1142 -pc_fieldsplit_type schur\ 1143 -pc_fieldsplit_schur_precondition full \ 1144 -field_petscspace_degree 2\ 1145 -field_petscfe_default_quadrature_order 1\ 1146 -field_petscspace_type sum \ 1147 -field_petscspace_variables 2 \ 1148 -field_petscspace_components 2 \ 1149 -field_petscspace_sum_spaces 2 \ 1150 -field_petscspace_sum_concatenate true \ 1151 -field_sumcomp_0_petscspace_variables 2 \ 1152 -field_sumcomp_0_petscspace_type tensor \ 1153 -field_sumcomp_0_petscspace_tensor_spaces 2 \ 1154 -field_sumcomp_0_petscspace_tensor_uniform false \ 1155 -field_sumcomp_0_tensorcomp_0_petscspace_degree 1 \ 1156 -field_sumcomp_0_tensorcomp_1_petscspace_degree 0 \ 1157 -field_sumcomp_1_petscspace_variables 2 \ 1158 -field_sumcomp_1_petscspace_type tensor \ 1159 -field_sumcomp_1_petscspace_tensor_spaces 2 \ 1160 -field_sumcomp_1_petscspace_tensor_uniform false \ 1161 -field_sumcomp_1_tensorcomp_0_petscspace_degree 0 \ 1162 -field_sumcomp_1_tensorcomp_1_petscspace_degree 1 \ 1163 -field_petscdualspace_form_degree -1 \ 1164 -field_petscdualspace_order 1 \ 1165 -field_petscdualspace_lagrange_trimmed true\ 1166 -ksp_gmres_restart 500 1167 1168 TEST*/ 1169