Skip to content

Commit

Permalink
Misc/update rerun example fdg (#548)
Browse files Browse the repository at this point in the history
Fun new example, Its a visualization of the lockfile

[Screencast from 2023-12-06
23-26-36.webm](https://github.com/prefix-dev/pixi/assets/12893423/9fc774d2-c0df-4673-9cfe-71494028c956)
  • Loading branch information
ruben-arts authored Dec 7, 2023
1 parent 9baf610 commit e048967
Show file tree
Hide file tree
Showing 5 changed files with 419 additions and 42 deletions.
File renamed without changes.
File renamed without changes.
108 changes: 108 additions & 0 deletions examples/rerun_example/force_driven_lockfile_graph.py
Original file line number Diff line number Diff line change
@@ -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)
Loading

0 comments on commit e048967

Please sign in to comment.