Added algo calc, spar_error and plot scripts
This commit is contained in:
413
src/plot.py
413
src/plot.py
@@ -2,73 +2,380 @@
|
||||
# Copyright (C) 2026 Association Exergie <association.exergie@gmail.com>
|
||||
# SPDX-License-Identifier: GPL-3.0-or-later
|
||||
|
||||
import pandas as pd
|
||||
from pathlib import Path
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
from tqdm import tqdm
|
||||
import sys
|
||||
from bisect import bisect_left
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import matplotlib.pyplot as plt
|
||||
import pandas as pd
|
||||
from matplotlib.ticker import MultipleLocator
|
||||
|
||||
parser = argparse.ArgumentParser()
|
||||
parser.add_argument("directory", type=Path, help="Source data directory")
|
||||
|
||||
args = parser.parse_args()
|
||||
GP_COLUMNS = ("gp0_falling", "gp1_falling")
|
||||
SPARK_COLUMNS = ("spark_1", "spark_2", "spark_3")
|
||||
PULSE_STYLES = {
|
||||
"gp0_falling": ("gp0 falling", "tab:green", 5.0, 0.18),
|
||||
"gp1_falling": ("gp1 falling", "tab:red", 5.0, 0.18),
|
||||
"spark_1": ("spark 1", "tab:blue", 2.4, 0.65),
|
||||
"spark_2": ("spark 2", "tab:orange", 2.4, 0.65),
|
||||
"spark_3": ("spark 3", "tab:purple", 2.4, 0.65),
|
||||
}
|
||||
|
||||
directory: Path = args.directory
|
||||
|
||||
if not directory.is_dir():
|
||||
parser.error(f"{directory} is not a valid directory")
|
||||
@dataclass(frozen=True)
|
||||
class EncoderPoint:
|
||||
time_us: int
|
||||
continuous_angle_deg: float
|
||||
|
||||
print(f"Processing data in: {directory}")
|
||||
|
||||
files: list[Path] = []
|
||||
def read_csv(path: Path) -> pd.DataFrame:
|
||||
df = pd.read_csv(path)
|
||||
if "time_us" not in df.columns:
|
||||
raise ValueError("missing column: time_us")
|
||||
df["time_us"] = df["time_us"].astype("int64")
|
||||
return df
|
||||
|
||||
for path in directory.glob("*.csv"):
|
||||
stem = path.stem
|
||||
|
||||
def encoder_points(df: pd.DataFrame) -> list[EncoderPoint]:
|
||||
required = {"time_us", "revolution_index", "crank_angle_deg"}
|
||||
if not required <= set(df.columns):
|
||||
return []
|
||||
|
||||
points: list[EncoderPoint] = []
|
||||
angle_df = df.dropna(subset=["revolution_index", "crank_angle_deg"])
|
||||
for row in angle_df.itertuples():
|
||||
revolution_index = int(row.revolution_index)
|
||||
if revolution_index < 0:
|
||||
continue
|
||||
continuous_angle = revolution_index * 360.0 + float(row.crank_angle_deg)
|
||||
points.append(EncoderPoint(int(row.time_us), continuous_angle))
|
||||
|
||||
deduped: dict[int, EncoderPoint] = {}
|
||||
for point in sorted(points, key=lambda item: item.time_us):
|
||||
deduped[point.time_us] = point
|
||||
return list(deduped.values())
|
||||
|
||||
|
||||
def interpolate_angle(points: list[EncoderPoint], time_us: int) -> float | None:
|
||||
if not points:
|
||||
return None
|
||||
|
||||
times = [point.time_us for point in points]
|
||||
index = bisect_left(times, time_us)
|
||||
if index < len(points) and points[index].time_us == time_us:
|
||||
return points[index].continuous_angle_deg
|
||||
if index == 0 or index == len(points):
|
||||
return None
|
||||
|
||||
before = points[index - 1]
|
||||
after = points[index]
|
||||
span = after.time_us - before.time_us
|
||||
if span <= 0:
|
||||
return before.continuous_angle_deg
|
||||
|
||||
fraction = (time_us - before.time_us) / span
|
||||
return before.continuous_angle_deg + fraction * (
|
||||
after.continuous_angle_deg - before.continuous_angle_deg
|
||||
)
|
||||
|
||||
|
||||
def iter_csv_paths(path: Path) -> list[Path]:
|
||||
if path.is_dir():
|
||||
return sorted(path.rglob("*.csv"))
|
||||
return [path]
|
||||
|
||||
|
||||
def has_angle_data(points: list[EncoderPoint]) -> bool:
|
||||
return len(points) >= 2
|
||||
|
||||
|
||||
def rpm_rows(df: pd.DataFrame) -> pd.DataFrame:
|
||||
if "rpm_raw" not in df.columns or "rpm_lpf" not in df.columns:
|
||||
return pd.DataFrame()
|
||||
return df.dropna(subset=["rpm_raw", "rpm_lpf"])
|
||||
|
||||
|
||||
def plot_time(df: pd.DataFrame, ax: plt.Axes) -> None:
|
||||
rpm_df = rpm_rows(df)
|
||||
if not rpm_df.empty:
|
||||
time_s = rpm_df["time_us"] / 1_000_000.0
|
||||
ax.plot(time_s, rpm_df["rpm_raw"], label="raw", linewidth=0.75, alpha=0.35)
|
||||
ax.plot(time_s, rpm_df["rpm_lpf"], label="lpf", linewidth=1.4)
|
||||
ax.set_xlabel("time (s)")
|
||||
|
||||
|
||||
def plot_angle_overlay(df: pd.DataFrame, ax: plt.Axes) -> None:
|
||||
rpm_df = rpm_rows(df)
|
||||
if rpm_df.empty:
|
||||
return
|
||||
|
||||
raw_label_used = False
|
||||
lpf_label_used = False
|
||||
for revolution in sorted(rpm_df["revolution_index"].dropna().astype(int).unique()):
|
||||
if revolution < 0:
|
||||
continue
|
||||
rev_df = rpm_df[rpm_df["revolution_index"].astype(int) == revolution]
|
||||
ax.plot(
|
||||
rev_df["crank_angle_deg"],
|
||||
rev_df["rpm_raw"],
|
||||
color="0.55",
|
||||
linewidth=0.5,
|
||||
alpha=0.12,
|
||||
label="raw" if not raw_label_used else None,
|
||||
)
|
||||
ax.plot(
|
||||
rev_df["crank_angle_deg"],
|
||||
rev_df["rpm_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(df: pd.DataFrame, ax: plt.Axes) -> None:
|
||||
rpm_df = rpm_rows(df)
|
||||
if rpm_df.empty:
|
||||
return
|
||||
|
||||
angle_df = rpm_df.dropna(subset=["revolution_index", "crank_angle_deg"])
|
||||
angle_df = angle_df[angle_df["revolution_index"].astype(int) >= 0]
|
||||
if angle_df.empty:
|
||||
return
|
||||
|
||||
continuous_angle = (
|
||||
angle_df["revolution_index"].astype(float) * 360.0
|
||||
+ angle_df["crank_angle_deg"].astype(float)
|
||||
)
|
||||
ax.plot(continuous_angle, angle_df["rpm_raw"], label="raw", linewidth=0.75, alpha=0.35)
|
||||
ax.plot(continuous_angle, angle_df["rpm_lpf"], label="lpf", linewidth=1.4)
|
||||
|
||||
max_angle = float(continuous_angle.max())
|
||||
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 selected_pulse_columns(
|
||||
df: pd.DataFrame,
|
||||
gp_mode: str,
|
||||
spark_columns: list[str],
|
||||
) -> list[str]:
|
||||
columns: list[str] = []
|
||||
if gp_mode in {"gp0", "both"} and "gp0_falling" in df.columns:
|
||||
columns.append("gp0_falling")
|
||||
if gp_mode in {"gp1", "both"} and "gp1_falling" in df.columns:
|
||||
columns.append("gp1_falling")
|
||||
columns.extend(column for column in spark_columns if column in df.columns)
|
||||
return columns
|
||||
|
||||
|
||||
def pulse_x_value(
|
||||
row: pd.Series,
|
||||
*,
|
||||
x_axis: str,
|
||||
points: list[EncoderPoint],
|
||||
) -> float | None:
|
||||
time_us = int(row["time_us"])
|
||||
if x_axis == "time":
|
||||
return time_us / 1_000_000.0
|
||||
|
||||
continuous_angle = interpolate_angle(points, time_us)
|
||||
if continuous_angle is None:
|
||||
return None
|
||||
if x_axis in {"angle", "angle-overlay"}:
|
||||
return continuous_angle % 360.0
|
||||
return continuous_angle
|
||||
|
||||
|
||||
def plot_pulse_bars(
|
||||
df: pd.DataFrame,
|
||||
ax: plt.Axes,
|
||||
*,
|
||||
x_axis: str,
|
||||
gp_mode: str,
|
||||
spark_columns: list[str],
|
||||
points: list[EncoderPoint],
|
||||
) -> None:
|
||||
for column in selected_pulse_columns(df, gp_mode, spark_columns):
|
||||
label, color, linewidth, alpha = PULSE_STYLES[column]
|
||||
pulse_df = df.loc[df[column].fillna(0).astype(int) == 1]
|
||||
plotted_label = False
|
||||
for _, row in pulse_df.iterrows():
|
||||
x_value = pulse_x_value(row, x_axis=x_axis, points=points)
|
||||
if x_value is None:
|
||||
continue
|
||||
ax.axvline(
|
||||
x_value,
|
||||
color=color,
|
||||
alpha=alpha,
|
||||
linewidth=linewidth,
|
||||
label=label if not plotted_label else None,
|
||||
)
|
||||
plotted_label = True
|
||||
|
||||
|
||||
def plot_file(
|
||||
csv_path: Path,
|
||||
output: Path | None,
|
||||
*,
|
||||
x_axis: str,
|
||||
gp_mode: str,
|
||||
spark_columns: list[str],
|
||||
) -> bool:
|
||||
try:
|
||||
base_name, channel = stem.rsplit("_", 1)
|
||||
except ValueError:
|
||||
print(f"Skipping badly named file: {path}")
|
||||
continue
|
||||
df = read_csv(csv_path)
|
||||
except (OSError, ValueError) as exc:
|
||||
print(f"Could not read {csv_path}: {exc}", file=sys.stderr)
|
||||
return False
|
||||
|
||||
if channel != "output":
|
||||
print(f"Skipping unknown file: {path}")
|
||||
continue
|
||||
points = encoder_points(df)
|
||||
if x_axis.startswith("angle") and not has_angle_data(points):
|
||||
print(f"Could not plot crank angle for {csv_path}: missing angle columns", file=sys.stderr)
|
||||
return False
|
||||
|
||||
files.append(path)
|
||||
fig, ax = plt.subplots(figsize=(12, 5))
|
||||
if x_axis in {"angle", "angle-overlay"}:
|
||||
plot_angle_overlay(df, ax)
|
||||
elif x_axis == "angle-continuous":
|
||||
plot_angle_continuous(df, ax)
|
||||
else:
|
||||
plot_time(df, ax)
|
||||
|
||||
|
||||
for file in files:
|
||||
output_df = pd.read_csv(file).set_index("time_us", drop=False)
|
||||
fig, ax = plt.subplots(figsize=(12, 6))
|
||||
|
||||
# RPM
|
||||
ax.plot(output_df["time_us"] / 1_000_000, output_df["crank"], label="crank")
|
||||
ax.plot(
|
||||
output_df["time_us"] / 1_000_000, output_df["cam"], label="cam", color="red"
|
||||
plot_pulse_bars(
|
||||
df,
|
||||
ax,
|
||||
x_axis=x_axis,
|
||||
gp_mode=gp_mode,
|
||||
spark_columns=spark_columns,
|
||||
points=points,
|
||||
)
|
||||
ax.plot(
|
||||
output_df["time_us"] / 1_000_000,
|
||||
output_df["spark_1"],
|
||||
label="spark_1",
|
||||
color="green",
|
||||
)
|
||||
ax.plot(
|
||||
output_df["time_us"] / 1_000_000,
|
||||
output_df["spark_2"],
|
||||
label="spark_2",
|
||||
color="yellow",
|
||||
)
|
||||
ax.plot(
|
||||
output_df["time_us"] / 1_000_000,
|
||||
output_df["spark_3"],
|
||||
label="spark_3",
|
||||
color="orange",
|
||||
)
|
||||
ax.set_xlabel("Time (s)")
|
||||
ax.set_ylim(0, 1)
|
||||
ax.grid(True)
|
||||
|
||||
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()
|
||||
plt.show()
|
||||
plt.close(fig)
|
||||
|
||||
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_spark_columns(value: str) -> list[str]:
|
||||
if value == "none":
|
||||
return []
|
||||
if value == "all":
|
||||
return list(SPARK_COLUMNS)
|
||||
|
||||
columns = [column.strip() for column in value.split(",") if column.strip()]
|
||||
if not columns:
|
||||
raise argparse.ArgumentTypeError("must specify at least one spark column")
|
||||
|
||||
invalid = sorted(set(columns) - set(SPARK_COLUMNS))
|
||||
if invalid:
|
||||
valid = ", ".join(("none", "all", *SPARK_COLUMNS))
|
||||
raise argparse.ArgumentTypeError(
|
||||
f"invalid spark column(s): {', '.join(invalid)}; choose from {valid}"
|
||||
)
|
||||
|
||||
deduped: list[str] = []
|
||||
for column in columns:
|
||||
if column not in deduped:
|
||||
deduped.append(column)
|
||||
return deduped
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description="Plot algo1 RPM output CSV files.")
|
||||
parser.add_argument("csv_path", type=Path, help="Algo1 CSV file or directory of 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="both",
|
||||
help="Overlay GP falling events as vertical bars.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--sparks",
|
||||
type=parse_spark_columns,
|
||||
default="all",
|
||||
help="Overlay spark events: none, all, one spark, or a comma-separated subset.",
|
||||
)
|
||||
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
|
||||
|
||||
spark_columns = (
|
||||
args.sparks if isinstance(args.sparks, list) else parse_spark_columns(args.sparks)
|
||||
)
|
||||
failures = 0
|
||||
for csv_path in csv_paths:
|
||||
output = output_path_for(csv_path, args.output, multiple)
|
||||
if not plot_file(
|
||||
csv_path,
|
||||
output,
|
||||
x_axis=args.x,
|
||||
gp_mode=args.gp,
|
||||
spark_columns=spark_columns,
|
||||
):
|
||||
failures += 1
|
||||
|
||||
return 1 if failures else 0
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
|
||||
Reference in New Issue
Block a user