diff --git a/openpmd_viewer/openpmd_timeseries/main.py b/openpmd_viewer/openpmd_timeseries/main.py index 42458d21..ac65ad1c 100644 --- a/openpmd_viewer/openpmd_timeseries/main.py +++ b/openpmd_viewer/openpmd_timeseries/main.py @@ -139,7 +139,7 @@ def get_particle(self, var_list=None, species=None, t=None, iteration=None, Parameters ---------- - var_list : list of string, optional + var_list : list of string A list of the particle variables to extract. If var_list is not provided, the available particle quantities are printed @@ -278,8 +278,8 @@ def get_particle(self, var_list=None, species=None, t=None, iteration=None, data_list = apply_selection( iteration, self.data_reader, data_list, select, species, self.extensions) elif isinstance( select, ParticleTracker ): - data_list = select.extract_tracked_particles( iteration, - self.data_reader, data_list, species, self.extensions ) + data_list = select.extract_tracked_particles(iteration, + self.data_reader, data_list, var_list, species, self.extensions) # Plotting if plot and len(var_list) in [1, 2]: diff --git a/openpmd_viewer/openpmd_timeseries/particle_tracker.py b/openpmd_viewer/openpmd_timeseries/particle_tracker.py index b51e6d98..8ad09731 100644 --- a/openpmd_viewer/openpmd_timeseries/particle_tracker.py +++ b/openpmd_viewer/openpmd_timeseries/particle_tracker.py @@ -106,7 +106,7 @@ def __init__(self, ts, species=None, t=None, def extract_tracked_particles( self, iteration, data_reader, data_list, - species, extensions ): + var_list, species, extensions ): """ Select the elements of each particle quantities in data_list, so as to only return those that correspond to the tracked particles @@ -123,6 +123,9 @@ def extract_tracked_particles( self, iteration, data_reader, data_list, A list of arrays with one element per macroparticle, that represent different particle quantities + var_list: list of string + The list of variable names that corresponds to `data_list` + species: string Name of the species being requested @@ -144,11 +147,11 @@ def extract_tracked_particles( self, iteration, data_reader, data_list, for i in range(len(data_list)): if len(data_list[i]) > 1: # Do not apply selection on scalars data_list[i] = self.extract_quantity( - data_list[i], selected_indices ) + data_list[i], var_list[i], selected_indices ) return( data_list ) - def extract_quantity( self, q, selected_indices ): + def extract_quantity( self, q, name, selected_indices ): """ Select the elements of the array `q`, so as to only return those that correspond to the tracked particles. @@ -158,6 +161,9 @@ def extract_quantity( self, q, selected_indices ): q: 1d array of floats or ints A particle quantity (one element per particle) + name: string + Name of the particle quantity `q` + selected_indices: 1d array of ints The indices (in array q) of the particles to be selected. If `preserve_particle_index` was selected to be True, this array @@ -174,14 +180,12 @@ def extract_quantity( self, q, selected_indices ): # Handle the absent particles if self.preserve_particle_index: - if q.dtype in [ np.float64, np.float32 ]: + if name == 'id': + selected_q = self.selected_pid + else: # Fill the position of absent particles by NaNs selected_q = np.where( selected_indices == -1, np.nan, selected_q) - else: - # The only non-float quantity in openPMD-viewer is particle id - selected_q = self.selected_pid - return( selected_q ) def get_extraction_indices( self, pid ):