Enable support for heterogeneous data and expose max_resolution_3d

openpmd_viewer/openpmd_timeseries/interactive.py

@@ -10,6 +10,7 @@
 """
 import math
 from functools import partial
+import numpy as np
 try:
     from ipywidgets import widgets, __version__
     ipywidgets_version = int(__version__[0])
@@ -272,26 +273,40 @@ def refresh_species(change=None):
         def change_iteration(change):
             "Plot the result at the required iteration"
             # Find the closest iteration
-            self._current_i = abs(self.iterations - change['new']).argmin()
-            self.current_iteration = self.iterations[ self._current_i ]
+            set_current_iteration(change['new'])
             refresh_field()
             refresh_ptcl()
 
         def step_fw(b):
             "Plot the result one iteration further"
             if self._current_i < len(self.t) - 1:
-                self.current_iteration = self.iterations[self._current_i + 1]
-            else:
-                self.current_iteration = self.iterations[self._current_i]
-            slider.value = self.current_iteration
+                slider.value += 1
 
         def step_bw(b):
             "Plot the result one iteration before"
             if self._current_i > 0:
-                self.current_iteration = self.iterations[self._current_i - 1]
-            else:
-                self.current_iteration = self.iterations[self._current_i]
-            slider.value = self.current_iteration
+                slider.value -= 1                
+
+        def set_current_iteration(requested_iteration):
+            "Set the current iteration to the requested value"
+            # Get iterations.
+            iterations = get_available_iterations()
+            # Find the closest available  iteration
+            closest_iteration = iterations[abs(iterations - requested_iteration).argmin()]
+            self._current_i = abs(self.iterations - closest_iteration).argmin()
+            self.current_iteration = self.iterations[ self._current_i ]
+
+        def get_available_iterations():
+            "Get iterations in which both the current field and species are available."
+            field_iterations = None
+            if self.avail_fields is not None:
+                current_field = fieldtype_button.value
+                field_iterations = self.fields_iterations[current_field]
+            species_iterations = None
+            if self.avail_species is not None:
+                current_species = ptcl_species_button.value
+                species_iterations = self.species_iterations[current_species]
+            return get_common_iterations(field_iterations, species_iterations)
 
         # ---------------
         # Define widgets
@@ -516,6 +531,12 @@ def step_bw(b):
                 children=[ptcl_refresh_toggle, ptcl_refresh_button])])
             set_widget_dimensions( container_ptcl, width=370 )
 
+        # Try to set the current iteration of the slider
+        # to the current iteration of the OpenPMDTimeSeries.
+        # If the displayed field/species are not available at this
+        # iteration, it will be set to the closest available one.
+        set_current_iteration(self.current_iteration)
+
         # Global container
         if (self.avail_fields is not None) and \
                 (self.avail_species is not None):
@@ -871,3 +892,26 @@ def create_checkbox( **kwargs ):
     else:
         c = widgets.Checkbox( **kwargs )
     return(c)
+
+
+def get_common_iterations(field_iterations, species_iterations):
+    """Get iterations in which both the field and the species are available.
+
+    Parameters
+    ----------
+    field_iterations : ndarray
+        The iterations at which the field is available.
+    species_iterations : ndarray
+        The iterations at which the species is available.
+
+    Returns
+    -------
+    ndarray
+        The iterations common to both.
+    """
+    if field_iterations is None:
+        return species_iterations
+    elif species_iterations is None:
+        return field_iterations
+    else:
+        return np.intersect1d(field_iterations, species_iterations)

openpmd_viewer/openpmd_timeseries/main.py

@@ -80,36 +80,7 @@ def __init__(self, path_to_dir, check_all_files=True, backend=None):
 
         # Go through the files of the series, extract the time
         # and a few parameters.
-        N_iterations = len(self.iterations)
-        self.t = np.zeros(N_iterations)
-
-        # - Extract parameters from the first file
-        t, params0 = self.data_reader.read_openPMD_params(self.iterations[0])
-        self.t[0] = t
-        self.extensions = params0['extensions']
-        self.avail_fields = params0['avail_fields']
-        if self.avail_fields is not None:
-            self.fields_metadata = params0['fields_metadata']
-            self.avail_geom = set( self.fields_metadata[field]['geometry']
-                                for field in self.avail_fields )
-        # Extract information of the particles
-        self.avail_species = params0['avail_species']
-        self.avail_record_components = \
-            params0['avail_record_components']
-
-        # - Extract the time for each file and, if requested, check
-        #   that the other files have the same parameters
-        for k in range(1, N_iterations):
-            t, params = self.data_reader.read_openPMD_params(
-                self.iterations[k], check_all_files)
-            self.t[k] = t
-            if check_all_files:
-                for key in params0.keys():
-                    if params != params0:
-                        print("Warning: File %s has different openPMD "
-                              "parameters than the rest of the time series."
-                              % self.iterations[k])
-                        break
+        self.determine_available_data(check_all_files)
 
         # - Set the current iteration and time
         self._current_i = 0
@@ -122,6 +93,86 @@ def __init__(self, path_to_dir, check_all_files=True, backend=None):
         # - Initialize a plotter object, which holds information about the time
         self.plotter = Plotter(self.t, self.iterations)
 
+    def determine_available_data(self, check_all_files=True):
+        """Find the available fields and species and their metadata.
+
+        Parameters
+        ----------
+        check_all_files : bool, optional
+            Whether to look for fields and species in every file. If `False`,
+            it assumes that all iterations have the same data (i.e., the fields
+            and species available in the first iteration). By default True.
+        """
+        N_iterations = len(self.iterations)
+        self.t = np.zeros(N_iterations)
+        self.avail_fields = []
+        self.avail_species = []
+        self.avail_record_components = {}
+        self.fields_metadata = {}
+        self.fields_t = {}
+        self.species_t = {}
+        self.fields_iterations = {}
+        self.species_iterations = {}
+        self.extensions = []
+        for i, it in enumerate(self.iterations):
+            check_file = (i == 0) or check_all_files
+            t, params = self.data_reader.read_openPMD_params(it, check_file)
+            self.t[i] = t
+            if check_file:
+                avail_fields_it = params['avail_fields']
+                avail_species_it = params['avail_species']
+                avail_record_components_it = params['avail_record_components']
+                extensions_it = params['extensions']
+                self.extensions.extend([ex for ex in extensions_it
+                                        if ex not in self.extensions])
+                if avail_fields_it:
+                    fields_metadata_it = params['fields_metadata']
+                    new_fields = [fld for fld in avail_fields_it
+                                  if fld not in self.avail_fields]
+                    self.avail_fields.extend(new_fields)
+                    for fld in new_fields:
+                        self.fields_t[fld] = []
+                        self.fields_iterations[fld] = []
+                        self.fields_metadata[fld] = fields_metadata_it[fld]
+                    for fld in avail_fields_it:
+                        self.fields_t[fld].append(t)
+                        self.fields_iterations[fld].append(it)
+                if avail_species_it:
+                    new_species = [sp for sp in avail_species_it
+                                   if sp not in self.avail_species]
+                    self.avail_species.extend(new_species)
+                    for sp in new_species:
+                        self.species_t[sp] = []
+                        self.species_iterations[sp] = []
+                        self.avail_record_components[sp] = avail_record_components_it[sp]
+                    for sp in avail_species_it:
+                        self.species_t[sp].append(t)
+                        self.species_iterations[sp].append(it)
+
+        if self.avail_fields:
+            self.avail_geom = set(self.fields_metadata[fld]['geometry']
+                                  for fld in self.avail_fields)
+
+        if check_all_files:
+            for fld in self.avail_fields:
+                self.fields_t[fld] = np.array(self.fields_t[fld])
+                self.fields_iterations[fld] = np.array(self.fields_iterations[fld])
+            for sp in self.avail_species:
+                self.species_t[sp] = np.array(self.species_t[sp])
+                self.species_iterations[sp] = np.array(self.species_iterations[sp])
+
+        else:
+            for fld in self.avail_fields:
+                self.fields_t[fld] = self.t
+                self.fields_iterations[fld] = self.iterations
+            for sp in self.avail_species:
+                self.species_t[sp] = self.t
+                self.species_iterations[sp] = self.iterations
+
+        # For backwards compatibility
+        if not self.avail_fields:
+            self.avail_fields = None
+
     def get_particle(self, var_list=None, species=None, t=None, iteration=None,
             select=None, plot=False, nbins=150,
             plot_range=[[None, None], [None, None]],
@@ -264,7 +315,8 @@ def get_particle(self, var_list=None, species=None, t=None, iteration=None,
 
         # Find the output that corresponds to the requested time/iteration
         # (Modifies self._current_i, self.current_iteration and self.current_t)
-        self._find_output(t, iteration)
+        self._find_output(t, iteration, self.species_t[species],
+                          self.species_iterations[species])
         # Get the corresponding iteration
         iteration = self.iterations[self._current_i]
 
@@ -362,8 +414,8 @@ def get_particle(self, var_list=None, species=None, t=None, iteration=None,
 
     def get_field(self, field=None, coord=None, t=None, iteration=None,
                   m='all', theta=0., slice_across=None,
-                  slice_relative_position=None, plot=False,
-                  plot_range=[[None, None], [None, None]], **kw):
+                  slice_relative_position=None, max_resolution_3d=None,
+                  plot=False, plot_range=[[None, None], [None, None]], **kw):
         """
         Extract a given field from a file in the openPMD format.
 
