zoom inset

particle_numbers.py

@@ -1,44 +1,88 @@
+import pickle
+import random
+import re
+from os.path import expanduser
+from pathlib import Path
+from random import shuffle
 from sys import argv
 
 from matplotlib import pyplot as plt
 from matplotlib.axes import Axes
 from matplotlib.figure import Figure
+from mpl_toolkits.axes_grid1.inset_locator import mark_inset, inset_axes
 from rebound import SimulationArchive, Simulation
 
 from extradata import ExtraData
 from utils import filename_from_argv, plot_settings, is_ci
 
+cache_file = Path("particle_numbers_cache.pickle.xz")
+if cache_file.exists():
+    with cache_file.open("rb") as f:
+        cache = pickle.load(f)
+else:
+    cache = {}
+
 plot_settings()
 
 fig: Figure = plt.figure()
 ax: Axes = plt.gca()
+# axins2 = zoomed_inset_axes(ax, zoom=1, loc="center")
+axins2 = inset_axes(ax, 2,2, loc="center right")  # no zoom
+random.seed(1)
 
-for file in argv[1:]:
+colors = {
+    "rbf": "C0",
+    "nn": "C1",
+    "lz": "C2",
+    "pm": "C3",
+}
+files = argv[1:]
+shuffle(files)
+for file in files:
     fn = filename_from_argv(file)
     print(fn)
-    try:
-        ed = ExtraData.load(fn)
-        sa = SimulationArchive(str(fn.with_suffix(".bin")))
-    except:
-        print("skipping")
+    if "bak" in str(fn):
         continue
-    ts = []
-    Ns = []
-    sim: Simulation
-    for sim in sa:
-        num_planetesimals = sum(ed.pd(p).type == "planetesimal" for p in sim.particles)
-        num_embryos = sum(ed.pd(p).type == "embryo" for p in sim.particles)
-        N = num_embryos + num_planetesimals
-        Ns.append(N)
-        ts.append(sim.t)
-    perfect_merging = "pm" in str(fn)
-    ax.step(ts, Ns, label=fn, where="post", linestyle="dashed" if perfect_merging else "solid", linewidth=0.7)
+    if fn.name in cache:
+        Ns, ts = cache[fn.name]
+    else:
+        try:
+            ed = ExtraData.load(fn)
+            sa = SimulationArchive(str(fn.with_suffix(".bin")))
+        except:
+            print("skipping")
+            continue
+        ts = []
+        Ns = []
+        sim: Simulation
+        for sim in sa:
+            num_planetesimals = sum(ed.pd(p).type == "planetesimal" for p in sim.particles)
+            num_embryos = sum(ed.pd(p).type == "embryo" for p in sim.particles)
+            N = num_embryos + num_planetesimals
+            Ns.append(N)
+            ts.append(sim.t)
+        cache[fn.name] = Ns, ts
+    mode = re.search(r"final_(\w+)_", str(fn)).group(1)
+    for a in [ax, axins2]:
+        a.step(ts, Ns, label=mode, where="post", color=colors[mode], linewidth=0.7)
+    # break
+
+with cache_file.open("wb") as f:
+    pickle.dump(cache, f)
+
 ax.set_xlabel("time [yr]")
 ax.set_ylabel("number of objects")
 # ax.set_xscale("log")
 
-plt.legend()
-plt.tight_layout()
+axins2.set_xlim(1e8, 2e8)
+axins2.set_ylim(0, 15)
+axins2.tick_params(labelleft=False, labelbottom=False)
+mark_inset(ax, axins2, loc1=2, loc2=4, fc="none", ec="0.5")
+
+handles, labels = ax.get_legend_handles_labels()
+by_label = dict(zip(labels, handles))
+ax.legend(by_label.values(), by_label.keys())
+# plt.tight_layout()
 if not is_ci():
-    plt.savefig("/home/lukas/tmp/particle_numbers.pdf", transparent=True)
+    plt.savefig(expanduser(f"~/tmp/particle_numbers.pdf"))
 plt.show()