2020-12-28 17:22:42 +01:00
|
|
|
import argparse
|
|
|
|
from collections import namedtuple
|
2021-03-23 13:21:20 +01:00
|
|
|
from math import log10, degrees
|
2020-12-28 17:22:42 +01:00
|
|
|
|
2020-12-27 13:48:38 +01:00
|
|
|
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
|
2021-05-03 16:16:28 +02:00
|
|
|
from scipy.constants import mega
|
2020-12-27 13:48:38 +01:00
|
|
|
|
|
|
|
from extradata import ExtraData, ParticleData
|
|
|
|
from utils import filename_from_argv
|
|
|
|
|
2021-05-03 16:16:28 +02:00
|
|
|
output_plots = False
|
|
|
|
if output_plots:
|
|
|
|
size_factor = 1
|
|
|
|
figsize = None
|
|
|
|
else:
|
|
|
|
size_factor = 3
|
|
|
|
figsize = (15, 10)
|
|
|
|
|
2020-12-28 17:22:42 +01:00
|
|
|
|
|
|
|
class MyProgramArgs(argparse.Namespace):
|
|
|
|
save_video: bool
|
2021-02-04 13:41:29 +01:00
|
|
|
log_time: bool
|
2020-12-28 17:22:42 +01:00
|
|
|
fps: int
|
|
|
|
duration: int
|
|
|
|
y_axis: str
|
|
|
|
|
|
|
|
|
|
|
|
plt.style.use("dark_background")
|
2020-12-27 13:48:38 +01:00
|
|
|
cmap: Colormap = matplotlib.cm.get_cmap('Blues')
|
|
|
|
|
2021-02-26 21:13:26 +01:00
|
|
|
mean_mass = None
|
|
|
|
|
2020-12-27 13:48:38 +01:00
|
|
|
|
2021-02-04 13:41:29 +01:00
|
|
|
def update_plot(num: int, args: MyProgramArgs, sa: SimulationArchive, ed: ExtraData, dots: PathCollection, title: Text):
|
2021-02-26 21:13:26 +01:00
|
|
|
global mean_mass
|
2020-12-28 17:22:42 +01:00
|
|
|
total_frames = args.fps * args.duration
|
2021-02-04 13:41:29 +01:00
|
|
|
if args.log_time:
|
|
|
|
log_timestep = (log10(ed.meta.current_time) - log10(50000)) / total_frames
|
|
|
|
time = 10 ** ((num + log10(50000) / log_timestep) * log_timestep)
|
2020-12-27 13:48:38 +01:00
|
|
|
else:
|
|
|
|
timestep = ed.meta.current_time / total_frames
|
2021-05-03 16:16:28 +02:00
|
|
|
print(num / total_frames)
|
2020-12-27 13:48:38 +01:00
|
|
|
time = num * timestep
|
2021-05-03 16:16:28 +02:00
|
|
|
print(f"{num / total_frames:.2f}, {time:.0f}")
|
2020-12-27 13:48:38 +01:00
|
|
|
sim = sa.getSimulation(t=time)
|
2020-12-28 17:22:42 +01:00
|
|
|
if time < 1e3:
|
|
|
|
timestr = f"{time:.0f}"
|
|
|
|
elif time < 1e6:
|
|
|
|
timestr = f"{time / 1e3:.2f}K"
|
|
|
|
else:
|
|
|
|
timestr = f"{time / 1e6:.2f}M"
|
|
|
|
title.set_text(f"({len(sim.particles)}) {timestr} Years")
|
2020-12-27 13:48:38 +01:00
|
|
|
|
|
|
|
p: Particle
|
|
|
|
|
|
|
|
water_fractions = []
|
2021-02-04 13:41:29 +01:00
|
|
|
for p in sim.particles[1:]:
|
2020-12-30 23:16:28 +01:00
|
|
|
pd: ParticleData = ed.pd(p)
|
2020-12-27 13:48:38 +01:00
|
|
|
wf = pd.water_mass_fraction
|
|
|
|
water_fractions.append(wf)
|
|
|
|
# a, e, i, M, M_rat = data[line]
|
|
|
|
# title.set_text(f"({len(a)}) {ages[line]:.2f}K Years")
|
2020-12-28 17:22:42 +01:00
|
|
|
a = [p.a for p in sim.particles[1:]]
|
2021-01-05 23:23:27 +01:00
|
|
|
m = np.array([p.m for p in sim.particles[1:]])
|
|
|
|
m[:2] /= 1e2
|
2020-12-28 17:22:42 +01:00
|
|
|
|
|
|
|
if args.y_axis == "e":
|
|
|
|
bla = np.array([a, [p.e for p in sim.particles[1:]]])
|
|
|
|
elif args.y_axis == "i":
|
2021-03-23 13:21:20 +01:00
|
|
|
bla = np.array([a, [degrees(p.inc) for p in sim.particles[1:]]])
|
2020-12-28 17:22:42 +01:00
|
|
|
elif args.y_axis == "Omega":
|
2021-03-23 13:21:20 +01:00
|
|
|
bla = np.array([a, [degrees(p.Omega) for p in sim.particles[1:]]])
|
2020-12-28 17:22:42 +01:00
|
|
|
else:
|
|
|
|
raise ValueError("invalid y-axis")
|
2020-12-27 13:48:38 +01:00
|
|
|
dots.set_offsets(bla.T)
|
2021-05-03 16:16:28 +02:00
|
|
|
with np.errstate(divide='ignore'): # allow 0 water (becomes -inf)
|
|
|
|
color_val = (np.log10(water_fractions) + 5) / 5
|
2020-12-27 13:48:38 +01:00
|
|
|
colors = cmap(color_val)
|
2021-02-26 21:13:26 +01:00
|
|
|
if not mean_mass:
|
|
|
|
mean_mass = np.mean(m[3:])
|
2021-05-03 16:16:28 +02:00
|
|
|
dots.set_sizes(size_factor * m / mean_mass)
|
2020-12-27 13:48:38 +01:00
|
|
|
dots.set_color(colors)
|
2021-05-03 16:16:28 +02:00
|
|
|
if output_plots:
|
|
|
|
plt.savefig(f"tmp/{num}_{time / mega}.pdf", transparent=True)
|
2020-12-27 13:48:38 +01:00
|
|
|
return dots, title
|
|
|
|
|
|
|
|
|
2020-12-28 17:22:42 +01:00
|
|
|
def main(args: MyProgramArgs):
|
|
|
|
total_frames = args.fps * args.duration
|
|
|
|
|
|
|
|
logtime = False
|
|
|
|
cmap: Colormap = matplotlib.cm.get_cmap("Blues")
|
2020-12-27 13:48:38 +01:00
|
|
|
|
2021-05-03 16:16:28 +02:00
|
|
|
fig1 = plt.figure(figsize=figsize)
|
2020-12-27 13:48:38 +01:00
|
|
|
|
2020-12-28 17:22:42 +01:00
|
|
|
l: PathCollection = plt.scatter([1], [1])
|
2020-12-27 13:48:38 +01:00
|
|
|
|
2020-12-28 17:22:42 +01:00
|
|
|
fn = filename_from_argv()
|
|
|
|
sa = SimulationArchive(str(fn.with_suffix(".bin")))
|
2021-02-04 17:29:01 +01:00
|
|
|
ed = ExtraData.load(fn)
|
2020-12-27 13:48:38 +01:00
|
|
|
|
2020-12-28 17:22:42 +01:00
|
|
|
plt.xlim(0, 10)
|
|
|
|
plt.xlabel("a")
|
|
|
|
title: Text = plt.title("0")
|
2020-12-27 13:48:38 +01:00
|
|
|
|
2020-12-28 17:22:42 +01:00
|
|
|
if args.y_axis == "e":
|
|
|
|
plt.ylim(-0.1, 1) # e
|
|
|
|
plt.ylabel("e")
|
|
|
|
elif args.y_axis == "i":
|
2021-03-23 13:21:20 +01:00
|
|
|
plt.ylim(0, 10) # i
|
2020-12-28 17:22:42 +01:00
|
|
|
plt.ylabel("i")
|
|
|
|
elif args.y_axis == "Omega":
|
|
|
|
plt.ylim(0, 360) # i
|
|
|
|
plt.ylabel("Omega")
|
2020-12-27 13:48:38 +01:00
|
|
|
|
2020-12-28 17:22:42 +01:00
|
|
|
# plt.yscale("log")
|
|
|
|
# plt.ylim(1e-7,1e-3)
|
2020-12-27 13:48:38 +01:00
|
|
|
|
2020-12-28 17:22:42 +01:00
|
|
|
# plt.ylim(-0.05, 0.2) # i
|
|
|
|
# plt.ylabel("water_fraction")
|
2020-12-27 13:48:38 +01:00
|
|
|
|
2020-12-28 17:22:42 +01:00
|
|
|
fig1.colorbar(ScalarMappable(norm=Normalize(vmin=-5, vmax=0), cmap=cmap), label="log(water fraction)")
|
2020-12-27 13:48:38 +01:00
|
|
|
|
2020-12-28 17:22:42 +01:00
|
|
|
plt.tight_layout()
|
2021-02-04 13:41:29 +01:00
|
|
|
line_ani = animation.FuncAnimation(fig1, update_plot, total_frames, fargs=(args, sa, ed, l, title),
|
2020-12-28 17:22:42 +01:00
|
|
|
interval=1000 / args.fps, repeat=False)
|
|
|
|
if args.save_video:
|
2021-02-26 21:13:26 +01:00
|
|
|
name = f"{args.y_axis}_{args.log_time}"
|
2021-02-04 13:41:29 +01:00
|
|
|
line_ani.save(str(fn.with_suffix(f".{name}.mp4")), dpi=200)
|
2020-12-28 17:22:42 +01:00
|
|
|
else:
|
|
|
|
plt.show()
|
|
|
|
|
|
|
|
|
|
|
|
if __name__ == "__main__":
|
|
|
|
parser = argparse.ArgumentParser(
|
|
|
|
description="create a video showing ",
|
|
|
|
formatter_class=argparse.ArgumentDefaultsHelpFormatter
|
|
|
|
)
|
|
|
|
parser.add_argument("file")
|
|
|
|
parser.add_argument("-v", "--save-video", action="store_true",
|
|
|
|
help="save video as .mp4 file")
|
2021-02-04 13:41:29 +01:00
|
|
|
parser.add_argument("-l", "--log-time", action="store_true", help="logarithmic time")
|
2020-12-28 17:22:42 +01:00
|
|
|
parser.add_argument("--fps", default=5, type=int,
|
|
|
|
help="frames per second")
|
|
|
|
parser.add_argument("--duration", default=10, type=int,
|
|
|
|
help="duration in seconds")
|
|
|
|
parser.add_argument("--y-axis", default="e", type=str, choices=["e", "i", "Omega"],
|
|
|
|
help="what to show on the y-axis")
|
|
|
|
ArgNamespace = namedtuple('ArgNamespace', ['some_arg', 'another_arg'])
|
|
|
|
args = parser.parse_args()
|
|
|
|
print(vars(args))
|
|
|
|
print(args.save_video)
|
|
|
|
# noinspection PyTypeChecker
|
|
|
|
main(args)
|