contextattr 0.1.0__py3-none-any.whl

contextattr/__init__.py

@@ -0,0 +1,22 @@
+"""ContextAttr: which parts of a prompt make a model give an output (docs/interface.md, docs/methods.md)."""
+
+from .cache import Cache
+from . import responses
+from .evaluate import faithfulness, intervene
+from .methods import LIME, Banzhaf, ContextCite, KernelSHAP, LeaveOneOut
+from .prompt import Group, Message, Prompt, Response, Unit, read_prompts, split, write_prompts
+from .responses import Verdict, sample
+from .result import Explanation
+from .run import Results, Run, attribute, explain
+from .targets import Choice, Contrast, LogProb, Pooled, Sampled
+from .value import Value
+from .wrapper import ModelReport, Sampling, Scored, ToyModel, WrapperError, check_model
+
+__version__ = "0.1.0"
+
+checks = responses   # ca.checks.regex(...), ca.checks.reference(...)
+
+__all__ = ["LIME", "Banzhaf", "Cache", "Choice", "ContextCite", "Contrast", "Explanation", "Group", "KernelSHAP",
+           "LeaveOneOut", "LogProb", "Message", "ModelReport", "Pooled", "Prompt", "Response", "Results", "Run", "Sampled", "Sampling",
+           "Scored", "ToyModel", "Unit", "Value", "Verdict", "WrapperError", "attribute", "check_model", "checks",
+           "explain", "faithfulness", "intervene", "read_prompts", "responses", "sample", "split", "write_prompts"]

contextattr/__main__.py

@@ -0,0 +1,7 @@
+"""python -m contextattr: the command line (cli.py)."""
+
+import sys
+
+from .cli import main
+
+sys.exit(main())

contextattr/cache.py

