make choice node configurable

2 files changed, 15 insertions, 6 deletions

diff --git a/bin/eval-kconfig.py b/bin/eval-kconfig.py
index 1f44b9e..7f48b52 100755
--- a/bin/eval-kconfig.py
+++ b/bin/eval-kconfig.py
@@ -10,6 +10,8 @@ import logging
 import os
 import sys
 
+import numpy as np
+
 from dfatool import kconfig, validation
 from dfatool.loader import KConfigAttributes
 from dfatool.model import KConfigModel
@@ -32,6 +34,9 @@ def main():
     parser.add_argument(
         "--attribute", choices=["rom", "ram"], default="rom", help="Model attribute"
     )
+    parser.add_argument(
+        "--with-choice-node", action="store_true", help="Use non-binary Choice Nodes"
+    )
     parser.add_argument("kconfig_path", type=str, help="Path to Kconfig file")
     parser.add_argument(
         "experiment_root", type=str, help="Experiment results directory"
@@ -53,14 +58,15 @@ def main():
     measures = list()
     for training_set, validation_set in partition_pairs:
         model = KConfigModel.from_benchmark(data, args.attribute, indices=training_set)
+        model.with_choice_node = args.with_choice_node
         model.build_tree()
         measures.append(model.assess_benchmark(data, indices=validation_set))
 
     aggregate = dict()
     for measure in measures[0].keys():
-        aggregate[measure] = np.mean(map(lambda m: m[measure], measures))
+        aggregate[measure] = np.mean(list(map(lambda m: m[measure], measures)))
     aggregate["unpredictable_count"] = np.sum(
-        map(lambda m: m["unpredictable_count"], measures)
+        list(map(lambda m: m["unpredictable_count"], measures))
     )
 
     print("10-fold Cross Validation:")
diff --git a/lib/model.py b/lib/model.py
index f422204..a75033c 100644
--- a/lib/model.py
+++ b/lib/model.py
@@ -1317,6 +1317,7 @@ class KConfigModel:
         self.choices = kconfig_benchmark.choice_names
         self.symbol = kconfig_benchmark.symbol
         self.choice = kconfig_benchmark.choice
+        self.with_choice_node = True
         self.max_loss = 10
         if callable(attribute):
             self.attribute = "custom"
@@ -1356,9 +1357,6 @@ class KConfigModel:
         return np.sum((model_value - values) ** 2, dtype=np.float64)
 
     def build_tree(self):
-        # without ChoiceNode:
-        # self.model = self._build_tree(self.symbols, list(), self.data, 0)
-
         standalone_symbols = list(
             filter(
                 lambda sym: self.symbol[sym].choice is None
@@ -1369,7 +1367,12 @@ class KConfigModel:
         tree_choices = list(
             filter(lambda choice: not self.choice[choice].is_optional, self.choices)
         )
-        self.model = self._build_tree(standalone_symbols, tree_choices, self.data, 0)
+        if self.with_choice_node:
+            self.model = self._build_tree(
+                standalone_symbols, tree_choices, self.data, 0
+            )
+        else:
+            self.model = self._build_tree(self.symbols, list(), self.data, 0)
 
     def value_for_config(self, kconf):
         return self.model.model(kconf)