324 lines
10 KiB
Python
324 lines
10 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import sys
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.ticker import MultipleLocator
|
|
|
|
|
|
REQUIRED_COLUMNS = {"time_us", "rpm_raw", "rpm_lpf"}
|
|
ANGLE_COLUMNS = {"revolution_index", "crank_angle_deg"}
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class GpEvent:
|
|
time_s: float
|
|
revolution_index: int
|
|
crank_angle_deg: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RpmData:
|
|
time_s: list[float]
|
|
revolution_index: list[int]
|
|
crank_angle_deg: list[float]
|
|
rpm_raw: list[float]
|
|
rpm_lpf: list[float]
|
|
turn_time_s: list[float]
|
|
gp0_falling: list[GpEvent]
|
|
gp1_falling: list[GpEvent]
|
|
|
|
|
|
def read_rpm(path: Path) -> RpmData:
|
|
time_s: list[float] = []
|
|
revolution_index: list[int] = []
|
|
crank_angle_deg: list[float] = []
|
|
rpm_raw: list[float] = []
|
|
rpm_lpf: list[float] = []
|
|
turn_time_s: list[float] = []
|
|
gp0_falling: list[GpEvent] = []
|
|
gp1_falling: list[GpEvent] = []
|
|
|
|
with path.open(newline="") as file:
|
|
reader = csv.DictReader(file)
|
|
fieldnames = set(reader.fieldnames or ())
|
|
missing = REQUIRED_COLUMNS - fieldnames
|
|
if missing:
|
|
raise ValueError(f"missing columns: {', '.join(sorted(missing))}")
|
|
|
|
has_turn = "turn" in fieldnames
|
|
has_angle = ANGLE_COLUMNS <= fieldnames
|
|
has_gp0 = "gp0_falling" in fieldnames
|
|
has_gp1 = "gp1_falling" in fieldnames
|
|
for row in reader:
|
|
sample_time_s = int(row["time_us"]) / 1_000_000
|
|
row_revolution_index = int(row["revolution_index"]) if has_angle and row["revolution_index"] else 0
|
|
row_crank_angle_deg = float(row["crank_angle_deg"]) if has_angle and row["crank_angle_deg"] else 0.0
|
|
|
|
if row["rpm_raw"] and row["rpm_lpf"]:
|
|
time_s.append(sample_time_s)
|
|
revolution_index.append(row_revolution_index)
|
|
crank_angle_deg.append(row_crank_angle_deg)
|
|
rpm_raw.append(float(row["rpm_raw"]))
|
|
rpm_lpf.append(float(row["rpm_lpf"]))
|
|
if has_turn and int(row["turn"]):
|
|
turn_time_s.append(sample_time_s)
|
|
if has_gp0 and int(row["gp0_falling"]):
|
|
gp0_falling.append(
|
|
GpEvent(
|
|
time_s=sample_time_s,
|
|
revolution_index=row_revolution_index,
|
|
crank_angle_deg=row_crank_angle_deg,
|
|
)
|
|
)
|
|
if has_gp1 and int(row["gp1_falling"]):
|
|
gp1_falling.append(
|
|
GpEvent(
|
|
time_s=sample_time_s,
|
|
revolution_index=row_revolution_index,
|
|
crank_angle_deg=row_crank_angle_deg,
|
|
)
|
|
)
|
|
|
|
return RpmData(
|
|
time_s=time_s,
|
|
revolution_index=revolution_index,
|
|
crank_angle_deg=crank_angle_deg,
|
|
rpm_raw=rpm_raw,
|
|
rpm_lpf=rpm_lpf,
|
|
turn_time_s=turn_time_s,
|
|
gp0_falling=gp0_falling,
|
|
gp1_falling=gp1_falling,
|
|
)
|
|
|
|
|
|
def iter_csv_paths(path: Path) -> list[Path]:
|
|
if path.is_dir():
|
|
return sorted(path.rglob("*.csv"))
|
|
return [path]
|
|
|
|
|
|
def plot_time(data: RpmData, ax: plt.Axes) -> None:
|
|
ax.plot(data.time_s, data.rpm_raw, label="raw", linewidth=0.75, alpha=0.35)
|
|
ax.plot(data.time_s, data.rpm_lpf, label="lpf", linewidth=1.4)
|
|
|
|
for turn_time in data.turn_time_s:
|
|
ax.axvline(turn_time, color="0.7", linewidth=0.5, alpha=0.35)
|
|
|
|
ax.set_xlabel("time (s)")
|
|
|
|
|
|
def selected_gp_events(data: RpmData, gp_mode: str) -> list[tuple[str, list[GpEvent], str]]:
|
|
selected: list[tuple[str, list[GpEvent], str]] = []
|
|
if gp_mode in {"gp0", "both"}:
|
|
selected.append(("gp0 falling", data.gp0_falling, "tab:green"))
|
|
if gp_mode in {"gp1", "both"}:
|
|
selected.append(("gp1 falling", data.gp1_falling, "tab:red"))
|
|
return selected
|
|
|
|
|
|
def plot_gp_bars(data: RpmData, ax: plt.Axes, *, x_axis: str, gp_mode: str) -> None:
|
|
for label, events, color in selected_gp_events(data, gp_mode):
|
|
plotted_label = False
|
|
for event in events:
|
|
if x_axis == "time":
|
|
x_value = event.time_s
|
|
elif x_axis in {"angle", "angle-overlay"}:
|
|
if event.revolution_index < 0:
|
|
continue
|
|
x_value = event.crank_angle_deg
|
|
else:
|
|
if event.revolution_index < 0:
|
|
continue
|
|
x_value = event.revolution_index * 360.0 + event.crank_angle_deg
|
|
|
|
ax.axvline(
|
|
x_value,
|
|
color=color,
|
|
alpha=0.18,
|
|
linewidth=5.0,
|
|
label=label if not plotted_label else None,
|
|
)
|
|
plotted_label = True
|
|
|
|
|
|
def has_angle_data(data: RpmData) -> bool:
|
|
return any(data.revolution_index) or any(data.crank_angle_deg)
|
|
|
|
|
|
def plot_angle_overlay(data: RpmData, ax: plt.Axes) -> None:
|
|
revolutions = sorted(set(data.revolution_index))
|
|
raw_label_used = False
|
|
lpf_label_used = False
|
|
|
|
for revolution in revolutions:
|
|
indices = [
|
|
index
|
|
for index, value in enumerate(data.revolution_index)
|
|
if value == revolution and value >= 0
|
|
]
|
|
if len(indices) < 2:
|
|
continue
|
|
|
|
angles = [data.crank_angle_deg[index] for index in indices]
|
|
raw = [data.rpm_raw[index] for index in indices]
|
|
lpf = [data.rpm_lpf[index] for index in indices]
|
|
|
|
ax.plot(
|
|
angles,
|
|
raw,
|
|
color="0.55",
|
|
linewidth=0.5,
|
|
alpha=0.12,
|
|
label="raw" if not raw_label_used else None,
|
|
)
|
|
ax.plot(
|
|
angles,
|
|
lpf,
|
|
linewidth=0.8,
|
|
alpha=0.28,
|
|
label="lpf per rev" if not lpf_label_used else None,
|
|
)
|
|
raw_label_used = True
|
|
lpf_label_used = True
|
|
|
|
ax.set_xlim(0, 360)
|
|
ax.set_xlabel("crank angle (deg)")
|
|
|
|
|
|
def plot_angle_continuous(data: RpmData, ax: plt.Axes) -> None:
|
|
indices = [
|
|
index
|
|
for index, revolution in enumerate(data.revolution_index)
|
|
if revolution >= 0
|
|
]
|
|
if not indices:
|
|
return
|
|
|
|
continuous_angle = [
|
|
data.revolution_index[index] * 360.0 + data.crank_angle_deg[index]
|
|
for index in indices
|
|
]
|
|
rpm_raw = [data.rpm_raw[index] for index in indices]
|
|
rpm_lpf = [data.rpm_lpf[index] for index in indices]
|
|
|
|
ax.plot(continuous_angle, rpm_raw, label="raw", linewidth=0.75, alpha=0.35)
|
|
ax.plot(continuous_angle, rpm_lpf, label="lpf", linewidth=1.4)
|
|
|
|
max_angle = max(continuous_angle)
|
|
turn_angle = 0.0
|
|
while turn_angle <= max_angle:
|
|
ax.axvline(turn_angle, color="0.7", linewidth=0.5, alpha=0.35)
|
|
turn_angle += 360.0
|
|
|
|
ax.set_xlim(0, max_angle)
|
|
ax.xaxis.set_major_locator(MultipleLocator(360.0))
|
|
ax.xaxis.set_minor_locator(MultipleLocator(180.0))
|
|
ax.set_xlabel("continuous crank angle (deg)")
|
|
|
|
|
|
def plot_rpm(csv_path: Path, output: Path | None, *, x_axis: str, gp_mode: str) -> bool:
|
|
try:
|
|
data = read_rpm(csv_path)
|
|
except (OSError, ValueError) as exc:
|
|
print(f"Could not read {csv_path}: {exc}", file=sys.stderr)
|
|
return False
|
|
|
|
if not data.time_s:
|
|
print(f"Skipping empty RPM file: {csv_path}", file=sys.stderr)
|
|
return True
|
|
if x_axis.startswith("angle") and not has_angle_data(data):
|
|
print(f"Could not plot crank angle for {csv_path}: missing angle columns", file=sys.stderr)
|
|
return False
|
|
|
|
fig, ax = plt.subplots(figsize=(12, 5))
|
|
if x_axis in {"angle", "angle-overlay"}:
|
|
plot_angle_overlay(data, ax)
|
|
elif x_axis == "angle-continuous":
|
|
plot_angle_continuous(data, ax)
|
|
else:
|
|
plot_time(data, ax)
|
|
plot_gp_bars(data, ax, x_axis=x_axis, gp_mode=gp_mode)
|
|
|
|
ax.set_ylabel("RPM")
|
|
ax.set_title(str(csv_path))
|
|
ax.grid(alpha=0.25)
|
|
if x_axis == "angle-continuous":
|
|
ax.grid(which="minor", axis="x", alpha=0.15, linewidth=0.6)
|
|
ax.legend()
|
|
fig.tight_layout()
|
|
|
|
if output is not None:
|
|
output.parent.mkdir(parents=True, exist_ok=True)
|
|
fig.savefig(output, dpi=160)
|
|
print(f"Wrote {output}")
|
|
plt.close(fig)
|
|
else:
|
|
plt.show()
|
|
plt.close(fig)
|
|
|
|
return True
|
|
|
|
|
|
def output_path_for(csv_path: Path, output: Path | None, multiple: bool) -> Path | None:
|
|
if output is None:
|
|
return None
|
|
if not multiple:
|
|
return output
|
|
return output / f"{csv_path.stem}.png"
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(description="Plot raw and filtered RPM CSV files.")
|
|
parser.add_argument("csv_path", type=Path, help="RPM CSV file or directory of RPM CSV files.")
|
|
parser.add_argument("-o", "--output", type=Path, help="Optional image output path or directory.")
|
|
parser.add_argument(
|
|
"--x",
|
|
choices=("time", "angle", "angle-overlay", "angle-continuous"),
|
|
default="time",
|
|
help="Plot RPM against time, overlaid 0-360 crank angle, or continuous crank angle.",
|
|
)
|
|
parser.add_argument(
|
|
"--gp",
|
|
choices=("none", "gp0", "gp1", "both"),
|
|
default="none",
|
|
help="Overlay GP falling events as translucent vertical bars.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
|
|
if not args.csv_path.exists():
|
|
print(f"Path does not exist: {args.csv_path}", file=sys.stderr)
|
|
return 1
|
|
|
|
csv_paths = iter_csv_paths(args.csv_path)
|
|
if not csv_paths:
|
|
print(f"No CSV files found under {args.csv_path}", file=sys.stderr)
|
|
return 1
|
|
|
|
multiple = len(csv_paths) > 1
|
|
if multiple and args.output is not None and args.output.suffix:
|
|
print("When plotting a directory, --output must be a directory.", file=sys.stderr)
|
|
return 1
|
|
|
|
failures = 0
|
|
for csv_path in csv_paths:
|
|
output = output_path_for(csv_path, args.output, multiple)
|
|
if multiple and output is None:
|
|
print(f"Plotting {csv_path}")
|
|
if not plot_rpm(csv_path, output, x_axis=args.x, gp_mode=args.gp):
|
|
failures += 1
|
|
|
|
return 1 if failures else 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|