Add more OpenPMD particle components to read

AngelFP · Apr 7, 2020 · 96fc2c8 · 96fc2c8
1 parent 0aab813
commit 96fc2c8
Showing 1 changed file with 14 additions and 1 deletion.
diff --git a/visualpic/data_reading/particle_readers.py b/visualpic/data_reading/particle_readers.py
@@ -153,7 +153,10 @@ def __init__(self, *args, **kwargs):
                                'pz': 'momentum/z',
                                'px': 'momentum/x',
                                'py': 'momentum/y',
-                               'q': 'charge'}
+                               'q': 'charge',
+                               'm': 'mass',
+                               'tag': 'id',
+                               'w': 'weighting'}
         return super().__init__(*args, **kwargs)
 
     def _read_component_data(self, file_path, species, component):
@@ -179,6 +182,12 @@ def _read_component_data(self, file_path, species, component):
                 data = beam_species[component_to_read].attrs['value']
                 w = beam_species['weighting'][:]
                 data = data * w
+            elif component_to_read == 'mass':
+                data = beam_species[component_to_read].attrs['value']
+                w = beam_species['weighting'][:]
+                data = data * w
+            else:
+                data = beam_species[component_to_read][:]
             return data
 
     def _read_component_metadata(self, file_path, species, component):
@@ -193,6 +202,10 @@ def _read_component_metadata(self, file_path, species, component):
             metadata['units'] = 'm_e*c'
         elif component_to_read == 'charge':
             metadata['units'] = 'C'
+        elif component_to_read == 'mass':
+            metadata['units'] = 'kg'
+        else:
+            metadata['units'] = ''
         metadata['time'] = {}
         metadata['time']['value'] = t
         metadata['time']['units'] = 's'