scripts/surrogate_postprocess.py

"""
Post-process the surrogate-demo runs to compute the Laplace pressure jump
ΔP = P_inside - P_outside across the droplet interface, then re-train the
GP / RandomForest / MLP surrogates on this physically-meaningful target.

Reuses the existing dumps in work/surrogate_demo/run_*/dumps/ — no new
simulations are run.

Per-atom 'pressure' is approximated from the dumped stress/atom output as
    P_atom ≈ -(sxx + syy) / 2 * rho_atom
since stress/atom is in units of pressure*volume and per-atom volume ~ 1/rho.
This is an order-of-magnitude estimate; the *jump* between inside and outside
is what the surrogate learns, and that is robust to a global volume scale.
"""
import csv
import json
import re
import sys
import time
from pathlib import Path

import numpy as np
from sklearn.ensemble import RandomForestRegressor
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF, ConstantKernel, WhiteKernel
from sklearn.model_selection import LeaveOneOut
from sklearn.neural_network import MLPRegressor
from sklearn.pipeline import make_pipeline
from sklearn.preprocessing import StandardScaler

ROOT = Path(__file__).resolve().parent.parent
WORK = ROOT / "work" / "surrogate_demo"
N_TAIL_FRAMES = 5
SEED = 42


def parse_dump(path: Path):
    """Return numpy array of shape (n,4): type, density, sxx, syy."""
    with path.open() as f:
        lines = f.readlines()
    # find ITEM: ATOMS line
    for i, ln in enumerate(lines):
        if ln.startswith("ITEM: ATOMS"):
            header = ln.split()[2:]
            data_start = i + 1
            break
    else:
        raise RuntimeError(f"no ATOMS section in {path}")
    cols = {name: idx for idx, name in enumerate(header)}
    needed = ("type", "c_density", "c_peratom[1]", "c_peratom[2]")
    idxs = [cols[n] for n in needed]
    arr = np.empty((len(lines) - data_start, 4), dtype=float)
    for j, ln in enumerate(lines[data_start:]):
        parts = ln.split()
        for k, ic in enumerate(idxs):
            arr[j, k] = float(parts[ic])
    return arr  # columns: type, density, sxx, syy


def pressure_jump(run_dir: Path):
    """Average ΔP = P_inside - P_outside over the last N_TAIL_FRAMES dump frames."""
    dumps = sorted(run_dir.glob("dumps/dump.*.lammpstrj"), key=lambda p: int(p.stem.split(".")[1]))
    if len(dumps) < N_TAIL_FRAMES:
        return None, None, None
    tail = dumps[-N_TAIL_FRAMES:]
    p_in_list, p_out_list = [], []
    for d in tail:
        a = parse_dump(d)
        # P_atom ≈ -(sxx + syy)/2 * rho   (per-atom virial → pressure)
        p_atom = -0.5 * (a[:, 2] + a[:, 3]) * a[:, 1]
        is_in = a[:, 0] == 2
        if is_in.sum() == 0 or (~is_in).sum() == 0:
            continue
        p_in_list.append(p_atom[is_in].mean())
        p_out_list.append(p_atom[~is_in].mean())
    if not p_in_list:
        return None, None, None
    p_in = float(np.mean(p_in_list))
    p_out = float(np.mean(p_out_list))
    return p_in, p_out, p_in - p_out


def loo_score(model_factory, X, y):
    loo = LeaveOneOut()
    preds, truth = [], []
    for tr, te in loo.split(X):
        m = model_factory()
        m.fit(X[tr], y[tr])
        preds.append(float(m.predict(X[te])[0]))
        truth.append(float(y[te][0]))
    p = np.array(preds)
    t = np.array(truth)
    rmse = float(np.sqrt(np.mean((p - t) ** 2)))
    ss_res = float(np.sum((p - t) ** 2))
    ss_tot = float(np.sum((t - t.mean()) ** 2))
    r2 = 1.0 - ss_res / ss_tot if ss_tot > 0 else float("nan")
    return r2, rmse


