Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Quick fix for correlation array, scalebar length #547

Merged
merged 8 commits into from
Oct 31, 2023
Merged
Show file tree
Hide file tree
Changes from 3 commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
129 changes: 64 additions & 65 deletions py4DSTEM/process/diffraction/crystal_ACOM.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ def orientation_plan(
corr_kernel_size: float = 0.08,
radial_power: float = 1.0,
intensity_power: float = 0.25, # New default intensity power scaling
calculate_correlation_array=True,
tol_peak_delete=None,
tol_distance: float = 0.01,
fiber_axis=None,
Expand Down Expand Up @@ -61,6 +62,8 @@ def orientation_plan(
corr_kernel_size (float): Correlation kernel size length in Angstroms
radial_power (float): Power for scaling the correlation intensity as a function of the peak radius
intensity_power (float): Power for scaling the correlation intensity as a function of the peak intensity
calculate_correlation_array (bool): Set to false to skip calculating the correlation array.
This is useful when we only want the angular range / rotation matrices.
tol_peak_delete (float): Distance to delete peaks for multiple matches.
Default is kernel_size * 0.5
tol_distance (float): Distance tolerance for radial shell assignment [1/Angstroms]
Expand Down Expand Up @@ -598,21 +601,6 @@ def orientation_plan(
# 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_ref = np.zeros(
(
self.orientation_num_zones,
np.size(self.orientation_shell_radii),
self.orientation_in_plane_steps,
),
dtype="float",
)
# self.orientation_ref_1D = np.zeros(
# (
# self.orientation_num_zones,
# np.size(self.orientation_shell_radii),
# ),
# dtype="float",
# )

# If possible, Get symmetry operations for this spacegroup, store in matrix form
if self.pymatgen_available:
Expand Down Expand Up @@ -697,65 +685,76 @@ def orientation_plan(
k0 = np.array([0.0, 0.0, -1.0 / self.wavelength])
n = np.array([0.0, 0.0, -1.0])

for a0 in tqdmnd(
np.arange(self.orientation_num_zones),
desc="Orientation plan",
unit=" zone axes",
disable=not progress_bar,
):
# reciprocal lattice spots and excitation errors
g = self.orientation_rotation_matrices[a0, :, :].T @ self.g_vec_all
sg = self.excitation_errors(g)

# Keep only points that will contribute to this orientation plan slice
keep = np.abs(sg) < self.orientation_kernel_size

# in-plane rotation angle
phi = np.arctan2(g[1, :], g[0, :])

# Loop over all peaks
for a1 in np.arange(self.g_vec_all.shape[1]):
ind_radial = self.orientation_shell_index[a1]
if calculate_correlation_array:
# initialize empty correlation array
self.orientation_ref = np.zeros(
(
self.orientation_num_zones,
np.size(self.orientation_shell_radii),
self.orientation_in_plane_steps,
),
dtype="float",
)

if keep[a1] and ind_radial >= 0:
# 2D orientation plan
self.orientation_ref[a0, ind_radial, :] += (
np.power(self.orientation_shell_radii[ind_radial], radial_power)
* np.power(self.struct_factors_int[a1], intensity_power)
* np.maximum(
1
- np.sqrt(
sg[a1] ** 2
+ (
(
np.mod(
self.orientation_gamma - phi[a1] + np.pi,
2 * np.pi,
for a0 in tqdmnd(
np.arange(self.orientation_num_zones),
desc="Orientation plan",
unit=" zone axes",
disable=not progress_bar,
):
# reciprocal lattice spots and excitation errors
g = self.orientation_rotation_matrices[a0, :, :].T @ self.g_vec_all
sg = self.excitation_errors(g)

# Keep only points that will contribute to this orientation plan slice
keep = np.abs(sg) < self.orientation_kernel_size

# in-plane rotation angle
phi = np.arctan2(g[1, :], g[0, :])

# Loop over all peaks
for a1 in np.arange(self.g_vec_all.shape[1]):
ind_radial = self.orientation_shell_index[a1]

if keep[a1] and ind_radial >= 0:
# 2D orientation plan
self.orientation_ref[a0, ind_radial, :] += (
np.power(self.orientation_shell_radii[ind_radial], radial_power)
* np.power(self.struct_factors_int[a1], intensity_power)
* np.maximum(
1
- np.sqrt(
sg[a1] ** 2
+ (
(
np.mod(
self.orientation_gamma - phi[a1] + np.pi,
2 * np.pi,
)
- np.pi
)
- np.pi
* self.orientation_shell_radii[ind_radial]
)
* self.orientation_shell_radii[ind_radial]
** 2
)
** 2
/ self.orientation_kernel_size,
0,
)
/ self.orientation_kernel_size,
0,
)
)

orientation_ref_norm = np.sqrt(np.sum(self.orientation_ref[a0, :, :] ** 2))
if orientation_ref_norm > 0:
self.orientation_ref[a0, :, :] /= orientation_ref_norm
orientation_ref_norm = np.sqrt(np.sum(self.orientation_ref[a0, :, :] ** 2))
if orientation_ref_norm > 0:
self.orientation_ref[a0, :, :] /= orientation_ref_norm

# Maximum value
self.orientation_ref_max = np.max(np.real(self.orientation_ref))
# Maximum value
self.orientation_ref_max = np.max(np.real(self.orientation_ref))

# Fourier domain along angular axis
if self.CUDA:
self.orientation_ref = cp.asarray(self.orientation_ref)
self.orientation_ref = cp.conj(cp.fft.fft(self.orientation_ref))
else:
self.orientation_ref = np.conj(np.fft.fft(self.orientation_ref))
# Fourier domain along angular axis
if self.CUDA:
self.orientation_ref = cp.asarray(self.orientation_ref)
self.orientation_ref = cp.conj(cp.fft.fft(self.orientation_ref))
else:
self.orientation_ref = np.conj(np.fft.fft(self.orientation_ref))


def match_orientations(
Expand Down
9 changes: 8 additions & 1 deletion py4DSTEM/visualize/overlay.py
Original file line number Diff line number Diff line change
Expand Up @@ -832,7 +832,14 @@ def add_scalebar(ax, d):
labelpos_y = y0

# Add line
ax.plot((yi, yf), (xi, xf), lw=width, color=color, alpha=alpha)
ax.plot(
(yi, yf),
(xi, xf),
color=color,
alpha=alpha,
lw=width,
solid_capstyle="butt",
)

# Add label
if label:
Expand Down