Skip to content

Commit

Permalink
Black formatting
Browse files Browse the repository at this point in the history
  • Loading branch information
@cophus
cophus committed Dec 12, 2023
1 parent e533d07 commit 94e0aca
Showing 1 changed file with 105 additions and 103 deletions.
208 changes: 105 additions & 103 deletions py4DSTEM/process/diffraction/flowlines.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -1005,7 +1005,7 @@ def make_flowline_combined_image(
def orientation_correlation(
orient_hist,
radius_max=None,
progress_bar = True,
progress_bar=True,
):
"""
Take in the 4D orientation histogram, and compute the distance-angle (auto)correlations
Expand Down Expand Up @@ -1096,8 +1096,8 @@ def orientation_correlation(
for a0, a1 in tqdmnd(
range(size_input[0]),
range(size_input[0]),
desc='Calculate correlation plots',
unit=" probe positions",
desc="Calculate correlation plots",
unit=" probe positions",
disable=not progress_bar,
):
# for a0 in range(size_input[0]):
Expand Down Expand Up @@ -1143,49 +1143,49 @@ def orientation_correlation(

def plot_orientation_correlation(
orient_corr,
prob_range = [0.1, 10.0],
calculate_coefs = False,
fraction_coefs = 0.5,
length_fit_slope = 10,
plot_overlaid_coefs = True,
inds_plot = None,
pixel_size = None,
pixel_units = None,
fontsize = 10,
figsize = (8, 6),
returnfig = False,
prob_range=[0.1, 10.0],
calculate_coefs=False,
fraction_coefs=0.5,
length_fit_slope=10,
plot_overlaid_coefs=True,
inds_plot=None,
pixel_size=None,
pixel_units=None,
fontsize=10,
figsize=(8, 6),
returnfig=False,
):
"""
Plot the distance-angle (auto)correlations in orient_corr.
Parameters
----------
orient_corr (array):
orient_corr (array):
3D or 4D array containing correlation images as function of (dr,dtheta)
1st index represents each pair of rings.
prob_range (array):
prob_range (array):
Plotting range in units of "multiples of random distribution".
calculate_coefs (bool):
If this value is True, the 0.5 and 0.1 distribution fraction of the
If this value is True, the 0.5 and 0.1 distribution fraction of the
radial and annular correlations will be calculated and printed.
fraction_coefs (float):
What fraction to calculate the correlation distribution coefficients for.
length_fit_slope (int):
Number of pixels to fit the slope of angular vs radial intercept.
plot_overlaid_coefs (bool):
If this value is True, the 0.5 and 0.1 distribution fraction of the
If this value is True, the 0.5 and 0.1 distribution fraction of the
radial and annular correlations will be overlaid onto the plots.
inds_plot (float):
inds_plot (float):
Which indices to plot for orient_corr. Set to "None" to plot all pairs.
pixel_size (float):
pixel_size (float):
Pixel size for x axis.
pixel_units (str):
pixel_units (str):
units of pixels.
fontsize (float):
fontsize (float):
Font size. Title will be slightly larger, axis slightly smaller.
figsize (array):
figsize (array):
Size of the figure panels.
returnfig (bool):
returnfig (bool):
Set to True to return figure axes.
Returns
Expand Down Expand Up @@ -1250,7 +1250,7 @@ def fit_dist(x, *coefs):
int_bg = coefs[1]
sigma = coefs[2]
p = coefs[3]
return (int0-int_bg)*np.exp(np.abs(x)**p/(-1*sigma**p)) + int_bg
return (int0 - int_bg) * np.exp(np.abs(x) ** p / (-1 * sigma**p)) + int_bg

# plotting
num_plot = inds_plot.shape[0]
Expand Down Expand Up @@ -1284,100 +1284,98 @@ def fit_dist(x, *coefs):
t_lab.append(f"{10**t[a1]:.2g}")

cbar.set_ticks(t)
cbar.ax.set_yticklabels(t_lab, fontsize = fontsize*0.8)
cbar.ax.set_yticklabels(t_lab, fontsize=fontsize * 0.8)
cbar.ax.set_ylabel("Probability (m.r.d.)", fontsize=fontsize)

ind_0 = pair_inds[0, ind]
ind_1 = pair_inds[1, ind]

if ind_0 != ind_1:
ax_handle.set_title(
"Correlation of Rings " + str(ind_0) + " and " + str(ind_1), fontsize=fontsize*1.2
"Correlation of Rings " + str(ind_0) + " and " + str(ind_1),
fontsize=fontsize * 1.2,
)
else:
ax_handle.set_title("Autocorrelation of Ring " + str(ind_0), fontsize=fontsize*1.2)
ax_handle.set_title(
"Autocorrelation of Ring " + str(ind_0), fontsize=fontsize * 1.2
)

# x axis labels
# if pixel_size is not None:
x_t = ax_handle.get_xticks()
sub = np.logical_or(x_t < 0, x_t > orient_corr.shape[2])
x_t_new = np.delete(x_t, sub)
ax_handle.set_xticks(x_t_new)
ax_handle.set_xticklabels(x_t_new * pixel_size, fontsize = fontsize*0.8)
ax_handle.set_xticklabels(x_t_new * pixel_size, fontsize=fontsize * 0.8)
ax_handle.set_xlabel("Radial Distance (" + pixel_units + ")", fontsize=fontsize)

# y axis labels
ax_handle.invert_yaxis()
ax_handle.set_ylabel("Relative Orientation (°)", fontsize=fontsize)
y_ticks = np.linspace(0.0,orient_corr.shape[1]-1,10, endpoint = True)
y_ticks = np.linspace(0.0, orient_corr.shape[1] - 1, 10, endpoint=True)
ax_handle.set_yticks(y_ticks)
ax_handle.set_yticklabels(["0", "", "", "30", "", "", "60", "", "", "90"],
fontsize = fontsize*0.8)
ax_handle.set_yticklabels(
["0", "", "", "30", "", "", "60", "", "", "90"], fontsize=fontsize * 0.8
)

if calculate_coefs:
# Radial fractions
y = np.arange(orient_corr.shape[2])
if orient_corr[ind,0,0] > orient_corr[ind,-1,0]:
z = orient_corr[ind,0,:]
if orient_corr[ind, 0, 0] > orient_corr[ind, -1, 0]:
z = orient_corr[ind, 0, :]
else:
z = orient_corr[ind,-1,:]
coefs = [np.max(z),np.min(z),y[-1]*0.25,2]
bounds = (
(1e-3,0,1e-3,1.0),
(np.inf,np.inf,np.inf,np.inf))
coefs = curve_fit(
fit_dist,
y,
z,
p0=coefs,
bounds=bounds)[0]
coef_radial = coefs[2] * (np.log(1/fraction_coefs)**(1/coefs[3]))
z = orient_corr[ind, -1, :]
coefs = [np.max(z), np.min(z), y[-1] * 0.25, 2]
bounds = ((1e-3, 0, 1e-3, 1.0), (np.inf, np.inf, np.inf, np.inf))
coefs = curve_fit(fit_dist, y, z, p0=coefs, bounds=bounds)[0]
coef_radial = coefs[2] * (np.log(1 / fraction_coefs) ** (1 / coefs[3]))

# Annular fractions
x = np.arange(orient_corr.shape[1])
if orient_corr[ind,0,0] > orient_corr[ind,-1,0]:
z = orient_corr[ind,:,0]
if orient_corr[ind, 0, 0] > orient_corr[ind, -1, 0]:
z = orient_corr[ind, :, 0]
else:
z = np.flip(orient_corr[ind,:,0], axis=0)
z = np.flip(orient_corr[ind, :, 0], axis=0)
z = np.maximum(z, 1.0)
coefs = [np.max(z),np.min(z),x[-1]*0.25,2]
bounds = (
(1e-3,0,1e-3,1.0),
(np.inf,np.inf,np.inf,np.inf))
coefs = curve_fit(
fit_dist,
x,
z,
p0=coefs,
bounds=bounds)[0]
coef_annular = coefs[2] * (np.log(1/fraction_coefs)**(1/coefs[3]))
if orient_corr[ind,0,0] <= orient_corr[ind,-1,0]:
coefs = [np.max(z), np.min(z), x[-1] * 0.25, 2]
bounds = ((1e-3, 0, 1e-3, 1.0), (np.inf, np.inf, np.inf, np.inf))
coefs = curve_fit(fit_dist, x, z, p0=coefs, bounds=bounds)[0]
coef_annular = coefs[2] * (np.log(1 / fraction_coefs) ** (1 / coefs[3]))
if orient_corr[ind, 0, 0] <= orient_corr[ind, -1, 0]:
coef_annular = orient_corr.shape[1] - 1 - coef_annular
pixel_size_annular = 90 / (orient_corr.shape[1]-1)
pixel_size_annular = 90 / (orient_corr.shape[1] - 1)

# Slope of annular vs radial correlations as radius --> 0
x_slope = np.argmin(np.abs(orient_corr[ind,:,:length_fit_slope]-1.0),axis=0)
x_slope = np.argmin(
np.abs(orient_corr[ind, :, :length_fit_slope] - 1.0), axis=0
)
y_slope = np.arange(length_fit_slope)
coefs_slope = np.polyfit(y_slope, x_slope, 1)

# Print results
if ind_0 != ind_1:
print(
"Correlation of Rings " + str(ind_0) + " and " + str(ind_1)
)
print("Correlation of Rings " + str(ind_0) + " and " + str(ind_1))
else:
print("Autocorrelation of Ring " + str(ind_0))
print(str(np.round(fraction_coefs*100).astype('int')) \
+ '% probability radial distance = ' \
+ str(np.round(coef_radial * pixel_size,2)) \
+ ' ' + pixel_units)
print(str(np.round(fraction_coefs*100).astype('int')) \
+ '% probability annular distance = ' \
+ str(np.round(coef_annular * pixel_size_annular,2)) \
+ ' degrees')
print('slope = ' \
+ str(np.round(coefs_slope[0]*pixel_size_annular/pixel_size,2)) \
+ ' degrees/' + pixel_units)
print(
str(np.round(fraction_coefs * 100).astype("int"))
+ "% probability radial distance = "
+ str(np.round(coef_radial * pixel_size, 2))
+ " "
+ pixel_units
)
print(
str(np.round(fraction_coefs * 100).astype("int"))
+ "% probability annular distance = "
+ str(np.round(coef_annular * pixel_size_annular, 2))
+ " degrees"
)
print(
"slope = "
+ str(np.round(coefs_slope[0] * pixel_size_annular / pixel_size, 2))
+ " degrees/"
+ pixel_units
)
print()

if plot_overlaid_coefs:
Expand All @@ -1386,43 +1384,47 @@ def fit_dist(x, *coefs):
else:
ax_handle = ax

if orient_corr[ind,0,0] > orient_corr[ind,-1,0]:
if orient_corr[ind, 0, 0] > orient_corr[ind, -1, 0]:
ax_handle.plot(
np.array((coef_radial,coef_radial,0.0)),
np.array((0.0,coef_annular,coef_annular)),
color = (1.0,1.0,1.0),
np.array((coef_radial, coef_radial, 0.0)),
np.array((0.0, coef_annular, coef_annular)),
color=(1.0, 1.0, 1.0),
)
ax_handle.plot(
np.array((coef_radial,coef_radial,0.0)),
np.array((0.0,coef_annular,coef_annular)),
color = (0.0,0.0,0.0),
linestyle = '--',
np.array((coef_radial, coef_radial, 0.0)),
np.array((0.0, coef_annular, coef_annular)),
color=(0.0, 0.0, 0.0),
linestyle="--",
)
else:
ax_handle.plot(
np.array((coef_radial,coef_radial,0.0)),
np.array((
orient_corr.shape[1] - 1,
coef_annular,
coef_annular,
)),
color = (1.0,1.0,1.0),
np.array((coef_radial, coef_radial, 0.0)),
np.array(
(
orient_corr.shape[1] - 1,
coef_annular,
coef_annular,
)
),
color=(1.0, 1.0, 1.0),
)
ax_handle.plot(
np.array((coef_radial,coef_radial,0.0)),
np.array((
orient_corr.shape[1] - 1,
coef_annular,
coef_annular,
)),
color = (0.0,0.0,0.0),
linestyle = '--',
np.array((coef_radial, coef_radial, 0.0)),
np.array(
(
orient_corr.shape[1] - 1,
coef_annular,
coef_annular,
)
),
color=(0.0, 0.0, 0.0),
linestyle="--",
)
ax_handle.plot(
y_slope,
y_slope.astype('float') * coefs_slope[0] + coefs_slope[1],
color = (0.0,0.0,0.0),
linestyle = '--',
y_slope.astype("float") * coefs_slope[0] + coefs_slope[1],
color=(0.0, 0.0, 0.0),
linestyle="--",
)

# Fix spacing
Expand Down

0 comments on commit 94e0aca

Please sign in to comment.