Add timings with --verbose.

clusttraj/io.py

@@ -55,6 +55,7 @@ class ClustOptions:
     solute_natoms: int = None
     reorder_excl: np.ndarray = None
     optimal_cut: np.ndarray = None
+    verbose: bool = None
 
     def update(self, new: dict) -> None:
         """Update the instance with new values.
@@ -305,12 +306,17 @@ def configure_runtime(args_in: List[str]) -> ClustOptions:
         default="clusttraj.log",
         help="log file (default: clusttraj.log)",
     )
-
     parser.add_argument(
         "--metrics",
         action="store_true",
         help="compute metrics to evaluate the clustering procedure quality.",
     )
+    parser.add_argument(
+        "-v",
+        "--verbose",
+        action="store_true",
+        help="increase verbosity, printing the timings for each part of the program.",
+    )
 
     rmsd_criterion = parser.add_mutually_exclusive_group(required=True)
 
@@ -570,6 +576,7 @@ def parse_args(args: argparse.Namespace) -> ClustOptions:
         "final_kabsch": args.final_kabsch,
         "silhouette_score": args.silhouette_score,
         "metrics": args.metrics,
+        "verbose": args.verbose,
     }
 
     if args.reorder:

clusttraj/main.py

@@ -7,6 +7,7 @@
 import sys
 import numpy as np
 from typing import List
+import time
 from .io import Logger, configure_runtime, save_clusters_config
 from .distmat import get_distmat
 from .plot import plot_clust_evo, plot_dendrogram, plot_mds, plot_tsne
@@ -23,23 +24,37 @@ def main(args: List[str] = None) -> None:
     Returns:
         None
     """
+    global_start_time = time.monotonic()
+
+    # parse command-line arguments
     if args is None:
         args = sys.argv[1:]
 
     # get ClustOptions class with parsed arguments
     clust_opt = configure_runtime(args)
 
     # get the distance matrix
+    start_time = time.monotonic()
     distmat = get_distmat(clust_opt)
+    end_time = time.monotonic()
+    if clust_opt.verbose:
+        Logger.logger.info(
+            f"Time spent computing (or loading) distance matrix: {end_time - start_time:.6f} s\n"
+        )
 
     # perform the clustering
+    start_time = time.monotonic()
     if clust_opt.silhouette_score:
         Z, clusters = classify_structures_silhouette(clust_opt, distmat)
     else:
         Z, clusters = classify_structures(clust_opt, distmat)
+    end_time = time.monotonic()
+    if clust_opt.verbose:
+        Logger.logger.info(f"Time spent clustering: {end_time - start_time:.6f} s\n")
 
     # get the elements closest to the centroid (see https://stackoverflow.com/a/39870085/3254658)
     if clust_opt.save_confs:
+        start_time = time.monotonic()
         outconf = clust_opt.out_conf_name + "_*." + clust_opt.out_conf_fmt
         Logger.logger.info(
             f"Writing superposed configurations per cluster to files {outconf}\n"
@@ -57,16 +72,25 @@ def main(args: List[str] = None) -> None:
             clust_opt.final_kabsch,
             clust_opt.overwrite,
         )
+        end_time = time.monotonic()
+        if clust_opt.verbose:
+            Logger.logger.info(
+                f"Time spent saving configurations: {end_time - start_time:.6f} s\n"
+            )
 
     # generate plots
     if clust_opt.plot:
+        start_time = time.monotonic()
         plot_clust_evo(clust_opt, clusters)
 
         plot_dendrogram(clust_opt, Z)
 
         plot_mds(clust_opt, clusters, distmat)
 
         plot_tsne(clust_opt, clusters, distmat)
+        end_time = time.monotonic()
+        if clust_opt.verbose:
+            Logger.logger.info(f"Time spent plotting: {end_time - start_time:.6f} s\n")
 
     # print the cluster sizes
     outclust_str = f"A total {len(clusters)} snapshots were read and {max(clusters)} cluster(s) was(were) found.\n"
@@ -79,18 +103,28 @@ def main(args: List[str] = None) -> None:
 
     # Compute the evaluation metrics
     if clust_opt.metrics:
+        start_time = time.monotonic()
         ss, ch, db, cpcc = compute_metrics(clust_opt, distmat, Z, clusters)
+        end_time = time.monotonic()
 
         outclust_str += f"\nSilhouette score: {ss:.3f}\n"
         outclust_str += f"Calinski Harabsz score: {ch:.3f}\n"
         outclust_str += f"Davies-Bouldin score: {db:.3f}\n"
         outclust_str += f"Cophenetic correlation coefficient: {cpcc:.3f}\n\n"
 
+        if clust_opt.verbose:
+            Logger.logger.info(
+                f"Time spent computing metrics: {end_time - start_time:.6f} s\n"
+            )
+
     # save summary
     with open(clust_opt.summary_name, "w") as f:
         f.write(str(clust_opt))
         f.write(outclust_str)
 
+    global_end_time = time.monotonic()
+    Logger.logger.info(f"Total wall time: {global_end_time - global_start_time:.6f} s\n")
+
 
 if __name__ == "__main__":
     main()

test/test_io.py

@@ -83,6 +83,7 @@ def test_parse_args():
         final_kabsch=True,
         silhouette_score=False,
         metrics=False,
+        verbose=False,
     )
     clust_opt = parse_args(args)
 
@@ -109,6 +110,7 @@ def test_parse_args():
         final_kabsch=True,
         silhouette_score=False,
         metrics=False,
+        verbose=False,
     )
     clust_opt = parse_args(args)