@@ -0,0 +1,64 @@
+"""The cache: every model answer stored once, under a key made from its request (docs/interface.md §7).
+
+One SQLite file per run, or memory when no path is given. A run that stops resumes from it, another method on the
+same masks reuses its answers, and nothing is asked of the model twice. The first answer stored for a key stays.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+import sqlite3
+from collections.abc import Iterable, Mapping
+
+_CHUNK = 500  # keys per SELECT, under SQLite's limit on query parameters
+
+
+class Cache:
+    """Answers by key. `key(request)` makes the key from a JSON-able request; answers are JSON-able values."""
+
+    def __init__(self, path: str | os.PathLike | None = None):
+        self.path = None if path is None else os.fspath(path)
+        if self.path is not None and os.path.dirname(self.path):
+            os.makedirs(os.path.dirname(self.path), exist_ok=True)
+        self._db = sqlite3.connect(":memory:" if self.path is None else self.path)
+        self._db.execute("CREATE TABLE IF NOT EXISTS answers "
+                         "(key TEXT PRIMARY KEY, kind TEXT NOT NULL, answer TEXT NOT NULL)")
+        self._db.commit()
+
+    @staticmethod
+    def key(request: Mapping) -> str:
+        """sha256 of the request as canonical JSON: the same request always gives the same key."""
+        text = json.dumps(request, sort_keys=True, ensure_ascii=False, separators=(",", ":"), allow_nan=False)
+        return hashlib.sha256(text.encode("utf-8")).hexdigest()
+
+    def get_many(self, keys: Iterable[str]) -> dict[str, object]:
+        """The stored answers among `keys`."""
+        keys = list(dict.fromkeys(keys))
+        found = {}
+        for i in range(0, len(keys), _CHUNK):
+            chunk = keys[i:i + _CHUNK]
+            rows = self._db.execute(f"SELECT key, answer FROM answers WHERE key IN ({','.join('?' * len(chunk))})",
+                                    chunk)
+            found.update((k, json.loads(a)) for k, a in rows)
+        return found
+
+    def put_many(self, kind: str, answers: Mapping[str, object]) -> None:
+        """Store answers and commit at once; a key already stored keeps its first answer."""
+        self._db.executemany("INSERT OR IGNORE INTO answers (key, kind, answer) VALUES (?, ?, ?)",
+                             [(k, kind, json.dumps(a, ensure_ascii=False, allow_nan=False))
+                              for k, a in answers.items()])
+        self._db.commit()
+
+    def __len__(self) -> int:
+        return self._db.execute("SELECT COUNT(*) FROM answers").fetchone()[0]
+
+    def close(self) -> None:
+        self._db.close()
+
+    def __enter__(self) -> Cache:
+        return self
+
+    def __exit__(self, *exc) -> None:
+        self.close()

contextattr/cli.py

@@ -0,0 +1,181 @@
+"""The command line (docs/interface.md §8):
+
+    contextattr plan   run.toml [--measure K]   masks, requests and tokens before any call (K real requests timed)
+    contextattr run    run.toml                 explain every prompt into the run directory; resumes a stopped run
+    contextattr report <run directory>          the HTML report again, and the summary by tag
+    contextattr template [my_wrapper.py]        a wrapper to fill in (wrapper/template.py), written or printed
+
+run.toml names the wrapper by import path, so one run file works wherever its wrapper does:
+
+    prompts = "prompts.jsonl"               # paths are relative to this file
+    out = "runs/rio"
+    [model]
+    use = "my_wrappers.py:MyModel"          # or "package.module:Class", or "contextattr:ToyModel.demo"
+    args = { model_id = "…" }               # given to it as keyword arguments
+    [run]                                   # optional: method, budget, test_masks, seed, replace, batch_size,
+    method = "contextcite"                  #           noise_repeats, removal
+    [target]                                # optional: kind = logprob | pooled | contrast | choice | sampled, with
+    kind = "contrast"                       #           response, reduce, label, a, b, samples, temperature, top_p,
+                                            #           max_tokens, seed
+    [check]                                 # optional: regex = "…" (negative = true) or reference = "contains"
+    [sample]                                # optional: n, temperature, top_p, max_tokens, seed
+"""
+
+from __future__ import annotations
+
+import argparse
+import importlib
+import importlib.util
+import json
+import sys
+import tomllib
+from pathlib import Path
+
+from .responses import reference, regex
+from .run import Results, Run
+from .targets import Choice, Contrast, LogProb, Pooled, Sampled
+from .wrapper.contract import Sampling
+
+TOP = {"prompts", "out", "model", "run", "target", "check", "sample"}
+RUN = {"method", "budget", "test_masks", "seed", "replace", "batch_size", "noise_repeats", "removal"}
+SAMPLING = {"n", "temperature", "top_p", "max_tokens", "seed", "stop"}
+
+
+def load(path) -> tuple[Run, Path]:
+    """The Run a run.toml describes, and its run directory."""
+    path = Path(path)
+    with open(path, "rb") as f:
+        cfg = tomllib.load(f)
+    _keys(cfg, TOP, {"prompts", "out", "model"}, path.name)
+    base = path.parent
+    model = build(cfg["model"], base)
+    check = None
+    if "check" in cfg:
+        c = dict(cfg["check"])
+        if "regex" in c:
+            _keys(c, {"regex", "negative"}, {"regex"}, "[check]")
+            negative = c.get("negative", True)
+            if not isinstance(c["regex"], str) or not isinstance(negative, bool):
+                raise ValueError("[check]: regex must be a string, and negative true or false")
+            check = regex(c["regex"], negative=negative)
+        elif "reference" in c:
+            _keys(c, {"reference"}, {"reference"}, "[check]")
+            if c["reference"] not in ("contains", "exact") or not isinstance(c["reference"], str):
+                raise ValueError(f"[check]: reference must be \"contains\" or \"exact\", not {c['reference']!r}")
+            check = reference(contains=c["reference"] == "contains")
+        else:
+            raise ValueError("[check] needs regex = \"…\" or reference = \"contains\" | \"exact\"")
+    sample = None
+    if "sample" in cfg:
+        _keys(cfg["sample"], SAMPLING, set(), "[sample]")
+        sample = Sampling(**{**{"n": 8, "temperature": 0.7, "top_p": 0.95}, **cfg["sample"]})
+    target = _target(cfg.get("target"), check)
+    run_args = dict(cfg.get("run", {}))
+    _keys(run_args, RUN, set(), "[run]")
+    if "removal" in run_args:
+        run_args["removal"] = tuple(run_args["removal"])
+    run = Run(model, base / cfg["prompts"], target=target, sample=sample, check=check, **run_args)
+    return run, base / cfg["out"]
+
+
+def build(spec: dict, base: Path = Path(".")):
+    """The wrapper `use` names ("module:attr" or "file.py:attr", attr may be dotted), called with `args`."""
+    _keys(spec, {"use", "args"}, {"use"}, "[model]")
+    where, _, attr = spec["use"].partition(":")
+    if not attr:
+        raise ValueError(f"[model] use must be \"module:attr\" or \"file.py:attr\", not {spec['use']!r}")
+    if where.endswith(".py"):
+        file = (base / where).resolve()
+        loader = importlib.util.spec_from_file_location(file.stem, file)
+        if loader is None:
+            raise ValueError(f"cannot load {file}")
+        module = importlib.util.module_from_spec(loader)
+        sys.modules.setdefault(file.stem, module)
+        loader.loader.exec_module(module)
+    else:
+        module = importlib.import_module(where)
+    obj = module
+    for part in attr.split("."):
+        obj = getattr(obj, part)
+    return obj(**spec.get("args", {}))
+
+
+def _target(spec, check):
+    if spec is None:
+        return None
+    spec = dict(spec)
+    kind = spec.pop("kind", None)
+    if kind == "logprob":
+        _keys(spec, {"response", "reduce"}, set(), "[target]")
+        return LogProb(spec.get("response"), reduce=spec.get("reduce", "mean"))
+    if kind == "pooled":
+        _keys(spec, {"label", "reduce"}, set(), "[target]")
+        return Pooled(spec.get("label", "negative"), reduce=spec.get("reduce", "mean"))
+    if kind == "contrast":
+        _keys(spec, {"reduce"}, set(), "[target]")
+        return Contrast(reduce=spec.get("reduce", "mean"))
+    if kind == "choice":
+        _keys(spec, {"a", "b"}, {"a", "b"}, "[target]")
+        return Choice(spec["a"], spec["b"])
+    if kind == "sampled":
+        if check is None:
+            raise ValueError("a sampled target needs a [check]")
+        _keys(spec, {"samples"} | SAMPLING - {"n"}, set(), "[target]")
+        samples = spec.pop("samples", 4)
+        sampling = Sampling(**{**{"temperature": 0.7, "top_p": 0.95}, **spec, "n": samples}) if spec else None
+        return Sampled(check, samples=samples, sampling=sampling)
+    raise ValueError(f"[target] kind must be logprob, pooled, contrast, choice or sampled, not {kind!r}")
+
+
+def _keys(d: dict, allowed: set, required: set, where: str) -> None:
+    unknown, missing = sorted(set(d) - allowed), sorted(required - set(d))
+    if unknown:
+        raise ValueError(f"{where}: unknown key {', '.join(unknown)} (allowed: {', '.join(sorted(allowed))})")
+    if missing:
+        raise ValueError(f"{where}: missing {', '.join(missing)}")
+
+
+def main(argv=None) -> int:
+    ap = argparse.ArgumentParser(prog="contextattr", description=__doc__.split("\n\n")[0])
+    sub = ap.add_subparsers(dest="command", required=True)
+    plan = sub.add_parser("plan", help="masks, requests and tokens before any call")
+    plan.add_argument("config")
+    plan.add_argument("--measure", type=int, default=0, help="time this many real requests of the first prompt")
+    run = sub.add_parser("run", help="explain every prompt into the run directory")
+    run.add_argument("config")
+    report = sub.add_parser("report", help="the HTML report and the summary of a run directory")
+    report.add_argument("directory")
+    template = sub.add_parser("template", help="a wrapper to fill in around your model: written to FILE, or printed")
+    template.add_argument("file", nargs="?")
+    args = ap.parse_args(argv)
+    if args.command == "template":
+        from importlib.resources import files  # noqa: PLC0415
+
+        text = files("contextattr.wrapper").joinpath("template.py").read_text(encoding="utf-8")
+        if args.file is None:
+            print(text, end="")
+            return 0
+        try:
+            with open(args.file, "x", encoding="utf-8") as f:    # "x": an existing file is never overwritten
+                f.write(text)
+        except FileExistsError:
+            print(f"contextattr template: {args.file} exists; give another name", file=sys.stderr)
+            return 1
+        print(f"wrote {args.file}: fill in score and/or generate, then run it (python {args.file}) to check it")
+        return 0
+    if args.command == "plan":
+        r, _ = load(args.config)
+        print(json.dumps(r.plan(measure=args.measure), indent=2))
+    elif args.command == "run":
+        r, out = load(args.config)
+        results = r.execute(out)
+        print(f"{len(results)} prompts explained in {out}; report: {out / 'report.html'}")
+    else:
+        results = Results.load(args.directory)
+        results.to_html(Path(args.directory) / "report.html")
+        print(json.dumps(results.summary(), indent=2))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

contextattr/evaluate.py

@@ -0,0 +1,137 @@
+"""Checks of an explanation (docs/methods.md; SPEX §6.1): how well its surrogate predicts the value on held-out
+masks (faithfulness R², eq. 7), how much the value changes when the units the surrogate ranks highest are removed
+(top-r removal, eq. 8), how well its largest terms fall on human-labelled units (Recovery@r, eq. 9), and what removing
+units does to the behaviour itself, sampled afresh (intervention, as SPEX B.5.2 did on the trolley prompt)."""
+
+from __future__ import annotations
+
+import itertools
+from math import comb
+
+import numpy as np
+
+from .methods.base import Surrogate
+from .methods.indices import terms as surrogate_terms
+
+TEST_STREAM = 7919   # the test masks' own random stream, apart from every method's
+EXACT_REMOVAL = 20_000
+
+
+def held_out_masks(n: int, count: int, seed: int = 0) -> np.ndarray:
+    """`count` masks keeping each unit with probability ½, from a stream of their own: the same for every method run
+    on a prompt with this seed, so their values are asked of the model once and then come from the cache."""
+    return np.random.default_rng([seed, TEST_STREAM]).random((count, n)) < 0.5
+
+
+def faithfulness(predicted, values) -> float:
+    """R² = 1 − Σ(f − f̂)² / Σ(f − f̄)² over held-out masks (SPEX eq. 7): 1 when the surrogate predicts every value,
+    0 for the mean alone, below 0 for worse; nan when every value is the same."""
+    predicted, values = np.asarray(predicted, dtype=float), np.asarray(values, dtype=float)
+    spread = ((values - values.mean()) ** 2).sum()
+    if spread == 0:
+        return float("nan")
+    return float(1 - ((values - predicted) ** 2).sum() / spread)
+
+
+def removal_mask(surrogate: Surrogate, r: int) -> np.ndarray:
+    """m* = argmax over masks with r units removed of |f̂(1) − f̂(m)| (SPEX eq. 8): exact over all C(n, r) masks up
+    to 20,000 of them, else greedy, removing one unit at a time."""
+    n = surrogate.n
+    if not 0 <= r <= n:
+        raise ValueError(f"r must lie in 0..{n}")
+    full = np.ones(n, dtype=bool)
+    top = surrogate.predict(full[None])[0]
+    if comb(n, r) <= EXACT_REMOVAL:
+        best, best_change = full, -1.0
+        for chunk in _chunks(itertools.combinations(range(n), r), 4096):
+            masks = np.ones((len(chunk), n), dtype=bool)
+            for row, removed in zip(masks, chunk):
+                row[list(removed)] = False
+            change = np.abs(top - surrogate.predict(masks))
+            k = int(change.argmax())
+            if change[k] > best_change:
+                best, best_change = masks[k], float(change[k])
+        return best
+    mask = full.copy()
+    for _ in range(r):
+        options = [i for i in range(n) if mask[i]]
+        trial = np.repeat(mask[None], len(options), axis=0)
+        trial[np.arange(len(options)), options] = False
+        mask = trial[int(np.abs(top - surrogate.predict(trial)).argmax())]
+    return mask
+
+
+def removal(surrogate: Surrogate, value, rs=(1, 2, 3)) -> list[dict]:
+    """Top-r removal (SPEX eq. 8): Rem(r) = |f(1) − f(m*)| / |f(1)|, with the value function asked for f(1) and each
+    f(m*). `change` is f(m*) − f(1), how the value moves when the units go, and `predicted_change` the surrogate's
+    f̂(m*) − f̂(1). Rem divides by |f(1)|, as SPEX's log-probabilities allow; for a value near 0 (a contrast) read
+    the change instead."""
+    n = surrogate.n
+    masks = [removal_mask(surrogate, r) for r in rs]
+    values = value(np.vstack([np.ones((1, n), dtype=bool), *[m[None] for m in masks]]))
+    f1 = values[0]
+    out = []
+    for r, m, f in zip(rs, masks, values[1:]):
+        out.append({"r": r, "removed": np.flatnonzero(~m).tolist(), "value": float(f), "change": float(f - f1),
+                    "predicted_change": float(surrogate.predict(m[None])[0] - surrogate.predict(np.ones((1, n)))[0]),
+                    "rem": float(abs(f1 - f) / abs(f1)) if f1 != 0 else float("nan")})
+    return out
+
+
+def recovery(surrogate: Surrogate, relevant, r: int) -> float:
+    """Recovery@r (SPEX eq. 9): (1/r) Σ_{i ≤ r} |S* ∩ S_i| / |S_i|, the S_i being the surrogate's r most important
+    terms, by absolute value (single units and interactions alike), and S* the human-labelled units. Only a non-zero
+    term is one the method found: with fewer than r of them the slots left add 0 (a sparse fit is not credited with
+    terms it gave no weight, in the order of the units). Terms tied in size at the r-th place share the slots left to
+    them, each adding its mean, so the order of the units never decides the result."""
+    if isinstance(r, bool) or not isinstance(r, (int, np.integer)) or r < 1:
+        raise ValueError(f"r must be a positive integer, not {r!r}")
+    relevant = {int(i) for i in relevant}
+    terms = surrogate_terms(surrogate)
+    if not all(np.isfinite(w) for w in terms.values()):
+        raise ValueError("recovery needs finite terms: the surrogate has a NaN or infinite one")
+    found = sorted(((abs(w), len(relevant & set(s)) / len(s)) for s, w in terms.items()
+                    if s and w != 0), key=lambda t: -t[0])
+    if len(found) <= r:
+        return float(sum(h for _, h in found) / r)
+    edge = found[r - 1][0]
+    above = [h for size, h in found if size > edge]
+    tied = [h for size, h in found if size == edge]
+    return float((sum(above) + (r - len(above)) * sum(tied) / len(tied)) / r)
+
+
+def _chunks(iterable, size):
+    chunk = []
+    for item in iterable:
+        chunk.append(item)
+        if len(chunk) == size:
+            yield chunk
+            chunk = []
+    if chunk:
+        yield chunk
+
+
+def intervene(prompt, model, check, drop, *, samples: int = 64, sampling=None, replace: str | None = None,
+              cache=None) -> dict:
+    """Remove units and sample again: the mean of the check's values (the share flagged negative, for a labelling
+    check) over `samples` fresh replies to the full prompt and to the prompt without `drop` (unit or group ids), the
+    change, and its standard error. Needs only the wrapper's generate."""
+    import math  # noqa: PLC0415
+
+    from .responses.checks import verdict  # noqa: PLC0415
+    from .targets.sampled import Sampled  # noqa: PLC0415
+    from .value import Value  # noqa: PLC0415
+
+    v = Value(prompt, model, Sampled(check, samples=samples, sampling=sampling), replace=replace, cache=cache)
+    replies = v.replies(np.vstack([np.ones(prompt.n, dtype=bool), prompt.mask(drop=drop)]))
+    rates = []
+    for rs in replies:
+        values = [x for x in (verdict(check(prompt, r)).value for r in rs) if x is not None]
+        if not values:
+            raise ValueError("the check labelled none of the replies")
+        mean = math.fsum(values) / len(values)
+        spread = float(np.std(values, ddof=1)) if len(values) > 1 else 0.0
+        rates.append((mean, spread / math.sqrt(len(values)), len(values)))
+    (before, se0, k0), (after, se1, k1) = rates
+    return {"dropped": [drop] if isinstance(drop, str) else list(drop), "before": before, "after": after,
+            "change": after - before, "se": math.hypot(se0, se1), "replies": [k0, k1], "requests": dict(v.stats)}

