diff --git a/openpmd_viewer/openpmd_timeseries/interactive.py b/openpmd_viewer/openpmd_timeseries/interactive.py index 3dfb5db5..de7c37e8 100644 --- a/openpmd_viewer/openpmd_timeseries/interactive.py +++ b/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) diff --git a/openpmd_viewer/openpmd_timeseries/main.py b/openpmd_viewer/openpmd_timeseries/main.py index 55e2077a..ef764e71 100644 --- a/openpmd_viewer/openpmd_timeseries/main.py +++ b/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))