Convert star file to imod mod for viz

star2mod.py

import starfile
import typer
from pathlib import Path
import subprocess
import os


star2mod = typer.Typer(add_completion=False)

@star2mod.command(no_args_is_help=True)

def main(star_file: Path = typer.Option(None, "--i", "--input-star", help="input star file"),
         tomo_bname: str = typer.Option(None, "--tn", "--tomo-name", help="Name of tomogram refernced in the star file"),
         tomo_path: Path = typer.Option(None, "--tp", "--tomo-path", help="Path to tomograms"),
         tomo_ang: float = typer.Option(None, "--tang","--tomo-angpix", help="tomogram angpix"),
         out_folder: Path = typer.Option(os.getcwd(), "--o", "--output-folder", help="Path to folder to save mod files in (folder must exist)"),
         combine: bool = typer.Option(True, "--combine/--dont-combine", help="Combine all mod files into one"),
         remove_tmp: bool = typer.Option(True, "--cleanup/--dont-cleanup", help="Remove text files?")):

    df = starfile.read(star_file)
    subtomo_ang = df['optics']['rlnImagePixelSize'].values[0]
    scaling = subtomo_ang/tomo_ang
    some_micrograph_name = f"{tomo_bname}.tomostar"
    num_classes = df['particles']['rlnClassNumber'].max()
    tomo_name = list(Path(tomo_path).glob(f'{tomo_bname}*.mrc'))[0]
    mod_files = []

    for i in range(1, num_classes + 1):
        coords = df['particles'][(df['particles']['rlnClassNumber'] == i) & (df['particles']['rlnMicrographName'] == some_micrograph_name)]
        coords_scaled = coords[['rlnCoordinateX','rlnCoordinateY','rlnCoordinateZ']].multiply(scaling)
        coords_scaled.to_csv(f'{out_folder}/{tomo_bname}_class{i:03}.txt', sep='\t', header=False, index=False)
        mod_files.append(f'{tomo_bname}_class{i:03}.mod')
        cmd = f"point2model -inp {out_folder}/{tomo_bname}_class{i:03}.txt -ou {out_folder}/{tomo_bname}_class{i:03}.mod -image {tomo_name} -scat -sphere 5"
        cmd_list = cmd.split()
        subprocess.run(cmd_list, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)

    if combine and len(mod_files) > 1:
        combined_mod = Path(out_folder) / f'{tomo_bname}_all_classes.mod'
        imodjoin_cmd = ["imodjoin","-c"] + [f'{out_folder}/{w}' for w in mod_files] + [str(combined_mod)]
        subprocess.run(imodjoin_cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, text=True, check=True)
        print(f"Combined mod file created: {combined_mod}")


    if remove_tmp:
        for i in range(1, num_classes + 1):
            tmp_file = f'{out_folder}/{tomo_bname}_class{i:03}.txt'
            if os.path.exists(tmp_file):
                os.remove(tmp_file)


if __name__ == '__main__':
    star2mod()