From fc7fb172e13d325b229541126aa2cf6a6de7fcda Mon Sep 17 00:00:00 2001 From: serinlee1065 Date: Tue, 2 Sep 2025 13:04:25 -0700 Subject: [PATCH 01/11] Enhance similarity calculation with optional Gaussian smoothing and improve cluster mapping logic --- py4DSTEM/process/utils/cluster.py | 131 +++++++++++++++++------------- 1 file changed, 75 insertions(+), 56 deletions(-) diff --git a/py4DSTEM/process/utils/cluster.py b/py4DSTEM/process/utils/cluster.py index aed0d213c..bb0088880 100644 --- a/py4DSTEM/process/utils/cluster.py +++ b/py4DSTEM/process/utils/cluster.py @@ -28,6 +28,7 @@ def __init__( def find_similarity( self, mask=None, # by default + smooth_sigma = 0, ): # Which neighbors to search # (-1,-1) will be equivalent to (1,1) @@ -54,10 +55,17 @@ def find_similarity( range(self.datacube.shape[0]), range(self.datacube.shape[1]), ): - if mask is None: - diff_ref = self.datacube[rx, ry] - else: - diff_ref = self.datacube[rx, ry][mask] + diff_ref = self.datacube[rx, ry].copy().astype('float') + diff_ref -= diff_ref.mean() + + if smooth_sigma > 0: + diff_ref = gaussian_filter(diff_ref,smooth_sigma) + + if mask is not None: + diff_ref = diff_ref[mask] + + norm_diff_ref = np.sqrt(np.sum(diff_ref * diff_ref)) + # diff_ref_mean = np.mean(diff_ref) # loop over neighbors for ind in range(self.dxy.shape[0]): @@ -69,26 +77,24 @@ def find_similarity( and x_ind < self.datacube.shape[0] and y_ind < self.datacube.shape[1] ): - - if mask is None: - diff = self.datacube[x_ind, y_ind] - else: - diff = self.datacube[x_ind, y_ind][mask] - - # # image self.similarity with mean abs difference - # self.similarity[rx,ry,ind] = np.mean( - # np.abs( - # diff - diff_ref - # ) - # ) - + diff = self.datacube[x_ind, y_ind].copy().astype('float') + diff -= diff.mean() + + if smooth_sigma > 0: + diff = gaussian_filter(diff,smooth_sigma) + + if mask is not None: + diff = diff[mask] + # image self.similarity with normalized corr: cosine self.similarity? self.similarity[rx, ry, ind] = ( np.sum(diff * diff_ref) / np.sqrt(np.sum(diff * diff)) - / np.sqrt(np.sum(diff_ref * diff_ref)) + / norm_diff_ref ) + # self.similarity[rx, ry, ind] = np.mean(np.abs(diff - diff_ref)) / diff_ref_mean + # Create a function to map cluster index to color def get_color(self, cluster_index): colors = [ @@ -108,28 +114,31 @@ def get_color(self, cluster_index): # Find the pixel with the highest self.similarity and start the clustering from there def indexing_clusters_all( self, - mask, + # mask, threshold, ): - self.dxy = np.array( - ( - (-1, -1), - (-1, 0), - (-1, 1), - (0, -1), - (1, 1), - (1, 0), - (1, -1), - (0, 1), - ) - ) + # self.dxy = np.array( + # ( + # (-1, -1), + # (-1, 0), + # (-1, 1), + # (0, -1), + # (1, 1), + # (1, 0), + # (1, -1), + # (0, 1), + # ) + # ) sim_averaged = np.mean(self.similarity, axis=2) # color the pixels with the cluster index # map_cluster = np.zeros((sim_averaged.shape[0],sim_averaged.shape[1])) - self.cluster_map = np.zeros( + self.cluster_map = -1 * np.ones( + (sim_averaged.shape[0], sim_averaged.shape[1]), dtype=np.float64 + ) + self.cluster_map_rgb = np.zeros( (sim_averaged.shape[0], sim_averaged.shape[1], 4), dtype=np.float64 ) @@ -155,8 +164,10 @@ def indexing_clusters_all( ) # map_cluster[rx0, ry0] = cluster_count_ind+1 + self.cluster_map[rx0, ry0] = cluster_count_ind + color = self.get_color(cluster_count_ind + 1) - self.cluster_map[rx0, ry0] = plt.cm.colors.to_rgba(color) + self.cluster_map_rgb[rx0, ry0] = plt.cm.colors.to_rgba(color) # Clustering: one cluster per while loop(until it breaks) # Marching algorithm: find a new position and search the nearest neighbor @@ -178,34 +189,42 @@ def indexing_clusters_all( x_ind = rx0 + self.dxy[ind, 0] y_ind = ry0 + self.dxy[ind, 1] - # add if the neighbor is similar, but don't add if the neighbor is already in a cluster - if self.similarity[ - rx0, ry0, ind - ] > threshold and np.array_equal( - self.cluster_map[x_ind, y_ind], [0, 0, 0, 0] - ): - - cluster_indices = np.append( - cluster_indices, [[x_ind, y_ind]], axis=0 - ) - # self.cluster_map[x_ind, y_ind] = cluster_count_ind+1 - color = self.get_color(cluster_count_ind + 1) - self.cluster_map[x_ind, y_ind] = plt.cm.colors.to_rgba( - color - ) + if x_ind > 1 and \ + y_ind > 1 and \ + x_ind < self.similarity.shape[0] - 2 and \ + y_ind < self.similarity.shape[1] - 2: + + # add if the neighbor is similar, but don't add if the neighbor is already in a cluster + if self.similarity[rx0, ry0, ind] >= threshold \ + and self.cluster_map[x_ind, y_ind] == -1: + + # print(cluster_indices) + # print([[x_ind, y_ind]]) + cluster_indices = np.append( + cluster_indices, [[x_ind, y_ind]], axis=0 + ) + + self.cluster_map[x_ind, y_ind] = cluster_count_ind + + + # self.cluster_map[x_ind, y_ind] = cluster_count_ind+1 + color = self.get_color(cluster_count_ind + 1) + self.cluster_map_rgb[x_ind, y_ind] = plt.cm.colors.to_rgba( + color + ) # if no new pixel is checked for NN then break if counting_added_pixel == 0: break - # single pixel cluster - if cluster_indices.shape[0] == 1: - self.cluster_map[cluster_indices[0, 0], cluster_indices[0, 1]] = [ - 0, - 0, - 0, - 1, - ] + # # single pixel cluster + # if cluster_indices.shape[0] == 1: + # self.cluster_map[cluster_indices[0, 0], cluster_indices[0, 1]] = [ + # 0, + # 0, + # 0, + # 1, + # ] self.cluster_list.append(cluster_indices) cluster_count_ind += 1 From e3e918a534019766500a6baf5538facce635566a Mon Sep 17 00:00:00 2001 From: Serin Lee Date: Thu, 4 Sep 2025 11:56:59 -0700 Subject: [PATCH 02/11] Replacing np.integer with np.int32 --- py4DSTEM/process/diffraction/crystal_ACOM.py | 12 ++++++++---- 1 file changed, 8 insertions(+), 4 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index f7a4a20db..d034c70a8 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -335,7 +335,8 @@ def orientation_plan( ) self.orientation_zone_axis_steps = ( np.round(step / self.orientation_refine_ratio) * self.orientation_refine_ratio - ).astype(np.integer) + ).astype(np.int32) + # ).astype(np.integer) if self.orientation_fiber and self.orientation_fiber_angles[0] == 0: self.orientation_num_zones = int(1) @@ -370,7 +371,8 @@ def orientation_plan( (self.orientation_zone_axis_steps + 1) * (self.orientation_zone_axis_steps + 2) / 2 - ).astype(np.integer) + ).astype(np.int32) + # ).astype(np.integer) self.orientation_vecs = np.zeros((self.orientation_num_zones, 3)) self.orientation_vecs[0, :] = self.orientation_zone_axis_range[0, :] self.orientation_inds = np.zeros((self.orientation_num_zones, 3), dtype="int") @@ -379,7 +381,8 @@ def orientation_plan( # or circular arc SLERP for fiber texture for a0 in np.arange(1, self.orientation_zone_axis_steps + 1): inds = np.arange(a0 * (a0 + 1) / 2, a0 * (a0 + 1) / 2 + a0 + 1).astype( - np.integer + np.int32 + # np.integer ) p0 = pv[a0, :] @@ -617,7 +620,8 @@ def orientation_plan( # Solve for number of angular steps along in-plane rotation direction self.orientation_in_plane_steps = np.round(360 / angle_step_in_plane).astype( - np.integer + np.int32 + # np.integer ) # Calculate -z angles (Euler angle 3) From 2ca59016c3725b83da9d894002f4ae0bd0bb9e38 Mon Sep 17 00:00:00 2001 From: Serin Lee Date: Tue, 23 Sep 2025 10:18:15 -0700 Subject: [PATCH 03/11] updating cluster module and ACOM to use cluster dataset for strain mapping --- py4DSTEM/process/diffraction/crystal_ACOM.py | 117 +++++++++---------- py4DSTEM/process/utils/cluster.py | 103 ++++++++++------ 2 files changed, 125 insertions(+), 95 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index d034c70a8..bb50d7598 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -2212,8 +2212,6 @@ def calculate_strain( deformation tensor which transforms the simulated diffraction pattern into the experimental pattern, for all probe positons. - TODO: add robust fitting? - Parameters ---------- bragg_peaks_array (PointListArray): @@ -2334,71 +2332,72 @@ def calculate_strain( inds_match[a0] = ind_min keep[a0] = True - # Get all paired peaks - qxy = np.vstack((p.data["qx"][keep], p.data["qy"][keep])).T - qxy_ref = np.vstack( - (p_ref.data["qx"][inds_match[keep]], p_ref.data["qy"][inds_match[keep]]) - ).T + if np.sum(keep) >= min_num_peaks: + # Get all paired peaks + qxy = np.vstack((p.data["qx"][keep], p.data["qy"][keep])).T + qxy_ref = np.vstack( + (p_ref.data["qx"][inds_match[keep]], p_ref.data["qy"][inds_match[keep]]) + ).T - # Fit transformation matrix - # Note - not sure about transpose here - # (though it might not matter if rotation isn't included) - if intensity_weighting: - weights = np.sqrt(p.data["intensity"][keep, None]) * 0 + 1 - m = lstsq( - qxy_ref * weights, - qxy * weights, - rcond=None, - )[0].T - else: - m = lstsq( - qxy_ref, - qxy, - rcond=None, - )[0].T - - # Robust fitting - if robust: - for a0 in range(5): - # calculate new weights - qxy_fit = qxy_ref @ m - diff2 = np.sum((qxy_fit - qxy) ** 2, axis=1) - - weights = np.exp( - diff2 / ((-2 * robust_thresh**2) * np.median(diff2)) - )[:, None] - if intensity_weighting: - weights *= np.sqrt(p.data["intensity"][keep, None]) - - # calculate new fits + # Fit transformation matrix + # Note - not sure about transpose here + # (though it might not matter if rotation isn't included) + if intensity_weighting: + weights = np.sqrt(p.data["intensity"][keep, None]) * 0 + 1 m = lstsq( qxy_ref * weights, qxy * weights, rcond=None, )[0].T + else: + m = lstsq( + qxy_ref, + qxy, + rcond=None, + )[0].T - # Set values into the infinitesimal strain matrix - strain_map.get_slice("e_xx").data[rx, ry] = 1 - m[0, 0] - strain_map.get_slice("e_yy").data[rx, ry] = 1 - m[1, 1] - strain_map.get_slice("e_xy").data[rx, ry] = -(m[0, 1] + m[1, 0]) / 2.0 - strain_map.get_slice("theta").data[rx, ry] = (m[0, 1] - m[1, 0]) / 2.0 - - # Add finite rotation from ACOM orientation map. - # I am not sure about the relative signs here. - # Also, maybe I need to add in the mirror operator? - if orientation_map.mirror[rx, ry, 0]: - strain_map.get_slice("theta").data[rx, ry] += ( - orientation_map.angles[rx, ry, 0, 0] - + orientation_map.angles[rx, ry, 0, 2] - ) - else: - strain_map.get_slice("theta").data[rx, ry] -= ( - orientation_map.angles[rx, ry, 0, 0] - + orientation_map.angles[rx, ry, 0, 2] - ) + # Robust fitting + if robust: + for a0 in range(5): + # calculate new weights + qxy_fit = qxy_ref @ m + diff2 = np.sum((qxy_fit - qxy) ** 2, axis=1) + + weights = np.exp( + diff2 / ((-2 * robust_thresh**2) * np.median(diff2)) + )[:, None] + if intensity_weighting: + weights *= np.sqrt(p.data["intensity"][keep, None]) + + # calculate new fits + m = lstsq( + qxy_ref * weights, + qxy * weights, + rcond=None, + )[0].T + + # Set values into the infinitesimal strain matrix + strain_map.get_slice("e_xx").data[rx, ry] = 1 - m[0, 0] + strain_map.get_slice("e_yy").data[rx, ry] = 1 - m[1, 1] + strain_map.get_slice("e_xy").data[rx, ry] = -(m[0, 1] + m[1, 0]) / 2.0 + strain_map.get_slice("theta").data[rx, ry] = (m[0, 1] - m[1, 0]) / 2.0 + + # Add finite rotation from ACOM orientation map. + # I am not sure about the relative signs here. + # Also, maybe I need to add in the mirror operator? + if orientation_map.mirror[rx, ry, 0]: + strain_map.get_slice("theta").data[rx, ry] += ( + orientation_map.angles[rx, ry, 0, 0] + + orientation_map.angles[rx, ry, 0, 2] + ) + else: + strain_map.get_slice("theta").data[rx, ry] -= ( + orientation_map.angles[rx, ry, 0, 0] + + orientation_map.angles[rx, ry, 0, 2] + ) - else: - strain_map.get_slice("mask").data[rx, ry] = 0.0 + else: + strain_map.get_slice("mask").data[rx, ry] = 0.0 if rotation_range is not None: strain_map.get_slice("theta").data[:] = np.mod( diff --git a/py4DSTEM/process/utils/cluster.py b/py4DSTEM/process/utils/cluster.py index bb0088880..17eec9b43 100644 --- a/py4DSTEM/process/utils/cluster.py +++ b/py4DSTEM/process/utils/cluster.py @@ -15,22 +15,49 @@ class Cluster: def __init__( self, datacube, + r_space_mask, ): """ Args: - datacube (py4DSTEM.DataCube): 4D-STEM data - + datacube (py4DSTEM.DataCube): 4D-STEM data + r_space_mask (np.ndarray): Mask in real space to apply background thresholding on the similarity array. """ - self.datacube = datacube + self.r_space_mask = r_space_mask + self.similarity = None + self.similarity_raw = None + + def bg_thresholding(self, r_space_mask,): + self.r_space_mask = np.asarray(r_space_mask) + + # if similarity is already computed, apply the thresholding + if self.similarity_raw is not None: + self.similarity = self._apply_bg_mask(self.similarity_raw) + + def _apply_bg_mask(self, similarity): + if self.r_space_mask is None: + return similarity + return similarity * self.r_space_mask[..., None] + def find_similarity( self, - mask=None, # by default + q_space_mask=None, smooth_sigma = 0, + return_similarity = False ): - # Which neighbors to search + + """ + Args: + q_space_mask : annular boolean q_space_mask to apply on the diffraction patterns + smooth_sigma : sigma for Gaussian smoothing of the diffraction patterns before calculating similarity + return_similarity : if True, return the similarity array + """ + if self.r_space_mask is None: + self.set_mask(r_space_mask) + + # List of neighbors to search # (-1,-1) will be equivalent to (1,1) self.dxy = np.array( ( @@ -61,8 +88,8 @@ def find_similarity( if smooth_sigma > 0: diff_ref = gaussian_filter(diff_ref,smooth_sigma) - if mask is not None: - diff_ref = diff_ref[mask] + if q_space_mask is not None: + diff_ref = diff_ref[q_space_mask] norm_diff_ref = np.sqrt(np.sum(diff_ref * diff_ref)) # diff_ref_mean = np.mean(diff_ref) @@ -83,18 +110,23 @@ def find_similarity( if smooth_sigma > 0: diff = gaussian_filter(diff,smooth_sigma) - if mask is not None: - diff = diff[mask] + if q_space_mask is not None: + diff = diff[q_space_mask] - # image self.similarity with normalized corr: cosine self.similarity? + # image self.similarity with normalized cosine correlation self.similarity[rx, ry, ind] = ( np.sum(diff * diff_ref) / np.sqrt(np.sum(diff * diff)) / norm_diff_ref ) - # self.similarity[rx, ry, ind] = np.mean(np.abs(diff - diff_ref)) / diff_ref_mean + + self.similarity_raw = self.similarity.copy() + self.similarity = self._apply_bg_mask(self.similarity) + if return_similarity: + return self.similarity + # Create a function to map cluster index to color def get_color(self, cluster_index): colors = [ @@ -114,33 +146,27 @@ def get_color(self, cluster_index): # Find the pixel with the highest self.similarity and start the clustering from there def indexing_clusters_all( self, - # mask, threshold, ): - - # self.dxy = np.array( - # ( - # (-1, -1), - # (-1, 0), - # (-1, 1), - # (0, -1), - # (1, 1), - # (1, 0), - # (1, -1), - # (0, 1), - # ) - # ) - + """ + Args: + threshold : threshold for similarity to consider two pixels as part of the same cluster + """ + sim_averaged = np.mean(self.similarity, axis=2) + # Assigning the background as 'counted' + sim_averaged[~self.r_space_mask] = -1.0 + # color the pixels with the cluster index - # map_cluster = np.zeros((sim_averaged.shape[0],sim_averaged.shape[1])) self.cluster_map = -1 * np.ones( (sim_averaged.shape[0], sim_averaged.shape[1]), dtype=np.float64 ) self.cluster_map_rgb = np.zeros( (sim_averaged.shape[0], sim_averaged.shape[1], 4), dtype=np.float64 ) + + self.cluster_map_rgb[..., 3] = 1.0 #start as opaque black # store arrays of cluster_indices in a list self.cluster_list = [] @@ -156,14 +182,18 @@ def indexing_clusters_all( # finding the pixel that has the highest self.similarity among the pixel that hasn't been clustered yet # this will be the 'starting pixel' of a new cluster rx0, ry0 = np.unravel_index(sim_averaged.argmax(), sim_averaged.shape) - # print(rx0, ry0) + + # Guarding to check if the seed is background + if self.r_space_mask is not None and not self.r_space_mask[rx0, ry0]: + sim_averaged[rx0, ry0] = -1 # mark processed so we don't pick it again + continue + cluster_indices = np.empty((0, 2)) cluster_indices = (np.append(cluster_indices, [[rx0, ry0]], axis=0)).astype( np.int32 ) - # map_cluster[rx0, ry0] = cluster_count_ind+1 self.cluster_map[rx0, ry0] = cluster_count_ind color = self.get_color(cluster_count_ind + 1) @@ -182,7 +212,7 @@ def indexing_clusters_all( # counter to check if pixel in the cluster are checked for NN counting_added_pixel += 1 - # set to -1 as its NN will be checked + # set to -1 since now its NN will be checked sim_averaged[rx0, ry0] = -1 for ind in range(self.dxy.shape[0]): @@ -194,12 +224,13 @@ def indexing_clusters_all( x_ind < self.similarity.shape[0] - 2 and \ y_ind < self.similarity.shape[1] - 2: + r_ok = True if self.r_space_mask is None else bool(self.r_space_mask[x_ind, y_ind]) + # add if the neighbor is similar, but don't add if the neighbor is already in a cluster if self.similarity[rx0, ry0, ind] >= threshold \ - and self.cluster_map[x_ind, y_ind] == -1: + and self.cluster_map[x_ind, y_ind] == -1 and r_ok: - # print(cluster_indices) - # print([[x_ind, y_ind]]) + cluster_indices = np.append( cluster_indices, [[x_ind, y_ind]], axis=0 ) @@ -217,9 +248,9 @@ def indexing_clusters_all( if counting_added_pixel == 0: break - # # single pixel cluster + # single pixel cluster # if cluster_indices.shape[0] == 1: - # self.cluster_map[cluster_indices[0, 0], cluster_indices[0, 1]] = [ + # self.cluster_map_rgb[cluster_indices[0, 0], cluster_indices[0, 1]] = [ # 0, # 0, # 0, @@ -229,7 +260,7 @@ def indexing_clusters_all( self.cluster_list.append(cluster_indices) cluster_count_ind += 1 - # return cluster_count_ind, self.cluster_list, map_cluster, sim_averaged + # return cluster_count_ind, self.cluster_list, self.cluster_map, self.cluster_map_rgb def create_cluster_cube( self, From c21eca47b878cb30ea59a972bf7babe74bab8f83 Mon Sep 17 00:00:00 2001 From: smribet Date: Mon, 29 Sep 2025 13:54:58 -0700 Subject: [PATCH 04/11] doc strings update and delete extra code --- py4DSTEM/process/utils/cluster.py | 152 +++++++++++++++++------------- 1 file changed, 86 insertions(+), 66 deletions(-) diff --git a/py4DSTEM/process/utils/cluster.py b/py4DSTEM/process/utils/cluster.py index 17eec9b43..89bc123f1 100644 --- a/py4DSTEM/process/utils/cluster.py +++ b/py4DSTEM/process/utils/cluster.py @@ -8,55 +8,57 @@ class Cluster: """ - Clustering 4D data - + Class for clustering data in 4D-STEM DataCube based on + similarity of neighboring diffraction patterns. """ def __init__( self, datacube, - r_space_mask, + r_space_mask=None, ): """ - Args: - datacube (py4DSTEM.DataCube): 4D-STEM data - r_space_mask (np.ndarray): Mask in real space to apply background thresholding on the similarity array. - + Parameters + ---------- + datacube: DataCube + 4D-STEM data + r_space_mask: np.ndarray + Mask in real space to apply background thresholding on the similarity array. """ self.datacube = datacube self.r_space_mask = r_space_mask self.similarity = None self.similarity_raw = None - def bg_thresholding(self, r_space_mask,): - self.r_space_mask = np.asarray(r_space_mask) - - # if similarity is already computed, apply the thresholding - if self.similarity_raw is not None: - self.similarity = self._apply_bg_mask(self.similarity_raw) - def _apply_bg_mask(self, similarity): if self.r_space_mask is None: return similarity return similarity * self.r_space_mask[..., None] - def find_similarity( - self, - q_space_mask=None, - smooth_sigma = 0, - return_similarity = False + self, q_space_mask=None, smooth_sigma=0, return_similarity=False ): - """ - Args: - q_space_mask : annular boolean q_space_mask to apply on the diffraction patterns - smooth_sigma : sigma for Gaussian smoothing of the diffraction patterns before calculating similarity - return_similarity : if True, return the similarity array + Find similarity to neighboring pixels + + Parameters + ---------- + q_space_mask : np.ndarray, optional + boolean q_space_mask to apply on the diffraction patterns + smooth_sigma : float, optional + sigma for Gaussian smoothing of the diffraction patterns + before calculating similarity + return_similarity : bool, optinal + if True, return the similarity array + + Returns + -------- + similarity: np.ndarray + similarity scores for each pixel """ if self.r_space_mask is None: self.set_mask(r_space_mask) - + # List of neighbors to search # (-1,-1) will be equivalent to (1,1) self.dxy = np.array( @@ -82,12 +84,12 @@ def find_similarity( range(self.datacube.shape[0]), range(self.datacube.shape[1]), ): - diff_ref = self.datacube[rx, ry].copy().astype('float') + diff_ref = self.datacube[rx, ry].copy().astype("float") diff_ref -= diff_ref.mean() if smooth_sigma > 0: - diff_ref = gaussian_filter(diff_ref,smooth_sigma) - + diff_ref = gaussian_filter(diff_ref, smooth_sigma) + if q_space_mask is not None: diff_ref = diff_ref[q_space_mask] @@ -104,15 +106,15 @@ def find_similarity( and x_ind < self.datacube.shape[0] and y_ind < self.datacube.shape[1] ): - diff = self.datacube[x_ind, y_ind].copy().astype('float') + diff = self.datacube[x_ind, y_ind].copy().astype("float") diff -= diff.mean() if smooth_sigma > 0: - diff = gaussian_filter(diff,smooth_sigma) - + diff = gaussian_filter(diff, smooth_sigma) + if q_space_mask is not None: diff = diff[q_space_mask] - + # image self.similarity with normalized cosine correlation self.similarity[rx, ry, ind] = ( np.sum(diff * diff_ref) @@ -120,13 +122,12 @@ def find_similarity( / norm_diff_ref ) - self.similarity_raw = self.similarity.copy() self.similarity = self._apply_bg_mask(self.similarity) if return_similarity: return self.similarity - + # Create a function to map cluster index to color def get_color(self, cluster_index): colors = [ @@ -149,14 +150,19 @@ def indexing_clusters_all( threshold, ): """ - Args: - threshold : threshold for similarity to consider two pixels as part of the same cluster + Index all pixsl in a cluster + + Parameters + ---------- + threshold: float + similarity score threshold to consider pixels as part + of the same cluster """ - + sim_averaged = np.mean(self.similarity, axis=2) - # Assigning the background as 'counted' - sim_averaged[~self.r_space_mask] = -1.0 + # Assigning the background as 'counted' + sim_averaged[~self.r_space_mask] = -1.0 # color the pixels with the cluster index self.cluster_map = -1 * np.ones( @@ -165,8 +171,8 @@ def indexing_clusters_all( self.cluster_map_rgb = np.zeros( (sim_averaged.shape[0], sim_averaged.shape[1], 4), dtype=np.float64 ) - - self.cluster_map_rgb[..., 3] = 1.0 #start as opaque black + + self.cluster_map_rgb[..., 3] = 1.0 # start as opaque black # store arrays of cluster_indices in a list self.cluster_list = [] @@ -182,13 +188,12 @@ def indexing_clusters_all( # finding the pixel that has the highest self.similarity among the pixel that hasn't been clustered yet # this will be the 'starting pixel' of a new cluster rx0, ry0 = np.unravel_index(sim_averaged.argmax(), sim_averaged.shape) - + # Guarding to check if the seed is background if self.r_space_mask is not None and not self.r_space_mask[rx0, ry0]: sim_averaged[rx0, ry0] = -1 # mark processed so we don't pick it again continue - cluster_indices = np.empty((0, 2)) cluster_indices = (np.append(cluster_indices, [[rx0, ry0]], axis=0)).astype( np.int32 @@ -219,54 +224,69 @@ def indexing_clusters_all( x_ind = rx0 + self.dxy[ind, 0] y_ind = ry0 + self.dxy[ind, 1] - if x_ind > 1 and \ - y_ind > 1 and \ - x_ind < self.similarity.shape[0] - 2 and \ - y_ind < self.similarity.shape[1] - 2: + if ( + x_ind > 1 + and y_ind > 1 + and x_ind < self.similarity.shape[0] - 2 + and y_ind < self.similarity.shape[1] - 2 + ): - r_ok = True if self.r_space_mask is None else bool(self.r_space_mask[x_ind, y_ind]) + r_ok = ( + True + if self.r_space_mask is None + else bool(self.r_space_mask[x_ind, y_ind]) + ) # add if the neighbor is similar, but don't add if the neighbor is already in a cluster - if self.similarity[rx0, ry0, ind] >= threshold \ - and self.cluster_map[x_ind, y_ind] == -1 and r_ok: + if ( + self.similarity[rx0, ry0, ind] >= threshold + and self.cluster_map[x_ind, y_ind] == -1 + and r_ok + ): - cluster_indices = np.append( cluster_indices, [[x_ind, y_ind]], axis=0 ) self.cluster_map[x_ind, y_ind] = cluster_count_ind - - # self.cluster_map[x_ind, y_ind] = cluster_count_ind+1 color = self.get_color(cluster_count_ind + 1) - self.cluster_map_rgb[x_ind, y_ind] = plt.cm.colors.to_rgba( - color + self.cluster_map_rgb[x_ind, y_ind] = ( + plt.cm.colors.to_rgba(color) ) # if no new pixel is checked for NN then break if counting_added_pixel == 0: break - # single pixel cluster - # if cluster_indices.shape[0] == 1: - # self.cluster_map_rgb[cluster_indices[0, 0], cluster_indices[0, 1]] = [ - # 0, - # 0, - # 0, - # 1, - # ] - self.cluster_list.append(cluster_indices) cluster_count_ind += 1 - # return cluster_count_ind, self.cluster_list, self.cluster_map, self.cluster_map_rgb - def create_cluster_cube( self, min_cluster_size, return_cluster_datacube=False, ): + """ + Create dataset (N, 1, qx, qy), where N is the number of clusters + that contains diffraction patterns that are averaged across pixels + in each cluster + + Parameters + ---------- + min_cluster_size: int + minimum size for a clsuter to be included in dataset + return_cluster_datacube: bool + if True, returns clustered dataset and list of indicies + of clusters + + Returns + -------- + cluster_cube: np.ndarray + dataset with clsutered diffraction patterns + filtered_cluster_list: list + list of indicies in real space of each pixel of each cluster + """ self.filtered_cluster_list = [ arr for arr in self.cluster_list if arr.shape[0] >= min_cluster_size From da828372c1a4088fc7de8d7c9505f9740dba2ab9 Mon Sep 17 00:00:00 2001 From: serinlee1065 Date: Tue, 28 Oct 2025 12:53:30 -0700 Subject: [PATCH 05/11] Adding lines to check the r_space_mask condition --- py4DSTEM/process/utils/cluster.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/py4DSTEM/process/utils/cluster.py b/py4DSTEM/process/utils/cluster.py index 89bc123f1..535c3d863 100644 --- a/py4DSTEM/process/utils/cluster.py +++ b/py4DSTEM/process/utils/cluster.py @@ -162,7 +162,8 @@ def indexing_clusters_all( sim_averaged = np.mean(self.similarity, axis=2) # Assigning the background as 'counted' - sim_averaged[~self.r_space_mask] = -1.0 + if self.r_space_mask.dtype == bool: + sim_averaged[~self.r_space_mask] = -1.0 # color the pixels with the cluster index self.cluster_map = -1 * np.ones( From e59820ce0d9567eece5d53145579dc41bcafdff9 Mon Sep 17 00:00:00 2001 From: smribet Date: Fri, 31 Oct 2025 13:55:39 -0700 Subject: [PATCH 06/11] a few small changes --- py4DSTEM/process/utils/cluster.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/py4DSTEM/process/utils/cluster.py b/py4DSTEM/process/utils/cluster.py index 535c3d863..8916bccf1 100644 --- a/py4DSTEM/process/utils/cluster.py +++ b/py4DSTEM/process/utils/cluster.py @@ -56,9 +56,6 @@ def find_similarity( similarity: np.ndarray similarity scores for each pixel """ - if self.r_space_mask is None: - self.set_mask(r_space_mask) - # List of neighbors to search # (-1,-1) will be equivalent to (1,1) self.dxy = np.array( @@ -162,7 +159,7 @@ def indexing_clusters_all( sim_averaged = np.mean(self.similarity, axis=2) # Assigning the background as 'counted' - if self.r_space_mask.dtype == bool: + if self.r_space_mask is not None: sim_averaged[~self.r_space_mask] = -1.0 # color the pixels with the cluster index From 67a19e2eb57a144dcf079926224eee8a087ca623 Mon Sep 17 00:00:00 2001 From: serinlee1065 Date: Tue, 24 Mar 2026 08:59:24 -0700 Subject: [PATCH 07/11] Black formatting --- py4DSTEM/__init__.py | 6 ------ py4DSTEM/process/__init__.py | 1 - 2 files changed, 7 deletions(-) diff --git a/py4DSTEM/__init__.py b/py4DSTEM/__init__.py index 0f4490eb2..e26f24c78 100644 --- a/py4DSTEM/__init__.py +++ b/py4DSTEM/__init__.py @@ -25,7 +25,6 @@ from py4DSTEM import io from py4DSTEM.io import import_file, read, save - ### basic data classes # data @@ -40,7 +39,6 @@ # datacube from py4DSTEM.datacube import DataCube, VirtualImage, VirtualDiffraction - ### visualization from py4DSTEM import visualize @@ -64,15 +62,12 @@ # diffraction from py4DSTEM.process.diffraction import Crystal, Orientation - # ptycho from py4DSTEM.process import phase - # polar from py4DSTEM.process.polar import PolarDatacube - # strain from py4DSTEM.process.strain.strain import StrainMap @@ -89,7 +84,6 @@ from py4DSTEM import preprocess from py4DSTEM import process - ### utilities # config diff --git a/py4DSTEM/process/__init__.py b/py4DSTEM/process/__init__.py index 00512f706..a1d43401a 100644 --- a/py4DSTEM/process/__init__.py +++ b/py4DSTEM/process/__init__.py @@ -21,4 +21,3 @@ raise exc from py4DSTEM.process.utils import Cluster - From c630fe42ee3551e1b7002f272c98bf1861817eca Mon Sep 17 00:00:00 2001 From: serinlee1065 Date: Tue, 24 Mar 2026 09:07:44 -0700 Subject: [PATCH 08/11] test_strain blackformatting --- test/test_strain.py | 1 - 1 file changed, 1 deletion(-) diff --git a/test/test_strain.py b/test/test_strain.py index d309c5ffb..1144d1c17 100644 --- a/test/test_strain.py +++ b/test/test_strain.py @@ -3,7 +3,6 @@ from os.path import join from numpy import zeros - # set filepath path = join(py4DSTEM._TESTPATH, "strain/downsample_Si_SiGe_analysis_braggdisks_cal.h5") From 914766de2bf4c36772594911b5ace90467ee2cc1 Mon Sep 17 00:00:00 2001 From: serinlee1065 Date: Tue, 24 Mar 2026 09:22:34 -0700 Subject: [PATCH 09/11] Black formatting on multiple files --- .../diskdetection_parallel_new.py | 1 - py4DSTEM/data/propagating_calibration.py | 1 - py4DSTEM/datacube/virtualimage.py | 1 - py4DSTEM/io/__init__.py | 1 - py4DSTEM/io/google_drive_downloader.py | 1 - .../io/legacy/legacy13/v13_emd_classes/io.py | 1 - .../legacy13/v13_py4dstem_classes/io.py | 1 - py4DSTEM/io/legacy/legacy13/v13_to_14.py | 2 - py4DSTEM/process/diffraction/crystal_ACOM.py | 5 +- py4DSTEM/process/diffraction/crystal_phase.py | 9 +-- py4DSTEM/process/polar/polar_analysis.py | 75 ++++++++++--------- test/gettestdata.py | 1 - test/test_native_io/test_realslice_read.py | 1 - test/test_native_io/test_v0_13.py | 1 - test/test_native_io/test_v0_14.py | 1 - test/test_nonnative_io/test_arina.py | 1 - test/test_nonnative_io/test_dm.py | 1 - 17 files changed, 47 insertions(+), 57 deletions(-) diff --git a/py4DSTEM/braggvectors/diskdetection_parallel_new.py b/py4DSTEM/braggvectors/diskdetection_parallel_new.py index dccc0dd4b..637237a46 100644 --- a/py4DSTEM/braggvectors/diskdetection_parallel_new.py +++ b/py4DSTEM/braggvectors/diskdetection_parallel_new.py @@ -21,7 +21,6 @@ from emdfile import PointListArray, PointList from py4DSTEM.braggvectors.diskdetection import _find_Bragg_disks_single_DP_FK - #### SERIALISERS #### # Define Serialiser # these are functions which allow the hdf5 objects to be passed. May not be required anymore diff --git a/py4DSTEM/data/propagating_calibration.py b/py4DSTEM/data/propagating_calibration.py index 4de0c8d96..774b6ab86 100644 --- a/py4DSTEM/data/propagating_calibration.py +++ b/py4DSTEM/data/propagating_calibration.py @@ -3,7 +3,6 @@ import warnings - # This is the abstract pattern: diff --git a/py4DSTEM/datacube/virtualimage.py b/py4DSTEM/datacube/virtualimage.py index d4fe15241..98c71f3d9 100644 --- a/py4DSTEM/datacube/virtualimage.py +++ b/py4DSTEM/datacube/virtualimage.py @@ -13,7 +13,6 @@ from py4DSTEM.preprocess import get_shifted_ar from py4DSTEM.visualize import show - # Virtual image container class diff --git a/py4DSTEM/io/__init__.py b/py4DSTEM/io/__init__.py index fa7cd099e..a8cdcd700 100644 --- a/py4DSTEM/io/__init__.py +++ b/py4DSTEM/io/__init__.py @@ -3,6 +3,5 @@ from py4DSTEM.io.read import read from py4DSTEM.io.save import save - # google downloader from py4DSTEM.io.google_drive_downloader import gdrive_download, get_sample_file_ids diff --git a/py4DSTEM/io/google_drive_downloader.py b/py4DSTEM/io/google_drive_downloader.py index a50fb2c09..bc41e764d 100644 --- a/py4DSTEM/io/google_drive_downloader.py +++ b/py4DSTEM/io/google_drive_downloader.py @@ -2,7 +2,6 @@ import os import warnings - ### File IDs # single files diff --git a/py4DSTEM/io/legacy/legacy13/v13_emd_classes/io.py b/py4DSTEM/io/legacy/legacy13/v13_emd_classes/io.py index 0452a5854..4cdab267a 100644 --- a/py4DSTEM/io/legacy/legacy13/v13_emd_classes/io.py +++ b/py4DSTEM/io/legacy/legacy13/v13_emd_classes/io.py @@ -6,7 +6,6 @@ from numbers import Number from emdfile import tqdmnd - # Define the EMD group types EMD_group_types = { diff --git a/py4DSTEM/io/legacy/legacy13/v13_py4dstem_classes/io.py b/py4DSTEM/io/legacy/legacy13/v13_py4dstem_classes/io.py index 2556ebe8f..d50c4e872 100644 --- a/py4DSTEM/io/legacy/legacy13/v13_py4dstem_classes/io.py +++ b/py4DSTEM/io/legacy/legacy13/v13_py4dstem_classes/io.py @@ -18,7 +18,6 @@ _read_metadata, ) - # Calibration diff --git a/py4DSTEM/io/legacy/legacy13/v13_to_14.py b/py4DSTEM/io/legacy/legacy13/v13_to_14.py index 650529b22..444293daa 100644 --- a/py4DSTEM/io/legacy/legacy13/v13_to_14.py +++ b/py4DSTEM/io/legacy/legacy13/v13_to_14.py @@ -3,7 +3,6 @@ import numpy as np from emdfile import tqdmnd - # v13 imports from py4DSTEM.io.legacy.legacy13.v13_emd_classes import ( @@ -25,7 +24,6 @@ BraggVectors as BraggVectors13, ) - # v14 imports from emdfile import Root, Metadata, Array, PointList, PointListArray diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index bb50d7598..c1fc46c9e 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -2336,7 +2336,10 @@ def calculate_strain( # Get all paired peaks qxy = np.vstack((p.data["qx"][keep], p.data["qy"][keep])).T qxy_ref = np.vstack( - (p_ref.data["qx"][inds_match[keep]], p_ref.data["qy"][inds_match[keep]]) + ( + p_ref.data["qx"][inds_match[keep]], + p_ref.data["qy"][inds_match[keep]], + ) ).T # Fit transformation matrix diff --git a/py4DSTEM/process/diffraction/crystal_phase.py b/py4DSTEM/process/diffraction/crystal_phase.py index 61a671f16..9edfa24e3 100644 --- a/py4DSTEM/process/diffraction/crystal_phase.py +++ b/py4DSTEM/process/diffraction/crystal_phase.py @@ -102,7 +102,7 @@ def quantify_single_pattern( plot_correlation_radius=False, scale_markers_experiment=40, scale_markers_calculated=200, - max_marker_size = 400, + max_marker_size=400, crystal_inds_plot=None, phase_colors=None, figsize=(10, 7), @@ -629,8 +629,7 @@ def quantify_single_pattern( qx, # s = scale_markers_experiment * intensity, s=np.minimum( - scale_markers_experiment - * bragg_peaks.data["intensity"][keep], + scale_markers_experiment * bragg_peaks.data["intensity"][keep], max_marker_size, ), marker="o", @@ -1230,7 +1229,7 @@ def plot_dominant_phase( self, use_correlation_scores=False, reliability_range=(0.0, 1.0), - normalize_exp_intensity = True, + normalize_exp_intensity=True, sigma=0.0, phase_colors=None, ticks=True, @@ -1319,7 +1318,6 @@ def plot_dominant_phase( # if normalize_exp_intensity: # phase_sig /= self.int_total[None,:,:] - # find highest correlation score for each crystal and match index for a0 in range(self.num_crystals): sub = phase_sig[a0] > phase_corr @@ -1349,7 +1347,6 @@ def plot_dominant_phase( if normalize_exp_intensity: phase_rel /= self.int_total - phase_scale = np.clip( (phase_rel - reliability_range[0]) / (reliability_range[1] - reliability_range[0]), diff --git a/py4DSTEM/process/polar/polar_analysis.py b/py4DSTEM/process/polar/polar_analysis.py index 120fa2386..02b6afeda 100644 --- a/py4DSTEM/process/polar/polar_analysis.py +++ b/py4DSTEM/process/polar/polar_analysis.py @@ -331,7 +331,7 @@ def calculate_pair_dist_function( k_width = np.array(k_width) if k_width.size == 1: - k_width = k_width*np.ones(2) + k_width = k_width * np.ones(2) # set up coordinates and scaling k = self.qq @@ -350,35 +350,44 @@ def calculate_pair_dist_function( # sub_fit = np.logical_and( # k >= k_min - # Calculate structure factor mask + # Calculate structure factor mask if k_max is None: k_max = np.max(k) - mask_low = np.sin( - np.clip( - (k - k_min) / k_width[0], - 0, - 1, - ) * np.pi/2.0, - )**2 - mask_high = np.sin( - np.clip( - (k_max - k) / k_width[1], - 0, - 1, - ) * np.pi/2.0, - )**2 + mask_low = ( + np.sin( + np.clip( + (k - k_min) / k_width[0], + 0, + 1, + ) + * np.pi + / 2.0, + ) + ** 2 + ) + mask_high = ( + np.sin( + np.clip( + (k_max - k) / k_width[1], + 0, + 1, + ) + * np.pi + / 2.0, + ) + ** 2 + ) mask = mask_low * mask_high # weighting function for fitting atomic scattering factors weights_fit = np.divide( 1, mask_low, - where = mask_low > 1e-4, + where=mask_low > 1e-4, ) weights_fit[mask_low <= 1e-4] = np.inf # Scale weighting to favour high k values - weights_fit *= ((k[-1] - 0.9*k + dk)) - + weights_fit *= k[-1] - 0.9 * k + dk # fig,ax = plt.subplots() # ax.plot( @@ -391,7 +400,6 @@ def calculate_pair_dist_function( # mask_high, # ) - # initial guesses for background coefs const_bg = np.min(Ik) / int_mean int0 = np.median(Ik) / int_mean - const_bg @@ -440,7 +448,7 @@ def calculate_pair_dist_function( # fk = scattering_model(k2, coefs_fk) bg = scattering_model(k2, coefs) fk = bg - coefs[0] - + # Estimate the reduced structure factor S(k) Sk = (Ik - bg) * k / fk @@ -486,20 +494,19 @@ def calculate_pair_dist_function( # m = pdf_thresh / r_thresh # pdf_reduced[r < r_thresh] = r[r < r_thresh] * m - # pdf_reduced[r < r[ind_thresh]] = m * r[ind_thresh] + # pdf_reduced[r < r[ind_thresh]] = m * r[ind_thresh] - # r_ind_max = r[ind_max] - # r_mask = np.minimum(r / r_ind_max, 1.0) - # r_mask = np.sin(r_mask * np.pi / 2) ** 2 - # pdf_reduced *= r_mask + # r_ind_max = r[ind_max] + # r_mask = np.minimum(r / r_ind_max, 1.0) + # r_mask = np.sin(r_mask * np.pi / 2) ** 2 + # pdf_reduced *= r_mask - - # original version - # ind_max = np.argmax(pdf_reduced) - # r_ind_max = r[ind_max] - # r_mask = np.minimum(r / r_ind_max, 1.0) - # r_mask = np.sin(r_mask * np.pi / 2) ** 2 - # pdf_reduced *= r_mask + # original version + # ind_max = np.argmax(pdf_reduced) + # r_ind_max = r[ind_max] + # r_mask = np.minimum(r / r_ind_max, 1.0) + # r_mask = np.sin(r_mask * np.pi / 2) ** 2 + # pdf_reduced *= r_mask # Store results self.pdf_r = r @@ -517,7 +524,6 @@ def calculate_pair_dist_function( pdf[1:] /= 4 * np.pi * density * r[1:] * (r[1] - r[0]) pdf += 1 - # Damp the unphysical fluctuations at the PDF origin if damp_origin_fluctuations: # Find radial value of primary peak @@ -536,7 +542,6 @@ def calculate_pair_dist_function( m = pdf_thresh / r_thresh pdf[r < r_thresh] = r[r < r_thresh] * m - # damp and clip values below zero # if damp_origin_fluctuations: # pdf *= r_mask diff --git a/test/gettestdata.py b/test/gettestdata.py index 0cb6cb964..376b5737b 100644 --- a/test/gettestdata.py +++ b/test/gettestdata.py @@ -8,7 +8,6 @@ from py4DSTEM import _TESTPATH as testpath from py4DSTEM.io import gdrive_download as download - # Make the argument parser parser = argparse.ArgumentParser( description="A command line tool for downloading data to run the py4DSTEM test suite" diff --git a/test/test_native_io/test_realslice_read.py b/test/test_native_io/test_realslice_read.py index eeee5217d..18727129c 100644 --- a/test/test_native_io/test_realslice_read.py +++ b/test/test_native_io/test_realslice_read.py @@ -4,7 +4,6 @@ import py4DSTEM from os.path import join - # Set filepaths filepath = join(py4DSTEM._TESTPATH, "test_io/test_realslice_io.h5") diff --git a/test/test_native_io/test_v0_13.py b/test/test_native_io/test_v0_13.py index e1d91bba4..c9e851f72 100644 --- a/test/test_native_io/test_v0_13.py +++ b/test/test_native_io/test_v0_13.py @@ -1,7 +1,6 @@ from py4DSTEM import read, print_h5_tree, _TESTPATH from os.path import join - # Set filepaths filepath = join(_TESTPATH, "test_io/legacy_v0.13.h5") diff --git a/test/test_native_io/test_v0_14.py b/test/test_native_io/test_v0_14.py index c71638c8e..77ca8ca7e 100644 --- a/test/test_native_io/test_v0_14.py +++ b/test/test_native_io/test_v0_14.py @@ -1,7 +1,6 @@ import py4DSTEM from os.path import join, exists - path = join(py4DSTEM._TESTPATH, "test_io/legacy_v0.14.h5") diff --git a/test/test_nonnative_io/test_arina.py b/test/test_nonnative_io/test_arina.py index c02964cf8..556fd1a13 100644 --- a/test/test_nonnative_io/test_arina.py +++ b/test/test_nonnative_io/test_arina.py @@ -2,7 +2,6 @@ import emdfile from os.path import join - # Set filepaths filepath = join(py4DSTEM._TESTPATH, "test_arina/STO_STEM_bench_20us_master.h5") diff --git a/test/test_nonnative_io/test_dm.py b/test/test_nonnative_io/test_dm.py index ee6f1b2eb..82964d547 100644 --- a/test/test_nonnative_io/test_dm.py +++ b/test/test_nonnative_io/test_dm.py @@ -2,7 +2,6 @@ import emdfile from os.path import join - # Set filepaths filepath_dm4_datacube = join(py4DSTEM._TESTPATH, "small_datacube.dm4") filepath_dm3_3Dstack = join(py4DSTEM._TESTPATH, "test_io/small_dm3_3Dstack.dm3") From fc1f9f36d8db083f6d89f056e62f684f3e6aae10 Mon Sep 17 00:00:00 2001 From: serinlee1065 Date: Tue, 24 Mar 2026 09:22:52 -0700 Subject: [PATCH 10/11] Adding histogram to crystal_viz and black formatting --- py4DSTEM/process/diffraction/crystal_viz.py | 656 +++++++++++++------- 1 file changed, 431 insertions(+), 225 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_viz.py b/py4DSTEM/process/diffraction/crystal_viz.py index b88fbb35b..dd6ecb066 100644 --- a/py4DSTEM/process/diffraction/crystal_viz.py +++ b/py4DSTEM/process/diffraction/crystal_viz.py @@ -1115,12 +1115,14 @@ def plot_orientation_maps( self, orientation_map=None, orientation_ind: int = 0, - mask = None, + mask=None, dir_in_plane_degrees: float = 0.0, corr_range: np.ndarray = np.array([0, 5]), corr_normalize: bool = True, show_legend: bool = True, scale_legend: bool = None, + legend_hist: bool = False, + scale_hist: float = 2, figsize: Union[list, tuple, np.ndarray] = (16, 5), figbound: Union[list, tuple, np.ndarray] = (0.01, 0.005), show_axes: bool = True, @@ -1128,6 +1130,8 @@ def plot_orientation_maps( plot_limit=None, plot_layout=0, swap_axes_xy_limits=False, + crop_range=None, + figax=None, returnfig: bool = False, progress_bar=False, ): @@ -1144,6 +1148,8 @@ def plot_orientation_maps( corr_normalize (bool): If true, set mean correlation to 1. show_legend (bool): Show the legend scale_legend (float): 2 elements, x and y scaling of legend panel + legend_hist(bool): If True, adds histogram of detected orientations to legend + scale_hist (float): Amount to scale histogram values by figsize (array): 2 elements defining figure size figbound (array): 2 elements defining figure boundary show_axes (bool): Flag setting whether orienation map axes are visible. @@ -1152,6 +1158,8 @@ def plot_orientation_maps( plot_layout (int): subplot layout: 0 - 1 row, 3 col 1 - 3 row, 1 col swap_axes_xy_limits (bool): swap x and y boundaries for legend (not sure why we need this in some cases) + crop_range (4-tuple): region to crop for making orientation maps + figax (matplotlib figure): figure and axes for plotting (3 axes needed) returnfig (bool): set to True to return figure and axes handles progress_bar (bool): Enable progressbar when calculating orientation images. @@ -1217,17 +1225,17 @@ def plot_orientation_maps( dir_in_plane = np.deg2rad(dir_in_plane_degrees) ct = np.cos(dir_in_plane) st = np.sin(dir_in_plane) - + rgb_x = np.zeros((orientation_map.num_x, orientation_map.num_y, 3)) rgb_z = np.zeros((orientation_map.num_x, orientation_map.num_y, 3)) - if self.pointgroup.get_crystal_system() == 'monoclinic': + if self.pointgroup.get_crystal_system() == "monoclinic": basis_x = np.zeros((orientation_map.num_x, orientation_map.num_y, 6)) # basis_y = np.zeros((orientation_map.num_x, orientation_map.num_y, 4)) basis_z = np.zeros((orientation_map.num_x, orientation_map.num_y, 6)) # Basis for fitting orientation projections A = np.linalg.inv(self.orientation_zone_axis_range).T - A = np.hstack((A,-A)) + A = np.hstack((A, -A)) # Testing # print(np.round(self.orientation_zone_axis_range)) @@ -1251,6 +1259,8 @@ def plot_orientation_maps( mask = (corr - corr_range[0]) / (corr_range[1] - corr_range[0]) mask = np.clip(mask, 0, 1) + self.mask_corr = mask + # Generate images for rx, ry in tqdmnd( orientation_map.num_x, @@ -1260,15 +1270,16 @@ def plot_orientation_maps( disable=not progress_bar, ): - if self.pointgroup.get_crystal_system() == 'monoclinic': - dir_x = orientation_map.matrix[rx, ry, orientation_ind, :, 0] * ct \ + if self.pointgroup.get_crystal_system() == "monoclinic": + dir_x = ( + orientation_map.matrix[rx, ry, orientation_ind, :, 0] * ct + orientation_map.matrix[rx, ry, orientation_ind, :, 1] * st - dir_z = orientation_map.matrix[rx, ry, orientation_ind, :, 2] - basis_x[rx,ry,:] = nnls(A,dir_x)[0] - basis_z[rx,ry,:] = nnls(A,dir_z)[0] - + ) + dir_z = orientation_map.matrix[rx, ry, orientation_ind, :, 2] + basis_x[rx, ry, :] = nnls(A, dir_x)[0] + basis_z[rx, ry, :] = nnls(A, dir_z)[0] - else: + else: if self.pymatgen_available: basis_x[rx, ry, :] = ( @@ -1288,17 +1299,20 @@ def plot_orientation_maps( ) basis_x = np.clip(basis_x, 0, 1) basis_z = np.clip(basis_z, 0, 1) + self.basis_z = basis_z + self.A = A # Convert to RGB images basis_x_max = np.max(basis_x, axis=2) sub = basis_x_max > 0 basis_x_scale = basis_x * mask[:, :, None] - if self.pointgroup.get_crystal_system() == 'monoclinic': + basis_x_scale_binary = basis_x * (mask > 0)[:, :, None] + if self.pointgroup.get_crystal_system() == "monoclinic": rgb_x = ( - basis_x_scale[:, :, 0][:, :, None] * np.array((1,1,1))[None, None, :] + basis_x_scale[:, :, 0][:, :, None] * np.array((1, 1, 1))[None, None, :] + basis_x_scale[:, :, 1][:, :, None] * color_basis[2, :][None, None, :] + basis_x_scale[:, :, 2][:, :, None] * color_basis[1, :][None, None, :] - + basis_x_scale[:, :, 3][:, :, None] * np.array((1,1,1))[None, None, :] + + basis_x_scale[:, :, 3][:, :, None] * np.array((1, 1, 1))[None, None, :] + basis_x_scale[:, :, 4][:, :, None] * color_basis[2, :][None, None, :] + basis_x_scale[:, :, 5][:, :, None] * color_basis[0, :][None, None, :] ) @@ -1306,6 +1320,8 @@ def plot_orientation_maps( for a0 in range(3): basis_x_scale[:, :, a0][sub] /= basis_x_max[sub] basis_x_scale[:, :, a0][np.logical_not(sub)] = 0 + basis_x_scale_binary[:, :, a0][sub] /= basis_x_max[sub] + basis_x_scale_binary[:, :, a0][np.logical_not(sub)] = 0 rgb_x = ( basis_x_scale[:, :, 0][:, :, None] * color_basis[0, :][None, None, :] + basis_x_scale[:, :, 1][:, :, None] * color_basis[1, :][None, None, :] @@ -1315,12 +1331,13 @@ def plot_orientation_maps( basis_z_max = np.max(basis_z, axis=2) sub = basis_z_max > 0 basis_z_scale = basis_z * mask[:, :, None] - if self.pointgroup.get_crystal_system() == 'monoclinic': + basis_z_scale_binary = basis_z * (mask > 0)[:, :, None] + if self.pointgroup.get_crystal_system() == "monoclinic": rgb_z = ( - basis_z_scale[:, :, 0][:, :, None] * np.array((1,1,1))[None, None, :] + basis_z_scale[:, :, 0][:, :, None] * np.array((1, 1, 1))[None, None, :] + basis_z_scale[:, :, 1][:, :, None] * color_basis[2, :][None, None, :] + basis_z_scale[:, :, 2][:, :, None] * color_basis[1, :][None, None, :] - + basis_z_scale[:, :, 3][:, :, None] * np.array((1,1,1))[None, None, :] + + basis_z_scale[:, :, 3][:, :, None] * np.array((1, 1, 1))[None, None, :] + basis_z_scale[:, :, 4][:, :, None] * color_basis[2, :][None, None, :] + basis_z_scale[:, :, 5][:, :, None] * color_basis[0, :][None, None, :] ) @@ -1328,14 +1345,17 @@ def plot_orientation_maps( for a0 in range(3): basis_z_scale[:, :, a0][sub] /= basis_z_max[sub] basis_z_scale[:, :, a0][np.logical_not(sub)] = 0 + basis_z_scale_binary[:, :, a0][sub] /= basis_z_max[sub] + basis_z_scale_binary[:, :, a0][np.logical_not(sub)] = 0 rgb_z = ( basis_z_scale[:, :, 0][:, :, None] * color_basis[0, :][None, None, :] + basis_z_scale[:, :, 1][:, :, None] * color_basis[1, :][None, None, :] + basis_z_scale[:, :, 2][:, :, None] * color_basis[2, :][None, None, :] ) - rgb_x = np.clip(rgb_x,0,1) - rgb_z = np.clip(rgb_z,0,1) + rgb_x = np.clip(rgb_x, 0, 1) + rgb_z = np.clip(rgb_z, 0, 1) + self.basis_z_scale = basis_z_scale # if np.abs(self.cell[4] - 120.0) < 1e-6 or np.abs(self.cell[5] - 120.0) or np.abs(self.cell[6] - 120.0): # label_0 = self.rational_ind( @@ -1377,39 +1397,34 @@ def plot_orientation_maps( # Legend coordinates # TODO - potentially replace approx. coordinates with stereographic projection power_radial = 0.9 - num_points = 90+1 - r = np.linspace(0,1,num_points) - xa,ya = np.meshgrid(np.flip(r),r,indexing='ij') - if self.pointgroup.get_crystal_system() == 'monoclinic': - xa = np.hstack(( - xa, - xa[:,1:] - )) - ya = np.hstack(( - ya - 1, - ya[:,1:] - )) + num_points = 90 + 1 + r = np.linspace(0, 1, num_points) + xa, ya = np.meshgrid(np.flip(r), r, indexing="ij") + if self.pointgroup.get_crystal_system() == "monoclinic": + xa = np.hstack((xa, xa[:, 1:])) + ya = np.hstack((ya - 1, ya[:, 1:])) ra = np.sqrt(xa**2 + ya**2) rmask = np.clip( - (1-ra)/(r[1]-r[0]) + 0.5, + (1 - ra) / (r[1] - r[0]) + 0.5, 0, 1, ) # angular range to span - assume zone axis range vectors are normalized - ta = np.arctan2(xa,ya) # Note theta direction is opposite from standard - if self.pointgroup.get_crystal_system() == 'monoclinic': - tspan = np.pi/2 + ta = np.arctan2(xa, ya) # Note theta direction is opposite from standard + if self.pointgroup.get_crystal_system() == "monoclinic": + tspan = np.pi / 2 tmask = np.ones_like(rmask) else: tspan = np.arccos( np.clip( - self.orientation_zone_axis_range[1] @ self.orientation_zone_axis_range[2], + self.orientation_zone_axis_range[1] + @ self.orientation_zone_axis_range[2], 0, 1, ) ) tmask = np.clip( - (tspan - ta)/(r[1]-r[0]) + 0.5, + (tspan - ta) / (r[1] - r[0]) + 0.5, 0, 1, ) @@ -1417,82 +1432,86 @@ def plot_orientation_maps( rscale = ra**power_radial mask_leg = tmask * rmask - if self.pointgroup.get_crystal_system() == 'monoclinic': + if self.pointgroup.get_crystal_system() == "monoclinic": scale = 1.5 # Right side coloring - weight_0 = np.clip(1-rscale, 0, 1) + weight_0 = np.clip(1 - rscale, 0, 1) weight_1 = np.clip(ya, 0, 1) weight_2 = np.clip(xa, 0, 1) weight_total = np.maximum(weight_0 + weight_1 + weight_2, 0) + 1e-8 weight_0 /= weight_total weight_1 /= weight_total weight_2 /= weight_total - weight_0 = np.minimum(weight_0,1) - weight_1 = np.minimum(weight_1,1) - weight_2 = np.minimum(weight_2,1) - weight_max = (weight_0**scale + weight_1**scale + weight_2**scale)**(1/scale) + 1e-8 + weight_0 = np.minimum(weight_0, 1) + weight_1 = np.minimum(weight_1, 1) + weight_2 = np.minimum(weight_2, 1) + weight_max = (weight_0**scale + weight_1**scale + weight_2**scale) ** ( + 1 / scale + ) + 1e-8 weight_0 /= weight_max weight_1 /= weight_max weight_2 /= weight_max rgb_leg_right = ( - weight_0[:, :, None] * np.array((1.0,1.0,1.0))[None, None, :] + weight_0[:, :, None] * np.array((1.0, 1.0, 1.0))[None, None, :] + weight_1[:, :, None] * color_basis[1, :][None, None, :] + weight_2[:, :, None] * color_basis[2, :][None, None, :] - ) + ) # left side coloring - weight_0 = np.clip(1-rscale, 0, 1) + weight_0 = np.clip(1 - rscale, 0, 1) weight_1 = np.clip(rscale * (ta - tspan) / tspan, 0, 1) weight_2 = np.clip(xa, 0, 1) weight_total = np.maximum(weight_0 + weight_1 + weight_2, 0) + 1e-8 weight_0 /= weight_total weight_1 /= weight_total weight_2 /= weight_total - weight_0 = np.minimum(weight_0,1) - weight_1 = np.minimum(weight_1,1) - weight_2 = np.minimum(weight_2,1) - weight_max = (weight_0**scale + weight_1**scale + weight_2**scale)**(1/scale) + 1e-8 + weight_0 = np.minimum(weight_0, 1) + weight_1 = np.minimum(weight_1, 1) + weight_2 = np.minimum(weight_2, 1) + weight_max = (weight_0**scale + weight_1**scale + weight_2**scale) ** ( + 1 / scale + ) + 1e-8 weight_0 /= weight_max weight_1 /= weight_max weight_2 /= weight_max rgb_leg_left = ( - weight_0[:, :, None] * np.array((1.0,1.0,1.0))[None, None, :] + weight_0[:, :, None] * np.array((1.0, 1.0, 1.0))[None, None, :] + weight_1[:, :, None] * color_basis[0, :][None, None, :] + weight_2[:, :, None] * color_basis[2, :][None, None, :] - ) + ) # combined legend - rgb_leg = np.zeros((num_points,2*num_points-1,3)) + rgb_leg = np.zeros((num_points, 2 * num_points - 1, 3)) sub_right = ya >= 0 sub_left = np.logical_not(sub_right) for a0 in range(3): - rgb_leg[:,:,a0][sub_right] = rgb_leg_right[:,:,a0][sub_right] - rgb_leg[:,:,a0][sub_left] = rgb_leg_left[:,:,a0][sub_left] - + rgb_leg[:, :, a0][sub_right] = rgb_leg_right[:, :, a0][sub_right] + rgb_leg[:, :, a0][sub_left] = rgb_leg_left[:, :, a0][sub_left] + else: - weight_0 = np.maximum(1-rscale, 0) + weight_0 = np.maximum(1 - rscale, 0) weight_1 = np.maximum(rscale * (tspan - ta) / tspan, 0) weight_2 = np.maximum(rscale * ta / tspan, 0) weight_total = np.maximum(weight_0 + weight_1 + weight_2, 0) + 1e-8 weight_0 /= weight_total weight_1 /= weight_total weight_2 /= weight_total - weight_0 = np.minimum(weight_0,1) - weight_1 = np.minimum(weight_1,1) - weight_2 = np.minimum(weight_2,1) + weight_0 = np.minimum(weight_0, 1) + weight_1 = np.minimum(weight_1, 1) + weight_2 = np.minimum(weight_2, 1) # Generate rgb legend image - hkl = weight_0[:,:,None] * self.orientation_zone_axis_range[0][None,None] \ - + weight_1[:,:,None] * self.orientation_zone_axis_range[1][None,None] \ - + weight_2[:,:,None] * self.orientation_zone_axis_range[2][None,None] - hkl /= np.linalg.norm(hkl,axis=2)[:,:,None] - basis_leg = np.zeros((num_points,num_points,3)) + hkl = ( + weight_0[:, :, None] * self.orientation_zone_axis_range[0][None, None] + + weight_1[:, :, None] * self.orientation_zone_axis_range[1][None, None] + + weight_2[:, :, None] * self.orientation_zone_axis_range[2][None, None] + ) + hkl /= np.linalg.norm(hkl, axis=2)[:, :, None] + basis_leg = np.zeros((num_points, num_points, 3)) for rx in range(num_points): for ry in range(num_points): - basis_leg[rx, ry, :] = ( - A @ hkl[rx,ry,:] - ) + basis_leg[rx, ry, :] = A @ hkl[rx, ry, :] basis_leg_max = np.max(basis_leg, axis=2) basis_leg_scale = basis_leg for a0 in range(3): @@ -1501,15 +1520,14 @@ def plot_orientation_maps( basis_leg_scale[:, :, 0][:, :, None] * color_basis[0, :][None, None, :] + basis_leg_scale[:, :, 1][:, :, None] * color_basis[1, :][None, None, :] + basis_leg_scale[:, :, 2][:, :, None] * color_basis[2, :][None, None, :] - ) + ) rgb_leg = np.clip( - rgb_leg * mask_leg[:,:,None] + (1-mask_leg[:,:,None]), + rgb_leg * mask_leg[:, :, None] + (1 - mask_leg[:, :, None]), 0, 1, ) - # rgb_leg = weight_0 # w_scale = np.maximum(np.maximum(weight_0, w1), w2) # w_scale = 1 - np.exp(-w_scale) @@ -1531,7 +1549,6 @@ def plot_orientation_maps( # ax[2].imshow(weight_2) # ax[3].imshow(im_leg) - # # projection vector # cam_dir = np.mean(self.orientation_zone_axis_range, axis=0) # cam_dir = cam_dir / np.linalg.norm(cam_dir) @@ -1560,26 +1577,35 @@ def plot_orientation_maps( # plotting frame # fig, ax = plt.subplots(1, 3, figsize=figsize) - fig = plt.figure(figsize=figsize) - if plot_layout == 0: - ax_x = fig.add_axes([0.0 + figbound[0], 0.0, 0.4 - 2 * +figbound[0], 1.0]) - ax_z = fig.add_axes([0.4 + figbound[0], 0.0, 0.4 - 2 * +figbound[0], 1.0]) - ax_l = fig.add_axes( - [0.8 + figbound[0], 0.0, 0.2 - 2 * +figbound[0], 1.0], - # projection="3d", - # elev=el, - # azim=az, - ) - elif plot_layout == 1: - ax_x = fig.add_axes([0.0, 0.0 + figbound[0], 1.0, 0.4 - 2 * +figbound[0]]) - ax_z = fig.add_axes([0.0, 0.4 + figbound[0], 1.0, 0.4 - 2 * +figbound[0]]) - ax_l = fig.add_axes( - [0.0, 0.8 + figbound[0], 1.0, 0.2 - 2 * +figbound[0]], - # projection="3d", - # elev=el, - # azim=az, - ) - + if figax is None: + fig = plt.figure(figsize=figsize) + if plot_layout == 0: + ax_x = fig.add_axes([0.0 + figbound[0], 0.0, 0.4 - 2 * +figbound[0], 1.0]) + ax_z = fig.add_axes([0.4 + figbound[0], 0.0, 0.4 - 2 * +figbound[0], 1.0]) + ax_l = fig.add_axes( + [0.8 + figbound[0], 0.0, 0.2 - 2 * +figbound[0], 1.0], + # projection="3d", + # elev=el, + # azim=az, + ) + elif plot_layout == 1: + ax_x = fig.add_axes([0.0, 0.0 + figbound[0], 1.0, 0.4 - 2 * +figbound[0]]) + ax_z = fig.add_axes([0.0, 0.4 + figbound[0], 1.0, 0.4 - 2 * +figbound[0]]) + ax_l = fig.add_axes( + [0.0, 0.8 + figbound[0], 1.0, 0.2 - 2 * +figbound[0]], + # projection="3d", + # elev=el, + # azim=az, + ) + else: + fig = figax[0] + ax_x = figax[1] + ax_z = figax[2] + ax_l = figax[3] + + if crop_range is not None: + rgb_x = rgb_x[crop_range[0] : crop_range[1], crop_range[2] : crop_range[3]] + rgb_z = rgb_z[crop_range[0] : crop_range[1], crop_range[2] : crop_range[3]] # orientation images if self.pymatgen_available: ax_x.imshow(rgb_x) @@ -1617,174 +1643,355 @@ def plot_orientation_maps( # Legend if show_legend: - ax_l.imshow(rgb_leg) - + if legend_hist is True: + A = np.array([0, 0]) + B = np.array([1, 0]) + C = np.array([0.5, np.sqrt(3) / 2]) + vertices = np.array([A, B, C]) + ax_l.scatter(vertices[:, 0], vertices[:, 1], c=color_basis) + points = [] + colors = [] + + for a0 in range(basis_z_scale.shape[0]): + for a1 in range(basis_z_scale.shape[1]): + if mask[a0, a1] > 0: + w = basis_z_scale_binary[a0, a1] + rgb_z = ( + w * color_basis[0, :] + + w * color_basis[1, :] + + w * color_basis[2, :] + ) + colors.append(rgb_z) + w = w / np.sum(w) + pt = w[0] * A + w[1] * B + w[2] * C + points.append(pt) - # Add text labels - text_scale_pos = 0.1 - text_params = { - "va": "center", - "family": "sans-serif", - "fontweight": "normal", - "color": "k", - "size": 12, - } - format_labels = "{0:.2g}" + points = np.asarray(points) + colors = np.asarray(colors) + points_and_colors = np.hstack([points, colors]) + points_and_colors_unique, counts = np.unique( + points_and_colors, axis=0, return_counts=True + ) - bound = num_points*0.25 - shift = num_points*0.10 - if self.pointgroup.get_crystal_system() == 'monoclinic': - p0 = np.array((1,1))*num_points - p1 = np.array((0,1))*num_points - p2 = np.array((1,2))*num_points - p3 = np.array((1,0))*num_points - else: - p0 = np.array((1,0))*num_points - p1 = np.array((1,1))*num_points - p2 = np.array((1-np.cos(np.pi/2-tspan),np.sin(np.pi/2-tspan)))*num_points + ax_l.scatter( + points_and_colors_unique[:, 0], + points_and_colors_unique[:, 1], + c=points_and_colors_unique[:, 2:], + s=counts * scale_hist, + ) + v0 = self.orientation_zone_axis_range[0, :].copy() + v0 /= np.max(np.abs(v0)) + v0 = np.round(v0, 2) + + v1 = self.orientation_zone_axis_range[1, :].copy() + v1 /= np.max(np.abs(v1)) + v1 = np.round(v1, 2) + + v2 = self.orientation_zone_axis_range[2, :].copy() + v2 /= np.max(np.abs(v2)) + v2 = np.round(v2, 2) + + text_params = { + "va": "center", + "family": "sans-serif", + "fontweight": "normal", + "color": "k", + "size": 12, + } + shift = 0.15 + ax_l.text( + A[0], + A[1] - shift, + v0, + ha="center", + **text_params, + ) - if self.pointgroup.get_crystal_system() == 'monoclinic': - v = np.double(self.orientation_zone_axis_range[0,:]).copy() - v /= np.max(np.abs(v)) - v = np.round(v,2) ax_l.text( - p0[1], - p0[0]+shift, - "[" - + format_labels.format(v[0]) - + " " - + format_labels.format(v[1]) - + " " - + format_labels.format(v[2]) - + "]", - None, - zorder=11, + B[0], + B[1] - shift, + v1, ha="center", **text_params, ) - v = np.double(self.orientation_zone_axis_range[1,:]).copy() - v /= np.max(np.abs(v)) - v = np.round(v,2) + ax_l.text( - p1[1], - p1[0]-shift, - "[" - + format_labels.format(v[0]) - + " " - + format_labels.format(v[1]) - + " " - + format_labels.format(v[2]) - + "]", - None, - zorder=11, + C[0] + shift * 2, + C[1], + v2, ha="center", **text_params, ) - v = np.double(self.orientation_zone_axis_range[2,:]).copy() - v /= np.max(np.abs(v)) - v = np.round(v,2) + ax_l.text( - p2[1], - p2[0]+shift, - "[" - + format_labels.format(v[0]) - + " " - + format_labels.format(v[1]) - + " " - + format_labels.format(v[2]) - + "]", - None, - zorder=11, + C[0], + C[1] + 0.1, + "out-of-plane orientation", ha="center", **text_params, ) - v = -1*np.double(self.orientation_zone_axis_range[2,:].copy()) + 0 - v /= np.max(np.abs(v)) - v = np.round(v,2) + + A = np.array([0, 1.5]) + B = np.array([1, 1.5]) + C = np.array([0.5, np.sqrt(3) / 2 + 1.5]) + vertices = np.array([A, B, C]) + points = [] + colors = [] + ax_l.scatter(vertices[:, 0], vertices[:, 1], c=color_basis) + + for a0 in range(basis_x_scale_binary.shape[0]): + for a1 in range(basis_x_scale_binary.shape[1]): + if mask[a0, a1] > 0: + w = basis_x_scale_binary[a0, a1] + rgb_z = ( + w * color_basis[0, :] + + w * color_basis[1, :] + + w * color_basis[2, :] + ) + colors.append(rgb_z) + w = w / np.sum(w) + pt = w[0] * A + w[1] * B + w[2] * C + points.append(pt) + + points = np.asarray(points) + colors = np.asarray(colors) + + points_and_colors = np.hstack([points, colors]) + points_and_colors_unique, counts = np.unique( + points_and_colors, axis=0, return_counts=True + ) + + ax_l.scatter( + points_and_colors_unique[:, 0], + points_and_colors_unique[:, 1], + c=points_and_colors_unique[:, 2:], + s=counts * scale_hist, + ) + + v0 = self.orientation_zone_axis_range[0, :].copy() + v0 /= np.max(np.abs(v0)) + v0 = np.round(v0, 2) + + v1 = self.orientation_zone_axis_range[1, :].copy() + v1 /= np.max(np.abs(v1)) + v1 = np.round(v1, 2) + + v2 = self.orientation_zone_axis_range[2, :].copy() + v2 /= np.max(np.abs(v2)) + v2 = np.round(v2, 2) + + text_params = { + "va": "center", + "family": "sans-serif", + "fontweight": "normal", + "color": "k", + "size": 12, + } ax_l.text( - p3[1], - p3[0]+shift, - "[" - + format_labels.format(v[0]) - + " " - + format_labels.format(v[1]) - + " " - + format_labels.format(v[2]) - + "]", - None, - zorder=11, + A[0], + A[1] - shift, + v0, ha="center", **text_params, ) - - else: - v = self.orientation_zone_axis_range[0,:].copy() - v /= np.max(np.abs(v)) - v = np.round(v,2) ax_l.text( - p0[1], - p0[0]+shift, - "[" - + format_labels.format(v[0]) - + " " - + format_labels.format(v[1]) - + " " - + format_labels.format(v[2]) - + "]", - None, - zorder=11, + B[0], + B[1] - shift, + v1, ha="center", **text_params, ) - v = self.orientation_zone_axis_range[1,:].copy() - v /= np.max(np.abs(v)) - v = np.round(v,2) + ax_l.text( - p1[1], - p1[0]+shift, - "[" - + format_labels.format(v[0]) - + " " - + format_labels.format(v[1]) - + " " - + format_labels.format(v[2]) - + "]", - None, - zorder=11, + C[0] + shift * 2, + C[1], + v2, ha="center", **text_params, ) - v = self.orientation_zone_axis_range[2,:].copy() - v /= np.max(np.abs(v)) - v = np.round(v,2) + ax_l.text( - p2[1], - p2[0]-shift, - "[" - + format_labels.format(v[0]) - + " " - + format_labels.format(v[1]) - + " " - + format_labels.format(v[2]) - + "]", - None, - zorder=11, + C[0], + C[1] + 0.1, + "in-plane orientation", ha="center", **text_params, ) - - + ax_l.set_aspect("equal") + ax_l.axis("off") + + ax_l.set_xlim([-0.3, 1.3]) + # ax_l.set_ylim([-0.3, C[1] + 0.3]) - if self.pointgroup.get_crystal_system() == 'monoclinic': - ax_l.set_xlim((-bound,num_points*2-1+bound)) else: - ax_l.set_xlim((-bound,num_points+bound)) - ax_l.set_ylim((num_points+bound,-bound)) - ax_l.axis("off") + ax_l.imshow(rgb_leg) + + # Add text labels + text_scale_pos = 0.1 + text_params = { + "va": "center", + "family": "sans-serif", + "fontweight": "normal", + "color": "k", + "size": 12, + } + format_labels = "{0:.2g}" + + bound = num_points * 0.25 + shift = num_points * 0.10 + if self.pointgroup.get_crystal_system() == "monoclinic": + p0 = np.array((1, 1)) * num_points + p1 = np.array((0, 1)) * num_points + p2 = np.array((1, 2)) * num_points + p3 = np.array((1, 0)) * num_points + else: + p0 = np.array((1, 0)) * num_points + p1 = np.array((1, 1)) * num_points + p2 = ( + np.array((1 - np.cos(np.pi / 2 - tspan), np.sin(np.pi / 2 - tspan))) + * num_points + ) + + if self.pointgroup.get_crystal_system() == "monoclinic": + v = np.double(self.orientation_zone_axis_range[0, :]).copy() + v /= np.max(np.abs(v)) + v = np.round(v, 2) + ax_l.text( + p0[1], + p0[0] + shift, + "[" + + format_labels.format(v[0]) + + " " + + format_labels.format(v[1]) + + " " + + format_labels.format(v[2]) + + "]", + None, + zorder=11, + ha="center", + **text_params, + ) + v = np.double(self.orientation_zone_axis_range[1, :]).copy() + v /= np.max(np.abs(v)) + v = np.round(v, 2) + ax_l.text( + p1[1], + p1[0] - shift, + "[" + + format_labels.format(v[0]) + + " " + + format_labels.format(v[1]) + + " " + + format_labels.format(v[2]) + + "]", + None, + zorder=11, + ha="center", + **text_params, + ) + v = np.double(self.orientation_zone_axis_range[2, :]).copy() + v /= np.max(np.abs(v)) + v = np.round(v, 2) + ax_l.text( + p2[1], + p2[0] + shift, + "[" + + format_labels.format(v[0]) + + " " + + format_labels.format(v[1]) + + " " + + format_labels.format(v[2]) + + "]", + None, + zorder=11, + ha="center", + **text_params, + ) + v = -1 * np.double(self.orientation_zone_axis_range[2, :].copy()) + 0 + v /= np.max(np.abs(v)) + v = np.round(v, 2) + ax_l.text( + p3[1], + p3[0] + shift, + "[" + + format_labels.format(v[0]) + + " " + + format_labels.format(v[1]) + + " " + + format_labels.format(v[2]) + + "]", + None, + zorder=11, + ha="center", + **text_params, + ) + plt.tight_layout() + else: + v = self.orientation_zone_axis_range[0, :].copy() + v /= np.max(np.abs(v)) + v = np.round(v, 2) + ax_l.text( + p0[1], + p0[0] + shift, + "[" + + format_labels.format(v[0]) + + " " + + format_labels.format(v[1]) + + " " + + format_labels.format(v[2]) + + "]", + None, + zorder=11, + ha="center", + **text_params, + ) + v = self.orientation_zone_axis_range[1, :].copy() + v /= np.max(np.abs(v)) + v = np.round(v, 2) + ax_l.text( + p1[1], + p1[0] + shift, + "[" + + format_labels.format(v[0]) + + " " + + format_labels.format(v[1]) + + " " + + format_labels.format(v[2]) + + "]", + None, + zorder=11, + ha="center", + **text_params, + ) + v = self.orientation_zone_axis_range[2, :].copy() + v /= np.max(np.abs(v)) + v = np.round(v, 2) + ax_l.text( + p2[1], + p2[0] - shift, + "[" + + format_labels.format(v[0]) + + " " + + format_labels.format(v[1]) + + " " + + format_labels.format(v[2]) + + "]", + None, + zorder=11, + ha="center", + **text_params, + ) + + if self.pointgroup.get_crystal_system() == "monoclinic": + ax_l.set_xlim((-bound, num_points * 2 - 1 + bound)) + else: + ax_l.set_xlim((-bound, num_points + bound)) + ax_l.set_ylim((num_points + bound, -bound)) + ax_l.axis("off") # if np.abs(self.cell[5] - 120.0) > 1e-6: # ax_l.text( @@ -1823,7 +2030,6 @@ def plot_orientation_maps( # **text_params, # ) - # # Triangulate faces # p = self.orientation_vecs[:, (1, 0, 2)] # tri = mtri.Triangulation( @@ -2006,7 +2212,7 @@ def plot_orientation_maps( # **text_params, # ) - images_orientation = np.zeros((orientation_map.num_x, orientation_map.num_y, 3, 2)) + images_orientation = np.zeros((rgb_x.shape[0], rgb_x.shape[1], 3, 2)) if self.pymatgen_available: images_orientation[:, :, :, 0] = rgb_x images_orientation[:, :, :, 1] = rgb_z From 379825a744faee3191192ade667a1ec6ba2bbe6a Mon Sep 17 00:00:00 2001 From: smribet Date: Thu, 26 Mar 2026 15:42:39 -0700 Subject: [PATCH 11/11] cleaning up legend --- py4DSTEM/process/diffraction/crystal_viz.py | 221 +------------------- 1 file changed, 2 insertions(+), 219 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_viz.py b/py4DSTEM/process/diffraction/crystal_viz.py index dd6ecb066..9532aa41a 100644 --- a/py4DSTEM/process/diffraction/crystal_viz.py +++ b/py4DSTEM/process/diffraction/crystal_viz.py @@ -1992,225 +1992,8 @@ def plot_orientation_maps( ax_l.set_xlim((-bound, num_points + bound)) ax_l.set_ylim((num_points + bound, -bound)) ax_l.axis("off") - - # if np.abs(self.cell[5] - 120.0) > 1e-6: - # ax_l.text( - # self.orientation_vecs[inds_legend[0], 1] + vec[1] * text_scale_pos, - # self.orientation_vecs[inds_legend[0], 0] + vec[0] * text_scale_pos, - # self.orientation_vecs[inds_legend[0], 2] + vec[2] * text_scale_pos, - # "[" - # + format_labels.format(label_0[0]) - # + " " - # + format_labels.format(label_0[1]) - # + " " - # + format_labels.format(label_0[2]) - # + "]", - # None, - # zorder=11, - # ha="center", - # **text_params, - # ) - # else: - # ax_l.text( - # self.orientation_vecs[inds_legend[0], 1] + vec[1] * text_scale_pos, - # self.orientation_vecs[inds_legend[0], 0] + vec[0] * text_scale_pos, - # self.orientation_vecs[inds_legend[0], 2] + vec[2] * text_scale_pos, - # "[" - # + format_labels.format(label_0[0]) - # + " " - # + format_labels.format(label_0[1]) - # + " " - # + format_labels.format(label_0[2]) - # + " " - # + format_labels.format(label_0[3]) - # + "]", - # None, - # zorder=11, - # ha="center", - # **text_params, - # ) - - # # Triangulate faces - # p = self.orientation_vecs[:, (1, 0, 2)] - # tri = mtri.Triangulation( - # self.orientation_inds[:, 1] - self.orientation_inds[:, 0] * 1e-3, - # self.orientation_inds[:, 0] - self.orientation_inds[:, 1] * 1e-3, - # ) - # convert rgb values of pixels to faces - # rgb_faces = ( - # rgb_legend[tri.triangles[:, 0], :] - # + rgb_legend[tri.triangles[:, 1], :] - # + rgb_legend[tri.triangles[:, 2], :] - # ) / 3 - # Add triangulated surface plot to axes - # pc = art3d.Poly3DCollection( - # p[tri.triangles], - # facecolors=rgb_faces, - # alpha=1, - # ) - # pc.set_antialiased(False) - # ax_l.add_collection(pc) - - # if plot_limit is None: - # plot_limit = np.array( - # [ - # [np.min(p[:, 0]), np.min(p[:, 1]), np.min(p[:, 2])], - # [np.max(p[:, 0]), np.max(p[:, 1]), np.max(p[:, 2])], - # ] - # ) - # # plot_limit = (plot_limit - np.mean(plot_limit, axis=0)) * 1.5 + np.mean( - # # plot_limit, axis=0 - # # ) - # plot_limit[:, 0] = ( - # plot_limit[:, 0] - np.mean(plot_limit[:, 0]) - # ) * 1.5 + np.mean(plot_limit[:, 0]) - # plot_limit[:, 1] = ( - # plot_limit[:, 2] - np.mean(plot_limit[:, 1]) - # ) * 1.5 + np.mean(plot_limit[:, 1]) - # plot_limit[:, 2] = ( - # plot_limit[:, 1] - np.mean(plot_limit[:, 2]) - # ) * 1.1 + np.mean(plot_limit[:, 2]) - - # # ax_l.view_init(elev=el, azim=az) - # # Appearance - # ax_l.invert_yaxis() - # if swap_axes_xy_limits: - # ax_l.axes.set_xlim3d(left=plot_limit[0, 0], right=plot_limit[1, 0]) - # ax_l.axes.set_ylim3d(bottom=plot_limit[0, 1], top=plot_limit[1, 1]) - # ax_l.axes.set_zlim3d(bottom=plot_limit[0, 2], top=plot_limit[1, 2]) - # else: - # ax_l.axes.set_xlim3d(left=plot_limit[0, 1], right=plot_limit[1, 1]) - # ax_l.axes.set_ylim3d(bottom=plot_limit[0, 0], top=plot_limit[1, 0]) - # ax_l.axes.set_zlim3d(bottom=plot_limit[0, 2], top=plot_limit[1, 2]) - # axisEqual3D(ax_l) - # if camera_dist is not None: - # ax_l.dist = camera_dist - # ax_l.axis("off") - - # # Add text labels - # text_scale_pos = 0.1 - # text_params = { - # "va": "center", - # "family": "sans-serif", - # "fontweight": "normal", - # "color": "k", - # "size": 14, - # } - # format_labels = "{0:.2g}" - # vec = self.orientation_vecs[inds_legend[0], :] - cam_dir - # vec = vec / np.linalg.norm(vec) - # if np.abs(self.cell[5] - 120.0) > 1e-6: - # ax_l.text( - # self.orientation_vecs[inds_legend[0], 1] + vec[1] * text_scale_pos, - # self.orientation_vecs[inds_legend[0], 0] + vec[0] * text_scale_pos, - # self.orientation_vecs[inds_legend[0], 2] + vec[2] * text_scale_pos, - # "[" - # + format_labels.format(label_0[0]) - # + " " - # + format_labels.format(label_0[1]) - # + " " - # + format_labels.format(label_0[2]) - # + "]", - # None, - # zorder=11, - # ha="center", - # **text_params, - # ) - # else: - # ax_l.text( - # self.orientation_vecs[inds_legend[0], 1] + vec[1] * text_scale_pos, - # self.orientation_vecs[inds_legend[0], 0] + vec[0] * text_scale_pos, - # self.orientation_vecs[inds_legend[0], 2] + vec[2] * text_scale_pos, - # "[" - # + format_labels.format(label_0[0]) - # + " " - # + format_labels.format(label_0[1]) - # + " " - # + format_labels.format(label_0[2]) - # + " " - # + format_labels.format(label_0[3]) - # + "]", - # None, - # zorder=11, - # ha="center", - # **text_params, - # ) - # vec = self.orientation_vecs[inds_legend[1], :] - cam_dir - # vec = vec / np.linalg.norm(vec) - # if np.abs(self.cell[5] - 120.0) > 1e-6: - # ax_l.text( - # self.orientation_vecs[inds_legend[1], 1] + vec[1] * text_scale_pos, - # self.orientation_vecs[inds_legend[1], 0] + vec[0] * text_scale_pos, - # self.orientation_vecs[inds_legend[1], 2] + vec[2] * text_scale_pos, - # "[" - # + format_labels.format(label_1[0]) - # + " " - # + format_labels.format(label_1[1]) - # + " " - # + format_labels.format(label_1[2]) - # + "]", - # None, - # zorder=12, - # ha=ha_1, - # **text_params, - # ) - # else: - # ax_l.text( - # self.orientation_vecs[inds_legend[1], 1] + vec[1] * text_scale_pos, - # self.orientation_vecs[inds_legend[1], 0] + vec[0] * text_scale_pos, - # self.orientation_vecs[inds_legend[1], 2] + vec[2] * text_scale_pos, - # "[" - # + format_labels.format(label_1[0]) - # + " " - # + format_labels.format(label_1[1]) - # + " " - # + format_labels.format(label_1[2]) - # + " " - # + format_labels.format(label_1[3]) - # + "]", - # None, - # zorder=12, - # ha=ha_1, - # **text_params, - # ) - # vec = self.orientation_vecs[inds_legend[2], :] - cam_dir - # vec = vec / np.linalg.norm(vec) - # if np.abs(self.cell[5] - 120.0) > 1e-6: - # ax_l.text( - # self.orientation_vecs[inds_legend[2], 1] + vec[1] * text_scale_pos, - # self.orientation_vecs[inds_legend[2], 0] + vec[0] * text_scale_pos, - # self.orientation_vecs[inds_legend[2], 2] + vec[2] * text_scale_pos, - # "[" - # + format_labels.format(label_2[0]) - # + " " - # + format_labels.format(label_2[1]) - # + " " - # + format_labels.format(label_2[2]) - # + "]", - # None, - # zorder=13, - # ha=ha_2, - # **text_params, - # ) - # else: - # ax_l.text( - # self.orientation_vecs[inds_legend[2], 1] + vec[1] * text_scale_pos, - # self.orientation_vecs[inds_legend[2], 0] + vec[0] * text_scale_pos, - # self.orientation_vecs[inds_legend[2], 2] + vec[2] * text_scale_pos, - # "[" - # + format_labels.format(label_2[0]) - # + " " - # + format_labels.format(label_2[1]) - # + " " - # + format_labels.format(label_2[2]) - # + " " - # + format_labels.format(label_2[3]) - # + "]", - # None, - # zorder=13, - # ha=ha_2, - # **text_params, - # ) + else: + ax_l.remove() images_orientation = np.zeros((rgb_x.shape[0], rgb_x.shape[1], 3, 2)) if self.pymatgen_available: