perturb-interp 0.1.0__py3-none-any.whl

perturb/__init__.py

@@ -0,0 +1,22 @@
+"""Perturb: Gaussian perturbation interpretability for any black-box function.
+
+The package follows the four steps of the method, with one module for each (and the perturbation split over two):
+
+    kernel.py    step 1, choosing the kernel width from how quickly the correlation between features falls off
+    ood.py       step 2, choosing how far to perturb, by checking how much of each perturbation the retrieval fits
+    grid.py      step 3, building the magnitude grid and the perturbed samples
+    perturb.py   step 3, running the black-box over the perturbed samples (Perturber) to get the response, its
+                 gradient and the importance
+    plotting.py  step 4, plotting the results
+"""
+from .grid import magnitude_grid, perturbation_grid, perturbed_samples
+from .kernel import correlation_matrix, fit_kernel_width, gaussian_window
+from .ood import fraction_fitted, in_distribution_curve, max_magnitude
+from .perturb import PerturbResult, Perturber
+
+__all__ = [
+    "magnitude_grid", "perturbation_grid", "perturbed_samples",
+    "correlation_matrix", "fit_kernel_width", "gaussian_window",
+    "fraction_fitted", "in_distribution_curve", "max_magnitude",
+    "PerturbResult", "Perturber",
+]

perturb/grid.py

@@ -0,0 +1,141 @@
+"""Step 3: building the perturbation grid.
+
+For a sample x* with noise sigma, we perturb the region around feature m by adding a Gaussian window w of FWHM eps,
+scaled by each magnitude xi_k in turn (in units of sigma_m),
+
+    X'_m = 1_K x* + sigma_m xi w(eps, m).
+
+This gives one perturbed copy of x* per magnitude. Everything here is pure numpy.
+"""
+import numpy as np
+
+from .kernel import gaussian_window
+
+
+def magnitude_grid(xi_max=3.0, n=101):
+    """Build the grid of perturbation magnitudes.
+
+    Parameters
+    ----------
+    xi_max : float, optional
+        The largest magnitude, in units of sigma. The grid runs from -xi_max to xi_max.
+    n : int, optional
+        The number of magnitudes. An even n is bumped up by one, and anything below 3 becomes 3, so that xi = 0 is
+        always on the grid.
+
+    Returns
+    -------
+    ndarray, shape (K,)
+        The magnitudes, linearly spaced.
+    """
+    if n % 2 == 0:
+        n += 1
+    n = max(n, 3)
+    return np.round(np.linspace(-xi_max, xi_max, n), 10)
+
+
+def perturbation_rows(x, sigma, x_axis, xi, fwhm, scale="centre", centres=None, rows=None):
+    """Build some or all of the rows of the flattened perturbation grid.
+
+    The full grid has shape (n_centres * K, M), where row ``j * K + k`` is centre j perturbed by magnitude xi_k.
+    Asking for only a subset of rows means the grid can be built a batch at a time, rather than all at once.
+
+    Parameters
+    ----------
+    x : array_like, shape (M,)
+        The sample being perturbed.
+    sigma : float or array_like, shape (M,)
+        The noise on each feature of x.
+    x_axis : array_like, shape (M,)
+        The position of each feature (e.g. wavelength).
+    xi : array_like, shape (K,)
+        The perturbation magnitudes, in units of sigma.
+    fwhm : float
+        The FWHM of the Gaussian window, in the units of x_axis.
+    scale : {"centre", "own"}, optional
+        How each feature in the window is scaled. "centre" moves every feature by the noise of the centre feature,
+        which keeps the smooth shape of the window. "own" moves each feature by its own noise instead, so that the
+        perturbation follows the noise from one feature to the next. This is worth using if the noise jumps sharply
+        between neighbouring features.
+    centres : array_like of int, optional
+        The indices of the centre features. By default every feature is used.
+    rows : array_like of int, optional
+        The flat row indices to build. By default every row is built.
+
+    Returns
+    -------
+    ndarray, shape (len(rows), M)
+        The perturbed samples.
+
+    Raises
+    ------
+    ValueError
+        If scale is not "centre" or "own".
+    """
+    if scale not in ("centre", "own"):
+        raise ValueError(f"scale must be 'centre' or 'own', not {scale!r}")
+
+    x = np.asarray(x, float)
+    x_axis = np.asarray(x_axis, float)
+    xi = np.atleast_1d(np.asarray(xi, float))
+    sigma = np.broadcast_to(np.asarray(sigma, float), x.shape)
+
+    if centres is None:
+        centres = np.arange(len(x))
+    else:
+        centres = np.atleast_1d(centres)
+
+    if rows is None:
+        rows = np.arange(len(centres) * len(xi))
+    else:
+        rows = np.asarray(rows)
+
+    # split each flat row index into its centre and its magnitude
+    j, k = np.divmod(rows, len(xi))
+    m = centres[j]
+
+    window = gaussian_window(x_axis[None, :], x_axis[m][:, None], fwhm)
+    if scale == "centre":
+        s = sigma[m][:, None]
+    else:
+        s = sigma[None, :]
+    return x[None, :] + xi[k][:, None] * s * window
+
+
+def perturbed_samples(x, sigma, x_axis, centre_index, xi, fwhm, scale="centre"):
+    """Perturb x around a single centre feature m, giving X'_m.
+
+    Parameters
+    ----------
+    x, sigma, x_axis, xi, fwhm, scale
+        As in `perturbation_rows`.
+    centre_index : int
+        The index of the centre feature m.
+
+    Returns
+    -------
+    ndarray, shape (K, M)
+        One perturbed copy of x per magnitude.
+    """
+    return perturbation_rows(x, sigma, x_axis, xi, fwhm, scale, [centre_index])
+
+
+def perturbation_grid(x, sigma, x_axis, xi, fwhm, scale="centre", centres=None):
+    """Perturb x around every centre feature.
+
+    Parameters
+    ----------
+    x, sigma, x_axis, xi, fwhm, scale, centres
+        As in `perturbation_rows`.
+
+    Returns
+    -------
+    ndarray, shape (n_centres, K, M)
+        The output of `perturbed_samples` for each centre in turn.
+    """
+    if centres is None:
+        n_centres = len(x)
+    else:
+        n_centres = len(np.atleast_1d(centres))
+    rows = perturbation_rows(x, sigma, x_axis, xi, fwhm, scale, centres)
+    return rows.reshape(n_centres, -1, len(x))

