train_model.py

import os
import sys
from pathlib import Path

from modules.calibration import calibrate_young
from modules.lab_utils import load_dataset, LAB_PATH, fix_path
from modules.pytorch_mlp import PytorchMLPReg

DEFAULT = "pytorch"


def train_pytorch_model(dataset_path, from_real=False):
    x_train, y_train, x_test, y_test = load_dataset(dataset_path, from_real)

    dataset_fname = dataset_path.parts[-1].strip(".csv")
    fname = f"{LAB_PATH}/data/results/{dataset_fname}.pth"

    mlp = PytorchMLPReg()

    mlp.train(x_train, y_train, x_test, y_test, n_epochs=2_000)

    mlp.save(fname)
    print(f"Trained model saved at {os.path.abspath(fname)}")


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Train model using dataset")

    parser.add_argument(
        "--model-type",
        type=str,
        choices=["pytorch", "calibrated"],
        default=DEFAULT,
        help="Model type: pytorch or calibrated",
    )
    parser.add_argument(
        "--dataset-path",
        type=Path,
        default=Path(
            "data/results/blueleg_beam_sphere515.csv"
        ),
        help="Path to dataset CSV",
    )
    parser.add_argument(
        "-r",
        "--from-real",
        action="store_true",
        help="Use real-world dataset instead of synthetic",
    )
    args = parser.parse_args()

    dataset_path = args.dataset_path
    learn_from_real = args.from_real
    model_type = args.model_type

    dataset_path = fix_path(dataset_path)
    if dataset_path is None:
        print(f"Dataset file not found: {dataset_path}")
        sys.exit(1)

    if model_type == "calibrated":
        calibrate_young(dataset_path, learn_from_real)
    elif model_type == "pytorch":
        train_pytorch_model(dataset_path, learn_from_real)
    else:
        print(f"Unknown model type: {model_type}")
        sys.exit(1)