assembly-economics/simulations/gadm02/gsa.py

"""Global sensitivity analysis of final VFAC: Morris screening + Sobol indices.

Morris (elementary effects, p = 4 levels) screens all parameters. Sobol
first-order (S1, Saltelli 2010 estimator) and total (ST, Jansen estimator)
indices with bootstrap
95 % intervals, are then computed for the top Morris parameters plus the
reuse bundle (rho0, IX, lin_open, std), so the "reuse works as a bundle"
reading in models/GADM-0.2.md can be tested directly.

Ranges: each parameter varies ±50 % around its scenario value; levers are
clipped to [0, 1]; zero-valued parameters vary in [0, 0.1].

Usage: python3 -B simulations/gadm02/gsa.py [--quick] [--model gadm04] [--jobs 8]

For models with power channels (gadm03, gadm03c, gadm04), the channel counter
coefficients c_* are not varied but recalibrated for every sample, so each
sample starts in equilibrium; the wealth_on switch is not varied.
"""

from __future__ import annotations

import importlib
import multiprocessing as mp
import random
import statistics
import sys
from pathlib import Path

sys.dont_write_bytecode = True

import gadm03  # noqa: E402
from scenarios import load_all  # noqa: E402

MODEL = sys.argv[sys.argv.index("--model") + 1] if "--model" in sys.argv else "gadm02"
JOBS = int(sys.argv[sys.argv.index("--jobs") + 1]) if "--jobs" in sys.argv else 1
g = importlib.import_module(MODEL)
CHANNELS = "Pm" in g.STOCKS
FIXED = {"wealth_on"} | ({c for c in gadm03.COEF.values()} if CHANNELS else set())

OUT = Path(__file__).resolve().parent / "results"
LEVERS = {"m", "L", "lin_open", "std", "disc", "expl", "x_att", "auto", "dshare", "cp", "od"}
BUNDLE = ["rho0", "IX", "lin_open", "std"]
SEED = 20261012


def ranges(base: dict) -> dict[str, tuple[float, float]]:
    out = {}
    for k, v in base.items():
        lo, hi = (0.0, 0.1) if v == 0 else (v * 0.5, v * 1.5)
        if k in LEVERS:
            lo, hi = max(0.0, lo), min(1.0, hi)
        out[k] = (lo, hi)
    return out


def model(base: dict, rng: dict, names: list[str], u: list[float], init: dict) -> float:
    p = dict(base)
    for k, x in zip(names, u):
        lo, hi = rng[k]
        p[k] = lo + x * (hi - lo)
    if CHANNELS:
        p.update(getattr(g, "calibrate", gadm03.calibrate)(p))
    return g.simulate(p, init)[-1]["VFAC"]


_CTX: dict = {}


def _init(ctx: dict) -> None:
    _CTX.update(ctx)


def _eval(u: list[float]) -> float:
    return model(_CTX["base"], _CTX["rng"], _CTX["names"], u, _CTX["init"])


def evaluate(base, rng, names, init, points: list[list[float]]) -> list[float]:
    """Evaluate many sample points, in parallel when --jobs > 1."""
    if JOBS <= 1:
        return [model(base, rng, names, u, init) for u in points]
    ctx = {"base": base, "rng": rng, "names": names, "init": init}
    with mp.Pool(JOBS, initializer=_init, initargs=(ctx,)) as pool:
        return pool.map(_eval, points, chunksize=8)


def morris(base, init, rng, r, rnd):
    names = [n for n in base if n not in FIXED]
    k, delta = len(names), 2 / 3
    effects = {n: [] for n in names}
    trajectories = []
    for _ in range(r):           # draw all trajectories first (same RNG order as before)
        x = [rnd.choice((0.0, 1 / 3)) for _ in range(k)]
        order = rnd.sample(range(k), k)
        pts = [list(x)]
        for i in order:
            x[i] += delta
            pts.append(list(x))
        trajectories.append((order, pts))
    flat = [pt for _o, pts in trajectories for pt in pts]
    ys = evaluate(base, rng, names, init, flat)
    pos = 0
    for order, pts in trajectories:
        y = ys[pos:pos + len(pts)]
        pos += len(pts)
        for j, i in enumerate(order):
            effects[names[i]].append((y[j + 1] - y[j]) / delta)
    res = []
    for n, ee in effects.items():
        res.append((n, statistics.fmean(abs(e) for e in ee), statistics.fmean(ee), statistics.pstdev(ee)))
    res.sort(key=lambda t: -t[1])
    return res