perturb/kernel.py

@@ -0,0 +1,181 @@
+"""Step 1: choosing the width of the Gaussian kernel.
+
+Neighbouring input features (e.g. adjacent wavelength bins) tend to move together, and the Pearson correlation between
+two features, taken across a set of samples drawn from the prior, falls off with their separation. Fitting a Gaussian
+to this fall-off gives us the FWHM of the perturbation window, i.e. the scale over which neighbouring features move
+together in real data.
+"""
+import numpy as np
+from scipy.optimize import curve_fit
+
+
+def gaussian_window(x_axis, centre, fwhm):
+    """The Gaussian window w(eps, m).
+
+    The window is defined on the feature axis itself rather than on the bin index, so unevenly spaced features are
+    handled properly.
+
+    Parameters
+    ----------
+    x_axis : array_like
+        The feature positions to evaluate the window at.
+    centre : float or array_like
+        The position of the peak, which broadcasts against x_axis.
+    fwhm : float
+        The full width at half maximum, in the units of x_axis.
+
+    Returns
+    -------
+    ndarray
+        The window, equal to 1 at the centre.
+    """
+    offset = np.asarray(x_axis, float) - centre
+    return np.exp(-4 * np.log(2) * offset ** 2 / fwhm ** 2)
+
+
+def minmax(X):
+    """Min-max normalise each sample over its features.
+
+    Parameters
+    ----------
+    X : ndarray, shape (..., M)
+        The samples, one per row.
+
+    Returns
+    -------
+    ndarray, shape (..., M)
+        Each sample rescaled to run from 0 to 1. A flat sample comes out as NaN, as it cannot be rescaled.
+    """
+    lo = X.min(axis=-1, keepdims=True)
+    hi = X.max(axis=-1, keepdims=True)
+    return (X - lo) / (hi - lo)
+
+
+def correlation_matrix(X, normalise=True):
+    """Calculate the Pearson correlation of every pair of features across a set of samples.
+
+    Parameters
+    ----------
+    X : array_like, shape (N, M)
+        N samples of M features.
+    normalise : bool, optional
+        Min-max normalise each sample first. This is needed when the overall level of a sample (e.g. the depth of
+        a spectrum) varies far more than its features do, as otherwise every pair of features would look
+        correlated. Flat samples, which cannot be normalised, are dropped.
+
+    Returns
+    -------
+    ndarray, shape (M, M)
+        The correlation matrix.
+    """
+    X = np.asarray(X, float)
+    if normalise:
+        with np.errstate(invalid="ignore", divide="ignore"):
+            X = minmax(X)
+        X = X[np.all(np.isfinite(X), axis=1)]
+    return np.corrcoef(X, rowvar=False)
+
+
+def correlation_vs_separation(corr, x_axis, bin_width, min_count=5):
+    """Bin the correlation of every pair of features by their separation.
+
+    Parameters
+    ----------
+    corr : array_like, shape (M, M)
+        The correlation matrix, from `correlation_matrix`.
+    x_axis : array_like, shape (M,)
+        The feature positions.
+    bin_width : float
+        The width of the separation bins, which are centred on zero.
+    min_count : int, optional
+        A bin needs more than this many pairs to be given an error bar.
+
+    Returns
+    -------
+    dict
+        Arrays over the non-empty bins, holding
+
+        - ``separation``: the bin centres
+        - ``rho``: the mean correlation
+        - ``rho_err``: the standard error of the mean, which is NaN for bins with min_count pairs or fewer
+        - ``n``: the number of pairs
+        - ``fit``: bool, the bins that can be used in the fit (those with a finite, non-zero error)
+    """
+    x_axis = np.asarray(x_axis, float)
+    sep = (x_axis[None, :] - x_axis[:, None]).ravel()
+    rho = np.asarray(corr, float).ravel()
+
+    keep = np.isfinite(rho)  # a constant feature has no correlation
+    sep = sep[keep]
+    rho = rho[keep]
+
+    n_side = np.ceil(np.abs(sep).max() / bin_width)
+    centres = np.arange(-n_side, n_side + 1) * bin_width
+    n_bins = len(centres)
+    idx = np.round(sep / bin_width).astype(int) + int(n_side)
+    idx = np.clip(idx, 0, n_bins - 1)
+
+    n = np.bincount(idx, minlength=n_bins)
+    mean = np.bincount(idx, rho, n_bins) / np.maximum(n, 1)
+    sum_sq = np.bincount(idx, (rho - mean[idx]) ** 2, n_bins)
+    with np.errstate(invalid="ignore", divide="ignore"):
+        # we use the standard error rather than the standard deviation, as the number of pairs falls with separation
+        std = np.sqrt(sum_sq / (n - 1))
+        err = np.where(n > min_count, std / np.sqrt(n), np.nan)
+
+    has = n > 0
+    out = dict(separation=centres[has], rho=mean[has], rho_err=err[has], n=n[has])
+    # the self-pairs (rho = 1, with no spread) fill the zero bin, but the Gaussian is 1 there for any FWHM anyway
+    out["fit"] = np.isfinite(out["rho_err"]) & (out["rho_err"] > 0)
+    return out
+
+
+def fit_kernel_width(X, x_axis, bin_width=None, min_count=5, normalise=True, p0=None):
+    """Fit the FWHM of the Gaussian kernel to a set of samples drawn from the prior.
+
+    Parameters
+    ----------
+    X : array_like, shape (N, M)
+        The samples from the prior.
+    x_axis : array_like, shape (M,)
+        The feature positions.
+    bin_width : float, optional
+        The width of the separation bins, which defaults to the median feature spacing.
+    min_count : int, optional
+        See `correlation_vs_separation`.
+    normalise : bool, optional
+        See `correlation_matrix`.
+    p0 : float, optional
+        A starting guess for the FWHM, which defaults to five bin widths.
+
+    Returns
+    -------
+    dict
+        The ``fwhm`` and ``fwhm_err``, the correlation matrix ``corr``, the binned curve ``binned`` (from
+        `correlation_vs_separation`), and the ``bin_width`` and ``x_axis`` used. This is what `plotting.kernel_fit`
+        and `plotting.correlation_matrix` expect.
+    """
+    x_axis = np.asarray(x_axis, float)
+    if bin_width is None:
+        bin_width = float(np.median(np.diff(np.sort(x_axis))))
+    if p0 is None:
+        p0 = 5 * bin_width
+
+    corr = correlation_matrix(X, normalise)
+    binned = correlation_vs_separation(corr, x_axis, bin_width, min_count)
+
+    def model(separation, fwhm):
+        return gaussian_window(separation, 0, fwhm)
+
+    use = binned["fit"]
+    popt, pcov = curve_fit(model, binned["separation"][use], binned["rho"][use], p0=[p0],
+                           sigma=binned["rho_err"][use], absolute_sigma=True)
+
+    return dict(
+        fwhm=abs(float(popt[0])),
+        fwhm_err=float(np.sqrt(pcov[0, 0])),
+        corr=corr,
+        binned=binned,
+        bin_width=bin_width,
+        x_axis=x_axis,
+    )

