|
14 | 14 | # add the |
15 | 15 | parser.add_argument("-h", "--help", action="help", default=ap.SUPPRESS, help="This help message.") |
16 | 16 | parser.add_argument("-o", "--output", type=str, help="The output file to save the plot.") |
| 17 | +parser.add_argument("-a", "--animate", type=int, help="Animate the plot.") |
17 | 18 |
|
18 | 19 | # add the plotting arguments |
19 | 20 | parser.add_argument("--colormap", type=str, default="tab10", help="The colormap to use for the plot.") |
20 | 21 | parser.add_argument("--dpi", type=int, default=96, help="The DPI of the plot.") |
21 | 22 | parser.add_argument("--figsize", type=int, nargs=2, default=[6, 8], help="The dimensions of the plot in inches.") |
| 23 | +parser.add_argument("--fps", type=int, default=30, help="Frames per second for the animation.") |
22 | 24 | parser.add_argument("--legend", type=str, nargs="+", help="Add a legend to the plot.") |
23 | 25 | parser.add_argument("--title", type=str, help="The title of the plot.") |
24 | 26 | parser.add_argument("--xlabel", type=str, help="The an x-axis label.") |
|
47 | 49 | data_errors, cols_errors = zip(*[load(file) for file in args.errors[1:]]) if args.errors else ((), ()) |
48 | 50 |
|
49 | 51 | # get the line indices for the error bars |
50 | | -lind_errors = (list(map(lambda x: int(x), args.errors [0].split(","))) if args.errors [0] != "all" else np.arange(sum([len(col) for col in cols_errors]))) if args.errors else [] |
51 | | -lind_offsets = (list(map(lambda x: int(x), args.offsets[0].split(","))) if args.offsets[0] != "all" else np.arange(len(args.offsets) - 1 )) if args.offsets else [] |
52 | | -lind_scales = (list(map(lambda x: int(x), args.scales [0].split(","))) if args.scales [0] != "all" else np.arange(len(args.scales ) - 1 )) if args.scales else [] |
| 52 | +lind_errors = (list(map(lambda x: int(x), args.errors [0].split(","))) if args.errors [0] != "all" else range(sum([len(col) for col in cols_errors]))) if args.errors else [] |
| 53 | +lind_offsets = (list(map(lambda x: int(x), args.offsets[0].split(","))) if args.offsets[0] != "all" else range(len(args.offsets) - 1 )) if args.offsets else [] |
| 54 | +lind_scales = (list(map(lambda x: int(x), args.scales [0].split(","))) if args.scales [0] != "all" else range(len(args.scales ) - 1 )) if args.scales else [] |
53 | 55 |
|
54 | | -# add the scales and offsets to the data |
55 | | -for i, (data_i, j) in enumerate(dcit(data, data_cols)): data_i[:, j + 1] *= float(args.scales [lind_scales .index(i) + 1]) if i in lind_scales else 1 |
56 | | -for i, (data_i, j) in enumerate(dcit(data, data_cols)): data_i[:, j + 1] += float(args.offsets[lind_offsets.index(i) + 1]) if i in lind_offsets else 0 |
| 56 | +# add the scales and offsets to the initial lines in the data |
| 57 | +for (frame, step) in ((frame, args.animate if args.animate else 0) for frame in range((data[0].shape[1] - 1) // args.animate if args.animate else 1)): |
| 58 | + for i, (data_i, j) in enumerate(dcit(data, data_cols)): data_i[:, j + frame * step + 1] *= float(args.scales [lind_scales .index(i) + 1]) if i in lind_scales else 1 |
| 59 | +for (frame, step) in ((frame, args.animate if args.animate else 0) for frame in range((data[0].shape[1] - 1) // args.animate if args.animate else 1)): |
| 60 | + for i, (data_i, j) in enumerate(dcit(data, data_cols)): data_i[:, j + frame * step + 1] += float(args.offsets[lind_offsets.index(i) + 1]) if i in lind_offsets else 0 |
57 | 61 |
|
58 | 62 | # create the figure and the container for plots and error bars |
59 | 63 | fig, ax = plt.subplots(dpi=args.dpi, figsize=(args.figsize[1], args.figsize[0])); plots, ebars = [], [] |
|
78 | 82 | ebars.append(ax.fill_between(data_errors_i[:, 0], bot, top, color=cmap[lind_errors[len(ebars)] % len(cmap)], alpha=0.2)) |
79 | 83 |
|
80 | 84 | # set the default axis limits |
81 | | -ax.set_xlim(min([data_i[:, 0].min() for data_i in data]), max([data_i[:, 0].max() for data_i in data])) |
| 85 | +ax.set_xlim(min([data_i[:, 0 ].min() for data_i in data]), max([data_i[:, 0 ].max() for data_i in data])) |
| 86 | +ax.set_ylim(min([data_i[:, 1:].min() for data_i in data]), max([data_i[:, 1:].max() for data_i in data])) |
| 87 | + |
| 88 | +# enlarge the y limit by 5 percent |
| 89 | +ax.set_ylim(ax.get_ylim()[0] - 0.05 * np.diff(ax.get_ylim()), ax.get_ylim()[1] + 0.05 * np.diff(ax.get_ylim())) |
82 | 90 |
|
83 | 91 | # add the legend |
84 | 92 | if args.legend: ax.legend(args.legend, frameon=False) |
|
97 | 105 | # set the layout |
98 | 106 | fig.tight_layout() |
99 | 107 |
|
| 108 | +# update function |
| 109 | +def update(frame): |
| 110 | + |
| 111 | + # loop over the data and its columns to plot |
| 112 | + for i, (data_i, j) in enumerate(dcit(data, data_cols)): |
| 113 | + |
| 114 | + # update the data |
| 115 | + plots[i][0].set_ydata(data_i[:, j + frame * args.animate + 1]) |
| 116 | + |
| 117 | + # loop over the error bars and their columns |
| 118 | + for i, (data_errors_i, j) in enumerate(dcit(data_errors, cols_errors)): |
| 119 | + |
| 120 | + # update the error bars |
| 121 | + top = plots[lind_errors[i]][0].get_ydata() + data_errors_i[:, j + 1] |
| 122 | + bot = plots[lind_errors[i]][0].get_ydata() - data_errors_i[:, j + 1] |
| 123 | + |
| 124 | + # fill the area between the top and bottom line |
| 125 | + ebars[i].remove(); ebars[i] = ax.fill_between(data_errors_i[:, 0], bot, top, color=cmap[lind_errors[i] % len(cmap)], alpha=0.2) |
| 126 | + |
| 127 | +# create the animation |
| 128 | +anm = anm.FuncAnimation(fig, update, frames=(data[0].shape[1] - 1) // args.animate, init_func=lambda: None, interval=1000 // args.fps) if args.animate else None |
| 129 | + |
100 | 130 | # save or show the plot |
101 | | -plt.savefig(args.output) if args.output else plt.show() |
| 131 | +(anm.save(args.output, fps=args.fps) if args.animate else plt.savefig(args.output)) if args.output else plt.show() |
0 commit comments