From 715a565fa0a573c5d14eb398c6680d27e67e8017 Mon Sep 17 00:00:00 2001 From: Patrick Avery Date: Thu, 9 Nov 2023 17:19:40 -0600 Subject: [PATCH 1/4] Refactor overlay drawing to use fewer artists For all overlay types, we were using more artists than we needed to. For powder overlays, for instance, we would have one artist per line. For Laue and rotation series overlays, we would have one artist per spot and one artist per range! This could easily result in hundreds of artists. However, if everything about the artists (including the style) are identical, except for the data, we are able to merge them together into a single artist. This can make matplotlib run much faster. For line artists, we can have a single artist draw every powder overlay line - we just insert a `[nan, nan]` row in between lines in the data. Different artists are used for different styles (i. e., merged ranges are red instead of green, so they get a different artist, and highlighted data/ranges are a different color too, so they get a different artist). But we end up using only about 5 line artists in total for the main canvas, per overlay and per detector (each detector is currently rendered separately). For scatter (path) artists, we ought to just pass every single spot to a single call to `scatter()`. These changes speed up rendering of the overlays significantly, especially when there were many lines or spots being drawn. This is a great step toward better interactivity. Signed-off-by: Patrick Avery --- hexrdgui/image_canvas.py | 302 ++++++++++++------- hexrdgui/overlays/laue_overlay.py | 11 +- hexrdgui/overlays/powder_overlay.py | 18 +- hexrdgui/overlays/rotation_series_overlay.py | 11 +- hexrdgui/utils/array.py | 17 ++ 5 files changed, 234 insertions(+), 125 deletions(-) create mode 100644 hexrdgui/utils/array.py diff --git a/hexrdgui/image_canvas.py b/hexrdgui/image_canvas.py index 89e332df2..f71a597ca 100644 --- a/hexrdgui/image_canvas.py +++ b/hexrdgui/image_canvas.py @@ -5,12 +5,13 @@ from PySide6.QtWidgets import QFileDialog, QMessageBox from matplotlib.backends.backend_qt5agg import FigureCanvas - from matplotlib.figure import Figure from matplotlib.lines import Line2D from matplotlib.patches import Circle from matplotlib.ticker import AutoLocator, FuncFormatter + import matplotlib.pyplot as plt +import matplotlib.transforms as tx import numpy as np @@ -23,6 +24,7 @@ from hexrdgui.hexrd_config import HexrdConfig from hexrdgui.snip_viewer_dialog import SnipViewerDialog from hexrdgui import utils +from hexrdgui.utils.array import split_array from hexrdgui.utils.conversions import ( angles_to_stereo, cart_to_angles, cart_to_pixels, q_to_tth, tth_to_q, ) @@ -232,9 +234,14 @@ def remove_overlay_artists(self, key): if key not in self.overlay_artists: return - artists = self.overlay_artists[key] - while artists: - artists.pop(0).remove() + for det_key, artist_dict in self.overlay_artists[key].items(): + for artist_name, artist in artist_dict.items(): + if isinstance(artist, list): + while artist: + artist.pop(0).remove() + else: + artist.remove() + del self.overlay_artists[key] def prune_overlay_artists(self): @@ -262,11 +269,12 @@ def overlay_axes_data(self, overlay): if not all(x in overlay.data for x in self.raw_axes): return [] - return [(self.raw_axes[x], overlay.data[x]) for x in self.raw_axes] + return [(self.raw_axes[x], x, overlay.data[x]) + for x in self.raw_axes] # If it is anything else, there is only one axis # Use the same axis for all of the data - return [(self.axis, x) for x in overlay.data.values()] + return [(self.axis, k, v) for k, v in overlay.data.items()] def overlay_draw_func(self, type): overlay_funcs = { @@ -306,9 +314,10 @@ def draw_overlay(self, overlay): type = overlay.type style = overlay.style highlight_style = overlay.highlight_style - for axis, data in self.overlay_axes_data(overlay): + for axis, det_key, data in self.overlay_axes_data(overlay): kwargs = { 'artist_key': overlay.name, + 'det_key': det_key, 'axis': axis, 'data': data, 'style': style, @@ -316,74 +325,100 @@ def draw_overlay(self, overlay): } self.overlay_draw_func(type)(**kwargs) - def draw_powder_overlay(self, artist_key, axis, data, style, + def draw_powder_overlay(self, artist_key, det_key, axis, data, style, highlight_style): rings = data['rings'] - rbnds = data['rbnds'] + ranges = data['rbnds'] rbnd_indices = data['rbnd_indices'] data_style = style['data'] ranges_style = style['ranges'] - highlight_indices = [i for i, x in enumerate(rings) - if id(x) in self.overlay_highlight_ids] - - artists = self.overlay_artists.setdefault(artist_key, []) - for i, pr in enumerate(rings): - current_style = data_style - if i in highlight_indices: - # Override with highlight style - current_style = highlight_style['data'] - - x, y = pr.T - artist, = axis.plot(x, y, **current_style) - artists.append(artist) - - # Add the rbnds too - for ind, pr in zip(rbnd_indices, rbnds): - x, y = pr.T - current_style = copy.deepcopy(ranges_style) - if any(x in highlight_indices for x in ind): - # Override with highlight style - current_style = highlight_style['ranges'] - elif len(ind) > 1: - # If ranges are combined, override the color to red - current_style['c'] = 'r' - artist, = axis.plot(x, y, **current_style) - artists.append(artist) - - if self.azimuthal_integral_axis is not None: - az_axis = self.azimuthal_integral_axis - for pr in rings: - x = pr[:, 0] - if len(x) == 0: - # Skip over rings that are out of bounds - continue - - # Average the points together for the vertical line - x = np.nanmean(x) - artist = az_axis.axvline(x, **data_style) - artists.append(artist) - - # Add the rbnds too - for ind, pr in zip(rbnd_indices, rbnds): - x = pr[:, 0] - if len(x) == 0: - # Skip over rbnds that are out of bounds - continue - - # Average the points together for the vertical line - x = np.nanmean(x) - - current_style = copy.deepcopy(ranges_style) - if len(ind) > 1: - # If rbnds are combined, override the color to red - current_style['c'] = 'r' - - artist = az_axis.axvline(x, **current_style) - artists.append(artist) - - def draw_laue_overlay(self, artist_key, axis, data, style, + merged_ranges_style = copy.deepcopy(ranges_style) + merged_ranges_style['c'] = 'r' + + overlay_artists = self.overlay_artists.setdefault(artist_key, {}) + artists = overlay_artists.setdefault(det_key, {}) + + highlight_indices = [] + + if self.overlay_highlight_ids: + # Split up highlighted and non-highlighted components for all + highlight_indices = [i for i, x in enumerate(rings) + if id(x) in self.overlay_highlight_ids] + + def split(data): + if not highlight_indices or len(data) == 0: + return [], data + + return split_array(data, highlight_indices) + + h_rings, rings = split(rings) + + # Find merged ranges and highlighted ranges + # Some ranges will be both "merged" and "highlighted" + merged_ranges = [] + h_ranges = [] + reg_ranges = [] + + found = False + for i, ind in enumerate(rbnd_indices): + if len(ind) > 1: + merged_ranges.append(ranges[i]) + found = True + + if highlight_indices and any(x in highlight_indices for x in ind): + h_ranges.append(ranges[i]) + found = True + + if not found: + # Not highlighted or merged + reg_ranges.append(ranges[i]) + else: + found = False + + def plot(data, key, kwargs): + # This logic was repeated + if len(data) != 0: + artists[key], = axis.plot(*np.vstack(data).T, **kwargs) + + plot(rings, 'rings', data_style) + plot(h_rings, 'h_rings', highlight_style['data']) + + plot(reg_ranges, 'ranges', ranges_style) + plot(merged_ranges, 'merged_ranges', merged_ranges_style) + # Highlighting goes after merged ranges to get precedence + plot(h_ranges, 'h_ranges', highlight_style['ranges']) + + az_axis = self.azimuthal_integral_axis + if az_axis: + trans = tx.blended_transform_factory(az_axis.transData, + az_axis.transAxes) + + def az_plot(data, key, kwargs): + if len(data) == 0: + return + + xmeans = np.array([np.nanmean(x[:, 0]) for x in data]) + + x = np.repeat(xmeans, 3) + y = np.tile([0, 1, np.nan], len(xmeans)) + + artists[key], = az_axis.plot(x, y, transform=trans, **kwargs) + + az_plot(rings, 'az_rings', data_style) + # NOTE: we still use the data_style for az_axis highlighted rings + az_plot(h_rings, 'az_h_rings', data_style) + + az_plot(reg_ranges, 'az_ranges', ranges_style) + + # NOTE: we still use the ranges_style for az_axis highlighted rings + az_plot(h_ranges, 'az_h_ranges', ranges_style) + + # Give merged ranges style precedence + az_plot(merged_ranges, 'az_merged_ranges', merged_ranges_style) + + def draw_laue_overlay(self, artist_key, det_key, axis, data, style, highlight_style): spots = data['spots'] ranges = data['ranges'] @@ -394,58 +429,90 @@ def draw_laue_overlay(self, artist_key, axis, data, style, ranges_style = style['ranges'] label_style = style['labels'] - highlight_indices = [i for i, x in enumerate(spots) - if id(x) in self.overlay_highlight_ids] + highlight_indices = [] + + if self.overlay_highlight_ids: + # Split up highlighted and non-highlighted components for all + highlight_indices = [i for i, x in enumerate(spots) + if id(x) in self.overlay_highlight_ids] + + def split(data): + if not highlight_indices or len(data) == 0: + return [], data + + return split_array(data, highlight_indices) + + h_spots, spots = split(spots) + h_ranges, ranges = split(ranges) + h_labels, labels = split(labels) - artists = self.overlay_artists.setdefault(artist_key, []) - for i, (x, y) in enumerate(spots): - current_style = data_style - if i in highlight_indices: - current_style = highlight_style['data'] + overlay_artists = self.overlay_artists.setdefault(artist_key, {}) + artists = overlay_artists.setdefault(det_key, {}) - artist = axis.scatter(x, y, **current_style) - artists.append(artist) + def scatter(data, key, kwargs): + # This logic was repeated + if len(data) != 0: + artists[key] = axis.scatter(*np.asarray(data).T, **kwargs) - if labels: - current_label_style = label_style - if i in highlight_indices: - current_label_style = highlight_style['labels'] + def plot(data, key, kwargs): + # This logic was repeated + if len(data) != 0: + artists[key], = axis.plot(*np.vstack(data).T, **kwargs) + # Draw spots and highlighted spots + scatter(spots, 'spots', data_style) + scatter(h_spots, 'h_spots', highlight_style['data']) + + # Draw ranges and highlighted ranges + plot(ranges, 'ranges', ranges_style) + plot(h_ranges, 'h_ranges', highlight_style['ranges']) + + # Draw labels and highlighted labels + if len(labels) or len(h_labels): + def plot_label(x, y, label, style): kwargs = { 'x': x + label_offsets[0], 'y': y + label_offsets[1], - 's': labels[i], + 's': label, 'clip_on': True, - **current_label_style, + **style, } - artist = axis.text(**kwargs) - artists.append(artist) - - for i, box in enumerate(ranges): - current_style = ranges_style - if i in highlight_indices: - current_style = highlight_style['ranges'] - - x, y = zip(*box) - artist, = axis.plot(x, y, **current_style) - artists.append(artist) - - def draw_rotation_series_overlay(self, artist_key, axis, data, style, - highlight_style): + return axis.text(**kwargs) + + # I don't know of a way to use a single artist for all labels. + # FIXME: figure out how to make this faster, if needed. + artists.setdefault('labels', []) + for label, (x, y) in zip(labels, spots): + artists['labels'].append(plot_label(x, y, label, label_style)) + + # I don't know of a way to use a single artist for all labels. + # FIXME: figure out how to make this faster, if needed. + artists.setdefault('h_labels', []) + style = highlight_style['labels'] + for label, (x, y) in zip(h_labels, h_spots): + artists['h_labels'].append(plot_label(x, y, label, style)) + + def draw_rotation_series_overlay(self, artist_key, det_key, axis, data, + style, highlight_style): is_aggregated = HexrdConfig().is_aggregated ome_range = HexrdConfig().omega_ranges aggregated = data['aggregated'] or is_aggregated or ome_range is None - if not aggregated: - ome_width = data['omega_width'] - ome_mean = np.mean(ome_range) - full_range = (ome_mean - ome_width / 2, ome_mean + ome_width / 2) - - def in_range(x): - return aggregated or full_range[0] <= x <= full_range[1] # Compute the indices that are in range for the current omega value ome_points = data['omegas'] - indices_in_range = [i for i, x in enumerate(ome_points) if in_range(x)] + + if aggregated: + # This means we will keep all + slicer = slice(None) + else: + ome_width = data['omega_width'] + ome_mean = np.mean(ome_range) + ome_min = ome_mean - ome_width / 2 + ome_max = ome_mean + ome_width / 2 + + in_range = np.logical_and(ome_min <= ome_points, + ome_points <= ome_max) + slicer = np.where(in_range) data_points = data['data'] ranges = data['ranges'] @@ -453,20 +520,21 @@ def in_range(x): data_style = style['data'] ranges_style = style['ranges'] - artists = self.overlay_artists.setdefault(artist_key, []) - for i in indices_in_range: - # data - x, y = data_points[i] - artist = axis.scatter(x, y, **data_style) - artists.append(artist) + if len(data_points) == 0: + return - # ranges - if i >= len(ranges): - continue + sliced_data = data_points[slicer] + if len(sliced_data) == 0: + return + + overlay_artists = self.overlay_artists.setdefault(artist_key, {}) + artists = overlay_artists.setdefault(det_key, {}) + + artists['data'] = axis.scatter(*sliced_data.T, **data_style) - x, y = zip(*ranges[i]) - artist, = axis.plot(x, y, **ranges_style) - artists.append(artist) + sliced_ranges = np.asarray(ranges)[slicer] + artists['ranges'], = axis.plot(*np.vstack(sliced_ranges).T, + **ranges_style) def redraw_overlay(self, overlay): # Remove the artists for this overlay diff --git a/hexrdgui/overlays/laue_overlay.py b/hexrdgui/overlays/laue_overlay.py index 8168ff909..d42d880a5 100644 --- a/hexrdgui/overlays/laue_overlay.py +++ b/hexrdgui/overlays/laue_overlay.py @@ -329,7 +329,13 @@ def range_data(self, spots, display_mode, panel): if self.width_shape not in range_func: raise Exception(f'Unknown range shape: {self.width_shape}') - return range_func[self.width_shape](spots, display_mode, panel) + data = range_func[self.width_shape](spots, display_mode, panel) + + # Add a nans row at the end of each range + # This makes it easier to vstack them for plotting + data = [np.append(x, nans_row, axis=0) for x in data] + + return data def rectangular_range_data(self, spots, display_mode, panel): range_corners = self.range_corners(spots) @@ -456,3 +462,6 @@ class LaueRangeShape(str, Enum): class LaueLabelType(str, Enum): hkls = 'hkls' energy = 'energy' + +# Constants +nans_row = np.nan * np.ones((1, 2)) diff --git a/hexrdgui/overlays/powder_overlay.py b/hexrdgui/overlays/powder_overlay.py index e176ac949..1c05d5b70 100644 --- a/hexrdgui/overlays/powder_overlay.py +++ b/hexrdgui/overlays/powder_overlay.py @@ -223,16 +223,22 @@ def generate_overlay(self): if plane_data.tThWidth is not None: # Generate the ranges too - lower_pts, _ = self.generate_ring_points( + lower_pts, lower_skipped = self.generate_ring_points( instr, r_lower, etas, panel, display_mode ) - upper_pts, _ = self.generate_ring_points( + upper_pts, upper_skipped = self.generate_ring_points( instr, r_upper, etas, panel, display_mode ) - for lpts, upts in zip(lower_pts, upper_pts): - point_groups[det_key]['rbnds'] += [lpts, upts] - for ind in indices: - point_groups[det_key]['rbnd_indices'] += [ind, ind] + lower_indices = [x for i, x in enumerate(indices) + if i not in lower_skipped] + upper_indices = [x for i, x in enumerate(indices) + if i not in upper_skipped] + + point_groups[det_key]['rbnds'] += lower_pts + point_groups[det_key]['rbnd_indices'] += lower_indices + + point_groups[det_key]['rbnds'] += upper_pts + point_groups[det_key]['rbnd_indices'] += upper_indices return point_groups diff --git a/hexrdgui/overlays/rotation_series_overlay.py b/hexrdgui/overlays/rotation_series_overlay.py index b282488e2..8831e32d5 100644 --- a/hexrdgui/overlays/rotation_series_overlay.py +++ b/hexrdgui/overlays/rotation_series_overlay.py @@ -282,7 +282,13 @@ def range_corners(self, spots): return ranges def range_data(self, spots, display_mode, panel): - return self.rectangular_range_data(spots, display_mode, panel) + data = self.rectangular_range_data(spots, display_mode, panel) + + # Add a nans row at the end of each range + # This makes it easier to vstack them for plotting + data = [np.append(x, nans_row, axis=0) for x in data] + + return data def rectangular_range_data(self, spots, display_mode, panel): from hexrdgui.hexrd_config import HexrdConfig @@ -371,3 +377,6 @@ def sync_omegas(self): if self.update_needed: HexrdConfig().overlay_config_changed.emit() HexrdConfig().update_overlay_editor.emit() + +# Constants +nans_row = np.nan * np.ones((1, 2)) diff --git a/hexrdgui/utils/array.py b/hexrdgui/utils/array.py new file mode 100644 index 000000000..247a7f660 --- /dev/null +++ b/hexrdgui/utils/array.py @@ -0,0 +1,17 @@ +import numpy as np + + +def split_array(x, indices, axis=0): + # Split an array into two subarrays: + # the first containing the values at the indices, + # and the second containing the values *not* at the indices. + try: + # This is faster than the list version + taken = np.take(x, indices, axis=axis) + deleted = np.delete(x, indices, axis=axis) + except ValueError: + # The list version might take longer + taken = [x for i, x in enumerate(x) if i in indices] + deleted = [x for i, x in enumerate(x) if i not in indices] + + return taken, deleted From cf9b726f9b5cf8f9199e54b7abc559e8fec11688 Mon Sep 17 00:00:00 2001 From: Patrick Avery Date: Fri, 10 Nov 2023 13:50:52 -0600 Subject: [PATCH 2/4] Fix applying powder masks Since the `rbnds` are no longer sorted in start/stop pairs, we have to add some logic to handle it better. We don't want to go back to the start/stop pair setup because it was error-prone and would make overlay rendering a little slower. Signed-off-by: Patrick Avery --- hexrdgui/main_window.py | 26 ++++++++++++++++++++++++-- 1 file changed, 24 insertions(+), 2 deletions(-) diff --git a/hexrdgui/main_window.py b/hexrdgui/main_window.py index 98dc4d673..d67295050 100644 --- a/hexrdgui/main_window.py +++ b/hexrdgui/main_window.py @@ -815,8 +815,30 @@ def action_edit_apply_powder_mask_to_polar(self): data = [] for overlay in powder_overlays: for _, val in overlay.data.items(): - a = iter(val['rbnds']) - for start, end in zip(a, a): + # We will only apply masks for ranges that have both a + # start and a stop. + start_end_pairs = {} + for i, indices in enumerate(val['rbnd_indices']): + for idx in indices: + # Get the pair for this HKL index + pairs = start_end_pairs.setdefault(idx, []) + if len(pairs) == 2: + # We already got this one (this shouldn't happen) + continue + + pairs.append(val['rbnds'][i]) + + # We only want to use the ranges once each. + # Since we found a use of this range already, + # just break. + break + + for key in list(start_end_pairs): + # Remove any ranges that have a missing half + if len(start_end_pairs[key]) < 2: + del start_end_pairs[key] + + for start, end in start_end_pairs.values(): ranges = np.append(start, np.flip(end, axis=0), axis=0) ranges = np.append(ranges, [ranges[0]], axis=0) data.append(ranges[~np.isnan(ranges).any(axis=1)]) From 5f367771c4e2a9974339ae60e527bfd90186af09 Mon Sep 17 00:00:00 2001 From: Patrick Avery Date: Fri, 10 Nov 2023 13:53:07 -0600 Subject: [PATCH 3/4] Fix flake8 issues Signed-off-by: Patrick Avery --- hexrdgui/overlays/laue_overlay.py | 1 + hexrdgui/overlays/rotation_series_overlay.py | 1 + 2 files changed, 2 insertions(+) diff --git a/hexrdgui/overlays/laue_overlay.py b/hexrdgui/overlays/laue_overlay.py index d42d880a5..34222a665 100644 --- a/hexrdgui/overlays/laue_overlay.py +++ b/hexrdgui/overlays/laue_overlay.py @@ -463,5 +463,6 @@ class LaueLabelType(str, Enum): hkls = 'hkls' energy = 'energy' + # Constants nans_row = np.nan * np.ones((1, 2)) diff --git a/hexrdgui/overlays/rotation_series_overlay.py b/hexrdgui/overlays/rotation_series_overlay.py index 8831e32d5..2dd797ab9 100644 --- a/hexrdgui/overlays/rotation_series_overlay.py +++ b/hexrdgui/overlays/rotation_series_overlay.py @@ -378,5 +378,6 @@ def sync_omegas(self): HexrdConfig().overlay_config_changed.emit() HexrdConfig().update_overlay_editor.emit() + # Constants nans_row = np.nan * np.ones((1, 2)) From 47481ec9ab0e917677850aa2107b5e9bde31bbf2 Mon Sep 17 00:00:00 2001 From: Patrick Avery Date: Fri, 10 Nov 2023 14:04:23 -0600 Subject: [PATCH 4/4] Apply tvec_s when convert masks to raw This doesn't matter for most examples, which have all zeros for their tvec_s. However, it is important to include this for the examples that do. Signed-off-by: Patrick Avery --- hexrdgui/create_raw_mask.py | 5 +++-- hexrdgui/utils/conversions.py | 6 ++++-- 2 files changed, 7 insertions(+), 4 deletions(-) diff --git a/hexrdgui/create_raw_mask.py b/hexrdgui/create_raw_mask.py index ea185f6aa..fd18eca42 100644 --- a/hexrdgui/create_raw_mask.py +++ b/hexrdgui/create_raw_mask.py @@ -47,9 +47,10 @@ def convert_polar_to_raw(line_data, reverse_tth_distortion=True): line_data[i] = line raw_line_data = [] + instr = create_hedm_instrument() for line in line_data: - for key, panel in create_hedm_instrument().detectors.items(): - raw = angles_to_pixels(line, panel) + for key, panel in instr.detectors.items(): + raw = angles_to_pixels(line, panel, tvec_s=instr.tvec) if all([np.isnan(x) for x in raw.flatten()]): continue diff --git a/hexrdgui/utils/conversions.py b/hexrdgui/utils/conversions.py index facd80a0e..4f5ed9740 100644 --- a/hexrdgui/utils/conversions.py +++ b/hexrdgui/utils/conversions.py @@ -46,8 +46,10 @@ def angles_to_cart(angles, panel, tvec_s=None, tvec_c=None, return panel.angles_to_cart(**kwargs) -def angles_to_pixels(angles, panel): - xys = angles_to_cart(angles, panel) +def angles_to_pixels(angles, panel, tvec_s=None, tvec_c=None, + apply_distortion=True): + xys = angles_to_cart(angles, panel, tvec_s=tvec_s, tvec_c=tvec_c, + apply_distortion=apply_distortion) return cart_to_pixels(xys, panel)