tradeoff

Author: sdat2

adbo/exp.py

@@ -41,6 +41,7 @@
 import matplotlib.pyplot as plt
 from adforce.wrap import idealized_tc_observe, get_default_config
 from adforce.constants import NEW_ORLEANS
+from .wrap_utils import build_wrap_config
 from .ani import plot_gps
 from .rescale import rescale_inverse
 from .constants import FIGURE_PATH, EXP_PATH, CONFIG_PATH, PROJECT_PATH
@@ -199,44 +200,12 @@ def obj(x: tf.Tensor) -> tf.Tensor:
                 name: float(real_queries[i][j])
                 for j, name in enumerate(cfg["constraints"]["order"])
             }
-            wrap_cfg = get_default_config()
-
-            print("read config", wrap_cfg)
             # I want to generalize this so it's relative to the observation point not New Orleans
             tmp_dir = temp_dir()
-            wrap_cfg.files.run_folder = tmp_dir
-            # delete fort.22.nc after run
-            wrap_cfg.files.low_storage = True
-            if tradeoff_curve is not None:
-                # storm on the size-intensity tradeoff curve: generate the
-                # CLE15 profile for this sampled intensity and store it in
-                # the run folder (self-describing provenance; fort22 accepts
-                # a direct .json path in place of a profile name).
-                profile_json = os.path.join(tmp_dir, "profile.json")
-                tradeoff_curve.profile(inputs["vmax"], out_path=profile_json)
-                wrap_cfg.tc.profile_name.value = profile_json
-            else:
-                wrap_cfg.tc.profile_name.value = cfg["profile_name"]
-            wrap_cfg.adcirc.attempted_observation_location.value = [
-                cfg["obs_lon"],
-                cfg["obs_lat"],
-            ]
-            wrap_cfg.adcirc["resolution"].value = cfg["resolution"]
-            for inp in inputs:
-                if inp == "vmax":
-                    # handled above via the tradeoff-curve profile; not a
-                    # track parameter in the adforce tc config
-                    continue
-                if inp != "displacement":
-                    if inp == "trans_speed":
-                        wrap_cfg.tc["translation_speed"].value = inputs[inp]
-                    else:
-                        wrap_cfg.tc[inp].value = inputs[inp]
-                if inp == "displacement":
-                    wrap_cfg.tc.impact_location.value = [
-                        cfg["obs_lon"] + inputs["displacement"],
-                        cfg["obs_lat"],
-                    ]
+            wrap_cfg = build_wrap_config(
+                cfg, inputs, tmp_dir, tradeoff_curve=tradeoff_curve
+            )
+            print("read config", wrap_cfg)
             # if "displacement" in inputs:
             #    # assume impact lon relative to New Orleans
             #    inputs["impact_lon"] = NEW_ORLEANS.lon + inputs["displacement"]

adbo/exp_4d.py

@@ -0,0 +1,80 @@
+"""adbo.exp_4d.py
+
+Run a 4D Bayesian-optimization experiment varying the angle, displacement,
+translation speed AND intensity of the storm, with each sampled intensity
+placed on the size--intensity tradeoff curve r(V) of the potential size
+model (see ``w22.tradeoff``). Requires a precomputed curve netCDF:
+
+    from w22.tradeoff import env_from_point_ds, generate_curve
+    env = env_from_point_ds("w22/data/new_orleans_august_ssp585_CESM2_r4i1p1f1_isothermal_pi4new.nc", 2025)
+    generate_curve(env, "2025_new_orleans_r4i1p1f1")
+
+Example::
+
+    python -m adbo.exp_4d --curve_nc w22/data/curves/2025_new_orleans_r4i1p1f1.nc \
+        --exp_name no-4d-2025 --init_steps 35 --daf_steps 35
+"""
+
+import argparse
+import os
+
+import yaml
+
+from adforce.constants import NEW_ORLEANS
+from .constants import CONFIG_PATH
+
+CONSTRAINTS_4D_PATH = os.path.join(CONFIG_PATH, "4d_constraints.yaml")
+
+
+def run_4d_exp() -> None:
+    """Run the 4D (track + on-curve intensity) BO experiment from the CLI."""
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--curve_nc", type=str, required=True)
+    parser.add_argument("--test", action="store_true", help="wrap_test mode")
+    parser.add_argument("--seed", type=int, default=10)
+    parser.add_argument("--obs_lon", type=float, default=NEW_ORLEANS.lon)
+    parser.add_argument("--obs_lat", type=float, default=NEW_ORLEANS.lat)
+    parser.add_argument("--init_steps", type=int, default=35)
+    parser.add_argument("--daf_steps", type=int, default=35)
+    parser.add_argument("--resolution", type=str, default="mid")
+    parser.add_argument("--exp_name", type=str, default="bo-4d")
+    parser.add_argument("--kernel", type=str, default="Matern52")
+    parser.add_argument("--daf", type=str, default="mes")
+
+    # constraint overrides (vmax bounds default to the curve domain)
+    constraints = yaml.safe_load(open(CONSTRAINTS_4D_PATH))
+    for dim in ("angle", "trans_speed", "displacement", "vmax"):
+        for bound in ("min", "max"):
+            parser.add_argument(
+                f"--{dim}_{bound}",
+                type=float,
+                default=constraints[dim][bound],
+            )
+    args = parser.parse_args()
+    print(args)
+
+    for dim in ("angle", "trans_speed", "displacement", "vmax"):
+        constraints[dim]["min"] = getattr(args, f"{dim}_min")
+        constraints[dim]["max"] = getattr(args, f"{dim}_max")
+
+    # heavy import (tensorflow/trieste) deferred until actually running
+    from .exp import run_bayesopt_exp
+
+    run_bayesopt_exp(
+        constraints=constraints,
+        seed=args.seed,
+        exp_name=args.exp_name,
+        resolution=args.resolution,
+        obs_lon=args.obs_lon,
+        obs_lat=args.obs_lat,
+        init_steps=args.init_steps,
+        daf_steps=args.daf_steps,
+        daf=args.daf,
+        kernel=args.kernel,
+        wrap_test=args.test,
+        curve_path=args.curve_nc,
+    )
+
+
+if __name__ == "__main__":
+    run_4d_exp()

