From 589313cde42d8a030a57a469bab32c0b02488fe3 Mon Sep 17 00:00:00 2001 From: Lukas Winkler Date: Mon, 3 May 2021 16:16:28 +0200 Subject: [PATCH] plotting improvements (for presentation) --- analyze.py | 8 ++++++-- collisionhistory.py | 17 ++++++++++++++--- collisionstats.py | 20 ++++++++++++-------- particle_numbers.py | 10 +++++++--- timeplot.py | 23 +++++++++++++++++------ utils/plotting.py | 4 ++++ visualize_collision.py | 6 ++++-- 7 files changed, 64 insertions(+), 24 deletions(-) diff --git a/analyze.py b/analyze.py index 8572f62..125a901 100644 --- a/analyze.py +++ b/analyze.py @@ -2,7 +2,9 @@ import matplotlib.pyplot as plt from rebound import SimulationArchive, Particle, Simulation from extradata import ExtraData -from utils import filename_from_argv +from utils import filename_from_argv, plot_settings + +plot_settings() fn = filename_from_argv() sa = SimulationArchive(str(fn.with_suffix(".bin"))) @@ -28,7 +30,9 @@ for name, d in data.items(): if False: plt.scatter(times, values, label=name, s=.9) else: - plt.plot(times, values, label=name) + plt.plot(times, values, label=name, linewidth=0.6) # plt.legend() # OrbitPlot(sim, slices=1) +plt.tight_layout() +plt.savefig("/home/lukas/tmp/time.pdf", transparent=True) plt.show() diff --git a/collisionhistory.py b/collisionhistory.py index d21a228..a346c90 100644 --- a/collisionhistory.py +++ b/collisionhistory.py @@ -4,7 +4,9 @@ from matplotlib.figure import Figure from rebound import SimulationArchive, Simulation from extradata import ExtraData, CollisionMeta -from utils import filename_from_argv +from utils import filename_from_argv, earth_mass, earth_water_mass, plot_settings + +plot_settings() fn = filename_from_argv() ed = ExtraData.load(fn) @@ -15,8 +17,8 @@ print([p.hash.value for p in last_sim.particles]) print(last_sim.t) fig: Figure = plt.figure() -ax_masses: Axes = fig.add_subplot(2,1, 1) -ax_wmfs: Axes = fig.add_subplot(2,1, 2) +ax_masses: Axes = fig.add_subplot(2, 1, 1) +ax_wmfs: Axes = fig.add_subplot(2, 1, 2) for particle in last_sim.particles: if ed.pd(particle).type in ["sun", "gas giant"]: @@ -37,6 +39,7 @@ for particle in last_sim.particles: break meta: CollisionMeta = collision["meta"] parents = collision["parents"] + print("mass:", ed.pdata[hash].total_mass / earth_mass) masses.append(ed.pdata[hash].total_mass) objects.append(ed.pdata[hash]) times.append(meta.time) @@ -57,9 +60,17 @@ for particle in last_sim.particles: ax_wmfs.set_ylabel("water mass fraction") ax_masses.set_ylabel("masses [kg]") for ax in [ax_wmfs, ax_masses]: + ax.set_xlim(1e4, ed.meta.current_time) ax.set_xlabel("time [yr]") ax.set_xscale("log") ax.set_yscale("log") ax.legend() + twin_ax = ax_masses.twinx() + mn, mx = ax_masses.get_ylim() + twin_ax.set_ylim(mn / earth_mass, mx / earth_mass) + twin_ax.set_ylabel('[$M_\\oplus$]') + twin_ax.set_yscale("log") + ax_wmfs.axhline(earth_water_mass/earth_mass,linestyle="dotted") fig.tight_layout() +fig.savefig("/home/lukas/tmp/collisionhistory.pdf", transparent=True) plt.show() diff --git a/collisionstats.py b/collisionstats.py index 4f56536..5984c8f 100644 --- a/collisionstats.py +++ b/collisionstats.py @@ -5,7 +5,9 @@ from matplotlib import pyplot as plt from scipy.constants import mega from extradata import ExtraData, CollisionMeta -from utils import filename_from_argv, create_figure +from utils import filename_from_argv, create_figure, plot_settings + +plot_settings() dotsize = 1.5 @@ -36,7 +38,7 @@ ax4.set_yscale("log") ax4.set_xlabel(time_label) ax4.set_ylabel("water loss") - +sum = 0 for file in argv[1:]: fn = filename_from_argv(file) print(fn) @@ -51,6 +53,7 @@ for file in argv[1:]: times = [] for collision in ed.tree.get_tree().values(): + sum += 1 meta: CollisionMeta = collision["meta"] vs.append(meta.input.velocity_esc) angles.append(meta.input.alpha) @@ -60,14 +63,15 @@ for file in argv[1:]: water_loss.append(loss) times.append(meta.time / mega) - ax1.scatter(angles, vs,s=dotsize) - ax2.scatter(times, angles,s=dotsize) - ax3.scatter(times, vs,s=dotsize) - ax4.scatter(times, water_loss,s=dotsize) + ax1.scatter(angles, vs, s=dotsize) + ax2.scatter(times, angles, s=dotsize) + ax3.scatter(times, vs, s=dotsize) + ax4.scatter(times, water_loss, s=dotsize) ax4.autoscale(enable=True, axis='y') for i, fig in enumerate([fig1, fig2, fig3, fig4]): - fig.savefig(Path("plots") / fn.with_suffix(f".collision{i}.pdf").name) - + fig.tight_layout() + fig.savefig(Path("plots") / fn.with_suffix(f".collision{i}.pdf").name, transparent=True) +print(sum) plt.show() diff --git a/particle_numbers.py b/particle_numbers.py index ab3e211..f0fa46a 100644 --- a/particle_numbers.py +++ b/particle_numbers.py @@ -6,7 +6,9 @@ from matplotlib.figure import Figure from rebound import SimulationArchive, Simulation from extradata import ExtraData -from utils import filename_from_argv +from utils import filename_from_argv, plot_settings + +plot_settings() fig: Figure = plt.figure() ax: Axes = plt.gca() @@ -29,11 +31,13 @@ for file in argv[1:]: N = num_embryos + num_planetesimals Ns.append(N) ts.append(sim.t) - - ax.step(ts, Ns, label=fn, where="post") + perfect_merging = "pm" in str(fn) + ax.step(ts, Ns, label=fn, where="post", linestyle="dashed" if perfect_merging else "solid",linewidth=0.7) ax.set_xlabel("time [yr]") ax.set_ylabel("number of objects") # ax.set_xscale("log") plt.legend() +plt.tight_layout() +plt.savefig("/home/lukas/tmp/particle_numbers.pdf", transparent=True) plt.show() diff --git a/timeplot.py b/timeplot.py index ded0520..f62598a 100644 --- a/timeplot.py +++ b/timeplot.py @@ -11,10 +11,19 @@ from matplotlib.collections import PathCollection from matplotlib.colors import Normalize, Colormap from matplotlib.text import Text from rebound import SimulationArchive, Particle +from scipy.constants import mega from extradata import ExtraData, ParticleData from utils import filename_from_argv +output_plots = False +if output_plots: + size_factor = 1 + figsize = None +else: + size_factor = 3 + figsize = (15, 10) + class MyProgramArgs(argparse.Namespace): save_video: bool @@ -38,8 +47,9 @@ def update_plot(num: int, args: MyProgramArgs, sa: SimulationArchive, ed: ExtraD time = 10 ** ((num + log10(50000) / log_timestep) * log_timestep) else: timestep = ed.meta.current_time / total_frames + print(num / total_frames) time = num * timestep - print(time) + print(f"{num / total_frames:.2f}, {time:.0f}") sim = sa.getSimulation(t=time) if time < 1e3: timestr = f"{time:.0f}" @@ -71,14 +81,15 @@ def update_plot(num: int, args: MyProgramArgs, sa: SimulationArchive, ed: ExtraD else: raise ValueError("invalid y-axis") dots.set_offsets(bla.T) - water_fractions = np.array(water_fractions) - color_val = (np.log10(water_fractions) + 5) / 5 + with np.errstate(divide='ignore'): # allow 0 water (becomes -inf) + color_val = (np.log10(water_fractions) + 5) / 5 colors = cmap(color_val) if not mean_mass: mean_mass = np.mean(m[3:]) - dots.set_sizes(3 * m / mean_mass) + dots.set_sizes(size_factor * m / mean_mass) dots.set_color(colors) - # plt.savefig("tmp/" + str(num) + ".pdf",transparent=True) + if output_plots: + plt.savefig(f"tmp/{num}_{time / mega}.pdf", transparent=True) return dots, title @@ -88,7 +99,7 @@ def main(args: MyProgramArgs): logtime = False cmap: Colormap = matplotlib.cm.get_cmap("Blues") - fig1 = plt.figure(figsize=(15, 10)) + fig1 = plt.figure(figsize=figsize) l: PathCollection = plt.scatter([1], [1]) diff --git a/utils/plotting.py b/utils/plotting.py index 91a2a55..89e0b2e 100644 --- a/utils/plotting.py +++ b/utils/plotting.py @@ -12,3 +12,7 @@ def create_figure() -> Tuple[Figure, Axes]: fig: Figure = plt.figure() ax: Axes = fig.gca() return fig, ax + + +def plot_settings() -> None: + plt.style.use("dark_background") diff --git a/visualize_collision.py b/visualize_collision.py index 280379d..bda434c 100644 --- a/visualize_collision.py +++ b/visualize_collision.py @@ -5,11 +5,13 @@ from matplotlib.patches import FancyArrowPatch from mpl_toolkits.mplot3d import Axes3D, proj3d from extradata import ExtraData, CollisionMeta -from utils import filename_from_argv +from utils import filename_from_argv, plot_settings fn = filename_from_argv() ed = ExtraData.load(fn) +plot_settings() + def get_circle(dx, dy, dz, r): u, v = np.mgrid[0:2 * np.pi:40j, 0:np.pi:20j] @@ -64,5 +66,5 @@ for collision in ed.tree.get_tree().values(): ax.set_zlim3d((zmin, zmin - diff)) # ax.set_ylim3d(XYZlim) # ax.set_zlim3d(XYZlim * 3/4) - + plt.savefig("/home/lukas/tmp/3d.pdf") plt.show()