contextattr/methods/__init__.py

@@ -0,0 +1,27 @@
+"""Which masks to evaluate and which fit (docs/methods.md). A method is named or given as an object."""
+
+from .base import Method, Surrogate, all_masks
+from .indices import banzhaf, shapley, shapley_taylor, terms
+from .interactions import SPEX, FaithBanzhaf, FaithShap, ProxySPEX, ShapiqMethod
+from .kernelshap import KernelSHAP
+from .linear import LIME, Banzhaf, ContextCite
+from .loo import LeaveOneOut
+
+METHODS = {"loo": LeaveOneOut, "banzhaf": Banzhaf, "lime": LIME, "contextcite": ContextCite,
+           "kernelshap": KernelSHAP, "faithbanzhaf": FaithBanzhaf, "faithshap": FaithShap, "spex": SPEX,
+           "proxyspex": ProxySPEX}
+
+
+def get_method(method) -> Method:
+    """A Method from its name (loo, banzhaf, lime, contextcite, kernelshap, faithbanzhaf, faithshap, spex,
+    proxyspex) or the object itself."""
+    if isinstance(method, Method):
+        return method
+    if isinstance(method, str) and method in METHODS:
+        return METHODS[method]()
+    raise ValueError(f"method must be one of {', '.join(METHODS)} or a Method, not {method!r}")
+
+
+__all__ = ["LIME", "METHODS", "SPEX", "Banzhaf", "ContextCite", "FaithBanzhaf", "FaithShap", "KernelSHAP",
+           "LeaveOneOut", "Method", "ProxySPEX", "ShapiqMethod", "Surrogate", "all_masks", "banzhaf", "get_method",
+           "shapley", "shapley_taylor", "terms"]

