Better cli for analysis

benchmarks/scripts/analyze.py

@@ -2,7 +2,7 @@
 
 import re
 import cub
-import sys
+import argparse
 import itertools
 import functools
 import pandas as pd
@@ -144,16 +144,7 @@ def extract_scores(df):
     return result.sort_values(by=['score'], ascending=False)
 
 
-def main():
-    if not cub.bench.Storage().exists():
-        raise Exception("Storage does not exist")
-
-    args = cub.bench.parse_arguments()
-
-    if args.list_benches:
-        cub.bench.list_benches()
-        return
-
+def iterate_case_dfs(args, callable):
     storage = cub.bench.Storage()
     pattern = re.compile(args.R)
 
@@ -162,8 +153,6 @@ def main():
             continue
 
         df = storage.alg_to_df(algname)
-        num_variants = cub.bench.Config().variant_space_size(algname)
-
         for _, row in df[['ctk', 'cub']].drop_duplicates().iterrows():
             ctk_version = row['ctk']
             cub_version = row['cub']
@@ -176,14 +165,62 @@ def main():
                 target_df = compute_speedup(target_df)
 
                 for ct_point in ct_space(target_df):
-                    point_str = ", ".join(
-                        ["{}={}".format(k, ct_point[k]) for k in ct_point])
+                    point_str = ", ".join(["{}={}".format(k, ct_point[k]) for k in ct_point])
                     case_df = extract_complete_variants(extract_case(target_df, ct_point))
-                    num_covered_variants = len(case_df['variant'].unique())
-                    coverage = (num_covered_variants / num_variants) * 100
-                    print("{}[{}] coverage: {} / {} ({:.4f}%)".format(
-                        algname, point_str, num_covered_variants, num_variants, coverage))
-                    print(extract_scores(case_df).head(7))
+                    callable(algname, point_str, case_df)
+
+
+def case_top(N, algname, ct_point_name, case_df):
+    print("{}[{}]:".format(algname, ct_point_name))
+    print(extract_scores(case_df).head(N))
+
+
+def top(args):
+    iterate_case_dfs(args, functools.partial(case_top, args.top))
+
+
+def case_coverage(algname, ct_point_name, case_df):
+    num_variants = cub.bench.Config().variant_space_size(algname)
+    num_covered_variants = len(case_df['variant'].unique())
+    coverage = (num_covered_variants / num_variants) * 100
+    case_str = "{}[{}]".format(algname, ct_point_name)
+    print("{} coverage: {} / {} ({:.4f}%)".format(
+          case_str, num_covered_variants, num_variants, coverage))
+
+
+def coverage(args):
+    iterate_case_dfs(args, case_coverage)
+
+
+def parse_arguments():
+    parser = argparse.ArgumentParser(
+        description="Analyze benchmark results.")
+    parser.add_argument('-R', type=str, default='.*',
+                        help="Regex for benchmarks selection.")
+    parser.add_argument(
+        '--list-benches', action=argparse.BooleanOptionalAction, help="Show available benchmarks.")
+    parser.add_argument(
+        '--coverage', action=argparse.BooleanOptionalAction, help="Show variant space coverage.")
+    parser.add_argument(
+        '--top', default=7, type=int, action='store', nargs='?', help="Show top N variants with highest score.")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_arguments()
+
+    if not cub.bench.Storage().exists():
+        raise Exception("Storage does not exist")
+
+    if args.list_benches:
+        cub.bench.list_benches()
+        return
+
+    if args.coverage:
+        coverage(args)
+        return
+
+    top(args)
 
 
 if __name__ == "__main__":