import rerun as rr
import networkx as nx
import yaml
import numpy as np
import hashlib
import sys
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)
color_cache = {}
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) return color_cache[node]
def apply_forces_and_log(graph, pos):
damping = 0.9
max_force = 1
degree_scale = 2 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}
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 force[node1] += np.clip(force_vector, -max_force, max_force)
force[node2] -= np.clip(force_vector, -max_force, max_force)
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
for node in graph:
pos[node] += force[node] * damping
position = np.array(pos[node])
color = get_color_for_node(node) 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()])
rr.log("graph_nodes/graph_edges",
rr.LineStrips3D(edges_array, radii=0.02, colors=[1, 1, 1, 0.1]))
return pos
central_node = max(graph.degree, key=lambda x: x[1])[0]
initial_pos = nx.spring_layout(graph, dim=3)
initial_pos[central_node] = np.array([0.5, 0.5, 0.5])
final_pos = apply_forces_and_log(graph, initial_pos)