mirror of
https://github.com/Findus23/rebound-collisions.git
synced 2024-09-19 15:53:48 +02:00
add argparse to timeplot
This commit is contained in:
parent
a685205c56
commit
c984716b15
2 changed files with 100 additions and 37 deletions
127
timeplot.py
127
timeplot.py
|
@ -1,3 +1,6 @@
|
||||||
|
import argparse
|
||||||
|
from collections import namedtuple
|
||||||
|
|
||||||
import matplotlib
|
import matplotlib
|
||||||
import matplotlib.animation as animation
|
import matplotlib.animation as animation
|
||||||
import matplotlib.pyplot as plt
|
import matplotlib.pyplot as plt
|
||||||
|
@ -11,32 +14,43 @@ from rebound import SimulationArchive, Particle
|
||||||
from extradata import ExtraData, ParticleData
|
from extradata import ExtraData, ParticleData
|
||||||
from utils import filename_from_argv
|
from utils import filename_from_argv
|
||||||
|
|
||||||
matplotlib.use('Qt4Agg')
|
|
||||||
plt.style.use('dark_background')
|
class MyProgramArgs(argparse.Namespace):
|
||||||
fps = 5
|
save_video: bool
|
||||||
duration = 10 # s
|
fps: int
|
||||||
total_frames = fps * duration
|
duration: int
|
||||||
logtime = False
|
y_axis: str
|
||||||
|
|
||||||
|
|
||||||
|
matplotlib.use("Qt5Agg")
|
||||||
|
plt.style.use("dark_background")
|
||||||
cmap: Colormap = matplotlib.cm.get_cmap('Blues')
|
cmap: Colormap = matplotlib.cm.get_cmap('Blues')
|
||||||
|
# fps = 5
|
||||||
|
# duration = 10 # s
|
||||||
|
logtime = False # TODO: support logarithmic time
|
||||||
|
|
||||||
|
|
||||||
def update_line(num: int, sa: SimulationArchive, ed: ExtraData, dots: PathCollection, title: Text):
|
def update_line(num: int, args: MyProgramArgs, sa: SimulationArchive, ed: ExtraData, dots: PathCollection, title: Text):
|
||||||
# line = num * int((len(data) / total_frames))
|
# line = num * int((len(data) / total_frames))
|
||||||
|
total_frames = args.fps * args.duration
|
||||||
# print(line, num)
|
# print(line, num)
|
||||||
if logtime:
|
if logtime:
|
||||||
timestep = total_frames
|
timestep = total_frames
|
||||||
|
|
||||||
else:
|
else:
|
||||||
timestep = ed.meta.current_time / total_frames
|
timestep = ed.meta.current_time / total_frames
|
||||||
# timestep = 10e6 / total_frames
|
# timestep = 10e6 / total_frames
|
||||||
time = num * timestep
|
time = num * timestep
|
||||||
print(time)
|
print(time)
|
||||||
sim = sa.getSimulation(t=time)
|
sim = sa.getSimulation(t=time)
|
||||||
title.set_text(f"({len(sim.particles)}) {time / 1e6:.2f}M Years")
|
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")
|
||||||
|
|
||||||
p: Particle
|
p: Particle
|
||||||
a = [p.a for p in sim.particles[1:]]
|
|
||||||
e = [p.e for p in sim.particles[1:]]
|
|
||||||
|
|
||||||
water_fractions = []
|
water_fractions = []
|
||||||
for p in sim.particles[3:]:
|
for p in sim.particles[3:]:
|
||||||
|
@ -49,7 +63,16 @@ def update_line(num: int, sa: SimulationArchive, ed: ExtraData, dots: PathCollec
|
||||||
water_fractions.append(wf)
|
water_fractions.append(wf)
|
||||||
# a, e, i, M, M_rat = data[line]
|
# a, e, i, M, M_rat = data[line]
|
||||||
# title.set_text(f"({len(a)}) {ages[line]:.2f}K Years")
|
# title.set_text(f"({len(a)}) {ages[line]:.2f}K Years")
|
||||||
bla = np.array([a, e])
|
a = [p.a for p in sim.particles[1:]]
|
||||||
|
|
||||||
|
if args.y_axis == "e":
|
||||||
|
bla = np.array([a, [p.e for p in sim.particles[1:]]])
|
||||||
|
elif args.y_axis == "i":
|
||||||
|
bla = np.array([a, [p.inc for p in sim.particles[1:]]])
|
||||||
|
elif args.y_axis == "Omega":
|
||||||
|
bla = np.array([a, [p.Omega for p in sim.particles[1:]]])
|
||||||
|
else:
|
||||||
|
raise ValueError("invalid y-axis")
|
||||||
dots.set_offsets(bla.T)
|
dots.set_offsets(bla.T)
|
||||||
water_fractions = np.array(water_fractions)
|
water_fractions = np.array(water_fractions)
|
||||||
color_val = (np.log10(water_fractions) + 5) / 5
|
color_val = (np.log10(water_fractions) + 5) / 5
|
||||||
|
@ -57,38 +80,72 @@ def update_line(num: int, sa: SimulationArchive, ed: ExtraData, dots: PathCollec
|
||||||
colors = cmap(color_val)
|
colors = cmap(color_val)
|
||||||
print(colors)
|
print(colors)
|
||||||
dots.set_color(colors)
|
dots.set_color(colors)
|
||||||
plt.savefig("tmp/" + str(num) + ".pdf",transparent=True)
|
# plt.savefig("tmp/" + str(num) + ".pdf",transparent=True)
|
||||||
return dots, title
|
return dots, title
|
||||||
|
|
||||||
|
|
||||||
fig1 = plt.figure()
|
def main(args: MyProgramArgs):
|
||||||
|
total_frames = args.fps * args.duration
|
||||||
|
|
||||||
l: PathCollection = plt.scatter([1], [1])
|
logtime = False
|
||||||
|
cmap: Colormap = matplotlib.cm.get_cmap("Blues")
|
||||||
|
|
||||||
fn = filename_from_argv()
|
fig1 = plt.figure()
|
||||||
sa = SimulationArchive(str(fn.with_suffix(".bin")))
|
|
||||||
ed = ExtraData.load(fn.with_suffix(".extra.json"))
|
|
||||||
|
|
||||||
plt.xlim(0, 10)
|
l: PathCollection = plt.scatter([1], [1])
|
||||||
plt.xlabel("a")
|
|
||||||
title: Text = plt.title("0")
|
|
||||||
plt.ylim(-0.1, 1) # e
|
|
||||||
plt.ylabel("e")
|
|
||||||
|
|
||||||
# plt.ylim(0, 90) # i
|
fn = filename_from_argv()
|
||||||
# plt.ylabel("i")
|
sa = SimulationArchive(str(fn.with_suffix(".bin")))
|
||||||
|
ed = ExtraData.load(fn.with_suffix(".extra.json"))
|
||||||
|
|
||||||
# plt.yscale('log')
|
plt.xlim(0, 10)
|
||||||
# plt.ylim(1e-7,1e-3)
|
plt.xlabel("a")
|
||||||
|
title: Text = plt.title("0")
|
||||||
|
|
||||||
# plt.ylim(-0.05, 0.2) # i
|
if args.y_axis == "e":
|
||||||
# plt.ylabel("water_fraction")
|
plt.ylim(-0.1, 1) # e
|
||||||
|
plt.ylabel("e")
|
||||||
|
elif args.y_axis == "i":
|
||||||
|
plt.ylim(0, 90) # i
|
||||||
|
plt.ylabel("i")
|
||||||
|
elif args.y_axis == "Omega":
|
||||||
|
plt.ylim(0, 360) # i
|
||||||
|
plt.ylabel("Omega")
|
||||||
|
|
||||||
fig1.colorbar(ScalarMappable(norm=Normalize(vmin=-5, vmax=0), cmap=cmap), label="log(water fraction)")
|
# plt.yscale("log")
|
||||||
|
# plt.ylim(1e-7,1e-3)
|
||||||
|
|
||||||
plt.tight_layout()
|
# plt.ylim(-0.05, 0.2) # i
|
||||||
|
# plt.ylabel("water_fraction")
|
||||||
|
|
||||||
line_ani = animation.FuncAnimation(fig1, update_line, total_frames, fargs=(sa, ed, l, title),
|
fig1.colorbar(ScalarMappable(norm=Normalize(vmin=-5, vmax=0), cmap=cmap), label="log(water fraction)")
|
||||||
interval=1000 / fps, repeat=False)
|
|
||||||
line_ani.save(str(fn.with_suffix(".mp4")), dpi=300)
|
plt.tight_layout()
|
||||||
# plt.show()
|
line_ani = animation.FuncAnimation(fig1, update_line, total_frames, fargs=(args, sa, ed, l, title),
|
||||||
|
interval=1000 / args.fps, repeat=False)
|
||||||
|
if args.save_video:
|
||||||
|
line_ani.save(str(fn.with_suffix(".mp4")), dpi=300)
|
||||||
|
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")
|
||||||
|
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)
|
||||||
|
|
10
water_sim.py
10
water_sim.py
|
@ -58,7 +58,13 @@ with open("initcon/conditions_many.input") as f:
|
||||||
if len(columns) > 7:
|
if len(columns) > 7:
|
||||||
# print(columns[7:])
|
# print(columns[7:])
|
||||||
cmf, mmf, wmf = columns[7:]
|
cmf, mmf, wmf = columns[7:]
|
||||||
assert cmf + mmf + wmf == 1
|
total_fractions = cmf + mmf + wmf
|
||||||
|
if total_fractions != 1:
|
||||||
|
diff = 1 - total_fractions
|
||||||
|
print(f"fractions don't add up by {diff}")
|
||||||
|
print("adding rest to cmf")
|
||||||
|
cmf += diff
|
||||||
|
assert cmf + mmf + wmf - 1 <= 1e-10
|
||||||
extradata.pdata[hash.value] = ParticleData(water_mass_fraction=wmf)
|
extradata.pdata[hash.value] = ParticleData(water_mass_fraction=wmf)
|
||||||
else:
|
else:
|
||||||
wmf = 0
|
wmf = 0
|
||||||
|
@ -113,7 +119,7 @@ for i, t in enumerate(times, start=1):
|
||||||
except Collision:
|
except Collision:
|
||||||
merge_particles(sim, extradata)
|
merge_particles(sim, extradata)
|
||||||
except Escape:
|
except Escape:
|
||||||
handle_escape(sim,extradata)
|
handle_escape(sim, extradata)
|
||||||
except NoParticles:
|
except NoParticles:
|
||||||
print("No Particles left")
|
print("No Particles left")
|
||||||
abort = True
|
abort = True
|
||||||
|
|
Loading…
Reference in a new issue