Misc/update rerun example fdg

examples/rerun_example/force_driven_lockfile_graph.py

@@ -0,0 +1,108 @@
+import rerun as rr
+import networkx as nx
+import yaml
+import numpy as np
+import hashlib
+import sys
+
+# Give relative path or default to local pixi.lock
+lockfile_path = sys.argv[1] if len(sys.argv) > 1 else 'pixi.lock'
+
+with open(lockfile_path, 'r') as file:
+    lockfile_data = yaml.safe_load(file)
+
+package_data = lockfile_data['package']
+package_names = [package['name'] for package in package_data]
+
+graph = nx.DiGraph()
+for package in package_data:
+    package_name = package['name']
+    dependencies = package.get('dependencies', [])
+    graph.add_node(package_name)
+    for i, dep in enumerate(dependencies):
+        graph.add_edge(package_name, dep.split(" ")[0])
+
+rr.init("fdg", spawn=True)
+rr.connect()
+
+def hash_string_to_int(string):
+    return int(hashlib.sha256(string.encode('utf-8')).hexdigest(), 16) % (10 ** 8)
+
+
+# Memoization dictionary
+color_cache = {}
+
+
+# Function to get color
+def get_color_for_node(node):
+    if node not in color_cache:
+        np.random.seed(hash_string_to_int(node))
+        color_cache[node] = np.random.rand(3)  # Generate and store color
+    return color_cache[node]
+
+
+def apply_forces_and_log(graph, pos):
+    damping = 0.9
+    max_force = 1
+    degree_scale = 2  # Scale factor for degree-based forces
+    dist_scale = 0.5
+
+    iterations = 1000
+    repulsive_force = 0.01
+    attractive_force = 0.005
+
+    for iteration in range(iterations):
+        force = {node: np.zeros(3) for node in graph}
+
+        # Degree-based repulsive forces
+        for i, node1 in enumerate(graph):
+            for node2 in list(graph)[i + 1:]:
+                diff = pos[node1] - pos[node2]
+                dist = (np.linalg.norm(diff) + 1e-9) * dist_scale
+                degree_factor = (
+                        (graph.degree(node1) + graph.degree(node2)) * degree_scale)
+                repel = repulsive_force * degree_factor / dist ** 2
+                force_vector = repel * diff  # / dist
+                force[node1] += np.clip(force_vector, -max_force, max_force)
+                force[node2] -= np.clip(force_vector, -max_force, max_force)
+
+        # Degree-based attractive forces
+        for edge in graph.edges():
+            u, v = edge
+            diff = pos[u] - pos[v]
+            dist = dist = (np.linalg.norm(diff) + 1e-9) * dist_scale
+            if dist > 0:
+                degree_factor = (graph.degree(u) + graph.degree(v)) * degree_scale
+                attract = (attractive_force * dist ** 2) / degree_factor
+                force[u] -= attract * diff / dist
+                force[v] += attract * diff / dist
+
+        # Update positions with damping
+        for node in graph:
+            pos[node] += force[node] * damping
+            position = np.array(pos[node])
+            color = get_color_for_node(node)  # Retrieve color, memoized
+            rr.log(f"graph_nodes/{node}",
+                   rr.Points3D([position],
+                               colors=[color],
+                               radii=max(graph.degree(node) / 20, 0.5)),
+                   rr.AnyValues(node))
+
+        edges_array = np.array([[pos[u], pos[v]] for u, v in graph.edges()])
+
+        # Log the edges array
+        rr.log("graph_nodes/graph_edges",
+               rr.LineStrips3D(edges_array, radii=0.02, colors=[1, 1, 1, 0.1]))
+
+    return pos
+
+
+# Identify the node with the highest degree
+central_node = max(graph.degree, key=lambda x: x[1])[0]
+
+# Initial positions with the central node at the center
+initial_pos = nx.spring_layout(graph, dim=3)
+initial_pos[central_node] = np.array([0.5, 0.5, 0.5])  # Center position
+
+# Apply the force-directed simulation
+final_pos = apply_forces_and_log(graph, initial_pos)