Add bandwidth simulations

dsg-uwaterloo · Jun 16, 2024 · 85d022d · 85d022d
1 parent 4580d16
commit 85d022d
Show file tree

Hide file tree

Showing 2 changed files with 102 additions and 0 deletions.
diff --git a/simulation/simulate.py b/simulation/simulate.py
@@ -0,0 +1,50 @@
+import random
+
+class Tree:
+    def __init__(self, tree_path_count):
+        self.tree_path_count = tree_path_count
+
+    def get_random_paths(self, batch_size):
+        if batch_size > self.tree_path_count:
+            raise ValueError("Batch size cannot be greater than the number of paths in the tree")
+        return [random.randint(1, self.tree_path_count) for _ in range(batch_size)]
+
+    def get_unique_buckets_for_random_paths(self, random_paths):
+        unique_buckets = set()
+        for path_id in random_paths:
+            leaf_id = self.tree_path_count + path_id - 1
+            bucket_id = leaf_id
+            while bucket_id > 0:
+                unique_buckets.add(bucket_id)
+                bucket_id = bucket_id >> 1
+        return unique_buckets
+
+EXPERIMETNS = 1000
+
+class Simulator:
+    def __init__(self, num_requests, batch_size, tree_path_count):
+        self.num_requests = num_requests
+        self.tree = Tree(tree_path_count)
+        self.batch_size = batch_size
+
+    def run_once(self):
+        retrieved_buckets_count = 0
+        for i in range(0, self.num_requests, self.batch_size):
+            random_paths = self.tree.get_random_paths(self.batch_size)
+            unique_buckets = self.tree.get_unique_buckets_for_random_paths(random_paths)
+            retrieved_buckets_count += len(unique_buckets)
+        return retrieved_buckets_count
+
+    def run(self):
+        total_buckets_count = 0
+        for _ in range(EXPERIMETNS):
+            total_buckets_count += self.run_once()
+        print("Average number of buckets retrieved per request:", total_buckets_count / (self.num_requests * EXPERIMETNS))
+
+
+if __name__ == '__main__':
+    batch_size = 2
+    num_requests = 1000
+    tree_path_count = 2**20
+    sim = Simulator(num_requests, batch_size, tree_path_count)
+    sim.run()
diff --git a/simulation/simulate_test.py b/simulation/simulate_test.py
@@ -0,0 +1,52 @@
+import unittest
+from simulate import Tree
+
+
+class TestTree(unittest.TestCase):
+    def test_get_random_paths_is_in_range(self):
+        tree = Tree(tree_path_count=10)
+        batch_size = 5
+        paths = tree.get_random_paths(batch_size)
+        for path in paths:
+            self.assertTrue(1 <= path <= 10)
+
+    def test_get_random_paths_returns_correct_number_of_paths(self):
+        tree = Tree(tree_path_count=10)
+        batch_size = 5
+        paths = tree.get_random_paths(batch_size)
+        self.assertEqual(len(paths), batch_size)
+
+    def test_get_unique_buckets_for_random_paths_without_duplicates(self):
+        tree = Tree(tree_path_count=2)
+        #                 1
+        #                / \
+        #               2   3
+        # path number:  1   2
+        buckets = tree.get_unique_buckets_for_random_paths(random_paths=[1])
+        expected_buckets = {1, 2}
+        self.assertSetEqual(buckets, expected_buckets)
+
+    def test_get_unique_buckets_for_random_paths_with_duplicates(self):
+        tree = Tree(tree_path_count=2)
+        #                 1
+        #                / \
+        #               2   3
+        # path number:  1   2
+        buckets = tree.get_unique_buckets_for_random_paths(random_paths=[1, 2])
+        expected_buckets = {1, 2, 3}
+        self.assertSetEqual(buckets, expected_buckets)
+
+    def test_get_unique_buckets_for_random_paths_with_multiple_paths(self):
+        tree = Tree(tree_path_count=4)
+        #                    1
+        #                  /   \
+        #                 2     3
+        #                / \   / \
+        #               4   5 6   7
+        # path number:  1   2 3   4
+        buckets = tree.get_unique_buckets_for_random_paths(random_paths=[1, 2, 3])
+        expected_buckets = {1, 2, 4, 5, 3, 6}
+        self.assertSetEqual(buckets, expected_buckets)
+
+if __name__ == '__main__':
+    unittest.main()