470 lines
14 KiB
Python
470 lines
14 KiB
Python
from __future__ import annotations
|
|
|
|
import argparse
|
|
import csv
|
|
import statistics
|
|
import sys
|
|
from bisect import bisect_right
|
|
from dataclasses import dataclass
|
|
from pathlib import Path
|
|
|
|
|
|
INPUT_COLUMNS = {"time_us", "pulse"}
|
|
GP_COLUMNS = ("gp0_falling", "gp1_falling")
|
|
OUTPUT_COLUMNS = (
|
|
"time_us",
|
|
"pulse_index",
|
|
"revolution_index",
|
|
"crank_angle_deg",
|
|
"turn",
|
|
*GP_COLUMNS,
|
|
"rpm_raw",
|
|
"rpm_lpf",
|
|
)
|
|
DEFAULT_PPR = 256
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class Events:
|
|
pulse_times: list[int]
|
|
turn_times: list[int]
|
|
gp0_falling: list[int]
|
|
gp1_falling: list[int]
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class RpmSample:
|
|
time_us: int
|
|
pulse_index: int
|
|
revolution_index: int
|
|
crank_angle_deg: float
|
|
turn: int
|
|
rpm_raw: float
|
|
rpm_lpf: float
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class OutputRow:
|
|
time_us: int
|
|
pulse_index: int | None
|
|
revolution_index: int | None
|
|
crank_angle_deg: float | None
|
|
turn: int
|
|
gp0_falling: int
|
|
gp1_falling: int
|
|
rpm_raw: float | None
|
|
rpm_lpf: float | None
|
|
|
|
|
|
def read_events(path: Path) -> Events:
|
|
pulse_times: list[int] = []
|
|
turn_times: list[int] = []
|
|
gp0_falling: list[int] = []
|
|
gp1_falling: list[int] = []
|
|
|
|
with path.open(newline="") as file:
|
|
reader = csv.DictReader(file)
|
|
fieldnames = set(reader.fieldnames or ())
|
|
missing = INPUT_COLUMNS - fieldnames
|
|
if missing:
|
|
raise ValueError(f"missing columns: {', '.join(sorted(missing))}")
|
|
|
|
has_turn = "turn" in fieldnames
|
|
for row in reader:
|
|
time_us = int(row["time_us"])
|
|
if int(row["pulse"]):
|
|
pulse_times.append(time_us)
|
|
if has_turn and int(row["turn"]):
|
|
turn_times.append(time_us)
|
|
if "gp0_falling" in fieldnames and int(row["gp0_falling"]):
|
|
gp0_falling.append(time_us)
|
|
if "gp1_falling" in fieldnames and int(row["gp1_falling"]):
|
|
gp1_falling.append(time_us)
|
|
|
|
return Events(
|
|
pulse_times=sorted(pulse_times),
|
|
turn_times=sorted(turn_times),
|
|
gp0_falling=sorted(gp0_falling),
|
|
gp1_falling=sorted(gp1_falling),
|
|
)
|
|
|
|
|
|
def median_or_none(values: list[float]) -> float | None:
|
|
if not values:
|
|
return None
|
|
return float(statistics.median(values))
|
|
|
|
|
|
def despike(values: list[float], *, window: int, sigma: float) -> tuple[list[float], int]:
|
|
if len(values) < 3:
|
|
return values[:], 0
|
|
|
|
radius = max(1, window // 2)
|
|
cleaned = values[:]
|
|
replaced = 0
|
|
|
|
for index, value in enumerate(values):
|
|
start = max(0, index - radius)
|
|
end = min(len(values), index + radius + 1)
|
|
local = values[start:end]
|
|
median = median_or_none(local)
|
|
if median is None:
|
|
continue
|
|
|
|
deviations = [abs(sample - median) for sample in local]
|
|
mad = median_or_none(deviations) or 0.0
|
|
threshold = sigma * 1.4826 * mad
|
|
|
|
if threshold <= 0.0:
|
|
threshold = max(1.0, abs(median) * 0.25)
|
|
|
|
if abs(value - median) > threshold:
|
|
cleaned[index] = median
|
|
replaced += 1
|
|
|
|
return cleaned, replaced
|
|
|
|
|
|
def ema(values: list[float], alpha: float) -> list[float]:
|
|
if not values:
|
|
return []
|
|
|
|
filtered = [values[0]]
|
|
previous = values[0]
|
|
for value in values[1:]:
|
|
previous = alpha * value + (1.0 - alpha) * previous
|
|
filtered.append(previous)
|
|
|
|
return filtered
|
|
|
|
|
|
def nearest_index(times: list[int], target: int) -> int | None:
|
|
if not times:
|
|
return None
|
|
|
|
lo = 0
|
|
hi = len(times)
|
|
while lo < hi:
|
|
mid = (lo + hi) // 2
|
|
if times[mid] < target:
|
|
lo = mid + 1
|
|
else:
|
|
hi = mid
|
|
|
|
candidates = []
|
|
if lo < len(times):
|
|
candidates.append(lo)
|
|
if lo > 0:
|
|
candidates.append(lo - 1)
|
|
|
|
return min(candidates, key=lambda index: abs(times[index] - target))
|
|
|
|
|
|
def robust_median_interval(times: list[int]) -> float | None:
|
|
intervals = [b - a for a, b in zip(times, times[1:]) if b > a]
|
|
if not intervals:
|
|
return None
|
|
|
|
intervals = sorted(intervals)
|
|
if len(intervals) >= 20:
|
|
trim = len(intervals) // 20
|
|
intervals = intervals[trim : len(intervals) - trim]
|
|
|
|
return float(statistics.median(intervals))
|
|
|
|
|
|
def choose_turn_phase(pulse_times: list[int], turn_times: list[int], ppr: int) -> int:
|
|
if not pulse_times or not turn_times:
|
|
return 0
|
|
|
|
pulse_interval = robust_median_interval(pulse_times)
|
|
if pulse_interval is None:
|
|
return 0
|
|
|
|
max_distance_us = max(1_000, int(pulse_interval * 10))
|
|
scores = [0.0] * ppr
|
|
|
|
for turn_time in turn_times:
|
|
pulse_index = nearest_index(pulse_times, turn_time)
|
|
if pulse_index is None:
|
|
continue
|
|
distance = abs(pulse_times[pulse_index] - turn_time)
|
|
if distance <= max_distance_us:
|
|
scores[pulse_index % ppr] += 1.0 - distance / max_distance_us
|
|
|
|
best_score = max(scores)
|
|
if best_score <= 0.0:
|
|
return 0
|
|
|
|
return scores.index(best_score)
|
|
|
|
|
|
def revolution_index_for(pulse_index: int, phase: int, ppr: int) -> int:
|
|
return (pulse_index - phase) // ppr
|
|
|
|
|
|
def calculate_rpm(
|
|
events: Events,
|
|
*,
|
|
ppr: int,
|
|
filter_window: int,
|
|
filter_alpha: float,
|
|
hampel_sigma: float,
|
|
) -> tuple[list[RpmSample], int]:
|
|
raw_rows: list[tuple[int, int, int, float, int, float]] = []
|
|
pulse_times = events.pulse_times
|
|
phase = choose_turn_phase(pulse_times, events.turn_times, ppr)
|
|
degrees_per_pulse = 360.0 / ppr
|
|
|
|
for pulse_index, (previous_time, current_time) in enumerate(zip(pulse_times, pulse_times[1:]), start=1):
|
|
delta_us = current_time - previous_time
|
|
if delta_us <= 0:
|
|
continue
|
|
|
|
rpm = 60_000_000.0 / (delta_us * ppr)
|
|
angle_pulse = (pulse_index - phase) % ppr
|
|
angle_deg = angle_pulse * degrees_per_pulse
|
|
revolution_index = revolution_index_for(pulse_index, phase, ppr)
|
|
turn = 1 if angle_pulse == 0 else 0
|
|
raw_rows.append((current_time, pulse_index, revolution_index, angle_deg, turn, rpm))
|
|
|
|
raw_rpm = [row[5] for row in raw_rows]
|
|
despiked_rpm, replaced = despike(raw_rpm, window=filter_window, sigma=hampel_sigma)
|
|
filtered_rpm = ema(despiked_rpm, filter_alpha)
|
|
|
|
samples = [
|
|
RpmSample(
|
|
time_us=time_us,
|
|
pulse_index=pulse_index,
|
|
revolution_index=revolution_index,
|
|
crank_angle_deg=angle_deg,
|
|
turn=turn,
|
|
rpm_raw=rpm_raw,
|
|
rpm_lpf=rpm_lpf,
|
|
)
|
|
for (time_us, pulse_index, revolution_index, angle_deg, turn, rpm_raw), rpm_lpf in zip(
|
|
raw_rows,
|
|
filtered_rpm,
|
|
)
|
|
]
|
|
|
|
return samples, replaced
|
|
|
|
|
|
def pulse_position_for_time(
|
|
pulse_times: list[int],
|
|
time_us: int,
|
|
*,
|
|
phase: int,
|
|
ppr: int,
|
|
) -> tuple[int, int, float]:
|
|
pulse_index = bisect_right(pulse_times, time_us) - 1
|
|
if pulse_index < 0:
|
|
pulse_index = 0
|
|
|
|
angle_pulse = (pulse_index - phase) % ppr
|
|
revolution_index = revolution_index_for(pulse_index, phase, ppr)
|
|
angle_deg = angle_pulse * (360.0 / ppr)
|
|
|
|
return pulse_index, revolution_index, angle_deg
|
|
|
|
|
|
def build_output_rows(
|
|
events: Events,
|
|
samples: list[RpmSample],
|
|
*,
|
|
ppr: int,
|
|
) -> list[OutputRow]:
|
|
phase = choose_turn_phase(events.pulse_times, events.turn_times, ppr)
|
|
rows: dict[tuple[int, int], OutputRow] = {}
|
|
|
|
for sample in samples:
|
|
rows[(sample.time_us, 0)] = OutputRow(
|
|
time_us=sample.time_us,
|
|
pulse_index=sample.pulse_index,
|
|
revolution_index=sample.revolution_index,
|
|
crank_angle_deg=sample.crank_angle_deg,
|
|
turn=sample.turn,
|
|
gp0_falling=0,
|
|
gp1_falling=0,
|
|
rpm_raw=sample.rpm_raw,
|
|
rpm_lpf=sample.rpm_lpf,
|
|
)
|
|
|
|
for channel_index, gp_times in enumerate((events.gp0_falling, events.gp1_falling), start=1):
|
|
for gp_time in gp_times:
|
|
pulse_index, revolution_index, angle_deg = pulse_position_for_time(
|
|
events.pulse_times,
|
|
gp_time,
|
|
phase=phase,
|
|
ppr=ppr,
|
|
)
|
|
existing_key = (gp_time, 0)
|
|
if existing_key in rows:
|
|
existing = rows[existing_key]
|
|
rows[existing_key] = OutputRow(
|
|
time_us=existing.time_us,
|
|
pulse_index=existing.pulse_index,
|
|
revolution_index=existing.revolution_index,
|
|
crank_angle_deg=existing.crank_angle_deg,
|
|
turn=existing.turn,
|
|
gp0_falling=1 if channel_index == 1 else existing.gp0_falling,
|
|
gp1_falling=1 if channel_index == 2 else existing.gp1_falling,
|
|
rpm_raw=existing.rpm_raw,
|
|
rpm_lpf=existing.rpm_lpf,
|
|
)
|
|
continue
|
|
|
|
rows[(gp_time, channel_index)] = OutputRow(
|
|
time_us=gp_time,
|
|
pulse_index=pulse_index,
|
|
revolution_index=revolution_index,
|
|
crank_angle_deg=angle_deg,
|
|
turn=0,
|
|
gp0_falling=1 if channel_index == 1 else 0,
|
|
gp1_falling=1 if channel_index == 2 else 0,
|
|
rpm_raw=None,
|
|
rpm_lpf=None,
|
|
)
|
|
|
|
return [rows[key] for key in sorted(rows)]
|
|
|
|
|
|
def write_rpm(path: Path, rows: list[OutputRow]) -> None:
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
with path.open("w", newline="") as file:
|
|
writer = csv.writer(file)
|
|
writer.writerow(OUTPUT_COLUMNS)
|
|
for row in rows:
|
|
writer.writerow(
|
|
(
|
|
row.time_us,
|
|
"" if row.pulse_index is None else row.pulse_index,
|
|
"" if row.revolution_index is None else row.revolution_index,
|
|
"" if row.crank_angle_deg is None else f"{row.crank_angle_deg:.6f}",
|
|
row.turn,
|
|
row.gp0_falling,
|
|
row.gp1_falling,
|
|
"" if row.rpm_raw is None else f"{row.rpm_raw:.3f}",
|
|
"" if row.rpm_lpf is None else f"{row.rpm_lpf:.3f}",
|
|
)
|
|
)
|
|
|
|
|
|
def output_path_for(input_path: Path, input_root: Path, output_root: Path) -> Path:
|
|
relative = input_path.relative_to(input_root)
|
|
stem = relative.stem
|
|
if stem.endswith("_cleaned"):
|
|
stem = stem[: -len("_cleaned")]
|
|
return output_root / relative.with_name(f"{stem}_rpm.csv")
|
|
|
|
|
|
def process_file(
|
|
input_path: Path,
|
|
output_path: Path,
|
|
*,
|
|
ppr: int,
|
|
filter_window: int,
|
|
filter_alpha: float,
|
|
hampel_sigma: float,
|
|
) -> tuple[int, int]:
|
|
events = read_events(input_path)
|
|
samples, replaced = calculate_rpm(
|
|
events,
|
|
ppr=ppr,
|
|
filter_window=filter_window,
|
|
filter_alpha=filter_alpha,
|
|
hampel_sigma=hampel_sigma,
|
|
)
|
|
rows = build_output_rows(events, samples, ppr=ppr)
|
|
write_rpm(output_path, rows)
|
|
return len(samples), replaced
|
|
|
|
|
|
def parse_args() -> argparse.Namespace:
|
|
parser = argparse.ArgumentParser(
|
|
description="Convert cleaned pulse event CSVs into raw and filtered RPM CSVs."
|
|
)
|
|
parser.add_argument(
|
|
"--input-root",
|
|
type=Path,
|
|
default=Path("recordings_cleaned"),
|
|
help="Folder to search recursively for cleaned CSV files.",
|
|
)
|
|
parser.add_argument(
|
|
"--output-root",
|
|
type=Path,
|
|
default=Path("recordings_rpm"),
|
|
help="Folder where RPM CSV files are written.",
|
|
)
|
|
parser.add_argument("--ppr", type=int, default=DEFAULT_PPR, help="Pulse encoder pulses per revolution.")
|
|
parser.add_argument(
|
|
"--filter-window",
|
|
type=int,
|
|
default=5,
|
|
help="Odd sample count for local despiking before low-pass filtering.",
|
|
)
|
|
parser.add_argument(
|
|
"--filter-alpha",
|
|
type=float,
|
|
default=0.45,
|
|
help="EMA alpha for the low-pass RPM column. Higher follows transients faster.",
|
|
)
|
|
parser.add_argument(
|
|
"--hampel-sigma",
|
|
type=float,
|
|
default=4.0,
|
|
help="Local median absolute deviation threshold for replacing single-sample outliers.",
|
|
)
|
|
return parser.parse_args()
|
|
|
|
|
|
def main() -> int:
|
|
args = parse_args()
|
|
|
|
if not 0.0 < args.filter_alpha <= 1.0:
|
|
print("--filter-alpha must be in the range (0, 1].", file=sys.stderr)
|
|
return 1
|
|
if args.filter_window < 3:
|
|
print("--filter-window must be at least 3.", file=sys.stderr)
|
|
return 1
|
|
|
|
input_root = args.input_root
|
|
output_root = args.output_root
|
|
if not input_root.exists():
|
|
print(f"Input folder does not exist: {input_root}", file=sys.stderr)
|
|
return 1
|
|
|
|
csv_paths = sorted(path for path in input_root.rglob("*.csv") if not path.name.endswith("_rpm.csv"))
|
|
if not csv_paths:
|
|
print(f"No CSV files found under {input_root}", file=sys.stderr)
|
|
return 1
|
|
|
|
failures = 0
|
|
for input_path in csv_paths:
|
|
output_path = output_path_for(input_path, input_root, output_root)
|
|
try:
|
|
samples, replaced = process_file(
|
|
input_path,
|
|
output_path,
|
|
ppr=args.ppr,
|
|
filter_window=args.filter_window,
|
|
filter_alpha=args.filter_alpha,
|
|
hampel_sigma=args.hampel_sigma,
|
|
)
|
|
except (OSError, ValueError) as exc:
|
|
failures += 1
|
|
print(f"Skipping {input_path}: {exc}", file=sys.stderr)
|
|
continue
|
|
|
|
print(
|
|
f"{input_path} -> {output_path} "
|
|
f"samples={samples} despiked={replaced}"
|
|
)
|
|
|
|
return 1 if failures else 0
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sys.exit(main())
|