diff options
author | 2023-03-23 10:52:29 -0700 | |
---|---|---|
committer | 2023-03-23 17:52:29 +0000 | |
commit | f7b05278348b006aa8b54d58fb243b5536076918 (patch) | |
tree | e0560a4be1abbe17a8a30ef0b220054bbc2033a0 /Python | |
parent | fedca01aa44d4a1883f623601d11cf85fb353d2a (diff) | |
download | WarpX-f7b05278348b006aa8b54d58fb243b5536076918.tar.gz WarpX-f7b05278348b006aa8b54d58fb243b5536076918.tar.zst WarpX-f7b05278348b006aa8b54d58fb243b5536076918.zip |
Fix add_particles to account for theta attribute with RZ (#3776)
* Fix add_particles to account for theta attribute with RZ
* Add assertions in the number of attributes specified
Diffstat (limited to 'Python')
-rwxr-xr-x | Python/pywarpx/_libwarpx.py | 16 |
1 files changed, 12 insertions, 4 deletions
diff --git a/Python/pywarpx/_libwarpx.py b/Python/pywarpx/_libwarpx.py index a08b063ab..02ca22afb 100755 --- a/Python/pywarpx/_libwarpx.py +++ b/Python/pywarpx/_libwarpx.py @@ -683,14 +683,22 @@ class LibWarpX(): maxlen, val, self._numpy_particlereal_dtype ) - # --- The -3 is because the comps include the velocites - nattr = self.get_nattr_species(species_name) - 3 + # --- The number of built in attributes + # --- The three velocities + built_in_attrs = 3 + if self.geometry_dim == 'rz': + # --- With RZ, there is also theta + built_in_attrs += 1 + + # --- The number of extra attributes (including the weight) + nattr = self.get_nattr_species(species_name) - built_in_attrs attr = np.zeros((maxlen, nattr), self._numpy_particlereal_dtype) attr[:,0] = w + # --- Note that the velocities are handled separately and not included in attr + # --- (even though they are stored as attributes in the C++) for key, vals in kwargs.items(): - # --- The -3 is because components 1 to 3 are velocities - attr[:,self.get_particle_comp_index(species_name, key)-3] = vals + attr[:,self.get_particle_comp_index(species_name, key) - built_in_attrs] = vals nattr_int = 0 attr_int = np.empty([0], ctypes.c_int) |