mirror/rebound-collisions
1
0
Fork 0
mirror of https://github.com/Findus23/rebound-collisions.git synced 2024-09-19 15:53:48 +02:00
Code Releases Activity

improve graph

Browse source
This commit is contained in:
Lukas Winkler 2022-01-30 19:51:14 +01:00
parent 1b26f4091e
commit 203d60a86d
Signed by: lukas
GPG key ID: 54DE4D798D244853
1 changed files with 24 additions and 9 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

33
graph.py
View file

@ -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"))