numpy version

jrzkaminski · Nov 1, 2023 · 910e6f7 · 910e6f7
1 parent 6c9bd27
commit 910e6f7
Showing 1 changed file with 4 additions and 15 deletions.
diff --git a/bamt/nodes/discrete_node.py b/bamt/nodes/discrete_node.py
@@ -74,24 +74,13 @@ def choose(self, node_info: Dict[str, Union[float, str]], pvals: List[str]) -> s
         node_info: nodes info from distributions
         pvals: parent values
         """
-        # NUMPY VERSION DO NOT DELETE
-        # vals = node_info["vals"]
-        # dist = np.array(self.get_dist(node_info, pvals))
-        #
-        # cumulative_dist = np.cumsum(dist)
-        #
-        # rand = np.random.random()
-        # rindex = np.searchsorted(cumulative_dist, rand)
-        #
-        # return vals[rindex]
-
         vals = node_info["vals"]
-        dist = self.get_dist(node_info, pvals)
+        dist = np.array(self.get_dist(node_info, pvals))
 
-        cumulative_dist = list(accumulate(dist))
+        cumulative_dist = np.cumsum(dist)
 
-        rand = random.random()
-        rindex = next((i for i, ubound in enumerate(cumulative_dist) if rand < ubound), len(vals) - 1)
+        rand = np.random.random()
+        rindex = np.searchsorted(cumulative_dist, rand)
 
         return vals[rindex]