@@ -416,6 +468,13 @@ def get_field(self, field=None, coord=None, t=None, iteration=None,
            Default: None, which results in slicing at 0 in all direction
            of `slice_across`.
 
+        max_resolution_3d : list of int or None
+            Maximum resolution that the 3D reconstruction of the field (when
+            `theta` is None) can have. The list should contain two values,
+            e.g. `[200, 100]`, indicating the maximum longitudinal and
+            transverse resolution, respectively. This is useful for
+            performance reasons, particularly for 3D visualization.
+
         plot : bool, optional
            Whether to plot the requested quantity
 
@@ -485,7 +544,8 @@ def get_field(self, field=None, coord=None, t=None, iteration=None,
 
         # Find the output that corresponds to the requested time/iteration
         # (Modifies self._current_i, self.current_iteration and self.current_t)
-        self._find_output(t, iteration)
+        self._find_output(t, iteration, self.fields_t[field],
+                          self.fields_iterations[field])
         # Get the corresponding iteration
         iteration = self.iterations[self._current_i]
 
@@ -510,16 +570,16 @@ def get_field(self, field=None, coord=None, t=None, iteration=None,
                 # For Cartesian components, combine r and t components
                 Fr, info = self.data_reader.read_field_circ(
                     iteration, field, 'r', slice_relative_position,
-                    slice_across, m, theta)
+                    slice_across, m, theta, max_resolution_3d)
                 Ft, info = self.data_reader.read_field_circ(
                     iteration, field, 't', slice_relative_position,
-                    slice_across, m, theta)
+                    slice_across, m, theta, max_resolution_3d)
                 F = combine_cylindrical_components(Fr, Ft, theta, coord, info)
             else:
                 # For cylindrical or scalar components, no special treatment
                 F, info = self.data_reader.read_field_circ(iteration,
                     field, coord, slice_relative_position,
-                    slice_across, m, theta)
+                    slice_across, m, theta, max_resolution_3d)
 
         # Plot the resulting field
         # Deactivate plotting when there is no slice selection
