add tree graph option, refactor, improve memory

generators mean no more numpy

generate.py

 import argparse
-import math
 import itertools
+import math
+import queue
 import random
-import numpy as np
 import string
-import random
-from pathlib import Path
+import sys
 from datetime import datetime
+from pathlib import Path
+from timeit import default_timer as timer
 
 def fully_connected(n):
     return int((n * (n - 1)) / 2)
@@ -15,63 +16,124 @@ def default_filename():
     # the current time in ISO format (ish)
     return datetime.now().strftime("%Y-%m-%d-%H-%M-%S-%f") + ".txt"
 
+def format_connection(a, b, distance):
+    # for writing connections to file
+    return str(a) + " " + str(b) + " " + str(distance) + "\n"
+
+class Graph:
+    def __init__(self, num_nodes, num_connections):
+        self.num_nodes = num_nodes
+        self.num_connections = num_connections
+
+        if num_connections > fully_connected(num_nodes):
+            print("too many connections (> complete graph)")
+            exit(1)
+
+        self.charging = set(random.sample(range(num_nodes), num_charging))
+
+    def node_names(self):
+        '''
+        generator for node names from the alphabet based on num_nodes
+        '''
+        alphabet = list(string.ascii_uppercase)
+        repeat = 1 + int(math.log(self.num_nodes, len(alphabet)))
+
+        join_string = lambda x: ("".join(y) for y in x)
+
+        return join_string(
+                itertools.islice(
+                    itertools.product(alphabet, repeat=repeat),
+                    self.num_nodes
+                    )
+                )
+
+    def nodes(self):
+        for i, name in enumerate(graph.node_names()):
+            yield (name, i in self.charging)
+
+    def connections(self):
+        connections = set()
+        while len(connections) < self.num_connections:
+            pair = tuple(random.sample(range(num_nodes), 2))
+            if pair not in connections:
+                connections.add(pair)
+                yield pair
+
+
+class TreeGraph(Graph):
+    def __init__(self, height, m):
+        num_nodes = self.number_of_nodes(m, height)
+        super().__init__(num_nodes, num_nodes - 1)
+        self.height = height
+        self.m = m
+
+    def number_of_nodes(self, m, height):
+        return int((m ** height - 1) / (m - 1))
+
+    def connections(self):
+        '''
+        breadth-first generation for a m-ary tree graph
+        '''
+        q = queue.Queue()
+        total = 0
+        q.put(total)
+        while not q.empty() and total < self.num_connections:
+            parent = q.get()
+            for _ in range(self.m):
+                total += 1
+                child = total
+                q.put(child)
+                yield (parent, child)
+
 
 parser = argparse.ArgumentParser()
-parser.add_argument("nodes", type=int)
+parser.add_argument("size", type=int, help="number of nodes or height of m-ary tree")
 parser.add_argument("-o", dest="file", default=default_filename(), help="output file")
 parser.add_argument("-c", "--charging", type=int, default=0, help="number of charging points")
-parser.add_argument("-e", "--edges", type=int, help="number of edges, defaults to a complete graph")
 parser.add_argument("--min", type=int, default=1, help="minimum distance")
 parser.add_argument("--max", type=int, default=10, help="maximum distance")
 
+graph_group = parser.add_mutually_exclusive_group()
+graph_group.add_argument("-e", "--edges", type=int, help="number of edges, defaults to a complete graph")
+# no default for children since normal graph is the default
+graph_group.add_argument("-m", "--children", type=int, help="number of children for m-ary tree")
+
 args = parser.parse_args()
 
 out = Path(args.file)
 
-num_nodes = args.nodes
 num_charging = args.charging
-num_connections = args.edges or fully_connected(num_nodes)
-
-if num_connections > fully_connected(num_nodes):
-    print("too many connections (> complete graph)")
-    exit(1)
 
 min_distance = args.min
 max_distance = args.max
 
-alphabet = list(string.ascii_uppercase)
-repeat = 1 + int(math.log(num_nodes, len(alphabet)))
-
-join_string = lambda x: ("".join(y) for y in x)
-
-node_names = join_string(
-        itertools.islice(
-            itertools.product(alphabet, repeat=repeat),
-            num_nodes
-            )
-        )
-
-# generate paths
-connections = set()
-while len(connections) < num_connections:
-    pair = tuple(random.sample(range(num_nodes), 2))
-    connections.add(pair)
+start = timer()
 
-charging = np.zeros(num_nodes, dtype=bool)
-selected_charging = np.random.choice(np.arange(num_nodes), num_charging, replace=False)
-charging[selected_charging] = True
+if args.children:
+    # m-ary tree graph
+    graph = TreeGraph(args.size, args.children)
+else:
+    # normal graph
+    num_nodes = args.size
+    graph = Graph(num_nodes, args.edges or fully_connected(num_nodes))
 
 with out.open("w") as f:
-    f.write(f"{num_nodes} {num_connections}\n")
+    f.write(f"{graph.num_nodes} {graph.num_connections}\n")
     # write nodes and charging status
-    for i, n in enumerate(node_names):
-        f.write(str(n) + " " + str(int(charging[i])) + "\n")
+    print("writing names...")
+    for i, (name, is_charging) in enumerate(graph.nodes()):
+        print(str(int(i / graph.num_nodes * 100)) + "%", end="\r")
+        sys.stdout.flush()
+        f.write(name + " " + str(int(is_charging)) + "\n")
+
     # write connections
-    for a, b in connections:
-        try:
-            distance = random.randint(min_distance, max_distance)
-            f.write(str(a) +  " " + str(b) + " " + str(distance) + "\n")
-        except ValueError:
-            print(f"{a=} {b=}")
-
-print("saved to", out.name)
+    print("writing connections...")
+    for i, (a, b) in enumerate(graph.connections()):
+        print(str(int(i / graph.num_connections * 100)) + "%", end="\r")
+        sys.stdout.flush()
+        distance = random.randint(min_distance, max_distance)
+        f.write(format_connection(a, b, distance))
+
+end = timer()
+
+print(f"in {str(round(end - start, 2))}s wrote {graph.num_nodes} nodes, {graph.num_connections} edges to {out.name}")

requirements.txt

@@ -4,7 +4,6 @@ fonttools==4.47.0
 kiwisolver==1.4.5
 matplotlib==3.8.2
 networkx==3.2.1
-numpy==1.26.3
 packaging==23.2
 pillow==10.2.0
 pyparsing==3.1.1