Add assertions on probability distribution

gbso/optimize.py

@@ -1,8 +1,7 @@
-from random import random
-from typing import List, Tuple, TypeVar
-
 from gbso.program.program import Program
 
+EPSILON = 0.00001
+
 DEFAULT_NUM_ITERS = 1_000_000
 
 DEFAULT_PROB_OPCODE = 0.25
@@ -22,4 +21,7 @@ def optimize(
     prob_insn: float = DEFAULT_PROB_INSN,
     prob_insn_unused: float = DEFAULT_PROB_INSN_UNUSED,
 ) -> Program:
+    prob_sum = sum([prob_opcode, prob_operand, prob_swap, prob_insn])
+    assert abs(1 - prob_sum) < EPSILON
+
     return prgm