"""Train the P-Alert intensity GMM and report metrics. Usage: python train.py --epochs 120 """ import argparse import json import os import time import numpy as np import pandas as pd import torch import torch.nn as nn from palert.features import INTENSITY_LEVELS, intensity_class, make_features from palert.model import GMM from palert.paths import DATA_DIR DATASET_PATH = os.path.join(DATA_DIR, "dataset.csv") MODEL_PATH = os.path.join(DATA_DIR, "gmm.pt") METRICS_PATH = os.path.join(DATA_DIR, "metrics.json") def log(msg): print(msg, flush=True) def r2_score(y_true, y_pred): ss_res = np.sum((y_true - y_pred) ** 2) ss_tot = np.sum((y_true - y_true.mean()) ** 2) return 1.0 - ss_res / ss_tot def split_by_event(df, test_frac, seed): rng = np.random.default_rng(seed) events = np.array(df["event"].unique()) rng.shuffle(events) n_test = max(1, int(len(events) * test_frac)) test_events = set(events[:n_test]) mask = df["event"].isin(test_events).to_numpy() return ~mask, mask, test_events def sample_weights(df, mode): n = df.groupby("event")["event"].transform("size").to_numpy(np.float64) if mode == "none": return np.ones(len(df), dtype=np.float64) if mode == "event": return 1.0 / n if mode == "sqrt": return 1.0 / np.sqrt(n) raise ValueError(mode) def evaluate(y_true_log, y_pred_log, pga_true, pgv_true, pga_pred, pgv_pred): lvl_true = intensity_class(pga_true, pgv_true) lvl_pred = intensity_class(pga_pred, pgv_pred) exact = (lvl_true == lvl_pred).mean() within1 = (np.abs(lvl_true - lvl_pred) <= 1).mean() m = { "n": int(len(y_true_log)), "pga_r2_log": float(r2_score(y_true_log[:, 0], y_pred_log[:, 0])), "pgv_r2_log": float(r2_score(y_true_log[:, 1], y_pred_log[:, 1])), "pga_mae_log": float(np.mean(np.abs(y_true_log[:, 0] - y_pred_log[:, 0]))), "pgv_mae_log": float(np.mean(np.abs(y_true_log[:, 1] - y_pred_log[:, 1]))), "intensity_acc": float(exact), "intensity_acc_pm1": float(within1), } per_class = {} for i, name in enumerate(INTENSITY_LEVELS): sel = lvl_true == i if sel.sum(): per_class[name] = { "n": int(sel.sum()), "acc": float((lvl_pred[sel] == i).mean()), "acc_pm1": float((np.abs(lvl_pred[sel] - i) <= 1).mean()), } m["per_class"] = per_class return m def main(): ap = argparse.ArgumentParser() ap.add_argument("--epochs", type=int, default=120) ap.add_argument("--batch", type=int, default=8192) ap.add_argument("--lr", type=float, default=2e-3) ap.add_argument("--hidden", type=int, default=256) ap.add_argument("--depth", type=int, default=3) ap.add_argument("--dropout", type=float, default=0.0) ap.add_argument("--weight", choices=["none", "event", "sqrt"], default="sqrt") ap.add_argument("--test-frac", type=float, default=0.15) ap.add_argument("--seed", type=int, default=42) ap.add_argument("--threads", type=int, default=0) ap.add_argument("--dataset", default=DATASET_PATH, help="CSV with the P-Alert or CWA columns") ap.add_argument("--tag", default="", help="suffix for the saved model/metrics, e.g. _cwa") args = ap.parse_args() model_path = MODEL_PATH.replace(".pt", f"{args.tag}.pt") metrics_path = METRICS_PATH.replace(".json", f"{args.tag}.json") if args.threads: torch.set_num_threads(args.threads) torch.manual_seed(args.seed) np.random.seed(args.seed) df = pd.read_csv(args.dataset) log(f"[data] {len(df)} rows, {df['event'].nunique()} events") X = make_features(df) pga = df["pga"].to_numpy(np.float64) pgv = df["pgv"].to_numpy(np.float64) Y = np.stack([np.log10(pga), np.log10(pgv)], axis=1).astype(np.float32) tr, te, test_events = split_by_event(df, args.test_frac, args.seed) log(f"[split] train {tr.sum()} rows / test {te.sum()} rows ({len(test_events)} held-out events)") mean = X[tr].mean(0) std = X[tr].std(0) std[std == 0] = 1.0 w = sample_weights(df, args.weight).astype(np.float32) Xt = torch.from_numpy(X) Yt = torch.from_numpy(Y) Wt = torch.from_numpy(w) tr_t = torch.from_numpy(np.where(tr)[0]) te_t = torch.from_numpy(np.where(te)[0]) model = GMM(args.hidden, args.depth, 2, args.dropout) model.set_norm(mean, std) opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=1e-4) sched = torch.optim.lr_scheduler.OneCycleLR( opt, max_lr=args.lr, total_steps=args.epochs * max(1, len(tr_t) // args.batch + 1) ) lossf = nn.MSELoss(reduction="none") best = None t0 = time.time() for epoch in range(1, args.epochs + 1): model.train() perm = tr_t[torch.randperm(len(tr_t))] tot = 0.0 nb = 0 for i in range(0, len(perm), args.batch): idx = perm[i:i + args.batch] opt.zero_grad() pred = model(Xt[idx]) loss = (lossf(pred, Yt[idx]).mean(1) * Wt[idx]).mean() loss.backward() opt.step() sched.step() tot += float(loss.detach()) nb += 1 if epoch % 10 == 0 or epoch == args.epochs: model.eval() with torch.no_grad(): pred_te = model(Xt[te_t]).numpy() m = evaluate( Y[te], pred_te, pga[te], pgv[te], np.power(10.0, pred_te[:, 0]), np.power(10.0, pred_te[:, 1]), ) log( f"epoch {epoch:3d} loss {tot/nb:.4f} " f"PGA R2(log) {m['pga_r2_log']:.4f} PGV R2(log) {m['pgv_r2_log']:.4f} " f"intensity acc {m['intensity_acc']:.4f}" ) if best is None or m["pga_r2_log"] + m["pgv_r2_log"] > best["score"]: best = { "score": m["pga_r2_log"] + m["pgv_r2_log"], "epoch": epoch, "metrics": m, "state": {k: v.clone() for k, v in model.state_dict().items()}, } log(f"[train] done in {time.time()-t0:.1f}s; best epoch {best['epoch']}") model.load_state_dict(best["state"]) torch.save( { "state_dict": model.state_dict(), "config": {"hidden": args.hidden, "depth": args.depth, "dropout": args.dropout}, "metrics": best["metrics"], }, model_path, ) with open(metrics_path, "w", encoding="utf-8") as fh: json.dump(best["metrics"], fh, ensure_ascii=False, indent=2) log(f"[save] {model_path}") log(json.dumps({k: v for k, v in best["metrics"].items() if k != "per_class"}, indent=2)) if __name__ == "__main__": main()