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