From 7e67f22ca1c1d23d2937e3224745e30556516c5a Mon Sep 17 00:00:00 2001 From: David A Minton Date: Tue, 31 Jan 2023 15:51:53 -0500 Subject: [PATCH] Updated scattermovie to be compatible with updated numpy --- examples/Chambers2013/scattermovie.py | 117 +++++++++++++------------- 1 file changed, 57 insertions(+), 60 deletions(-) diff --git a/examples/Chambers2013/scattermovie.py b/examples/Chambers2013/scattermovie.py index 2ef59c867..09b29072a 100755 --- a/examples/Chambers2013/scattermovie.py +++ b/examples/Chambers2013/scattermovie.py @@ -4,6 +4,7 @@ import matplotlib.pyplot as plt from matplotlib import animation import matplotlib.colors as mcolors +from collections import namedtuple plt.switch_backend('agg') titletext = "Chambers (2013)" @@ -22,6 +23,8 @@ framejump = 1 animation_file = f"Chambers2013-{plot_style}.mp4" +origin_types = ["Initial conditions", "Merger", "Disruption", "Supercatastrophic", "Hit and run fragmentation"] + class AnimatedScatter(object): """An animated scatter plot using matplotlib.animations.FuncAnimation.""" def __init__(self, ds, param): @@ -31,91 +34,85 @@ def __init__(self, ds, param): self.ds = ds self.param = param self.Rcb = self.ds['radius'].sel(name="Sun").isel(time=0).values[()] + colors = ["k", "xkcd:faded blue", "xkcd:marigold", "xkcd:shocking pink", "xkcd:baby poop green"] + self.clist = dict(zip(origin_types,colors)) - self.clist = {'Initial conditions' : 'k', - 'Merger' : 'xkcd:faded blue', - 'Disruption' : 'xkcd:marigold', - 'Supercatastrophic' : 'xkcd:shocking pink', - 'Hit and run fragmentation' : 'xkcd:baby poop green'} # Setup the figure and axes... fig = plt.figure(figsize=(8,4.5), dpi=300) plt.tight_layout(pad=0) # set up the figure self.ax = plt.Axes(fig, [0.1, 0.15, 0.8, 0.75]) - self.ax.set_xlim(xlim[plot_style]) - self.ax.set_ylim(ylim[plot_style]) fig.add_axes(self.ax) + + self.make_artists() + + self.ani = animation.FuncAnimation(fig, func=self.update, interval=1, frames=nframes, init_func=self.init_plot, blit=True) + self.ani.save(animation_file, fps=60, dpi=300, extra_args=['-vcodec', 'libx264']) + print(f'Finished writing {animation_file}') + + def make_artists(self): + scatter_names = [f"s{i}" for i,k in enumerate(origin_types)] + self.scatter_artist_names = dict(zip(origin_types,scatter_names)) + + animated_elements = [self.ax.text(0.50, 1.05, "", bbox={'facecolor': 'w', 'alpha': 0.5, 'pad': 5}, transform=self.ax.transAxes,ha="center", animated=True)] + element_names = ["title"] + for key, value in self.clist.items(): + animated_elements.append(self.ax.scatter([], [], marker='o', s=[], c=value, alpha=0.75, label=key, animated=True)) + element_names.append(self.scatter_artist_names[key]) - # First frame - """Initial drawing of the scatter plot.""" - t, npl, pl, radmarker, origin = next(self.data_stream(0)) + Artists = namedtuple("Artists",tuple(element_names)) + self.artists = Artists(*animated_elements) + return + def init_plot(self): + self.ax.set_xlim(xlim[plot_style]) + self.ax.set_ylim(ylim[plot_style]) + # set up the figure self.ax.margins(x=10, y=1) self.ax.set_xlabel(xlabel[plot_style], fontsize='16', labelpad=1) - self.ax.set_ylabel(ylabel[plot_style], fontsize='16', labelpad=1) - - title = self.ax.text(0.50, 1.05, "", bbox={'facecolor': 'w', 'alpha': 0.5, 'pad': 5}, transform=self.ax.transAxes, - ha="center", animated=True) - - title.set_text(f"{titletext} - Time = ${t*1e-6:6.2f}$ My with ${npl:4.0f}$ particles") - self.slist = self.scatters(pl, radmarker, origin) - self.slist.append(title) + self.ax.set_ylabel(ylabel[plot_style], fontsize='16', labelpad=1) leg = plt.legend(loc="upper left", scatterpoints=1, fontsize=10) for i,l in enumerate(leg.legendHandles): - leg.legendHandles[i]._sizes = [20] + leg.legendHandles[i]._sizes = [20] - self.ani = animation.FuncAnimation(fig, self.update, interval=1, frames=nframes, init_func=self.setup_plot, blit=True) - self.ani.save(animation_file, fps=60, dpi=300, extra_args=['-vcodec', 'libx264']) - print(f'Finished writing {animation_file}') + return self.artists + + def get_data(self, frame=0): + d = self.ds.isel(time = frame) + n=len(d['name']) + d = d.isel(name=range(1,n)) + d['radmarker'] = (d['radius'] / self.Rcb) * self.radscale + + t = d['time'].values + npl = d['npl'].values + radmarker = d['radmarker'].values + origin = d['origin_type'].values + + if plot_style == "aescatter": + pl = np.c_[d['a'].values,d['e'].values] + elif plot_style == "aiscatter": + pl = np.c_[d['a'].values,d['inc'].values] - def scatters(self, pl, radmarker, origin): - scat = [] - for key, value in self.clist.items(): - idx = origin == key - s = self.ax.scatter(pl[idx, 0], pl[idx, 1], marker='o', s=radmarker[idx], c=value, alpha=0.75, label=key, animated=True) - scat.append(s) - return scat - - def setup_plot(self): - return self.slist - - def data_stream(self, frame=0): - while True: - d = self.ds.isel(time = frame) - n=len(d['name']) - d = d.isel(name=range(1,n)) - d['radmarker'] = (d['radius'] / self.Rcb) * self.radscale - - t = d['time'].values - npl = d['npl'].values - radmarker = d['radmarker'].values - origin = d['origin_type'].values - - if plot_style == "aescatter": - pl = np.c_[d['a'].values,d['e'].values] - elif plot_style == "aiscatter": - pl = np.c_[d['a'].values,d['inc'].values] - - yield t, npl, pl, radmarker, origin + return t, npl, pl, radmarker, origin def update(self,frame): """Update the scatter plot.""" - t, npl, pl, radmarker, origin = next(self.data_stream(framejump * frame)) + t, npl, pl, radmarker, origin = self.get_data(framejump * frame) - self.slist[-1].set_text(f"{titletext} - Time = ${t*1e-6:6.3f}$ My with ${npl:4.0f}$ particles") + self.artists.title.set_text(f"{titletext} - Time = ${t*1e-6:6.3f}$ My with ${npl:4.0f}$ particles") - # We need to return the updated artist for FuncAnimation to draw.. - # Note that it expects a sequence of artists, thus the trailing comma. - for i, (key, value) in enumerate(self.clist.items()): + for key,name in self.scatter_artist_names.items(): idx = origin == key - self.slist[i].set_sizes(radmarker[idx]) - self.slist[i].set_offsets(pl[idx,:]) - self.slist[i].set_facecolor(value) + if any(idx) and any(~np.isnan(radmarker[idx])): + scatter = self.artists._asdict()[name] + scatter.set_sizes(radmarker[idx]) + scatter.set_offsets(pl[idx,:]) + scatter.set_facecolor(self.clist[key]) - return self.slist + return self.artists sim = swiftest.Simulation(read_data=True) print('Making animation')