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