def test_get_weight_and_index_off_center():
"""Test get_weight_and_index for off-centered points."""
pixel_position = xp.array([[0.1, 0.2, 0.3]])
voxel_length = 2.
voxel_num_1d = 5
index, weight = mapping.get_weight_and_index(
pixel_position, voxel_length, voxel_num_1d)
assert xp.all(index[0][0] == xp.array([2, 2, 2]))
assert xp.isclose(weight.sum(), 1.)
assert xp.allclose(xp.dot(weight[0], index[0]),
xp.array([ 2.05, 2.1 , 2.15]))
def test_take_slice_center_value(self):
"""Test take_slice on center value."""
orientation = cst.quat1
pixel_momentum = xp.array([[[[0., 0., 0.]]]])
slice_ = sk.take_slice(self.volume, self.voxel_length, pixel_momentum,
orientation)
assert xp.isclose(slice_[0, 0, 0], self.volume.mean())
def test_get_weight_and_index_center_even():
"""Test get_weight_and_index centering for even-sized meshes."""
pixel_position = xp.zeros((1,3), dtype=float)
voxel_length = 1.
voxel_num_1d = 4
index, weight = mapping.get_weight_and_index(
pixel_position, voxel_length, voxel_num_1d)
assert xp.all(index[0][0] == xp.array([1, 1, 1]))
assert xp.allclose(weight[0], 0.125)
def distance(self, value):
if not xp.allclose(self.orientation, xp.array([0, 0, 1])):
raise NotImplementedError(
"Detector distance setter only implemented for "
"detector orientations along the z axis.")
self.pixel_position[..., 2] *= value / self._distance
self._distance = value
if self._has_beam: # Update pixel_position_reciprocal & co
self.initialize_pixels_with_beam(beam=self._beam)
def __init__(self):
# Reciprocal space information is only available if detector has access to beam.
self._has_beam = False
# Define the hierarchy system. For simplicity, we only use two-layer structure.
self.panel_num = 1
# Define all properties the detector should have
self._distance = 1 # (m) detector distance
self.pixel_width = 0 # (m)
self.pixel_height = 0 # (m)
self.pixel_area = 0 # (m^2)
self.panel_pixel_num_x = 0 # number of pixels in x
self.panel_pixel_num_y = 0 # number of pixels in y
self.pixel_num_total = 0 # total number of pixels (px*py)
self.center_x = 0 # center of detector in x
self.center_y = 0 # center of detector in y
self.orientation = xp.array([0, 0, 1])
self.pixel_position = None # (m)
self.pixel_position_ideal = None #(m)
# pixel information in reciprocal space
self.pixel_position_reciprocal = None # (m^-1)
self.pixel_distance_reciprocal = None # (m^-1)
# Pixel map
self.pixel_index_map = None
self.detector_pixel_num_x = 1
self.detector_pixel_num_y = 1
# Corrections
self.solid_angle_per_pixel = None # solid angle
self.polarization_correction = None # Polarization correction
"""
The theoretical differential cross section of an electron ignoring the
polarization effect is,
do/dO = ( e^2/(4*Pi*epsilon0*m*c^2) )^2 * ( 1 + cos(xi)^2 )/2
Therefore, one needs to includes the leading constant factor which is the
following numerical value.
"""
# Tompson Scattering factor
self.Thomson_factor = 2.817895019671143 * 2.817895019671143 * 1e-30
# Total scaling and correction factor.
self.linear_correction = None
# Detector effects
self._pedestal = 0
self._pixel_rms = 0
self._pixel_bkgd = 0
self._pixel_status = 0
self._pixel_mask = 0
self._pixel_gain = 0
# self.geometry currently only work for the pre-defined detectors
self.geometry = None
def test_get_weight_and_index_center_odd():
"""Test get_weight_and_index centering for odd-sized meshes."""
pixel_position = xp.zeros((1,3), dtype=float)
voxel_length = 1.
voxel_num_1d = 5
index, weight = mapping.get_weight_and_index(
pixel_position, voxel_length, voxel_num_1d)
assert xp.all(index[0][0] == xp.array([2, 2, 2]))
assert weight[0][0] == 1
assert xp.all(weight[0][1:] == 0)
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 or an array of zeros is if dimensions are incompatible.
"""
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
try:
data_to_merge = volume_1d[index_2d]
except IndexError:
print("Requested slice and diffraction volume have incompatible dimensions")
return np.zeros(pattern_shape)
# 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 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 as s-vectors, where q=2*pi*s
: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))
qvectors_1d = 2*np.pi*reciprocal_space_1d
# Calculate atom form factor for the reciprocal space, passing in sin(theta)/lambda in per Angstrom
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(qvectors_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)
# Add another contribution if defined, e.g. virus void...
if particle.other_mask is not None:
other_position = np.ascontiguousarray(particle.mesh[particle.other_mask,:])
other_num = np.sum(particle.other_mask)
other_prefactor = particle.other_mean_electron_density * particle.mesh_voxel_size**3
cuda_other_position = cuda.to_device(other_position)
calculate_solvent_pattern_gpu_back_engine[(pixel_number + 511) // 512, 512](
cuda_reciprocal_position, cuda_other_position,
pattern_cos, pattern_sin, other_prefactor, other_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)