def main() -> int:
    csv_path = WORK / "data.csv"
    if not csv_path.exists():
        print(f"missing {csv_path} — run scripts/surrogate_demo.py first")
        return 1
    rows = list(csv.DictReader(csv_path.open()))

    enriched = []
    print(f"{'run':>5} {'sigmao':>7} {'dh':>7} {'rad':>6} {'P_in':>8} {'P_out':>8} {'ΔP':>8} {'σ/r':>7}")
    for i, r in enumerate(rows):
        run_dir = WORK / f"run_{i:03d}"
        log = (run_dir / "log.lammps").read_text()
        m = re.search(r"^Radio droplet:\s*([0-9.eE+-]+)", log, flags=re.M)
        rad = float(m.group(1)) if m else float("nan")
        p_in, p_out, dp = pressure_jump(run_dir)
        sigmao = float(r["sigmao"])
        dh = float(r["dh"])
        sr = sigmao / rad if rad > 0 else float("nan")
        if dp is None:
            print(f"{i:5d}  --- skipped ---")
            continue
        enriched.append({
            "sigmao": sigmao, "dh": dh, "rad": rad,
            "P_in": p_in, "P_out": p_out, "delta_P": dp,
            "sigma_over_r": sr,
        })
        print(f"{i:5d} {sigmao:7.3f} {dh:7.4f} {rad:6.3f} {p_in:8.4f} {p_out:8.4f} {dp:8.4f} {sr:7.4f}")

    out_csv = WORK / "data_with_dp.csv"
    with out_csv.open("w") as f:
        w = csv.DictWriter(f, fieldnames=list(enriched[0].keys()))
        w.writeheader()
        w.writerows(enriched)
    print(f"\nSaved enriched data to {out_csv}")

    # Correlations vs ΔP
    arr = {k: np.array([d[k] for d in enriched]) for k in enriched[0]}
    print(f"\nCorrelations with ΔP (n={len(enriched)}):")
    for k in ("sigmao", "dh", "rad", "sigma_over_r"):
        r = float(np.corrcoef(arr[k], arr["delta_P"])[0, 1])
        print(f"  {k:15s} r = {r:+.4f}")

    X = np.column_stack([arr["sigmao"], arr["dh"]])

    # ----- Primary target: P_inside (pressure inside the droplet) -----
    # P_in is governed by local surface-tension physics; P_out is
    # contaminated by box/wall artifacts and is not smooth in (σ, dh).
    y = arr["P_in"]
    print(f"\n=== Surrogate on (sigmao, dh) → P_inside ===")
    print(f"Target P_in: range [{y.min():.4f}, {y.max():.4f}], mean {y.mean():.4f}, std {y.std():.4f}")
    print(f"Baseline (predict mean): LOO R² = 0.000  RMSE = {y.std():.4f}")

    def gp_factory():
        kernel = ConstantKernel(1.0, (1e-3, 1e3)) * RBF(length_scale=[0.5, 0.05]) + WhiteKernel(1e-5, (1e-12, 1e-1))
        return GaussianProcessRegressor(kernel=kernel, normalize_y=True, n_restarts_optimizer=4, random_state=SEED)

    def rf_factory():
        return RandomForestRegressor(n_estimators=300, min_samples_leaf=1, random_state=SEED)

    def mlp_factory():
        return make_pipeline(
            StandardScaler(),
            MLPRegressor(hidden_layer_sizes=(16, 16), activation="tanh", max_iter=8000, random_state=SEED, tol=1e-7),
        )

    results = {}
    for name, factory in [("GP (RBF)", gp_factory), ("RandomForest", rf_factory), ("MLP (16,16) tanh", mlp_factory)]:
        r2, rmse = loo_score(factory, X, y)
        results[name] = {"r2": r2, "rmse": rmse}
        verdict = "  ← beats mean" if r2 > 0 else "  worse than mean"
        print(f"{name:24s}  LOO R² = {r2:+.4f}  RMSE = {rmse:.4f}{verdict}")

    # Linear baseline: P_in = a*sigmao + b*dh + c
    from sklearn.linear_model import LinearRegression
    def lin_factory():
        return LinearRegression()
    r2_lin, rmse_lin = loo_score(lin_factory, X, y)
    verdict = "  ← beats mean" if r2_lin > 0 else "  worse than mean"
    print(f"{'Linear (2-feature)':24s}  LOO R² = {r2_lin:+.4f}  RMSE = {rmse_lin:.4f}{verdict}")

    # ----- Secondary: ΔP across droplet interface (less clean signal) -----
    y_dp = arr["delta_P"]
    print(f"\n=== Surrogate on (sigmao, dh) → ΔP  (P_in − P_out, secondary)  ===")
    print(f"Target ΔP: range [{y_dp.min():.4f}, {y_dp.max():.4f}], std {y_dp.std():.4f}")
    print(f"Baseline (predict mean): LOO R² = 0.000  RMSE = {y_dp.std():.4f}")
    for name, factory in [("GP (RBF)", gp_factory), ("RandomForest", rf_factory), ("Linear", lin_factory)]:
        r2, rmse = loo_score(factory, X, y_dp)
        verdict = "  ← beats mean" if r2 > 0 else "  worse than mean"
        print(f"{name:24s}  LOO R² = {r2:+.4f}  RMSE = {rmse:.4f}{verdict}")

    # Demo speedup: time the surrogate for many predictions
    n_predict = 1000
    rng = np.random.default_rng(SEED)
    test_X = rng.uniform([arr["sigmao"].min(), arr["dh"].min()],
                         [arr["sigmao"].max(), arr["dh"].max()],
                         (n_predict, 2))
    best_name = max(results, key=lambda k: results[k]["r2"])
    factories = {"GP (RBF)": gp_factory, "RandomForest": rf_factory, "MLP (16,16) tanh": mlp_factory}
    m = factories[best_name]()
    m.fit(X, y)
    t0 = time.time()
    m.predict(test_X)
    pred_ms = (time.time() - t0) * 1000
    print(f"\n=== Speedup demo ===")
    print(f"Best model: {best_name} (LOO R² = {results[best_name]['r2']:+.4f})")
    print(f"Surrogate predicts {n_predict} new (sigmao,dh) points in {pred_ms:.1f} ms")
    print(f"Each LAMMPS run takes ~65 s")
    print(f"Speedup per query: ~{(65000 * n_predict) / pred_ms:,.0f}×")

    summary = {"n": len(enriched), "models_xy_for_P_in": results}
    (WORK / "summary_dp.json").write_text(json.dumps(summary, indent=2))
    return 0


if __name__ == "__main__":
    sys.exit(main())