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/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/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)]) diff --git a/hexrdgui/overlays/laue_overlay.py b/hexrdgui/overlays/laue_overlay.py index 8168ff909..34222a665 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,7 @@ 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..2dd797ab9 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,7 @@ 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 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)