From d75e2458ed1b85e495f0c4b7d15fab22768a2cdf Mon Sep 17 00:00:00 2001
From: Forest Belton <forest@homolo.gy>
Date: Tue, 3 Aug 2021 12:15:27 -0400
Subject: [PATCH] Add assertions on probability distribution

---
 gbso/optimize.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/gbso/optimize.py b/gbso/optimize.py
index de60d82..52485e2 100644
--- a/gbso/optimize.py
+++ b/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