scripts/train.py

#!/usr/bin/env python3
"""Train the paper-faithful PointCFD main experiment."""

from __future__ import annotations

import argparse
import copy
import sys
import time
from pathlib import Path
from typing import Any, Dict, Optional

import numpy as np
import torch
import yaml
from torch import nn


PROJECT_ROOT = Path(__file__).resolve().parents[1]
project_root_string = str(PROJECT_ROOT)
if project_root_string in sys.path:
    sys.path.remove(project_root_string)
sys.path.insert(0, project_root_string)

from models import PointNetCFD, count_trainable_parameters  # noqa: E402
from scripts.common import (  # noqa: E402
    AVAILABLE_SAMPLE_COUNT,
    PAPER_SAMPLE_COUNT,
    append_jsonl,
    checkpoint_metadata_matches,
    choose_device,
    configured_paths,
    evaluate_model,
    load_checkpoint,
    load_config,
    make_loader,
    prepare_datasets,
    resolve_path,
    restore_rng_state,
    rng_state,
    save_checkpoint,
    set_deterministic_seed,
    write_json,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument(
        "--config",
        type=Path,
        default=PROJECT_ROOT / "config" / "config.yaml",
        help="Experiment YAML (default: project config/config.yaml)",
    )
    parser.add_argument("--device", default="auto", help="auto, cpu, cuda, or cuda:N")
    parser.add_argument("--epochs", type=int, default=None, help="Override configured epochs")
    parser.add_argument(
        "--batch-size", type=int, default=None, help="Override configured batch size"
    )
    parser.add_argument(
        "--num-workers", type=int, default=None, help="Override DataLoader workers"
    )
    parser.add_argument("--seed", type=int, default=None, help="Override configured seed")
    parser.add_argument(
        "--resume",
        type=Path,
        default=None,
        help="Resume model, optimizer, epoch, metrics, and RNG state",
    )
    parser.add_argument(
        "--smoke-test",
        action="store_true",
        help="Use tiny fixed splits and isolated smoke output paths",
    )
    return parser.parse_args()


def build_effective_config(config: Dict[str, Any], args: argparse.Namespace) -> Dict[str, Any]:
    effective = copy.deepcopy(config)
    training = effective["training"]
    if args.epochs is not None:
        training["epochs"] = args.epochs
    elif args.smoke_test:
        training["epochs"] = 1
    if args.batch_size is not None:
        training["batch_size"] = args.batch_size
    elif args.smoke_test:
        training["batch_size"] = 2
    if args.num_workers is not None:
        training["num_workers"] = args.num_workers
    if args.seed is not None:
        training["seed"] = args.seed
    if int(training["epochs"]) <= 0:
        raise ValueError("epochs must be positive")
    if int(training["batch_size"]) <= 0:
        raise ValueError("batch_size must be positive")
    if int(training["num_workers"]) < 0:
        raise ValueError("num_workers cannot be negative")
    if int(training["validation_interval"]) != 1:
        raise ValueError("The paper validates after every epoch")
    if str(training["optimizer"]).lower() != "adam":
        raise ValueError("The paper uses Adam")
    if str(training["precision"]).lower() != "float32":
        raise ValueError("This reproduction uses the paper-compatible float32 path")
    return effective


def metric_line(metrics: Dict[str, Any]) -> str:
    rmse = metrics["rmse"]
    relative = metrics["relative_l2"]
    return (
        f"val_mse={metrics['normalized_mse']:.9e} "
        f"rmse_u={rmse['u']:.9e} rmse_v={rmse['v']:.9e} rmse_p={rmse['p']:.9e} "
        f"rel_l2_u={relative['u']['mean']:.9e} "
        f"rel_l2_v={relative['v']['mean']:.9e} "
        f"rel_l2_p={relative['p']['mean']:.9e}"
    )


def main() -> None:
    args = parse_args()
    base_config = load_config(args.config)
    config = build_effective_config(base_config, args)
    training = config["training"]
    seed = int(training["seed"])
    set_deterministic_seed(seed)
    device = choose_device(args.device)

    datasets, target_min, target_max, resolved_paths, split_counts = prepare_datasets(
        config, PROJECT_ROOT, smoke_test=args.smoke_test
    )
    batch_size = int(training["batch_size"])
    num_workers = int(training["num_workers"])
    pin_memory = device.type == "cuda"
    train_loader = make_loader(
        datasets["train"], batch_size, True, num_workers, seed, pin_memory
    )
    validation_loader = make_loader(
        datasets["validation"], batch_size, False, num_workers, seed, pin_memory
    )

    model = PointNetCFD(
        input_dim=int(config["model"]["input_dim"]),
        output_dim=int(config["model"]["output_dim"]),
    ).to(device=device, dtype=torch.float32)
    parameter_count = count_trainable_parameters(model)
    paper_parameter_count = int(config["model"]["expected_paper_parameters"])
    print(
        f"device={device} trainable_parameters={parameter_count} "
        f"paper_reference_parameters={paper_parameter_count}",
        flush=True,
    )
    optimizer = torch.optim.Adam(
        model.parameters(),
        lr=float(training["learning_rate"]),
        betas=(float(training["beta1"]), float(training["beta2"])),
        eps=float(training["epsilon"]),
        weight_decay=float(training["weight_decay"]),
    )
    criterion = nn.MSELoss(reduction="mean")

    if args.smoke_test:
        checkpoint_path = PROJECT_ROOT / "weight" / "smoke_best_model.pth"
        results_dir = PROJECT_ROOT / "results" / "smoke"
    else:
        checkpoint_path = resolved_paths["checkpoint"]
        results_dir = resolved_paths["results_dir"]
    results_dir.mkdir(parents=True, exist_ok=True)
    checkpoint_path.parent.mkdir(parents=True, exist_ok=True)
    history_path = results_dir / "train_history.jsonl"
    if args.resume is None and history_path.exists():
        history_path.unlink()

    effective_config_path = results_dir / "effective_config.yaml"
    with effective_config_path.open("w", encoding="utf-8") as handle:
        yaml.safe_dump(config, handle, sort_keys=False)

    start_epoch = 1
    best_validation_mse = float("inf")
    if args.resume is not None:
        resume_path = resolve_path(PROJECT_ROOT, str(args.resume))
        checkpoint = load_checkpoint(resume_path, device)
        checkpoint_metadata_matches(checkpoint, config, target_min, target_max)
        model.load_state_dict(checkpoint["model_state_dict"], strict=True)
        optimizer.load_state_dict(checkpoint["optimizer_state_dict"])
        start_epoch = int(checkpoint["epoch"]) + 1
        best_validation_mse = float(checkpoint["best_validation_mse"])
        restore_rng_state(checkpoint.get("rng_state"))
        if checkpoint.get("train_loader_generator_state") is not None:
            train_loader.generator.set_state(checkpoint["train_loader_generator_state"])
        print(
            f"resumed_from={resume_path} start_epoch={start_epoch} "
            f"best_val_mse={best_validation_mse:.9e}",
            flush=True,
        )

    final_epoch = int(training["epochs"])
    if start_epoch > final_epoch:
        raise ValueError(
            f"Resume checkpoint epoch {start_epoch - 1} already reaches requested epoch {final_epoch}"
        )
    target_names = list(config["data"]["target_names"])
    relative_l2_epsilon = float(config["evaluation"]["relative_l2_epsilon"])
    log_every = int(training["log_every_batches"])
    if log_every <= 0:
        raise ValueError("log_every_batches must be positive")

    run_started = time.time()
    for epoch in range(start_epoch, final_epoch + 1):
        epoch_started = time.time()
        model.train()
        squared_error_sum = 0.0
        element_count = 0
        for batch_number, (coordinates, targets, _) in enumerate(train_loader, start=1):
            coordinates = coordinates.to(device, non_blocking=True)
            targets = targets.to(device, non_blocking=True)
            optimizer.zero_grad(set_to_none=True)
            predictions = model(coordinates)
            loss = criterion(predictions, targets)
            loss.backward()
            optimizer.step()

            batch_elements = targets.numel()
            squared_error_sum += float(loss.detach().item()) * batch_elements
            element_count += batch_elements
            if batch_number % log_every == 0 or batch_number == len(train_loader):
                print(
                    f"epoch={epoch}/{final_epoch} "
                    f"batch={batch_number}/{len(train_loader)} "
                    f"train_loss={loss.detach().item():.9e}",
                    flush=True,
                )
        train_mse = squared_error_sum / element_count

        validation_metrics, _ = evaluate_model(
            model,
            validation_loader,
            device,
            target_min,
            target_max,
            target_names,
            relative_l2_epsilon,
        )
        validation_mse = float(validation_metrics["normalized_mse"])
        elapsed = time.time() - epoch_started
        print(
            f"epoch={epoch}/{final_epoch} train_mse={train_mse:.9e} "
            f"{metric_line(validation_metrics)} epoch_seconds={elapsed:.3f}",
            flush=True,
        )

        history_record = {
            "epoch": epoch,
            "train_mse": train_mse,
            "validation": validation_metrics,
            "epoch_seconds": elapsed,
            "learning_rate": float(optimizer.param_groups[0]["lr"]),
            "seed": seed,
            "smoke_test": bool(args.smoke_test),
        }
        append_jsonl(history_path, history_record)

        if validation_mse < best_validation_mse:
            best_validation_mse = validation_mse
            checkpoint_payload: Dict[str, Any] = {
                "format_version": "pointcfd-checkpoint-v1",
                "epoch": epoch,
                "model_state_dict": model.state_dict(),
                "optimizer_state_dict": optimizer.state_dict(),
                "best_validation_mse": best_validation_mse,
                "target_min": target_min,
                "target_max": target_max,
                "source_channels": list(config["data"]["source_channels"]),
                "input_names": list(config["data"]["input_names"]),
                "target_names": target_names,
                "input_indices": list(config["data"]["input_indices"]),
                "target_indices": list(config["data"]["target_indices"]),
                "model_config": copy.deepcopy(config["model"]),
                "training_config": copy.deepcopy(training),
                "data_paths": {key: str(value) for key, value in resolved_paths.items()},
                "split_counts": split_counts,
                "seed": seed,
                "available_sample_count": AVAILABLE_SAMPLE_COUNT,
                "paper_sample_count": PAPER_SAMPLE_COUNT,
                "trainable_parameters": parameter_count,
                "rng_state": rng_state(),
                "train_loader_generator_state": train_loader.generator.get_state(),
                "smoke_test": bool(args.smoke_test),
            }
            save_checkpoint(checkpoint_path, checkpoint_payload)
            print(
                f"saved_best_checkpoint={checkpoint_path} "
                f"best_val_mse={best_validation_mse:.9e}",
                flush=True,
            )

    summary = {
        "status": "completed",
        "start_epoch": start_epoch,
        "final_epoch": final_epoch,
        "best_validation_mse": best_validation_mse,
        "checkpoint": str(checkpoint_path),
        "history": str(history_path),
        "effective_config": str(effective_config_path),
        "trainable_parameters": parameter_count,
        "paper_reference_parameters": paper_parameter_count,
        "split_counts": split_counts,
        "available_sample_count": AVAILABLE_SAMPLE_COUNT,
        "paper_sample_count": PAPER_SAMPLE_COUNT,
        "dataset_limitation": (
            "The supplied dataset has 2215 cases rather than the paper's 2595; "
            "this is a best-available subset reproduction."
        ),
        "elapsed_seconds": time.time() - run_started,
        "device": str(device),
        "smoke_test": bool(args.smoke_test),
    }
    write_json(results_dir / "training_summary.json", summary)
    print(f"training_complete summary={results_dir / 'training_summary.json'}", flush=True)


if __name__ == "__main__":
    main()