diff --git a/openpmd_viewer/openpmd_timeseries/data_reader/data_reader.py b/openpmd_viewer/openpmd_timeseries/data_reader/data_reader.py index f7c1d1a6..dd1c2c4b 100644 --- a/openpmd_viewer/openpmd_timeseries/data_reader/data_reader.py +++ b/openpmd_viewer/openpmd_timeseries/data_reader/data_reader.py @@ -280,7 +280,7 @@ def read_species_data( self, iteration, species, record_comp, extensions): record_comp: string The record component to extract - Either 'x', 'y', 'z', 'ux', 'uy', 'uz', or 'w' + Either 'x', 'y', 'z', 'r', 'ux', 'uy', 'uz', 'ur', or 'w' extensions: list of strings The extensions that the current OpenPMDTimeSeries complies with diff --git a/openpmd_viewer/openpmd_timeseries/data_reader/h5py_reader/params_reader.py b/openpmd_viewer/openpmd_timeseries/data_reader/h5py_reader/params_reader.py index ffdcfdbf..196d1a4d 100644 --- a/openpmd_viewer/openpmd_timeseries/data_reader/h5py_reader/params_reader.py +++ b/openpmd_viewer/openpmd_timeseries/data_reader/h5py_reader/params_reader.py @@ -207,6 +207,10 @@ def simplify_record(record_comps): record_comps.remove('position/z') record_comps.remove('positionOffset/z') record_comps.append('z') + if ('position/r' in record_comps) and ('positionOffset/r' in record_comps): + record_comps.remove('position/r') + record_comps.remove('positionOffset/r') + record_comps.append('r') # Replace the names of the momenta if 'momentum/x' in record_comps: @@ -218,6 +222,9 @@ def simplify_record(record_comps): if 'momentum/z' in record_comps: record_comps.remove('momentum/z') record_comps.append('uz') + if 'momentum/r' in record_comps: + record_comps.remove('momentum/r') + record_comps.append('ur') # Replace the name for 'weights' if 'weighting' in record_comps: diff --git a/openpmd_viewer/openpmd_timeseries/data_reader/h5py_reader/particle_reader.py b/openpmd_viewer/openpmd_timeseries/data_reader/h5py_reader/particle_reader.py index 82fddd20..a1a5f3e6 100644 --- a/openpmd_viewer/openpmd_timeseries/data_reader/h5py_reader/particle_reader.py +++ b/openpmd_viewer/openpmd_timeseries/data_reader/h5py_reader/particle_reader.py @@ -32,7 +32,7 @@ def read_species_data(filename, iteration, species, record_comp, extensions): record_comp: string The record component to extract - Either 'x', 'y', 'z', 'ux', 'uy', 'uz', or 'w' + Either 'x', 'y', 'z', 'r', 'ux', 'uy', 'uz', 'ur', or 'w' extensions: list of strings The extensions that the current OpenPMDTimeSeries complies with @@ -43,9 +43,11 @@ def read_species_data(filename, iteration, species, record_comp, extensions): dict_record_comp = {'x': 'position/x', 'y': 'position/y', 'z': 'position/z', + 'r': 'position/r', 'ux': 'momentum/x', 'uy': 'momentum/y', 'uz': 'momentum/z', + 'ur': 'momentum/r', 'w': 'weighting'} if record_comp in dict_record_comp: opmd_record_comp = dict_record_comp[record_comp] @@ -78,11 +80,11 @@ def read_species_data(filename, iteration, species, record_comp, extensions): data *= w ** (-weighting_power) # - Return positions, with an offset - if record_comp in ['x', 'y', 'z']: + if record_comp in ['x', 'y', 'z', 'r']: offset = get_data(species_grp['positionOffset/%s' % record_comp]) data += offset # - Return momentum in normalized units - elif record_comp in ['ux', 'uy', 'uz' ]: + elif record_comp in ['ux', 'uy', 'uz', 'ur']: m = get_data(species_grp['mass']) # Normalize only if the particle mass is non-zero if np.all( m != 0 ): diff --git a/openpmd_viewer/openpmd_timeseries/data_reader/io_reader/params_reader.py b/openpmd_viewer/openpmd_timeseries/data_reader/io_reader/params_reader.py index d80f7b32..81ac248e 100644 --- a/openpmd_viewer/openpmd_timeseries/data_reader/io_reader/params_reader.py +++ b/openpmd_viewer/openpmd_timeseries/data_reader/io_reader/params_reader.py @@ -184,6 +184,10 @@ def simplify_record(record_comps): record_comps.remove('position/z') record_comps.remove('positionOffset/z') record_comps.append('z') + if ('position/r' in record_comps) and ('positionOffset/r' in record_comps): + record_comps.remove('position/r') + record_comps.remove('positionOffset/r') + record_comps.append('r') # Replace the names of the momenta if 'momentum/x' in record_comps: @@ -195,6 +199,9 @@ def simplify_record(record_comps): if 'momentum/z' in record_comps: record_comps.remove('momentum/z') record_comps.append('uz') + if 'momentum/r' in record_comps: + record_comps.remove('momentum/r') + record_comps.append('ur') # Replace the name for 'weights' if 'weighting' in record_comps: diff --git a/openpmd_viewer/openpmd_timeseries/data_reader/io_reader/particle_reader.py b/openpmd_viewer/openpmd_timeseries/data_reader/io_reader/particle_reader.py index 17119d6f..e9c51ade 100644 --- a/openpmd_viewer/openpmd_timeseries/data_reader/io_reader/particle_reader.py +++ b/openpmd_viewer/openpmd_timeseries/data_reader/io_reader/particle_reader.py @@ -32,7 +32,7 @@ def read_species_data(series, iteration, species_name, component_name, component_name: string The record component to extract - Either 'x', 'y', 'z', 'ux', 'uy', 'uz', or 'w' + Either 'x', 'y', 'z', 'r', 'ux', 'uy', 'uz', 'ur', or 'w' extensions: list of strings The extensions that the current OpenPMDTimeSeries complies with @@ -43,9 +43,11 @@ def read_species_data(series, iteration, species_name, component_name, dict_record_comp = {'x': ['position', 'x'], 'y': ['position', 'y'], 'z': ['position', 'z'], + 'r': ['position', 'r'], 'ux': ['momentum', 'x'], 'uy': ['momentum', 'y'], 'uz': ['momentum', 'z'], + 'ur': ['momentum', 'r'], 'w': ['weighting', None]} if component_name in dict_record_comp: @@ -84,11 +86,11 @@ def read_species_data(series, iteration, species_name, component_name, data *= w ** (-weighting_power) # - Return positions, with an offset - if component_name in ['x', 'y', 'z']: + if component_name in ['x', 'y', 'z', 'r']: offset = get_data(series, species['positionOffset'][component_name]) data += offset # - Return momentum in normalized units - elif component_name in ['ux', 'uy', 'uz' ]: + elif component_name in ['ux', 'uy', 'uz', 'ur']: mass_component = next(species['mass'].items())[1] m = get_data(series, mass_component) # Normalize only if the particle mass is non-zero