contextattr/methods/base.py

@@ -0,0 +1,86 @@
+"""What every method is (docs/methods.md): a design of masks to evaluate, and a fit of the surrogate whose terms
+are the scores. The surrogate predicts the value of any mask, which is how its faithfulness is measured (eq. 7)."""
+
+from __future__ import annotations
+
+import itertools
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from types import MappingProxyType
+
+import numpy as np
+
+
+@dataclass(frozen=True)
+class Surrogate:
+    """f̂(m) = intercept + Σ_i scores[i]·m_i + Σ_S interactions[S]·Π_{i∈S} m_i, for masks of 0/1. A unit's score is
+    its effect alone; an interaction S, a tuple of two or more unit positions, is what the units do together."""
+
+    intercept: float
+    scores: np.ndarray
+    interactions: Mapping[tuple[int, ...], float] = field(default_factory=dict)
+
+    def __post_init__(self) -> None:
+        scores = np.asarray(self.scores, dtype=float).copy()
+        scores.flags.writeable = False
+        object.__setattr__(self, "scores", scores)
+        object.__setattr__(self, "intercept", float(self.intercept))
+        object.__setattr__(self, "interactions", MappingProxyType(
+            {tuple(int(i) for i in s): float(w) for s, w in dict(self.interactions).items()}))
+
+    @property
+    def n(self) -> int:
+        return len(self.scores)
+
+    def predict(self, masks) -> np.ndarray:
+        m = np.asarray(masks, dtype=float)
+        out = self.intercept + m @ self.scores
+        for s, w in self.interactions.items():
+            out = out + w * m[:, list(s)].prod(axis=1)
+        return out
+
+
+class Method:
+    """A method: `masks(n, budget, rng)` designs the masks to evaluate, `fit(masks, values)` gives the surrogate, and
+    `run(game, n, budget, rng)` does both against a game (masks → values). A method that chooses its masks as it
+    goes (the shapiq ones) overrides `run` instead."""
+
+    name = "method"
+
+    def masks(self, n: int, budget: int, rng: np.random.Generator) -> np.ndarray:
+        raise NotImplementedError
+
+    def fit(self, masks: np.ndarray, values: np.ndarray) -> Surrogate:
+        raise NotImplementedError
+
+    def run(self, game, n: int, budget: int, rng: np.random.Generator) -> tuple[Surrogate, int]:
+        """The surrogate, and how many masks it took."""
+        masks = self.masks(n, budget, rng)
+        return self.fit(masks, game(masks)), len(masks)
+
+    def __repr__(self) -> str:
+        return f"{type(self).__name__}()"
+
+
+def all_masks(n: int) -> np.ndarray:
+    """Every mask of n units (2^n rows), the empty one first and the full one last."""
+    return np.array(list(itertools.product([False, True], repeat=n)), dtype=bool).reshape(2 ** n, n)
+
+
+def check_budget(name: str, budget: int, least: int) -> None:
+    if isinstance(budget, bool) or not isinstance(budget, (int, np.integer)) or budget < least:
+        raise ValueError(f"{name} needs a budget of at least {least} masks, not {budget!r}")
+
+
+def least_squares(x: np.ndarray, y: np.ndarray, w: np.ndarray | None = None, *, refuse: str | None = None
+                  ) -> np.ndarray:
+    """The (weighted) least-squares coefficients of y on the columns of x. With `refuse`, a ValueError of that message
+    when the rows leave the coefficients undetermined (numerical rank below the number of columns), where lstsq would
+    silently give one fit of many."""
+    if w is not None:
+        root = np.sqrt(w)
+        x, y = x * root[:, None], y * root
+    coef, _, rank, _ = np.linalg.lstsq(x, y, rcond=None)
+    if refuse is not None and rank < x.shape[1]:
+        raise ValueError(refuse)
+    return coef

