From fc7fb172e13d325b229541126aa2cf6a6de7fcda Mon Sep 17 00:00:00 2001 From: serinlee1065 Date: Tue, 2 Sep 2025 13:04:25 -0700 Subject: [PATCH 1/6] 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 2/6] 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 3/6] 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 4/6] 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 5/6] 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 6/6] 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