def take_n_slices(volume, voxel_length, pixel_momentum, orientations,
inverse=False):
"""
Take several slices.
:param volume: The volume to slice from.
:param voxel_length: The length unit of the voxel
:param pixel_momentum: The coordinate of each pixel in the
reciprocal space measured in A.
:param orientations: The orientations of the slices.
:param inverse: Whether to use the inverse of the rotation or not.
:return: n slices.
"""
# Preprocess
slice_num = orientations.shape[0]
pattern_shape = pixel_momentum.shape[:-1]
# Create variable to hold the slices
slices_holder = xp.zeros((slice_num,) + pattern_shape, dtype=volume.dtype)
for l in range(slice_num):
slices_holder[l] = take_slice(volume, voxel_length, pixel_momentum,
orientations[l], inverse)
return slices_holder
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 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 assemble_image_stack(self, image_stack):
"""
Assemble the image stack into a 2D diffraction pattern.
For this specific object, since it only has one panel, the result is to remove the
first dimension.
:param image_stack: The [1, num_x, num_y] numpy array.
:return: The [num_x, num_y] numpy array.
"""
if self.pixel_index_map is None:
raise RuntimeError(
"This detector does not have pixel mapping information.")
# construct the image holder:
image = xp.zeros(
(self.detector_pixel_num_x, self.detector_pixel_num_y))
for l in range(self.panel_num):
image[self.pixel_index_map[l, :, :, 0],
self.pixel_index_map[l, :, :, 1]] = image_stack[l, :, :]
return image
def assemble_image_stack_batch(self, image_stack_batch):
"""
Assemble the image stack batch into a stack of 2D diffraction patterns.
For this specific object, since it has only one panel, the result is a simple reshape.
:param image_stack_batch: The [stack_num, 1, num_x, num_y] numpy array
:return: The [stack_num, num_x, num_y] numpy array
"""
if self.pixel_index_map is None:
raise RuntimeError(
"This detector does not have pixel mapping information.")
stack_num = image_stack_batch.shape[0]
# construct the image holder:
image = xp.zeros(
(stack_num, self.detector_pixel_num_x, self.detector_pixel_num_y))
for l in range(self.panel_num):
idx_map_1 = self.pixel_index_map[l, :, :, 0]
idx_map_2 = self.pixel_index_map[l, :, :, 1]
image[:, idx_map_1, idx_map_2] = image_stack_batch[:, l]
return image
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)
def get_weight_and_index(pixel_position, voxel_length, voxel_num_1d):
"""
Obtain the weight of the pixel for adjacent voxels.
In this function, pixel position is first cast to the shape [pixel number,3].
:param pixel_position: The position of each pixel in the space.
:param voxel_length:
:param voxel_num_1d:
:return:
"""
# Extract the detector shape
detector_shape = pixel_position.shape[:-1]
pixel_num = np.prod(detector_shape)
# Cast the position infor to the shape [pixel number, 3]
pixel_position = xp.asarray(pixel_position)
pixel_position_1d = xp.reshape(pixel_position, (pixel_num, 3))
# convert_to_voxel_unit
pixel_position_1d_voxel_unit = pixel_position_1d / voxel_length
# shift the center position
shift = (voxel_num_1d - 1) / 2.
pixel_position_1d_voxel_unit += shift
# Get one nearest neighbor
tmp_index = xp.floor(pixel_position_1d_voxel_unit).astype(np.int64)
# Generate the holders
indexes = xp.zeros((pixel_num, 8, 3), dtype=np.int64)
weight = xp.ones((pixel_num, 8), dtype=np.float64)
# Calculate the floors and the ceilings
dfloor = pixel_position_1d_voxel_unit - tmp_index
dceiling = 1 - dfloor
# Assign the correct values to the indexes
indexes[:, 0, :] = tmp_index
indexes[:, 1, 0] = tmp_index[:, 0]
indexes[:, 1, 1] = tmp_index[:, 1]
indexes[:, 1, 2] = tmp_index[:, 2] + 1
indexes[:, 2, 0] = tmp_index[:, 0]
indexes[:, 2, 1] = tmp_index[:, 1] + 1
indexes[:, 2, 2] = tmp_index[:, 2]
indexes[:, 3, 0] = tmp_index[:, 0]
indexes[:, 3, 1] = tmp_index[:, 1] + 1
indexes[:, 3, 2] = tmp_index[:, 2] + 1
indexes[:, 4, 0] = tmp_index[:, 0] + 1
indexes[:, 4, 1] = tmp_index[:, 1]
indexes[:, 4, 2] = tmp_index[:, 2]
indexes[:, 5, 0] = tmp_index[:, 0] + 1
indexes[:, 5, 1] = tmp_index[:, 1]
indexes[:, 5, 2] = tmp_index[:, 2] + 1
indexes[:, 6, 0] = tmp_index[:, 0] + 1
indexes[:, 6, 1] = tmp_index[:, 1] + 1
indexes[:, 6, 2] = tmp_index[:, 2]
indexes[:, 7, :] = tmp_index + 1
# Assign the correct values to the weight
weight[:, 0] = xp.prod(dceiling, axis=-1)
weight[:, 1] = dceiling[:, 0] * dceiling[:, 1] * dfloor[:, 2]
weight[:, 2] = dceiling[:, 0] * dfloor[:, 1] * dceiling[:, 2]
weight[:, 3] = dceiling[:, 0] * dfloor[:, 1] * dfloor[:, 2]
weight[:, 4] = dfloor[:, 0] * dceiling[:, 1] * dceiling[:, 2]
weight[:, 5] = dfloor[:, 0] * dceiling[:, 1] * dfloor[:, 2]
weight[:, 6] = dfloor[:, 0] * dfloor[:, 1] * dceiling[:, 2]
weight[:, 7] = xp.prod(dfloor, axis=-1)
# Change the shape of the index and weight variable
indexes = xp.reshape(indexes, detector_shape + (8, 3))
weight = xp.reshape(weight, detector_shape + (8, ))
return indexes, weight
def __init__(self, N_pixel, det_size, det_distance, beam=None):
"""
Initialize the detector
:param N_pixel: Number of pixels per dimension.
:param det_size: Length of detector sides (m).
:param det_distance: Sample-Detector distance (m).
"""
super(SimpleSquareDetector, self).__init__()
ar = xp.arange(N_pixel)
sep = float(det_size) / N_pixel
x = -det_size / 2 + sep / 2 + ar * sep
y = -det_size / 2 + sep / 2 + ar * sep
X, Y = xp.meshgrid(x, y, indexing='xy')
Xar, Yar = xp.meshgrid(ar, ar, indexing='xy')
self.panel_num = 1
self.panel_pixel_num_x = N_pixel
self.panel_pixel_num_y = N_pixel
self._shape = (1, N_pixel, N_pixel)
# Define all properties the detector should have
self._distance = det_distance
self.center_z = self._distance * xp.ones(self._shape, dtype=xp.float64)
p_center_x = xp.stack((X, ))
p_center_y = xp.stack((Y, ))
self.pixel_width = sep * xp.ones(self._shape, dtype=xp.float64)
self.pixel_height = sep * xp.ones(self._shape, dtype=xp.float64)
self.center_x = p_center_x
self.center_y = p_center_y
# construct the the pixel position array
self.pixel_position = xp.zeros(self._shape + (3, ))
self.pixel_position[..., 0] = self.center_x
self.pixel_position[..., 1] = self.center_y
self.pixel_position[..., 2] = self.center_z
# Pixel map
p_map_x = xp.stack((Xar, ))
p_map_y = xp.stack((Yar, ))
# [panel number, pixel num x, pixel num y]
self.pixel_index_map = xp.stack((p_map_x, p_map_y), axis=-1)
# Detector pixel number info
self.detector_pixel_num_x = asnumpy(
xp.max(self.pixel_index_map[..., 0]) + 1)
self.detector_pixel_num_y = asnumpy(
xp.max(self.pixel_index_map[..., 1]) + 1)
# Panel pixel number info
# number of pixels in each panel in x/y direction
self.panel_pixel_num_x = self.pixel_index_map.shape[1]
self.panel_pixel_num_y = self.pixel_index_map.shape[2]
# total number of pixels (px*py)
self.pixel_num_total = np.prod(self._shape)
# Calculate the pixel area
self.pixel_area = xp.multiply(self.pixel_height, self.pixel_width)
# Get reciprocal space configurations and corrections.
self.initialize_pixels_with_beam(beam=beam)
def initialize(self, geom, beam):
"""
Initialize the detector with user defined parameters
:param geom: The dictionary containing all the necessary information to initialized the
detector.
:param beam: The beam object
:return: None
"""
"""
Doc:
To use this class, the user has to provide the necessary information to initialize
the detector.
All the necessary entries are listed in the example notebook.
"""
# 'detector distance': detector distance in (m)
##########################################################################################
# Extract necessary information
##########################################################################################
# Define the hierarchy system. For simplicity, we only use two-layer structure.
for key in {'panel number', 'panel pixel num x', 'panel pixel num y'}:
if key not in geom:
raise KeyError("Missing required '{}' key.".format(key))
self.panel_num = int(geom['panel number'])
self.panel_pixel_num_x = int(geom['panel pixel num x'])
self.panel_pixel_num_y = int(geom['panel pixel num y'])
self._shape = (self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y)
# Define all properties the detector should have
self._distance = None
if 'pixel center z' in geom:
if 'detector distance' in geom:
raise ValueError("Please provide one of "
"'pixel center z' or 'detector distance'.")
self.center_z = xp.asarray(geom['pixel center z'],
dtype=xp.float64)
self._distance = float(self.center_z.mean())
else:
if 'detector distance' not in geom:
KeyError("Missing required 'detector distance' key.")
self._distance = float(geom['detector distance'])
self.center_z = self._distance * xp.ones(self._shape,
dtype=xp.float64)
# Below: [panel number, pixel num x, pixel num y] in (m)
# Change dtype and make numpy/cupy array
self.pixel_width = xp.asarray(geom['pixel width'], dtype=xp.float64)
self.pixel_height = xp.asarray(geom['pixel height'], dtype=xp.float64)
self.center_x = xp.asarray(geom['pixel center x'], dtype=xp.float64)
self.center_y = xp.asarray(geom['pixel center y'], dtype=xp.float64)
self.orientation = np.array([0, 0, 1])
# construct the the pixel position array
self.pixel_position = xp.zeros(self._shape + (3, ))
self.pixel_position[..., 0] = self.center_x
self.pixel_position[..., 1] = self.center_y
self.pixel_position[..., 2] = self.center_z
# Pixel map
if 'pixel map' in geom:
# [panel number, pixel num x, pixel num y]
self.pixel_index_map = xp.asarray(geom['pixel map'],
dtype=xp.int64)
# Detector pixel number info
self.detector_pixel_num_x = asnumpy(
xp.max(self.pixel_index_map[..., 0]) + 1)
self.detector_pixel_num_y = asnumpy(
xp.max(self.pixel_index_map[..., 1]) + 1)
# Panel pixel number info
# number of pixels in each panel in x/y direction
self.panel_pixel_num_x = self.pixel_index_map.shape[1]
self.panel_pixel_num_y = self.pixel_index_map.shape[2]
# total number of pixels (px*py)
self.pixel_num_total = np.prod(self._shape)
###########################################################################################
# Do necessary calculation to finishes the initialization
###########################################################################################
# self.geometry currently only work for the pre-defined detectors
self.geometry = geom
# Calculate the pixel area
self.pixel_area = xp.multiply(self.pixel_height, self.pixel_width)
# Get reciprocal space configurations and corrections.
self.initialize_pixels_with_beam(beam=beam)
##########################################################################################
# Do necessary calculation to finishes the initialization
##########################################################################################
# Detector effects
if 'pedestal' in geom:
self._pedestal = xp.asarray(geom['pedestal'], dtype=xp.float64)
else:
self._pedestal = xp.zeros((self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
if 'pixel rms' in geom:
self._pixel_rms = xp.asarray(geom['pixel rms'], dtype=xp.float64)
else:
self._pixel_rms = xp.zeros((self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
if 'pixel bkgd' in geom:
self._pixel_bkgd = xp.asarray(geom['pixel bkgd'], dtype=xp.float64)
else:
self._pixel_bkgd = xp.zeros(
(self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
if 'pixel status' in geom:
self._pixel_status = xp.asarray(geom['pixel status'],
dtype=xp.float64)
else:
self._pixel_status = xp.zeros(
(self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
if 'pixel mask' in geom:
self._pixel_mask = xp.asarray(geom['pixel mask'], dtype=xp.float64)
else:
self._pixel_mask = xp.zeros(
(self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
if 'pixel gain' in geom:
self._pixel_gain = xp.asarray(geom['pixel gain'], dtype=xp.float64)
else:
self._pixel_gain = xp.ones((self.panel_num, self.panel_pixel_num_x,
self.panel_pixel_num_y))
def initialize(self, geom, run_num=0, cframe=0):
"""
Initialize the detector
:param geom: The *-end.data file which characterizes the geometry profile.
:param run_num: The run_num containing the background, rms and gain and the other
pixel pixel properties.
:param cframe: The desired coordinate frame, 0 for psana and 1 for lab conventions.
:return: None
"""
# Redirect the output stream
old_stdout = sys.stdout
f = six.StringIO()
# f = open('Detector_initialization.log', 'w')
sys.stdout = f
###########################################################################################
# Initialize the geometry configuration
############################################################################################
self.geometry = GeometryAccess(geom, cframe=cframe)
self.run_num = run_num
# Set coordinate in real space (convert to m)
temp = [xp.asarray(t) * 1e-6 for t in self.geometry.get_pixel_coords(cframe=cframe)]
temp_index = [xp.asarray(t)
for t in self.geometry.get_pixel_coord_indexes(cframe=cframe)]
self.panel_num = np.prod(temp[0].shape[:-2])
self._distance = float(temp[2].mean())
self._shape = (self.panel_num, temp[0].shape[-2], temp[0].shape[-1])
self.pixel_position = xp.zeros(self._shape + (3,))
self.pixel_index_map = xp.zeros(self._shape + (2,))
for n in range(3):
self.pixel_position[..., n] = temp[n].reshape(self._shape)
for n in range(2):
self.pixel_index_map[..., n] = temp_index[n].reshape(self._shape)
self.pixel_index_map = self.pixel_index_map.astype(xp.int64)
# Get the range of the pixel index
self.detector_pixel_num_x = asnumpy(
xp.max(self.pixel_index_map[..., 0]) + 1)
self.detector_pixel_num_y = asnumpy(
xp.max(self.pixel_index_map[..., 1]) + 1)
self.panel_pixel_num_x = np.array([self.pixel_index_map.shape[1], ] * self.panel_num)
self.panel_pixel_num_y = np.array([self.pixel_index_map.shape[2], ] * self.panel_num)
self.pixel_num_total = np.sum(np.multiply(self.panel_pixel_num_x, self.panel_pixel_num_y))
tmp = float(self.geometry.get_pixel_scale_size() * 1e-6) # Convert to m
self.pixel_width = xp.ones(
(self.panel_num, self.panel_pixel_num_x[0], self.panel_pixel_num_y[0])) * tmp
self.pixel_height = xp.ones(
(self.panel_num, self.panel_pixel_num_x[0], self.panel_pixel_num_y[0])) * tmp
# Calculate the pixel area
self.pixel_area = xp.multiply(self.pixel_height, self.pixel_width)
###########################################################################################
# Initialize the pixel effects
###########################################################################################
# first we should parse the path
parsed_path = geom.split('/')
self.exp = parsed_path[-5]
if self.exp == 'calib':
self.exp = parsed_path[-6]
self.group = parsed_path[-4]
self.source = parsed_path[-3]
self._pedestals = None
self._pixel_rms = None
self._pixel_mask = None
self._pixel_bkgd = None
self._pixel_status = None
self._pixel_gain = None
if psana_version==1:
try:
cbase = self._get_cbase()
self.calibdir = '/'.join(parsed_path[:-4])
pbits = 255
gcp = GenericCalibPars(cbase, self.calibdir, self.group, self.source, run_num, pbits)
self._pedestals = gcp.pedestals()
self._pixel_rms = gcp.pixel_rms()
self._pixel_mask = gcp.pixel_mask()
self._pixel_bkgd = gcp.pixel_bkgd()
self._pixel_status = gcp.pixel_status()
self._pixel_gain = gcp.pixel_gain()
except NotImplementedError:
# No GenericCalibPars information.
pass
else:
try:
self.det = self._get_det_id(self.group)
except NotImplementedError:
# No GenericCalibPars information.
self.det = None
# Redirect the output stream
sys.stdout = old_stdout