vmax=1,
cmap=my_cmap,
title='Input amplitude')
#################################
# pad data to the original size #
#################################
print("Initial data size:", amp.shape)
if len(original_size) == 0:
original_size = amp.shape
print("FFT size before accounting for binning", original_size)
original_size = tuple(
[original_size[index] // binning[index] for index in range(len(binning))])
print("Binning used during phasing:", binning)
print("Padding back to original FFT size", original_size, '\n')
amp = pu.crop_pad(array=amp, output_shape=original_size)
nz, ny, nx = amp.shape
##########################################################
# calculate q for later regridding in the detector frame #
##########################################################
kin = 2 * np.pi / setup.wavelength * beam_direction # in laboratory frame z downstream, y vertical, x outboard
kout = setup.exit_wavevector(
) # in laboratory frame z downstream, y vertical, x outboard
q = kout - kin
Qnorm = np.linalg.norm(q)
q = q / Qnorm
Qnorm = Qnorm * 1e-10 # switch to angstroms
planar_dist = 2 * np.pi / Qnorm # Qnorm should be in angstroms
print("Wavevector transfer [z, y, x]:", q * Qnorm)
print("Wavevector transfer: (angstroms)", str('{:.4f}'.format(Qnorm)))
##################################################
# pad arrays to obtain the desired field of view #
##################################################
z_pixel_FOV = int(np.rint((field_of_view[0] / voxel_size) /
2)) # half-number of pixels corresponding to the FOV
y_pixel_FOV = int(np.rint((field_of_view[1] / voxel_size) /
2)) # half-number of pixels corresponding to the FOV
x_pixel_FOV = int(np.rint((field_of_view[2] / voxel_size) /
2)) # half-number of pixels corresponding to the FOV
new_shape = [
max(numz, 2 * z_pixel_FOV),
max(numy, 2 * y_pixel_FOV),
max(numx, 2 * x_pixel_FOV)
]
amp = pu.crop_pad(array=amp, output_shape=new_shape, debugging=False)
phase = pu.crop_pad(array=phase, output_shape=new_shape, debugging=False)
strain = pu.crop_pad(array=strain, output_shape=new_shape, debugging=False)
numz, numy, numx = amp.shape
print("Cropped/padded data size: (", numz, ',', numy, ',', numx, ')')
##########################################################
# set the strain and phase to NAN outside of the support #
##########################################################
support = np.zeros((numz, numy, numx))
support[np.nonzero(amp)] = 1
strain[support == 0] = np.nan
phase[support == 0] = np.nan
############################################
#############################
file_path = filedialog.askopenfilename(initialdir=detector.savedir, title="Select reconstructions (prtf)",
filetypes=[("NPZ", "*.npz"), ("NPY", "*.npy"),
("CXI", "*.cxi"), ("HDF5", "*.h5")])
obj, extension = util.load_file(file_path)
print('Opening ', file_path)
if extension == '.h5':
comment = comment + '_mode'
# check if the shape of the real space object is the same as the measured diffraction pattern
# the real space object may have been further cropped to a tight support, to save memory space.
if obj.shape != diff_pattern.shape:
print('Reconstructed object shape = ', obj.shape, 'different from the experimental diffraction pattern: crop/pad')
obj = pu.crop_pad(array=obj, output_shape=diff_pattern.shape, debugging=False)
# calculate the retrieved diffraction amplitude
phased_fft = fftshift(fftn(obj)) / (np.sqrt(numz)*np.sqrt(numy)*np.sqrt(numx)) # complex amplitude
del obj
gc.collect()
if debug:
plt.figure()
plt.imshow(np.log10(abs(phased_fft).sum(axis=0)), cmap=my_cmap, vmin=0, vmax=3.5)
plt.colorbar()
plt.title('abs(retrieved amplitude).sum(axis=0) before alignment')
plt.pause(0.1)
if align_pattern:
# align the reconstruction with the initial diffraction data
obj, extension = util.load_file(file_path[0])
# obj[0:40,:,:] = 0
# obj[65:,:,:] = 0
if extension == '.h5':
comment = comment + '_mode'
nz, ny, nx = obj.shape
print("Initial data size: (", nz, ',', ny, ',', nx, ')')
if len(original_size) == 0:
original_size = obj.shape
print("FFT size before accounting for binning", original_size)
original_size = tuple([original_size[index] // binning[index] for index in range(len(binning))])
print("Binning used during phasing:", binning)
print("Padding back to original FFT size", original_size)
obj = pu.crop_pad(array=obj, output_shape=original_size)
nz, ny, nx = obj.shape
###########################################################################
# define range for orthogonalization and plotting - speed up calculations #
###########################################################################
zrange, yrange, xrange =\
pu.find_datarange(array=obj, plot_margin=plot_margin, amplitude_threshold=0.1, keep_size=keep_size)
numz = zrange * 2
numy = yrange * 2
numx = xrange * 2
print("Data shape used for orthogonalization and plotting: (", numz, ',', numy, ',', numx, ')')
####################################################################################
# find the best reconstruction from the list, based on mean amplitude and variance #
file_path = filedialog.askopenfilename(initialdir=datadir,
title="Select 3D data",
filetypes=[("NPZ", "*.npz")])
data, _ = util.load_file(file_path)
nbz, nby, nbx = data.shape
print('data shape:', data.shape)
if crop_shape:
assert len(
crop_shape) == 3, 'crop should be a sequence of 3 voxels numbers'
assert np.all(np.asarray(origin) - np.asarray(crop_shape) // 2 >= 0
), 'origin incompatible with crop_shape'
assert origin[0] + crop_shape[0] // 2 <= nbz and origin[1] + crop_shape[1] // 2 <= nby \
and origin[2] + crop_shape[2] // 2 <= nbx, 'origin incompatible with crop_shape'
data = pu.crop_pad(array=data, output_shape=crop_shape, crop_center=origin)
gc.collect()
nbz, nby, nbx = data.shape
print('data shape after cropping:', data.shape)
# calculate the new position of origin
new_origin = (crop_shape[0] // 2, crop_shape[1] // 2, crop_shape[2] // 2)
else:
new_origin = origin
if plots:
gu.multislices_plot(data,
sum_frames=True,
scale='log',
plot_colorbar=True,
title='S' + str(scan) + '\n Data before rotation',
vmin=0,
print('Rocking curve ', idx + 1, ': Pearson correlation coefficient = ',
str('{:.2f}'.format(correlation)))
corr_coeff.append(str('{:.2f}'.format(correlation)))
if correlation >= correlation_threshold:
scanlist.append(scan_list[idx])
sumdata = sumdata + data
summask = summask + mask
else:
print('Scan ', scan_list[idx],
', correlation below threshold, skip concatenation')
summask[np.nonzero(summask)] = 1 # mask should be 0 or 1
sumdata = sumdata / len(scanlist)
summask = pu.crop_pad(array=summask, output_shape=output_shape)
sumdata = pu.crop_pad(array=sumdata, output_shape=output_shape)
savedir = homedir + sample_name + 'sum_S' + str(scanlist[0]) + '_to_S' + str(
scanlist[-1]) + '/'
template = '_S'+str(scanlist[0]) + '_to_S' + str(scanlist[-1]) + '_' +\
str(output_shape[0]) + '_' + str(output_shape[1]) + '_' + str(output_shape[2]) + '.npz'
pathlib.Path(savedir).mkdir(parents=True, exist_ok=True)
np.savez_compressed(savedir + 'pynx' + template, obj=sumdata)
np.savez_compressed(savedir + 'maskpynx' + template, obj=summask)
print('Sum of ', len(scanlist), 'scans')
fig, _, _ = gu.multislices_plot(sumdata,
sum_frames=True,
scale='log',
pixel_y=pixel_y)
newvoxelsize_z, newvoxelsize_y, newvoxelsize_x = setup.voxel_sizes(
unbinned_shape,
tilt_angle=tilt_angle,
pixel_x=pixel_x,
pixel_y=pixel_y)
print('Original voxel sizes zyx (nm):', str('{:.2f}'.format(voxelsize_z)),
str('{:.2f}'.format(voxelsize_y)), str('{:.2f}'.format(voxelsize_x)))
print('Output voxel sizes zyx (nm):', str('{:.2f}'.format(newvoxelsize_z)),
str('{:.2f}'.format(newvoxelsize_y)),
str('{:.2f}'.format(newvoxelsize_x)))
# Interpolate the support
print('\nInterpolating the support...')
data = pu.crop_pad(
data, pynx_shape) # the data could be cropped near the support
fig, _, _ = gu.multislices_plot(data,
sum_frames=True,
scale='linear',
plot_colorbar=True,
vmin=0,
title='Support before interpolation\n',
is_orthogonal=True,
reciprocal_space=False)
rgi = RegularGridInterpolator(
(np.arange(-pynx_shape[0] // 2, pynx_shape[0] // 2, 1) * voxelsize_z,
np.arange(-pynx_shape[1] // 2, pynx_shape[1] // 2, 1) * voxelsize_y,
np.arange(-pynx_shape[2] // 2, pynx_shape[2] // 2, 1) * voxelsize_x),
data,
method='linear',
numz, numy, numx = amp.shape
print("Initial data size: (", numz, ',', numy, ',', numx, ')')
#############################
# pad arrays to obtain the desired field of view
#############################
pixel_spacing = tick_spacing / voxel_size
pixel_FOV = int(np.rint((field_of_view / voxel_size) /
2)) # half-number of pixels corresponding to the FOV
new_shape = [
max(numz, 2 * pixel_FOV),
max(numy, 2 * pixel_FOV),
max(numx, 2 * pixel_FOV)
]
support = pu.crop_pad(array=support, output_shape=new_shape, debugging=False)
strain = pu.crop_pad(array=strain, output_shape=new_shape, debugging=False)
phase = pu.crop_pad(array=phase, output_shape=new_shape, debugging=False)
amp = pu.crop_pad(array=amp, output_shape=new_shape, debugging=False)
bulk = pu.crop_pad(array=bulk, output_shape=new_shape, debugging=False)
numz, numy, numx = amp.shape
print("Cropped/padded data size: (", numz, ',', numy, ',', numx, ')')
#############################
# center arrays based on the support
#############################
zcom, ycom, xcom = center_of_mass(support)
zcom, ycom, xcom = [int(np.rint(zcom)), int(np.rint(ycom)), int(np.rint(xcom))]
support = np.roll(support,
(numz // 2 - zcom, numy // 2 - ycom, numx // 2 - xcom),
axis=(0, 1, 2))
amp = amp * np.exp(1j * phase) # amp is the complex amplitude
del phase
gc.collect()
else:
comment = comment + "_support"
gu.multislices_plot(abs(amp),
sum_frames=False,
reciprocal_space=False,
is_orthogonal=True,
title='abs(amp)')
####################################################
# pad array to improve reciprocal space resolution #
####################################################
amp = pu.crop_pad(amp, (nz1, ny1, nx1))
nz, ny, nx = amp.shape
print('CDI data shape after padding', amp.shape)
####################################################
# interpolate the array with isotropic voxel sizes #
####################################################
if len(voxel_size) == 0:
print('Using isotropic voxel size of 1 nm')
voxel_size = [1, 1, 1] # nm
if not all(voxsize == voxel_size[0] for voxsize in voxel_size):
newvoxelsize = min(voxel_size) # nm
# size of the original object
rgi = RegularGridInterpolator(
(np.arange(-nz // 2, nz // 2, 1) * voxel_size[0],
("HDF5", "*.h5")])
obj, _ = util.load_file(file_path)
if obj.ndim != 3:
print('a 3D reconstruction array is expected')
sys.exit()
obj = abs(obj)
numz, numy, numx = obj.shape
print("Initial data size: (", numz, ',', numy, ',', numx, ')')
#############################
# rotate the reconstruction #
#############################
new_shape = [int(1.2 * numz), int(1.2 * numy), int(1.2 * numx)]
obj = pu.crop_pad(array=obj, output_shape=new_shape, debugging=False)
numz, numy, numx = obj.shape
print("Cropped/padded data size before rotating: (", numz, ',', numy, ',',
numx, ')')
print('Rotating object to have the crystallographic axes along array axes')
obj = pu.rotate_crystal(array=obj,
axis_to_align=axis_outofplane,
reference_axis=np.array([0, 1, 0]),
debugging=True) # out of plane alignement
obj = pu.rotate_crystal(array=obj,
axis_to_align=axis_inplane,
reference_axis=np.array([1, 0, 0]),
debugging=True) # inplane alignement
#################################################
plt.close(fig)
data = data / data.max() # normalize
data[data < support_threshold] = 0
if binary_support:
data[np.nonzero(data)] = 1 # change data into a support
############################################
# go back to original shape before binning #
############################################
if reload_support:
binned_shape = [int(output_shape[idx] / binning_output[idx]) for idx in range(0, len(binning_output))]
else:
binned_shape = [int(original_shape[idx] * binning_original[idx]) for idx in range(0, len(binning_original))]
nz, ny, nx = binned_shape
data = pu.crop_pad(data, binned_shape)
print('Data shape after considering original binning and shape:', data.shape)
######################
# center the support #
######################
if center:
data = pu.center_com(data)
# Use user-defined roll when the center by COM is not optimal
data = np.roll(data, roll_centering, axis=(0, 1, 2))
#################################
# rescale the support if needed #
#################################
nbz, nby, nbx = output_shape
if ((nbz != nz) or (nby != ny) or (nbx != nx)) and not reload_support:
f"({new_voxelsizes[0]:.2f} nm, {new_voxelsizes[1]:.2f} nm, {new_voxelsizes[2]:.2f} nm)"
)
obj = pu.regrid(array=obj,
old_voxelsize=voxel_sizes,
new_voxelsize=new_voxelsizes)
#######################################
# pad the object to the desired shape #
#######################################
if qvalues_shape != padding_shape:
print(
f'\nThe shape defined by q_values {qvalues_shape} is different from padding_shape {padding_shape}'
)
print(f"Overriding padding_shape with {qvalues_shape}")
padding_shape = qvalues_shape
obj = pu.crop_pad(array=obj, output_shape=padding_shape, debugging=debug)
#####################################
# calculate the diffraction pattern #
#####################################
data = fftshift(fftn(obj)) / np.sqrt(reduce(
lambda x, y: x * y, padding_shape)) # complex diffraction amplitude
data = abs(np.multiply(data, np.conjugate(data))) # diffraction intensity
zcom, ycom, xcom = com = tuple(
map(lambda x: int(np.rint(x)), center_of_mass(data)))
print('Center of mass of the diffraction pattern at pixel:', zcom, ycom, xcom)
print(f"\nintensity in a ROI of 7x7x7 voxels centered on the COM:"
f" {int(data[zcom-3:zcom+4, ycom-3:ycom+4, xcom-3:xcom+4].sum())}")
if peak_value is not None:
print('Loading ', nbfiles, 'objects')
if nbfiles == 1:
print('More than one array is needed.')
sys.exit()
##################################################################
# align objects against the first one and stack it in a 4D array #
##################################################################
obj0, _ = util.load_file(file_path[0])
ndim = obj0.ndim
if ndim != 3:
print('3D objects are expected')
sys.exit()
nz, ny, nx = obj0.shape
obj0 = pu.crop_pad(array=obj0, output_shape=(nz+10, ny+10, nx+10))
nz, ny, nx = obj0.shape
print('Array shape', obj0.shape)
amp0 = abs(obj0)
sum0 = amp0.sum()
phase0 = np.angle(obj0)
if alignment_method in ['modulus', 'skip']:
piz0, piy0, pix0 = center_of_mass(amp0)
else: # 'support'
support = np.zeros(amp0.shape)
support[amp0 > support_threshold * amp0.max()] = 1
piz0, piy0, pix0 = center_of_mass(support)
piz0, piy0, pix0 = int(piz0), int(piy0), int(pix0)
phase0 = phase0 - phase0[piz0, piy0, pix0] # set the phase to 0 at the COM of the support
file_path = filedialog.askopenfilenames(initialdir=datadir,
filetypes=[("NPZ", "*.npz"),
("NPY", "*.npy"),
("CXI", "*.cxi"),
("HDF5", "*.h5")])
nbfiles = len(file_path)
print(nbfiles, 'files selected')
#################################################################
# loop through files and calculate the correlation coefficients #
#################################################################
correlation = np.zeros((nbfiles, nbfiles))
for raw in range(nbfiles):
reference_obj, _ = util.load_file(file_path[raw])
reference_obj = abs(reference_obj) / abs(reference_obj).max()
nbz, nby, nbx = reference_obj.shape
reference_obj = pu.crop_pad(array=reference_obj,
output_shape=[nbz + 10, nby + 10, nbx + 10])
correlation[raw, raw] = 1
for col in range(raw + 1, nbfiles):
test_obj, _ = util.load_file(file_path[col]) # which index?
test_obj = abs(test_obj) / abs(test_obj).max()
test_obj = pu.crop_pad(array=test_obj,
output_shape=[nbz + 10, nby + 10, nbx + 10])
# align reconstructions
shiftz, shifty, shiftx = reg.getimageregistration(abs(reference_obj),
abs(test_obj),
precision=100)
test_obj = reg.subpixel_shift(test_obj, shiftz, shifty, shiftx)
print('\nReference =', raw, ' Test =', col)
print('z shift', str('{:.2f}'.format(shiftz)), ', y shift',
str('{:.2f}'.format(shifty)), ', x shift',
str('{:.2f}'.format(shiftx)))
obj, extension = util.load_file(file_path)
obj = abs(obj)
obj = obj / obj.max()
if extension == '.h5':
comment = comment + '_mode'
else:
file_path = filedialog.askopenfilename(initialdir=datadir, title="Select the reconstructed object",
filetypes=[("NPZ", "*.npz")])
obj = np.load(file_path)[field_name]
nbz, nby, nbx = obj.shape
#################
# rotate object #
#################
new_shape = [int(1.2*nbz), int(1.2*nby), int(1.2*nbx)]
obj = pu.crop_pad(array=obj, output_shape=new_shape, debugging=False)
nbz, nby, nbx = obj.shape
print("Cropped/padded object size before rotating: (", nbz, ',', nby, ',', nbx, ')')
print('Rotating object to have the crystallographic axes along array axes')
axis_to_align = np.array([0.2, 1, 0.02]) # in order x y z for rotate_crystal()
obj = pu.rotate_crystal(array=obj, axis_to_align=axis_to_align, reference_axis=np.array([0, 1, 0]),
debugging=True) # out of plane alignement
axis_to_align = np.array([1, 0, -0.1]) # in order x y z for rotate_crystal()
obj = pu.rotate_crystal(array=obj, axis_to_align=axis_to_align, reference_axis=np.array([1, 0, 0]),
debugging=True) # inplane alignement
###################
# apply threshold #
###################
if not np.isnan(threshold):
and crop_center[2]+output_shape[2]//2 <= nbx, 'crop_center incompatible with output_shape'
corr_roi = corr_roi or [0, nbz, 0, nby, 0, nbx
] # if None, use the full array for corr_roi
assert len(corr_roi) == 6, 'corr_roi should be a tuple or list of lenght 6'
if not 0 <= corr_roi[0] < corr_roi[1] <= nbz\
or not 0 <= corr_roi[2] < corr_roi[3] <= nby\
or not 0 <= corr_roi[4] < corr_roi[5] <= nbx:
print('Incorrect value for the parameter corr_roi')
sys.exit()
# crop the data directly to output_shape if no alignment is required, update corr_roi accordingly
if alignement_method == 'skip':
refmask = pu.crop_pad(array=refmask,
output_shape=output_shape,
crop_center=crop_center)
refdata = pu.crop_pad(array=refdata,
output_shape=output_shape,
crop_center=crop_center)
# correct for the offset due to cropping
corr_roi = [
corr_roi[0] - (crop_center[0] - output_shape[0] // 2),
corr_roi[1] - (crop_center[0] - output_shape[0] // 2),
corr_roi[2] - (crop_center[1] - output_shape[1] // 2),
corr_roi[3] - (crop_center[1] - output_shape[1] // 2),
corr_roi[4] - (crop_center[2] - output_shape[2] // 2),
corr_roi[5] - (crop_center[2] - output_shape[2] // 2)
]
# check if this still fits the cropped data, default to the data range otherwise
corr_roi[0] = max(0, corr_roi[0])
