recentering/reorienting RELION particles for Josh Hutchings

josh.py

"""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)