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 simulations/gadm02/gsa.py [--quick]
"""

from __future__ import annotations

import random
import statistics
import sys
from pathlib import Path

import gadm02 as g
from scenarios import load_all

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)
    return g.simulate(p, init)[-1]["VFAC"]


def morris(base, init, rng, r, rnd):
    names = list(base)
    k, delta = len(names), 2 / 3
    effects = {n: [] for n in names}
    for _ in range(r):
        x = [rnd.choice((0.0, 1 / 3)) for _ in range(k)]
        y = model(base, rng, names, x, init)
        for i in rnd.sample(range(k), k):
            x[i] += delta
            y2 = model(base, rng, names, x, init)
            effects[names[i]].append((y2 - y) / delta)
            y = y2
    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)]
    fA = [model(base, rng, names, a, init) for a in A]
    fB = [model(base, rng, names, b, init) for b in B]
    fAB = []
    for i in range(k):
        fAB.append([model(base, rng, names, A[j][:i] + [B[j][i]] + A[j][i + 1:], init) for j in range(n)])

    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", "",
             "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)
    (OUT / ("gsa-quick.md" if quick else "gsa.md")).write_text("\n".join(lines))
    print("\n".join(lines))


if __name__ == "__main__":
    main()