configurable maximum stddev for kconfig model

2 files changed, 39 insertions, 8 deletions

diff --git a/bin/analyze-kconfig.py b/bin/analyze-kconfig.py
index 2f2e973..f532a5b 100755
--- a/bin/analyze-kconfig.py
+++ b/bin/analyze-kconfig.py
@@ -21,25 +21,53 @@ def main():
         formatter_class=argparse.RawDescriptionHelpFormatter, description=__doc__
     )
     parser.add_argument(
-        "--export-tree", type=str, help="Export decision tree model to file"
+        "--export-tree",
+        type=str,
+        help="Export decision tree model to file",
+        metavar="FILE",
     )
     parser.add_argument(
-        "--config", type=str, help="Show model results for symbols in .config file"
+        "--config",
+        type=str,
+        help="Show model results for symbols in .config file",
+        metavar="FILE",
     )
     parser.add_argument(
         "--attribute", choices=["rom", "ram"], default="rom", help="Model attribute"
     )
+    parser.add_argument(
+        "--max-stddev",
+        type=float,
+        help="Maximum acceptable model standard deviation for DecisionTree Leaves",
+        default=10,
+    )
+    parser.add_argument(
+        "--log-level",
+        default=logging.INFO,
+        type=lambda level: getattr(logging, level.upper()),
+        help="Set log level",
+    )
     parser.add_argument("kconfig_path", type=str, help="Path to Kconfig file")
     parser.add_argument("experiment_root", type=str, help="Path to experiment results")
 
     args = parser.parse_args()
 
+    if isinstance(args.log_level, int):
+        logging.basicConfig(level=args.log_level)
+    else:
+        print(f"Invalid log level. Setting log level to INFO.", file=sys.stderr)
+
     data = KConfigAttributes(args.kconfig_path, args.experiment_root)
     model = KConfigModel(data, args.attribute)
 
+    if args.max_stddev:
+        model.max_stddev = args.max_stddev
+
+    model.build_tree()
+
     if args.export_tree:
         with open(args.export_tree, "w") as f:
-            json.dump(model.get_tree(), f)
+            json.dump(model.to_json(), f)
 
     if args.config:
         kconf = kconfiglib.Kconfig(args.kconfig_path)
diff --git a/lib/model.py b/lib/model.py
index 118a755..c74ebd1 100644
--- a/lib/model.py
+++ b/lib/model.py
@@ -1211,6 +1211,7 @@ class KConfigModel:
     def __init__(self, kconfig_benchmark, attribute):
         self.data = kconfig_benchmark.data
         self.symbols = kconfig_benchmark.symbols
+        self.max_stddev = 10
         if callable(attribute):
             self.attribute = "custom"
             self.attr_function = lambda x: attribute(x[1])
@@ -1222,7 +1223,9 @@ class KConfigModel:
             self.attr_function = lambda x: x[1]["total"]["RAM"]
         else:
             raise ValueError("attribute must be a a function, 'rom', or 'ram'")
-        self.model = self.build_tree(self.symbols, self.data, 0)
+
+    def build_tree(self):
+        self.model = self._build_tree(self.symbols, self.data, 0)
 
     def value_for_config(self, kconf):
         return self.model.model(kconf)
@@ -1235,11 +1238,11 @@ class KConfigModel:
         }
         return output
 
-    def build_tree(self, this_symbols, this_data, level):
+    def _build_tree(self, this_symbols, this_data, level):
 
         rom_sizes = list(map(self.attr_function, this_data))
 
-        if np.std(rom_sizes) < 100 or len(this_symbols) == 0:
+        if np.std(rom_sizes) < self.max_stddev or len(this_symbols) == 0:
             return self.Leaf(np.mean(rom_sizes), np.std(rom_sizes))
 
         mean_stds = list()
@@ -1280,8 +1283,8 @@ class KConfigModel:
             f"Level {level} split on {symbol} has {len(enabled)} children when enabled and {len(disabled)} children when disabled"
         )
         if len(enabled):
-            node.set_child_y(self.build_tree(new_symbols, enabled, level + 1))
+            node.set_child_y(self._build_tree(new_symbols, enabled, level + 1))
         if len(disabled):
-            node.set_child_n(self.build_tree(new_symbols, disabled, level + 1))
+            node.set_child_n(self._build_tree(new_symbols, disabled, level + 1))
 
         return node