mirror of
https://github.com/Findus23/rebound-collisions.git
synced 2024-09-19 15:53:48 +02:00
add timeplot
This commit is contained in:
parent
3fe9ef163a
commit
a685205c56
2 changed files with 95 additions and 0 deletions
1
.gitignore
vendored
1
.gitignore
vendored
|
@ -11,3 +11,4 @@ __pycache__/
|
|||
data/
|
||||
*.svg
|
||||
*.mp4
|
||||
tmp/
|
||||
|
|
94
timeplot.py
Normal file
94
timeplot.py
Normal file
|
@ -0,0 +1,94 @@
|
|||
import matplotlib
|
||||
import matplotlib.animation as animation
|
||||
import matplotlib.pyplot as plt
|
||||
import numpy as np
|
||||
from matplotlib.cm import ScalarMappable
|
||||
from matplotlib.collections import PathCollection
|
||||
from matplotlib.colors import Normalize, Colormap
|
||||
from matplotlib.text import Text
|
||||
from rebound import SimulationArchive, Particle
|
||||
|
||||
from extradata import ExtraData, ParticleData
|
||||
from utils import filename_from_argv
|
||||
|
||||
matplotlib.use('Qt4Agg')
|
||||
plt.style.use('dark_background')
|
||||
fps = 5
|
||||
duration = 10 # s
|
||||
total_frames = fps * duration
|
||||
logtime = False
|
||||
cmap: Colormap = matplotlib.cm.get_cmap('Blues')
|
||||
|
||||
|
||||
def update_line(num: int, sa: SimulationArchive, ed: ExtraData, dots: PathCollection, title: Text):
|
||||
# line = num * int((len(data) / total_frames))
|
||||
# print(line, num)
|
||||
if logtime:
|
||||
timestep = total_frames
|
||||
|
||||
else:
|
||||
timestep = ed.meta.current_time / total_frames
|
||||
# timestep = 10e6 / total_frames
|
||||
time = num * timestep
|
||||
print(time)
|
||||
sim = sa.getSimulation(t=time)
|
||||
title.set_text(f"({len(sim.particles)}) {time / 1e6:.2f}M Years")
|
||||
|
||||
p: Particle
|
||||
a = [p.a for p in sim.particles[1:]]
|
||||
e = [p.e for p in sim.particles[1:]]
|
||||
|
||||
water_fractions = []
|
||||
for p in sim.particles[3:]:
|
||||
# try:
|
||||
pd: ParticleData = ed.pdata[p.hash.value]
|
||||
wf = pd.water_mass_fraction
|
||||
# except KeyError: # gas planet
|
||||
# print(p.hash.value)
|
||||
# wf = 0
|
||||
water_fractions.append(wf)
|
||||
# a, e, i, M, M_rat = data[line]
|
||||
# title.set_text(f"({len(a)}) {ages[line]:.2f}K Years")
|
||||
bla = np.array([a, e])
|
||||
dots.set_offsets(bla.T)
|
||||
water_fractions = np.array(water_fractions)
|
||||
color_val = (np.log10(water_fractions) + 5) / 5
|
||||
print(color_val)
|
||||
colors = cmap(color_val)
|
||||
print(colors)
|
||||
dots.set_color(colors)
|
||||
plt.savefig("tmp/" + str(num) + ".pdf",transparent=True)
|
||||
return dots, title
|
||||
|
||||
|
||||
fig1 = plt.figure()
|
||||
|
||||
l: PathCollection = plt.scatter([1], [1])
|
||||
|
||||
fn = filename_from_argv()
|
||||
sa = SimulationArchive(str(fn.with_suffix(".bin")))
|
||||
ed = ExtraData.load(fn.with_suffix(".extra.json"))
|
||||
|
||||
plt.xlim(0, 10)
|
||||
plt.xlabel("a")
|
||||
title: Text = plt.title("0")
|
||||
plt.ylim(-0.1, 1) # e
|
||||
plt.ylabel("e")
|
||||
|
||||
# plt.ylim(0, 90) # i
|
||||
# plt.ylabel("i")
|
||||
|
||||
# plt.yscale('log')
|
||||
# plt.ylim(1e-7,1e-3)
|
||||
|
||||
# plt.ylim(-0.05, 0.2) # i
|
||||
# plt.ylabel("water_fraction")
|
||||
|
||||
fig1.colorbar(ScalarMappable(norm=Normalize(vmin=-5, vmax=0), cmap=cmap), label="log(water fraction)")
|
||||
|
||||
plt.tight_layout()
|
||||
|
||||
line_ani = animation.FuncAnimation(fig1, update_line, total_frames, fargs=(sa, ed, l, title),
|
||||
interval=1000 / fps, repeat=False)
|
||||
line_ani.save(str(fn.with_suffix(".mp4")), dpi=300)
|
||||
# plt.show()
|
Loading…
Reference in a new issue