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_fxs_photons_slices(self,
particles,
beam_focus_radius,
jet_radius,
mesh_length,
device=None):
mesh, voxel_length = self.get_reciprocal_mesh(
voxel_number_1d=mesh_length)
state, coords = distribute_particles(particles, beam_focus_radius,
jet_radius)
count = 0
field_acc = xp.zeros(self.pixel_position_reciprocal.shape[:3],
dtype=xp.complex128)
for particle in particles:
if particles[particle] > 0:
volume = pgd.calculate_diffraction_pattern_gpu(
mesh, particle, return_type="complex_field")
orientations = ps.geometry.get_random_quat(
num_pts=particles[particle])
slices = ps.geometry.take_n_slices(
volume=volume,
voxel_length=voxel_length,
pixel_momentum=self.pixel_position_reciprocal,
orientations=orientations)
for i in range(particles[particle]):
field_acc += self.add_phase_shift(slices[i], coords[count])
count += 1
return self.add_correction_and_quantization(
xp.square(xp.abs(field_acc)))
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 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 get_fxs_photons_unittest(self,
particles,
beam_focus_radius,
jet_radius,
device=None):
raw_data = None
state, coords = distribute_particles(particles, beam_focus_radius,
jet_radius)
for i in range(len(state)):
this_data = self.get_fxs_pattern_without_corrections(
particle=state[i],
coords=coords[i],
return_type="complex_field")
if raw_data is None:
raw_data = this_data
else:
raw_data += this_data
return self.add_correction_and_quantization(xp.square(
xp.abs(raw_data)))
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)