adbo/plot_tradeoff.py

@@ -0,0 +1,191 @@
+"""Plots for the size--intensity tradeoff experiments.
+
+- :func:`plot_vmax_sweep`: surge height against intensity from a 1D sweep
+  (``adbo.sweep_vmax``), with the curve's r0/rmax context panels.
+- :func:`plot_4d_samples`: surge against sampled intensity from a 4D BO
+  experiment (``adbo.exp_4d``), split into initial-design vs acquisition
+  samples.
+
+Both read the ``experiments.json`` ledger, so they work on partial runs.
+
+Example::
+
+    python -m adbo.plot_tradeoff --exp_dir <exp>/no-sweep-2025 --mode sweep
+"""
+
+import argparse
+import json
+import os
+from typing import Optional
+
+import matplotlib
+
+matplotlib.use("Agg")
+import matplotlib.pyplot as plt
+import numpy as np
+
+
+def _load_ledger(exp_dir: str) -> dict:
+    with open(os.path.join(exp_dir, "experiments.json"), "r", encoding="utf-8") as f:
+        return json.load(f)
+
+
+def _fig_path(exp_dir: str, name: str) -> str:
+    img_dir = os.path.join(exp_dir, "img")
+    os.makedirs(img_dir, exist_ok=True)
+    return os.path.join(img_dir, name)
+
+
+def _curve_context(curve_nc: Optional[str]):
+    if curve_nc is None or not os.path.exists(curve_nc):
+        return None
+    from w22.tradeoff import TradeoffCurve
+
+    return TradeoffCurve.from_file(curve_nc)
+
+
+def plot_vmax_sweep(exp_dir: str, curve_nc: Optional[str] = None) -> str:
+    """Plot surge height vs intensity for a 1D sweep experiment.
+
+    Args:
+        exp_dir (str): Sweep experiment directory (contains experiments.json).
+        curve_nc (str, optional): Tradeoff curve for context panels; defaults
+            to the path recorded in sweep-config.json.
+
+    Returns:
+        str: Path of the written figure.
+    """
+    ledger = _load_ledger(exp_dir)
+    if curve_nc is None:
+        cfg_path = os.path.join(exp_dir, "sweep-config.json")
+        if os.path.exists(cfg_path):
+            with open(cfg_path, "r", encoding="utf-8") as f:
+                curve_nc = json.load(f).get("curve_path")
+    curve = _curve_context(curve_nc)
+
+    v = np.array([rec["vmax"] for rec in ledger.values()])
+    z = np.array([rec["res"] for rec in ledger.values()])
+    order = np.argsort(v)
+    v, z = v[order], z[order]
+    ok = np.isfinite(z)
+
+    nrows = 2 if curve is not None else 1
+    fig, axs = plt.subplots(
+        nrows, 1, sharex=True, figsize=(6, 3 * nrows), squeeze=False
+    )
+    ax = axs[0][0]
+    ax.plot(v[ok], z[ok], "o-", color="tab:blue")
+    if (~ok).any():
+        for vf in v[~ok]:
+            ax.axvline(vf, color="red", alpha=0.3, linestyle=":")
+    if len(z[ok]):
+        i_best = int(np.nanargmax(z[ok]))
+        ax.plot(v[ok][i_best], z[ok][i_best], "*", color="tab:orange", markersize=15)
+        ax.annotate(
+            f"max {z[ok][i_best]:.2f} m @ {v[ok][i_best]:.1f} m/s",
+            (v[ok][i_best], z[ok][i_best]),
+            textcoords="offset points",
+            xytext=(5, -12),
+        )
+    ax.set_ylabel("Max SSH at observation point [m]")
+    ax.set_title("Surge along the size-intensity tradeoff curve")
+
+    if curve is not None:
+        ax.axvline(curve.v_max, color="green", linestyle="--", alpha=0.7)
+        ax.annotate(
+            "$V_p$", (curve.v_max, ax.get_ylim()[0]), color="green", ha="right"
+        )
+        ax2 = axs[1][0]
+        vv = np.linspace(curve.v_min, curve.v_max, 200)
+        ax2.plot(vv, [curve.rmax(x) / 1000 for x in vv], color="tab:purple")
+        ax2.set_ylabel("$r_{\\mathrm{max}}(V)$ [km]", color="tab:purple")
+        ax2b = ax2.twinx()
+        ax2b.plot(vv, [curve.r0(x) / 1000 for x in vv], color="tab:gray")
+        ax2b.set_ylabel("$r_0(V)$ [km]", color="tab:gray")
+        ax2.set_xlabel("Maximum gradient wind speed $V$ [m s$^{-1}$]")
+    else:
+        ax.set_xlabel("Maximum gradient wind speed $V$ [m s$^{-1}$]")
+
+    out = _fig_path(exp_dir, "vmax_sweep.pdf")
+    plt.tight_layout()
+    plt.savefig(out)
+    plt.close(fig)
+    print(f"plot_vmax_sweep: wrote {out}")
+    return out
+
+
+def plot_4d_samples(
+    exp_dir: str, init_steps: Optional[int] = None, curve_nc: Optional[str] = None
+) -> str:
+    """Plot surge vs sampled intensity for a 4D BO experiment.
+
+    Args:
+        exp_dir (str): BO experiment directory.
+        init_steps (int, optional): Number of initial-design samples; defaults
+            to the value in bo-config.json.
+        curve_nc (str, optional): Curve for the V_p marker; defaults to the
+            path in bo-config.json.
+
+    Returns:
+        str: Path of the written figure.
+    """
+    ledger = _load_ledger(exp_dir)
+    bo_cfg_path = os.path.join(exp_dir, "bo-config.json")
+    if os.path.exists(bo_cfg_path):
+        with open(bo_cfg_path, "r", encoding="utf-8") as f:
+            bo_cfg = json.load(f)
+        init_steps = init_steps if init_steps is not None else bo_cfg.get("init_steps")
+        curve_nc = curve_nc if curve_nc is not None else bo_cfg.get("curve_path")
+    curve = _curve_context(curve_nc)
+
+    calls = np.array(sorted(int(k) for k in ledger))
+    v = np.array([ledger[str(c)]["vmax"] for c in calls])
+    z = np.array([ledger[str(c)]["res"] for c in calls])
+    is_init = (
+        calls < init_steps if init_steps is not None else np.ones_like(calls, bool)
+    )
+
+    fig, ax = plt.subplots(figsize=(6, 4))
+    ax.scatter(
+        v[is_init], z[is_init], marker="x", color="tab:blue", label="initial design"
+    )
+    if (~is_init).any():
+        sc = ax.scatter(
+            v[~is_init],
+            z[~is_init],
+            c=calls[~is_init],
+            marker="o",
+            cmap="viridis",
+            label="acquisition",
+        )
+        plt.colorbar(sc, ax=ax, label="evaluation index")
+    if len(z):
+        i_best = int(np.nanargmax(z))
+        ax.plot(v[i_best], z[i_best], "*", color="tab:orange", markersize=15)
+    if curve is not None:
+        ax.axvline(curve.v_max, color="green", linestyle="--", alpha=0.7)
+        ax.annotate(
+            "$V_p$", (curve.v_max, ax.get_ylim()[0]), color="green", ha="right"
+        )
+    ax.set_xlabel("Maximum gradient wind speed $V$ [m s$^{-1}$]")
+    ax.set_ylabel("Max SSH at observation point [m]")
+    ax.legend()
+
+    out = _fig_path(exp_dir, "bo_4d_samples.pdf")
+    plt.tight_layout()
+    plt.savefig(out)
+    plt.close(fig)
+    print(f"plot_4d_samples: wrote {out}")
+    return out
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--exp_dir", type=str, required=True)
+    parser.add_argument("--mode", type=str, choices=["sweep", "4d"], default="sweep")
+    parser.add_argument("--curve_nc", type=str, default=None)
+    args = parser.parse_args()
+    if args.mode == "sweep":
+        plot_vmax_sweep(args.exp_dir, curve_nc=args.curve_nc)
+    else:
+        plot_4d_samples(args.exp_dir, curve_nc=args.curve_nc)