Building/scripts_tree/train.py

from __future__ import annotations

import pickle
from dataclasses import asdict
from pathlib import Path
from typing import Any
import sys

import pandas as pd
from sklearn.ensemble import ExtraTreesRegressor, GradientBoostingRegressor, RandomForestRegressor

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_datasets, report_to_text

try:
    from .config import FEATURE_NAMES, ExperimentConfig, checkpoint_dir, make_experiment_config
except ImportError:
    from config import FEATURE_NAMES, ExperimentConfig, checkpoint_dir, make_experiment_config


def build_regressor(model_name: str, random_state: int):
    if model_name == "extra_trees":
        return ExtraTreesRegressor(
            n_estimators=600,
            max_depth=None,
            min_samples_leaf=1,
            min_samples_split=2,
            random_state=random_state,
        )
    if model_name == "random_forest":
        return RandomForestRegressor(
            n_estimators=500,
            max_depth=None,
            min_samples_leaf=1,
            min_samples_split=2,
            random_state=random_state,
        )
    if model_name == "gradient_boosting":
        return GradientBoostingRegressor(
            n_estimators=300,
            learning_rate=0.03,
            max_depth=3,
            random_state=random_state,
            loss="huber",
        )
    raise ValueError(f"Unsupported model: {model_name}")


def records_to_frame(records: list[Any]) -> pd.DataFrame:
    rows = []
    for record in records:
        row = {
            "file_name": record.file_path.name,
            "frequency_hz": float(record.frequency_hz),
            "x_rms": float(record.x_rms.item()),
            "y_rms": float(record.y_rms.item()),
            "target_tr": float(record.target_y_rms.item()),
        }
        for name, value in zip(FEATURE_NAMES, record.features.tolist()):
            row[name] = float(value)
        rows.append(row)
    return pd.DataFrame(rows)


def evaluate_frame(model, frame: pd.DataFrame) -> dict[str, float]:
    feature_matrix = frame.loc[:, FEATURE_NAMES].to_numpy()
    pred_tr = model.predict(feature_matrix)
    pred_tr = pred_tr.clip(min=1e-6)
    pred_rms = pred_tr * frame["x_rms"].to_numpy()
    true_rms = frame["y_rms"].to_numpy()
    relative_error = abs(pred_rms - true_rms) / true_rms.clip(min=1e-6)
    return {
        "mean_rms_error": float(relative_error.mean()),
        "median_rms_error": float(pd.Series(relative_error).median()),
        "max_rms_error": float(relative_error.max()),
    }


def serialize_for_checkpoint(value: Any) -> Any:
    if isinstance(value, Path):
        return str(value)
    if isinstance(value, dict):
        return {key: serialize_for_checkpoint(sub_value) for key, sub_value in value.items()}
    if isinstance(value, tuple):
        return [serialize_for_checkpoint(item) for item in value]
    if isinstance(value, list):
        return [serialize_for_checkpoint(item) for item in value]
    return value


def train() -> None:
    config = make_experiment_config()
    datasets, raw_records, reports = build_datasets(config.data)
    del datasets
    print(report_to_text(reports))

    train_df = records_to_frame(raw_records["train"])
    val_df = records_to_frame(raw_records["val"])

    feature_matrix = train_df.loc[:, FEATURE_NAMES].to_numpy()
    target = train_df["target_tr"].to_numpy()

    results: list[dict[str, Any]] = []
    best_bundle: dict[str, Any] | None = None
    best_val_error = float("inf")

    for model_name in config.model.candidate_models:
        model = build_regressor(model_name, config.model.random_state)
        model.fit(feature_matrix, target)

        train_metrics = evaluate_frame(model, train_df)
        val_metrics = evaluate_frame(model, val_df)
        row = {
            "model_name": model_name,
            "train_mean_rms_error": train_metrics["mean_rms_error"],
            "train_median_rms_error": train_metrics["median_rms_error"],
            "train_max_rms_error": train_metrics["max_rms_error"],
            "val_mean_rms_error": val_metrics["mean_rms_error"],
            "val_median_rms_error": val_metrics["median_rms_error"],
            "val_max_rms_error": val_metrics["max_rms_error"],
        }
        results.append(row)
        print(
            f"{model_name}: train_mean={train_metrics['mean_rms_error']:.6f} | "
            f"val_mean={val_metrics['mean_rms_error']:.6f}"
        )

        if val_metrics["mean_rms_error"] < best_val_error:
            best_val_error = val_metrics["mean_rms_error"]
            best_bundle = {
                "model_name": model_name,
                "model": model,
                "config": serialize_for_checkpoint(asdict(config)),
                "feature_names": list(FEATURE_NAMES),
                "val_metrics": val_metrics,
                "train_metrics": train_metrics,
            }

    summary_df = pd.DataFrame(results).sort_values("val_mean_rms_error").reset_index(drop=True)
    ckpt_dir = checkpoint_dir(config)
    summary_path = ckpt_dir / config.train.summary_name
    summary_df.to_csv(summary_path, index=False)
    print(f"Model selection saved to: {summary_path}")
    print(summary_df.to_string(index=False))

    if best_bundle is None:
        raise RuntimeError("No valid model trained.")

    best_path = ckpt_dir / config.train.best_model_name
    with best_path.open("wb") as handle:
        pickle.dump(best_bundle, handle)
    print(f"Best tree model saved to: {best_path}")
    print(f"Best model: {best_bundle['model_name']} | val_mean_rms_error={best_val_error:.6f}")


if __name__ == "__main__":
    train()