From a4ab3a141d1d3af60c947b87cfdffd1807cc2d44 Mon Sep 17 00:00:00 2001 From: Shinjan Dutta Date: Tue, 26 Mar 2024 12:39:22 -0500 Subject: [PATCH 1/7] Fix for the no-correlation flag to the ACOM orientation plan --- py4DSTEM/process/diffraction/crystal_ACOM.py | 263 ++++++++++--------- 1 file changed, 132 insertions(+), 131 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index 7ac33957e..5bec155f2 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -550,142 +550,143 @@ def orientation_plan( ).astype("int") self.orientation_inds[:, 2] = orientation_sector - # If needed, create coarse orientation sieve - if self.orientation_refine: - self.orientation_sieve = np.logical_and( - np.mod(self.orientation_inds[:, 0], self.orientation_refine_ratio) == 0, - np.mod(self.orientation_inds[:, 1], self.orientation_refine_ratio) == 0, - ) - if self.CUDA: - self.orientation_sieve_CUDA = cp.asarray(self.orientation_sieve) - - # Convert to spherical coordinates - elev = np.arctan2( - np.hypot(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]), - self.orientation_vecs[:, 2], - ) - # azim = np.pi / 2 + np.arctan2( - # self.orientation_vecs[:, 1], self.orientation_vecs[:, 0] - # ) - azim = np.arctan2(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]) - - # 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 - ) - - # Calculate -z angles (Euler angle 3) - self.orientation_gamma = np.linspace( - 0, 2 * np.pi, self.orientation_in_plane_steps, endpoint=False - ) - - # Determine the radii of all spherical shells - radii_test = np.round(self.g_vec_leng / tol_distance) * tol_distance - radii = np.unique(radii_test) - # Remove zero beam - keep = np.abs(radii) > tol_distance - self.orientation_shell_radii = radii[keep] - - # init - self.orientation_shell_index = -1 * np.ones(self.g_vec_all.shape[1], dtype="int") - self.orientation_shell_count = np.zeros(self.orientation_shell_radii.size) - - # Assign each structure factor point to a radial shell - for a0 in range(self.orientation_shell_radii.size): - sub = np.abs(self.orientation_shell_radii[a0] - radii_test) <= tol_distance / 2 - - self.orientation_shell_index[sub] = a0 - self.orientation_shell_count[a0] = np.sum(sub) - self.orientation_shell_radii[a0] = np.mean(self.g_vec_leng[sub]) - - # init storage arrays - self.orientation_rotation_angles = np.zeros((self.orientation_num_zones, 3)) - self.orientation_rotation_matrices = np.zeros((self.orientation_num_zones, 3, 3)) - - # If possible, Get symmetry operations for this spacegroup, store in matrix form - if self.pymatgen_available: - # get operators - ops = self.pointgroup.get_point_group_operations() - - # Inverse of lattice - zone_axis_range_inv = np.linalg.inv(self.orientation_zone_axis_range) - - # init - num_sym = len(ops) - self.symmetry_operators = np.zeros((num_sym, 3, 3)) - self.symmetry_reduction = np.zeros((num_sym, 3, 3)) - - # calculate symmetry and reduction matrices - for a0 in range(num_sym): - self.symmetry_operators[a0] = ( - self.lat_inv.T @ ops[a0].rotation_matrix.T @ self.lat_real + # Calculate rotation matrices for zone axes + for a0 in np.arange(self.orientation_num_zones): + m1z = np.array( + [ + [np.cos(azim[a0]), np.sin(azim[a0]), 0], + [-np.sin(azim[a0]), np.cos(azim[a0]), 0], + [0, 0, 1], + ] ) - self.symmetry_reduction[a0] = ( - zone_axis_range_inv.T @ self.symmetry_operators[a0] - ).T - - # Remove duplicates - keep = np.ones(num_sym, dtype="bool") - for a0 in range(num_sym): - if keep[a0]: - diff = np.sum( - np.abs(self.symmetry_operators - self.symmetry_operators[a0]), - axis=(1, 2), - ) - sub = diff < 1e-3 - sub[: a0 + 1] = False - keep[sub] = False - self.symmetry_operators = self.symmetry_operators[keep] - self.symmetry_reduction = self.symmetry_reduction[keep] - - if ( - self.orientation_fiber_angles is not None - and np.abs(self.orientation_fiber_angles[0] - 180.0) < 1e-3 - ): - zone_axis_range_flip = self.orientation_zone_axis_range.copy() - zone_axis_range_flip[0, :] = -1 * zone_axis_range_flip[0, :] - zone_axis_range_inv = np.linalg.inv(zone_axis_range_flip) - - num_sym = self.symmetry_operators.shape[0] - self.symmetry_operators = np.tile(self.symmetry_operators, [2, 1, 1]) - self.symmetry_reduction = np.tile(self.symmetry_reduction, [2, 1, 1]) - - for a0 in range(num_sym): - self.symmetry_reduction[a0 + num_sym] = ( - zone_axis_range_inv.T @ self.symmetry_operators[a0 + num_sym] - ).T - - # Calculate rotation matrices for zone axes - for a0 in np.arange(self.orientation_num_zones): - m1z = np.array( - [ - [np.cos(azim[a0]), np.sin(azim[a0]), 0], - [-np.sin(azim[a0]), np.cos(azim[a0]), 0], - [0, 0, 1], - ] + m2x = np.array( + [ + [1, 0, 0], + [0, np.cos(elev[a0]), np.sin(elev[a0])], + [0, -np.sin(elev[a0]), np.cos(elev[a0])], + ] + ) + m3z = np.array( + [ + [np.cos(azim[a0]), -np.sin(azim[a0]), 0], + [np.sin(azim[a0]), np.cos(azim[a0]), 0], + [0, 0, 1], + ] + ) + self.orientation_rotation_matrices[a0, :, :] = m1z @ m2x @ m3z + self.orientation_rotation_angles[a0, :] = [azim[a0], elev[a0], -azim[a0]] + + # Rest of this function is only required for calculating orientation arrays + if calculate_correlation_array: + # If needed, create coarse orientation sieve + if self.orientation_refine: + self.orientation_sieve = np.logical_and( + np.mod(self.orientation_inds[:, 0], self.orientation_refine_ratio) == 0, + np.mod(self.orientation_inds[:, 1], self.orientation_refine_ratio) == 0, + ) + if self.CUDA: + self.orientation_sieve_CUDA = cp.asarray(self.orientation_sieve) + + # Convert to spherical coordinates + elev = np.arctan2( + np.hypot(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]), + self.orientation_vecs[:, 2], ) - m2x = np.array( - [ - [1, 0, 0], - [0, np.cos(elev[a0]), np.sin(elev[a0])], - [0, -np.sin(elev[a0]), np.cos(elev[a0])], - ] + # azim = np.pi / 2 + np.arctan2( + # self.orientation_vecs[:, 1], self.orientation_vecs[:, 0] + # ) + azim = np.arctan2(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]) + + # 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 ) - m3z = np.array( - [ - [np.cos(azim[a0]), -np.sin(azim[a0]), 0], - [np.sin(azim[a0]), np.cos(azim[a0]), 0], - [0, 0, 1], - ] + + # Calculate -z angles (Euler angle 3) + self.orientation_gamma = np.linspace( + 0, 2 * np.pi, self.orientation_in_plane_steps, endpoint=False ) - self.orientation_rotation_matrices[a0, :, :] = m1z @ m2x @ m3z - self.orientation_rotation_angles[a0, :] = [azim[a0], elev[a0], -azim[a0]] - - # Calculate reference arrays for all orientations - k0 = np.array([0.0, 0.0, -1.0 / self.wavelength]) - n = np.array([0.0, 0.0, -1.0]) + + # Determine the radii of all spherical shells + radii_test = np.round(self.g_vec_leng / tol_distance) * tol_distance + radii = np.unique(radii_test) + # Remove zero beam + keep = np.abs(radii) > tol_distance + self.orientation_shell_radii = radii[keep] + + # init + self.orientation_shell_index = -1 * np.ones(self.g_vec_all.shape[1], dtype="int") + self.orientation_shell_count = np.zeros(self.orientation_shell_radii.size) + + # Assign each structure factor point to a radial shell + for a0 in range(self.orientation_shell_radii.size): + sub = np.abs(self.orientation_shell_radii[a0] - radii_test) <= tol_distance / 2 + + self.orientation_shell_index[sub] = a0 + self.orientation_shell_count[a0] = np.sum(sub) + self.orientation_shell_radii[a0] = np.mean(self.g_vec_leng[sub]) + + # init storage arrays + self.orientation_rotation_angles = np.zeros((self.orientation_num_zones, 3)) + self.orientation_rotation_matrices = np.zeros((self.orientation_num_zones, 3, 3)) + + # If possible, Get symmetry operations for this spacegroup, store in matrix form + if self.pymatgen_available: + # get operators + ops = self.pointgroup.get_point_group_operations() + + # Inverse of lattice + zone_axis_range_inv = np.linalg.inv(self.orientation_zone_axis_range) + + # init + num_sym = len(ops) + self.symmetry_operators = np.zeros((num_sym, 3, 3)) + self.symmetry_reduction = np.zeros((num_sym, 3, 3)) + + # calculate symmetry and reduction matrices + for a0 in range(num_sym): + self.symmetry_operators[a0] = ( + self.lat_inv.T @ ops[a0].rotation_matrix.T @ self.lat_real + ) + self.symmetry_reduction[a0] = ( + zone_axis_range_inv.T @ self.symmetry_operators[a0] + ).T + + # Remove duplicates + keep = np.ones(num_sym, dtype="bool") + for a0 in range(num_sym): + if keep[a0]: + diff = np.sum( + np.abs(self.symmetry_operators - self.symmetry_operators[a0]), + axis=(1, 2), + ) + sub = diff < 1e-3 + sub[: a0 + 1] = False + keep[sub] = False + self.symmetry_operators = self.symmetry_operators[keep] + self.symmetry_reduction = self.symmetry_reduction[keep] + + if ( + self.orientation_fiber_angles is not None + and np.abs(self.orientation_fiber_angles[0] - 180.0) < 1e-3 + ): + zone_axis_range_flip = self.orientation_zone_axis_range.copy() + zone_axis_range_flip[0, :] = -1 * zone_axis_range_flip[0, :] + zone_axis_range_inv = np.linalg.inv(zone_axis_range_flip) + + num_sym = self.symmetry_operators.shape[0] + self.symmetry_operators = np.tile(self.symmetry_operators, [2, 1, 1]) + self.symmetry_reduction = np.tile(self.symmetry_reduction, [2, 1, 1]) + + for a0 in range(num_sym): + self.symmetry_reduction[a0 + num_sym] = ( + zone_axis_range_inv.T @ self.symmetry_operators[a0 + num_sym] + ).T + + # Calculate reference arrays for all orientations + k0 = np.array([0.0, 0.0, -1.0 / self.wavelength]) + n = np.array([0.0, 0.0, -1.0]) - if calculate_correlation_array: # initialize empty correlation array self.orientation_ref = np.zeros( ( From 91aacb49ff28dcfd842747caaf20ca603dc028c9 Mon Sep 17 00:00:00 2001 From: Shinjan Dutta Date: Tue, 26 Mar 2024 12:40:32 -0500 Subject: [PATCH 2/7] black formatting --- py4DSTEM/process/diffraction/crystal_ACOM.py | 49 +++++++++++--------- 1 file changed, 28 insertions(+), 21 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index 5bec155f2..f1cc65ce3 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -575,7 +575,7 @@ def orientation_plan( ) self.orientation_rotation_matrices[a0, :, :] = m1z @ m2x @ m3z self.orientation_rotation_angles[a0, :] = [azim[a0], elev[a0], -azim[a0]] - + # Rest of this function is only required for calculating orientation arrays if calculate_correlation_array: # If needed, create coarse orientation sieve @@ -586,7 +586,7 @@ def orientation_plan( ) if self.CUDA: self.orientation_sieve_CUDA = cp.asarray(self.orientation_sieve) - + # Convert to spherical coordinates elev = np.arctan2( np.hypot(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]), @@ -596,53 +596,60 @@ def orientation_plan( # self.orientation_vecs[:, 1], self.orientation_vecs[:, 0] # ) azim = np.arctan2(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]) - + # 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 ) - + # Calculate -z angles (Euler angle 3) self.orientation_gamma = np.linspace( 0, 2 * np.pi, self.orientation_in_plane_steps, endpoint=False ) - + # Determine the radii of all spherical shells radii_test = np.round(self.g_vec_leng / tol_distance) * tol_distance radii = np.unique(radii_test) # Remove zero beam keep = np.abs(radii) > tol_distance self.orientation_shell_radii = radii[keep] - + # init - self.orientation_shell_index = -1 * np.ones(self.g_vec_all.shape[1], dtype="int") + self.orientation_shell_index = -1 * np.ones( + self.g_vec_all.shape[1], dtype="int" + ) self.orientation_shell_count = np.zeros(self.orientation_shell_radii.size) - + # Assign each structure factor point to a radial shell for a0 in range(self.orientation_shell_radii.size): - sub = np.abs(self.orientation_shell_radii[a0] - radii_test) <= tol_distance / 2 - + sub = ( + np.abs(self.orientation_shell_radii[a0] - radii_test) + <= tol_distance / 2 + ) + self.orientation_shell_index[sub] = a0 self.orientation_shell_count[a0] = np.sum(sub) self.orientation_shell_radii[a0] = np.mean(self.g_vec_leng[sub]) - + # init storage arrays self.orientation_rotation_angles = np.zeros((self.orientation_num_zones, 3)) - self.orientation_rotation_matrices = np.zeros((self.orientation_num_zones, 3, 3)) - + self.orientation_rotation_matrices = np.zeros( + (self.orientation_num_zones, 3, 3) + ) + # If possible, Get symmetry operations for this spacegroup, store in matrix form if self.pymatgen_available: # get operators ops = self.pointgroup.get_point_group_operations() - + # Inverse of lattice zone_axis_range_inv = np.linalg.inv(self.orientation_zone_axis_range) - + # init num_sym = len(ops) self.symmetry_operators = np.zeros((num_sym, 3, 3)) self.symmetry_reduction = np.zeros((num_sym, 3, 3)) - + # calculate symmetry and reduction matrices for a0 in range(num_sym): self.symmetry_operators[a0] = ( @@ -651,7 +658,7 @@ def orientation_plan( self.symmetry_reduction[a0] = ( zone_axis_range_inv.T @ self.symmetry_operators[a0] ).T - + # Remove duplicates keep = np.ones(num_sym, dtype="bool") for a0 in range(num_sym): @@ -665,7 +672,7 @@ def orientation_plan( keep[sub] = False self.symmetry_operators = self.symmetry_operators[keep] self.symmetry_reduction = self.symmetry_reduction[keep] - + if ( self.orientation_fiber_angles is not None and np.abs(self.orientation_fiber_angles[0] - 180.0) < 1e-3 @@ -673,16 +680,16 @@ def orientation_plan( zone_axis_range_flip = self.orientation_zone_axis_range.copy() zone_axis_range_flip[0, :] = -1 * zone_axis_range_flip[0, :] zone_axis_range_inv = np.linalg.inv(zone_axis_range_flip) - + num_sym = self.symmetry_operators.shape[0] self.symmetry_operators = np.tile(self.symmetry_operators, [2, 1, 1]) self.symmetry_reduction = np.tile(self.symmetry_reduction, [2, 1, 1]) - + for a0 in range(num_sym): self.symmetry_reduction[a0 + num_sym] = ( zone_axis_range_inv.T @ self.symmetry_operators[a0 + num_sym] ).T - + # Calculate reference arrays for all orientations k0 = np.array([0.0, 0.0, -1.0 / self.wavelength]) n = np.array([0.0, 0.0, -1.0]) From 2a673f16ad2353c36444aba356294892a6779f9b Mon Sep 17 00:00:00 2001 From: Shinjan Dutta Date: Wed, 27 Mar 2024 16:20:24 -0500 Subject: [PATCH 3/7] Moved up the definition of elev and azim, before the function to calculate rotation matrices for zone axes, to avoid error --- py4DSTEM/process/diffraction/crystal.py | 155 +++++++++++++++--------- 1 file changed, 99 insertions(+), 56 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index 97d596ec4..6e99c38b9 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -1060,16 +1060,52 @@ def generate_projected_potential( # Vector projected along optic axis m_proj = np.squeeze(np.delete(lat_real, inds_tile, axis=0)) + # Thickness + if thickness_angstroms > 0: + num_proj = np.round(thickness_angstroms / np.abs(m_proj[2])).astype("int") + if num_proj > 1: + vec_proj = m_proj[:2] / pixel_size_angstroms + shifts = np.arange(num_proj).astype("float") + shifts -= np.mean(shifts) + x_proj = shifts * vec_proj[0] + y_proj = shifts * vec_proj[1] + else: + num_proj = 1 + else: + num_proj = 1 + # Determine tiling range - p_corners = np.array( - [ - [-im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, 0.0], - [im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, 0.0], - [-im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, 0.0], - [im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, 0.0], - ] - ) - ab = np.linalg.lstsq(m_tile[:, :2].T, p_corners[:, :2].T, rcond=None)[0] + if thickness_angstroms > 0: + # dx = m_proj[0] * num_proj * 0.5 + # dy = m_proj[1] * num_proj * 0.5 + # include the cell height + dz = m_proj[2] * num_proj * 0.5 + p_corners = np.array( + [ + [-im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, dz], + [im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, dz], + [-im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, dz], + [im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, dz], + [-im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, -dz], + [im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, -dz], + [-im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, -dz], + [im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, -dz], + ] + ) + else: + p_corners = np.array( + [ + [-im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, 0.0], + [im_size_Ang[0] * 0.5, -im_size_Ang[1] * 0.5, 0.0], + [-im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, 0.0], + [im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, 0.0], + ] + ) + + ab = np.linalg.lstsq( + m_tile[:, :2].T, + p_corners[:, :2].T, + rcond=None)[0] ab = np.floor(ab) a_range = np.array((np.min(ab[0]) - 1, np.max(ab[0]) + 2)) b_range = np.array((np.min(ab[1]) - 1, np.max(ab[1]) + 2)) @@ -1084,11 +1120,24 @@ def generate_projected_potential( abc_atoms[:, inds_tile[0]] += a_ind.ravel() abc_atoms[:, inds_tile[1]] += b_ind.ravel() xyz_atoms_ang = abc_atoms @ lat_real - atoms_ID_all = self.numbers[atoms_ind.ravel()] + atoms_ID_all_0 = self.numbers[atoms_ind.ravel()] # Center atoms on image plane - x = xyz_atoms_ang[:, 0] / pixel_size_angstroms + im_size[0] / 2.0 - y = xyz_atoms_ang[:, 1] / pixel_size_angstroms + im_size[1] / 2.0 + x0 = xyz_atoms_ang[:, 0] / pixel_size_angstroms + im_size[0] / 2.0 + y0 = xyz_atoms_ang[:, 1] / pixel_size_angstroms + im_size[1] / 2.0 + + # if needed, tile atoms in the projection direction + if num_proj > 1: + x = (x0[:,None] + x_proj[None,:]).ravel() + y = (y0[:,None] + y_proj[None,:]).ravel() + atoms_ID_all = np.tile(atoms_ID_all_0,(num_proj,1)) + else: + x = x0 + y = y0 + atoms_ID_all = atoms_ID_all_0 + # print(x.shape, y.shape) + + # delete atoms outside the field of view atoms_del = np.logical_or.reduce( ( x <= -potential_radius_angstroms / 2, @@ -1121,18 +1170,6 @@ def generate_projected_potential( atoms_lookup[a0, :, :] = atom_sf.projected_potential(atoms_ID[a0], R_2D) atoms_lookup **= power_scale - # Thickness - if thickness_angstroms > 0: - num_proj = np.round(thickness_angstroms / np.abs(m_proj[2])).astype("int") - if num_proj > 1: - vec_proj = m_proj[:2] / pixel_size_angstroms - shifts = np.arange(num_proj).astype("float") - shifts -= np.mean(shifts) - x_proj = shifts * vec_proj[0] - y_proj = shifts * vec_proj[1] - else: - num_proj = 1 - # initialize potential im_potential = np.zeros(im_size) @@ -1140,38 +1177,38 @@ def generate_projected_potential( for a0 in range(atoms_ID_all.shape[0]): ind = np.argmin(np.abs(atoms_ID - atoms_ID_all[a0])) - if num_proj > 1: - for a1 in range(num_proj): - x_ind = np.round(x[a0] + x_proj[a1]).astype("int") + R_ind - y_ind = np.round(y[a0] + y_proj[a1]).astype("int") + R_ind - x_sub = np.logical_and( - x_ind >= 0, - x_ind < im_size[0], - ) - y_sub = np.logical_and( - y_ind >= 0, - y_ind < im_size[1], - ) - - im_potential[ - x_ind[x_sub][:, None], y_ind[y_sub][None, :] - ] += atoms_lookup[ind][x_sub, :][:, y_sub] - - else: - x_ind = np.round(x[a0]).astype("int") + R_ind - y_ind = np.round(y[a0]).astype("int") + R_ind - x_sub = np.logical_and( - x_ind >= 0, - x_ind < im_size[0], - ) - y_sub = np.logical_and( - y_ind >= 0, - y_ind < im_size[1], - ) + # if num_proj > 1: + # for a1 in range(num_proj): + # x_ind = np.round(x[a0] + x_proj[a1]).astype("int") + R_ind + # y_ind = np.round(y[a0] + y_proj[a1]).astype("int") + R_ind + # x_sub = np.logical_and( + # x_ind >= 0, + # x_ind < im_size[0], + # ) + # y_sub = np.logical_and( + # y_ind >= 0, + # y_ind < im_size[1], + # ) + + # im_potential[ + # x_ind[x_sub][:, None], y_ind[y_sub][None, :] + # ] += atoms_lookup[ind][x_sub, :][:, y_sub] + + # else: + x_ind = np.round(x[a0]).astype("int") + R_ind + y_ind = np.round(y[a0]).astype("int") + R_ind + x_sub = np.logical_and( + x_ind >= 0, + x_ind < im_size[0], + ) + y_sub = np.logical_and( + y_ind >= 0, + y_ind < im_size[1], + ) - im_potential[ - x_ind[x_sub][:, None], y_ind[y_sub][None, :] - ] += atoms_lookup[ind][x_sub, :][:, y_sub] + im_potential[ + x_ind[x_sub][:, None], y_ind[y_sub][None, :] + ] += atoms_lookup[ind][x_sub, :][:, y_sub] if thickness_angstroms > 0: im_potential /= num_proj @@ -1198,10 +1235,16 @@ def generate_projected_potential( fig, ax = plt.subplots(figsize=figsize) ax.imshow( im_potential, - cmap="turbo", + cmap="gray", vmin=int_range[0], vmax=int_range[1], ) + # if num_proj > 1: + # for a1 in range(x_proj.size): + # ax.scatter(y+y_proj[a1],x+x_proj[a1], c='r') + + # ax.scatter(y+y_proj[0],x+x_proj[0], c='r') + # ax.scatter(y+y_proj[-1],x+x_proj[-1], c='g') ax.set_axis_off() ax.set_aspect("equal") From a6b99537bf72cec644520a55ff02740eb7d24580 Mon Sep 17 00:00:00 2001 From: Shinjan Dutta Date: Wed, 27 Mar 2024 17:28:58 -0500 Subject: [PATCH 4/7] Added warning when calculating orientation plan before structure factors --- py4DSTEM/process/diffraction/crystal.py | 25 +++++++++----------- py4DSTEM/process/diffraction/crystal_ACOM.py | 21 ++++++++++------ 2 files changed, 25 insertions(+), 21 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index 6e99c38b9..eae5e30e4 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -1073,11 +1073,11 @@ def generate_projected_potential( num_proj = 1 else: num_proj = 1 - + # Determine tiling range if thickness_angstroms > 0: - # dx = m_proj[0] * num_proj * 0.5 - # dy = m_proj[1] * num_proj * 0.5 + # dx = m_proj[0] * num_proj * 0.5 + # dy = m_proj[1] * num_proj * 0.5 # include the cell height dz = m_proj[2] * num_proj * 0.5 p_corners = np.array( @@ -1101,11 +1101,8 @@ def generate_projected_potential( [im_size_Ang[0] * 0.5, im_size_Ang[1] * 0.5, 0.0], ] ) - - ab = np.linalg.lstsq( - m_tile[:, :2].T, - p_corners[:, :2].T, - rcond=None)[0] + + ab = np.linalg.lstsq(m_tile[:, :2].T, p_corners[:, :2].T, rcond=None)[0] ab = np.floor(ab) a_range = np.array((np.min(ab[0]) - 1, np.max(ab[0]) + 2)) b_range = np.array((np.min(ab[1]) - 1, np.max(ab[1]) + 2)) @@ -1128,9 +1125,9 @@ def generate_projected_potential( # if needed, tile atoms in the projection direction if num_proj > 1: - x = (x0[:,None] + x_proj[None,:]).ravel() - y = (y0[:,None] + y_proj[None,:]).ravel() - atoms_ID_all = np.tile(atoms_ID_all_0,(num_proj,1)) + x = (x0[:, None] + x_proj[None, :]).ravel() + y = (y0[:, None] + y_proj[None, :]).ravel() + atoms_ID_all = np.tile(atoms_ID_all_0, (num_proj, 1)) else: x = x0 y = y0 @@ -1206,9 +1203,9 @@ def generate_projected_potential( y_ind < im_size[1], ) - im_potential[ - x_ind[x_sub][:, None], y_ind[y_sub][None, :] - ] += atoms_lookup[ind][x_sub, :][:, y_sub] + im_potential[x_ind[x_sub][:, None], y_ind[y_sub][None, :]] += atoms_lookup[ + ind + ][x_sub, :][:, y_sub] if thickness_angstroms > 0: im_potential /= num_proj diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index f1cc65ce3..a24f9bf0e 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -76,6 +76,13 @@ def orientation_plan( progress_bar (bool): If false no progress bar is displayed """ + # Check to make sure structure factors have been calculated if we want to calculate the correlation array. + if calculate_correlation_array is True: + if not hasattr(self, "g_vec_leng"): + raise ValueError( + "Run .calculate_structure_factors() before calculating orientation plan" + ) + # Store inputs self.accel_voltage = np.asarray(accel_voltage) self.orientation_kernel_size = np.asarray(corr_kernel_size) @@ -550,6 +557,13 @@ def orientation_plan( ).astype("int") self.orientation_inds[:, 2] = orientation_sector + # Convert to spherical coordinates + elev = np.arctan2( + np.hypot(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]), + self.orientation_vecs[:, 2], + ) + azim = np.arctan2(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]) + # Calculate rotation matrices for zone axes for a0 in np.arange(self.orientation_num_zones): m1z = np.array( @@ -587,16 +601,9 @@ def orientation_plan( if self.CUDA: self.orientation_sieve_CUDA = cp.asarray(self.orientation_sieve) - # Convert to spherical coordinates - elev = np.arctan2( - np.hypot(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]), - self.orientation_vecs[:, 2], - ) # azim = np.pi / 2 + np.arctan2( # self.orientation_vecs[:, 1], self.orientation_vecs[:, 0] # ) - azim = np.arctan2(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]) - # 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 From 490bb2387b5ed6c241011e506248b9b9238865a8 Mon Sep 17 00:00:00 2001 From: Colin Date: Wed, 27 Mar 2024 16:19:08 -0700 Subject: [PATCH 5/7] Fix to orientation plan calculation --- py4DSTEM/process/diffraction/crystal.py | 2 - py4DSTEM/process/diffraction/crystal_ACOM.py | 242 +++++++++---------- 2 files changed, 117 insertions(+), 127 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index eae5e30e4..0639df4fa 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -1076,8 +1076,6 @@ def generate_projected_potential( # Determine tiling range if thickness_angstroms > 0: - # dx = m_proj[0] * num_proj * 0.5 - # dy = m_proj[1] * num_proj * 0.5 # include the cell height dz = m_proj[2] * num_proj * 0.5 p_corners = np.array( diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index a24f9bf0e..a77f42bbd 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -76,12 +76,10 @@ def orientation_plan( progress_bar (bool): If false no progress bar is displayed """ - # Check to make sure structure factors have been calculated if we want to calculate the correlation array. - if calculate_correlation_array is True: - if not hasattr(self, "g_vec_leng"): - raise ValueError( - "Run .calculate_structure_factors() before calculating orientation plan" - ) + # Check to make sure user has calculated the structure factors if needed + if calculate_correlation_array: + if not hasattr(self,'g_vec_leng'): + raise ValueError('Run .calculate_structure_factors() before calculating an orientation plan.') # Store inputs self.accel_voltage = np.asarray(accel_voltage) @@ -557,63 +555,124 @@ def orientation_plan( ).astype("int") self.orientation_inds[:, 2] = orientation_sector - # Convert to spherical coordinates - elev = np.arctan2( - np.hypot(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]), - self.orientation_vecs[:, 2], + # If needed, create coarse orientation sieve + if self.orientation_refine: + self.orientation_sieve = np.logical_and( + np.mod(self.orientation_inds[:, 0], self.orientation_refine_ratio) == 0, + np.mod(self.orientation_inds[:, 1], self.orientation_refine_ratio) == 0, ) - azim = np.arctan2(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]) + if self.CUDA: + self.orientation_sieve_CUDA = cp.asarray(self.orientation_sieve) - # Calculate rotation matrices for zone axes - for a0 in np.arange(self.orientation_num_zones): - m1z = np.array( - [ - [np.cos(azim[a0]), np.sin(azim[a0]), 0], - [-np.sin(azim[a0]), np.cos(azim[a0]), 0], - [0, 0, 1], - ] - ) - m2x = np.array( - [ - [1, 0, 0], - [0, np.cos(elev[a0]), np.sin(elev[a0])], - [0, -np.sin(elev[a0]), np.cos(elev[a0])], - ] - ) - m3z = np.array( - [ - [np.cos(azim[a0]), -np.sin(azim[a0]), 0], - [np.sin(azim[a0]), np.cos(azim[a0]), 0], - [0, 0, 1], - ] - ) - self.orientation_rotation_matrices[a0, :, :] = m1z @ m2x @ m3z - self.orientation_rotation_angles[a0, :] = [azim[a0], elev[a0], -azim[a0]] + # Convert to spherical coordinates + elev = np.arctan2( + np.hypot(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]), + self.orientation_vecs[:, 2], + ) + # azim = np.pi / 2 + np.arctan2( + # self.orientation_vecs[:, 1], self.orientation_vecs[:, 0] + # ) + azim = np.arctan2(self.orientation_vecs[:, 0], self.orientation_vecs[:, 1]) + + # 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 + ) - # Rest of this function is only required for calculating orientation arrays - if calculate_correlation_array: - # If needed, create coarse orientation sieve - if self.orientation_refine: - self.orientation_sieve = np.logical_and( - np.mod(self.orientation_inds[:, 0], self.orientation_refine_ratio) == 0, - np.mod(self.orientation_inds[:, 1], self.orientation_refine_ratio) == 0, + # Calculate -z angles (Euler angle 3) + self.orientation_gamma = np.linspace( + 0, 2 * np.pi, self.orientation_in_plane_steps, endpoint=False + ) + + # init storage arrays + self.orientation_rotation_angles = np.zeros((self.orientation_num_zones, 3)) + self.orientation_rotation_matrices = np.zeros((self.orientation_num_zones, 3, 3)) + + # If possible, Get symmetry operations for this spacegroup, store in matrix form + if self.pymatgen_available: + # get operators + ops = self.pointgroup.get_point_group_operations() + + # Inverse of lattice + zone_axis_range_inv = np.linalg.inv(self.orientation_zone_axis_range) + + # init + num_sym = len(ops) + self.symmetry_operators = np.zeros((num_sym, 3, 3)) + self.symmetry_reduction = np.zeros((num_sym, 3, 3)) + + # calculate symmetry and reduction matrices + for a0 in range(num_sym): + self.symmetry_operators[a0] = ( + self.lat_inv.T @ ops[a0].rotation_matrix.T @ self.lat_real ) - if self.CUDA: - self.orientation_sieve_CUDA = cp.asarray(self.orientation_sieve) - - # azim = np.pi / 2 + np.arctan2( - # self.orientation_vecs[:, 1], self.orientation_vecs[:, 0] - # ) - # 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 - ) + self.symmetry_reduction[a0] = ( + zone_axis_range_inv.T @ self.symmetry_operators[a0] + ).T + + # Remove duplicates + keep = np.ones(num_sym, dtype="bool") + for a0 in range(num_sym): + if keep[a0]: + diff = np.sum( + np.abs(self.symmetry_operators - self.symmetry_operators[a0]), + axis=(1, 2), + ) + sub = diff < 1e-3 + sub[: a0 + 1] = False + keep[sub] = False + self.symmetry_operators = self.symmetry_operators[keep] + self.symmetry_reduction = self.symmetry_reduction[keep] + + if ( + self.orientation_fiber_angles is not None + and np.abs(self.orientation_fiber_angles[0] - 180.0) < 1e-3 + ): + zone_axis_range_flip = self.orientation_zone_axis_range.copy() + zone_axis_range_flip[0, :] = -1 * zone_axis_range_flip[0, :] + zone_axis_range_inv = np.linalg.inv(zone_axis_range_flip) - # Calculate -z angles (Euler angle 3) - self.orientation_gamma = np.linspace( - 0, 2 * np.pi, self.orientation_in_plane_steps, endpoint=False + num_sym = self.symmetry_operators.shape[0] + self.symmetry_operators = np.tile(self.symmetry_operators, [2, 1, 1]) + self.symmetry_reduction = np.tile(self.symmetry_reduction, [2, 1, 1]) + + for a0 in range(num_sym): + self.symmetry_reduction[a0 + num_sym] = ( + zone_axis_range_inv.T @ self.symmetry_operators[a0 + num_sym] + ).T + + # Calculate rotation matrices for zone axes + for a0 in np.arange(self.orientation_num_zones): + m1z = np.array( + [ + [np.cos(azim[a0]), np.sin(azim[a0]), 0], + [-np.sin(azim[a0]), np.cos(azim[a0]), 0], + [0, 0, 1], + ] + ) + m2x = np.array( + [ + [1, 0, 0], + [0, np.cos(elev[a0]), np.sin(elev[a0])], + [0, -np.sin(elev[a0]), np.cos(elev[a0])], + ] ) + m3z = np.array( + [ + [np.cos(azim[a0]), -np.sin(azim[a0]), 0], + [np.sin(azim[a0]), np.cos(azim[a0]), 0], + [0, 0, 1], + ] + ) + self.orientation_rotation_matrices[a0, :, :] = m1z @ m2x @ m3z + self.orientation_rotation_angles[a0, :] = [azim[a0], elev[a0], -azim[a0]] + + # Calculate reference arrays for all orientations + k0 = np.array([0.0, 0.0, -1.0 / self.wavelength]) + n = np.array([0.0, 0.0, -1.0]) + # Remaining calculations are only required if we are computing the correlation array. + if calculate_correlation_array: # Determine the radii of all spherical shells radii_test = np.round(self.g_vec_leng / tol_distance) * tol_distance radii = np.unique(radii_test) @@ -622,84 +681,17 @@ def orientation_plan( self.orientation_shell_radii = radii[keep] # init - self.orientation_shell_index = -1 * np.ones( - self.g_vec_all.shape[1], dtype="int" - ) + self.orientation_shell_index = -1 * np.ones(self.g_vec_all.shape[1], dtype="int") self.orientation_shell_count = np.zeros(self.orientation_shell_radii.size) # Assign each structure factor point to a radial shell for a0 in range(self.orientation_shell_radii.size): - sub = ( - np.abs(self.orientation_shell_radii[a0] - radii_test) - <= tol_distance / 2 - ) + sub = np.abs(self.orientation_shell_radii[a0] - radii_test) <= tol_distance / 2 self.orientation_shell_index[sub] = a0 self.orientation_shell_count[a0] = np.sum(sub) self.orientation_shell_radii[a0] = np.mean(self.g_vec_leng[sub]) - # init storage arrays - self.orientation_rotation_angles = np.zeros((self.orientation_num_zones, 3)) - self.orientation_rotation_matrices = np.zeros( - (self.orientation_num_zones, 3, 3) - ) - - # If possible, Get symmetry operations for this spacegroup, store in matrix form - if self.pymatgen_available: - # get operators - ops = self.pointgroup.get_point_group_operations() - - # Inverse of lattice - zone_axis_range_inv = np.linalg.inv(self.orientation_zone_axis_range) - - # init - num_sym = len(ops) - self.symmetry_operators = np.zeros((num_sym, 3, 3)) - self.symmetry_reduction = np.zeros((num_sym, 3, 3)) - - # calculate symmetry and reduction matrices - for a0 in range(num_sym): - self.symmetry_operators[a0] = ( - self.lat_inv.T @ ops[a0].rotation_matrix.T @ self.lat_real - ) - self.symmetry_reduction[a0] = ( - zone_axis_range_inv.T @ self.symmetry_operators[a0] - ).T - - # Remove duplicates - keep = np.ones(num_sym, dtype="bool") - for a0 in range(num_sym): - if keep[a0]: - diff = np.sum( - np.abs(self.symmetry_operators - self.symmetry_operators[a0]), - axis=(1, 2), - ) - sub = diff < 1e-3 - sub[: a0 + 1] = False - keep[sub] = False - self.symmetry_operators = self.symmetry_operators[keep] - self.symmetry_reduction = self.symmetry_reduction[keep] - - if ( - self.orientation_fiber_angles is not None - and np.abs(self.orientation_fiber_angles[0] - 180.0) < 1e-3 - ): - zone_axis_range_flip = self.orientation_zone_axis_range.copy() - zone_axis_range_flip[0, :] = -1 * zone_axis_range_flip[0, :] - zone_axis_range_inv = np.linalg.inv(zone_axis_range_flip) - - num_sym = self.symmetry_operators.shape[0] - self.symmetry_operators = np.tile(self.symmetry_operators, [2, 1, 1]) - self.symmetry_reduction = np.tile(self.symmetry_reduction, [2, 1, 1]) - - for a0 in range(num_sym): - self.symmetry_reduction[a0 + num_sym] = ( - zone_axis_range_inv.T @ self.symmetry_operators[a0 + num_sym] - ).T - - # Calculate reference arrays for all orientations - k0 = np.array([0.0, 0.0, -1.0 / self.wavelength]) - n = np.array([0.0, 0.0, -1.0]) # initialize empty correlation array self.orientation_ref = np.zeros( From 21a0cd5efc6bdac10feae04aa329ad836e790332 Mon Sep 17 00:00:00 2001 From: Colin Date: Wed, 27 Mar 2024 16:20:05 -0700 Subject: [PATCH 6/7] black formatting --- py4DSTEM/process/diffraction/crystal_ACOM.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal_ACOM.py b/py4DSTEM/process/diffraction/crystal_ACOM.py index a77f42bbd..fc6e691c3 100644 --- a/py4DSTEM/process/diffraction/crystal_ACOM.py +++ b/py4DSTEM/process/diffraction/crystal_ACOM.py @@ -78,8 +78,10 @@ def orientation_plan( # Check to make sure user has calculated the structure factors if needed if calculate_correlation_array: - if not hasattr(self,'g_vec_leng'): - raise ValueError('Run .calculate_structure_factors() before calculating an orientation plan.') + if not hasattr(self, "g_vec_leng"): + raise ValueError( + "Run .calculate_structure_factors() before calculating an orientation plan." + ) # Store inputs self.accel_voltage = np.asarray(accel_voltage) @@ -681,18 +683,22 @@ def orientation_plan( self.orientation_shell_radii = radii[keep] # init - self.orientation_shell_index = -1 * np.ones(self.g_vec_all.shape[1], dtype="int") + self.orientation_shell_index = -1 * np.ones( + self.g_vec_all.shape[1], dtype="int" + ) self.orientation_shell_count = np.zeros(self.orientation_shell_radii.size) # Assign each structure factor point to a radial shell for a0 in range(self.orientation_shell_radii.size): - sub = np.abs(self.orientation_shell_radii[a0] - radii_test) <= tol_distance / 2 + sub = ( + np.abs(self.orientation_shell_radii[a0] - radii_test) + <= tol_distance / 2 + ) self.orientation_shell_index[sub] = a0 self.orientation_shell_count[a0] = np.sum(sub) self.orientation_shell_radii[a0] = np.mean(self.g_vec_leng[sub]) - # initialize empty correlation array self.orientation_ref = np.zeros( ( From 940609096ea2013172de61957df65f295cbfa5dd Mon Sep 17 00:00:00 2001 From: Shinjan Dutta Date: Wed, 27 Mar 2024 19:02:22 -0500 Subject: [PATCH 7/7] bugs fixed --- py4DSTEM/process/diffraction/crystal.py | 56 +++++++------------------ 1 file changed, 16 insertions(+), 40 deletions(-) diff --git a/py4DSTEM/process/diffraction/crystal.py b/py4DSTEM/process/diffraction/crystal.py index eae5e30e4..601be0867 100644 --- a/py4DSTEM/process/diffraction/crystal.py +++ b/py4DSTEM/process/diffraction/crystal.py @@ -1135,12 +1135,13 @@ def generate_projected_potential( # print(x.shape, y.shape) # delete atoms outside the field of view + bound = potential_radius_angstroms / pixel_size_angstroms atoms_del = np.logical_or.reduce( ( - x <= -potential_radius_angstroms / 2, - y <= -potential_radius_angstroms / 2, - x >= im_size[0] + potential_radius_angstroms / 2, - y >= im_size[1] + potential_radius_angstroms / 2, + x <= -bound, + y <= -bound, + x >= im_size[0] + bound, + y >= im_size[1] + bound, ) ) x = np.delete(x, atoms_del) @@ -1165,6 +1166,14 @@ def generate_projected_potential( for a0 in range(atoms_ID.shape[0]): atom_sf = single_atom_scatter([atoms_ID[a0]]) atoms_lookup[a0, :, :] = atom_sf.projected_potential(atoms_ID[a0], R_2D) + + # if needed, apply gaussian blurring to each atom + if sigma_image_blur_angstroms > 0: + atoms_lookup[a0, :, :] = gaussian_filter( + atoms_lookup[a0, :, :], + sigma_image_blur_angstroms / pixel_size_angstroms, + mode="nearest", + ) atoms_lookup **= power_scale # initialize potential @@ -1174,24 +1183,6 @@ def generate_projected_potential( for a0 in range(atoms_ID_all.shape[0]): ind = np.argmin(np.abs(atoms_ID - atoms_ID_all[a0])) - # if num_proj > 1: - # for a1 in range(num_proj): - # x_ind = np.round(x[a0] + x_proj[a1]).astype("int") + R_ind - # y_ind = np.round(y[a0] + y_proj[a1]).astype("int") + R_ind - # x_sub = np.logical_and( - # x_ind >= 0, - # x_ind < im_size[0], - # ) - # y_sub = np.logical_and( - # y_ind >= 0, - # y_ind < im_size[1], - # ) - - # im_potential[ - # x_ind[x_sub][:, None], y_ind[y_sub][None, :] - # ] += atoms_lookup[ind][x_sub, :][:, y_sub] - - # else: x_ind = np.round(x[a0]).astype("int") + R_ind y_ind = np.round(y[a0]).astype("int") + R_ind x_sub = np.logical_and( @@ -1202,30 +1193,20 @@ def generate_projected_potential( y_ind >= 0, y_ind < im_size[1], ) - im_potential[x_ind[x_sub][:, None], y_ind[y_sub][None, :]] += atoms_lookup[ ind - ][x_sub, :][:, y_sub] + ][x_sub][:, y_sub] if thickness_angstroms > 0: im_potential /= num_proj - # if needed, apply gaussian blurring - if sigma_image_blur_angstroms > 0: - sigma_image_blur = sigma_image_blur_angstroms / pixel_size_angstroms - im_potential = gaussian_filter( - im_potential, - sigma_image_blur, - mode="nearest", - ) - if plot_result: # quick plotting of the result int_vals = np.sort(im_potential.ravel()) int_range = np.array( ( int_vals[np.round(0.02 * int_vals.size).astype("int")], - int_vals[np.round(0.98 * int_vals.size).astype("int")], + int_vals[np.round(0.999 * int_vals.size).astype("int")], ) ) @@ -1236,12 +1217,7 @@ def generate_projected_potential( vmin=int_range[0], vmax=int_range[1], ) - # if num_proj > 1: - # for a1 in range(x_proj.size): - # ax.scatter(y+y_proj[a1],x+x_proj[a1], c='r') - - # ax.scatter(y+y_proj[0],x+x_proj[0], c='r') - # ax.scatter(y+y_proj[-1],x+x_proj[-1], c='g') + # ax.scatter(y,x,c='r') # for testing ax.set_axis_off() ax.set_aspect("equal")