yum/vfft

git clone https://git.yummers.dev/yum/vfft

yumFix broken axis in ratio plotb279c5c

master
branchesmasterall refs →
3.0 KiB115 linesraw
1#!/usr/bin/env -S uv run --script
2#
3# /// script
4# requires-python = ">=3.12"
5# dependencies = ["matplotlib", "PySide6"]
6# ///
7
8import argparse
9import csv
10import matplotlib
11matplotlib.use("QtAgg")
12import matplotlib.pyplot as plt
13from pathlib import Path
14
15def parse_args():
16    parser = argparse.ArgumentParser(description="Plot three CSV columns.")
17    parser.add_argument("csv_path", help="Path to the CSV file")
18    parser.add_argument("--x", default=0, help="Column for the x axis")
19    parser.add_argument("--y0", default=1, help="Column 0 for the y axis")
20    parser.add_argument("--y1", default=2, help="Column 1 for the y axis")
21    parser.add_argument("--save", action="store_true", help="Save the chart instead of displaying it")
22    args = parser.parse_args()
23    return args
24
25def to_ns(datum):
26    if datum.endswith(" ns"):
27        return float(datum[:-3])
28    if datum.endswith(" us"):
29        return float(datum[:-3]) * 1000
30    if datum.endswith(" ms"):
31        return float(datum[:-3]) * 1000 * 1000
32    # Assume it's already in ns
33    return float(datum)
34
35def get_basic_save_path(args):
36    return Path(args.csv_path).with_suffix(".png")
37
38def get_ratio_save_path(args):
39    path = Path(args.csv_path)
40    stem = path.stem
41    return Path(args.csv_path).with_stem(stem + "_ratio").with_suffix(".png")
42
43def parse_data(args):
44    data = ([], [], [])
45    with open(args.csv_path, newline="", encoding="utf-8") as file:
46      reader = (row for row in csv.reader(file))
47      header = next(reader)
48
49      for row in reader:
50          x = float(row[args.x])
51          y0 = to_ns(row[args.y0])
52          y1 = to_ns(row[args.y1])
53
54          data[0].append(x)
55          data[1].append(y0)
56          data[2].append(y1)
57
58    return data
59
60def plot_basic(args, data):
61    x, y0, y1 = data
62    fig, ax = plt.subplots()
63    ax.plot(x, y0, marker="o", label="rustfft")
64    ax.plot(x, y1, marker="o", label="vfft")
65    ax.set(xlabel="Input length", ylabel="Time (ns)")
66    ax.legend()
67    ax.grid(True, alpha=0.3)
68    fig.tight_layout()
69
70    if args.save:
71      fig.savefig(get_basic_save_path(args), dpi=150)
72    else:
73      plt.show()
74
75def plot_ratio(args, data):
76    x, y0, y1 = data
77    ratio = [ x1 / x0 for x0, x1 in zip(y0, y1) ]
78    fig, ax = plt.subplots()
79    ax.plot(x, ratio, marker="o", label="vfft / rustfft")
80    ax.set_xlim(left=0)
81    ax.set_ylim(bottom=0)
82    ax.axhline(y=1, color="red", linestyle="--")
83    ax.set(xlabel="Input length", ylabel="Execution time ratio (vfft / rustfft)")
84    ax.legend()
85    ax.grid(True, alpha=0.3)
86    fig.tight_layout()
87
88    if args.save:
89      fig.savefig(get_ratio_save_path(args), dpi=150)
90    else:
91      plt.show()
92
93def print_stats(data):
94    x, y0, y1 = data
95
96    num_wins = 0
97    for x0, x1 in zip(y0, y1):
98        if x1 < x0:
99            num_wins += 1
100
101    print(f"num trials: {len(x)}")
102    print(f"num wins:   {num_wins}")
103    print(f"win rate:   {num_wins / len(x)}")
104
105def main():
106    args = parse_args()
107    data = parse_data(args)
108    plot_basic(args, data)
109    plot_ratio(args, data)
110    print_stats(data)
111
112
113if __name__ == "__main__":
114    main()
115