creating utilities command from utils/analyze_convergence

cryodrgn/command_line.py

@@ -112,6 +112,7 @@ def util_commands() -> None:
     _get_commands(
         cmd_dir=os.path.join(os.path.dirname(__file__), "commands_utils"),
         cmds=[
+            "analyze_convergence",
             "add_psize",
             "clean",
             "concat_pkls",

utils/analyze_convergence.py → cryodrgn/commands_utils/analyze_convergence.py

@@ -1,9 +1,12 @@
 """
 Visualize convergence and training dynamics
-
 (BETA -- contributed by Barrett Powell bmp@mit.edu)
-"""
 
+Example usage
+-------------
+$ cryodrgn_utils analyze_convergence 003_train-vae_out/ 25
+
+"""
 import argparse
 import itertools
 import multiprocessing
@@ -34,7 +37,7 @@
 logger = logging.getLogger(__name__)
 
 
-def add_args(parser):
+def add_args(parser: argparse.ArgumentParser) -> None:
     parser.add_argument(
         "workdir", type=os.path.abspath, help="Directory with cryoDRGN results"
     )
@@ -184,8 +187,6 @@ def add_args(parser):
         help="Number of voxels over which to apply a soft cosine falling edge from dilated mask boundary",
     )
 
-    return parser
-
 
 def plot_loss(logfile, outdir, E):
     """
@@ -1039,7 +1040,7 @@ def calc_fsc(vol1_path: str, vol2_path: str):
     )
 
 
-def main(args):
+def main(args: argparse.Namespace) -> None:
     t1 = dt.now()
 
     # Configure paths
@@ -1238,13 +1239,3 @@ def main(args):
     calculate_FSCs(outdir, epochs, labels, img_size, chimerax_colors)
 
     logger.info(f"Finished in {dt.now() - t1}")
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser(
-        description=__doc__,
-        epilog="Example usage: $ python analyze_convergence.py [workdir] [epoch]",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-    add_args(parser)
-    main(parser.parse_args())