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.

Sign up for GitHub

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

Quick fix for correlation array, scalebar length #547

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
8 commits
Select commit Hold shift + click to select a range
3ec7b0b
Fixing scale bar being plotted as slightly too long
cophus Oct 27, 2023
77762d4
Add option to skip calculating correlation array in crystal.orientati…
cophus Oct 27, 2023
0c84848
Black formatting
cophus Oct 27, 2023
958642e
adding assert statement
alex-rakowski Oct 27, 2023
3290342
black
alex-rakowski Oct 27, 2023
a23839e
changing assert to warning
alex-rakowski Oct 30, 2023
7ca7edb
black
alex-rakowski Oct 30, 2023
7bd2ceb
change Warning to ValueError
sezelt Oct 30, 2023
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
137 changes: 71 additions & 66 deletions py4DSTEM/process/diffraction/crystal_ACOM.py
View file
Open in desktop
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 Expand Up @@ -903,7 +902,13 @@ def match_single_pattern(
Figure handles for the plotting output
"""

# init orientation output
# adding assert statement for checking self.orientation_ref is present
# adding assert statement for checking self.orientation_ref is present
if not hasattr(self, "orientation_ref"):
raise ValueError(
"orientation_plan must be run with 'calculate_correlation_array=True'"
)

orientation = Orientation(num_matches=num_matches_return)
if bragg_peaks.data.shape[0] < min_number_peaks:
return orientation
Expand Down
9 changes: 8 additions & 1 deletion py4DSTEM/visualize/overlay.py
View file
Open in desktop
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