def sobol(base, init, rng, names, n, rnd, boot=200):
    k = len(names)
    A = [[rnd.random() for _ in range(k)] for _ in range(n)]
    B = [[rnd.random() for _ in range(k)] for _ in range(n)]
    pts = A + B + [A[j][:i] + [B[j][i]] + A[j][i + 1:] for i in range(k) for j in range(n)]
    ys = evaluate(base, rng, names, init, pts)
    fA, fB = ys[:n], ys[n:2 * n]
    fAB = [ys[2 * n + i * n: 2 * n + (i + 1) * n] for i in range(k)]

    def indices(idx):
        a = [fA[j] for j in idx]
        b = [fB[j] for j in idx]
        var = statistics.pvariance(a + b) or 1e-12
        mu = statistics.fmean(a + b)
        out = []
        for i in range(k):
            ab = [fAB[i][j] for j in idx]
            # Saltelli 2010 with centred outputs: VFAC has a large mean relative to its
            # spread, and centring removes that mean from the product's variance.
            s1 = sum((bb - mu) * (x - aa) for aa, bb, x in zip(a, b, ab)) / len(idx) / var
            st = sum((aa - x) ** 2 for aa, x in zip(a, ab)) / (2 * len(idx)) / var
            out.append((s1, st))
        return out

    point = indices(range(n))
    reps = [indices([rnd.randrange(n) for _ in range(n)]) for _ in range(boot)]

    def ci(i, which):
        vals = sorted(r[i][which] for r in reps)
        return vals[int(0.025 * boot)], vals[int(0.975 * boot) - 1]

    return [(names[i], point[i][0], ci(i, 0), point[i][1], ci(i, 1)) for i in range(k)]


def main() -> None:
    quick = "--quick" in sys.argv
    r, n, top = (4, 32, 6) if quick else (20, 256, 10)
    OUT.mkdir(exist_ok=True)
    scs = load_all()
    lines = [f"# {g.MODEL_VERSION} global sensitivity analysis", "",
             *([("Channel counters c_* are recalibrated per sample (not varied); "
                 "wealth_on is fixed."), ""] if CHANNELS else []),
             "Generated by `simulations/gadm02/gsa.py`" + (" (--quick)" if quick else "") + ".",
             f"Output: final VFAC (t = 100 {g.TIME_UNIT}s). Parameters vary ±50 % around the scenario",
             "value (levers clipped to [0, 1]; zero-valued parameters in [0, 0.1]).",
             f"Morris: {r} trajectories, 4 levels. Sobol: S1 Saltelli (2010, centred), ST Jansen, N = {n},",
             f"bootstrap 95 % intervals, seed {SEED}. Parameters outside the Sobol set stay at",
             "their scenario values.", "",
             "μ* ranks importance. A σ that is large relative to μ* indicates",
             "interactions or non-linearity. ST − S1 is the share of variance due to",
             "interactions involving that parameter.", ""]
    for name in ("baseline", "generative_regime"):
        rnd = random.Random(SEED)
        sc = scs[name]
        base = {**g.PARAMS, **sc["params"]}
        rng = ranges(base)
        mres = morris(base, sc["init"], rng, r, rnd)
        lines += [f"## {sc['title']}", "", "### Morris screening (top 15)", "",
                  "| Rank | Parameter | μ* | μ | σ |", "|---|---|---|---|---|"]
        lines += [f"| {i + 1} | `{p}` | {ms:.4f} | {mu:+.4f} | {sd:.4f} |"
                  for i, (p, ms, mu, sd) in enumerate(mres[:15])]
        mrank = {p: i + 1 for i, (p, *_rest) in enumerate(mres)}
        lines += ["", "Reuse bundle Morris ranks: " + ", ".join(f"`{b}` {mrank[b]}" for b in BUNDLE), ""]
        names = [p for p, *_rest in mres[:top]] + [b for b in BUNDLE if b not in [p for p, *_r in mres[:top]]]
        sres = sobol(base, sc["init"], rng, names, n, rnd)
        lines += ["### Sobol indices", "",
                  "| Parameter | S1 | S1 95 % CI | ST | ST 95 % CI | ST − S1 |", "|---|---|---|---|---|---|"]
        for p, s1, c1, st, ct in sorted(sres, key=lambda t: -t[3]):
            lines.append(f"| `{p}` | {s1:.3f} | [{c1[0]:.3f}, {c1[1]:.3f}] | {st:.3f} | "
                         f"[{ct[0]:.3f}, {ct[1]:.3f}] | {st - s1:+.3f} |")
        b1 = sum(s1 for p, s1, _c, _st, _ct in sres if p in BUNDLE)
        bt = sum(st for p, _s1, _c, st, _ct in sres if p in BUNDLE)
        tot1 = sum(s1 for _p, s1, *_r in sres)
        lines += ["", f"Sum of S1 over analysed parameters: {tot1:.3f}. Reuse bundle: "
                  f"ΣS1 = {b1:.3f}, ΣST = {bt:.3f}.", ""]
        print(f"done {name}", file=sys.stderr)
    stem = "gsa" if MODEL == "gadm02" else f"gsa_{MODEL}"
    (OUT / (f"{stem}-quick.md" if quick else f"{stem}.md")).write_text("\n".join(lines))
    print("\n".join(lines))


if __name__ == "__main__":
    main()