def update_plot(num: int, args: MyProgramArgs, sa: SimulationArchive,
ed: ExtraData, dots: PathCollection, title: Text):
global mean_mass
total_frames = args.fps * args.duration
if args.log_time:
log_timestep = (log10(ed.meta.current_time) -
log10(50000)) / total_frames
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(f"{num / total_frames:.2f}, {time:.0f}")
sim = sa.getSimulation(t=time)
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
water_fractions = []
for p in sim.particles[1:]:
pd: ParticleData = ed.pd(p)
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")
a = [p.a for p in sim.particles[1:]]
m = np.array([p.m for p in sim.particles[1:]])
m[:2] /= 1e2
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, [degrees(p.inc) for p in sim.particles[1:]]])
elif args.y_axis == "Omega":
bla = np.array([a, [degrees(p.Omega) for p in sim.particles[1:]]])
else:
raise ValueError("invalid y-axis")
dots.set_offsets(bla.T)
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(size_factor * m / mean_mass)
dots.set_color(colors)
if output_plots:
plt.savefig(f"tmp/{num}_{time / mega}.pdf", transparent=True)
return dots, title