perturb/ood.py

@@ -0,0 +1,150 @@
+"""Step 2: how far can a sample be perturbed before it leaves the distribution? We ask the retrieval itself.
+
+For each perturbed sample, we feed f's answer back through the forward model g (outputs -> inputs), and compare this
+best fit with the perturbation that was applied. The fraction fitted,
+
+    F = dfit . dx / dx . dx    (in units of the noise),
+
+is the fraction of the perturbation dx that the best fit reproduces. By default the best fit is g(f(X')), but you can
+pass in your own best-fit spectra as best_fit instead (e.g. the maximum a posteriori spectra from a sampler).
+
+At small magnitudes, F is the retrieval's linear response to the perturbation. We count a perturbed sample as 'in
+distribution' while F stays within tol of that value, i.e. while the retrieval is still treating the perturbation the
+same way it treats a small one (`in_distribution_curve`). Everything here works from the PerturbResult, so f never
+needs to be called again.
+"""
+import numpy as np
+
+
+def fraction_fitted(result, forward=None, batch_size=None, best_fit=None):
+    """Calculate the fraction of each perturbation that the retrieval's best fit reproduces.
+
+    With dx = X' - x* and dfit = g(f(X')) - g(f(x*)), both in units of sigma, F = dfit . dx / dx . dx. F is 1 when
+    the best fit follows the perturbation exactly, and 0 when the retrieval treats it as noise.
+
+    By default the best fit is g(f(X')), the forward model at f's answer. However, if f returns a posterior mean,
+    this is not necessarily the best-fitting spectrum (the maximum a posteriori spectrum from a sampler would be, for
+    example). In that case, pass those spectra in as best_fit instead.
+
+    Parameters
+    ----------
+    result : PerturbResult
+        The perturbation run.
+    forward : callable, optional
+        The forward model g, which maps outputs (N, P) to inputs (N, M). This is needed unless best_fit is given.
+    batch_size : int, optional
+        The maximum number of rows per call to g, which also caps how many perturbed samples are held in memory at
+        once. By default everything is done in one call.
+    best_fit : array_like, shape (n_centres, K, M), optional
+        Your own best-fit spectrum for every perturbed sample, used in place of g(f(X')). The xi = 0 column must
+        hold the best fit to x* itself, which is the same for every centre, so only the first centre's is used.
+
+    Returns
+    -------
+    ndarray, shape (n_centres, K)
+        The fraction fitted, which is NaN at xi = 0.
+
+    Raises
+    ------
+    ValueError
+        If neither forward nor best_fit is given, or if best_fit has the wrong shape or no xi = 0 column.
+    """
+    n_outputs = result.response.shape[-1]
+    params = (result.baseline + result.response).reshape(-1, n_outputs)
+    if best_fit is not None:
+        best_fit = np.asarray(best_fit, float)
+        if best_fit.shape != result.response.shape[:2] + result.x.shape:
+            raise ValueError(f"best_fit must have shape (n_centres, K, M) = "
+                             f"{result.response.shape[:2] + result.x.shape}, not {best_fit.shape}")
+        if not (result.xi == 0).any():
+            raise ValueError("best_fit needs xi = 0 on the grid, as that column is the best fit to x* itself")
+        fit0 = best_fit[0, result.xi == 0][0]
+        best_fit = best_fit.reshape(len(params), -1)
+    elif forward is None:
+        raise ValueError("either forward or best_fit must be given")
+    else:
+        fit0 = np.asarray(forward(result.baseline[None]))[0]
+
+    n = len(params)
+    batch = batch_size or n
+    F = np.empty(n)
+    for start in range(0, n, batch):
+        rows = np.arange(start, min(start + batch, n))
+        dx = (result.samples(rows) - result.x) / result.sigma
+        fit = best_fit[rows] if best_fit is not None else np.asarray(forward(params[rows]))
+        dfit = (fit - fit0) / result.sigma
+        with np.errstate(invalid="ignore", divide="ignore"):
+            F[rows] = np.sum(dfit * dx, axis=-1) / np.sum(dx * dx, axis=-1)
+    return F.reshape(result.response.shape[:2])
+
+
+def in_distribution_curve(result, forward=None, tol=0.1, batch_size=None, best_fit=None):
+    """Find which of the perturbed samples are still in distribution.
+
+    A perturbed sample is in distribution if its fraction fitted is within tol of the value at the smallest non-zero
+    magnitudes (+-xi_min) for the same centre.
+
+    Parameters
+    ----------
+    result : PerturbResult
+        The perturbation run.
+    forward : callable, optional
+        The forward model g (see `fraction_fitted`). This is needed unless best_fit is given.
+    tol : float, optional
+        How far the fraction fitted is allowed to drift before a sample counts as out of distribution.
+    batch_size : int, optional
+        See `fraction_fitted`.
+    best_fit : array_like, shape (n_centres, K, M), optional
+        Your own best-fit spectra, used in place of forward (see `fraction_fitted`).
+
+    Returns
+    -------
+    dict
+        - ``xi``: (K,) the magnitudes
+        - ``values``: (n_centres, K) the fraction fitted
+        - ``inside``: (n_centres, K) bool, whether each sample is in distribution
+        - ``fraction_in``: (K,) the fraction of centres inside at each magnitude
+        - ``tol``: the tolerance used
+
+        `max_magnitude` reads the limit on |xi| from this.
+    """
+    xi = result.xi
+    values = fraction_fitted(result, forward, batch_size, best_fit)
+
+    nonzero = np.flatnonzero(xi != 0)
+    abs_xi = np.abs(xi[nonzero])
+    smallest = nonzero[abs_xi == abs_xi.min()]
+    ref = np.nanmean(values[:, smallest], axis=1, keepdims=True)
+
+    inside = (np.abs(values - ref) <= tol) | (xi == 0)
+    return dict(xi=xi, values=values, inside=inside, fraction_in=inside.mean(axis=0), tol=tol)
+
+
+def max_magnitude(curve, level=0.95):
+    """Find the largest |xi| that keeps enough of the perturbed samples in distribution.
+
+    At every magnitude up to the returned value, on both sides of zero, at least `level` of the perturbed samples
+    must be in distribution.
+
+    Parameters
+    ----------
+    curve : dict
+        From `in_distribution_curve`.
+    level : float, optional
+        The fraction of samples required to be in distribution.
+
+    Returns
+    -------
+    float
+        The limit on |xi|, or 0 if even the smallest magnitude fails.
+    """
+    abs_xi = np.abs(curve["xi"])
+    ok = curve["fraction_in"] >= level
+    if ok.all():
+        return float(abs_xi.max())
+
+    first_failure = abs_xi[~ok].min()
+    below = abs_xi[abs_xi < first_failure]
+    if below.size == 0:
+        return 0.0
+    return float(below.max())