select classes from relion STAR files

relion_class_select.py

#!/usr/bin/env python
import click
import starfile

@click.command()
@click.option('--star_file', prompt='Input STAR file')
@click.option('--class_idx', type=int, prompt='Class number')
def class_select(star_file, class_idx):
	star = starfile.read(star_file)
	subset = star[star['rlnClassNumber']==class_idx]
	out_fname = star_file.replace(f'.star', f'_class{class_idx}.star')
	starfile.write(subset, out_fname)
	click.echo(f'Before: {star.shape[0]} particles')
	click.echo(f'After: {subset.shape[0]} particles')
	return
	
if __name__ == '__main__':
	class_select()