| """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() |
|
|
