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

improve graph

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

View file

@ -1,16 +1,18 @@
from os import environ from os import environ
import numpy as np
from graphviz import Digraph from graphviz import Digraph
from matplotlib.colors import to_hex
from extradata import ExtraData, CollisionMeta from extradata import ExtraData, CollisionMeta
from utils import filename_from_argv from utils import filename_from_argv, get_water_cmap
fn = filename_from_argv() fn = filename_from_argv()
ed = ExtraData.load(fn) ed = ExtraData.load(fn)
dot = Digraph(comment='Collisions') dot = Digraph(comment='Collisions')
dot.engine = "neato" # dot.engine = "neato"
if dot.engine == "neato": if dot.engine == "neato":
dot.attr("graph", overlap="false") 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}" label = f"{water_ret:.2f}/{mantle_ret:.2f}/{core_ret:.2f}"
first_parent = False first_parent = False
dot.node(label=label, name=f"{merged}-collision", shape="diamond")
dot.edge(f"{merged}-collision", str(merged))
else: else:
label = None label = None
dot.edge(str(parent), str(merged), xlabel=label) dot.edge(str(parent), f"{merged}-collision")
interacting_objects.add(parent) interacting_objects.add(parent)
interacting_objects.add(int(merged)) interacting_objects.add(int(merged))
cmap = get_water_cmap()
for name in ed.pdata.keys(): for name in ed.pdata.keys():
object = ed.pdata[name] object = ed.pdata[name]
if object.type == "sun": if object.type == "sun":
@ -44,12 +50,21 @@ for name in ed.pdata.keys():
mass = object.total_mass mass = object.total_mass
except KeyError: except KeyError:
mass = 0 mass = 0
dot.node(name=str(name), label=f"{displayname} ({object.water_mass_fraction:.2e}; {mass:.2e})", with np.errstate(divide='ignore'): # allow 0 water (becomes -inf)
shape='box' if object.type == "planetesimal" else "ellipse") color_val = (np.log10(object.water_mass_fraction) + 5) / 5
if object.escaped: color = cmap(color_val)
dot.edge(str(name), str("escaped")) if color_val > 0.55:
if object.collided_with_sun: textcolor = "white"
dot.edge(str(name), str("collided with sun")) 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"): if environ.get("CI"):
dot.save(fn.with_suffix(".gv")) dot.save(fn.with_suffix(".gv"))