From 203d60a86d17d77fa5b8a2b27a2ed1550977b78a Mon Sep 17 00:00:00 2001
From: Lukas Winkler <git@lw1.at>
Date: Sun, 30 Jan 2022 19:51:14 +0100
Subject: [PATCH] improve graph

---
 graph.py | 33 ++++++++++++++++++++++++---------
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/graph.py b/graph.py
index 6fc7378..fb882a0 100644
--- a/graph.py
+++ b/graph.py
@@ -1,16 +1,18 @@
 from os import environ
 
+import numpy as np
 from graphviz import Digraph
+from matplotlib.colors import to_hex
 
 from extradata import ExtraData, CollisionMeta
-from utils import filename_from_argv
+from utils import filename_from_argv, get_water_cmap
 
 fn = filename_from_argv()
 ed = ExtraData.load(fn)
 
 dot = Digraph(comment='Collisions')
 
-dot.engine = "neato"
+# dot.engine = "neato"
 
 if dot.engine == "neato":
     dot.attr("graph", overlap="false")
@@ -26,12 +28,16 @@ for merged, originals in ed.tree.get_tree().items():
 
             label = f"{water_ret:.2f}/{mantle_ret:.2f}/{core_ret:.2f}"
             first_parent = False
+            dot.node(label=label, name=f"{merged}-collision", shape="diamond")
+            dot.edge(f"{merged}-collision", str(merged))
         else:
             label = None
-        dot.edge(str(parent), str(merged), xlabel=label)
+        dot.edge(str(parent), f"{merged}-collision")
         interacting_objects.add(parent)
         interacting_objects.add(int(merged))
 
+cmap = get_water_cmap()
+
 for name in ed.pdata.keys():
     object = ed.pdata[name]
     if object.type == "sun":
@@ -44,12 +50,21 @@ for name in ed.pdata.keys():
         mass = object.total_mass
     except KeyError:
         mass = 0
-    dot.node(name=str(name), label=f"{displayname} ({object.water_mass_fraction:.2e}; {mass:.2e})",
-             shape='box' if object.type == "planetesimal" else "ellipse")
-    if object.escaped:
-        dot.edge(str(name), str("escaped"))
-    if object.collided_with_sun:
-        dot.edge(str(name), str("collided with sun"))
+    with np.errstate(divide='ignore'):  # allow 0 water (becomes -inf)
+        color_val = (np.log10(object.water_mass_fraction) + 5) / 5
+    color = cmap(color_val)
+    if color_val > 0.55:
+        textcolor = "white"
+    else:
+        textcolor = "black"
+    print(to_hex(color), color_val)
+    dot.node(name=str(name), label=f"m={mass:.1e}\nwmf={object.water_mass_fraction:.1e}",
+             shape='box' if object.type == "planetesimal" else "ellipse", style="filled", fillcolor=to_hex(color),
+             fontcolor=textcolor)
+    # if object.escaped:
+    #     dot.edge(str(name), str("escaped"))
+    # if object.collided_with_sun:
+    #     dot.edge(str(name), str("collided with sun"))
 
 if environ.get("CI"):
     dot.save(fn.with_suffix(".gv"))