def get_reciprocal_space_pixel_position(pixel_center, wave_vector):
"""
Obtain the coordinate of each pixel in the reciprocal space.
:param pixel_center: The coordinate of the pixel in real space.
:param wave_vector: The wavevector.
:return: The array containing the pixel coordinates.
"""
wave_vector = xp.asarray(wave_vector)
# reshape the array into a 1d position array
pixel_center_1d = reshape_pixels_position_arrays_to_1d(pixel_center)
# Calculate the reciprocal position of each pixel
wave_vector_norm = xp.sqrt(xp.sum(xp.square(wave_vector)))
wave_vector_direction = wave_vector / wave_vector_norm
pixel_center_norm = xp.sqrt(xp.sum(xp.square(pixel_center_1d), axis=1))
pixel_center_direction = pixel_center_1d / pixel_center_norm[:, np.newaxis]
pixel_position_reciprocal_1d = wave_vector_norm * (pixel_center_direction -
wave_vector_direction)
# restore the pixels shape
pixel_position_reciprocal = xp.reshape(pixel_position_reciprocal_1d,
pixel_center.shape)
return pixel_position_reciprocal
def get_reciprocal_position_and_correction(pixel_position, pixel_area,
wave_vector, polarization,
orientation):
"""
Calculate the pixel positions in reciprocal space and all the related corrections.
:param pixel_position: The position of the pixel in real space.
:param wave_vector: The wavevector.
:param polarization: The polarization vector.
:param orientation: The normal direction of the detector.
:param pixel_area: The pixel area for each pixel. In pixel stack format.
:return: pixel_position_reciprocal, pixel_position_reciprocal_norm, polarization_correction,
geometry_correction
"""
# Calculate the position and distance in reciprocal space
pixel_position_reciprocal = get_reciprocal_space_pixel_position(
pixel_center=pixel_position, wave_vector=wave_vector)
pixel_position_reciprocal_norm = xp.sqrt(
xp.sum(xp.square(pixel_position_reciprocal), axis=-1))
# Calculate the corrections.
polarization_correction = get_polarization_correction(
pixel_center=pixel_position, polarization=polarization)
# Because the pixel area in this function is measured in m^2,
# therefore,the distance has to be in m
solid_angle_array = solid_angle(pixel_center=pixel_position,
pixel_area=pixel_area,
orientation=orientation)
return (pixel_position_reciprocal, pixel_position_reciprocal_norm,
polarization_correction, solid_angle_array)
def extract_slice(local_index, local_weight, volume):
"""
Take one slice from the volume given the index and weight map.
:param local_index: The index containing values to take.
:param local_weight: The weight for each index.
:param volume: The volume to slice from.
:return: The slice.
"""
local_index = xp.asarray(local_index)
local_weight = xp.asarray(local_weight)
volume = xp.asarray(volume)
# Convert the index of the 3D diffraction volume to 1D
pattern_shape = local_index.shape[:3]
pixel_num = int(np.prod(pattern_shape))
volume_num_1d = volume.shape[0]
convertion_factor = xp.array(
[volume_num_1d * volume_num_1d, volume_num_1d, 1], dtype=np.int64)
index_2d = xp.reshape(local_index, [pixel_num, 8, 3])
index_2d = xp.matmul(index_2d, convertion_factor)
volume_1d = xp.reshape(volume, volume_num_1d**3)
weight_2d = xp.reshape(local_weight, [pixel_num, 8])
# Expand the data to merge
data_to_merge = volume_1d[index_2d]
# Merge the data
data_merged = xp.sum(xp.multiply(weight_2d, data_to_merge), axis=-1)
data = xp.reshape(data_merged, pattern_shape)
return data
def solid_angle(pixel_center, pixel_area, orientation):
"""
Calculate the solid angle for each pixel.
:param pixel_center: The position of each pixel in real space. In pixel stack format.
:param orientation: The orientation of the detector.
:param pixel_area: The pixel area for each pixel. In pixel stack format.
:return: Solid angle of each pixel.
"""
orientation = xp.asarray(orientation)
# Use 1D format
pixel_center_1d = reshape_pixels_position_arrays_to_1d(pixel_center)
pixel_center_norm_1d = xp.sqrt(xp.sum(xp.square(pixel_center_1d), axis=-1))
pixel_area_1d = xp.reshape(pixel_area, np.prod(pixel_area.shape))
# Calculate the direction of each pixel.
pixel_center_direction_1d = pixel_center_1d / pixel_center_norm_1d[:, xp.
newaxis]
# Normalize the orientation vector
orientation_norm = xp.sqrt(xp.sum(xp.square(orientation)))
orientation_normalized = orientation / orientation_norm
# The correction induced by projection which is a factor of cosine.
cosine_1d = xp.abs(
xp.dot(pixel_center_direction_1d, orientation_normalized))
# Calculate the solid angle ignoring the projection
solid_angle_1d = xp.divide(pixel_area_1d, xp.square(pixel_center_norm_1d))
solid_angle_1d = xp.multiply(cosine_1d, solid_angle_1d)
# Restore the pixel stack format
solid_angle_stack = xp.reshape(solid_angle_1d, pixel_area.shape)
return solid_angle_stack
def preferred_voxel_length(self, wave_vector):
"""
If one want to put the diffraction pattern into 3D reciprocal space, then one needs to
select a proper voxel length for the reciprocal space. This function gives a reasonable
estimation of this length
:param wave_vector: The wavevector of in this experiment.
:return: voxel_length.
"""
# Notice that this voxel length has nothing to do with the voxel length
# utilized in dragonfly.
voxel_length = xp.sqrt(xp.sum(xp.square(wave_vector)))
voxel_length /= self.distance * xp.min(self.pixel_width, self.pixel_height)
return voxel_length
def calculate_diffraction_pattern_gpu(reciprocal_space,
particle,
return_type='intensity'):
"""
Calculate the diffraction field of the specified reciprocal space.
:param reciprocal_space: The reciprocal space over which to calculate the diffraction field.
:param particle: The particle object to calculate the diffraction field.
:param return_type: 'intensity' to return the intensity field. 'complex_field' to return the full diffraction field.
:return: The diffraction field.
"""
"""This function can be used to calculate the diffraction field for
arbitrary reciprocal space """
# convert the reciprocal space into a 1d series.
shape = reciprocal_space.shape
pixel_number = int(np.prod(shape[:-1]))
reciprocal_space_1d = xp.reshape(reciprocal_space, [pixel_number, 3])
reciprocal_norm_1d = xp.sqrt(
xp.sum(xp.square(reciprocal_space_1d), axis=-1))
# Calculate atom form factor for the reciprocal space
form_factor = pd.calculate_atomic_factor(
particle=particle,
q_space=reciprocal_norm_1d * (1e-10 / 2.), # For unit compatibility
pixel_num=pixel_number)
# Get atom position
atom_position = np.ascontiguousarray(particle.atom_pos[:])
atom_type_num = len(particle.split_idx) - 1
# create
pattern_cos = xp.zeros(pixel_number, dtype=xp.float64)
pattern_sin = xp.zeros(pixel_number, dtype=xp.float64)
# atom_number = atom_position.shape[0]
split_index = xp.array(particle.split_idx)
cuda_split_index = cuda.to_device(split_index)
cuda_atom_position = cuda.to_device(atom_position)
cuda_reciprocal_position = cuda.to_device(reciprocal_space_1d)
cuda_form_factor = cuda.to_device(form_factor)
# Calculate the pattern
calculate_pattern_gpu_back_engine[(pixel_number + 511) // 512,
512](cuda_form_factor,
cuda_reciprocal_position,
cuda_atom_position, pattern_cos,
pattern_sin, atom_type_num,
cuda_split_index, pixel_number)
# Add the hydration layer
if particle.mesh is not None:
water_position = np.ascontiguousarray(
particle.mesh[particle.solvent_mask, :])
water_num = np.sum(particle.solvent_mask)
water_prefactor = particle.solvent_mean_electron_density * particle.mesh_voxel_size**3
cuda_water_position = cuda.to_device(water_position)
calculate_solvent_pattern_gpu_back_engine[(pixel_number + 511) // 512,
512](
cuda_reciprocal_position,
cuda_water_position,
pattern_cos, pattern_sin,
water_prefactor,
water_num, pixel_number)
if return_type == "intensity":
pattern = np.reshape(np.square(np.abs(pattern_cos + 1j * pattern_sin)),
shape[:-1])
return xp.asarray(pattern)
elif return_type == "complex_field":
pattern = np.reshape(pattern_cos + 1j * pattern_sin, shape[:-1])
return xp.asarray(pattern)
else:
print(
"Please set the parameter return_type = 'intensity' or 'complex_field'"
)
print("This time, this program return the complex field.")
pattern = np.reshape(pattern_cos + 1j * pattern_sin, shape[:-1])
return xp.asarray(pattern)
