convert a set of IMOD models to a RELION 3 star file for warp (for matteo)

imod_model_to_rln3_star_file.py

from pathlib import Path
from typing import Tuple

import eulerangles
import imodmodel
import numpy as np
import pandas as pd
import starfile as starfile
from pydantic import BaseModel


class Dipole(BaseModel):
    center: Tuple[float, float, float]
    north: Tuple[float, float, float]


def relion_euler_angles_from_dipole(dipole: Dipole) -> Tuple[float, float, float]:
    arbitrary_vec = np.random.random(3)
    z = np.array(dipole.north) - np.array(dipole.center)
    z /= np.linalg.norm(z)
    y = np.cross(z, arbitrary_vec)
    y /= np.linalg.norm(y)
    x = np.cross(y, z)
    x /= np.linalg.norm(x)
    rotation_matrix = np.zeros((3, 3))
    rotation_matrix[:, 0] = x
    rotation_matrix[:, 1] = y
    rotation_matrix[:, 2] = z
    rotation_matrix = np.linalg.pinv(rotation_matrix)  # relion expects ref2particle
    euler_angles = eulerangles.matrix2euler(
        rotation_matrices=rotation_matrix,
        axes='zyz',
        intrinsic=True,
        right_handed_rotation=True,
    )
    return np.squeeze(euler_angles)


def mod_file_to_rln3_df(model_file) -> pd.DataFrame:
    df = imodmodel.read(model_file)

    dipoles = []
    for contour_idx, group in df.groupby('contour_idx'):
        assert (len(group) == 2)
        c = group.iloc[0, :][['x', 'y', 'z']]
        n = group.iloc[1, :][['x', 'y', 'z']]
        dipoles.append(Dipole(center=tuple(c), north=tuple(n)))

    xyz = np.array([(d.center[0], d.center[1], d.center[2]) for d in dipoles])
    euler_angles = np.array([relion_euler_angles_from_dipole(d) for d in dipoles])
    df = pd.DataFrame(
        {
            'rlnCoordinateX': xyz[:, 0],
            'rlnCoordinateY': xyz[:, 1],
            'rlnCoordinateZ': xyz[:, 2],
            'rlnAngleRot': euler_angles[:, 0],
            'rlnAngleTilt': euler_angles[:, 1],
            'rlnAnglePsi': euler_angles[:, 2],
        }
    )
    return df


def mod2tomogram(mod_file_name: str) -> str:
    return f'TS_{mod_file_name}.tif.tif.mrc_17.41Apx.mrc'


if __name__ == '__main__':
    mod_files = list(Path().glob('???.mod'))
    dfs = []
    for mod_file in mod_files:
        df = mod_file_to_rln3_df(mod_file)
        df['rlnMicrographName'] = mod2tomogram(mod_file.stem)
        dfs.append(dfs)
    df = pd.concat(dfs)
    starfile.write(df, 'particles.star')