From a0ea19973a4769f2a515d5f8e2d53b2a5d3127b8 Mon Sep 17 00:00:00 2001 From: Colin Date: Fri, 12 May 2023 11:48:34 -0700 Subject: [PATCH 01/13] Adding grain clustering function --- py4DSTEM/process/diffraction/crystal.py | 1 + py4DSTEM/process/diffraction/crystal_ACOM.py | 241 ++++++++++++++++++- py4DSTEM/process/diffraction/crystal_viz.py | 6 +- 3 files changed, 243 insertions(+), 5 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index e0b171875..5bd931b1f 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -36,6 +36,7 @@ class Crystal: orientation_plan, match_orientations, match_single_pattern, + cluster_grains, calculate_strain, save_ang_file, symmetry_reduce_directions, diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index bb4623491..67c0fe0b4 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -2,6 +2,7 @@ import matplotlib.pyplot as plt import os from typing import Union, Optional +import time, sys from py4DSTEM.process.diffraction.utils import Orientation, OrientationMap, axisEqual3D from py4DSTEM.process.utils import electron_wavelength_angstrom @@ -753,9 +754,10 @@ def match_orientations( self, bragg_peaks_array: PointListArray, num_matches_return: int = 1, - min_number_peaks = 3, - inversion_symmetry = True, - multiple_corr_reset = True, + min_number_peaks: int = 3, + inversion_symmetry: bool = True, + multiple_corr_reset: bool = True, + return_orientation: bool = True, progress_bar: bool = True, ): ''' @@ -786,7 +788,13 @@ def match_orientations( orientation_map.set_orientation(orientation,rx,ry) - return orientation_map + # assign and return + self.orientation_map = orientation_map + + if return_orientation: + return orientation_map + else: + return def match_single_pattern( self, @@ -1563,6 +1571,207 @@ def match_single_pattern( + + +def cluster_grains( + self, + threshold_add = 1.0, + threshold_grow = 0.0, + angle_tolerance_deg = 5.0, + progress_bar = True, + ): + """ + Cluster grains from a specific radial bin + + Parameters + -------- + corr_threshold_add: float + Minimum signal required for a probe position to initialize a cluster. + corr_threshold_grow: float + Minimum signal required for a probe position to be added to a cluster. + angle_tolerance_deg: float + Rotation rolerance for clustering grains. + progress_bar: bool + Turns on the progress bar for the polar transformation + + Returns + -------- + + + """ + + # symmetry operators + sym = self.symmetry_operators + + # Get data + matrix = self.orientation_map.matrix.copy() + corr = self.orientation_map.corr.copy() + corr_init = corr.copy() + mark = corr >= threshold_grow + corr[np.logical_not(mark)] = 0 + + # phi = np.squeeze(self.radial_peaks[:,:,radial_index,:,0]).copy() + # sig = np.squeeze(self.radial_peaks[:,:,radial_index,:,1]).copy() + # sig_init = sig.copy() + # mark = sig >= threshold_grow + # sig[np.logical_not(mark)] = 0 + + # init + self.cluster_sizes = np.array((), dtype='int') + self.cluster_sig = np.array(()) + self.cluster_inds = [] + inds_all = np.zeros_like(corr, dtype='int') + inds_all.ravel()[:] = np.arange(inds_all.size) + + # Tolerance + tol = np.deg2rad(angle_tolerance_deg) + + # Main loop + search = True + while search is True: + inds_grain = np.argmax(corr) + val = corr.ravel()[inds_grain] + + if val < threshold_add: + search = False + + else: + # progressbar + if progress_bar: + comp = 1 - np.mean(np.max(mark,axis = 2)) + update_progress(comp) + + # # Start cluster + # x,y,z = np.unravel_index(inds_grain, sig.shape) + # mark[x,y,z] = False + # sig[x,y,z] = 0 + # phi_cluster = phi[x,y,z] + + # # Neighbors to search + # xr = np.clip(x + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1) + # yr = np.clip(y + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1) + # inds_cand = inds_all[xr,yr,:].ravel() + # inds_cand = np.delete(inds_cand, mark.ravel()[inds_cand] == False) + # # [mark[xr,yr,:].ravel()] + + # if inds_cand.size == 0: + # grow = False + # else: + # grow = True + + # # grow the cluster + # while grow is True: + # inds_new = np.array((),dtype='int') + + # keep = np.zeros(inds_cand.size, dtype='bool') + # for a0 in range(inds_cand.size): + # xc,yc,zc = np.unravel_index(inds_cand[a0], sig.shape) + + # phi_test = phi[xc,yc,zc] + # dphi = np.mod(phi_cluster - phi_test + np.pi/2.0, np.pi) - np.pi/2.0 + + # if np.abs(dphi) < tol: + # keep[a0] = True + + # sig[xc,yc,zc] = 0 + # mark[xc,yc,zc] = False + + # xr = np.clip(xc + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1) + # yr = np.clip(yc + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1) + # inds_add = inds_all[xr,yr,:].ravel() + # inds_new = np.append(inds_new, inds_add) + + + # inds_grain = np.append(inds_grain, inds_cand[keep]) + # inds_cand = np.unique(np.delete(inds_new, mark.ravel()[inds_new] == False)) + + # if inds_cand.size == 0: + # grow = False + + # # convert grain to x,y coordinates, add = list + # xg,yg,zg = np.unravel_index(inds_grain, sig.shape) + # xyg = np.unique(np.vstack((xg,yg)), axis = 1) + # sig_mean = np.mean(sig_init.ravel()[inds_grain]) + # self.cluster_sizes = np.append(self.cluster_sizes, xyg.shape[1]) + # self.cluster_sig = np.append(self.cluster_sig, sig_mean) + # self.cluster_inds.append(xyg) + search = False + + # finish progressbar + if progress_bar: + update_progress(1) + +# def cluster_plot_size( +# self, +# area_max = None, +# weight_intensity = False, +# pixel_area = 1.0, +# pixel_area_units = 'px^2', +# figsize = (8,6), +# returnfig = False, +# ): +# """ +# Plot the cluster sizes + +# Parameters +# -------- +# area_max: int (optional) +# Max area bin in pixels +# weight_intensity: bool +# Weight histogram by the peak intensity. +# pixel_area: float +# Size of pixel area unit square +# pixel_area_units: string +# Units of the pixel area +# figsize: tuple +# Size of the figure panel +# returnfig: bool +# Setting this to true returns the figure and axis handles + +# Returns +# -------- +# fig, ax (optional) +# Figure and axes handles + +# """ + +# if area_max is None: +# area_max = np.max(self.cluster_sizes) +# area = np.arange(area_max) +# sub = self.cluster_sizes.astype('int') < area_max +# if weight_intensity: +# hist = np.bincount( +# self.cluster_sizes[sub], +# weights = self.cluster_sig[sub], +# minlength = area_max, +# ) +# else: +# hist = np.bincount( +# self.cluster_sizes[sub], +# minlength = area_max, +# ) + + +# # plotting +# fig,ax = plt.subplots(figsize = figsize) +# ax.bar( +# area * pixel_area, +# hist, +# width = 0.8 * pixel_area, +# ) + +# ax.set_xlabel('Grain Area [' + pixel_area_units + ']') +# if weight_intensity: +# ax.set_ylabel('Total Signal [arb. units]') +# else: +# ax.set_ylabel('Number of Grains') + +# if returnfig: +# return fig,ax + + + + def calculate_strain( self, bragg_peaks_array: PointListArray, @@ -2066,3 +2275,27 @@ def symmetry_reduce_directions( # "-3m": ["fiber", [0, 0, 1], [90.0, 60.0]], # "-3m": ["fiber", [0, 0, 1], [180.0, 30.0]], + + +# Progressbar taken from stackexchange: +# https://stackoverflow.com/questions/3160699/python-progress-bar +def update_progress(progress): + barLength = 60 # Modify this to change the length of the progress bar + status = "" + if isinstance(progress, int): + progress = float(progress) + if not isinstance(progress, float): + progress = 0 + status = "error: progress var must be float\r\n" + if progress < 0: + progress = 0 + status = "Halt...\r\n" + if progress >= 1: + progress = 1 + status = "Done...\r\n" + block = int(round(barLength*progress)) + text = "\rPercent: [{0}] {1}% {2}".format( "#"*block + "-"*(barLength-block), + np.round(progress*100,4), + status) + sys.stdout.write(text) + sys.stdout.flush() diff --git a/py4DSTEM/process/diffraction/crystal_viz.py b/py4DSTEM/process/diffraction/crystal_viz.py index f0ed95663..a2247b3e0 100644 --- a/py4DSTEM/process/diffraction/crystal_viz.py +++ b/py4DSTEM/process/diffraction/crystal_viz.py @@ -955,7 +955,7 @@ def overline(x): def plot_orientation_maps( self, - orientation_map, + orientation_map = None, orientation_ind: int = 0, dir_in_plane_degrees: float = 0.0, corr_range: np.ndarray = np.array([0, 5]), @@ -976,6 +976,7 @@ def plot_orientation_maps( Args: orientation_map (OrientationMap): Class containing orientation matrices, correlation values, etc. + Optional - can reference internally stored OrientationMap. orientation_ind (int): Which orientation match to plot if num_matches > 1 dir_in_plane_degrees (float): In-plane angle to plot in degrees. Default is 0 / x-axis / vertical down. corr_range (np.ndarray): Correlation intensity range for the plot @@ -1003,6 +1004,9 @@ def plot_orientation_maps( """ # Inputs + if orientation_map is None: + orientation_map = self.orientation_map + # Legend size leg_size = np.array([300, 300], dtype="int") From 234e372058ad14f1504fd10e6a918d8624cf0701 Mon Sep 17 00:00:00 2001 From: Colin Date: Fri, 12 May 2023 14:18:39 -0700 Subject: [PATCH 02/13] Fixing bugs in the grain clustering --- py4DSTEM/process/diffraction/crystal.py | 2 + py4DSTEM/process/diffraction/crystal_ACOM.py | 214 +++++++------------ py4DSTEM/process/diffraction/crystal_viz.py | 90 ++++++++ 3 files changed, 166 insertions(+), 140 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index 5bd931b1f..f93702c5d 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -52,6 +52,8 @@ class Crystal: plot_orientation_plan, plot_orientation_maps, plot_fiber_orientation_maps, + plot_clusters, + plot_cluster_size, ) from py4DSTEM.process.diffraction.crystal_calibrate import ( diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index 67c0fe0b4..6a8fe5280 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -1576,7 +1576,7 @@ def match_single_pattern( def cluster_grains( self, threshold_add = 1.0, - threshold_grow = 0.0, + threshold_grow = 0.1, angle_tolerance_deg = 5.0, progress_bar = True, ): @@ -1604,23 +1604,19 @@ def cluster_grains( sym = self.symmetry_operators # Get data + # Correlation data = signal to cluster with + sig = self.orientation_map.corr.copy() + sig_init = sig.copy() + mark = sig >= threshold_grow + sig[np.logical_not(mark)] = 0 + # orientation matrix used for angle tolerance matrix = self.orientation_map.matrix.copy() - corr = self.orientation_map.corr.copy() - corr_init = corr.copy() - mark = corr >= threshold_grow - corr[np.logical_not(mark)] = 0 - - # phi = np.squeeze(self.radial_peaks[:,:,radial_index,:,0]).copy() - # sig = np.squeeze(self.radial_peaks[:,:,radial_index,:,1]).copy() - # sig_init = sig.copy() - # mark = sig >= threshold_grow - # sig[np.logical_not(mark)] = 0 - + # init self.cluster_sizes = np.array((), dtype='int') self.cluster_sig = np.array(()) self.cluster_inds = [] - inds_all = np.zeros_like(corr, dtype='int') + inds_all = np.zeros_like(sig, dtype='int') inds_all.ravel()[:] = np.arange(inds_all.size) # Tolerance @@ -1629,8 +1625,9 @@ def cluster_grains( # Main loop search = True while search is True: - inds_grain = np.argmax(corr) - val = corr.ravel()[inds_grain] + inds_grain = np.argmax(sig) + + val = sig.ravel()[inds_grain] if val < threshold_add: search = False @@ -1641,136 +1638,73 @@ def cluster_grains( comp = 1 - np.mean(np.max(mark,axis = 2)) update_progress(comp) - # # Start cluster - # x,y,z = np.unravel_index(inds_grain, sig.shape) - # mark[x,y,z] = False - # sig[x,y,z] = 0 - # phi_cluster = phi[x,y,z] - - # # Neighbors to search - # xr = np.clip(x + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1) - # yr = np.clip(y + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1) - # inds_cand = inds_all[xr,yr,:].ravel() - # inds_cand = np.delete(inds_cand, mark.ravel()[inds_cand] == False) - # # [mark[xr,yr,:].ravel()] - - # if inds_cand.size == 0: - # grow = False - # else: - # grow = True - - # # grow the cluster - # while grow is True: - # inds_new = np.array((),dtype='int') - - # keep = np.zeros(inds_cand.size, dtype='bool') - # for a0 in range(inds_cand.size): - # xc,yc,zc = np.unravel_index(inds_cand[a0], sig.shape) - - # phi_test = phi[xc,yc,zc] - # dphi = np.mod(phi_cluster - phi_test + np.pi/2.0, np.pi) - np.pi/2.0 - - # if np.abs(dphi) < tol: - # keep[a0] = True + # Start cluster + x,y,z = np.unravel_index(inds_grain, sig.shape) + mark[x,y,z] = False + sig[x,y,z] = 0 + matrix_cluster = matrix[x,y,z] - # sig[xc,yc,zc] = 0 - # mark[xc,yc,zc] = False + # Neighbors to search + xr = np.clip(x + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1) + yr = np.clip(y + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1) + inds_cand = inds_all[xr[:,None],yr[None],:].ravel() + inds_cand = np.delete(inds_cand, mark.ravel()[inds_cand] == False) - # xr = np.clip(xc + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1) - # yr = np.clip(yc + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1) - # inds_add = inds_all[xr,yr,:].ravel() - # inds_new = np.append(inds_new, inds_add) - - - # inds_grain = np.append(inds_grain, inds_cand[keep]) - # inds_cand = np.unique(np.delete(inds_new, mark.ravel()[inds_new] == False)) - - # if inds_cand.size == 0: - # grow = False - - # # convert grain to x,y coordinates, add = list - # xg,yg,zg = np.unravel_index(inds_grain, sig.shape) - # xyg = np.unique(np.vstack((xg,yg)), axis = 1) - # sig_mean = np.mean(sig_init.ravel()[inds_grain]) - # self.cluster_sizes = np.append(self.cluster_sizes, xyg.shape[1]) - # self.cluster_sig = np.append(self.cluster_sig, sig_mean) - # self.cluster_inds.append(xyg) - search = False + if inds_cand.size == 0: + grow = False + else: + grow = True + + # grow the cluster + while grow is True: + inds_new = np.array((),dtype='int') + + keep = np.zeros(inds_cand.size, dtype='bool') + for a0 in range(inds_cand.size): + xc,yc,zc = np.unravel_index(inds_cand[a0], sig.shape) + + # Angle test between orientation matrices + dphi = np.min(np.arccos(np.clip((np.trace( + self.symmetry_operators @ matrix[xc,yc,zc] \ + @ np.transpose(matrix_cluster), + axis1=1, + axis2=2)-1)/2,-1,1))) + # dphi = np.min(np.arccos(np.clip((np.trace( + # matrix_cluster @ \ + # np.transpose(self.symmetry_operators @ matrix[xc,yc,zc],(0,2,1)), + # axis1=1, + # axis2=2)-1)/2,-1,1))) + + if np.abs(dphi) < tol: + keep[a0] = True + + sig[xc,yc,zc] = 0 + mark[xc,yc,zc] = False + + xr = np.clip(xc + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1) + yr = np.clip(yc + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1) + inds_add = inds_all[xr[:,None],yr[None],:].ravel() + inds_new = np.append(inds_new, inds_add) + + inds_grain = np.append(inds_grain, inds_cand[keep]) + inds_cand = np.unique(np.delete(inds_new, mark.ravel()[inds_new] == False)) + # print(inds_cand) + + if inds_cand.size == 0: + grow = False + + # convert grain to x,y coordinates, add = list + xg,yg,zg = np.unravel_index(inds_grain, sig.shape) + xyg = np.unique(np.vstack((xg,yg)), axis = 1) + sig_mean = np.mean(sig_init.ravel()[inds_grain]) + self.cluster_sizes = np.append(self.cluster_sizes, xyg.shape[1]) + self.cluster_sig = np.append(self.cluster_sig, sig_mean) + self.cluster_inds.append(xyg) # finish progressbar if progress_bar: update_progress(1) -# def cluster_plot_size( -# self, -# area_max = None, -# weight_intensity = False, -# pixel_area = 1.0, -# pixel_area_units = 'px^2', -# figsize = (8,6), -# returnfig = False, -# ): -# """ -# Plot the cluster sizes - -# Parameters -# -------- -# area_max: int (optional) -# Max area bin in pixels -# weight_intensity: bool -# Weight histogram by the peak intensity. -# pixel_area: float -# Size of pixel area unit square -# pixel_area_units: string -# Units of the pixel area -# figsize: tuple -# Size of the figure panel -# returnfig: bool -# Setting this to true returns the figure and axis handles - -# Returns -# -------- -# fig, ax (optional) -# Figure and axes handles - -# """ - -# if area_max is None: -# area_max = np.max(self.cluster_sizes) -# area = np.arange(area_max) -# sub = self.cluster_sizes.astype('int') < area_max -# if weight_intensity: -# hist = np.bincount( -# self.cluster_sizes[sub], -# weights = self.cluster_sig[sub], -# minlength = area_max, -# ) -# else: -# hist = np.bincount( -# self.cluster_sizes[sub], -# minlength = area_max, -# ) - - -# # plotting -# fig,ax = plt.subplots(figsize = figsize) -# ax.bar( -# area * pixel_area, -# hist, -# width = 0.8 * pixel_area, -# ) - -# ax.set_xlabel('Grain Area [' + pixel_area_units + ']') -# if weight_intensity: -# ax.set_ylabel('Total Signal [arb. units]') -# else: -# ax.set_ylabel('Number of Grains') - -# if returnfig: -# return fig,ax - - - def calculate_strain( self, @@ -2292,10 +2226,10 @@ def update_progress(progress): status = "Halt...\r\n" if progress >= 1: progress = 1 - status = "Done...\r\n" + status = "Done\r\n" block = int(round(barLength*progress)) text = "\rPercent: [{0}] {1}% {2}".format( "#"*block + "-"*(barLength-block), - np.round(progress*100,4), + np.round(progress*100,2), status) sys.stdout.write(text) sys.stdout.flush() diff --git a/py4DSTEM/process/diffraction/crystal_viz.py b/py4DSTEM/process/diffraction/crystal_viz.py index a2247b3e0..b4cc7e65e 100644 --- a/py4DSTEM/process/diffraction/crystal_viz.py +++ b/py4DSTEM/process/diffraction/crystal_viz.py @@ -1623,6 +1623,96 @@ def plot_fiber_orientation_maps( else: return images_orientation +def plot_clusters( + self, + returnfig = False, + ): + 1+1 + + +def plot_cluster_size( + self, + area_min = None, + area_max = None, + weight_intensity = False, + pixel_area = 1.0, + pixel_area_units = 'px^2', + figsize = (8,6), + returnfig = False, + ): + """ + Plot the cluster sizes + + Parameters + -------- + area_min: int (optional) + Min area to include + area_max: int (optional) + Max area bin in pixels + weight_intensity: bool + Weight histogram by the peak intensity. + pixel_area: float + Size of pixel area unit square + pixel_area_units: string + Units of the pixel area + figsize: tuple + Size of the figure panel + returnfig: bool + Setting this to true returns the figure and axis handles + + Returns + -------- + fig, ax (optional) + Figure and axes handles + + """ + + if area_max is None: + area_max = np.max(self.cluster_sizes) + area = np.arange(area_max) + if area_min is None: + sub = self.cluster_sizes.astype('int') < area_max + else: + sub = np.logical_and( + self.cluster_sizes.astype('int') < area_max, + self.cluster_sizes.astype('int') >= area_min, + ) + + + if weight_intensity: + hist = np.bincount( + self.cluster_sizes[sub], + weights = self.cluster_sig[sub], + minlength = area_max, + ) + else: + hist = np.bincount( + self.cluster_sizes[sub], + minlength = area_max, + ) + + + # plotting + fig,ax = plt.subplots(figsize = figsize) + ax.bar( + area * pixel_area, + hist, + width = 0.8 * pixel_area, + ) + + ax.set_xlabel('Grain Area [' + pixel_area_units + ']') + if weight_intensity: + ax.set_ylabel('Total Signal [arb. units]') + else: + ax.set_ylabel('Number of Grains') + + if returnfig: + return fig,ax + + + + + def axisEqual3D(ax): extents = np.array([getattr(ax, "get_{}lim".format(dim))() for dim in "xyz"]) From 7c89abbdc08fd0390bbf2024c9e4b50fc8c5258f Mon Sep 17 00:00:00 2001 From: Colin Date: Fri, 12 May 2023 18:03:11 -0700 Subject: [PATCH 03/13] More plotting tools --- py4DSTEM/process/diffraction/crystal.py | 1 + py4DSTEM/process/diffraction/crystal_ACOM.py | 98 +++++++++++++-- py4DSTEM/process/diffraction/crystal_viz.py | 120 ++++++++++++++++++- py4DSTEM/process/diffraction/utils.py | 11 ++ 4 files changed, 222 insertions(+), 8 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index f93702c5d..709093f4d 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -37,6 +37,7 @@ class Crystal: match_orientations, match_single_pattern, cluster_grains, + cluster_orientation_map, calculate_strain, save_ang_file, symmetry_reduce_directions, diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index 6a8fe5280..8440b4836 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -1581,7 +1581,7 @@ def cluster_grains( progress_bar = True, ): """ - Cluster grains from a specific radial bin + Cluster grains using rotation criterion, and correlation values. Parameters -------- @@ -1616,6 +1616,7 @@ def cluster_grains( self.cluster_sizes = np.array((), dtype='int') self.cluster_sig = np.array(()) self.cluster_inds = [] + self.cluster_orientation = [] inds_all = np.zeros_like(sig, dtype='int') inds_all.ravel()[:] = np.arange(inds_all.size) @@ -1643,6 +1644,7 @@ def cluster_grains( mark[x,y,z] = False sig[x,y,z] = 0 matrix_cluster = matrix[x,y,z] + orientation_cluster = self.orientation_map.get_orientation_single(x,y,z) # Neighbors to search xr = np.clip(x + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1) @@ -1669,11 +1671,6 @@ def cluster_grains( @ np.transpose(matrix_cluster), axis1=1, axis2=2)-1)/2,-1,1))) - # dphi = np.min(np.arccos(np.clip((np.trace( - # matrix_cluster @ \ - # np.transpose(self.symmetry_operators @ matrix[xc,yc,zc],(0,2,1)), - # axis1=1, - # axis2=2)-1)/2,-1,1))) if np.abs(dphi) < tol: keep[a0] = True @@ -1688,7 +1685,6 @@ def cluster_grains( inds_grain = np.append(inds_grain, inds_cand[keep]) inds_cand = np.unique(np.delete(inds_new, mark.ravel()[inds_new] == False)) - # print(inds_cand) if inds_cand.size == 0: grow = False @@ -1699,6 +1695,7 @@ def cluster_grains( sig_mean = np.mean(sig_init.ravel()[inds_grain]) self.cluster_sizes = np.append(self.cluster_sizes, xyg.shape[1]) self.cluster_sig = np.append(self.cluster_sig, sig_mean) + self.cluster_orientation.append(orientation_cluster) self.cluster_inds.append(xyg) # finish progressbar @@ -1706,6 +1703,93 @@ def cluster_grains( update_progress(1) + + +def cluster_orientation_map( + self, + stripe_width = 1, + area_min = 2, + ): + """ + Produce a new orientation map from the clustered grains. + Use a stripe pattern for the overlapping grains. + + Parameters + -------- + stripe_width: int + Width of strips in the overlapping regions. + + Returns + -------- + + orientation_map + + """ + + # init + orientation_map = OrientationMap( + num_x = self.orientation_map.num_x, + num_y = self.orientation_map.num_y, + num_matches=1) + im_grain = np.zeros(( + self.orientation_map.num_x, + self.orientation_map.num_y), dtype='bool') + im_count = np.zeros(( + self.orientation_map.num_x, + self.orientation_map.num_y)) + im_mark = np.zeros(( + self.orientation_map.num_x, + self.orientation_map.num_y)) + + # coordinates + xa,ya = np.meshgrid( + range(self.orientation_map.num_x), + range(self.orientation_map.num_y), + indexing = 'ij') + + # Loop over grains to determine number in each pixel + for a0 in range(self.cluster_sizes.shape[0]): + if self.cluster_sizes[a0] >= area_min: + im_grain[:] = False + im_grain[ + self.cluster_inds[a0][0,:], + self.cluster_inds[a0][1,:], + ] = True + im_count += im_grain + im_stripe = im_count >= 2 + im_single = np.logical_not(im_stripe) + + # loop over grains + for a0 in range(1): + if self.cluster_sizes[a0] >= area_min: + im_grain[:] = False + im_grain[ + self.cluster_inds[a0][0,:], + self.cluster_inds[a0][1,:], + ] = True + + sub = np.logical_and( + im_grain, + im_single) + print(np.sum(sub)) + + sub = np.logical_and( + im_grain, + im_stripe) + print(np.sum(sub)) + + + + + fig,ax = plt.subplots(figsize=(8,8)) + ax.imshow(im_stripe) + + + + + return orientation_map + + def calculate_strain( self, bragg_peaks_array: PointListArray, diff --git a/py4DSTEM/process/diffraction/crystal_viz.py b/py4DSTEM/process/diffraction/crystal_viz.py index b4cc7e65e..19bdc62ab 100644 --- a/py4DSTEM/process/diffraction/crystal_viz.py +++ b/py4DSTEM/process/diffraction/crystal_viz.py @@ -5,6 +5,8 @@ from mpl_toolkits.mplot3d import Axes3D, art3d from scipy.signal import medfilt from scipy.ndimage import gaussian_filter +from scipy.ndimage.morphology import distance_transform_edt +from skimage.morphology import dilation, erosion import warnings import numpy as np @@ -1625,9 +1627,125 @@ def plot_fiber_orientation_maps( def plot_clusters( self, + area_min = 2, + outline_grains = True, + outline_thickness = 1, + fill_grains = 0.25, + smooth_grains = 1.0, + cmap = 'viridis', + figsize = (8,8), returnfig = False, ): - 1+1 + """ + Plot the clusters as an image. + + Parameters + -------- + area_min: int (optional) + Min cluster size to include, in units of probe positions. + outline_grains: bool (optional) + Set to True to draw grains with outlines + outline_thickness: int (optional) + Thickenss of the grain outline + fill_grains: float (optional) + Outlined grains are filled with this value in pixels. + smooth_grains: float (optional) + Grain boundaries are smoothed by this value in pixels. + figsize: tuple + Size of the figure panel + returnfig: bool + Setting this to true returns the figure and axis handles + + Returns + -------- + fig, ax (optional) + Figure and axes handles + + """ + + # init + im_plot = np.zeros(( + self.orientation_map.num_x, + self.orientation_map.num_y, + )) + im_grain = np.zeros(( + self.orientation_map.num_x, + self.orientation_map.num_y, + ), dtype='bool') + + # make plotting image + + for a0 in range(self.cluster_sizes.shape[0]): + if self.cluster_sizes[a0] >= area_min: + if outline_grains: + im_grain[:] = False + im_grain[ + self.cluster_inds[a0][0,:], + self.cluster_inds[a0][1,:], + ] = True + + im_dist = \ + distance_transform_edt( + erosion(np.invert(im_grain),footprint=np.ones((3,3),dtype='bool')) + ) \ + - \ + distance_transform_edt( + im_grain + ) + im_dist = gaussian_filter( + im_dist, + sigma = smooth_grains, + mode = 'nearest') + im_add = np.exp(im_dist**2 / (-0.5*outline_thickness**2)) + + if fill_grains > 0: + im_dist = \ + distance_transform_edt( + erosion(np.invert(im_grain),footprint=np.ones((3,3),dtype='bool')) + ) + im_dist = gaussian_filter( + im_dist, + sigma = smooth_grains, + mode = 'nearest') + im_add += fill_grains*np.exp(im_dist**2 / (-0.5*outline_thickness**2)) + + + # im_add = 1 - np.exp( + # distance_transform_edt(im_grain)**2 \ + # / (-2*outline_thickness**2)) + im_plot += im_add + # im_plot = np.minimum(im_plot, im_add) + else: + # xg,yg = np.unravel_index(self.cluster_inds[a0], im_plot.shape) + im_grain[:] = False + im_grain[ + self.cluster_inds[a0][0,:], + self.cluster_inds[a0][1,:], + ] = True + im_plot += gaussian_filter( + im_grain.astype('float'), + sigma = smooth_grains, + mode = 'nearest') + + # im_plot[ + # self.cluster_inds[a0][0,:], + # self.cluster_inds[a0][1,:], + # ] += 1 + + if outline_grains: + im_plot = np.clip(im_plot,0,2) + + + # plotting + fig,ax = plt.subplots(figsize=figsize) + ax.imshow( + im_plot, + # vmin = -3, + # vmax = 3, + cmap = cmap, + ) + + def plot_cluster_size( diff --git a/py4DSTEM/process/diffraction/utils.py b/py4DSTEM/process/diffraction/utils.py index c3ed3085a..a390e99c8 100644 --- a/py4DSTEM/process/diffraction/utils.py +++ b/py4DSTEM/process/diffraction/utils.py @@ -25,6 +25,7 @@ def __post_init__(self): self.angles = np.zeros((self.num_matches, 3)) + @dataclass class OrientationMap: """ @@ -66,6 +67,16 @@ def get_orientation(self, ind_x, ind_y): orientation.angles = self.angles[ind_x, ind_y] return orientation + def get_orientation_single(self, ind_x, ind_y, ind_match): + orientation = Orientation(num_matches=1) + orientation.matrix = self.matrix[ind_x, ind_y, ind_match] + orientation.family = self.family[ind_x, ind_y, ind_match] + orientation.corr = self.corr[ind_x, ind_y, ind_match] + orientation.inds = self.inds[ind_x, ind_y, ind_match] + orientation.mirror = self.mirror[ind_x, ind_y, ind_match] + orientation.angles = self.angles[ind_x, ind_y, ind_match] + return orientation + # def __copy__(self): # return OrientationMap(self.name) # def __deepcopy__(self, memo): From e0e0a12e52fc7903b4a93811be725cfd127799c1 Mon Sep 17 00:00:00 2001 From: Colin Date: Fri, 12 May 2023 18:30:22 -0700 Subject: [PATCH 04/13] Updated viz --- py4DSTEM/process/diffraction/crystal_ACOM.py | 66 +++++++++++++++----- 1 file changed, 50 insertions(+), 16 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index 8440b4836..6e11960f6 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -1707,7 +1707,7 @@ def cluster_grains( def cluster_orientation_map( self, - stripe_width = 1, + stripe_width = (2,2), area_min = 2, ): """ @@ -1741,11 +1741,11 @@ def cluster_orientation_map( self.orientation_map.num_x, self.orientation_map.num_y)) - # coordinates - xa,ya = np.meshgrid( - range(self.orientation_map.num_x), - range(self.orientation_map.num_y), - indexing = 'ij') + # # coordinates + # xa,ya = np.meshgrid( + # range(self.orientation_map.num_x), + # range(self.orientation_map.num_y), + # indexing = 'ij') # Loop over grains to determine number in each pixel for a0 in range(self.cluster_sizes.shape[0]): @@ -1759,8 +1759,19 @@ def cluster_orientation_map( im_stripe = im_count >= 2 im_single = np.logical_not(im_stripe) + # prefactor for stripes + if stripe_width[0] == 0: + dx = 0 + else: + dx = 1/stripe_width[0] + if stripe_width[1] == 0: + dy = 0 + else: + dy = 1/stripe_width[1] + + # loop over grains - for a0 in range(1): + for a0 in range(self.cluster_sizes.shape[0]): if self.cluster_sizes[a0] >= area_min: im_grain[:] = False im_grain[ @@ -1768,21 +1779,44 @@ def cluster_orientation_map( self.cluster_inds[a0][1,:], ] = True + # non-overlapping grains sub = np.logical_and( im_grain, im_single) - print(np.sum(sub)) - + x,y = np.squeeze(np.unravel_index(np.where(sub.ravel()), im_grain.shape)) + for a1 in range(x.size): + orientation_map.set_orientation( + self.cluster_orientation[a0], + x[a1], + y[a1]) + + # overlapping grains sub = np.logical_and( im_grain, im_stripe) - print(np.sum(sub)) - - - - - fig,ax = plt.subplots(figsize=(8,8)) - ax.imshow(im_stripe) + x,y = np.squeeze(np.unravel_index(np.where(sub.ravel()), im_grain.shape)) + for a1 in range(x.size): + d = np.mod( + x[a1]*dx + \ + y[a1]*dy + \ + im_mark[x[a1],y[a1]] + \ + + 0.5, + im_count[x[a1],y[a1]]) + if d < 1.0: + orientation_map.set_orientation( + self.cluster_orientation[a0], + x[a1], + y[a1]) + im_mark[x[a1],y[a1]] += 1 + + # print(np.sum(sub)) + # fig,ax = plt.subplots(figsize=(12,12)) + # ax.imshow(im_stripe) + # ax.scatter( + # y, + # x, + # s = 1, + # c = 'r') From a5cfd363811ef9acd874a25ca7052f15e0f76588 Mon Sep 17 00:00:00 2001 From: Colin Date: Fri, 12 May 2023 18:39:57 -0700 Subject: [PATCH 05/13] Cleaning up --- py4DSTEM/process/diffraction/crystal_ACOM.py | 27 +++++++++----------- 1 file changed, 12 insertions(+), 15 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index 6e11960f6..54b233ac4 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -1716,13 +1716,17 @@ def cluster_orientation_map( Parameters -------- - stripe_width: int + stripe_width: (int,ind) Width of strips in the overlapping regions. + Kind of janky but it mostly works! + area_min: (int) + Minimum size of grains to include Returns -------- orientation_map + The clustered orientation map """ @@ -1783,7 +1787,9 @@ def cluster_orientation_map( sub = np.logical_and( im_grain, im_single) - x,y = np.squeeze(np.unravel_index(np.where(sub.ravel()), im_grain.shape)) + x,y = np.unravel_index(np.where(sub.ravel()), im_grain.shape) + x = np.atleast_1d(np.squeeze(x)) + y = np.atleast_1d(np.squeeze(y)) for a1 in range(x.size): orientation_map.set_orientation( self.cluster_orientation[a0], @@ -1794,7 +1800,9 @@ def cluster_orientation_map( sub = np.logical_and( im_grain, im_stripe) - x,y = np.squeeze(np.unravel_index(np.where(sub.ravel()), im_grain.shape)) + x,y = np.unravel_index(np.where(sub.ravel()), im_grain.shape) + x = np.atleast_1d(np.squeeze(x)) + y = np.atleast_1d(np.squeeze(y)) for a1 in range(x.size): d = np.mod( x[a1]*dx + \ @@ -1802,24 +1810,13 @@ def cluster_orientation_map( im_mark[x[a1],y[a1]] + \ + 0.5, im_count[x[a1],y[a1]]) + if d < 1.0: orientation_map.set_orientation( self.cluster_orientation[a0], x[a1], y[a1]) im_mark[x[a1],y[a1]] += 1 - - # print(np.sum(sub)) - # fig,ax = plt.subplots(figsize=(12,12)) - # ax.imshow(im_stripe) - # ax.scatter( - # y, - # x, - # s = 1, - # c = 'r') - - - return orientation_map From 2ebe41473ceb62ed5fafa229e8fb46dee8e2fac9 Mon Sep 17 00:00:00 2001 From: Colin Date: Mon, 15 May 2023 09:56:21 -0700 Subject: [PATCH 06/13] Updating plotting function with new args --- py4DSTEM/process/diffraction/crystal_viz.py | 32 +++++++++------------ 1 file changed, 14 insertions(+), 18 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_viz.py b/py4DSTEM/process/diffraction/crystal_viz.py index 19bdc62ab..328e3ebb8 100644 --- a/py4DSTEM/process/diffraction/crystal_viz.py +++ b/py4DSTEM/process/diffraction/crystal_viz.py @@ -1752,6 +1752,7 @@ def plot_cluster_size( self, area_min = None, area_max = None, + area_step = 1, weight_intensity = False, pixel_area = 1.0, pixel_area_units = 'px^2', @@ -1764,9 +1765,11 @@ def plot_cluster_size( Parameters -------- area_min: int (optional) - Min area to include + Min area to include in pixels^2 area_max: int (optional) - Max area bin in pixels + Max area bin in pixels^2 + area_step: int (optional) + Step size of the histogram bin in pixels^2 weight_intensity: bool Weight histogram by the peak intensity. pixel_area: float @@ -1787,37 +1790,34 @@ def plot_cluster_size( if area_max is None: area_max = np.max(self.cluster_sizes) - area = np.arange(area_max) + area = np.arange(0,area_max,area_step) if area_min is None: sub = self.cluster_sizes.astype('int') < area_max else: sub = np.logical_and( - self.cluster_sizes.astype('int') < area_max, self.cluster_sizes.astype('int') >= area_min, + self.cluster_sizes.astype('int') < area_max ) - - if weight_intensity: hist = np.bincount( - self.cluster_sizes[sub], + self.cluster_sizes[sub] // area_step, weights = self.cluster_sig[sub], - minlength = area_max, + minlength = area.shape[0], ) else: hist = np.bincount( - self.cluster_sizes[sub], - minlength = area_max, + self.cluster_sizes[sub] // area_step, + minlength = area.shape[0], ) - - + # plotting fig,ax = plt.subplots(figsize = figsize) ax.bar( area * pixel_area, hist, - width = 0.8 * pixel_area, + width = 0.8 * pixel_area * area_step, ) - + ax.set_xlim((0,area_max*pixel_area)) ax.set_xlabel('Grain Area [' + pixel_area_units + ']') if weight_intensity: ax.set_ylabel('Total Signal [arb. units]') @@ -1828,10 +1828,6 @@ def plot_cluster_size( return fig,ax - - - - def axisEqual3D(ax): extents = np.array([getattr(ax, "get_{}lim".format(dim))() for dim in "xyz"]) sz = extents[:, 1] - extents[:, 0] From 5db507f3f511111690705e117bbd089f7d59fd17 Mon Sep 17 00:00:00 2001 From: Colin Date: Mon, 15 May 2023 09:58:20 -0700 Subject: [PATCH 07/13] Fix duplicate arg --- py4DSTEM/process/diffraction/crystal_ACOM.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index 5c0506457..29aa1251b 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -759,8 +759,7 @@ def match_orientations( multiple_corr_reset: bool = True, return_orientation: bool = True, progress_bar: bool = True, - return_orientation: bool = True, -): + ): ''' This function computes the orientation of any number of PointLists stored in a PointListArray, and returns an OrienationMap. From d621c19afc2fe664d83c0d1838782c3a901ff321 Mon Sep 17 00:00:00 2001 From: cophus Date: Sun, 24 Sep 2023 13:01:11 -0700 Subject: [PATCH 08/13] Adding min angle between matches. --- py4DSTEM/process/diffraction/crystal_ACOM.py | 103 +++++++++++++++---- 1 file changed, 83 insertions(+), 20 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index f8ff10bf5..aa1fdf093 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -763,14 +763,16 @@ def match_orientations( self, bragg_peaks_array: PointListArray, num_matches_return: int = 1, + min_angle_between_matches_deg = None, min_number_peaks: int = 3, inversion_symmetry: bool = True, multiple_corr_reset: bool = True, return_orientation: bool = True, progress_bar: bool = True, ): - ''' - This function computes the orientation of any number of PointLists stored in a PointListArray, and returns an OrienationMap. + """ + This function computes the orientation of any number of PointLists stored in a PointListArray, + and returns an OrienationMap. Options are the same as match_single_pattern(). """ orientation_map = OrientationMap( @@ -809,6 +811,7 @@ def match_orientations( orientation = self.match_single_pattern( bragg_peaks=vectors, num_matches_return=num_matches_return, + min_angle_between_matches_deg = min_angle_between_matches_deg, min_number_peaks=min_number_peaks, inversion_symmetry=inversion_symmetry, multiple_corr_reset=multiple_corr_reset, @@ -836,6 +839,7 @@ def match_single_pattern( self, bragg_peaks: PointList, num_matches_return: int = 1, + min_angle_between_matches_deg = None, min_number_peaks=3, inversion_symmetry=True, multiple_corr_reset=True, @@ -849,23 +853,42 @@ def match_single_pattern( """ Solve for the best fit orientation of a single diffraction pattern. - Args: - bragg_peaks (PointList): numpy array containing the Bragg positions and intensities ('qx', 'qy', 'intensity') - num_matches_return (int): return these many matches as 3th dim of orient (matrix) - min_number_peaks (int): Minimum number of peaks required to perform ACOM matching - inversion_symmetry (bool): check for inversion symmetry in the matches - multiple_corr_reset (bool): keep original correlation score for multiple matches - subpixel_tilt (bool): set to false for faster matching, returning the nearest corr point - plot_polar (bool): set to true to plot the polar transform of the diffraction pattern - plot_corr (bool): set to true to plot the resulting correlogram - returnfig (bool): Return figure handles - figsize (list): size of figure - verbose (bool): Print the fitted zone axes, correlation scores - CUDA (bool): Enable CUDA for the FFT steps + Parameters + -------- + bragg_peaks: PointList + numpy array containing the Bragg positions and intensities ('qx', 'qy', 'intensity') + num_matches_return: int + return these many matches as 3th dim of orient (matrix) + min_angle_between_matches_deg: int + Minimum angle between zone axis of multiple matches, in degrees. + Note that I haven't thought how to handle in-plane rotations, since multiple matches are possible. + min_number_peaks: int + Minimum number of peaks required to perform ACOM matching + inversion_symmetry bool + check for inversion symmetry in the matches + multiple_corr_reset bool + keep original correlation score for multiple matches + subpixel_tilt: bool + set to false for faster matching, returning the nearest corr point + plot_polar: bool + set to true to plot the polar transform of the diffraction pattern + plot_corr: bool + set to true to plot the resulting correlogram + returnfig: bool + return figure handles + figsize: list + size of figure + verbose: bool + Print the fitted zone axes, correlation scores + CUDA: bool + Enable CUDA for the FFT steps - Returns: - orientation (Orientation): Orientation class containing all outputs - fig, ax (handles): Figure handles for the plotting output + Returns + -------- + orientation: Orientation + Orientation class containing all outputs + fig, ax: handles + Figure handles for the plotting output """ # init orientation output @@ -1036,6 +1059,24 @@ def match_single_pattern( 0, ) + # If minimum angle is specified and we're on a match later than the first, + # we zero correlation values within the given range. + if min_angle_between_matches_deg is not None: + if match_ind > 0: + inds_previous = orientation.inds[:match_ind, 0] + for a0 in range(inds_previous.size): + mask_zero = np.arccos( + np.clip( + np.sum( + self.orientation_vecs * self.orientation_vecs[inds_previous[a0], :], + axis=1, + ), + -1, + 1, + ) + ) < np.deg2rad(min_angle_between_matches_deg) + corr_full[mask_zero,:] = 0.0 + # Get maximum (non inverted) correlation value ind_phi = np.argmax(corr_full, axis=1) @@ -1103,6 +1144,25 @@ def match_single_pattern( ), 0, ) + + # If minimum angle is specified and we're on a match later than the first, + # we zero correlation values within the given range. + if min_angle_between_matches_deg is not None: + if match_ind > 0: + inds_previous = orientation.inds[:match_ind, 0] + for a0 in range(inds_previous.size): + mask_zero = np.arccos( + np.clip( + np.sum( + self.orientation_vecs * self.orientation_vecs[inds_previous[a0], :], + axis=1, + ), + -1, + 1, + ) + ) < np.deg2rad(min_angle_between_matches_deg) + corr_full_inv[mask_zero,:] = 0.0 + ind_phi_inv = np.argmax(corr_full_inv, axis=1) corr_inv = np.zeros(self.orientation_num_zones, dtype="bool") @@ -1746,6 +1806,7 @@ def cluster_grains( # Main loop search = True + comp = 0.0 while search is True: inds_grain = np.argmax(sig) @@ -1757,8 +1818,10 @@ def cluster_grains( else: # progressbar if progress_bar: - comp = 1 - np.mean(np.max(mark,axis = 2)) - update_progress(comp) + new_comp = 1 - np.mean(np.max(mark,axis = 2)) + if new_comp > comp + 0.001: + comp = new_comp + update_progress(comp) # Start cluster x,y,z = np.unravel_index(inds_grain, sig.shape) From 65206a1f284769adf395d2080cb442d9a4c3c5e1 Mon Sep 17 00:00:00 2001 From: Steven Zeltmann Date: Wed, 27 Sep 2023 11:46:26 -0400 Subject: [PATCH 09/13] format with black --- py4DSTEM/process/diffraction/crystal_ACOM.py | 232 ++++++++++--------- py4DSTEM/process/diffraction/crystal_viz.py | 167 +++++++------ py4DSTEM/process/diffraction/utils.py | 1 - 3 files changed, 200 insertions(+), 200 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index aa1fdf093..e7b54f712 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -763,15 +763,15 @@ def match_orientations( self, bragg_peaks_array: PointListArray, num_matches_return: int = 1, - min_angle_between_matches_deg = None, + min_angle_between_matches_deg=None, min_number_peaks: int = 3, inversion_symmetry: bool = True, multiple_corr_reset: bool = True, return_orientation: bool = True, progress_bar: bool = True, - ): +): """ - This function computes the orientation of any number of PointLists stored in a PointListArray, + This function computes the orientation of any number of PointLists stored in a PointListArray, and returns an OrienationMap. Options are the same as match_single_pattern(). """ @@ -811,7 +811,7 @@ def match_orientations( orientation = self.match_single_pattern( bragg_peaks=vectors, num_matches_return=num_matches_return, - min_angle_between_matches_deg = min_angle_between_matches_deg, + min_angle_between_matches_deg=min_angle_between_matches_deg, min_number_peaks=min_number_peaks, inversion_symmetry=inversion_symmetry, multiple_corr_reset=multiple_corr_reset, @@ -835,11 +835,12 @@ def match_orientations( else: return + def match_single_pattern( self, bragg_peaks: PointList, num_matches_return: int = 1, - min_angle_between_matches_deg = None, + min_angle_between_matches_deg=None, min_number_peaks=3, inversion_symmetry=True, multiple_corr_reset=True, @@ -855,39 +856,39 @@ def match_single_pattern( Parameters -------- - bragg_peaks: PointList + bragg_peaks: PointList numpy array containing the Bragg positions and intensities ('qx', 'qy', 'intensity') - num_matches_return: int + num_matches_return: int return these many matches as 3th dim of orient (matrix) min_angle_between_matches_deg: int Minimum angle between zone axis of multiple matches, in degrees. Note that I haven't thought how to handle in-plane rotations, since multiple matches are possible. - min_number_peaks: int + min_number_peaks: int Minimum number of peaks required to perform ACOM matching - inversion_symmetry bool + inversion_symmetry bool check for inversion symmetry in the matches - multiple_corr_reset bool + multiple_corr_reset bool keep original correlation score for multiple matches - subpixel_tilt: bool + subpixel_tilt: bool set to false for faster matching, returning the nearest corr point - plot_polar: bool + plot_polar: bool set to true to plot the polar transform of the diffraction pattern - plot_corr: bool + plot_corr: bool set to true to plot the resulting correlogram - returnfig: bool + returnfig: bool return figure handles - figsize: list + figsize: list size of figure - verbose: bool + verbose: bool Print the fitted zone axes, correlation scores - CUDA: bool + CUDA: bool Enable CUDA for the FFT steps Returns -------- - orientation: Orientation + orientation: Orientation Orientation class containing all outputs - fig, ax: handles + fig, ax: handles Figure handles for the plotting output """ @@ -1068,14 +1069,15 @@ def match_single_pattern( mask_zero = np.arccos( np.clip( np.sum( - self.orientation_vecs * self.orientation_vecs[inds_previous[a0], :], + self.orientation_vecs + * self.orientation_vecs[inds_previous[a0], :], axis=1, ), -1, 1, ) ) < np.deg2rad(min_angle_between_matches_deg) - corr_full[mask_zero,:] = 0.0 + corr_full[mask_zero, :] = 0.0 # Get maximum (non inverted) correlation value ind_phi = np.argmax(corr_full, axis=1) @@ -1146,7 +1148,7 @@ def match_single_pattern( ) # If minimum angle is specified and we're on a match later than the first, - # we zero correlation values within the given range. + # we zero correlation values within the given range. if min_angle_between_matches_deg is not None: if match_ind > 0: inds_previous = orientation.inds[:match_ind, 0] @@ -1154,14 +1156,15 @@ def match_single_pattern( mask_zero = np.arccos( np.clip( np.sum( - self.orientation_vecs * self.orientation_vecs[inds_previous[a0], :], + self.orientation_vecs + * self.orientation_vecs[inds_previous[a0], :], axis=1, ), -1, 1, ) ) < np.deg2rad(min_angle_between_matches_deg) - corr_full_inv[mask_zero,:] = 0.0 + corr_full_inv[mask_zero, :] = 0.0 ind_phi_inv = np.argmax(corr_full_inv, axis=1) corr_inv = np.zeros(self.orientation_num_zones, dtype="bool") @@ -1750,17 +1753,13 @@ def match_single_pattern( return orientation - - - - def cluster_grains( self, - threshold_add = 1.0, - threshold_grow = 0.1, - angle_tolerance_deg = 5.0, - progress_bar = True, - ): + threshold_add=1.0, + threshold_grow=0.1, + angle_tolerance_deg=5.0, + progress_bar=True, +): """ Cluster grains using rotation criterion, and correlation values. @@ -1777,7 +1776,7 @@ def cluster_grains( Returns -------- - + """ @@ -1794,11 +1793,11 @@ def cluster_grains( matrix = self.orientation_map.matrix.copy() # init - self.cluster_sizes = np.array((), dtype='int') + self.cluster_sizes = np.array((), dtype="int") self.cluster_sig = np.array(()) self.cluster_inds = [] self.cluster_orientation = [] - inds_all = np.zeros_like(sig, dtype='int') + inds_all = np.zeros_like(sig, dtype="int") inds_all.ravel()[:] = np.arange(inds_all.size) # Tolerance @@ -1814,26 +1813,26 @@ def cluster_grains( if val < threshold_add: search = False - + else: # progressbar if progress_bar: - new_comp = 1 - np.mean(np.max(mark,axis = 2)) + new_comp = 1 - np.mean(np.max(mark, axis=2)) if new_comp > comp + 0.001: comp = new_comp update_progress(comp) # Start cluster - x,y,z = np.unravel_index(inds_grain, sig.shape) - mark[x,y,z] = False - sig[x,y,z] = 0 - matrix_cluster = matrix[x,y,z] - orientation_cluster = self.orientation_map.get_orientation_single(x,y,z) + x, y, z = np.unravel_index(inds_grain, sig.shape) + mark[x, y, z] = False + sig[x, y, z] = 0 + matrix_cluster = matrix[x, y, z] + orientation_cluster = self.orientation_map.get_orientation_single(x, y, z) # Neighbors to search - xr = np.clip(x + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1) - yr = np.clip(y + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1) - inds_cand = inds_all[xr[:,None],yr[None],:].ravel() + xr = np.clip(x + np.arange(-1, 2, dtype="int"), 0, sig.shape[0] - 1) + yr = np.clip(y + np.arange(-1, 2, dtype="int"), 0, sig.shape[1] - 1) + inds_cand = inds_all[xr[:, None], yr[None], :].ravel() inds_cand = np.delete(inds_cand, mark.ravel()[inds_cand] == False) if inds_cand.size == 0: @@ -1843,39 +1842,59 @@ def cluster_grains( # grow the cluster while grow is True: - inds_new = np.array((),dtype='int') + inds_new = np.array((), dtype="int") - keep = np.zeros(inds_cand.size, dtype='bool') + keep = np.zeros(inds_cand.size, dtype="bool") for a0 in range(inds_cand.size): - xc,yc,zc = np.unravel_index(inds_cand[a0], sig.shape) + xc, yc, zc = np.unravel_index(inds_cand[a0], sig.shape) # Angle test between orientation matrices - dphi = np.min(np.arccos(np.clip((np.trace( - self.symmetry_operators @ matrix[xc,yc,zc] \ - @ np.transpose(matrix_cluster), - axis1=1, - axis2=2)-1)/2,-1,1))) + dphi = np.min( + np.arccos( + np.clip( + ( + np.trace( + self.symmetry_operators + @ matrix[xc, yc, zc] + @ np.transpose(matrix_cluster), + axis1=1, + axis2=2, + ) + - 1 + ) + / 2, + -1, + 1, + ) + ) + ) if np.abs(dphi) < tol: keep[a0] = True - sig[xc,yc,zc] = 0 - mark[xc,yc,zc] = False + sig[xc, yc, zc] = 0 + mark[xc, yc, zc] = False - xr = np.clip(xc + np.arange(-1,2,dtype='int'), 0, sig.shape[0] - 1) - yr = np.clip(yc + np.arange(-1,2,dtype='int'), 0, sig.shape[1] - 1) - inds_add = inds_all[xr[:,None],yr[None],:].ravel() + xr = np.clip( + xc + np.arange(-1, 2, dtype="int"), 0, sig.shape[0] - 1 + ) + yr = np.clip( + yc + np.arange(-1, 2, dtype="int"), 0, sig.shape[1] - 1 + ) + inds_add = inds_all[xr[:, None], yr[None], :].ravel() inds_new = np.append(inds_new, inds_add) inds_grain = np.append(inds_grain, inds_cand[keep]) - inds_cand = np.unique(np.delete(inds_new, mark.ravel()[inds_new] == False)) + inds_cand = np.unique( + np.delete(inds_new, mark.ravel()[inds_new] == False) + ) if inds_cand.size == 0: grow = False # convert grain to x,y coordinates, add = list - xg,yg,zg = np.unravel_index(inds_grain, sig.shape) - xyg = np.unique(np.vstack((xg,yg)), axis = 1) + xg, yg, zg = np.unravel_index(inds_grain, sig.shape) + xyg = np.unique(np.vstack((xg, yg)), axis=1) sig_mean = np.mean(sig_init.ravel()[inds_grain]) self.cluster_sizes = np.append(self.cluster_sizes, xyg.shape[1]) self.cluster_sig = np.append(self.cluster_sig, sig_mean) @@ -1887,13 +1906,11 @@ def cluster_grains( update_progress(1) - - def cluster_orientation_map( self, - stripe_width = (2,2), - area_min = 2, - ): + stripe_width=(2, 2), + area_min=2, +): """ Produce a new orientation map from the clustered grains. Use a stripe pattern for the overlapping grains. @@ -1908,7 +1925,7 @@ def cluster_orientation_map( Returns -------- - + orientation_map The clustered orientation map @@ -1916,18 +1933,15 @@ def cluster_orientation_map( # init orientation_map = OrientationMap( - num_x = self.orientation_map.num_x, - num_y = self.orientation_map.num_y, - num_matches=1) - im_grain = np.zeros(( - self.orientation_map.num_x, - self.orientation_map.num_y), dtype='bool') - im_count = np.zeros(( - self.orientation_map.num_x, - self.orientation_map.num_y)) - im_mark = np.zeros(( - self.orientation_map.num_x, - self.orientation_map.num_y)) + num_x=self.orientation_map.num_x, + num_y=self.orientation_map.num_y, + num_matches=1, + ) + im_grain = np.zeros( + (self.orientation_map.num_x, self.orientation_map.num_y), dtype="bool" + ) + im_count = np.zeros((self.orientation_map.num_x, self.orientation_map.num_y)) + im_mark = np.zeros((self.orientation_map.num_x, self.orientation_map.num_y)) # # coordinates # xa,ya = np.meshgrid( @@ -1940,8 +1954,8 @@ def cluster_orientation_map( if self.cluster_sizes[a0] >= area_min: im_grain[:] = False im_grain[ - self.cluster_inds[a0][0,:], - self.cluster_inds[a0][1,:], + self.cluster_inds[a0][0, :], + self.cluster_inds[a0][1, :], ] = True im_count += im_grain im_stripe = im_count >= 2 @@ -1951,56 +1965,47 @@ def cluster_orientation_map( if stripe_width[0] == 0: dx = 0 else: - dx = 1/stripe_width[0] + dx = 1 / stripe_width[0] if stripe_width[1] == 0: dy = 0 else: - dy = 1/stripe_width[1] - + dy = 1 / stripe_width[1] # loop over grains for a0 in range(self.cluster_sizes.shape[0]): if self.cluster_sizes[a0] >= area_min: im_grain[:] = False im_grain[ - self.cluster_inds[a0][0,:], - self.cluster_inds[a0][1,:], + self.cluster_inds[a0][0, :], + self.cluster_inds[a0][1, :], ] = True # non-overlapping grains - sub = np.logical_and( - im_grain, - im_single) - x,y = np.unravel_index(np.where(sub.ravel()), im_grain.shape) + sub = np.logical_and(im_grain, im_single) + x, y = np.unravel_index(np.where(sub.ravel()), im_grain.shape) x = np.atleast_1d(np.squeeze(x)) y = np.atleast_1d(np.squeeze(y)) for a1 in range(x.size): orientation_map.set_orientation( - self.cluster_orientation[a0], - x[a1], - y[a1]) - + self.cluster_orientation[a0], x[a1], y[a1] + ) + # overlapping grains - sub = np.logical_and( - im_grain, - im_stripe) - x,y = np.unravel_index(np.where(sub.ravel()), im_grain.shape) + sub = np.logical_and(im_grain, im_stripe) + x, y = np.unravel_index(np.where(sub.ravel()), im_grain.shape) x = np.atleast_1d(np.squeeze(x)) y = np.atleast_1d(np.squeeze(y)) for a1 in range(x.size): d = np.mod( - x[a1]*dx + \ - y[a1]*dy + \ - im_mark[x[a1],y[a1]] + \ - + 0.5, - im_count[x[a1],y[a1]]) + x[a1] * dx + y[a1] * dy + im_mark[x[a1], y[a1]] + +0.5, + im_count[x[a1], y[a1]], + ) if d < 1.0: orientation_map.set_orientation( - self.cluster_orientation[a0], - x[a1], - y[a1]) - im_mark[x[a1],y[a1]] += 1 + self.cluster_orientation[a0], x[a1], y[a1] + ) + im_mark[x[a1], y[a1]] += 1 return orientation_map @@ -2545,11 +2550,10 @@ def symmetry_reduce_directions( # "-3m": ["fiber", [0, 0, 1], [180.0, 30.0]], - # Progressbar taken from stackexchange: # https://stackoverflow.com/questions/3160699/python-progress-bar def update_progress(progress): - barLength = 60 # Modify this to change the length of the progress bar + barLength = 60 # Modify this to change the length of the progress bar status = "" if isinstance(progress, int): progress = float(progress) @@ -2562,9 +2566,9 @@ def update_progress(progress): if progress >= 1: progress = 1 status = "Done\r\n" - block = int(round(barLength*progress)) - text = "\rPercent: [{0}] {1}% {2}".format( "#"*block + "-"*(barLength-block), - np.round(progress*100,2), - status) + block = int(round(barLength * progress)) + text = "\rPercent: [{0}] {1}% {2}".format( + "#" * block + "-" * (barLength - block), np.round(progress * 100, 2), status + ) sys.stdout.write(text) sys.stdout.flush() diff --git a/py4DSTEM/process/diffraction/crystal_viz.py b/py4DSTEM/process/diffraction/crystal_viz.py index 1c68d5ba6..e17e87b93 100644 --- a/py4DSTEM/process/diffraction/crystal_viz.py +++ b/py4DSTEM/process/diffraction/crystal_viz.py @@ -991,7 +991,7 @@ def overline(x): def plot_orientation_maps( self, - orientation_map = None, + orientation_map=None, orientation_ind: int = 0, dir_in_plane_degrees: float = 0.0, corr_range: np.ndarray = np.array([0, 5]), @@ -1725,17 +1725,18 @@ def plot_fiber_orientation_maps( else: return images_orientation + def plot_clusters( self, - area_min = 2, - outline_grains = True, - outline_thickness = 1, - fill_grains = 0.25, - smooth_grains = 1.0, - cmap = 'viridis', - figsize = (8,8), - returnfig = False, - ): + area_min=2, + outline_grains=True, + outline_thickness=1, + fill_grains=0.25, + smooth_grains=1.0, + cmap="viridis", + figsize=(8, 8), + returnfig=False, +): """ Plot the clusters as an image. @@ -1761,17 +1762,22 @@ def plot_clusters( fig, ax (optional) Figure and axes handles - """ - + """ + # init - im_plot = np.zeros(( - self.orientation_map.num_x, - self.orientation_map.num_y, - )) - im_grain = np.zeros(( - self.orientation_map.num_x, - self.orientation_map.num_y, - ), dtype='bool') + im_plot = np.zeros( + ( + self.orientation_map.num_x, + self.orientation_map.num_y, + ) + ) + im_grain = np.zeros( + ( + self.orientation_map.num_x, + self.orientation_map.num_y, + ), + dtype="bool", + ) # make plotting image @@ -1780,35 +1786,30 @@ def plot_clusters( if outline_grains: im_grain[:] = False im_grain[ - self.cluster_inds[a0][0,:], - self.cluster_inds[a0][1,:], + self.cluster_inds[a0][0, :], + self.cluster_inds[a0][1, :], ] = True - im_dist = \ - distance_transform_edt( - erosion(np.invert(im_grain),footprint=np.ones((3,3),dtype='bool')) - ) \ - - \ - distance_transform_edt( - im_grain - ) - im_dist = gaussian_filter( - im_dist, - sigma = smooth_grains, - mode = 'nearest') - im_add = np.exp(im_dist**2 / (-0.5*outline_thickness**2)) - + im_dist = distance_transform_edt( + erosion( + np.invert(im_grain), footprint=np.ones((3, 3), dtype="bool") + ) + ) - distance_transform_edt(im_grain) + im_dist = gaussian_filter(im_dist, sigma=smooth_grains, mode="nearest") + im_add = np.exp(im_dist**2 / (-0.5 * outline_thickness**2)) + if fill_grains > 0: - im_dist = \ - distance_transform_edt( - erosion(np.invert(im_grain),footprint=np.ones((3,3),dtype='bool')) + im_dist = distance_transform_edt( + erosion( + np.invert(im_grain), footprint=np.ones((3, 3), dtype="bool") + ) ) im_dist = gaussian_filter( - im_dist, - sigma = smooth_grains, - mode = 'nearest') - im_add += fill_grains*np.exp(im_dist**2 / (-0.5*outline_thickness**2)) - + im_dist, sigma=smooth_grains, mode="nearest" + ) + im_add += fill_grains * np.exp( + im_dist**2 / (-0.5 * outline_thickness**2) + ) # im_add = 1 - np.exp( # distance_transform_edt(im_grain)**2 \ @@ -1819,13 +1820,12 @@ def plot_clusters( # xg,yg = np.unravel_index(self.cluster_inds[a0], im_plot.shape) im_grain[:] = False im_grain[ - self.cluster_inds[a0][0,:], - self.cluster_inds[a0][1,:], + self.cluster_inds[a0][0, :], + self.cluster_inds[a0][1, :], ] = True im_plot += gaussian_filter( - im_grain.astype('float'), - sigma = smooth_grains, - mode = 'nearest') + im_grain.astype("float"), sigma=smooth_grains, mode="nearest" + ) # im_plot[ # self.cluster_inds[a0][0,:], @@ -1833,32 +1833,29 @@ def plot_clusters( # ] += 1 if outline_grains: - im_plot = np.clip(im_plot,0,2) - + im_plot = np.clip(im_plot, 0, 2) # plotting - fig,ax = plt.subplots(figsize=figsize) + fig, ax = plt.subplots(figsize=figsize) ax.imshow( im_plot, # vmin = -3, # vmax = 3, - cmap = cmap, - ) - - + cmap=cmap, + ) def plot_cluster_size( self, - area_min = None, - area_max = None, - area_step = 1, - weight_intensity = False, - pixel_area = 1.0, - pixel_area_units = 'px^2', - figsize = (8,6), - returnfig = False, - ): + area_min=None, + area_max=None, + area_step=1, + weight_intensity=False, + pixel_area=1.0, + pixel_area_units="px^2", + figsize=(8, 6), + returnfig=False, +): """ Plot the cluster sizes @@ -1875,7 +1872,7 @@ def plot_cluster_size( pixel_area: float Size of pixel area unit square pixel_area_units: string - Units of the pixel area + Units of the pixel area figsize: tuple Size of the figure panel returnfig: bool @@ -1890,42 +1887,42 @@ def plot_cluster_size( if area_max is None: area_max = np.max(self.cluster_sizes) - area = np.arange(0,area_max,area_step) + area = np.arange(0, area_max, area_step) if area_min is None: - sub = self.cluster_sizes.astype('int') < area_max + sub = self.cluster_sizes.astype("int") < area_max else: sub = np.logical_and( - self.cluster_sizes.astype('int') >= area_min, - self.cluster_sizes.astype('int') < area_max - ) + self.cluster_sizes.astype("int") >= area_min, + self.cluster_sizes.astype("int") < area_max, + ) if weight_intensity: hist = np.bincount( self.cluster_sizes[sub] // area_step, - weights = self.cluster_sig[sub], - minlength = area.shape[0], - ) + weights=self.cluster_sig[sub], + minlength=area.shape[0], + ) else: hist = np.bincount( self.cluster_sizes[sub] // area_step, - minlength = area.shape[0], - ) - + minlength=area.shape[0], + ) + # plotting - fig,ax = plt.subplots(figsize = figsize) + fig, ax = plt.subplots(figsize=figsize) ax.bar( area * pixel_area, hist, - width = 0.8 * pixel_area * area_step, - ) - ax.set_xlim((0,area_max*pixel_area)) - ax.set_xlabel('Grain Area [' + pixel_area_units + ']') + width=0.8 * pixel_area * area_step, + ) + ax.set_xlim((0, area_max * pixel_area)) + ax.set_xlabel("Grain Area [" + pixel_area_units + "]") if weight_intensity: - ax.set_ylabel('Total Signal [arb. units]') + ax.set_ylabel("Total Signal [arb. units]") else: - ax.set_ylabel('Number of Grains') + ax.set_ylabel("Number of Grains") if returnfig: - return fig,ax + return fig, ax def axisEqual3D(ax): diff --git a/py4DSTEM/process/diffraction/utils.py b/py4DSTEM/process/diffraction/utils.py index aaf4c023c..cfb11f044 100644 --- a/py4DSTEM/process/diffraction/utils.py +++ b/py4DSTEM/process/diffraction/utils.py @@ -26,7 +26,6 @@ def __post_init__(self): self.angles = np.zeros((self.num_matches, 3)) - @dataclass class OrientationMap: """ From 23a2e946bc7c39e9da68377d6776682710aab5df Mon Sep 17 00:00:00 2001 From: Steven Zeltmann Date: Mon, 2 Oct 2023 15:18:05 -0400 Subject: [PATCH 10/13] update docstrings --- py4DSTEM/process/diffraction/crystal_ACOM.py | 28 ++++++++++++++------ 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index e7b54f712..54d78409a 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -771,8 +771,26 @@ def match_orientations( progress_bar: bool = True, ): """ - This function computes the orientation of any number of PointLists stored in a PointListArray, - and returns an OrienationMap. Options are the same as match_single_pattern(). + Parameters + -------- + bragg_peaks_array: PointListArray + PointListArray containing the Bragg peaks and intensities, with calibrations applied + num_matches_return: int + return these many matches as 3th dim of orient (matrix) + min_angle_between_matches_deg: int + Minimum angle between zone axis of multiple matches, in degrees. + Note that I haven't thought how to handle in-plane rotations, since multiple matches are possible. + min_number_peaks: int + Minimum number of peaks required to perform ACOM matching + inversion_symmetry: bool + check for inversion symmetry in the matches + multiple_corr_reset: bool + keep original correlation score for multiple matches + return_orientation: bool + Return orientation map from function for inspection. + The map is always stored in the Crystal object. + progress_bar: bool + Show or hide the progress bar """ orientation_map = OrientationMap( @@ -820,12 +838,6 @@ def match_orientations( ) orientation_map.set_orientation(orientation, rx, ry) - self.orientation_map = orientation_map - - if return_orientation: - return orientation_map - else: - return # assign and return self.orientation_map = orientation_map From 17710bd5c5c7b6fe0063cac51e91fce04da4f379 Mon Sep 17 00:00:00 2001 From: Steven Zeltmann Date: Mon, 2 Oct 2023 15:56:03 -0400 Subject: [PATCH 11/13] replace pbar with tqdm --- py4DSTEM/process/diffraction/crystal_ACOM.py | 39 ++++++++------------ 1 file changed, 15 insertions(+), 24 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index 54d78409a..75c729865 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -3,6 +3,7 @@ import os from typing import Union, Optional import time, sys +from tqdm import tqdm from emdfile import tqdmnd, PointList, PointListArray from py4DSTEM.data import RealSlice @@ -15,7 +16,7 @@ try: import cupy as cp -except: +except ModuleNotFoundError: cp = None @@ -1777,19 +1778,15 @@ def cluster_grains( Parameters -------- - corr_threshold_add: float + threshold_add: float Minimum signal required for a probe position to initialize a cluster. - corr_threshold_grow: float + threshold_grow: float Minimum signal required for a probe position to be added to a cluster. angle_tolerance_deg: float Rotation rolerance for clustering grains. progress_bar: bool Turns on the progress bar for the polar transformation - Returns - -------- - - """ # symmetry operators @@ -1818,6 +1815,7 @@ def cluster_grains( # Main loop search = True comp = 0.0 + pbar = tqdm(total=N, display = not progress_bar) while search is True: inds_grain = np.argmax(sig) @@ -1828,11 +1826,11 @@ def cluster_grains( else: # progressbar - if progress_bar: - new_comp = 1 - np.mean(np.max(mark, axis=2)) - if new_comp > comp + 0.001: - comp = new_comp - update_progress(comp) + # if progress_bar: + # new_comp = 1 - np.mean(np.max(mark, axis=2)) + # if new_comp > comp + 0.001: + # comp = new_comp + # update_progress(comp) # Start cluster x, y, z = np.unravel_index(inds_grain, sig.shape) @@ -1896,6 +1894,8 @@ def cluster_grains( inds_add = inds_all[xr[:, None], yr[None], :].ravel() inds_new = np.append(inds_new, inds_add) + pbar.update(inds_add.size) + inds_grain = np.append(inds_grain, inds_cand[keep]) inds_cand = np.unique( np.delete(inds_new, mark.ravel()[inds_new] == False) @@ -1913,9 +1913,7 @@ def cluster_grains( self.cluster_orientation.append(orientation_cluster) self.cluster_inds.append(xyg) - # finish progressbar - if progress_bar: - update_progress(1) + pbar.close() def cluster_orientation_map( @@ -1929,9 +1927,8 @@ def cluster_orientation_map( Parameters -------- - stripe_width: (int,ind) - Width of strips in the overlapping regions. - Kind of janky but it mostly works! + stripe_width: (int,int) + Width of stripes for plotting maps with overlapping grains area_min: (int) Minimum size of grains to include @@ -1955,12 +1952,6 @@ def cluster_orientation_map( im_count = np.zeros((self.orientation_map.num_x, self.orientation_map.num_y)) im_mark = np.zeros((self.orientation_map.num_x, self.orientation_map.num_y)) - # # coordinates - # xa,ya = np.meshgrid( - # range(self.orientation_map.num_x), - # range(self.orientation_map.num_y), - # indexing = 'ij') - # Loop over grains to determine number in each pixel for a0 in range(self.cluster_sizes.shape[0]): if self.cluster_sizes[a0] >= area_min: From 221b34b95d98942fce60f9e55687a06585a33b1a Mon Sep 17 00:00:00 2001 From: Steven Zeltmann Date: Wed, 4 Oct 2023 12:32:37 -0400 Subject: [PATCH 12/13] tqdm-based progress bar in clustering --- py4DSTEM/process/diffraction/crystal_ACOM.py | 40 +++----------------- 1 file changed, 6 insertions(+), 34 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index 75c729865..115f8140c 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -788,7 +788,7 @@ def match_orientations( multiple_corr_reset: bool keep original correlation score for multiple matches return_orientation: bool - Return orientation map from function for inspection. + Return orientation map from function for inspection. The map is always stored in the Crystal object. progress_bar: bool Show or hide the progress bar @@ -1815,7 +1815,8 @@ def cluster_grains( # Main loop search = True comp = 0.0 - pbar = tqdm(total=N, display = not progress_bar) + mark_total = np.sum(np.max(mark, axis=2)) + pbar = tqdm(total=mark_total, disable=not progress_bar) while search is True: inds_grain = np.argmax(sig) @@ -1825,12 +1826,9 @@ def cluster_grains( search = False else: - # progressbar - # if progress_bar: - # new_comp = 1 - np.mean(np.max(mark, axis=2)) - # if new_comp > comp + 0.001: - # comp = new_comp - # update_progress(comp) + new_marks = mark_total - np.sum(np.max(mark, axis=2)) + pbar.update(new_marks) + mark_total -= new_marks # Start cluster x, y, z = np.unravel_index(inds_grain, sig.shape) @@ -1894,8 +1892,6 @@ def cluster_grains( inds_add = inds_all[xr[:, None], yr[None], :].ravel() inds_new = np.append(inds_new, inds_add) - pbar.update(inds_add.size) - inds_grain = np.append(inds_grain, inds_cand[keep]) inds_cand = np.unique( np.delete(inds_new, mark.ravel()[inds_new] == False) @@ -2551,27 +2547,3 @@ def symmetry_reduce_directions( # "-3m": ["fiber", [0, 0, 1], [90.0, 60.0]], # "-3m": ["fiber", [0, 0, 1], [180.0, 30.0]], - - -# Progressbar taken from stackexchange: -# https://stackoverflow.com/questions/3160699/python-progress-bar -def update_progress(progress): - barLength = 60 # Modify this to change the length of the progress bar - status = "" - if isinstance(progress, int): - progress = float(progress) - if not isinstance(progress, float): - progress = 0 - status = "error: progress var must be float\r\n" - if progress < 0: - progress = 0 - status = "Halt...\r\n" - if progress >= 1: - progress = 1 - status = "Done\r\n" - block = int(round(barLength * progress)) - text = "\rPercent: [{0}] {1}% {2}".format( - "#" * block + "-" * (barLength - block), np.round(progress * 100, 2), status - ) - sys.stdout.write(text) - sys.stdout.flush() From a48093f9a9b8bc398bc9b94f3dba4519be9dc31d Mon Sep 17 00:00:00 2001 From: Steven Zeltmann Date: Thu, 5 Oct 2023 09:23:51 -0400 Subject: [PATCH 13/13] move progress update step --- py4DSTEM/process/diffraction/crystal_ACOM.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index 115f8140c..da553456f 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -1826,10 +1826,6 @@ def cluster_grains( search = False else: - new_marks = mark_total - np.sum(np.max(mark, axis=2)) - pbar.update(new_marks) - mark_total -= new_marks - # Start cluster x, y, z = np.unravel_index(inds_grain, sig.shape) mark[x, y, z] = False @@ -1909,6 +1905,11 @@ def cluster_grains( self.cluster_orientation.append(orientation_cluster) self.cluster_inds.append(xyg) + # update progressbar + new_marks = mark_total - np.sum(np.max(mark, axis=2)) + pbar.update(new_marks) + mark_total -= new_marks + pbar.close()