new file: checkpoints_tree/task1_tr_tree/best_tr_tree.pkl

new file:   checkpoints_tree/task1_tr_tree/model_selection.csv
	new file:   evaluation_outputs/task1_tr_tree/evaluation_train_all_samples.csv
	new file:   evaluation_outputs/task1_tr_tree/evaluation_train_curve.png
	new file:   evaluation_outputs/task1_tr_tree/evaluation_val_all_samples.csv
	new file:   evaluation_outputs/task1_tr_tree/evaluation_val_curve.png
	new file:   evaluation_outputs/task1_tr_tree/harmonic_5mm_1.55Hz_prediction.png
	new file:   scripts_tree/__pycache__/config.cpython-310.pyc
	new file:   scripts_tree/config.py
	new file:   scripts_tree/evaluate.py
	new file:   scripts_tree/predict_single.py
	new file:   scripts_tree/train.py
This commit is contained in:
2026-05-06 12:32:49 +08:00
parent 484643409d
commit 02e488cd0d
12 changed files with 549 additions and 0 deletions

View File

@@ -0,0 +1,115 @@
from __future__ import annotations
import argparse
import pickle
from pathlib import Path
import sys
import matplotlib.pyplot as plt
import numpy as np
PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
sys.path.insert(0, str(PROJECT_ROOT))
from scripts.dataset import build_record_from_file
try:
from .config import FEATURE_NAMES, evaluation_dir, make_experiment_config
except ImportError:
from config import FEATURE_NAMES, evaluation_dir, make_experiment_config
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description="Predict task1 RMS from a single waveform CSV with tree model.")
parser.add_argument("--file", type=str, required=True)
parser.add_argument("--checkpoint", type=str, default=None)
return parser.parse_args()
def resolve_checkpoint_path(project_root: Path, checkpoint_arg: str | None) -> Path:
if checkpoint_arg:
return Path(checkpoint_arg).resolve()
return project_root / "checkpoints_tree" / "task1_tr_tree" / "best_tr_tree.pkl"
def save_prediction_figure(file_path: Path, record, pred_rms: float, true_rms: float | None, output_dir: Path) -> Path:
time_middle = record.time_middle.detach().cpu().numpy().reshape(-1)
x_middle = record.x_middle.detach().cpu().numpy().reshape(-1)
sampling_rate = record.sampling_rate
fft_values = np.fft.rfft(x_middle)
freqs = np.fft.rfftfreq(x_middle.size, d=1.0 / sampling_rate)
amplitudes = np.abs(fft_values)
fig, axes = plt.subplots(3, 1, figsize=(12, 10))
fig.suptitle(f"Task1 TR Tree Single-File Prediction | {file_path.name}")
axes[0].plot(time_middle, x_middle, color="tab:blue")
axes[0].set_title("Input Base Excitation (Middle Segment)")
axes[0].set_xlabel("Time")
axes[0].set_ylabel("Acceleration")
axes[0].grid(True, alpha=0.3)
axes[1].plot(freqs, amplitudes, color="tab:purple")
axes[1].axvline(record.frequency_hz, color="tab:red", linestyle="--", label=f"Dominant freq = {record.frequency_hz:.4f} Hz")
axes[1].set_xlim(0.0, 5.0)
axes[1].set_title("Input Spectrum")
axes[1].set_xlabel("Frequency (Hz)")
axes[1].set_ylabel("Amplitude")
axes[1].grid(True, alpha=0.3)
axes[1].legend()
labels = ["Pred RMS"] if true_rms is None else ["True RMS", "Pred RMS"]
values = [pred_rms] if true_rms is None else [true_rms, pred_rms]
colors = ["tab:orange"] if true_rms is None else ["tab:green", "tab:orange"]
axes[2].bar(labels, values, color=colors)
title = f"Predicted RMS = {pred_rms:.6f}"
if true_rms is not None:
error_percent = abs(pred_rms - true_rms) / max(abs(true_rms), 1e-12) * 100.0
title = f"True RMS = {true_rms:.6f} | Pred RMS = {pred_rms:.6f} | Error = {error_percent:.2f}%"
axes[2].set_title(title)
axes[2].set_ylabel("RMS")
axes[2].grid(True, axis="y", alpha=0.3)
plt.tight_layout()
output_dir.mkdir(parents=True, exist_ok=True)
figure_path = output_dir / f"{file_path.stem}_prediction.png"
plt.savefig(figure_path, dpi=180, bbox_inches="tight")
plt.close(fig)
return figure_path
def main() -> None:
args = parse_args()
config = make_experiment_config()
ckpt_path = resolve_checkpoint_path(config.data.project_root, args.checkpoint)
with ckpt_path.open("rb") as handle:
bundle = pickle.load(handle)
model = bundle["model"]
file_path = Path(args.file).resolve()
record = build_record_from_file(file_path, split="predict", config=config.data)
pred_tr = max(float(model.predict([record.features.tolist()])[0]), 1e-6)
pred_rms = pred_tr * float(record.x_rms.item())
true_rms = float(record.y_rms.item()) if record.y_rms.numel() > 0 else None
out_dir = evaluation_dir(config)
fig_path = save_prediction_figure(file_path, record, pred_rms, true_rms, out_dir)
print(f"Checkpoint: {ckpt_path}")
print(f"Model: {bundle['model_name']}")
print(f"Input file: {file_path}")
print(f"Extracted frequency (Hz): {record.frequency_hz:.6f}")
for feature_name, feature_value in zip(FEATURE_NAMES, record.features.tolist()):
print(f"{feature_name}: {feature_value:.6f}")
print(f"Predicted TR: {pred_tr:.6f}")
print(f"Predicted RMS: {pred_rms:.6f}")
if true_rms is not None:
error_percent = abs(pred_rms - true_rms) / max(abs(true_rms), 1e-12) * 100.0
print(f"True RMS: {true_rms:.6f}")
print(f"Relative RMS Error (%): {error_percent:.4f}")
print(f"Figure saved to: {fig_path}")
if __name__ == "__main__":
main()