diff --git a/cryodrgn/commands_utils/filter_star.py b/cryodrgn/commands_utils/filter_star.py index e0788ce7..7415b0f5 100644 --- a/cryodrgn/commands_utils/filter_star.py +++ b/cryodrgn/commands_utils/filter_star.py @@ -25,6 +25,7 @@ def add_args(parser): parser.add_argument( "--et", action="store_true", help="Set if .star file includes tilts" ) + parser.add_argument("--tomo", action="store_true", help="Split output by micrograph name into separate .star files") parser.add_argument("-o", type=os.path.abspath, help="Output .star file") @@ -41,6 +42,16 @@ def main(args): tilt_indices = TiltSeriesData.particles_to_tilts(particles_to_tilts, ind) df = s.df.loc[tilt_indices] + if args.tomo: + df = s.df.loc[ind] + os.makedirs(args.o, exist_ok=True) # Ensure the output directory exists + for micrograph_name, group_df in df.groupby('_rlnMicrographName'): + filename_without_extension = os.path.splitext(micrograph_name)[0] # Remove extension + output_path = os.path.join(args.o, f"{filename_without_extension}.star") + new_star = starfile.Starfile(headers=None, df=group_df) + new_star.write(output_path) + logger.info(f"Written .star file for {filename_without_extension} to {output_path}") + exit() else: df = s.df.loc[ind] @@ -52,3 +63,4 @@ def main(args): parser = argparse.ArgumentParser(description=__doc__) add_args(parser) main(parser.parse_args()) +