def from_translation_rotation(translation, rotation, com_reference,
com_moving):
transform = gemmi.Transform()
transform.vec.fromlist(translation.tolist())
transform.mat.fromlist(rotation.as_matrix().tolist())
return Transform(transform, com_reference, com_moving)
def sample_map(
gemmi_grid,
centroid,
rotation,
translation,
shape,
scale=0.5,
):
rotation = np.matmul(
rotation,
np.eye(3) * scale,
)
offset = -(np.array(shape) / 2) * scale
rotated_offset = np.matmul(rotation, offset) + translation
offset_translation = centroid + rotated_offset
arr = np.zeros(
shape,
dtype=np.float32,
)
tr = gemmi.Transform()
tr.mat.fromlist(rotation.tolist())
tr.vec.fromlist(offset_translation.tolist())
gemmi_grid.interpolate_values(
arr,
tr,
)
return arr
def test_rotation(structure, orientation):
"""Test gemmi rotation transformation does what we expect"""
sites = structure.get_all_unit_cell_sites()
tr = gemmi.Transform()
tr.mat.fromlist(orientation.as_list_of_lists())
orig_pos = [site.orth(structure.cell) for site in sites]
plot_atom_positions(orig_pos)
new_pos = [tr.apply(pos) for pos in orig_pos]
for vector, rotated in zip(orig_pos, new_pos):
rotated2 = orientation * matrix.col(vector.tolist())
for i in range(3):
assert rotated2[i] == pytest.approx(rotated[i])
def sample_dataset_global(unaligned_xmap, sample_region: SampleRegion):
tr = gemmi.Transform()
tr.mat.fromlist(sample_region.rotation.tolist())
tr.vec.fromlist(sample_region.origin.tolist())
arr = np.zeros(
(
int(sample_region.gridding[0]),
int(sample_region.gridding[1]),
int(sample_region.gridding[2]),
),
dtype=np.float32,
)
unaligned_xmap.interpolate_values(arr, tr)
return arr
def perturb_sample_region(transformed_sample_region, perturbation):
# Split out the components of the perturbation
transformation_perturbation = perturbation[0:3]
rotation_perturbation = perturbation[3:6]
#
rotation_perturbation_obj = scipy.spatial.transform.Rotation.from_euler(
"xyz",
[rotation_perturbation[0], rotation_perturbation[1], rotation_perturbation[2]], degrees=True)
rotation_perturbation_mat = rotation_perturbation_obj.as_matrix()
# Package them as a transform
transform = gemmi.Transform()
transform.vec.fromlist(transformation_perturbation.tolist())
transform.mat.fromlist(rotation_perturbation_mat.tolist())
# Get the perturbed sample region
perturbed_sample_region = get_transformed_sample_region(transformed_sample_region, Transform(transform))
return perturbed_sample_region
def sample_map(
gemmi_grid,
centroid,
rotation,
translation,
shape,
):
offset = translation - (np.array(shape) / 2)
rotated_offset = np.matmul(rotation, offset)
offset_translation = centroid + rotated_offset
arr = np.zeros(shape)
tr = gemmi.Transform()
tr.mat.fromlist(rotation.tolist())
tr.vec.fromlist(offset_translation.tolist())
gemmi_grid.interpolate_values(
arr,
tr,
)
return arr
def sample(xmap, parameters):
shape = 32
scale = 0.5
arr = np.zeros([32, 32, 32], dtype=np.float32)
tr = gemmi.Transform()
trans = [
parameters[0],
parameters[1],
parameters[2],
]
angles = [
parameters[3],
parameters[4],
parameters[5],
]
rotation_x = Rotation.from_euler("x", angles[0])
rotation_y = Rotation.from_euler("y", angles[1])
rotation_z = Rotation.from_euler("z", angles[2])
scale = np.eye(3) * 0.5
rotation = np.matmul(
rotation_x.as_matrix(),
np.matmul(
rotation_y.as_matrix(),
rotation_z.as_matrix(),
),
)
rotation = np.matmul(rotation, scale)
trans = trans - np.matmul(rotation, np.array([shape, shape, shape]) / 2)
tr.mat.fromlist(rotation.tolist())
tr.vec.fromlist(trans)
xmap.grid.interpolate_values(arr, tr)
return arr
def to_gemmi(self):
transform_gemmi = gemmi.Transform()
transform_gemmi.mat.fromlist(self.transform)
return transform_gemmi
def resample(
reference_xmap: gemmi.FloatGrid,
moving_xmap: gemmi.FloatGrid,
reference_structure: gemmi.Structure,
moving_structure: gemmi.Structure,
monomerized=False,
):
# Get transform: from ref to align
transform = get_alignment(moving_structure, reference_structure, monomerized=monomerized)
print(f"Transform: {transform}; {transform.transform.vec} {transform.transform.mat}")
interpolated_grid = gemmi.FloatGrid(
reference_xmap.nu,
reference_xmap.nv,
reference_xmap.nw,
)
interpolated_grid.set_unit_cell(reference_xmap.unit_cell)
interpolated_grid.spacegroup = reference_xmap.spacegroup
# points
mask = gemmi.FloatGrid(reference_xmap.nu,
reference_xmap.nv,
reference_xmap.nw, )
mask.set_unit_cell(reference_xmap.unit_cell)
mask.spacegroup = gemmi.find_spacegroup_by_name("P 1")
for model in reference_structure:
for chain in model:
for residue in chain.get_polymer():
for atom in residue:
mask.set_points_around(atom.pos, 3.0, 1.0)
mask_array = np.array(mask)
mask_indicies = np.hstack([x.reshape((len(x), 1)) for x in np.nonzero(mask)])
print(f"Mask indicies shape: {mask_indicies.shape}")
fractional_coords = []
for model in reference_structure:
for chain in model:
for residue in chain.get_polymer():
for atom in residue:
fractional = reference_xmap.unit_cell.fractionalize(atom.pos)
fractional_coords.append([fractional.x, fractional.y, fractional.z])
fractional_coords_array = np.array(fractional_coords)
max_coord = np.max(fractional_coords_array, axis=0)
min_coord = np.min(fractional_coords_array, axis=0)
min_index = np.floor(min_coord * np.array([interpolated_grid.nu, interpolated_grid.nv, interpolated_grid.nw]))
max_index = np.floor(max_coord * np.array([interpolated_grid.nu, interpolated_grid.nv, interpolated_grid.nw]))
points = itertools.product(range(int(min_index[0]), int(max_index[0])),
range(int(min_index[1]), int(max_index[1])),
range(int(min_index[2]), int(max_index[2])),
)
# Unpack the points, poitions and transforms
point_list: List[Tuple[int, int, int]] = []
position_list: List[Tuple[float, float, float]] = []
transform_list: List[gemmi.transform] = []
com_moving_list: List[np.array] = []
com_reference_list: List[np.array] = []
transform_rotate_reference_to_moving = transform.transform
transform_rotate_reference_to_moving.vec.fromlist([0.0, 0.0, 0.0])
transform_reference_to_centered = gemmi.Transform()
transform_reference_to_centered.vec.fromlist((-transform.com_reference).tolist())
transform_reference_to_centered.mat.fromlist(np.eye(3).tolist())
tranform_centered_to_moving = gemmi.Transform()
tranform_centered_to_moving.vec.fromlist(transform.com_moving.tolist())
tranform_centered_to_moving.mat.fromlist(np.eye(3).tolist())
# indicies to positions
for point in points:
if mask.get_value(*point) < 1.0:
continue
# get position
position = interpolated_grid.get_position(*point)
# Tranform to origin frame
position_origin_reference = gemmi.Position(transform_reference_to_centered.apply(position))
# Rotate
position_origin_moving = gemmi.Position(transform_rotate_reference_to_moving.apply(position_origin_reference))
# Transform to moving frame
position_moving = gemmi.Position(tranform_centered_to_moving.apply(position_origin_moving))
# Interpolate moving map
interpolated_map_value = moving_xmap.interpolate_value(position_moving)
# Set original point
interpolated_grid.set_value(point[0], point[1], point[2], interpolated_map_value)
interpolated_grid.symmetrize_max()
return interpolated_grid
def __setstate__(self, state):
self.transform = gemmi.Transform()
self.transform.mat.fromlist(state["mat"])
self.transform.vec.fromlist(state["vec"])