Last active
March 13, 2024 01:05
-
-
Save alisterburt/5fa3e867f8004171fff57579b379ffa9 to your computer and use it in GitHub Desktop.
recentering/reorienting RELION particles for Josh Hutchings
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
"""Recenter and reorient RELION particles | |
This shows the gist of how to manually move/reorient RELION particles in 3D. | |
You define a set of shifts and orientations in the coordinate system of the | |
reference for all 'subparticles' you want to generate from each existing particle. | |
""" | |
import pandas as pd | |
import starfile | |
import numpy as np | |
from scipy.spatial.transform import Rotation as R | |
import einops | |
star = starfile.read('jh_star.star') | |
# get current poses (positions and orientations) | |
apix_image = float(star['optics']['rlnImagePixelSize'][0]) | |
particle_positions_xyz = star['particles'][['rlnCoordinateX', 'rlnCoordinateY', 'rlnCoordinateZ']].to_numpy() | |
shifts_xyz_ang = star['particles'][['rlnOriginXAngst', 'rlnOriginYAngst', 'rlnOriginZAngst']].to_numpy() | |
shifts_xyz_px = shifts_xyz_ang / apix_image | |
particle_positions_xyz -= shifts_xyz_px # weird RELION convention | |
eulers = star['particles'][['rlnAngleRot', 'rlnAngleTilt', 'rlnAnglePsi']].to_numpy() | |
particle_orientations = R.from_euler(seq='ZYZ', angles=eulers, | |
degrees=True).inv().as_matrix() | |
# define subparticle positions/orientations relative to reference | |
subparticle_xyz = np.array( | |
[[0, 20, 0], # shift -20 in y | |
[0, 0, 20]], # shift 20 in z | |
) | |
subparticle_orientations = R.random(num=2).as_matrix() | |
subparticle_x = subparticle_orientations[:, :, 0] | |
subparticle_y = subparticle_orientations[:, :, 1] | |
subparticle_z = subparticle_orientations[:, :, 2] | |
# calculate new particle positions | |
# orient the shifts and add them | |
oriented_shifts_xyz = particle_orientations @ subparticle_xyz.reshape((-1, 1, 3, 1)) | |
oriented_shifts_xyz = oriented_shifts_xyz.squeeze() # (2, 88, 3, 1) -> (2, 88, 3) | |
new_particle_positions_xyz = particle_positions_xyz + oriented_shifts_xyz | |
# reorient the particles | |
new_orientations = particle_orientations @ subparticle_orientations.reshape((-1, 1, 3, 3)) | |
new_orientations = einops.rearrange(new_orientations, 'sp p i j -> (sp p) i j') | |
new_eulers = R.from_matrix(new_orientations).inv().as_euler(seq='ZYZ', degrees=True) | |
new_eulers = einops.rearrange( | |
new_eulers, '(sp p) i -> sp p i', | |
sp=len(subparticle_orientations), # number of subparticles per particle (can be 1) | |
p=len(particle_orientations) # number of particles | |
) | |
# express new positions relative to existing coordinates in star file | |
new_shifts_xyz_px = new_particle_positions_xyz - particle_positions_xyz | |
new_shifts_xyz_ang = new_shifts_xyz_px * apix_image | |
# write out star file | |
# put a list of the values for the shift/orientation of each subparticle in every row | |
star['particles']['rlnOriginXAngst'] = list(einops.rearrange( | |
new_shifts_xyz_ang[..., 0], 'sp p -> p sp' | |
)) | |
star['particles']['rlnOriginYAngst'] = list(einops.rearrange( | |
new_shifts_xyz_ang[..., 1], 'sp p -> p sp' | |
)) | |
star['particles']['rlnOriginZAngst'] = list(einops.rearrange( | |
new_shifts_xyz_ang[..., 2], 'sp p -> p sp' | |
)) | |
star['particles']['rlnAngleRot'] = list(einops.rearrange( | |
new_eulers[..., 0], 'sp p -> p sp' | |
)) | |
star['particles']['rlnAngleTilt'] = list(einops.rearrange( | |
new_eulers[..., 1], 'sp p -> p sp' | |
)) | |
star['particles']['rlnAnglePsi'] = list(einops.rearrange( | |
new_eulers[..., 2], 'sp p -> p sp' | |
)) | |
# make a new row for every subparticle, maintaining data for the rest of the columns | |
star['particles'] = star['particles'].explode( | |
column=[ | |
'rlnOriginXAngst', | |
'rlnOriginYAngst', | |
'rlnOriginZAngst', | |
'rlnAngleRot', | |
'rlnAngleTilt', | |
'rlnAnglePsi', | |
] | |
) | |
star['particles'] = star['particles'].apply(pd.to_numeric, errors='ignore') | |
starfile.write(star, 'output.star', overwrite=True) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment