Created
September 18, 2024 22:08
-
-
Save alisterburt/10ea8f24fd14bbf400647bbc5b776155 to your computer and use it in GitHub Desktop.
R5 parsing for Barrett
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
| import ast | |
| import einops | |
| import numpy as np | |
| import starfile | |
| from scipy.spatial.transform import Rotation as R | |
| DO_VIS = True | |
| particle_star = starfile.read('matching.star') | |
| tomogram_star = starfile.read('matching_tomograms.star') | |
| particle_df = particle_star['particles'].merge(particle_star['optics'], on='rlnOpticsGroup') | |
| for key, particle_group_df in particle_df.groupby('rlnTomoName'): | |
| # get (n_particles, 3, 3) particle rotation_matrices | |
| particle_euler_angles = particle_group_df[['rlnAngleRot', 'rlnAngleTilt', 'rlnAnglePsi']].to_numpy() | |
| particle_rotation_matrices = R.from_euler( | |
| seq='ZYZ', angles=particle_euler_angles, degrees=True | |
| ).inv().as_matrix() | |
| # get table for specific table | |
| tomogram_df = tomogram_star[key] | |
| # get (n_tilts, 4) arrays for first three rows of projection matrices | |
| projection_matrices_x = np.stack(tomogram_df['rlnTomoProjX'].apply(lambda x: ast.literal_eval(x))) | |
| projection_matrices_y = np.stack(tomogram_df['rlnTomoProjY'].apply(lambda x: ast.literal_eval(x))) | |
| projection_matrices_z = np.stack(tomogram_df['rlnTomoProjZ'].apply(lambda x: ast.literal_eval(x))) | |
| # get (n_tilts, 3, 4) projection_matrices | |
| tilt_projection_matrices = einops.rearrange( | |
| [projection_matrices_x, projection_matrices_y, projection_matrices_z], | |
| pattern='projxyz b d -> b projxyz d' | |
| ) | |
| # get (n_tilts, 3, 3) rotation matrices | |
| tilt_rotation_matrices = tilt_projection_matrices[:, :, :-1] | |
| # get (n_particles, n_tilts, 3, 3) per particle per tilt rotation matrices | |
| particle_rotation_matrices = einops.rearrange(particle_rotation_matrices, 'b i j -> b 1 i j') | |
| particle_tilt_rotation_matrices = particle_rotation_matrices @ tilt_rotation_matrices # (n_particles, n_tilts, 3, 3) | |
| # get (n_particles * ntilts, 3) euler angles in RELION format | |
| particle_eulers = R.from_matrix( | |
| einops.rearrange(particle_tilt_rotation_matrices, 'nparticles ntilts i j -> (nparticles ntilts) i j') | |
| ).inv().as_euler(seq='ZYZ', degrees=True) | |
| # view first particle fourier planes rotated and unrotated | |
| if DO_VIS: | |
| unrotated = tilt_rotation_matrices | |
| rotated = particle_tilt_rotation_matrices[0] | |
| xy_plane = np.meshgrid( | |
| np.linspace(-1, 1, ), | |
| np.linspace(-1, 1, ) | |
| ) | |
| xy_plane = einops.rearrange(xy_plane, 'd h w -> (h w) d 1') # (b, 2, 1) | |
| xy_plane = np.pad(xy_plane, ((0, 0), (0, 1), (0, 0)), mode='constant', constant_values=0) # (b, 3, 1) | |
| unrotated = einops.rearrange(unrotated, 'ntilt i j -> ntilt 1 i j') | |
| unrotated_planes = unrotated @ xy_plane | |
| unrotated_planes = einops.rearrange(unrotated_planes, 'ntilt npoints d 1 -> (ntilt npoints) d') | |
| rotated = einops.rearrange(rotated, 'ntilt i j -> ntilt 1 i j') | |
| rotated_planes = rotated @ xy_plane | |
| rotated_planes = einops.rearrange(rotated_planes, 'ntilt npoints d 1 -> (ntilt npoints) d') | |
| import napari | |
| viewer = napari.Viewer(ndisplay=3) | |
| viewer.axes.visible = True | |
| viewer.add_points(unrotated_planes, size=0.01, face_color='grey') | |
| viewer.add_points(rotated_planes, size=0.01, face_color='purple') | |
| napari.run() | |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment