Conversation
|
I think this is it. Two stream instability is 3 times faster. @cwoolfo1 it is ready for review |
There was a problem hiding this comment.
Pull request overview
This PR switches the default particle backend to a new “flat” particle layout when the simulation configuration is compatible, with automatic fallback to the existing species-per-list backend when it is not. The goal is to improve runtime performance while keeping solver/field physics unchanged and preserving output metadata.
Changes:
- Default to the flat particle backend when compatible; otherwise fall back to the existing backend.
- Add
flat_particles.pyimplementing a flattened species container, compatibility checks, and flattening conversion. - Update Boris pusher to support per-particle charge/mass arrays and update TOML parameter dumping to preserve original species metadata.
Reviewed changes
Copilot reviewed 4 out of 4 changed files in this pull request and generated 5 comments.
| File | Description |
|---|---|
| PyPIC3D/initialization.py | Adds fast_backend selection logic and defaults to flat when compatible. |
| PyPIC3D/flat_particles.py | Introduces flattened particle species representation and flattening utilities. |
| PyPIC3D/boris.py | Extends vmap setup to handle scalar vs per-particle q/m. |
| PyPIC3D/utils.py | Uses species_meta to dump original species metadata when flat backend is active. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| def tree_flatten(self): | ||
| children = (self.x1, self.x2, self.x3, self.v1, self.v2, self.v3) | ||
| aux_data = ( | ||
| self.name, | ||
| self.N_particles, | ||
| self.charge, | ||
| self.mass, | ||
| self.weight, | ||
| self.T, | ||
| self.x_wind, | ||
| self.y_wind, | ||
| self.z_wind, | ||
| self.dx, | ||
| self.dy, | ||
| self.dz, | ||
| self.x_bc, | ||
| self.y_bc, | ||
| self.z_bc, | ||
| self.update_pos, | ||
| self.update_v, | ||
| self.update_x, | ||
| self.update_y, | ||
| self.update_z, | ||
| self.update_vx, | ||
| self.update_vy, | ||
| self.update_vz, | ||
| self.shape, | ||
| self.dt, | ||
| self.species_meta, | ||
| ) |
| self.shape, | ||
| self.dt, | ||
| self.species_meta, | ||
| ) |
| if particles[0].x_bc != "periodic" or particles[0].y_bc != "periodic" or particles[0].z_bc != "periodic": | ||
| return False | ||
| if not _same([p.update_pos for p in particles]) or not _same([p.update_v for p in particles]): | ||
| return False |
|
Updated openPMD particle diagnostics to account for flat particles |
|
I am getting a 1.26x speedup for the two stream instability |
|
WITHOUT FLATTENED BACKEND: Kernel Compile Time: 1.0667036710001412 s WITH FLATTENED BACKEND: Kernel Compile Time: 0.6272344169992721 s Biggest performance speedups appear to happen in the particle pusher and the position updates/bc. Current deposition remains about the same |
Summary
Changes (minimal core edits)
PyPIC3D/initialization.py: choose flat backend by default when compatible; fallback todefaultfor incompatible configs.PyPIC3D/flat_particles.py: flat species structure + compatibility normalization + momentum helper.PyPIC3D/boris.py: vmap handles per-particle charge/mass arrays.PyPIC3D/utils.py: output metadata uses original species list when flat backend is active.Benchmarks (long runs, CPU, JAX x64)
Configs:
t_wind = 1.6239e-7(3× longer)Nt = 18000(3× longer)Runtime summary
Raw results JSON (artifacts branch):
docs/benchmarks/results.jsonPlots (artifacts branch)
Runtime per step
Total wall time
Speedup per step
Speedup total
Two-stream energy
Weibel energy
Notes
codex/flat-default-artifactsto keep plots/results out of the main branch.