contextattr/methods/indices.py

@@ -0,0 +1,59 @@
+"""One surrogate read as the classic indices (SPEX App. C, Table 1). A surrogate's terms are Möbius coefficients a(S):
+f̂(m) = Σ_S a(S)·Π_{i∈S} m_i, the intercept being a(∅). From them, exactly:
+
+    Shapley value      φ_i = Σ_{S∋i} a(S) / |S|
+    Banzhaf value      β_i = Σ_{S∋i} a(S) / 2^(|S|−1)
+    Shapley-Taylor     of order k: a(S) for |S| < k, and Σ_{T⊇S} a(T) / C(|T|, k) for |S| = k  (App. C, eq. 58)
+"""
+
+from __future__ import annotations
+
+from math import comb
+
+import numpy as np
+
+from .base import Surrogate
+
+
+def terms(surrogate: Surrogate) -> dict[tuple[int, ...], float]:
+    """Every Möbius term of the surrogate: () → intercept, (i,) → score, S → interaction."""
+    out = {(): surrogate.intercept}
+    out.update({(i,): float(w) for i, w in enumerate(surrogate.scores)})
+    out.update(surrogate.interactions)
+    return out
+
+
+def shapley(surrogate: Surrogate) -> np.ndarray:
+    phi = np.zeros(surrogate.n)
+    for s, a in terms(surrogate).items():
+        for i in s:
+            phi[i] += a / len(s)
+    return phi
+
+
+def banzhaf(surrogate: Surrogate) -> np.ndarray:
+    beta = np.zeros(surrogate.n)
+    for s, a in terms(surrogate).items():
+        for i in s:
+            beta[i] += a / 2 ** (len(s) - 1)
+    return beta
+
+
+def shapley_taylor(surrogate: Surrogate, order: int = 2) -> dict[tuple[int, ...], float]:
+    """The Shapley-Taylor interaction index of the given order, every set of at most `order` units that has a value."""
+    if order < 1:
+        raise ValueError("order must be at least 1")
+    out: dict[tuple[int, ...], float] = {}
+    for t, a in terms(surrogate).items():
+        if len(t) < order:
+            out[t] = out.get(t, 0.0) + a
+        else:
+            share = a / comb(len(t), order)
+            for s in _subsets(t, order):
+                out[s] = out.get(s, 0.0) + share
+    return out
+
+
+def _subsets(t: tuple[int, ...], k: int):
+    from itertools import combinations
+    return combinations(t, k)