@@ -591,7 +651,7 @@ def iterate( self, called_method, *args, **kwargs ):
             accumulated_result = try_array( accumulated_result )
             return accumulated_result
 
-    def _find_output(self, t, iteration):
+    def _find_output(self, t, iteration, record_t, record_iterations):
         """
         Find the output that correspond to the requested `t` or `iteration`
         Modify self._current_i accordingly.
@@ -603,6 +663,12 @@ def _find_output(self, t, iteration):
 
         iteration : int
             Iteration requested
+
+        record_t : ndarray
+            The time steps at which the record (species/field) is available.
+
+        record_iterations : ndarray
+            The iterations at which the record (species/field) is available.
         """
         # Check the arguments
         if (t is not None) and (iteration is not None):
@@ -612,20 +678,21 @@ def _find_output(self, t, iteration):
         # If a time is requested
         elif (t is not None):
             # Make sure the time requested does not exceed the allowed bounds
-            if t < self.tmin:
-                self._current_i = 0
-            elif t > self.tmax:
-                self._current_i = len(self.t) - 1
+            if t < np.min(record_t):
+                self._current_i = abs(record_iterations[0] - self.iterations).argmin()
+            elif t > np.max(record_t):
+                self._current_i = abs(record_iterations[-1] - self.iterations).argmin()
             # Find the closest existing iteration
             else:
-                self._current_i = abs(self.t - t).argmin()
+                closest_record_iteration = record_iterations[abs(record_t - t).argmin()]
+                self._current_i = abs(closest_record_iteration - self.iterations).argmin()
         # If an iteration is requested
         elif (iteration is not None):
-            if (iteration in self.iterations):
+            if (iteration in record_iterations):
                 # Get the index that corresponds to this iteration
                 self._current_i = abs(iteration - self.iterations).argmin()
             else:
-                iter_list = '\n - '.join([str(it) for it in self.iterations])
+                iter_list = '\n - '.join([str(it) for it in record_iterations])
                 raise OpenPMDException(
                       "The requested iteration '%s' is not available.\nThe "
                       "available iterations are: \n - %s\n" % (iteration, iter_list))