-
Notifications
You must be signed in to change notification settings - Fork 180
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Misc/update rerun example fdg (#548)
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
1 parent
9baf610
commit e048967
Showing
5 changed files
with
419 additions
and
42 deletions.
There are no files selected for viewing
File renamed without changes.
File renamed without changes.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
Oops, something went wrong.