1 static char help[] = "Two-level system for Landau Damping using Vlasov-Poisson equations\n"; 2 3 /* 4 Moment Equations: 5 6 We will discretize the moment equations using finite elements, and we will project the moments into the finite element space We will use the PFAK method, which guarantees that our FE approximation is weakly equivalent to the true moment. The first moment, number density, is given by 7 8 \int dx \phi_i n_f = \int dx \phi_i n_p 9 \int dx \phi_i n_f = \int dx \phi_i \int dv f 10 \int dx \phi_i n_f = \int dx \phi_i \int dv \sum_p w_p \delta(x - x_p) \delta(v - v_p) 11 \int dx \phi_i n_f = \int dx \phi_i \sum_p w_p \delta(x - x_p) 12 M n_F = M_p w_p 13 14 where 15 16 (M_p){ip} = \phi_i(x_p) 17 18 which is just a scaled version of the charge density. The second moment, momentum density, is given by 19 20 \int dx \phi_i p_f = m \int dx \phi_i \int dv v f 21 \int dx \phi_i p_f = m \int dx \phi_i \sum_p w_p \delta(x - x_p) v_p 22 M p_F = M_p v_p w_p 23 24 And finally the third moment, pressure, is given by 25 26 \int dx \phi_i pr_f = m \int dx \phi_i \int dv (v - u)^2 f 27 \int dx \phi_i pr_f = m \int dx \phi_i \sum_p w_p \delta(x - x_p) (v_p - u)^2 28 M pr_F = M_p (v_p - u)^2 w_p 29 = M_p (v_p - p_F(x_p) / m n_F(x_p))^2 w_p 30 = M_p (v_p - (\sum_j p_F \phi_j(x_p)) / m (\sum_k n_F \phi_k(x_p)))^2 w_p 31 32 Here we need all FEM basis functions \phi_i that see that particle p. 33 34 To run the code with particles sinusoidally perturbed in x space use the test "pp_poisson_bsi_1d_4" or "pp_poisson_bsi_2d_4" 35 According to Lukas, good damping results come at ~16k particles per cell 36 37 Swarm CellDMs 38 ============= 39 Name: "space" 40 Fields: DMSwarmPICField_coor, "velocity" 41 Coordinates: DMSwarmPICField_coor 42 43 Name: "velocity" 44 Fields: "w_q" 45 Coordinates: "velocity" 46 47 Name: "moments" 48 Fields: "w_q" 49 Coordinates: DMSwarmPICField_coor 50 51 Name: "moment fields" 52 Fields: "velocity" 53 Coordinates: DMSwarmPICField_coor 54 55 To visualize the maximum electric field use 56 57 -efield_monitor 58 59 To monitor velocity moments of the distribution use 60 61 -ptof_pc_type lu -moments_monitor 62 63 To monitor the particle positions in phase space use 64 65 -positions_monitor 66 67 To monitor the charge density, E field, and potential use 68 69 -poisson_monitor 70 71 To monitor the remapping field use 72 73 -remap_uf_view draw 74 75 To visualize the swarm distribution use 76 77 -ts_monitor_hg_swarm 78 79 To visualize the particles, we can use 80 81 -ts_monitor_sp_swarm -ts_monitor_sp_swarm_retain 0 -ts_monitor_sp_swarm_phase 1 -draw_size 500,500 82 */ 83 #include <petsctao.h> 84 #include <petscts.h> 85 #include <petscdmplex.h> 86 #include <petscdmswarm.h> 87 #include <petscfe.h> 88 #include <petscds.h> 89 #include <petscbag.h> 90 #include <petscdraw.h> 91 #include <petsc/private/petscfeimpl.h> /* For interpolation */ 92 #include <petsc/private/dmswarmimpl.h> /* For swarm debugging */ 93 #include "petscdm.h" 94 #include "petscdmlabel.h" 95 96 PETSC_EXTERN PetscErrorCode stream(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *); 97 PETSC_EXTERN PetscErrorCode line(PetscInt, PetscReal, const PetscReal[], PetscInt, PetscScalar[], void *); 98 99 const char *EMTypes[] = {"primal", "mixed", "coulomb", "none", "EMType", "EM_", NULL}; 100 typedef enum { 101 EM_PRIMAL, 102 EM_MIXED, 103 EM_COULOMB, 104 EM_NONE 105 } EMType; 106 107 typedef enum { 108 V0, 109 X0, 110 T0, 111 M0, 112 Q0, 113 PHI0, 114 POISSON, 115 VLASOV, 116 SIGMA, 117 NUM_CONSTANTS 118 } ConstantType; 119 120 typedef enum { 121 E_MONITOR_NONE, 122 E_MONITOR_FULL, 123 E_MONITOR_QUIET 124 } EMonitorType; 125 const char *const EMonitorTypes[] = {"NONE", "FULL", "QUIET", "EMonitorType", "E_MONITOR_", NULL}; 126 127 typedef struct { 128 PetscScalar v0; /* Velocity scale, often the thermal velocity */ 129 PetscScalar t0; /* Time scale */ 130 PetscScalar x0; /* Space scale */ 131 PetscScalar m0; /* Mass scale */ 132 PetscScalar q0; /* Charge scale */ 133 PetscScalar kb; 134 PetscScalar epsi0; 135 PetscScalar phi0; /* Potential scale */ 136 PetscScalar poissonNumber; /* Non-Dimensional Poisson Number */ 137 PetscScalar vlasovNumber; /* Non-Dimensional Vlasov Number */ 138 PetscReal sigma; /* Nondimensional charge per length in x */ 139 } Parameter; 140 141 typedef struct { 142 PetscBag bag; // Problem parameters 143 PetscBool error; // Flag for printing the error 144 PetscInt remapFreq; // Number of timesteps between remapping 145 EMonitorType efield_monitor; // Flag to show electric field monitor 146 PetscBool moment_monitor; // Flag to show distribution moment monitor 147 PetscBool moment_field_monitor; // Flag to show moment field monitor 148 PetscBool positions_monitor; // Flag to show particle positins at each time step 149 PetscBool poisson_monitor; // Flag to display charge, E field, and potential at each solve 150 PetscBool initial_monitor; // Flag to monitor the initial conditions 151 PetscInt velocity_monitor; // Cell to monitor the velocity distribution for 152 PetscBool perturbed_weights; // Uniformly sample x,v space with gaussian weights 153 PetscInt ostep; // Print the energy at each ostep time steps 154 PetscInt numParticles; 155 PetscReal timeScale; /* Nondimensionalizing time scale */ 156 PetscReal charges[2]; /* The charges of each species */ 157 PetscReal masses[2]; /* The masses of each species */ 158 PetscReal thermal_energy[2]; /* Thermal Energy (used to get other constants)*/ 159 PetscReal cosine_coefficients[2]; /*(alpha, k)*/ 160 PetscReal totalWeight; 161 PetscReal stepSize; 162 PetscInt steps; 163 PetscReal initVel; 164 EMType em; // Type of electrostatic model 165 SNES snes; // EM solver 166 DM dmMom; // The DM for moment fields 167 DM dmN; // The DM for number density fields 168 IS isN; // The IS mapping dmN into dmMom 169 Mat MN; // The finite element mass matrix for number density 170 DM dmP; // The DM for momentum density fields 171 IS isP; // The IS mapping dmP into dmMom 172 Mat MP; // The finite element mass matrix for momentum density 173 DM dmE; // The DM for energy density (pressure) fields 174 IS isE; // The IS mapping dmE into dmMom 175 Mat ME; // The finite element mass matrix for energy density (pressure) 176 DM dmPot; // The DM for potential 177 Mat fftPot; // Fourier Transform operator for the potential 178 Vec fftX, fftY; // FFT vectors with phases added (complex parts) 179 IS fftReal; // The indices for real parts 180 IS isPot; // The IS for potential, or NULL in primal 181 Mat M; // The finite element mass matrix for potential 182 PetscFEGeom *fegeom; // Geometric information for the DM cells 183 PetscDrawHG drawhgic_x; // Histogram of the particle weight in each X cell 184 PetscDrawHG drawhgic_v; // Histogram of the particle weight in each X cell 185 PetscDrawHG drawhgcell_v; // Histogram of the particle weight in a given cell 186 PetscBool validE; // Flag to indicate E-field in swarm is valid 187 PetscReal drawlgEmin; // The minimum lg(E) to plot 188 PetscDrawLG drawlgE; // Logarithm of maximum electric field 189 PetscDrawSP drawspE; // Electric field at particle positions 190 PetscDrawSP drawspX; // Particle positions 191 PetscViewer viewerRho; // Charge density viewer 192 PetscViewer viewerRhoHat; // Charge density Fourier Transform viewer 193 PetscViewer viewerPhi; // Potential viewer 194 PetscViewer viewerN; // Number density viewer 195 PetscViewer viewerP; // Momentum density viewer 196 PetscViewer viewerE; // Energy density (pressure) viewer 197 PetscViewer viewerNRes; // Number density residual viewer 198 PetscViewer viewerPRes; // Momentum density residual viewer 199 PetscViewer viewerERes; // Energy density (pressure) residual viewer 200 PetscDrawLG drawlgMomRes; // Residuals for the moment equations 201 DM swarm; // The particle swarm 202 PetscRandom random; // Used for particle perturbations 203 PetscBool twostream; // Flag for activating 2-stream setup 204 PetscBool checkweights; // Check weight normalization 205 PetscInt checkVRes; // Flag to check/output velocity residuals for nightly tests 206 PetscBool checkLandau; // Check the Landau damping result 207 PetscReal gamma; // The damping rate for Landau damping 208 PetscReal omega; // The perturbed oscillation frequency for Landau damping 209 210 PetscLogEvent RhsXEvent, RhsVEvent, ESolveEvent, ETabEvent; 211 } AppCtx; 212 213 static PetscErrorCode ProcessOptions(MPI_Comm comm, AppCtx *options) 214 { 215 PetscFunctionBeginUser; 216 PetscInt d = 2; 217 PetscInt maxSpecies = 2; 218 options->error = PETSC_FALSE; 219 options->remapFreq = 1; 220 options->efield_monitor = E_MONITOR_NONE; 221 options->moment_monitor = PETSC_FALSE; 222 options->moment_field_monitor = PETSC_FALSE; 223 options->initial_monitor = PETSC_FALSE; 224 options->perturbed_weights = PETSC_FALSE; 225 options->poisson_monitor = PETSC_FALSE; 226 options->positions_monitor = PETSC_FALSE; 227 options->velocity_monitor = -1; 228 options->ostep = 100; 229 options->timeScale = 2.0e-14; 230 options->charges[0] = -1.0; 231 options->charges[1] = 1.0; 232 options->masses[0] = 1.0; 233 options->masses[1] = 1000.0; 234 options->thermal_energy[0] = 1.0; 235 options->thermal_energy[1] = 1.0; 236 options->cosine_coefficients[0] = 0.01; 237 options->cosine_coefficients[1] = 0.5; 238 options->initVel = 1; 239 options->totalWeight = 1.0; 240 options->drawhgic_x = NULL; 241 options->drawhgic_v = NULL; 242 options->drawhgcell_v = NULL; 243 options->drawlgEmin = -6; 244 options->drawlgE = NULL; 245 options->drawspE = NULL; 246 options->drawspX = NULL; 247 options->viewerRho = NULL; 248 options->viewerRhoHat = NULL; 249 options->viewerPhi = NULL; 250 options->viewerN = NULL; 251 options->viewerP = NULL; 252 options->viewerE = NULL; 253 options->viewerNRes = NULL; 254 options->viewerPRes = NULL; 255 options->viewerERes = NULL; 256 options->drawlgMomRes = NULL; 257 options->em = EM_COULOMB; 258 options->snes = NULL; 259 options->dmMom = NULL; 260 options->dmN = NULL; 261 options->MN = NULL; 262 options->dmP = NULL; 263 options->MP = NULL; 264 options->dmE = NULL; 265 options->ME = NULL; 266 options->dmPot = NULL; 267 options->fftPot = NULL; 268 options->fftX = NULL; 269 options->fftY = NULL; 270 options->fftReal = NULL; 271 options->isPot = NULL; 272 options->M = NULL; 273 options->numParticles = 32768; 274 options->twostream = PETSC_FALSE; 275 options->checkweights = PETSC_FALSE; 276 options->checkVRes = 0; 277 options->checkLandau = PETSC_FALSE; 278 279 PetscOptionsBegin(comm, "", "Landau Damping and Two Stream options", "DMSWARM"); 280 PetscCall(PetscOptionsBool("-error", "Flag to print the error", __FILE__, options->error, &options->error, NULL)); 281 PetscCall(PetscOptionsInt("-remap_freq", "Number", __FILE__, options->remapFreq, &options->remapFreq, NULL)); 282 PetscCall(PetscOptionsEnum("-efield_monitor", "Flag to record and plot log(max E) over time", __FILE__, EMonitorTypes, (PetscEnum)options->efield_monitor, (PetscEnum *)&options->efield_monitor, NULL)); 283 PetscCall(PetscOptionsReal("-efield_min_monitor", "Minimum E field to plot", __FILE__, options->drawlgEmin, &options->drawlgEmin, NULL)); 284 PetscCall(PetscOptionsBool("-moments_monitor", "Flag to show moments table", __FILE__, options->moment_monitor, &options->moment_monitor, NULL)); 285 PetscCall(PetscOptionsBool("-moment_field_monitor", "Flag to show moment fields", __FILE__, options->moment_field_monitor, &options->moment_field_monitor, NULL)); 286 PetscCall(PetscOptionsBool("-ics_monitor", "Flag to show initial condition histograms", __FILE__, options->initial_monitor, &options->initial_monitor, NULL)); 287 PetscCall(PetscOptionsBool("-positions_monitor", "The flag to show particle positions", __FILE__, options->positions_monitor, &options->positions_monitor, NULL)); 288 PetscCall(PetscOptionsBool("-poisson_monitor", "The flag to show charges, Efield and potential solve", __FILE__, options->poisson_monitor, &options->poisson_monitor, NULL)); 289 PetscCall(PetscOptionsInt("-velocity_monitor", "Cell to show velocity histograms", __FILE__, options->velocity_monitor, &options->velocity_monitor, NULL)); 290 PetscCall(PetscOptionsBool("-twostream", "Run two stream instability", __FILE__, options->twostream, &options->twostream, NULL)); 291 PetscCall(PetscOptionsBool("-perturbed_weights", "Flag to run uniform sampling with perturbed weights", __FILE__, options->perturbed_weights, &options->perturbed_weights, NULL)); 292 PetscCall(PetscOptionsBool("-check_weights", "Ensure all particle weights are positive", __FILE__, options->checkweights, &options->checkweights, NULL)); 293 PetscCall(PetscOptionsBool("-check_landau", "Check the decay from Landau damping", __FILE__, options->checkLandau, &options->checkLandau, NULL)); 294 PetscCall(PetscOptionsInt("-output_step", "Number of time steps between output", __FILE__, options->ostep, &options->ostep, NULL)); 295 PetscCall(PetscOptionsReal("-timeScale", "Nondimensionalizing time scale", __FILE__, options->timeScale, &options->timeScale, NULL)); 296 PetscCall(PetscOptionsInt("-check_vel_res", "Check particle velocity residuals for nightly tests", __FILE__, options->checkVRes, &options->checkVRes, NULL)); 297 PetscCall(PetscOptionsReal("-initial_velocity", "Initial velocity of perturbed particle", __FILE__, options->initVel, &options->initVel, NULL)); 298 PetscCall(PetscOptionsReal("-total_weight", "Total weight of all particles", __FILE__, options->totalWeight, &options->totalWeight, NULL)); 299 PetscCall(PetscOptionsRealArray("-cosine_coefficients", "Amplitude and frequency of cosine equation used in initialization", __FILE__, options->cosine_coefficients, &d, NULL)); 300 PetscCall(PetscOptionsRealArray("-charges", "Species charges", __FILE__, options->charges, &maxSpecies, NULL)); 301 PetscCall(PetscOptionsEnum("-em_type", "Type of electrostatic solver", __FILE__, EMTypes, (PetscEnum)options->em, (PetscEnum *)&options->em, NULL)); 302 PetscOptionsEnd(); 303 304 PetscCall(PetscLogEventRegister("RhsX", TS_CLASSID, &options->RhsXEvent)); 305 PetscCall(PetscLogEventRegister("RhsV", TS_CLASSID, &options->RhsVEvent)); 306 PetscCall(PetscLogEventRegister("ESolve", TS_CLASSID, &options->ESolveEvent)); 307 PetscCall(PetscLogEventRegister("ETab", TS_CLASSID, &options->ETabEvent)); 308 PetscFunctionReturn(PETSC_SUCCESS); 309 } 310 311 static PetscErrorCode SetupContext(DM dm, DM sw, AppCtx *user) 312 { 313 MPI_Comm comm; 314 315 PetscFunctionBeginUser; 316 PetscCall(PetscObjectGetComm((PetscObject)dm, &comm)); 317 if (user->efield_monitor) { 318 PetscDraw draw; 319 PetscDrawAxis axis; 320 321 if (user->efield_monitor == E_MONITOR_FULL) { 322 PetscCall(PetscDrawCreate(comm, NULL, "Max Electric Field", 0, 0, 400, 300, &draw)); 323 PetscCall(PetscDrawSetSave(draw, "ex2_Efield")); 324 PetscCall(PetscDrawSetFromOptions(draw)); 325 } else { 326 PetscCall(PetscDrawOpenNull(comm, &draw)); 327 } 328 PetscCall(PetscDrawLGCreate(draw, 1, &user->drawlgE)); 329 PetscCall(PetscDrawDestroy(&draw)); 330 PetscCall(PetscDrawLGGetAxis(user->drawlgE, &axis)); 331 PetscCall(PetscDrawAxisSetLabels(axis, "Max Electric Field", "time", "E_max")); 332 PetscCall(PetscDrawLGSetLimits(user->drawlgE, 0., user->steps * user->stepSize, user->drawlgEmin, 0.)); 333 } 334 335 if (user->initial_monitor) { 336 PetscDraw drawic_x, drawic_v; 337 PetscDrawAxis axis1, axis2; 338 PetscReal dmboxlower[2], dmboxupper[2]; 339 PetscInt dim, cStart, cEnd; 340 341 PetscCall(DMGetDimension(sw, &dim)); 342 PetscCall(DMGetBoundingBox(dm, dmboxlower, dmboxupper)); 343 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd)); 344 345 PetscCall(PetscDrawCreate(comm, NULL, "monitor_initial_conditions_x", 0, 300, 400, 300, &drawic_x)); 346 PetscCall(PetscDrawSetSave(drawic_x, "ex2_ic_x")); 347 PetscCall(PetscDrawSetFromOptions(drawic_x)); 348 PetscCall(PetscDrawHGCreate(drawic_x, (int)dim, &user->drawhgic_x)); 349 PetscCall(PetscDrawHGCalcStats(user->drawhgic_x, PETSC_TRUE)); 350 PetscCall(PetscDrawHGGetAxis(user->drawhgic_x, &axis1)); 351 PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_x, (int)(cEnd - cStart))); 352 PetscCall(PetscDrawAxisSetLabels(axis1, "Initial X Distribution", "X", "weight")); 353 PetscCall(PetscDrawAxisSetLimits(axis1, dmboxlower[0], dmboxupper[0], 0, 0)); 354 PetscCall(PetscDrawDestroy(&drawic_x)); 355 356 PetscCall(PetscDrawCreate(comm, NULL, "monitor_initial_conditions_v", 400, 300, 400, 300, &drawic_v)); 357 PetscCall(PetscDrawSetSave(drawic_v, "ex9_ic_v")); 358 PetscCall(PetscDrawSetFromOptions(drawic_v)); 359 PetscCall(PetscDrawHGCreate(drawic_v, (int)dim, &user->drawhgic_v)); 360 PetscCall(PetscDrawHGCalcStats(user->drawhgic_v, PETSC_TRUE)); 361 PetscCall(PetscDrawHGGetAxis(user->drawhgic_v, &axis2)); 362 PetscCall(PetscDrawHGSetNumberBins(user->drawhgic_v, 21)); 363 PetscCall(PetscDrawAxisSetLabels(axis2, "Initial V_x Distribution", "V", "weight")); 364 PetscCall(PetscDrawAxisSetLimits(axis2, -6, 6, 0, 0)); 365 PetscCall(PetscDrawDestroy(&drawic_v)); 366 } 367 368 if (user->velocity_monitor >= 0) { 369 DM vdm; 370 DMSwarmCellDM celldm; 371 PetscDraw drawcell_v; 372 PetscDrawAxis axis; 373 PetscReal dmboxlower[2], dmboxupper[2]; 374 PetscInt dim; 375 char title[PETSC_MAX_PATH_LEN]; 376 377 PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm)); 378 PetscCall(DMSwarmCellDMGetDM(celldm, &vdm)); 379 PetscCall(DMGetDimension(vdm, &dim)); 380 PetscCall(DMGetBoundingBox(vdm, dmboxlower, dmboxupper)); 381 382 PetscCall(PetscSNPrintf(title, PETSC_MAX_PATH_LEN, "Cell %" PetscInt_FMT ": Velocity Distribution", user->velocity_monitor)); 383 PetscCall(PetscDrawCreate(comm, NULL, title, 400, 300, 400, 300, &drawcell_v)); 384 PetscCall(PetscDrawSetSave(drawcell_v, "ex2_cell_v")); 385 PetscCall(PetscDrawSetFromOptions(drawcell_v)); 386 PetscCall(PetscDrawHGCreate(drawcell_v, (int)dim, &user->drawhgcell_v)); 387 PetscCall(PetscDrawHGCalcStats(user->drawhgcell_v, PETSC_TRUE)); 388 PetscCall(PetscDrawHGGetAxis(user->drawhgcell_v, &axis)); 389 PetscCall(PetscDrawHGSetNumberBins(user->drawhgcell_v, 21)); 390 PetscCall(PetscDrawAxisSetLabels(axis, "V_x Distribution", "V", "weight")); 391 PetscCall(PetscDrawAxisSetLimits(axis, dmboxlower[0], dmboxupper[0], 0, 0)); 392 PetscCall(PetscDrawDestroy(&drawcell_v)); 393 } 394 395 if (user->positions_monitor) { 396 PetscDraw draw; 397 PetscDrawAxis axis; 398 399 PetscCall(PetscDrawCreate(comm, NULL, "Particle Position", 0, 0, 400, 300, &draw)); 400 PetscCall(PetscDrawSetSave(draw, "ex9_pos")); 401 PetscCall(PetscDrawSetFromOptions(draw)); 402 PetscCall(PetscDrawSPCreate(draw, 10, &user->drawspX)); 403 PetscCall(PetscDrawDestroy(&draw)); 404 PetscCall(PetscDrawSPSetDimension(user->drawspX, 1)); 405 PetscCall(PetscDrawSPGetAxis(user->drawspX, &axis)); 406 PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "v")); 407 PetscCall(PetscDrawSPReset(user->drawspX)); 408 } 409 if (user->poisson_monitor) { 410 Vec rho, rhohat, phi; 411 PetscDraw draw; 412 PetscDrawAxis axis; 413 414 PetscCall(PetscDrawCreate(comm, NULL, "Electric_Field", 0, 0, 400, 300, &draw)); 415 PetscCall(PetscDrawSetFromOptions(draw)); 416 PetscCall(PetscDrawSetSave(draw, "ex9_E_spatial")); 417 PetscCall(PetscDrawSPCreate(draw, 10, &user->drawspE)); 418 PetscCall(PetscDrawDestroy(&draw)); 419 PetscCall(PetscDrawSPSetDimension(user->drawspE, 1)); 420 PetscCall(PetscDrawSPGetAxis(user->drawspE, &axis)); 421 PetscCall(PetscDrawAxisSetLabels(axis, "Particles", "x", "E")); 422 PetscCall(PetscDrawSPReset(user->drawspE)); 423 424 PetscCall(PetscViewerDrawOpen(comm, NULL, "Charge Density", 0, 0, 400, 300, &user->viewerRho)); 425 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerRho, "rho_")); 426 PetscCall(PetscViewerDrawGetDraw(user->viewerRho, 0, &draw)); 427 PetscCall(PetscDrawSetSave(draw, "ex9_rho_spatial")); 428 PetscCall(PetscViewerSetFromOptions(user->viewerRho)); 429 PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho)); 430 PetscCall(PetscObjectSetName((PetscObject)rho, "charge_density")); 431 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho)); 432 433 PetscInt dim, N; 434 435 PetscCall(DMGetDimension(user->dmPot, &dim)); 436 if (dim == 1) { 437 PetscCall(DMGetNamedGlobalVector(user->dmPot, "rhohat", &rhohat)); 438 PetscCall(VecGetSize(rhohat, &N)); 439 PetscCall(MatCreateFFT(comm, dim, &N, MATFFTW, &user->fftPot)); 440 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rhohat", &rhohat)); 441 PetscCall(MatCreateVecs(user->fftPot, &user->fftX, &user->fftY)); 442 PetscCall(ISCreateStride(PETSC_COMM_SELF, N, 0, 1, &user->fftReal)); 443 } 444 445 PetscCall(PetscViewerDrawOpen(comm, NULL, "rhohat: Charge Density FT", 0, 0, 400, 300, &user->viewerRhoHat)); 446 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerRhoHat, "rhohat_")); 447 PetscCall(PetscViewerDrawGetDraw(user->viewerRhoHat, 0, &draw)); 448 PetscCall(PetscDrawSetSave(draw, "ex9_rho_ft")); 449 PetscCall(PetscViewerSetFromOptions(user->viewerRhoHat)); 450 PetscCall(DMGetNamedGlobalVector(user->dmPot, "rhohat", &rhohat)); 451 PetscCall(PetscObjectSetName((PetscObject)rhohat, "charge_density_ft")); 452 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rhohat", &rhohat)); 453 454 PetscCall(PetscViewerDrawOpen(comm, NULL, "Potential", 400, 0, 400, 300, &user->viewerPhi)); 455 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerPhi, "phi_")); 456 PetscCall(PetscViewerDrawGetDraw(user->viewerPhi, 0, &draw)); 457 PetscCall(PetscDrawSetSave(draw, "ex9_phi_spatial")); 458 PetscCall(PetscViewerSetFromOptions(user->viewerPhi)); 459 PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi)); 460 PetscCall(PetscObjectSetName((PetscObject)phi, "potential")); 461 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi)); 462 } 463 if (user->moment_field_monitor) { 464 Vec n, p, e; 465 Vec nres, pres, eres; 466 PetscDraw draw; 467 468 PetscCall(PetscViewerDrawOpen(comm, NULL, "Number Density", 400, 0, 400, 300, &user->viewerN)); 469 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerN, "n_")); 470 PetscCall(PetscViewerDrawGetDraw(user->viewerN, 0, &draw)); 471 PetscCall(PetscDrawSetSave(draw, "ex4_n_spatial")); 472 PetscCall(PetscViewerSetFromOptions(user->viewerN)); 473 PetscCall(DMGetNamedGlobalVector(user->dmN, "n", &n)); 474 PetscCall(PetscObjectSetName((PetscObject)n, "Number Density")); 475 PetscCall(DMRestoreNamedGlobalVector(user->dmN, "n", &n)); 476 477 PetscCall(PetscViewerDrawOpen(comm, NULL, "Momentum Density", 800, 0, 400, 300, &user->viewerP)); 478 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerP, "p_")); 479 PetscCall(PetscViewerDrawGetDraw(user->viewerP, 0, &draw)); 480 PetscCall(PetscDrawSetSave(draw, "ex4_p_spatial")); 481 PetscCall(PetscViewerSetFromOptions(user->viewerP)); 482 PetscCall(DMGetNamedGlobalVector(user->dmP, "p", &p)); 483 PetscCall(PetscObjectSetName((PetscObject)p, "Momentum Density")); 484 PetscCall(DMRestoreNamedGlobalVector(user->dmP, "p", &p)); 485 486 PetscCall(PetscViewerDrawOpen(comm, NULL, "Emergy Density (Pressure)", 1200, 0, 400, 300, &user->viewerE)); 487 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerE, "e_")); 488 PetscCall(PetscViewerDrawGetDraw(user->viewerE, 0, &draw)); 489 PetscCall(PetscDrawSetSave(draw, "ex4_e_spatial")); 490 PetscCall(PetscViewerSetFromOptions(user->viewerE)); 491 PetscCall(DMGetNamedGlobalVector(user->dmE, "e", &e)); 492 PetscCall(PetscObjectSetName((PetscObject)e, "Energy Density (Pressure)")); 493 PetscCall(DMRestoreNamedGlobalVector(user->dmE, "e", &e)); 494 495 PetscDrawAxis axis; 496 497 PetscCall(PetscDrawCreate(comm, NULL, "Moment Residual", 0, 320, 400, 300, &draw)); 498 PetscCall(PetscDrawSetSave(draw, "ex4_moment_res")); 499 PetscCall(PetscDrawSetFromOptions(draw)); 500 PetscCall(PetscDrawLGCreate(draw, 3, &user->drawlgMomRes)); 501 PetscCall(PetscDrawDestroy(&draw)); 502 PetscCall(PetscDrawLGGetAxis(user->drawlgMomRes, &axis)); 503 PetscCall(PetscDrawAxisSetLabels(axis, "Moment Residial", "time", "Residual Norm")); 504 PetscCall(PetscDrawLGSetLimits(user->drawlgMomRes, 0., user->steps * user->stepSize, -8, 0)); 505 506 PetscCall(PetscViewerDrawOpen(comm, NULL, "Number Density Residual", 400, 300, 400, 300, &user->viewerNRes)); 507 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerNRes, "nres_")); 508 PetscCall(PetscViewerDrawGetDraw(user->viewerNRes, 0, &draw)); 509 PetscCall(PetscDrawSetSave(draw, "ex4_nres_spatial")); 510 PetscCall(PetscViewerSetFromOptions(user->viewerNRes)); 511 PetscCall(DMGetNamedGlobalVector(user->dmN, "nres", &nres)); 512 PetscCall(PetscObjectSetName((PetscObject)nres, "Number Density Residual")); 513 PetscCall(DMRestoreNamedGlobalVector(user->dmN, "nres", &nres)); 514 515 PetscCall(PetscViewerDrawOpen(comm, NULL, "Momentum Density Residual", 800, 300, 400, 300, &user->viewerPRes)); 516 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerPRes, "pres_")); 517 PetscCall(PetscViewerDrawGetDraw(user->viewerPRes, 0, &draw)); 518 PetscCall(PetscDrawSetSave(draw, "ex4_pres_spatial")); 519 PetscCall(PetscViewerSetFromOptions(user->viewerPRes)); 520 PetscCall(DMGetNamedGlobalVector(user->dmP, "pres", &pres)); 521 PetscCall(PetscObjectSetName((PetscObject)pres, "Momentum Density Residual")); 522 PetscCall(DMRestoreNamedGlobalVector(user->dmP, "pres", &pres)); 523 524 PetscCall(PetscViewerDrawOpen(comm, NULL, "Energy Density Residual", 1200, 300, 400, 300, &user->viewerERes)); 525 PetscCall(PetscObjectSetOptionsPrefix((PetscObject)user->viewerERes, "eres_")); 526 PetscCall(PetscViewerDrawGetDraw(user->viewerERes, 0, &draw)); 527 PetscCall(PetscDrawSetSave(draw, "ex4_eres_spatial")); 528 PetscCall(PetscViewerSetFromOptions(user->viewerERes)); 529 PetscCall(DMGetNamedGlobalVector(user->dmE, "eres", &eres)); 530 PetscCall(PetscObjectSetName((PetscObject)eres, "Energy Density Residual")); 531 PetscCall(DMRestoreNamedGlobalVector(user->dmE, "eres", &eres)); 532 } 533 PetscFunctionReturn(PETSC_SUCCESS); 534 } 535 536 static PetscErrorCode DestroyContext(AppCtx *user) 537 { 538 PetscFunctionBeginUser; 539 PetscCall(PetscDrawHGDestroy(&user->drawhgic_x)); 540 PetscCall(PetscDrawHGDestroy(&user->drawhgic_v)); 541 PetscCall(PetscDrawHGDestroy(&user->drawhgcell_v)); 542 543 PetscCall(PetscDrawLGDestroy(&user->drawlgE)); 544 PetscCall(PetscDrawSPDestroy(&user->drawspE)); 545 PetscCall(PetscDrawSPDestroy(&user->drawspX)); 546 PetscCall(PetscViewerDestroy(&user->viewerRho)); 547 PetscCall(PetscViewerDestroy(&user->viewerRhoHat)); 548 PetscCall(MatDestroy(&user->fftPot)); 549 PetscCall(VecDestroy(&user->fftX)); 550 PetscCall(VecDestroy(&user->fftY)); 551 PetscCall(ISDestroy(&user->fftReal)); 552 PetscCall(PetscViewerDestroy(&user->viewerPhi)); 553 PetscCall(PetscViewerDestroy(&user->viewerN)); 554 PetscCall(PetscViewerDestroy(&user->viewerP)); 555 PetscCall(PetscViewerDestroy(&user->viewerE)); 556 PetscCall(PetscViewerDestroy(&user->viewerNRes)); 557 PetscCall(PetscViewerDestroy(&user->viewerPRes)); 558 PetscCall(PetscViewerDestroy(&user->viewerERes)); 559 PetscCall(PetscDrawLGDestroy(&user->drawlgMomRes)); 560 561 PetscCall(PetscBagDestroy(&user->bag)); 562 PetscFunctionReturn(PETSC_SUCCESS); 563 } 564 565 static PetscErrorCode CheckNonNegativeWeights(DM sw, AppCtx *user) 566 { 567 const PetscScalar *w; 568 PetscInt Np; 569 570 PetscFunctionBeginUser; 571 if (!user->checkweights) PetscFunctionReturn(PETSC_SUCCESS); 572 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&w)); 573 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 574 for (PetscInt p = 0; p < Np; ++p) PetscCheck(w[p] >= 0.0, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " has negative weight %g", p, w[p]); 575 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&w)); 576 PetscFunctionReturn(PETSC_SUCCESS); 577 } 578 579 static PetscErrorCode computeVelocityFEMMoments(DM sw, PetscReal moments[], AppCtx *user) 580 { 581 DMSwarmCellDM celldm; 582 DM vdm; 583 Vec u[1]; 584 const char *fields[1] = {"w_q"}; 585 586 PetscFunctionBegin; 587 PetscCall(DMSwarmSetCellDMActive(sw, "velocity")); 588 PetscCall(DMSwarmGetCellDMActive(sw, &celldm)); 589 PetscCall(DMSwarmCellDMGetDM(celldm, &vdm)); 590 PetscCall(DMGetGlobalVector(vdm, &u[0])); 591 PetscCall(DMSwarmProjectFields(sw, vdm, 1, fields, u, SCATTER_FORWARD)); 592 PetscCall(DMPlexComputeMoments(vdm, u[0], moments)); 593 PetscCall(DMRestoreGlobalVector(vdm, &u[0])); 594 PetscCall(DMSwarmSetCellDMActive(sw, "space")); 595 PetscFunctionReturn(PETSC_SUCCESS); 596 } 597 598 static void f0_grad_phi2(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[]) 599 { 600 f0[0] = 0.; 601 for (PetscInt d = 0; d < dim; ++d) f0[0] += PetscSqr(u_x[uOff_x[0] + d * dim + d]); 602 } 603 604 typedef struct { 605 PetscInt s; // Starting sample (we ignore some in the beginning) 606 PetscInt e; // Ending sample 607 const PetscReal *t; // Time for each sample 608 const PetscReal *Emax; // Emax for each sample 609 } EmaxCtx; 610 611 // Our model is E_max(t) = C e^{-gamma t} |cos(omega t - phi)| 612 static PetscErrorCode ComputeEmaxResidual(Tao tao, Vec x, Vec res, void *user) 613 { 614 EmaxCtx *ctx = (EmaxCtx *)user; 615 const PetscScalar *a; 616 PetscScalar *F; 617 PetscReal C, gamma, omega, phi; 618 619 PetscFunctionBegin; 620 PetscCall(VecGetArrayRead(x, &a)); 621 PetscCall(VecGetArray(res, &F)); 622 C = PetscRealPart(a[0]); 623 gamma = PetscRealPart(a[1]); 624 omega = PetscRealPart(a[2]); 625 phi = PetscRealPart(a[3]); 626 PetscCall(VecRestoreArrayRead(x, &a)); 627 for (PetscInt i = ctx->s; i < ctx->e; ++i) F[i - ctx->s] = PetscPowReal(10., ctx->Emax[i]) - C * PetscExpReal(-gamma * ctx->t[i]) * PetscAbsReal(PetscCosReal(omega * ctx->t[i] - phi)); 628 PetscCall(VecRestoreArray(res, &F)); 629 PetscFunctionReturn(PETSC_SUCCESS); 630 } 631 632 // The Jacobian of the residual J = dr(x)/dx 633 static PetscErrorCode ComputeEmaxJacobian(Tao tao, Vec x, Mat J, Mat Jpre, void *user) 634 { 635 EmaxCtx *ctx = (EmaxCtx *)user; 636 const PetscScalar *a; 637 PetscScalar *jac; 638 PetscReal C, gamma, omega, phi; 639 const PetscInt n = ctx->e - ctx->s; 640 641 PetscFunctionBegin; 642 PetscCall(VecGetArrayRead(x, &a)); 643 C = PetscRealPart(a[0]); 644 gamma = PetscRealPart(a[1]); 645 omega = PetscRealPart(a[2]); 646 phi = PetscRealPart(a[3]); 647 PetscCall(VecRestoreArrayRead(x, &a)); 648 PetscCall(MatDenseGetArray(J, &jac)); 649 for (PetscInt i = 0; i < n; ++i) { 650 const PetscInt k = i + ctx->s; 651 652 jac[i * 4 + 0] = -PetscExpReal(-gamma * ctx->t[k]) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi)); 653 jac[i * 4 + 1] = C * ctx->t[k] * PetscExpReal(-gamma * ctx->t[k]) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi)); 654 jac[i * 4 + 2] = C * ctx->t[k] * PetscExpReal(-gamma * ctx->t[k]) * (PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * PetscSinReal(omega * ctx->t[k] - phi); 655 jac[i * 4 + 3] = -C * PetscExpReal(-gamma * ctx->t[k]) * (PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * PetscSinReal(omega * ctx->t[k] - phi); 656 } 657 PetscCall(MatDenseRestoreArray(J, &jac)); 658 PetscFunctionReturn(PETSC_SUCCESS); 659 } 660 661 // Our model is log_10 E_max(t) = log_10 C - gamma t log_10 e + log_10 |cos(omega t - phi)| 662 static PetscErrorCode ComputeLogEmaxResidual(Tao tao, Vec x, Vec res, void *user) 663 { 664 EmaxCtx *ctx = (EmaxCtx *)user; 665 const PetscScalar *a; 666 PetscScalar *F; 667 PetscReal C, gamma, omega, phi; 668 669 PetscFunctionBegin; 670 PetscCall(VecGetArrayRead(x, &a)); 671 PetscCall(VecGetArray(res, &F)); 672 C = PetscRealPart(a[0]); 673 gamma = PetscRealPart(a[1]); 674 omega = PetscRealPart(a[2]); 675 phi = PetscRealPart(a[3]); 676 PetscCall(VecRestoreArrayRead(x, &a)); 677 for (PetscInt i = ctx->s; i < ctx->e; ++i) { 678 if (C < 0) { 679 F[i - ctx->s] = 1e10; 680 continue; 681 } 682 F[i - ctx->s] = ctx->Emax[i] - (PetscLog10Real(C) - gamma * ctx->t[i] * PetscLog10Real(PETSC_E) + PetscLog10Real(PetscAbsReal(PetscCosReal(omega * ctx->t[i] - phi)))); 683 } 684 PetscCall(VecRestoreArray(res, &F)); 685 PetscFunctionReturn(PETSC_SUCCESS); 686 } 687 688 // The Jacobian of the residual J = dr(x)/dx 689 static PetscErrorCode ComputeLogEmaxJacobian(Tao tao, Vec x, Mat J, Mat Jpre, void *user) 690 { 691 EmaxCtx *ctx = (EmaxCtx *)user; 692 const PetscScalar *a; 693 PetscScalar *jac; 694 PetscReal C, omega, phi; 695 const PetscInt n = ctx->e - ctx->s; 696 697 PetscFunctionBegin; 698 PetscCall(VecGetArrayRead(x, &a)); 699 C = PetscRealPart(a[0]); 700 omega = PetscRealPart(a[2]); 701 phi = PetscRealPart(a[3]); 702 PetscCall(VecRestoreArrayRead(x, &a)); 703 PetscCall(MatDenseGetArray(J, &jac)); 704 for (PetscInt i = 0; i < n; ++i) { 705 const PetscInt k = i + ctx->s; 706 707 jac[0 * n + i] = -1. / (PetscLog10Real(PETSC_E) * C); 708 jac[1 * n + i] = ctx->t[k] * PetscLog10Real(PETSC_E); 709 jac[2 * n + i] = (PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * ctx->t[k] * PetscSinReal(omega * ctx->t[k] - phi) / (PetscLog10Real(PETSC_E) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi))); 710 jac[3 * n + i] = -(PetscCosReal(omega * ctx->t[k] - phi) < 0. ? -1. : 1.) * PetscSinReal(omega * ctx->t[k] - phi) / (PetscLog10Real(PETSC_E) * PetscAbsReal(PetscCosReal(omega * ctx->t[k] - phi))); 711 } 712 PetscCall(MatDenseRestoreArray(J, &jac)); 713 PetscCall(MatViewFromOptions(J, NULL, "-emax_jac_view")); 714 PetscFunctionReturn(PETSC_SUCCESS); 715 } 716 717 static PetscErrorCode MonitorEField(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx) 718 { 719 AppCtx *user = (AppCtx *)ctx; 720 DM sw; 721 PetscScalar intESq; 722 PetscReal *E, *x, *weight; 723 PetscReal Enorm = 0., lgEnorm, lgEmax, sum = 0., Emax = 0., chargesum = 0.; 724 PetscReal pmoments[4]; /* \int f, \int v f, \int v^2 f */ 725 PetscInt *species, dim, Np, gNp; 726 MPI_Comm comm; 727 PetscMPIInt rank; 728 729 PetscFunctionBeginUser; 730 if (step < 0 || !user->validE) PetscFunctionReturn(PETSC_SUCCESS); 731 PetscCall(PetscObjectGetComm((PetscObject)ts, &comm)); 732 PetscCallMPI(MPI_Comm_rank(comm, &rank)); 733 PetscCall(TSGetDM(ts, &sw)); 734 PetscCall(DMGetDimension(sw, &dim)); 735 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 736 PetscCall(DMSwarmGetSize(sw, &gNp)); 737 PetscCall(DMSwarmSortGetAccess(sw)); 738 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x)); 739 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E)); 740 PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species)); 741 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight)); 742 743 for (PetscInt p = 0; p < Np; ++p) { 744 for (PetscInt d = 0; d < 1; ++d) { 745 PetscReal temp = PetscAbsReal(E[p * dim + d]); 746 if (temp > Emax) Emax = temp; 747 } 748 Enorm += PetscSqrtReal(E[p * dim] * E[p * dim]); 749 sum += E[p * dim]; 750 chargesum += user->charges[0] * weight[p]; 751 } 752 PetscCallMPI(MPIU_Allreduce(MPI_IN_PLACE, &Emax, 1, MPIU_REAL, MPIU_MAX, comm)); 753 lgEnorm = Enorm != 0 ? PetscLog10Real(Enorm) : -16.; 754 lgEmax = Emax != 0 ? PetscLog10Real(Emax) : user->drawlgEmin; 755 756 PetscDS ds; 757 Vec phi; 758 759 PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi)); 760 PetscCall(DMGetDS(user->dmPot, &ds)); 761 PetscCall(PetscDSSetObjective(ds, 0, &f0_grad_phi2)); 762 PetscCall(DMPlexComputeIntegralFEM(user->dmPot, phi, &intESq, user)); 763 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi)); 764 765 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x)); 766 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight)); 767 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E)); 768 PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species)); 769 PetscCall(PetscDrawLGAddPoint(user->drawlgE, &t, &lgEmax)); 770 if (user->efield_monitor == E_MONITOR_FULL) { 771 PetscDraw draw; 772 773 PetscCall(PetscDrawLGDraw(user->drawlgE)); 774 PetscCall(PetscDrawLGGetDraw(user->drawlgE, &draw)); 775 PetscCall(PetscDrawSave(draw)); 776 777 PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments)); 778 PetscCall(PetscPrintf(comm, "E: %f\t%+e\t%e\t%f\t%20.15e\t%f\t%f\t%f\t%20.15e\t%20.15e\t%20.15e\t%" PetscInt_FMT "\t(%" PetscInt_FMT ")\n", (double)t, (double)sum, (double)Enorm, (double)lgEnorm, (double)Emax, (double)lgEmax, (double)chargesum, (double)pmoments[0], (double)pmoments[1], (double)pmoments[1 + dim], (double)PetscSqrtReal(intESq), gNp, step)); 779 PetscCall(DMViewFromOptions(sw, NULL, "-sw_efield_view")); 780 } 781 782 // Compute decay rate and frequency 783 EmaxCtx ectx; 784 785 ectx.s = 50; 786 PetscCall(PetscDrawLGGetData(user->drawlgE, NULL, &ectx.e, &ectx.t, &ectx.Emax)); 787 if (!rank && !(ectx.e % 100)) { 788 Tao tao; 789 Mat J; 790 Vec x, r; 791 PetscScalar *a; 792 PetscBool fitLog = PETSC_TRUE, debug = PETSC_FALSE; 793 794 PetscCall(TaoCreate(PETSC_COMM_SELF, &tao)); 795 PetscCall(TaoSetOptionsPrefix(tao, "emax_")); 796 PetscCall(VecCreateSeq(PETSC_COMM_SELF, 4, &x)); 797 PetscCall(TaoSetSolution(tao, x)); 798 PetscCall(VecCreateSeq(PETSC_COMM_SELF, ectx.e - ectx.s, &r)); 799 if (fitLog) PetscCall(TaoSetResidualRoutine(tao, r, ComputeLogEmaxResidual, &ectx)); 800 else PetscCall(TaoSetResidualRoutine(tao, r, ComputeEmaxResidual, &ectx)); 801 PetscCall(VecDestroy(&r)); 802 PetscCall(MatCreateSeqDense(PETSC_COMM_SELF, ectx.e - ectx.s, 4, NULL, &J)); 803 if (fitLog) PetscCall(TaoSetJacobianResidualRoutine(tao, J, J, ComputeLogEmaxJacobian, &ectx)); 804 else PetscCall(TaoSetJacobianResidualRoutine(tao, J, J, ComputeEmaxJacobian, &ectx)); 805 PetscCall(MatDestroy(&J)); 806 PetscCall(TaoSetFromOptions(tao)); 807 PetscCall(VecGetArray(x, &a)); 808 a[0] = 0.02; 809 a[1] = 0.15; 810 a[2] = 1.4; 811 a[3] = 0.45; 812 PetscCall(VecRestoreArray(x, &a)); 813 PetscCall(TaoSolve(tao)); 814 if (debug) { 815 PetscCall(PetscPrintf(PETSC_COMM_SELF, "t = [")); 816 for (PetscInt i = 0; i < ectx.e; ++i) { 817 if (i > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", ")); 818 PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", ectx.t[i])); 819 } 820 PetscCall(PetscPrintf(PETSC_COMM_SELF, "]\n")); 821 PetscCall(PetscPrintf(PETSC_COMM_SELF, "Emax = [")); 822 for (PetscInt i = 0; i < ectx.e; ++i) { 823 if (i > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", ")); 824 PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", ectx.Emax[i])); 825 } 826 PetscCall(PetscPrintf(PETSC_COMM_SELF, "]\n")); 827 } 828 PetscDraw draw; 829 PetscDrawAxis axis; 830 char title[PETSC_MAX_PATH_LEN]; 831 832 PetscCall(VecGetArray(x, &a)); 833 user->gamma = a[1]; 834 user->omega = a[2]; 835 if (user->efield_monitor == E_MONITOR_FULL) { 836 PetscCall(PetscPrintf(PETSC_COMM_SELF, "Emax Fit: gamma %g omega %g C %g phi %g\n", a[1], a[2], a[0], a[3])); 837 PetscCall(PetscDrawLGGetDraw(user->drawlgE, &draw)); 838 PetscCall(PetscSNPrintf(title, PETSC_MAX_PATH_LEN, "Max Electric Field gamma %.4g omega %.4g", a[1], a[2])); 839 PetscCall(PetscDrawSetTitle(draw, title)); 840 PetscCall(PetscDrawLGGetAxis(user->drawlgE, &axis)); 841 PetscCall(PetscDrawAxisSetLabels(axis, title, "time", "E_max")); 842 } 843 PetscCall(VecRestoreArray(x, &a)); 844 PetscCall(VecDestroy(&x)); 845 PetscCall(TaoDestroy(&tao)); 846 } 847 PetscFunctionReturn(PETSC_SUCCESS); 848 } 849 850 static PetscErrorCode MonitorMoments(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx) 851 { 852 AppCtx *user = (AppCtx *)ctx; 853 DM sw; 854 PetscReal pmoments[4], fmoments[4]; /* \int f, \int v f, \int v^2 f */ 855 856 PetscFunctionBeginUser; 857 if (step < 0) PetscFunctionReturn(PETSC_SUCCESS); 858 PetscCall(TSGetDM(ts, &sw)); 859 860 PetscCall(DMSwarmComputeMoments(sw, "velocity", "w_q", pmoments)); 861 PetscCall(computeVelocityFEMMoments(sw, fmoments, user)); 862 863 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "%f\t%f\t%f\t%f\t%f\t%f\t%f\n", (double)t, (double)pmoments[0], (double)pmoments[1], (double)pmoments[3], (double)fmoments[0], (double)fmoments[1], (double)fmoments[2])); 864 PetscFunctionReturn(PETSC_SUCCESS); 865 } 866 867 static PetscErrorCode zero(PetscInt dim, PetscReal time, const PetscReal x[], PetscInt Nc, PetscScalar *u, void *ctx) 868 { 869 u[0] = 0.0; 870 return PETSC_SUCCESS; 871 } 872 873 /* 874 M_p w_p 875 - Make M_p with "moments" 876 - Get w_p from Swarm 877 M_p v_p w_p 878 - Get v_p from Swarm 879 - pointwise multiply v_p and w_p 880 M_p (v_p - (\sum_j p_F \phi_j(x_p)) / m (\sum_k n_F \phi_k(x_p)))^2 w_p 881 - ProjectField(sw, {n, p} U, {v_p} A, tmp_p) 882 - pointwise multiply tmp_p and w_p 883 884 Projection works fpr swarms 885 Fields are FE from the CellDM, and aux fields are the swarm fields 886 */ 887 static PetscErrorCode ComputeMomentFields(TS ts) 888 { 889 AppCtx *user; 890 DM sw; 891 KSP ksp; 892 Mat M_p, D_p; 893 Vec f, v, E, tmpMom; 894 Vec m, mold, mfluxold, mres, n, nrhs, nflux, nres, p, prhs, pflux, pres, e, erhs, eflux, eres; 895 PetscReal dt, t; 896 PetscInt Nts; 897 898 PetscFunctionBegin; 899 PetscCall(TSGetStepNumber(ts, &Nts)); 900 PetscCall(TSGetTimeStep(ts, &dt)); 901 PetscCall(TSGetTime(ts, &t)); 902 PetscCall(TSGetDM(ts, &sw)); 903 PetscCall(DMGetApplicationContext(sw, (void *)&user)); 904 PetscCall(DMSwarmSetCellDMActive(sw, "moment fields")); 905 PetscCall(DMSwarmMigrate(sw, PETSC_FALSE)); 906 // TODO In higher dimensions, we will have to create different M_p and D_p for each field 907 PetscCall(DMCreateMassMatrix(sw, user->dmN, &M_p)); 908 PetscCall(DMCreateGradientMatrix(sw, user->dmN, &D_p)); 909 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f)); 910 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &v)); 911 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "E_field", &E)); 912 PetscCall(PetscObjectSetName((PetscObject)f, "particle weight")); 913 914 PetscCall(MatViewFromOptions(user->MN, NULL, "-mn_view")); 915 PetscCall(MatViewFromOptions(user->MP, NULL, "-mp_view")); 916 PetscCall(MatViewFromOptions(user->ME, NULL, "-me_view")); 917 PetscCall(VecViewFromOptions(f, NULL, "-weights_view")); 918 919 PetscCall(DMGetGlobalVector(user->dmN, &nrhs)); 920 PetscCall(DMGetGlobalVector(user->dmN, &nflux)); 921 PetscCall(PetscObjectSetName((PetscObject)nrhs, "Weak number density")); 922 PetscCall(DMGetNamedGlobalVector(user->dmN, "n", &n)); 923 PetscCall(DMGetGlobalVector(user->dmP, &prhs)); 924 PetscCall(DMGetGlobalVector(user->dmP, &pflux)); 925 PetscCall(PetscObjectSetName((PetscObject)prhs, "Weak momentum density")); 926 PetscCall(DMGetNamedGlobalVector(user->dmP, "p", &p)); 927 PetscCall(DMGetGlobalVector(user->dmE, &erhs)); 928 PetscCall(DMGetGlobalVector(user->dmE, &eflux)); 929 PetscCall(PetscObjectSetName((PetscObject)erhs, "Weak energy density (pressure)")); 930 PetscCall(DMGetNamedGlobalVector(user->dmE, "e", &e)); 931 932 // Compute moments and fluxes 933 PetscCall(VecDuplicate(f, &tmpMom)); 934 935 PetscCall(MatMultTranspose(M_p, f, nrhs)); 936 937 PetscCall(VecPointwiseMult(tmpMom, f, v)); 938 PetscCall(MatMultTranspose(M_p, tmpMom, prhs)); 939 PetscCall(MatMultTranspose(D_p, tmpMom, nflux)); 940 941 PetscCall(VecPointwiseMult(tmpMom, tmpMom, v)); 942 PetscCall(MatMultTranspose(M_p, tmpMom, erhs)); 943 PetscCall(MatMultTranspose(D_p, tmpMom, pflux)); 944 945 PetscCall(VecPointwiseMult(tmpMom, tmpMom, v)); 946 PetscCall(MatMultTranspose(D_p, tmpMom, eflux)); 947 948 PetscCall(VecPointwiseMult(tmpMom, f, E)); 949 PetscCall(MatMultTransposeAdd(M_p, tmpMom, pflux, pflux)); 950 951 PetscCall(VecPointwiseMult(tmpMom, v, E)); 952 PetscCall(VecScale(tmpMom, 2.)); 953 PetscCall(MatMultTransposeAdd(M_p, tmpMom, eflux, eflux)); 954 955 PetscCall(VecDestroy(&tmpMom)); 956 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &v)); 957 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f)); 958 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "E_field", &E)); 959 960 PetscCall(MatDestroy(&M_p)); 961 PetscCall(MatDestroy(&D_p)); 962 963 PetscCall(KSPCreate(PetscObjectComm((PetscObject)sw), &ksp)); 964 PetscCall(KSPSetOptionsPrefix(ksp, "mom_proj_")); 965 PetscCall(KSPSetOperators(ksp, user->MN, user->MN)); 966 PetscCall(KSPSetFromOptions(ksp)); 967 PetscCall(KSPSolve(ksp, nrhs, n)); 968 PetscCall(KSPSetOperators(ksp, user->MP, user->MP)); 969 PetscCall(KSPSetFromOptions(ksp)); 970 PetscCall(KSPSolve(ksp, prhs, p)); 971 PetscCall(KSPSetOperators(ksp, user->ME, user->ME)); 972 PetscCall(KSPSetFromOptions(ksp)); 973 PetscCall(KSPSolve(ksp, erhs, e)); 974 PetscCall(KSPDestroy(&ksp)); 975 PetscCall(DMRestoreGlobalVector(user->dmN, &nrhs)); 976 PetscCall(DMRestoreGlobalVector(user->dmP, &prhs)); 977 PetscCall(DMRestoreGlobalVector(user->dmE, &erhs)); 978 979 // Check moment residual 980 // TODO Fix global2local here 981 PetscReal res[3], logres[3]; 982 983 PetscCall(DMGetGlobalVector(user->dmMom, &m)); 984 PetscCall(VecISCopy(m, user->isN, SCATTER_FORWARD, n)); 985 PetscCall(VecISCopy(m, user->isP, SCATTER_FORWARD, p)); 986 PetscCall(VecISCopy(m, user->isE, SCATTER_FORWARD, e)); 987 PetscCall(DMGetNamedGlobalVector(user->dmMom, "mold", &mold)); 988 PetscCall(DMGetNamedGlobalVector(user->dmMom, "mfluxold", &mfluxold)); 989 if (!Nts) goto end; 990 991 // e = \Tr{\tau} 992 // M_p w^{k+1} - M_p w^k - \Delta t D_p (w^k \vb{v}^k) = 0 993 // M_p \vb{p}^{k+1} - M_p \vb{p}^k - \Delta t D_p \tau - e \Delta t M_p \left( n \vb{E} \right) = 0 994 // M_p e^{k+1} - M_p e^k - \Delta t D_p \vb{Q} - 2 e \Delta t M_p \left( \vb{p} \cdot \vb{E} \right) = 0 995 PetscCall(DMGetGlobalVector(user->dmMom, &mres)); 996 PetscCall(VecCopy(mfluxold, mres)); 997 PetscCall(VecAXPBYPCZ(mres, 1. / dt, -1. / dt, -1., m, mold)); 998 999 PetscCall(DMGetNamedGlobalVector(user->dmN, "nres", &nres)); 1000 PetscCall(DMGetNamedGlobalVector(user->dmP, "pres", &pres)); 1001 PetscCall(DMGetNamedGlobalVector(user->dmE, "eres", &eres)); 1002 PetscCall(VecISCopy(mres, user->isN, SCATTER_REVERSE, nres)); 1003 PetscCall(VecISCopy(mres, user->isP, SCATTER_REVERSE, pres)); 1004 PetscCall(VecISCopy(mres, user->isE, SCATTER_REVERSE, eres)); 1005 PetscCall(VecNorm(nres, NORM_2, &res[0])); 1006 PetscCall(VecNorm(pres, NORM_2, &res[1])); 1007 PetscCall(VecNorm(eres, NORM_2, &res[2])); 1008 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)sw), "Mass Residual: %g\n", (double)res[0])); 1009 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)sw), "Momentum Residual: %g\n", (double)res[1])); 1010 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)sw), "Energy Residual: %g\n", (double)res[2])); 1011 PetscCall(DMRestoreNamedGlobalVector(user->dmN, "nres", &nres)); 1012 PetscCall(DMRestoreNamedGlobalVector(user->dmP, "pres", &pres)); 1013 PetscCall(DMRestoreNamedGlobalVector(user->dmE, "eres", &eres)); 1014 PetscCall(DMRestoreGlobalVector(user->dmMom, &mres)); 1015 1016 for (PetscInt i = 0; i < 3; ++i) logres[i] = PetscLog10Real(res[i]); 1017 PetscCall(PetscDrawLGAddCommonPoint(user->drawlgMomRes, t, logres)); 1018 PetscCall(PetscDrawLGDraw(user->drawlgMomRes)); 1019 { 1020 PetscDraw draw; 1021 1022 PetscCall(PetscDrawLGGetDraw(user->drawlgMomRes, &draw)); 1023 PetscCall(PetscDrawSave(draw)); 1024 } 1025 1026 end: 1027 PetscCall(VecCopy(m, mold)); 1028 PetscCall(DMRestoreGlobalVector(user->dmMom, &m)); 1029 PetscCall(DMRestoreNamedGlobalVector(user->dmMom, "mold", &mold)); 1030 PetscCall(VecISCopy(mfluxold, user->isN, SCATTER_FORWARD, nflux)); 1031 PetscCall(VecISCopy(mfluxold, user->isP, SCATTER_FORWARD, pflux)); 1032 PetscCall(VecISCopy(mfluxold, user->isE, SCATTER_FORWARD, eflux)); 1033 PetscCall(DMRestoreNamedGlobalVector(user->dmMom, "mfluxold", &mfluxold)); 1034 1035 PetscCall(DMRestoreGlobalVector(user->dmN, &nflux)); 1036 PetscCall(DMRestoreGlobalVector(user->dmP, &pflux)); 1037 PetscCall(DMRestoreGlobalVector(user->dmE, &eflux)); 1038 PetscCall(DMRestoreNamedGlobalVector(user->dmN, "n", &n)); 1039 PetscCall(DMRestoreNamedGlobalVector(user->dmP, "p", &p)); 1040 PetscCall(DMRestoreNamedGlobalVector(user->dmE, "e", &e)); 1041 PetscCall(DMSwarmSetCellDMActive(sw, "space")); 1042 PetscFunctionReturn(PETSC_SUCCESS); 1043 } 1044 1045 static PetscErrorCode MonitorMomentFields(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx) 1046 { 1047 AppCtx *user = (AppCtx *)ctx; 1048 Vec n, p, e; 1049 Vec nres, pres, eres; 1050 1051 PetscFunctionBeginUser; 1052 if (step < 0) PetscFunctionReturn(PETSC_SUCCESS); 1053 PetscCall(ComputeMomentFields(ts)); 1054 1055 PetscCall(DMGetNamedGlobalVector(user->dmN, "n", &n)); 1056 PetscCall(VecView(n, user->viewerN)); 1057 PetscCall(DMRestoreNamedGlobalVector(user->dmN, "n", &n)); 1058 1059 PetscCall(DMGetNamedGlobalVector(user->dmP, "p", &p)); 1060 PetscCall(VecView(p, user->viewerP)); 1061 PetscCall(DMRestoreNamedGlobalVector(user->dmP, "p", &p)); 1062 1063 PetscCall(DMGetNamedGlobalVector(user->dmE, "e", &e)); 1064 PetscCall(VecView(e, user->viewerE)); 1065 PetscCall(DMRestoreNamedGlobalVector(user->dmE, "e", &e)); 1066 1067 PetscCall(DMGetNamedGlobalVector(user->dmN, "nres", &nres)); 1068 PetscCall(VecView(nres, user->viewerNRes)); 1069 PetscCall(DMRestoreNamedGlobalVector(user->dmN, "nres", &nres)); 1070 1071 PetscCall(DMGetNamedGlobalVector(user->dmP, "pres", &pres)); 1072 PetscCall(VecView(pres, user->viewerPRes)); 1073 PetscCall(DMRestoreNamedGlobalVector(user->dmP, "pres", &pres)); 1074 1075 PetscCall(DMGetNamedGlobalVector(user->dmE, "eres", &eres)); 1076 PetscCall(VecView(eres, user->viewerERes)); 1077 PetscCall(DMRestoreNamedGlobalVector(user->dmE, "eres", &eres)); 1078 PetscFunctionReturn(PETSC_SUCCESS); 1079 } 1080 1081 PetscErrorCode MonitorInitialConditions(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx) 1082 { 1083 AppCtx *user = (AppCtx *)ctx; 1084 DM sw; 1085 PetscDraw drawic_x, drawic_v; 1086 PetscReal *weight, *pos, *vel; 1087 PetscInt dim, Np; 1088 1089 PetscFunctionBegin; 1090 if (step < 0) PetscFunctionReturn(PETSC_SUCCESS); /* -1 indicates interpolated solution */ 1091 if (step == 0) { 1092 PetscCall(TSGetDM(ts, &sw)); 1093 PetscCall(DMGetDimension(sw, &dim)); 1094 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 1095 1096 PetscCall(PetscDrawHGReset(user->drawhgic_x)); 1097 PetscCall(PetscDrawHGGetDraw(user->drawhgic_x, &drawic_x)); 1098 PetscCall(PetscDrawClear(drawic_x)); 1099 PetscCall(PetscDrawFlush(drawic_x)); 1100 1101 PetscCall(PetscDrawHGReset(user->drawhgic_v)); 1102 PetscCall(PetscDrawHGGetDraw(user->drawhgic_v, &drawic_v)); 1103 PetscCall(PetscDrawClear(drawic_v)); 1104 PetscCall(PetscDrawFlush(drawic_v)); 1105 1106 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos)); 1107 PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel)); 1108 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight)); 1109 for (PetscInt p = 0; p < Np; ++p) { 1110 PetscCall(PetscDrawHGAddWeightedValue(user->drawhgic_x, pos[p * dim], weight[p])); 1111 PetscCall(PetscDrawHGAddWeightedValue(user->drawhgic_v, vel[p * dim], weight[p])); 1112 } 1113 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&pos)); 1114 PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel)); 1115 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight)); 1116 1117 PetscCall(PetscDrawHGDraw(user->drawhgic_x)); 1118 PetscCall(PetscDrawHGSave(user->drawhgic_x)); 1119 PetscCall(PetscDrawHGDraw(user->drawhgic_v)); 1120 PetscCall(PetscDrawHGSave(user->drawhgic_v)); 1121 } 1122 PetscFunctionReturn(PETSC_SUCCESS); 1123 } 1124 1125 // Right now, make the complete velocity histogram 1126 PetscErrorCode MonitorVelocity(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx) 1127 { 1128 AppCtx *user = (AppCtx *)ctx; 1129 DM sw, dm; 1130 Vec ks; 1131 PetscProbFn *cdf; 1132 PetscDraw drawcell_v; 1133 PetscScalar *ksa; 1134 PetscReal *weight, *vel; 1135 PetscInt *pidx; 1136 PetscInt dim, Npc, cStart, cEnd, cell = user->velocity_monitor; 1137 1138 PetscFunctionBegin; 1139 PetscCall(TSGetDM(ts, &sw)); 1140 PetscCall(DMGetDimension(sw, &dim)); 1141 1142 PetscCall(DMSwarmGetCellDM(sw, &dm)); 1143 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd)); 1144 PetscCall(VecCreate(PetscObjectComm((PetscObject)dm), &ks)); 1145 PetscCall(PetscObjectSetName((PetscObject)ks, "KS Statistic by Cell")); 1146 PetscCall(VecSetSizes(ks, cEnd - cStart, PETSC_DETERMINE)); 1147 PetscCall(VecSetFromOptions(ks)); 1148 switch (dim) { 1149 case 1: 1150 //cdf = PetscCDFMaxwellBoltzmann1D; 1151 cdf = PetscCDFGaussian1D; 1152 break; 1153 case 2: 1154 cdf = PetscCDFMaxwellBoltzmann2D; 1155 break; 1156 case 3: 1157 cdf = PetscCDFMaxwellBoltzmann3D; 1158 break; 1159 default: 1160 SETERRQ(PetscObjectComm((PetscObject)dm), PETSC_ERR_ARG_OUTOFRANGE, "Dimension %" PetscInt_FMT " not supported", dim); 1161 } 1162 1163 PetscCall(PetscDrawHGReset(user->drawhgcell_v)); 1164 PetscCall(PetscDrawHGGetDraw(user->drawhgcell_v, &drawcell_v)); 1165 PetscCall(PetscDrawClear(drawcell_v)); 1166 PetscCall(PetscDrawFlush(drawcell_v)); 1167 1168 PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel)); 1169 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight)); 1170 PetscCall(DMSwarmSortGetAccess(sw)); 1171 PetscCall(VecGetArrayWrite(ks, &ksa)); 1172 for (PetscInt c = cStart; c < cEnd; ++c) { 1173 Vec cellv, cellw; 1174 PetscScalar *cella, *cellaw; 1175 PetscReal totWgt = 0.; 1176 1177 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx)); 1178 PetscCall(VecCreate(PETSC_COMM_SELF, &cellv)); 1179 PetscCall(VecSetBlockSize(cellv, dim)); 1180 PetscCall(VecSetSizes(cellv, Npc * dim, Npc)); 1181 PetscCall(VecSetFromOptions(cellv)); 1182 PetscCall(VecCreate(PETSC_COMM_SELF, &cellw)); 1183 PetscCall(VecSetSizes(cellw, Npc, Npc)); 1184 PetscCall(VecSetFromOptions(cellw)); 1185 PetscCall(VecGetArrayWrite(cellv, &cella)); 1186 PetscCall(VecGetArrayWrite(cellw, &cellaw)); 1187 for (PetscInt q = 0; q < Npc; ++q) { 1188 const PetscInt p = pidx[q]; 1189 if (c == cell) PetscCall(PetscDrawHGAddWeightedValue(user->drawhgcell_v, vel[p * dim], weight[p])); 1190 for (PetscInt d = 0; d < dim; ++d) cella[q * dim + d] = vel[p * dim + d]; 1191 cellaw[q] = weight[p]; 1192 totWgt += weight[p]; 1193 } 1194 PetscCall(VecRestoreArrayWrite(cellv, &cella)); 1195 PetscCall(VecRestoreArrayWrite(cellw, &cellaw)); 1196 PetscCall(VecScale(cellw, 1. / totWgt)); 1197 PetscCall(PetscProbComputeKSStatisticWeighted(cellv, cellw, cdf, &ksa[c - cStart])); 1198 PetscCall(VecDestroy(&cellv)); 1199 PetscCall(VecDestroy(&cellw)); 1200 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx)); 1201 } 1202 PetscCall(VecRestoreArrayWrite(ks, &ksa)); 1203 PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel)); 1204 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight)); 1205 PetscCall(DMSwarmSortRestoreAccess(sw)); 1206 1207 PetscReal minalpha, maxalpha; 1208 PetscInt mincell, maxcell; 1209 1210 PetscCall(VecFilter(ks, PETSC_SMALL)); 1211 PetscCall(VecMin(ks, &mincell, &minalpha)); 1212 PetscCall(VecMax(ks, &maxcell, &maxalpha)); 1213 PetscCall(PetscPrintf(PetscObjectComm((PetscObject)dm), "Step %" PetscInt_FMT ": Min/Max KS statistic %g/%g in cell %" PetscInt_FMT "/%" PetscInt_FMT "\n", step, minalpha, maxalpha, mincell, maxcell)); 1214 PetscCall(VecViewFromOptions(ks, NULL, "-ks_view")); 1215 PetscCall(VecDestroy(&ks)); 1216 1217 PetscCall(PetscDrawHGDraw(user->drawhgcell_v)); 1218 PetscCall(PetscDrawHGSave(user->drawhgcell_v)); 1219 PetscFunctionReturn(PETSC_SUCCESS); 1220 } 1221 1222 static PetscErrorCode MonitorPositions_2D(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx) 1223 { 1224 AppCtx *user = (AppCtx *)ctx; 1225 DM dm, sw; 1226 PetscScalar *x, *v, *weight; 1227 PetscReal lower[3], upper[3], speed; 1228 const PetscInt *s; 1229 PetscInt dim, cStart, cEnd, c; 1230 1231 PetscFunctionBeginUser; 1232 if (step > 0 && step % user->ostep == 0) { 1233 PetscCall(TSGetDM(ts, &sw)); 1234 PetscCall(DMSwarmGetCellDM(sw, &dm)); 1235 PetscCall(DMGetDimension(dm, &dim)); 1236 PetscCall(DMGetBoundingBox(dm, lower, upper)); 1237 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd)); 1238 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x)); 1239 PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v)); 1240 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight)); 1241 PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&s)); 1242 PetscCall(DMSwarmSortGetAccess(sw)); 1243 PetscCall(PetscDrawSPReset(user->drawspX)); 1244 PetscCall(PetscDrawSPSetLimits(user->drawspX, lower[0], upper[0], lower[1], upper[1])); 1245 PetscCall(PetscDrawSPSetLimits(user->drawspX, lower[0], upper[0], -12, 12)); 1246 for (c = 0; c < cEnd - cStart; ++c) { 1247 PetscInt *pidx, Npc, q; 1248 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx)); 1249 for (q = 0; q < Npc; ++q) { 1250 const PetscInt p = pidx[q]; 1251 if (s[p] == 0) { 1252 speed = 0.; 1253 for (PetscInt d = 0; d < dim; ++d) speed += PetscSqr(v[p * dim + d]); 1254 speed = PetscSqrtReal(speed); 1255 if (dim == 1) { 1256 PetscCall(PetscDrawSPAddPointColorized(user->drawspX, &x[p * dim], &v[p * dim], &speed)); 1257 } else { 1258 PetscCall(PetscDrawSPAddPointColorized(user->drawspX, &x[p * dim], &x[p * dim + 1], &speed)); 1259 } 1260 } else if (s[p] == 1) { 1261 PetscCall(PetscDrawSPAddPoint(user->drawspX, &x[p * dim], &v[p * dim])); 1262 } 1263 } 1264 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx)); 1265 } 1266 PetscCall(PetscDrawSPDraw(user->drawspX, PETSC_TRUE)); 1267 PetscDraw draw; 1268 PetscCall(PetscDrawSPGetDraw(user->drawspX, &draw)); 1269 PetscCall(PetscDrawSave(draw)); 1270 PetscCall(DMSwarmSortRestoreAccess(sw)); 1271 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x)); 1272 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight)); 1273 PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v)); 1274 PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&s)); 1275 } 1276 PetscFunctionReturn(PETSC_SUCCESS); 1277 } 1278 1279 static PetscErrorCode MonitorPoisson(TS ts, PetscInt step, PetscReal t, Vec U, void *ctx) 1280 { 1281 AppCtx *user = (AppCtx *)ctx; 1282 DM dm, sw; 1283 1284 PetscFunctionBeginUser; 1285 if (step > 0 && step % user->ostep == 0) { 1286 PetscCall(TSGetDM(ts, &sw)); 1287 PetscCall(DMSwarmGetCellDM(sw, &dm)); 1288 1289 if (user->validE) { 1290 PetscScalar *x, *E, *weight; 1291 PetscReal lower[3], upper[3], xval; 1292 PetscDraw draw; 1293 PetscInt dim, cStart, cEnd; 1294 1295 PetscCall(DMGetDimension(dm, &dim)); 1296 PetscCall(DMGetBoundingBox(dm, lower, upper)); 1297 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd)); 1298 1299 PetscCall(PetscDrawSPReset(user->drawspE)); 1300 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x)); 1301 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E)); 1302 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight)); 1303 1304 PetscCall(DMSwarmSortGetAccess(sw)); 1305 for (PetscInt c = 0; c < cEnd - cStart; ++c) { 1306 PetscReal Eavg = 0.0; 1307 PetscInt *pidx, Npc; 1308 1309 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx)); 1310 for (PetscInt q = 0; q < Npc; ++q) { 1311 const PetscInt p = pidx[q]; 1312 Eavg += E[p * dim]; 1313 } 1314 Eavg /= Npc; 1315 xval = (c + 0.5) * ((upper[0] - lower[0]) / (cEnd - cStart)); 1316 PetscCall(PetscDrawSPAddPoint(user->drawspE, &xval, &Eavg)); 1317 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx)); 1318 } 1319 PetscCall(PetscDrawSPDraw(user->drawspE, PETSC_TRUE)); 1320 PetscCall(PetscDrawSPGetDraw(user->drawspE, &draw)); 1321 PetscCall(PetscDrawSave(draw)); 1322 PetscCall(DMSwarmSortRestoreAccess(sw)); 1323 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x)); 1324 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight)); 1325 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E)); 1326 } 1327 1328 Vec rho, rhohat, phi; 1329 1330 PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho)); 1331 PetscCall(DMGetNamedGlobalVector(user->dmPot, "rhohat", &rhohat)); 1332 PetscCall(VecView(rho, user->viewerRho)); 1333 PetscCall(VecISCopy(user->fftX, user->fftReal, SCATTER_FORWARD, rho)); 1334 PetscCall(MatMult(user->fftPot, user->fftX, user->fftY)); 1335 PetscCall(VecFilter(user->fftY, PETSC_SMALL)); 1336 PetscCall(VecViewFromOptions(user->fftX, NULL, "-real_view")); 1337 PetscCall(VecViewFromOptions(user->fftY, NULL, "-fft_view")); 1338 PetscCall(VecISCopy(user->fftY, user->fftReal, SCATTER_REVERSE, rhohat)); 1339 PetscCall(VecSetValue(rhohat, 0, 0., INSERT_VALUES)); // Remove large DC component 1340 PetscCall(VecView(rhohat, user->viewerRhoHat)); 1341 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho)); 1342 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rhohat", &rhohat)); 1343 1344 PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi)); 1345 PetscCall(VecView(phi, user->viewerPhi)); 1346 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi)); 1347 } 1348 PetscFunctionReturn(PETSC_SUCCESS); 1349 } 1350 1351 static PetscErrorCode SetupParameters(MPI_Comm comm, AppCtx *ctx) 1352 { 1353 PetscBag bag; 1354 Parameter *p; 1355 1356 PetscFunctionBeginUser; 1357 /* setup PETSc parameter bag */ 1358 PetscCall(PetscBagGetData(ctx->bag, (void **)&p)); 1359 PetscCall(PetscBagSetName(ctx->bag, "par", "Vlasov-Poisson Parameters")); 1360 bag = ctx->bag; 1361 PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "v0", "Velocity scale, m/s")); 1362 PetscCall(PetscBagRegisterScalar(bag, &p->t0, 1.0, "t0", "Time scale, s")); 1363 PetscCall(PetscBagRegisterScalar(bag, &p->x0, 1.0, "x0", "Space scale, m")); 1364 PetscCall(PetscBagRegisterScalar(bag, &p->v0, 1.0, "phi0", "Potential scale, kg*m^2/A*s^3")); 1365 PetscCall(PetscBagRegisterScalar(bag, &p->q0, 1.0, "q0", "Charge Scale, A*s")); 1366 PetscCall(PetscBagRegisterScalar(bag, &p->m0, 1.0, "m0", "Mass Scale, kg")); 1367 PetscCall(PetscBagRegisterScalar(bag, &p->epsi0, 1.0, "epsi0", "Permittivity of Free Space, kg")); 1368 PetscCall(PetscBagRegisterScalar(bag, &p->kb, 1.0, "kb", "Boltzmann Constant, m^2 kg/s^2 K^1")); 1369 1370 PetscCall(PetscBagRegisterScalar(bag, &p->sigma, 1.0, "sigma", "Charge per unit area, C/m^3")); 1371 PetscCall(PetscBagRegisterScalar(bag, &p->poissonNumber, 1.0, "poissonNumber", "Non-Dimensional Poisson Number")); 1372 PetscCall(PetscBagRegisterScalar(bag, &p->vlasovNumber, 1.0, "vlasovNumber", "Non-Dimensional Vlasov Number")); 1373 PetscCall(PetscBagSetFromOptions(bag)); 1374 { 1375 PetscViewer viewer; 1376 PetscViewerFormat format; 1377 PetscBool flg; 1378 1379 PetscCall(PetscOptionsCreateViewer(comm, NULL, NULL, "-param_view", &viewer, &format, &flg)); 1380 if (flg) { 1381 PetscCall(PetscViewerPushFormat(viewer, format)); 1382 PetscCall(PetscBagView(bag, viewer)); 1383 PetscCall(PetscViewerFlush(viewer)); 1384 PetscCall(PetscViewerPopFormat(viewer)); 1385 PetscCall(PetscViewerDestroy(&viewer)); 1386 } 1387 } 1388 PetscFunctionReturn(PETSC_SUCCESS); 1389 } 1390 1391 static PetscErrorCode CreateMesh(MPI_Comm comm, AppCtx *user, DM *dm) 1392 { 1393 PetscFunctionBeginUser; 1394 PetscCall(DMCreate(comm, dm)); 1395 PetscCall(DMSetType(*dm, DMPLEX)); 1396 PetscCall(DMSetFromOptions(*dm)); 1397 PetscCall(PetscObjectSetName((PetscObject)*dm, "space")); 1398 PetscCall(DMViewFromOptions(*dm, NULL, "-dm_view")); 1399 1400 // Cache the mesh geometry 1401 DMField coordField; 1402 IS cellIS; 1403 PetscQuadrature quad; 1404 PetscReal *wt, *pt; 1405 PetscInt cdim, cStart, cEnd; 1406 1407 PetscCall(DMGetCoordinateField(*dm, &coordField)); 1408 PetscCheck(coordField, comm, PETSC_ERR_USER, "DM must have a coordinate field"); 1409 PetscCall(DMGetCoordinateDim(*dm, &cdim)); 1410 PetscCall(DMPlexGetHeightStratum(*dm, 0, &cStart, &cEnd)); 1411 PetscCall(ISCreateStride(PETSC_COMM_SELF, cEnd - cStart, cStart, 1, &cellIS)); 1412 PetscCall(PetscQuadratureCreate(PETSC_COMM_SELF, &quad)); 1413 PetscCall(PetscMalloc1(1, &wt)); 1414 PetscCall(PetscMalloc1(2, &pt)); 1415 wt[0] = 1.; 1416 pt[0] = -1.; 1417 pt[1] = -1.; 1418 PetscCall(PetscQuadratureSetData(quad, cdim, 1, 1, pt, wt)); 1419 PetscCall(DMFieldCreateFEGeom(coordField, cellIS, quad, PETSC_FEGEOM_BASIC, &user->fegeom)); 1420 PetscCall(PetscQuadratureDestroy(&quad)); 1421 PetscCall(ISDestroy(&cellIS)); 1422 PetscFunctionReturn(PETSC_SUCCESS); 1423 } 1424 1425 static void ion_f0(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[]) 1426 { 1427 f0[0] = -constants[SIGMA]; 1428 } 1429 1430 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[]) 1431 { 1432 PetscInt d; 1433 for (d = 0; d < dim; ++d) f1[d] = u_x[d]; 1434 } 1435 1436 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[]) 1437 { 1438 PetscInt d; 1439 for (d = 0; d < dim; ++d) g3[d * dim + d] = 1.0; 1440 } 1441 1442 /* 1443 / I -grad\ / q \ = /0\ 1444 \-div 0 / \phi/ \f/ 1445 */ 1446 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[]) 1447 { 1448 for (PetscInt d = 0; d < dim; ++d) f0[d] += u[uOff[0] + d]; 1449 } 1450 1451 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[]) 1452 { 1453 for (PetscInt d = 0; d < dim; ++d) f1[d * dim + d] = u[uOff[1]]; 1454 } 1455 1456 static void f0_phi_backgroundCharge(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[]) 1457 { 1458 f0[0] += constants[SIGMA]; 1459 for (PetscInt d = 0; d < dim; ++d) f0[0] += u_x[uOff_x[0] + d * dim + d]; 1460 } 1461 1462 /* Boundary residual. Dirichlet boundary for u means u_bdy=p*n */ 1463 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[]) 1464 { 1465 for (PetscInt d = 0; d < dim; ++d) g0[d * dim + d] = 1.0; 1466 } 1467 1468 static void g2_qphi(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[]) 1469 { 1470 for (PetscInt d = 0; d < dim; ++d) g2[d * dim + d] = 1.0; 1471 } 1472 1473 static void g1_phiq(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[]) 1474 { 1475 for (PetscInt d = 0; d < dim; ++d) g1[d * dim + d] = 1.0; 1476 } 1477 1478 static PetscErrorCode CreateFEM(DM dm, AppCtx *user) 1479 { 1480 PetscFE fephi, feq; 1481 PetscDS ds; 1482 PetscBool simplex; 1483 PetscInt dim; 1484 1485 PetscFunctionBeginUser; 1486 PetscCall(DMGetDimension(dm, &dim)); 1487 PetscCall(DMPlexIsSimplex(dm, &simplex)); 1488 if (user->em == EM_MIXED) { 1489 DMLabel label; 1490 const PetscInt id = 1; 1491 1492 PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, dim, simplex, "field_", PETSC_DETERMINE, &feq)); 1493 PetscCall(PetscObjectSetName((PetscObject)feq, "field")); 1494 PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, "potential_", PETSC_DETERMINE, &fephi)); 1495 PetscCall(PetscObjectSetName((PetscObject)fephi, "potential")); 1496 PetscCall(PetscFECopyQuadrature(feq, fephi)); 1497 PetscCall(DMSetField(dm, 0, NULL, (PetscObject)feq)); 1498 PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fephi)); 1499 PetscCall(DMCreateDS(dm)); 1500 PetscCall(PetscFEDestroy(&fephi)); 1501 PetscCall(PetscFEDestroy(&feq)); 1502 1503 PetscCall(DMGetLabel(dm, "marker", &label)); 1504 PetscCall(DMGetDS(dm, &ds)); 1505 1506 PetscCall(PetscDSSetResidual(ds, 0, f0_q, f1_q)); 1507 PetscCall(PetscDSSetResidual(ds, 1, f0_phi_backgroundCharge, NULL)); 1508 PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_qq, NULL, NULL, NULL)); 1509 PetscCall(PetscDSSetJacobian(ds, 0, 1, NULL, NULL, g2_qphi, NULL)); 1510 PetscCall(PetscDSSetJacobian(ds, 1, 0, NULL, g1_phiq, NULL, NULL)); 1511 1512 PetscCall(DMAddBoundary(dm, DM_BC_ESSENTIAL, "wall", label, 1, &id, 0, 0, NULL, (void (*)(void))zero, NULL, NULL, NULL)); 1513 1514 } else { 1515 MatNullSpace nullsp; 1516 PetscCall(PetscFECreateDefault(PETSC_COMM_SELF, dim, 1, simplex, NULL, PETSC_DETERMINE, &fephi)); 1517 PetscCall(PetscObjectSetName((PetscObject)fephi, "potential")); 1518 PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fephi)); 1519 PetscCall(DMCreateDS(dm)); 1520 PetscCall(DMGetDS(dm, &ds)); 1521 PetscCall(PetscDSSetResidual(ds, 0, ion_f0, laplacian_f1)); 1522 PetscCall(PetscDSSetJacobian(ds, 0, 0, NULL, NULL, NULL, laplacian_g3)); 1523 PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullsp)); 1524 PetscCall(PetscObjectCompose((PetscObject)fephi, "nullspace", (PetscObject)nullsp)); 1525 PetscCall(MatNullSpaceDestroy(&nullsp)); 1526 PetscCall(PetscFEDestroy(&fephi)); 1527 } 1528 PetscFunctionReturn(PETSC_SUCCESS); 1529 } 1530 1531 static PetscErrorCode CreatePoisson(DM dm, AppCtx *user) 1532 { 1533 SNES snes; 1534 Mat J; 1535 MatNullSpace nullSpace; 1536 1537 PetscFunctionBeginUser; 1538 PetscCall(CreateFEM(dm, user)); 1539 PetscCall(SNESCreate(PetscObjectComm((PetscObject)dm), &snes)); 1540 PetscCall(SNESSetOptionsPrefix(snes, "em_")); 1541 PetscCall(SNESSetDM(snes, dm)); 1542 PetscCall(DMPlexSetSNESLocalFEM(dm, PETSC_FALSE, user)); 1543 PetscCall(SNESSetFromOptions(snes)); 1544 1545 PetscCall(DMCreateMatrix(dm, &J)); 1546 PetscCall(MatNullSpaceCreate(PetscObjectComm((PetscObject)dm), PETSC_TRUE, 0, NULL, &nullSpace)); 1547 PetscCall(MatSetNullSpace(J, nullSpace)); 1548 PetscCall(MatNullSpaceDestroy(&nullSpace)); 1549 PetscCall(SNESSetJacobian(snes, J, J, NULL, NULL)); 1550 PetscCall(MatDestroy(&J)); 1551 if (user->em == EM_MIXED) { 1552 const PetscInt potential = 1; 1553 1554 PetscCall(DMCreateSubDM(dm, 1, &potential, &user->isPot, &user->dmPot)); 1555 } else { 1556 user->dmPot = dm; 1557 PetscCall(PetscObjectReference((PetscObject)user->dmPot)); 1558 } 1559 PetscCall(DMCreateMassMatrix(user->dmPot, user->dmPot, &user->M)); 1560 PetscCall(DMPlexCreateClosureIndex(dm, NULL)); 1561 user->snes = snes; 1562 PetscFunctionReturn(PETSC_SUCCESS); 1563 } 1564 1565 // Conservation of mass (m = 1.0) 1566 // n_t + 1/ m p_x = 0 1567 static void f0_mass(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[]) 1568 { 1569 for (PetscInt d = 0; d < dim; ++d) f0[0] += u_t[uOff[0]] + u_x[uOff_x[1] + d * dim + d]; 1570 } 1571 1572 // Conservation of momentum (m = 1, e = 1) 1573 // p_t + (u p)_x = -pr_x + e n E 1574 // p_t + (div u) p + u . grad p = -pr_x + e n E 1575 // p_t + (div p) p / n - (p . grad n) p / n^2 + p / n . grad p = -pr_x + e n E 1576 static void f0_momentum(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[]) 1577 { 1578 const PetscScalar n = u[uOff[0]]; 1579 1580 for (PetscInt d = 0; d < dim; ++d) { 1581 PetscReal divp = 0.; 1582 1583 f0[d] += u_t[uOff[1] + d]; 1584 for (PetscInt e = 0; e < dim; ++e) { 1585 f0[d] += u[uOff[1] + e] * u_x[uOff_x[1] + d * dim + e] / n; // p / n . grad p 1586 f0[d] -= (u[uOff[1] + e] * u_x[uOff_x[0] + e]) * u[uOff[1] + d] / PetscSqr(n); // -(p . grad n) p / n^2 1587 divp += u_x[uOff_x[1] + e * dim + e]; 1588 } 1589 f0[d] += divp * u[uOff[1] + d] / n; // (div p) p / n 1590 f0[d] += u_x[uOff_x[2] + d]; // pr_x 1591 f0[d] -= n * a[d]; // -e n E 1592 } 1593 } 1594 1595 // Conservation of energy 1596 // pr_t + (u pr)_x = -3 pr u_x - q_x 1597 // pr_t + (div u) pr + u . grad pr = -3 pr (div u) - q_x 1598 // pr_t + 4 (div u) pr + u . grad pr = -q_x 1599 // pr_t + 4 div p pr / n - 4 (p . grad n) pr / n^2 + p . grad pr / n = -q_x 1600 static void f0_energy(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[]) 1601 { 1602 const PetscScalar n = u[uOff[0]]; 1603 const PetscScalar pr = u[uOff[2]]; 1604 PetscReal divp = 0.; 1605 1606 f0[0] += u_t[uOff[2]]; 1607 for (PetscInt d = 0; d < dim; ++d) { 1608 f0[0] += u[uOff[1] + d] * u_x[uOff_x[2] + d] / n; // p . grad pr / n 1609 f0[0] -= 4. * u[uOff[1] + d] * u_x[uOff_x[0] + d] * pr / PetscSqr(n); // -4 (p . grad n) pr / n^2 1610 divp += u_x[uOff_x[1] + d * dim + d]; 1611 } 1612 f0[0] += 4. * divp * pr / n; // 4 div p pr / n 1613 } 1614 1615 static PetscErrorCode SetupMomentProblem(DM dm, AppCtx *ctx) 1616 { 1617 PetscDS ds; 1618 1619 PetscFunctionBegin; 1620 PetscCall(DMGetDS(dm, &ds)); 1621 PetscCall(PetscDSSetResidual(ds, 0, f0_mass, NULL)); 1622 PetscCall(PetscDSSetResidual(ds, 1, f0_momentum, NULL)); 1623 PetscCall(PetscDSSetResidual(ds, 2, f0_energy, NULL)); 1624 //PetscCall(PetscDSSetJacobian(ds, 0, 0, g0_mass_uu, NULL, NULL, NULL)); 1625 PetscFunctionReturn(PETSC_SUCCESS); 1626 } 1627 1628 static PetscErrorCode CreateMomentFields(DM odm, AppCtx *user) 1629 { 1630 DM dm; 1631 PetscFE fe; 1632 DMPolytopeType ct; 1633 PetscInt dim, cStart; 1634 1635 PetscFunctionBeginUser; 1636 PetscCall(DMClone(odm, &dm)); 1637 PetscCall(DMGetDimension(dm, &dim)); 1638 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, NULL)); 1639 PetscCall(DMPlexGetCellType(dm, cStart, &ct)); 1640 PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, NULL, PETSC_DETERMINE, &fe)); 1641 PetscCall(PetscObjectSetName((PetscObject)fe, "number density")); 1642 PetscCall(DMSetField(dm, 0, NULL, (PetscObject)fe)); 1643 PetscCall(PetscFEDestroy(&fe)); 1644 PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, dim, ct, NULL, PETSC_DETERMINE, &fe)); 1645 PetscCall(PetscObjectSetName((PetscObject)fe, "momentum density")); 1646 PetscCall(DMSetField(dm, 1, NULL, (PetscObject)fe)); 1647 PetscCall(PetscFEDestroy(&fe)); 1648 PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, NULL, PETSC_DETERMINE, &fe)); 1649 PetscCall(PetscObjectSetName((PetscObject)fe, "energy density")); 1650 PetscCall(DMSetField(dm, 2, NULL, (PetscObject)fe)); 1651 PetscCall(PetscFEDestroy(&fe)); 1652 PetscCall(DMCreateDS(dm)); 1653 PetscCall(SetupMomentProblem(dm, user)); 1654 user->dmMom = dm; 1655 PetscInt field; 1656 1657 field = 0; 1658 PetscCall(DMCreateSubDM(user->dmMom, 1, &field, &user->isN, &user->dmN)); 1659 PetscCall(DMCreateMassMatrix(user->dmN, user->dmN, &user->MN)); 1660 field = 1; 1661 PetscCall(DMCreateSubDM(user->dmMom, 1, &field, &user->isP, &user->dmP)); 1662 PetscCall(DMCreateMassMatrix(user->dmP, user->dmP, &user->MP)); 1663 field = 2; 1664 PetscCall(DMCreateSubDM(user->dmMom, 1, &field, &user->isE, &user->dmE)); 1665 PetscCall(DMCreateMassMatrix(user->dmE, user->dmE, &user->ME)); 1666 PetscFunctionReturn(PETSC_SUCCESS); 1667 } 1668 1669 PetscErrorCode PetscPDFPertubedConstant2D(const PetscReal x[], const PetscReal dummy[], PetscReal p[]) 1670 { 1671 p[0] = (1 + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI); 1672 p[1] = (1 + 0.01 * PetscCosReal(0.5 * x[1])) / (2 * PETSC_PI); 1673 return PETSC_SUCCESS; 1674 } 1675 PetscErrorCode PetscPDFPertubedConstant1D(const PetscReal x[], const PetscReal dummy[], PetscReal p[]) 1676 { 1677 p[0] = (1. + 0.01 * PetscCosReal(0.5 * x[0])) / (2 * PETSC_PI); 1678 return PETSC_SUCCESS; 1679 } 1680 1681 PetscErrorCode PetscPDFCosine1D(const PetscReal x[], const PetscReal scale[], PetscReal p[]) 1682 { 1683 const PetscReal alpha = scale ? scale[0] : 0.0; 1684 const PetscReal k = scale ? scale[1] : 1.; 1685 p[0] = (1 + alpha * PetscCosReal(k * x[0])); 1686 return PETSC_SUCCESS; 1687 } 1688 1689 PetscErrorCode PetscPDFCosine2D(const PetscReal x[], const PetscReal scale[], PetscReal p[]) 1690 { 1691 const PetscReal alpha = scale ? scale[0] : 0.; 1692 const PetscReal k = scale ? scale[0] : 1.; 1693 p[0] = (1 + alpha * PetscCosReal(k * (x[0] + x[1]))); 1694 return PETSC_SUCCESS; 1695 } 1696 1697 static PetscErrorCode CreateVelocityDM(DM sw, DM *vdm) 1698 { 1699 PetscFE fe; 1700 DMPolytopeType ct; 1701 PetscInt dim, cStart; 1702 const char *prefix = "v"; 1703 1704 PetscFunctionBegin; 1705 PetscCall(DMCreate(PETSC_COMM_SELF, vdm)); 1706 PetscCall(DMSetType(*vdm, DMPLEX)); 1707 PetscCall(DMPlexSetOptionsPrefix(*vdm, prefix)); 1708 PetscCall(DMSetFromOptions(*vdm)); 1709 PetscCall(PetscObjectSetName((PetscObject)*vdm, "velocity")); 1710 PetscCall(DMViewFromOptions(*vdm, NULL, "-dm_view")); 1711 1712 PetscCall(DMGetDimension(*vdm, &dim)); 1713 PetscCall(DMPlexGetHeightStratum(*vdm, 0, &cStart, NULL)); 1714 PetscCall(DMPlexGetCellType(*vdm, cStart, &ct)); 1715 PetscCall(PetscFECreateByCell(PETSC_COMM_SELF, dim, 1, ct, prefix, PETSC_DETERMINE, &fe)); 1716 PetscCall(PetscObjectSetName((PetscObject)fe, "distribution")); 1717 PetscCall(DMSetField(*vdm, 0, NULL, (PetscObject)fe)); 1718 PetscCall(DMCreateDS(*vdm)); 1719 PetscCall(PetscFEDestroy(&fe)); 1720 PetscFunctionReturn(PETSC_SUCCESS); 1721 } 1722 1723 /* 1724 InitializeParticles_Centroid - Initialize a regular grid of particles. 1725 1726 Input Parameters: 1727 + sw - The `DMSWARM` 1728 - force1D - Treat the spatial domain as 1D 1729 1730 Notes: 1731 This functions sets the species, cellid, spatial coordinate, and velocity fields for all particles. 1732 1733 It places one particle in the centroid of each cell in the implicit tensor product of the spatial 1734 and velocity meshes. 1735 */ 1736 static PetscErrorCode InitializeParticles_Centroid(DM sw) 1737 { 1738 DM_Swarm *swarm = (DM_Swarm *)sw->data; 1739 DMSwarmCellDM celldm; 1740 DM xdm, vdm; 1741 PetscReal vmin[3], vmax[3]; 1742 PetscReal *x, *v; 1743 PetscInt *species, *cellid; 1744 PetscInt dim, xcStart, xcEnd, vcStart, vcEnd, Ns, Np, Npc, debug; 1745 PetscBool flg; 1746 MPI_Comm comm; 1747 const char *cellidname; 1748 1749 PetscFunctionBegin; 1750 PetscCall(PetscObjectGetComm((PetscObject)sw, &comm)); 1751 1752 PetscOptionsBegin(comm, "", "DMSwarm Options", "DMSWARM"); 1753 PetscCall(DMSwarmGetNumSpecies(sw, &Ns)); 1754 PetscCall(PetscOptionsInt("-dm_swarm_num_species", "The number of species", "DMSwarmSetNumSpecies", Ns, &Ns, &flg)); 1755 if (flg) PetscCall(DMSwarmSetNumSpecies(sw, Ns)); 1756 PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_coords", "Debug output level for particle coordinate computations", "InitializeParticles", 0, &swarm->printCoords, NULL, 0)); 1757 PetscCall(PetscOptionsBoundedInt("-dm_swarm_print_weights", "Debug output level for particle weight computations", "InitializeWeights", 0, &swarm->printWeights, NULL, 0)); 1758 PetscOptionsEnd(); 1759 debug = swarm->printCoords; 1760 1761 PetscCall(DMGetDimension(sw, &dim)); 1762 PetscCall(DMSwarmGetCellDM(sw, &xdm)); 1763 PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd)); 1764 1765 PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm)); 1766 PetscCall(DMSwarmCellDMGetDM(celldm, &vdm)); 1767 PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd)); 1768 1769 // One particle per centroid on the tensor product grid 1770 Npc = (vcEnd - vcStart) * Ns; 1771 Np = (xcEnd - xcStart) * Npc; 1772 PetscCall(DMSwarmSetLocalSizes(sw, Np, 0)); 1773 if (debug) { 1774 PetscInt gNp, gNc, Nc = xcEnd - xcStart; 1775 PetscCallMPI(MPIU_Allreduce(&Np, &gNp, 1, MPIU_INT, MPIU_SUM, comm)); 1776 PetscCall(PetscPrintf(comm, "Global Np = %" PetscInt_FMT "\n", gNp)); 1777 PetscCallMPI(MPIU_Allreduce(&Nc, &gNc, 1, MPIU_INT, MPIU_SUM, comm)); 1778 PetscCall(PetscPrintf(comm, "Global X-cells = %" PetscInt_FMT "\n", gNc)); 1779 PetscCall(PetscPrintf(comm, "Global V-cells = %" PetscInt_FMT "\n", vcEnd - vcStart)); 1780 } 1781 1782 // Set species and cellid 1783 PetscCall(DMSwarmGetCellDMActive(sw, &celldm)); 1784 PetscCall(DMSwarmCellDMGetCellID(celldm, &cellidname)); 1785 PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species)); 1786 PetscCall(DMSwarmGetField(sw, cellidname, NULL, NULL, (void **)&cellid)); 1787 for (PetscInt c = 0, p = 0; c < xcEnd - xcStart; ++c) { 1788 for (PetscInt s = 0; s < Ns; ++s) { 1789 for (PetscInt q = 0; q < Npc / Ns; ++q, ++p) { 1790 species[p] = s; 1791 cellid[p] = c; 1792 } 1793 } 1794 } 1795 PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species)); 1796 PetscCall(DMSwarmRestoreField(sw, cellidname, NULL, NULL, (void **)&cellid)); 1797 1798 // Set particle coordinates 1799 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x)); 1800 PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&v)); 1801 PetscCall(DMSwarmSortGetAccess(sw)); 1802 PetscCall(DMGetBoundingBox(vdm, vmin, vmax)); 1803 PetscCall(DMGetCoordinatesLocalSetUp(xdm)); 1804 for (PetscInt c = 0; c < xcEnd - xcStart; ++c) { 1805 const PetscInt cell = c + xcStart; 1806 PetscInt *pidx, Npc; 1807 PetscReal centroid[3], volume; 1808 1809 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx)); 1810 PetscCall(DMPlexComputeCellGeometryFVM(xdm, cell, &volume, centroid, NULL)); 1811 for (PetscInt s = 0; s < Ns; ++s) { 1812 for (PetscInt q = 0; q < Npc / Ns; ++q) { 1813 const PetscInt p = pidx[q * Ns + s]; 1814 1815 for (PetscInt d = 0; d < dim; ++d) { 1816 x[p * dim + d] = centroid[d]; 1817 v[p * dim + d] = vmin[0] + (q + 0.5) * ((vmax[0] - vmin[0]) / (Npc / Ns)); 1818 } 1819 if (debug > 1) { 1820 PetscCall(PetscPrintf(PETSC_COMM_SELF, "Particle %4" PetscInt_FMT " ", p)); 1821 PetscCall(PetscPrintf(PETSC_COMM_SELF, " x: (")); 1822 for (PetscInt d = 0; d < dim; ++d) { 1823 if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", ")); 1824 PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", x[p * dim + d])); 1825 } 1826 PetscCall(PetscPrintf(PETSC_COMM_SELF, ") v:(")); 1827 for (PetscInt d = 0; d < dim; ++d) { 1828 if (d > 0) PetscCall(PetscPrintf(PETSC_COMM_SELF, ", ")); 1829 PetscCall(PetscPrintf(PETSC_COMM_SELF, "%g", v[p * dim + d])); 1830 } 1831 PetscCall(PetscPrintf(PETSC_COMM_SELF, ")\n")); 1832 } 1833 } 1834 } 1835 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx)); 1836 } 1837 PetscCall(DMSwarmSortRestoreAccess(sw)); 1838 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&x)); 1839 PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&v)); 1840 PetscFunctionReturn(PETSC_SUCCESS); 1841 } 1842 1843 /* 1844 InitializeWeights - Compute weight for each local particle 1845 1846 Input Parameters: 1847 + sw - The `DMSwarm` 1848 . totalWeight - The sum of all particle weights 1849 . func - The PDF for the particle spatial distribution 1850 - param - The PDF parameters 1851 1852 Notes: 1853 The PDF for velocity is assumed to be a Gaussian 1854 1855 The particle weights are returned in the `w_q` field of `sw`. 1856 */ 1857 static PetscErrorCode InitializeWeights(DM sw, PetscReal totalWeight, PetscProbFn *func, const PetscReal param[]) 1858 { 1859 DM xdm, vdm; 1860 DMSwarmCellDM celldm; 1861 PetscScalar *weight; 1862 PetscQuadrature xquad; 1863 const PetscReal *xq, *xwq; 1864 const PetscInt order = 5; 1865 PetscReal xi0[3]; 1866 PetscReal xwtot = 0., pwtot = 0.; 1867 PetscInt xNq; 1868 PetscInt dim, Ns, xcStart, xcEnd, vcStart, vcEnd, debug = ((DM_Swarm *)sw->data)->printWeights; 1869 MPI_Comm comm; 1870 PetscMPIInt rank; 1871 1872 PetscFunctionBegin; 1873 PetscCall(PetscObjectGetComm((PetscObject)sw, &comm)); 1874 PetscCallMPI(MPI_Comm_rank(comm, &rank)); 1875 PetscCall(DMGetDimension(sw, &dim)); 1876 PetscCall(DMSwarmGetCellDM(sw, &xdm)); 1877 PetscCall(DMSwarmGetNumSpecies(sw, &Ns)); 1878 PetscCall(DMPlexGetHeightStratum(xdm, 0, &xcStart, &xcEnd)); 1879 PetscCall(DMSwarmGetCellDMByName(sw, "velocity", &celldm)); 1880 PetscCall(DMSwarmCellDMGetDM(celldm, &vdm)); 1881 PetscCall(DMPlexGetHeightStratum(vdm, 0, &vcStart, &vcEnd)); 1882 1883 // Setup Quadrature for spatial and velocity weight calculations 1884 PetscCall(PetscDTGaussTensorQuadrature(dim, 1, order, -1.0, 1.0, &xquad)); 1885 PetscCall(PetscQuadratureGetData(xquad, NULL, NULL, &xNq, &xq, &xwq)); 1886 for (PetscInt d = 0; d < dim; ++d) xi0[d] = -1.0; 1887 1888 // Integrate the density function to get the weights of particles in each cell 1889 PetscCall(DMGetCoordinatesLocalSetUp(vdm)); 1890 PetscCall(DMSwarmSortGetAccess(sw)); 1891 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight)); 1892 for (PetscInt c = xcStart; c < xcEnd; ++c) { 1893 PetscReal xv0[3], xJ[9], xinvJ[9], xdetJ, xqr[3], xden, xw = 0.; 1894 PetscInt *pidx, Npc; 1895 PetscInt xNc; 1896 const PetscScalar *xarray; 1897 PetscScalar *xcoords = NULL; 1898 PetscBool xisDG; 1899 1900 PetscCall(DMPlexGetCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords)); 1901 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Npc, &pidx)); 1902 PetscCheck(Npc == (vcEnd - vcStart) * Ns, PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Number of particles %" PetscInt_FMT " in cell (rank %d) != %" PetscInt_FMT " number of velocity vertices", Npc, rank, (vcEnd - vcStart) * Ns); 1903 PetscCall(DMPlexComputeCellGeometryFEM(xdm, c, NULL, xv0, xJ, xinvJ, &xdetJ)); 1904 for (PetscInt q = 0; q < xNq; ++q) { 1905 // Transform quadrature points from ref space to real space 1906 CoordinatesRefToReal(dim, dim, xi0, xv0, xJ, &xq[q * dim], xqr); 1907 // Get probability density at quad point 1908 // No need to scale xqr since PDF will be periodic 1909 PetscCall((*func)(xqr, param, &xden)); 1910 xw += xden * (xwq[q] * xdetJ); 1911 } 1912 xwtot += xw; 1913 if (debug) { 1914 IS globalOrdering; 1915 const PetscInt *ordering; 1916 1917 PetscCall(DMPlexGetCellNumbering(xdm, &globalOrdering)); 1918 PetscCall(ISGetIndices(globalOrdering, &ordering)); 1919 PetscCall(PetscSynchronizedPrintf(comm, "c:%" PetscInt_FMT " [x_a,x_b] = %1.15f,%1.15f -> cell weight = %1.15f\n", ordering[c], (double)PetscRealPart(xcoords[0]), (double)PetscRealPart(xcoords[0 + dim]), (double)xw)); 1920 PetscCall(ISRestoreIndices(globalOrdering, &ordering)); 1921 } 1922 // Set weights to be Gaussian in velocity cells 1923 for (PetscInt vc = vcStart; vc < vcEnd; ++vc) { 1924 const PetscInt p = pidx[vc * Ns + 0]; 1925 PetscReal vw = 0.; 1926 PetscInt vNc; 1927 const PetscScalar *varray; 1928 PetscScalar *vcoords = NULL; 1929 PetscBool visDG; 1930 1931 PetscCall(DMPlexGetCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords)); 1932 // TODO: Fix 2 stream Ask Joe 1933 // Two stream function from 1/2pi v^2 e^(-v^2/2) 1934 // vw = 1. / (PetscSqrtReal(2 * PETSC_PI)) * (((coords_v[0] * PetscExpReal(-PetscSqr(coords_v[0]) / 2.)) - (coords_v[1] * PetscExpReal(-PetscSqr(coords_v[1]) / 2.)))) - 0.5 * PetscErfReal(coords_v[0] / PetscSqrtReal(2.)) + 0.5 * (PetscErfReal(coords_v[1] / PetscSqrtReal(2.))); 1935 vw = 0.5 * (PetscErfReal(vcoords[1] / PetscSqrtReal(2.)) - PetscErfReal(vcoords[0] / PetscSqrtReal(2.))); 1936 1937 weight[p] = totalWeight * vw * xw; 1938 pwtot += weight[p]; 1939 PetscCheck(weight[p] <= 10., PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Particle %" PetscInt_FMT " weight exceeded 1: %g, %g, %g", p, xw, vw, totalWeight); 1940 PetscCall(DMPlexRestoreCellCoordinates(vdm, vc, &visDG, &vNc, &varray, &vcoords)); 1941 if (debug > 1) PetscCall(PetscPrintf(comm, "particle %" PetscInt_FMT ": %g, vw: %g xw: %g\n", p, weight[p], vw, xw)); 1942 } 1943 PetscCall(DMPlexRestoreCellCoordinates(xdm, c, &xisDG, &xNc, &xarray, &xcoords)); 1944 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Npc, &pidx)); 1945 } 1946 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight)); 1947 PetscCall(DMSwarmSortRestoreAccess(sw)); 1948 PetscCall(PetscQuadratureDestroy(&xquad)); 1949 1950 if (debug) { 1951 PetscReal wtot[2] = {pwtot, xwtot}, gwtot[2]; 1952 1953 PetscCall(PetscSynchronizedFlush(comm, NULL)); 1954 PetscCallMPI(MPIU_Allreduce(wtot, gwtot, 2, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD)); 1955 PetscCall(PetscPrintf(comm, "particle weight sum = %1.10f cell weight sum = %1.10f\n", (double)gwtot[0], (double)gwtot[1])); 1956 } 1957 PetscFunctionReturn(PETSC_SUCCESS); 1958 } 1959 1960 static PetscErrorCode InitializeParticles_PerturbedWeights(DM sw, AppCtx *user) 1961 { 1962 PetscReal scale[2] = {user->cosine_coefficients[0], user->cosine_coefficients[1]}; 1963 PetscInt dim; 1964 1965 PetscFunctionBegin; 1966 PetscCall(DMGetDimension(sw, &dim)); 1967 PetscCall(InitializeParticles_Centroid(sw)); 1968 PetscCall(InitializeWeights(sw, user->totalWeight, dim == 1 ? PetscPDFCosine1D : PetscPDFCosine2D, scale)); 1969 PetscFunctionReturn(PETSC_SUCCESS); 1970 } 1971 1972 static PetscErrorCode InitializeConstants(DM sw, AppCtx *user) 1973 { 1974 DM dm; 1975 PetscInt *species; 1976 PetscReal *weight, totalCharge = 0., totalWeight = 0., gmin[3], gmax[3], global_charge, global_weight; 1977 PetscInt Np, dim; 1978 1979 PetscFunctionBegin; 1980 PetscCall(DMSwarmGetCellDM(sw, &dm)); 1981 PetscCall(DMGetDimension(sw, &dim)); 1982 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 1983 PetscCall(DMGetBoundingBox(dm, gmin, gmax)); 1984 PetscCall(DMSwarmGetField(sw, "w_q", NULL, NULL, (void **)&weight)); 1985 PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species)); 1986 for (PetscInt p = 0; p < Np; ++p) { 1987 totalWeight += weight[p]; 1988 totalCharge += user->charges[species[p]] * weight[p]; 1989 } 1990 PetscCall(DMSwarmRestoreField(sw, "w_q", NULL, NULL, (void **)&weight)); 1991 PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species)); 1992 { 1993 Parameter *param; 1994 PetscReal Area; 1995 1996 PetscCall(PetscBagGetData(user->bag, (void **)¶m)); 1997 switch (dim) { 1998 case 1: 1999 Area = (gmax[0] - gmin[0]); 2000 break; 2001 case 2: 2002 Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]); 2003 break; 2004 case 3: 2005 Area = (gmax[0] - gmin[0]) * (gmax[1] - gmin[1]) * (gmax[2] - gmin[2]); 2006 break; 2007 default: 2008 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "Dimension %" PetscInt_FMT " not supported", dim); 2009 } 2010 PetscCallMPI(MPIU_Allreduce(&totalWeight, &global_weight, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD)); 2011 PetscCallMPI(MPIU_Allreduce(&totalCharge, &global_charge, 1, MPIU_REAL, MPIU_SUM, PETSC_COMM_WORLD)); 2012 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "dim = %" PetscInt_FMT "\ttotalWeight = %f, user->charges[species[0]] = %f\ttotalCharge = %f, Total Area = %f\n", dim, (double)global_weight, (double)user->charges[0], (double)global_charge, (double)Area)); 2013 param->sigma = PetscAbsReal(global_charge / (Area)); 2014 2015 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "sigma: %g\n", (double)param->sigma)); 2016 PetscCall(PetscPrintf(PETSC_COMM_WORLD, "(x0,v0,t0,m0,q0,phi0): (%e, %e, %e, %e, %e, %e) - (P, V) = (%e, %e)\n", (double)param->x0, (double)param->v0, (double)param->t0, (double)param->m0, (double)param->q0, (double)param->phi0, (double)param->poissonNumber, 2017 (double)param->vlasovNumber)); 2018 } 2019 /* Setup Constants */ 2020 { 2021 PetscDS ds; 2022 Parameter *param; 2023 PetscCall(PetscBagGetData(user->bag, (void **)¶m)); 2024 PetscScalar constants[NUM_CONSTANTS]; 2025 constants[SIGMA] = param->sigma; 2026 constants[V0] = param->v0; 2027 constants[T0] = param->t0; 2028 constants[X0] = param->x0; 2029 constants[M0] = param->m0; 2030 constants[Q0] = param->q0; 2031 constants[PHI0] = param->phi0; 2032 constants[POISSON] = param->poissonNumber; 2033 constants[VLASOV] = param->vlasovNumber; 2034 PetscCall(DMGetDS(dm, &ds)); 2035 PetscCall(PetscDSSetConstants(ds, NUM_CONSTANTS, constants)); 2036 } 2037 PetscFunctionReturn(PETSC_SUCCESS); 2038 } 2039 2040 static PetscErrorCode CreateSwarm(DM dm, AppCtx *user, DM *sw) 2041 { 2042 DMSwarmCellDM celldm; 2043 DM vdm; 2044 PetscReal v0[2] = {1., 0.}; 2045 PetscInt dim; 2046 2047 PetscFunctionBeginUser; 2048 PetscCall(DMGetDimension(dm, &dim)); 2049 PetscCall(DMCreate(PetscObjectComm((PetscObject)dm), sw)); 2050 PetscCall(DMSetType(*sw, DMSWARM)); 2051 PetscCall(DMSetDimension(*sw, dim)); 2052 PetscCall(DMSwarmSetType(*sw, DMSWARM_PIC)); 2053 PetscCall(DMSetApplicationContext(*sw, user)); 2054 2055 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "w_q", 1, PETSC_SCALAR)); 2056 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "velocity", dim, PETSC_REAL)); 2057 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "species", 1, PETSC_INT)); 2058 PetscCall(DMSwarmRegisterPetscDatatypeField(*sw, "E_field", dim, PETSC_REAL)); 2059 2060 const char *fieldnames[2] = {DMSwarmPICField_coor, "velocity"}; 2061 2062 PetscCall(DMSwarmCellDMCreate(dm, 2, fieldnames, 1, fieldnames, &celldm)); 2063 PetscCall(DMSwarmAddCellDM(*sw, celldm)); 2064 PetscCall(DMSwarmCellDMDestroy(&celldm)); 2065 2066 const char *vfieldnames[2] = {"w_q"}; 2067 2068 PetscCall(CreateVelocityDM(*sw, &vdm)); 2069 PetscCall(DMSwarmCellDMCreate(vdm, 1, vfieldnames, 1, &fieldnames[1], &celldm)); 2070 PetscCall(DMSwarmAddCellDM(*sw, celldm)); 2071 PetscCall(DMSwarmCellDMDestroy(&celldm)); 2072 PetscCall(DMDestroy(&vdm)); 2073 2074 DM mdm; 2075 2076 PetscCall(DMClone(dm, &mdm)); 2077 PetscCall(PetscObjectSetName((PetscObject)mdm, "moments")); 2078 PetscCall(DMCopyDisc(dm, mdm)); 2079 PetscCall(DMSwarmCellDMCreate(mdm, 1, vfieldnames, 1, fieldnames, &celldm)); 2080 PetscCall(DMDestroy(&mdm)); 2081 PetscCall(DMSwarmAddCellDM(*sw, celldm)); 2082 PetscCall(DMSwarmCellDMDestroy(&celldm)); 2083 2084 DM mfdm; 2085 2086 PetscCall(DMClone(dm, &mfdm)); 2087 PetscCall(PetscObjectSetName((PetscObject)mfdm, "moment fields")); 2088 PetscCall(DMCopyDisc(dm, mfdm)); 2089 PetscCall(DMSwarmCellDMCreate(mfdm, 1, &fieldnames[1], 1, fieldnames, &celldm)); 2090 PetscCall(DMDestroy(&mfdm)); 2091 PetscCall(DMSwarmAddCellDM(*sw, celldm)); 2092 PetscCall(DMSwarmCellDMDestroy(&celldm)); 2093 2094 PetscCall(DMSetFromOptions(*sw)); 2095 PetscCall(DMSetUp(*sw)); 2096 2097 PetscCall(DMSwarmSetCellDMActive(*sw, "space")); 2098 user->swarm = *sw; 2099 // TODO: This is redundant init since it is done in InitializeSolveAndSwarm, however DMSetUp() requires the local size be set 2100 if (user->perturbed_weights) { 2101 PetscCall(InitializeParticles_PerturbedWeights(*sw, user)); 2102 } else { 2103 PetscCall(DMSwarmComputeLocalSizeFromOptions(*sw)); 2104 PetscCall(DMSwarmInitializeCoordinates(*sw)); 2105 PetscCall(DMSwarmInitializeVelocitiesFromOptions(*sw, v0)); 2106 } 2107 PetscCall(PetscObjectSetName((PetscObject)*sw, "Particles")); 2108 PetscCall(DMViewFromOptions(*sw, NULL, "-sw_view")); 2109 PetscFunctionReturn(PETSC_SUCCESS); 2110 } 2111 2112 static PetscErrorCode ComputeFieldAtParticles_Coulomb(SNES snes, DM sw, PetscReal E[]) 2113 { 2114 AppCtx *user; 2115 PetscReal *coords; 2116 PetscInt *species, dim, Np, Ns; 2117 PetscMPIInt size; 2118 2119 PetscFunctionBegin; 2120 PetscCallMPI(MPI_Comm_size(PetscObjectComm((PetscObject)snes), &size)); 2121 PetscCheck(size == 1, PetscObjectComm((PetscObject)snes), PETSC_ERR_SUP, "Coulomb code only works in serial"); 2122 PetscCall(DMGetDimension(sw, &dim)); 2123 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 2124 PetscCall(DMSwarmGetNumSpecies(sw, &Ns)); 2125 PetscCall(DMGetApplicationContext(sw, (void *)&user)); 2126 2127 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 2128 PetscCall(DMSwarmGetField(sw, "species", NULL, NULL, (void **)&species)); 2129 for (PetscInt p = 0; p < Np; ++p) { 2130 PetscReal *pcoord = &coords[p * dim]; 2131 PetscReal pE[3] = {0., 0., 0.}; 2132 2133 /* Calculate field at particle p due to particle q */ 2134 for (PetscInt q = 0; q < Np; ++q) { 2135 PetscReal *qcoord = &coords[q * dim]; 2136 PetscReal rpq[3], r, r3, q_q; 2137 2138 if (p == q) continue; 2139 q_q = user->charges[species[q]] * 1.; 2140 for (PetscInt d = 0; d < dim; ++d) rpq[d] = pcoord[d] - qcoord[d]; 2141 r = DMPlex_NormD_Internal(dim, rpq); 2142 if (r < PETSC_SQRT_MACHINE_EPSILON) continue; 2143 r3 = PetscPowRealInt(r, 3); 2144 for (PetscInt d = 0; d < dim; ++d) pE[d] += q_q * rpq[d] / r3; 2145 } 2146 for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = pE[d]; 2147 } 2148 PetscCall(DMSwarmRestoreField(sw, "species", NULL, NULL, (void **)&species)); 2149 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 2150 PetscFunctionReturn(PETSC_SUCCESS); 2151 } 2152 2153 static PetscErrorCode ComputeFieldAtParticles_Primal(SNES snes, DM sw, Mat M_p, PetscReal E[]) 2154 { 2155 DM dm; 2156 AppCtx *user; 2157 PetscDS ds; 2158 PetscFE fe; 2159 KSP ksp; 2160 Vec rhoRhs; // Weak charge density, \int phi_i rho 2161 Vec rho; // Charge density, M^{-1} rhoRhs 2162 Vec phi, locPhi; // Potential 2163 Vec f; // Particle weights 2164 PetscReal *coords; 2165 PetscInt dim, cStart, cEnd, Np; 2166 2167 PetscFunctionBegin; 2168 PetscCall(DMGetApplicationContext(sw, (void *)&user)); 2169 PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0)); 2170 PetscCall(DMGetDimension(sw, &dim)); 2171 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 2172 2173 PetscCall(SNESGetDM(snes, &dm)); 2174 PetscCall(DMGetGlobalVector(dm, &rhoRhs)); 2175 PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density")); 2176 PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho)); 2177 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f)); 2178 PetscCall(PetscObjectSetName((PetscObject)f, "particle weight")); 2179 2180 PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view")); 2181 PetscCall(MatViewFromOptions(user->M, NULL, "-m_view")); 2182 PetscCall(VecViewFromOptions(f, NULL, "-weights_view")); 2183 2184 PetscCall(MatMultTranspose(M_p, f, rhoRhs)); 2185 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f)); 2186 2187 PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp)); 2188 PetscCall(KSPSetOptionsPrefix(ksp, "em_proj_")); 2189 PetscCall(KSPSetOperators(ksp, user->M, user->M)); 2190 PetscCall(KSPSetFromOptions(ksp)); 2191 PetscCall(KSPSolve(ksp, rhoRhs, rho)); 2192 2193 PetscCall(VecScale(rhoRhs, -1.0)); 2194 2195 PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view")); 2196 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho)); 2197 PetscCall(KSPDestroy(&ksp)); 2198 2199 PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi)); 2200 PetscCall(VecSet(phi, 0.0)); 2201 PetscCall(SNESSolve(snes, rhoRhs, phi)); 2202 PetscCall(DMRestoreGlobalVector(dm, &rhoRhs)); 2203 PetscCall(VecViewFromOptions(phi, NULL, "-phi_view")); 2204 2205 PetscCall(DMGetLocalVector(dm, &locPhi)); 2206 PetscCall(DMGlobalToLocalBegin(dm, phi, INSERT_VALUES, locPhi)); 2207 PetscCall(DMGlobalToLocalEnd(dm, phi, INSERT_VALUES, locPhi)); 2208 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi)); 2209 PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0)); 2210 2211 PetscCall(DMGetDS(dm, &ds)); 2212 PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe)); 2213 PetscCall(DMSwarmSortGetAccess(sw)); 2214 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd)); 2215 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 2216 2217 PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0)); 2218 PetscTabulation tab; 2219 PetscReal *pcoord, *refcoord; 2220 PetscFEGeom *chunkgeom = NULL; 2221 PetscInt maxNcp = 0; 2222 2223 for (PetscInt c = cStart; c < cEnd; ++c) { 2224 PetscInt Ncp; 2225 2226 PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp)); 2227 maxNcp = PetscMax(maxNcp, Ncp); 2228 } 2229 PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord)); 2230 PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord)); 2231 // This can raise an FP_INEXACT in the dgemm inside 2232 PetscCall(PetscFPTrapPush(PETSC_FP_TRAP_OFF)); 2233 PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab)); 2234 PetscCall(PetscFPTrapPop()); 2235 for (PetscInt c = cStart; c < cEnd; ++c) { 2236 PetscScalar *clPhi = NULL; 2237 PetscInt *points; 2238 PetscInt Ncp; 2239 2240 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points)); 2241 for (PetscInt cp = 0; cp < Ncp; ++cp) 2242 for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d]; 2243 { 2244 PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom)); 2245 for (PetscInt i = 0; i < Ncp; ++i) { 2246 const PetscReal x0[3] = {-1., -1., -1.}; 2247 CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]); 2248 } 2249 } 2250 PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab)); 2251 PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi)); 2252 for (PetscInt cp = 0; cp < Ncp; ++cp) { 2253 const PetscReal *basisDer = tab->T[1]; 2254 const PetscInt p = points[cp]; 2255 2256 for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.; 2257 PetscCall(PetscFEFreeInterpolateGradient_Static(fe, basisDer, clPhi, dim, chunkgeom->invJ, NULL, cp, &E[p * dim])); 2258 for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0; 2259 } 2260 PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi)); 2261 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points)); 2262 } 2263 PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord)); 2264 PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord)); 2265 PetscCall(PetscTabulationDestroy(&tab)); 2266 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 2267 PetscCall(DMSwarmSortRestoreAccess(sw)); 2268 PetscCall(DMRestoreLocalVector(dm, &locPhi)); 2269 PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom)); 2270 PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0)); 2271 PetscFunctionReturn(PETSC_SUCCESS); 2272 } 2273 2274 static PetscErrorCode ComputeFieldAtParticles_Mixed(SNES snes, DM sw, Mat M_p, PetscReal E[]) 2275 { 2276 DM dm; 2277 AppCtx *user; 2278 PetscDS ds; 2279 PetscFE fe; 2280 KSP ksp; 2281 Vec rhoRhs, rhoRhsFull; // Weak charge density, \int phi_i rho, and embedding in mixed problem 2282 Vec rho; // Charge density, M^{-1} rhoRhs 2283 Vec phi, locPhi, phiFull; // Potential and embedding in mixed problem 2284 Vec f; // Particle weights 2285 PetscReal *coords; 2286 PetscInt dim, cStart, cEnd, Np; 2287 2288 PetscFunctionBegin; 2289 PetscCall(DMGetApplicationContext(sw, &user)); 2290 PetscCall(PetscLogEventBegin(user->ESolveEvent, snes, sw, 0, 0)); 2291 PetscCall(DMGetDimension(sw, &dim)); 2292 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 2293 2294 PetscCall(SNESGetDM(snes, &dm)); 2295 PetscCall(DMGetGlobalVector(user->dmPot, &rhoRhs)); 2296 PetscCall(PetscObjectSetName((PetscObject)rhoRhs, "Weak charge density")); 2297 PetscCall(DMGetGlobalVector(dm, &rhoRhsFull)); 2298 PetscCall(PetscObjectSetName((PetscObject)rhoRhsFull, "Weak charge density")); 2299 PetscCall(DMGetNamedGlobalVector(user->dmPot, "rho", &rho)); 2300 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "w_q", &f)); 2301 PetscCall(PetscObjectSetName((PetscObject)f, "particle weight")); 2302 2303 PetscCall(MatViewFromOptions(M_p, NULL, "-mp_view")); 2304 PetscCall(MatViewFromOptions(user->M, NULL, "-m_view")); 2305 PetscCall(VecViewFromOptions(f, NULL, "-weights_view")); 2306 2307 PetscCall(MatMultTranspose(M_p, f, rhoRhs)); 2308 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "w_q", &f)); 2309 2310 PetscCall(KSPCreate(PetscObjectComm((PetscObject)dm), &ksp)); 2311 PetscCall(KSPSetOptionsPrefix(ksp, "em_proj")); 2312 PetscCall(KSPSetOperators(ksp, user->M, user->M)); 2313 PetscCall(KSPSetFromOptions(ksp)); 2314 PetscCall(KSPSolve(ksp, rhoRhs, rho)); 2315 2316 PetscCall(VecISCopy(rhoRhsFull, user->isPot, SCATTER_FORWARD, rhoRhs)); 2317 //PetscCall(VecScale(rhoRhsFull, -1.0)); 2318 2319 PetscCall(VecViewFromOptions(rhoRhs, NULL, "-rho_view")); 2320 PetscCall(VecViewFromOptions(rhoRhsFull, NULL, "-rho_full_view")); 2321 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "rho", &rho)); 2322 PetscCall(DMRestoreGlobalVector(user->dmPot, &rhoRhs)); 2323 PetscCall(KSPDestroy(&ksp)); 2324 2325 PetscCall(DMGetGlobalVector(dm, &phiFull)); 2326 PetscCall(DMGetNamedGlobalVector(user->dmPot, "phi", &phi)); 2327 PetscCall(VecSet(phiFull, 0.0)); 2328 PetscCall(SNESSolve(snes, rhoRhsFull, phiFull)); 2329 PetscCall(DMRestoreGlobalVector(dm, &rhoRhsFull)); 2330 PetscCall(VecViewFromOptions(phi, NULL, "-phi_view")); 2331 2332 PetscCall(VecISCopy(phiFull, user->isPot, SCATTER_REVERSE, phi)); 2333 PetscCall(DMRestoreNamedGlobalVector(user->dmPot, "phi", &phi)); 2334 2335 PetscCall(DMGetLocalVector(dm, &locPhi)); 2336 PetscCall(DMGlobalToLocalBegin(dm, phiFull, INSERT_VALUES, locPhi)); 2337 PetscCall(DMGlobalToLocalEnd(dm, phiFull, INSERT_VALUES, locPhi)); 2338 PetscCall(DMRestoreGlobalVector(dm, &phiFull)); 2339 PetscCall(PetscLogEventEnd(user->ESolveEvent, snes, sw, 0, 0)); 2340 2341 PetscCall(DMGetDS(dm, &ds)); 2342 PetscCall(PetscDSGetDiscretization(ds, 0, (PetscObject *)&fe)); 2343 PetscCall(DMSwarmSortGetAccess(sw)); 2344 PetscCall(DMPlexGetHeightStratum(dm, 0, &cStart, &cEnd)); 2345 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 2346 2347 PetscCall(PetscLogEventBegin(user->ETabEvent, snes, sw, 0, 0)); 2348 PetscTabulation tab; 2349 PetscReal *pcoord, *refcoord; 2350 PetscFEGeom *chunkgeom = NULL; 2351 PetscInt maxNcp = 0; 2352 2353 for (PetscInt c = cStart; c < cEnd; ++c) { 2354 PetscInt Ncp; 2355 2356 PetscCall(DMSwarmSortGetNumberOfPointsPerCell(sw, c, &Ncp)); 2357 maxNcp = PetscMax(maxNcp, Ncp); 2358 } 2359 PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord)); 2360 PetscCall(DMGetWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord)); 2361 PetscCall(PetscFECreateTabulation(fe, 1, maxNcp, refcoord, 1, &tab)); 2362 for (PetscInt c = cStart; c < cEnd; ++c) { 2363 PetscScalar *clPhi = NULL; 2364 PetscInt *points; 2365 PetscInt Ncp; 2366 2367 PetscCall(DMSwarmSortGetPointsPerCell(sw, c, &Ncp, &points)); 2368 for (PetscInt cp = 0; cp < Ncp; ++cp) 2369 for (PetscInt d = 0; d < dim; ++d) pcoord[cp * dim + d] = coords[points[cp] * dim + d]; 2370 { 2371 PetscCall(PetscFEGeomGetChunk(user->fegeom, c - cStart, c - cStart + 1, &chunkgeom)); 2372 for (PetscInt i = 0; i < Ncp; ++i) { 2373 const PetscReal x0[3] = {-1., -1., -1.}; 2374 CoordinatesRealToRef(dim, dim, x0, chunkgeom->v, chunkgeom->invJ, &pcoord[dim * i], &refcoord[dim * i]); 2375 } 2376 } 2377 PetscCall(PetscFEComputeTabulation(fe, Ncp, refcoord, 1, tab)); 2378 PetscCall(DMPlexVecGetClosure(dm, NULL, locPhi, c, NULL, &clPhi)); 2379 for (PetscInt cp = 0; cp < Ncp; ++cp) { 2380 const PetscInt p = points[cp]; 2381 2382 for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] = 0.; 2383 PetscCall(PetscFEInterpolateAtPoints_Static(fe, tab, clPhi, chunkgeom, cp, &E[p * dim])); 2384 PetscCall(PetscFEPushforward(fe, chunkgeom, 1, &E[p * dim])); 2385 for (PetscInt d = 0; d < dim; ++d) E[p * dim + d] *= -1.0; 2386 } 2387 PetscCall(DMPlexVecRestoreClosure(dm, NULL, locPhi, c, NULL, &clPhi)); 2388 PetscCall(DMSwarmSortRestorePointsPerCell(sw, c, &Ncp, &points)); 2389 } 2390 PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &pcoord)); 2391 PetscCall(DMRestoreWorkArray(dm, maxNcp * dim, MPIU_REAL, &refcoord)); 2392 PetscCall(PetscTabulationDestroy(&tab)); 2393 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 2394 PetscCall(DMSwarmSortRestoreAccess(sw)); 2395 PetscCall(DMRestoreLocalVector(dm, &locPhi)); 2396 PetscCall(PetscFEGeomRestoreChunk(user->fegeom, 0, 1, &chunkgeom)); 2397 PetscCall(PetscLogEventEnd(user->ETabEvent, snes, sw, 0, 0)); 2398 PetscFunctionReturn(PETSC_SUCCESS); 2399 } 2400 2401 static PetscErrorCode ComputeFieldAtParticles(SNES snes, DM sw) 2402 { 2403 AppCtx *user; 2404 Mat M_p; 2405 PetscReal *E; 2406 PetscInt dim, Np; 2407 2408 PetscFunctionBegin; 2409 PetscValidHeaderSpecific(snes, SNES_CLASSID, 1); 2410 PetscValidHeaderSpecific(sw, DM_CLASSID, 2); 2411 PetscCall(DMGetDimension(sw, &dim)); 2412 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 2413 PetscCall(DMGetApplicationContext(sw, &user)); 2414 2415 PetscCall(DMSwarmSetCellDMActive(sw, "moments")); 2416 // TODO: Could share sort context with space cellDM 2417 PetscCall(DMSwarmMigrate(sw, PETSC_FALSE)); 2418 PetscCall(DMCreateMassMatrix(sw, user->dmPot, &M_p)); 2419 PetscCall(DMSwarmSetCellDMActive(sw, "space")); 2420 2421 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E)); 2422 PetscCall(PetscArrayzero(E, Np * dim)); 2423 user->validE = PETSC_TRUE; 2424 2425 switch (user->em) { 2426 case EM_COULOMB: 2427 PetscCall(ComputeFieldAtParticles_Coulomb(snes, sw, E)); 2428 break; 2429 case EM_PRIMAL: 2430 PetscCall(ComputeFieldAtParticles_Primal(snes, sw, M_p, E)); 2431 break; 2432 case EM_MIXED: 2433 PetscCall(ComputeFieldAtParticles_Mixed(snes, sw, M_p, E)); 2434 break; 2435 case EM_NONE: 2436 break; 2437 default: 2438 SETERRQ(PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "No solver for electrostatic model %s", EMTypes[user->em]); 2439 } 2440 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E)); 2441 PetscCall(MatDestroy(&M_p)); 2442 PetscFunctionReturn(PETSC_SUCCESS); 2443 } 2444 2445 static PetscErrorCode RHSFunction(TS ts, PetscReal t, Vec U, Vec G, void *ctx) 2446 { 2447 DM sw; 2448 SNES snes = ((AppCtx *)ctx)->snes; 2449 const PetscScalar *u; 2450 PetscScalar *g; 2451 PetscReal *E, m_p = 1., q_p = -1.; 2452 PetscInt dim, d, Np, p; 2453 2454 PetscFunctionBeginUser; 2455 PetscCall(TSGetDM(ts, &sw)); 2456 PetscCall(ComputeFieldAtParticles(snes, sw)); 2457 2458 PetscCall(DMGetDimension(sw, &dim)); 2459 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 2460 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E)); 2461 PetscCall(VecGetArrayRead(U, &u)); 2462 PetscCall(VecGetArray(G, &g)); 2463 Np /= 2 * dim; 2464 for (p = 0; p < Np; ++p) { 2465 for (d = 0; d < dim; ++d) { 2466 g[(p * 2 + 0) * dim + d] = u[(p * 2 + 1) * dim + d]; 2467 g[(p * 2 + 1) * dim + d] = q_p * E[p * dim + d] / m_p; 2468 } 2469 } 2470 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E)); 2471 PetscCall(VecRestoreArrayRead(U, &u)); 2472 PetscCall(VecRestoreArray(G, &g)); 2473 PetscFunctionReturn(PETSC_SUCCESS); 2474 } 2475 2476 /* J_{ij} = dF_i/dx_j 2477 J_p = ( 0 1) 2478 (-w^2 0) 2479 TODO Now there is another term with w^2 from the electric field. I think we will need to invert the operator. 2480 Perhaps we can approximate the Jacobian using only the cellwise P-P gradient from Coulomb 2481 */ 2482 static PetscErrorCode RHSJacobian(TS ts, PetscReal t, Vec U, Mat J, Mat P, void *ctx) 2483 { 2484 DM sw; 2485 const PetscReal *coords, *vel; 2486 PetscInt dim, d, Np, p, rStart; 2487 2488 PetscFunctionBeginUser; 2489 PetscCall(TSGetDM(ts, &sw)); 2490 PetscCall(DMGetDimension(sw, &dim)); 2491 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 2492 PetscCall(MatGetOwnershipRange(J, &rStart, NULL)); 2493 PetscCall(DMSwarmGetField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 2494 PetscCall(DMSwarmGetField(sw, "velocity", NULL, NULL, (void **)&vel)); 2495 Np /= 2 * dim; 2496 for (p = 0; p < Np; ++p) { 2497 const PetscReal x0 = coords[p * dim + 0]; 2498 const PetscReal vy0 = vel[p * dim + 1]; 2499 const PetscReal omega = vy0 / x0; 2500 PetscScalar vals[4] = {0., 1., -PetscSqr(omega), 0.}; 2501 2502 for (d = 0; d < dim; ++d) { 2503 const PetscInt rows[2] = {(p * 2 + 0) * dim + d + rStart, (p * 2 + 1) * dim + d + rStart}; 2504 PetscCall(MatSetValues(J, 2, rows, 2, rows, vals, INSERT_VALUES)); 2505 } 2506 } 2507 PetscCall(DMSwarmRestoreField(sw, DMSwarmPICField_coor, NULL, NULL, (void **)&coords)); 2508 PetscCall(DMSwarmRestoreField(sw, "velocity", NULL, NULL, (void **)&vel)); 2509 PetscCall(MatAssemblyBegin(J, MAT_FINAL_ASSEMBLY)); 2510 PetscCall(MatAssemblyEnd(J, MAT_FINAL_ASSEMBLY)); 2511 PetscFunctionReturn(PETSC_SUCCESS); 2512 } 2513 2514 static PetscErrorCode RHSFunctionX(TS ts, PetscReal t, Vec V, Vec Xres, void *ctx) 2515 { 2516 AppCtx *user = (AppCtx *)ctx; 2517 DM sw; 2518 const PetscScalar *v; 2519 PetscScalar *xres; 2520 PetscInt Np, p, d, dim; 2521 2522 PetscFunctionBeginUser; 2523 PetscCall(PetscLogEventBegin(user->RhsXEvent, ts, 0, 0, 0)); 2524 PetscCall(TSGetDM(ts, &sw)); 2525 PetscCall(DMGetDimension(sw, &dim)); 2526 PetscCall(VecGetLocalSize(Xres, &Np)); 2527 PetscCall(VecGetArrayRead(V, &v)); 2528 PetscCall(VecGetArray(Xres, &xres)); 2529 Np /= dim; 2530 for (p = 0; p < Np; ++p) { 2531 for (d = 0; d < dim; ++d) xres[p * dim + d] = v[p * dim + d]; 2532 } 2533 PetscCall(VecRestoreArrayRead(V, &v)); 2534 PetscCall(VecRestoreArray(Xres, &xres)); 2535 PetscCall(PetscLogEventEnd(user->RhsXEvent, ts, 0, 0, 0)); 2536 PetscFunctionReturn(PETSC_SUCCESS); 2537 } 2538 2539 static PetscErrorCode RHSFunctionV(TS ts, PetscReal t, Vec X, Vec Vres, void *ctx) 2540 { 2541 DM sw; 2542 AppCtx *user = (AppCtx *)ctx; 2543 SNES snes = ((AppCtx *)ctx)->snes; 2544 const PetscScalar *x; 2545 PetscScalar *vres; 2546 PetscReal *E, m_p, q_p; 2547 PetscInt Np, p, dim, d; 2548 Parameter *param; 2549 2550 PetscFunctionBeginUser; 2551 PetscCall(PetscLogEventBegin(user->RhsVEvent, ts, 0, 0, 0)); 2552 PetscCall(TSGetDM(ts, &sw)); 2553 PetscCall(ComputeFieldAtParticles(snes, sw)); 2554 2555 PetscCall(DMGetDimension(sw, &dim)); 2556 PetscCall(DMSwarmGetField(sw, "E_field", NULL, NULL, (void **)&E)); 2557 PetscCall(PetscBagGetData(user->bag, (void **)¶m)); 2558 m_p = user->masses[0] * param->m0; 2559 q_p = user->charges[0] * param->q0; 2560 PetscCall(VecGetLocalSize(Vres, &Np)); 2561 PetscCall(VecGetArrayRead(X, &x)); 2562 PetscCall(VecGetArray(Vres, &vres)); 2563 Np /= dim; 2564 for (p = 0; p < Np; ++p) { 2565 for (d = 0; d < dim; ++d) vres[p * dim + d] = q_p * E[p * dim + d] / m_p; 2566 } 2567 PetscCall(VecRestoreArrayRead(X, &x)); 2568 /* 2569 Synchronized, ordered output for parallel/sequential test cases. 2570 In the 1D (on the 2D mesh) case, every y component should be zero. 2571 */ 2572 if (user->checkVRes) { 2573 PetscBool pr = user->checkVRes > 1 ? PETSC_TRUE : PETSC_FALSE; 2574 PetscInt step; 2575 2576 PetscCall(TSGetStepNumber(ts, &step)); 2577 if (pr) PetscCall(PetscPrintf(PETSC_COMM_WORLD, "step: %" PetscInt_FMT "\n", step)); 2578 for (PetscInt p = 0; p < Np; ++p) { 2579 if (pr) PetscCall(PetscSynchronizedPrintf(PETSC_COMM_WORLD, "Residual: %.12g %.12g\n", (double)PetscRealPart(vres[p * dim + 0]), (double)PetscRealPart(vres[p * dim + 1]))); 2580 PetscCheck(PetscAbsScalar(vres[p * dim + 1]) < PETSC_SMALL, PETSC_COMM_SELF, PETSC_ERR_PLIB, "Y velocity should be 0., not %g", (double)PetscRealPart(vres[p * dim + 1])); 2581 } 2582 if (pr) PetscCall(PetscSynchronizedFlush(PETSC_COMM_WORLD, PETSC_STDOUT)); 2583 } 2584 PetscCall(VecRestoreArray(Vres, &vres)); 2585 PetscCall(DMSwarmRestoreField(sw, "E_field", NULL, NULL, (void **)&E)); 2586 PetscCall(PetscLogEventEnd(user->RhsVEvent, ts, 0, 0, 0)); 2587 PetscFunctionReturn(PETSC_SUCCESS); 2588 } 2589 2590 static PetscErrorCode CreateSolution(TS ts) 2591 { 2592 DM sw; 2593 Vec u; 2594 PetscInt dim, Np; 2595 2596 PetscFunctionBegin; 2597 PetscCall(TSGetDM(ts, &sw)); 2598 PetscCall(DMGetDimension(sw, &dim)); 2599 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 2600 PetscCall(VecCreate(PETSC_COMM_WORLD, &u)); 2601 PetscCall(VecSetBlockSize(u, dim)); 2602 PetscCall(VecSetSizes(u, 2 * Np * dim, PETSC_DECIDE)); 2603 PetscCall(VecSetUp(u)); 2604 PetscCall(TSSetSolution(ts, u)); 2605 PetscCall(VecDestroy(&u)); 2606 PetscFunctionReturn(PETSC_SUCCESS); 2607 } 2608 2609 static PetscErrorCode SetProblem(TS ts) 2610 { 2611 AppCtx *user; 2612 DM sw; 2613 2614 PetscFunctionBegin; 2615 PetscCall(TSGetDM(ts, &sw)); 2616 PetscCall(DMGetApplicationContext(sw, (void **)&user)); 2617 // Define unified system for (X, V) 2618 { 2619 Mat J; 2620 PetscInt dim, Np; 2621 2622 PetscCall(DMGetDimension(sw, &dim)); 2623 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 2624 PetscCall(MatCreate(PETSC_COMM_WORLD, &J)); 2625 PetscCall(MatSetSizes(J, 2 * Np * dim, 2 * Np * dim, PETSC_DECIDE, PETSC_DECIDE)); 2626 PetscCall(MatSetBlockSize(J, 2 * dim)); 2627 PetscCall(MatSetFromOptions(J)); 2628 PetscCall(MatSetUp(J)); 2629 PetscCall(TSSetRHSFunction(ts, NULL, RHSFunction, user)); 2630 PetscCall(TSSetRHSJacobian(ts, J, J, RHSJacobian, user)); 2631 PetscCall(MatDestroy(&J)); 2632 } 2633 /* Define split system for X and V */ 2634 { 2635 Vec u; 2636 IS isx, isv, istmp; 2637 const PetscInt *idx; 2638 PetscInt dim, Np, rstart; 2639 2640 PetscCall(TSGetSolution(ts, &u)); 2641 PetscCall(DMGetDimension(sw, &dim)); 2642 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 2643 PetscCall(VecGetOwnershipRange(u, &rstart, NULL)); 2644 PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 0, 2, &istmp)); 2645 PetscCall(ISGetIndices(istmp, &idx)); 2646 PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isx)); 2647 PetscCall(ISRestoreIndices(istmp, &idx)); 2648 PetscCall(ISDestroy(&istmp)); 2649 PetscCall(ISCreateStride(PETSC_COMM_WORLD, Np, (rstart / dim) + 1, 2, &istmp)); 2650 PetscCall(ISGetIndices(istmp, &idx)); 2651 PetscCall(ISCreateBlock(PETSC_COMM_WORLD, dim, Np, idx, PETSC_COPY_VALUES, &isv)); 2652 PetscCall(ISRestoreIndices(istmp, &idx)); 2653 PetscCall(ISDestroy(&istmp)); 2654 PetscCall(TSRHSSplitSetIS(ts, "position", isx)); 2655 PetscCall(TSRHSSplitSetIS(ts, "momentum", isv)); 2656 PetscCall(ISDestroy(&isx)); 2657 PetscCall(ISDestroy(&isv)); 2658 PetscCall(TSRHSSplitSetRHSFunction(ts, "position", NULL, RHSFunctionX, user)); 2659 PetscCall(TSRHSSplitSetRHSFunction(ts, "momentum", NULL, RHSFunctionV, user)); 2660 } 2661 PetscFunctionReturn(PETSC_SUCCESS); 2662 } 2663 2664 static PetscErrorCode DMSwarmTSRedistribute(TS ts) 2665 { 2666 DM sw; 2667 Vec u; 2668 PetscReal t, maxt, dt; 2669 PetscInt n, maxn; 2670 2671 PetscFunctionBegin; 2672 PetscCall(TSGetDM(ts, &sw)); 2673 PetscCall(TSGetTime(ts, &t)); 2674 PetscCall(TSGetMaxTime(ts, &maxt)); 2675 PetscCall(TSGetTimeStep(ts, &dt)); 2676 PetscCall(TSGetStepNumber(ts, &n)); 2677 PetscCall(TSGetMaxSteps(ts, &maxn)); 2678 2679 PetscCall(TSReset(ts)); 2680 PetscCall(TSSetDM(ts, sw)); 2681 PetscCall(TSSetFromOptions(ts)); 2682 PetscCall(TSSetTime(ts, t)); 2683 PetscCall(TSSetMaxTime(ts, maxt)); 2684 PetscCall(TSSetTimeStep(ts, dt)); 2685 PetscCall(TSSetStepNumber(ts, n)); 2686 PetscCall(TSSetMaxSteps(ts, maxn)); 2687 2688 PetscCall(CreateSolution(ts)); 2689 PetscCall(SetProblem(ts)); 2690 PetscCall(TSGetSolution(ts, &u)); 2691 PetscFunctionReturn(PETSC_SUCCESS); 2692 } 2693 2694 PetscErrorCode line(PetscInt dim, PetscReal time, const PetscReal dummy[], PetscInt p, PetscScalar x[], void *ctx) 2695 { 2696 DM sw, cdm; 2697 PetscInt Np; 2698 PetscReal low[2], high[2]; 2699 AppCtx *user = (AppCtx *)ctx; 2700 2701 sw = user->swarm; 2702 PetscCall(DMSwarmGetCellDM(sw, &cdm)); 2703 // Get the bounding box so we can equally space the particles 2704 PetscCall(DMGetLocalBoundingBox(cdm, low, high)); 2705 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 2706 // shift it by h/2 so nothing is initialized directly on a boundary 2707 x[0] = ((high[0] - low[0]) / Np) * (p + 0.5); 2708 x[1] = 0.; 2709 return PETSC_SUCCESS; 2710 } 2711 2712 /* 2713 InitializeSolveAndSwarm - Set the solution values to the swarm coordinates and velocities, and also possibly set the initial values. 2714 2715 Input Parameters: 2716 + ts - The TS 2717 - useInitial - Flag to also set the initial conditions to the current coordinates and velocities and setup the problem 2718 2719 Output Parameters: 2720 . u - The initialized solution vector 2721 2722 Level: advanced 2723 2724 .seealso: InitializeSolve() 2725 */ 2726 static PetscErrorCode InitializeSolveAndSwarm(TS ts, PetscBool useInitial) 2727 { 2728 DM sw; 2729 Vec u, gc, gv; 2730 IS isx, isv; 2731 PetscInt dim; 2732 AppCtx *user; 2733 2734 PetscFunctionBeginUser; 2735 PetscCall(TSGetDM(ts, &sw)); 2736 PetscCall(DMGetApplicationContext(sw, &user)); 2737 PetscCall(DMGetDimension(sw, &dim)); 2738 if (useInitial) { 2739 PetscReal v0[2] = {1., 0.}; 2740 if (user->perturbed_weights) { 2741 PetscCall(InitializeParticles_PerturbedWeights(sw, user)); 2742 } else { 2743 PetscCall(DMSwarmComputeLocalSizeFromOptions(sw)); 2744 PetscCall(DMSwarmInitializeCoordinates(sw)); 2745 PetscCall(DMSwarmInitializeVelocitiesFromOptions(sw, v0)); 2746 } 2747 PetscCall(DMSwarmMigrate(sw, PETSC_TRUE)); 2748 PetscCall(DMSwarmTSRedistribute(ts)); 2749 } 2750 PetscCall(DMSetUp(sw)); 2751 PetscCall(TSGetSolution(ts, &u)); 2752 PetscCall(TSRHSSplitGetIS(ts, "position", &isx)); 2753 PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv)); 2754 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc)); 2755 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv)); 2756 PetscCall(VecISCopy(u, isx, SCATTER_FORWARD, gc)); 2757 PetscCall(VecISCopy(u, isv, SCATTER_FORWARD, gv)); 2758 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc)); 2759 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv)); 2760 PetscFunctionReturn(PETSC_SUCCESS); 2761 } 2762 2763 static PetscErrorCode InitializeSolve(TS ts, Vec u) 2764 { 2765 PetscFunctionBegin; 2766 PetscCall(TSSetSolution(ts, u)); 2767 PetscCall(InitializeSolveAndSwarm(ts, PETSC_TRUE)); 2768 PetscFunctionReturn(PETSC_SUCCESS); 2769 } 2770 2771 static PetscErrorCode MigrateParticles(TS ts) 2772 { 2773 DM sw, cdm; 2774 const PetscReal *L; 2775 AppCtx *ctx; 2776 2777 PetscFunctionBeginUser; 2778 PetscCall(TSGetDM(ts, &sw)); 2779 PetscCall(DMGetApplicationContext(sw, &ctx)); 2780 PetscCall(DMViewFromOptions(sw, NULL, "-migrate_view_pre")); 2781 { 2782 Vec u, gc, gv, position, momentum; 2783 IS isx, isv; 2784 PetscReal *pos, *mom; 2785 2786 PetscCall(TSGetSolution(ts, &u)); 2787 PetscCall(TSRHSSplitGetIS(ts, "position", &isx)); 2788 PetscCall(TSRHSSplitGetIS(ts, "momentum", &isv)); 2789 PetscCall(VecGetSubVector(u, isx, &position)); 2790 PetscCall(VecGetSubVector(u, isv, &momentum)); 2791 PetscCall(VecGetArray(position, &pos)); 2792 PetscCall(VecGetArray(momentum, &mom)); 2793 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc)); 2794 PetscCall(DMSwarmCreateGlobalVectorFromField(sw, "velocity", &gv)); 2795 PetscCall(VecISCopy(u, isx, SCATTER_REVERSE, gc)); 2796 PetscCall(VecISCopy(u, isv, SCATTER_REVERSE, gv)); 2797 2798 PetscCall(DMSwarmGetCellDM(sw, &cdm)); 2799 PetscCall(DMGetPeriodicity(cdm, NULL, NULL, &L)); 2800 PetscCheck(L, PetscObjectComm((PetscObject)cdm), PETSC_ERR_ARG_WRONG, "Mesh must be periodic"); 2801 if ((L[0] || L[1]) >= 0.) { 2802 PetscReal *x, *v, upper[3], lower[3]; 2803 PetscInt Np, dim; 2804 2805 PetscCall(DMSwarmGetLocalSize(sw, &Np)); 2806 PetscCall(DMGetDimension(cdm, &dim)); 2807 PetscCall(DMGetBoundingBox(cdm, lower, upper)); 2808 PetscCall(VecGetArray(gc, &x)); 2809 PetscCall(VecGetArray(gv, &v)); 2810 for (PetscInt p = 0; p < Np; ++p) { 2811 for (PetscInt d = 0; d < dim; ++d) { 2812 if (pos[p * dim + d] < lower[d]) { 2813 x[p * dim + d] = pos[p * dim + d] + (upper[d] - lower[d]); 2814 } else if (pos[p * dim + d] > upper[d]) { 2815 x[p * dim + d] = pos[p * dim + d] - (upper[d] - lower[d]); 2816 } else { 2817 x[p * dim + d] = pos[p * dim + d]; 2818 } 2819 PetscCheck(x[p * dim + d] >= lower[d] && x[p * dim + d] <= upper[d], PETSC_COMM_SELF, PETSC_ERR_ARG_WRONG, "p: %" PetscInt_FMT "x[%" PetscInt_FMT "] %g", p, d, (double)x[p * dim + d]); 2820 v[p * dim + d] = mom[p * dim + d]; 2821 } 2822 } 2823 PetscCall(VecRestoreArray(gc, &x)); 2824 PetscCall(VecRestoreArray(gv, &v)); 2825 } 2826 PetscCall(VecRestoreArray(position, &pos)); 2827 PetscCall(VecRestoreArray(momentum, &mom)); 2828 PetscCall(VecRestoreSubVector(u, isx, &position)); 2829 PetscCall(VecRestoreSubVector(u, isv, &momentum)); 2830 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, "velocity", &gv)); 2831 PetscCall(DMSwarmDestroyGlobalVectorFromField(sw, DMSwarmPICField_coor, &gc)); 2832 } 2833 PetscCall(DMSwarmMigrate(sw, PETSC_TRUE)); 2834 PetscInt step; 2835 2836 PetscCall(TSGetStepNumber(ts, &step)); 2837 if (!(step % ctx->remapFreq)) { 2838 // Monitor electric field before we destroy it 2839 PetscReal ptime; 2840 PetscInt step; 2841 2842 PetscCall(TSGetStepNumber(ts, &step)); 2843 PetscCall(TSGetTime(ts, &ptime)); 2844 if (ctx->efield_monitor) PetscCall(MonitorEField(ts, step, ptime, NULL, ctx)); 2845 if (ctx->poisson_monitor) PetscCall(MonitorPoisson(ts, step, ptime, NULL, ctx)); 2846 PetscCall(DMSwarmRemap(sw)); 2847 ctx->validE = PETSC_FALSE; 2848 } 2849 // This MUST come last, since it recreates the subswarms and they must DMClone() the new swarm 2850 PetscCall(DMSwarmTSRedistribute(ts)); 2851 PetscCall(InitializeSolveAndSwarm(ts, PETSC_FALSE)); 2852 PetscFunctionReturn(PETSC_SUCCESS); 2853 } 2854 2855 int main(int argc, char **argv) 2856 { 2857 DM dm, sw; 2858 TS ts; 2859 Vec u; 2860 PetscReal dt; 2861 PetscInt maxn; 2862 AppCtx user; 2863 2864 PetscCall(PetscInitialize(&argc, &argv, NULL, help)); 2865 PetscCall(ProcessOptions(PETSC_COMM_WORLD, &user)); 2866 PetscCall(PetscBagCreate(PETSC_COMM_SELF, sizeof(Parameter), &user.bag)); 2867 PetscCall(CreateMesh(PETSC_COMM_WORLD, &user, &dm)); 2868 PetscCall(CreatePoisson(dm, &user)); 2869 PetscCall(CreateMomentFields(dm, &user)); 2870 PetscCall(CreateSwarm(dm, &user, &sw)); 2871 PetscCall(SetupParameters(PETSC_COMM_WORLD, &user)); 2872 PetscCall(InitializeConstants(sw, &user)); 2873 PetscCall(DMSetApplicationContext(sw, &user)); 2874 2875 PetscCall(TSCreate(PETSC_COMM_WORLD, &ts)); 2876 PetscCall(TSSetProblemType(ts, TS_NONLINEAR)); 2877 PetscCall(TSSetDM(ts, sw)); 2878 PetscCall(TSSetMaxTime(ts, 0.1)); 2879 PetscCall(TSSetTimeStep(ts, 0.00001)); 2880 PetscCall(TSSetMaxSteps(ts, 100)); 2881 PetscCall(TSSetExactFinalTime(ts, TS_EXACTFINALTIME_MATCHSTEP)); 2882 2883 if (user.efield_monitor) PetscCall(TSMonitorSet(ts, MonitorEField, &user, NULL)); 2884 if (user.moment_monitor) PetscCall(TSMonitorSet(ts, MonitorMoments, &user, NULL)); 2885 if (user.moment_field_monitor) PetscCall(TSMonitorSet(ts, MonitorMomentFields, &user, NULL)); 2886 if (user.initial_monitor) PetscCall(TSMonitorSet(ts, MonitorInitialConditions, &user, NULL)); 2887 if (user.positions_monitor) PetscCall(TSMonitorSet(ts, MonitorPositions_2D, &user, NULL)); 2888 if (user.poisson_monitor) PetscCall(TSMonitorSet(ts, MonitorPoisson, &user, NULL)); 2889 if (user.velocity_monitor >= 0) PetscCall(TSMonitorSet(ts, MonitorVelocity, &user, NULL)); 2890 2891 PetscCall(TSSetFromOptions(ts)); 2892 PetscCall(TSGetTimeStep(ts, &dt)); 2893 PetscCall(TSGetMaxSteps(ts, &maxn)); 2894 user.steps = maxn; 2895 user.stepSize = dt; 2896 PetscCall(SetupContext(dm, sw, &user)); 2897 PetscCall(TSSetComputeInitialCondition(ts, InitializeSolve)); 2898 PetscCall(TSSetPostStep(ts, MigrateParticles)); 2899 PetscCall(CreateSolution(ts)); 2900 PetscCall(TSGetSolution(ts, &u)); 2901 PetscCall(TSComputeInitialCondition(ts, u)); 2902 PetscCall(CheckNonNegativeWeights(sw, &user)); 2903 PetscCall(TSSolve(ts, NULL)); 2904 2905 if (user.checkLandau) { 2906 // We should get a lookup table based on charge density and \hat k 2907 const PetscReal gammaEx = -0.15336; 2908 const PetscReal omegaEx = 1.4156; 2909 const PetscReal tol = 1e-2; 2910 2911 PetscCheck(PetscAbsReal((user.gamma - gammaEx) / gammaEx) < tol, PETSC_COMM_WORLD, PETSC_ERR_LIB, "Invalid Landau gamma %g != %g", user.gamma, gammaEx); 2912 PetscCheck(PetscAbsReal((user.omega - omegaEx) / omegaEx) < tol, PETSC_COMM_WORLD, PETSC_ERR_LIB, "Invalid Landau omega %g != %g", user.omega, omegaEx); 2913 } 2914 2915 PetscCall(SNESDestroy(&user.snes)); 2916 PetscCall(DMDestroy(&user.dmN)); 2917 PetscCall(ISDestroy(&user.isN)); 2918 PetscCall(MatDestroy(&user.MN)); 2919 PetscCall(DMDestroy(&user.dmP)); 2920 PetscCall(ISDestroy(&user.isP)); 2921 PetscCall(MatDestroy(&user.MP)); 2922 PetscCall(DMDestroy(&user.dmE)); 2923 PetscCall(ISDestroy(&user.isE)); 2924 PetscCall(MatDestroy(&user.ME)); 2925 PetscCall(DMDestroy(&user.dmMom)); 2926 PetscCall(DMDestroy(&user.dmPot)); 2927 PetscCall(ISDestroy(&user.isPot)); 2928 PetscCall(MatDestroy(&user.M)); 2929 PetscCall(PetscFEGeomDestroy(&user.fegeom)); 2930 PetscCall(TSDestroy(&ts)); 2931 PetscCall(DMDestroy(&sw)); 2932 PetscCall(DMDestroy(&dm)); 2933 PetscCall(DestroyContext(&user)); 2934 PetscCall(PetscFinalize()); 2935 return 0; 2936 } 2937 2938 /*TEST 2939 2940 build: 2941 requires: !complex double 2942 2943 # This tests that we can compute the correct decay rate and frequency 2944 # For gold runs, use -dm_plex_box_faces 160 -vdm_plex_box_faces 450 -remap_dm_plex_box_faces 80,150 2945 test: 2946 suffix: 0 2947 args: -cosine_coefficients 0.01 -charges -1. -perturbed_weights -total_weight 1. \ 2948 -dm_plex_dim 1 -dm_plex_box_faces 80 -dm_plex_box_lower 0. -dm_plex_box_upper 12.5664 \ 2949 -dm_plex_box_bd periodic -dm_plex_hash_location \ 2950 -vdm_plex_dim 1 -vdm_plex_box_faces 220 -vdm_plex_box_lower -6 -vdm_plex_box_upper 6 \ 2951 -vpetscspace_degree 2 -vdm_plex_hash_location \ 2952 -remap_freq 100 -dm_swarm_remap_type pfak -remap_dm_plex_dim 2 -remap_dm_plex_simplex 0 \ 2953 -remap_dm_plex_box_faces 40,110 -remap_dm_plex_box_bd periodic,none \ 2954 -remap_dm_plex_box_lower 0.,-6. -remap_dm_plex_box_upper 12.5664,6. \ 2955 -remap_petscspace_degree 1 -remap_dm_plex_hash_location \ 2956 -ftop_ksp_type lsqr -ftop_pc_type none -ftop_ksp_rtol 1.e-14 -ptof_pc_type lu \ 2957 -em_type primal -petscspace_degree 1 -em_snes_atol 1.e-12 -em_snes_error_if_not_converged \ 2958 -em_ksp_error_if_not_converged -em_pc_type svd -em_proj_pc_type lu \ 2959 -ts_type basicsymplectic -ts_basicsymplectic_type 1 -ts_dt 0.03 -ts_max_steps 500 \ 2960 -ts_max_time 100 \ 2961 -emax_tao_type brgn -emax_tao_max_it 100 -emax_tao_brgn_regularization_type l2pure \ 2962 -emax_tao_brgn_regularizer_weight 1e-5 -tao_brgn_subsolver_tao_bnk_ksp_rtol 1e-12 \ 2963 -output_step 1 -efield_monitor quiet 2964 2965 TEST*/ 2966