Merge 4229f95 into ef1f72a

renatahodovan · Nov 6, 2022 · 5d497df · 5d497df
2 parents ef1f72a + 4229f95
commit 5d497df
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/grammarinator/runtime/cooldown_model.py b/grammarinator/runtime/cooldown_model.py
@@ -15,11 +15,18 @@ def __init__(self, model, weights=None, cooldown=1.0):
 
     def choice(self, node, idx, choices):
         i = self._model.choice(node, idx, [w * self._weights.get((node.name, idx, i), 1) for i, w in enumerate(choices)])
-        self._weights[(node.name, idx, i)] = self._weights.get((node.name, idx, i), 1) * self._cooldown
+        weights = CooldownModel.normalize([self._weights.get((node.name, idx, i), 1) * self._cooldown if iw == i else w for iw, w in enumerate(choices)])
+        for iw, w in enumerate(weights):
+            self._weights[(node.name, idx, iw)] = w
         return i
 
     def quantify(self, node, idx, min, max):
         yield from self._model.quantify(node, idx, min, max)
 
     def charset(self, node, idx, chars):
         return self._model.charset(node, idx, chars)
+
+    @staticmethod
+    def normalize(values):
+        vsum = sum(values)
+        return [float(value) / vsum for value in values]