batch reconstruction of tomograms in RELION 4

batch_reconstruct.py

import subprocess
from pathlib import Path

import starfile
import typer
from rich.console import Console
from rich.progress import track

prompt_kwargs = {'prompt': True, 'prompt_required': True}


def batch_reconstruct_tomo(
        optimisation_set: Path = typer.Option(..., **prompt_kwargs),
        output_directory: Path = typer.Option(..., **prompt_kwargs),
        binning: int = typer.Option(8, **prompt_kwargs),
        threads: int = typer.Option(8, **prompt_kwargs),
):
    console = Console(record=True)
    Path(output_directory).mkdir(parents=True, exist_ok=True)

    # get tomogram names
    star = starfile.read(optimisation_set, always_dict=True)
    df = star[0]
    tomogram_set = df['rlnTomoTomogramsFile'][0]
    console.log(f'extracted tomogram set {tomogram_set} from {optimisation_set}')

    tomogram_set_star = starfile.read(tomogram_set, always_dict=True)
    df = tomogram_set_star['global']
    tomogram_names = df['rlnTomoName']
    console.log(f'reconstructing {len(tomogram_names)} tomograms...')

    # loop over and reconstruct
    for tomogram in track(tomogram_names, description='reconstructing tomograms'):
        output_file = Path(output_directory) / f'{tomogram}_bin{binning}.mrc'
        command = [
            'relion_tomo_reconstruct_tomogram',
            '--i', f'{optimisation_set}',
            '--tn', f'{tomogram}',
            '--o', f'{output_file}',
            '--j', f'{threads}',
            '--bin', f'{binning}',
        ]
        completed_process = subprocess.run(
            command,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.STD_ERROR_HANDLE
        )
        if completed_process.returncode == 0:
            console.log(f'reconstructed {tomogram}')
        else:
            console.log(f'[bold red]failed to reconstruct {tomogram}')


if __name__ == '__main__':
    typer.run(batch_reconstruct_tomo)