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)