diff --git a/python/lines.py b/python/lines.py new file mode 100755 index 0000000..96db8c1 --- /dev/null +++ b/python/lines.py @@ -0,0 +1,63 @@ +#!/usr/bin/env python + +import argparse as ap, matplotlib.animation as anm, matplotlib.pyplot as plt, numpy as np + +import matplotlib + +# create the parser +parser = ap.ArgumentParser( + prog="Acorn 1D Plotter", description="One-dimensional plotting script for the Quantum Acorn package.", + formatter_class=lambda prog: ap.HelpFormatter(prog, max_help_position=128), + add_help=False, allow_abbrev=False +) + +# add the +parser.add_argument("-h", "--help", action="help", default=ap.SUPPRESS, help="This help message.") +parser.add_argument("-o", "--output", type=str, help="The output file to save the plot.") + +# add the plotting arguments +parser.add_argument("--colormap", type=str, default="tab10", help="The colormap to use for the plot.") +parser.add_argument("--dpi", type=int, default=96, help="The DPI of the plot.") +parser.add_argument("--figsize", type=int, nargs=2, default=[6, 8], help="The dimensions of the plot in inches.") +parser.add_argument("--legend", type=str, nargs="+", help="Add a legend to the plot.") +parser.add_argument("--linewidth", type=float, default=1.8, help="The linewidth of the plot.") +parser.add_argument("--title", type=str, help="The title of the plot.") +parser.add_argument("--xlabel", type=str, help="The an x-axis label.") +parser.add_argument("--ylabel", type=str, help="The an y-axis label.") + +# add positional arguments and parse the arguments +parser.add_argument("files", nargs="+", help="The data files to plot."); args = parser.parse_args() + +# set the matplotlib parameters +matplotlib.rcParams["lines.linewidth"] = args.linewidth + +# function to load the data from the input string as well as the columns to plot +load = lambda s: (d := np.loadtxt((p := s.split(":", 1))[0], ndmin=2, skiprows=1), np.array(p[1].split(","), int) if len(p) == 2 and p[1] else np.arange(d.shape[1] - 1) + 1) + +# load the data with plotted columns and get the total number of lines +data, cols = zip(*[load(file) for file in args.files]); nline = sum(len(cols[i]) for i in range(len(cols))) + +# set the colormap +cmap = plt.get_cmap(args.colormap)(np.linspace(0.1, 0.9, nline) if (ncolor := plt.get_cmap(args.colormap).N) > 20 else np.linspace(0, 1, ncolor)) + +# create the figure +fig, ax = plt.subplots(dpi=args.dpi, figsize=(args.figsize[1], args.figsize[0])) + +# plot the data +plots = [ax.plot(data[:, 0], data[:, j], color=cmap[i]) for i, (data, j) in enumerate((data, j) for data, cols in zip(data, cols) for j in cols)] + +# add the legend +if args.legend: ax.legend(args.legend, frameon=False) + +# set the title +if args.title: ax.set_title(args.title) + +# set the axis labels +if args.xlabel: ax.set_xlabel(args.xlabel) +if args.ylabel: ax.set_ylabel(args.ylabel) + +# set the layout +fig.tight_layout() + +# save or show the plot +plt.savefig(args.output) if args.output else plt.show()