Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

[WIP] Avoid assuming types in particle tracker #342

Open
wants to merge 1 commit into
base: dev
Choose a base branch
from
Open

[WIP] Avoid assuming types in particle tracker #342

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
6 changes: 3 additions & 3 deletions openpmd_viewer/openpmd_timeseries/main.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down Expand Up @@ -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]:
Expand Down
20 changes: 12 additions & 8 deletions openpmd_viewer/openpmd_timeseries/particle_tracker.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

Expand All @@ -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.
Expand All @@ -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
Expand All @@ -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 ):
Expand Down