Iterative Gaussian kernel smoothing .

gaussian_kernel_regression.py

import numpy as np


def gaussian_kernel_regression(
    x, y, sigma, yerr=None, w=None, x_new=None, niter=20, atol=1e-6, rtol=1e-3
):
    """
    Iterative Gaussian kernel smoothing.

    Parameters
    ----------
    x, y : arrays, known data points
    sigma : float, Gaussian kernel width
    yerr : array, optional, measurement errors of y
    w : array, optional, initial weights of y
    x_new : array, optional, points to evaluate smoothed values
    niter : int, maximum iterations
    atol, rtol : float, absolute and relative tolerance for convergence

    Returns
    -------
    y_smooth_out : array
        Smoothed values at x_new (if provided) or at x (if x_new is None).
    sigma_tot_out : array
        Local total scatter at x_new or at x.
    info : dict
        A dictionary containing iteration info:
        - "niter": (int) Number of iterations performed.
        - "converged": (bool) True if the iteration converged.
        - "y_smooth": (array) Smoothed values at x.
        - "sigma_tot": (array) Local total scatter at x.
        - "sigma_int": (array) Inferred intrinsic scatter at x.
    """
    x, y = np.asarray(x), np.asarray(y)
    N = len(x)

    if yerr is None:
        yerr = np.zeros(N)
        var_sum = np.ones(N)
    else:
        yerr = np.asarray(yerr)
        if (yerr <= 0).any():
            raise ValueError("yerr should be all positive")
        var_sum = yerr**2

    w = np.ones(N) if w is None else np.asarray(w)
    if (w < 0).any():
        raise ValueError("w should be all nonnegative")

    # Gaussian kernel matrix between data points
    dx = x[:, None] - x[None, :]
    K = np.exp(-0.5 * (dx / sigma) ** 2)

    y_smooth = y
    converged = False
    for it in range(niter):
        y_old = y_smooth

        W = K * (w / var_sum)
        W /= W.sum(axis=1, keepdims=True).clip(1e-300)

        y_smooth = W @ y
        resid_sq = (y - y_smooth) ** 2

        # Estimate total variance at x locations
        var_tot = (K @ resid_sq) / K.sum(axis=1).clip(1e-300)
        var_int = (var_tot - yerr**2).clip(0)
        var_sum = yerr**2 + var_int  # total variance used for weights

        if np.allclose(y_smooth, y_old, atol=atol, rtol=rtol):
            converged = True
            break

    if x_new is not None:
        x_new = np.asarray(x_new)
        dxq = x_new[:, None] - x[None, :]
        Kq = np.exp(-0.5 * (dxq / sigma) ** 2)

        Wq = Kq * (w / var_sum)
        Wq /= Wq.sum(axis=1, keepdims=True).clip(1e-300)

        y_smooth_out = Wq @ y

        var_tot_new = (Kq @ resid_sq) / Kq.sum(axis=1).clip(1e-300)
        sigma_tot_out = np.sqrt(var_tot_new)
    else:
        y_smooth_out = y_smooth
        sigma_tot_out = np.sqrt(var_tot)

    info = {
        "niter": it + 1,
        "converged": converged,
        "y_smooth": y_smooth,
        "sigma_tot": np.sqrt(var_tot),
        "sigma_int": np.sqrt(var_int),
    }
    return y_smooth_out, sigma_tot_out, info


if __name__ == "__main__":
    import matplotlib.pyplot as plt

    np.random.seed(42)
    N = 50
    x = np.linspace(0, 10, N)
    y_true = np.sin(x) + 0.5 * np.cos(2 * x)
    yerr = 0.1 + 0.2 * np.random.rand(N)
    y_obs = y_true + np.random.normal(0, yerr)
    w = np.ones_like(y_obs)
    x_new = np.linspace(0, 10, 200)

    y_new, sigma_tot_new, info_new = gaussian_kernel_regression(
        x, y_obs, sigma=0.25, yerr=yerr, w=w, x_new=x_new
    )

    plt.figure(figsize=(8, 5))
    plt.errorbar(x, y_obs, yerr=yerr, fmt="o", label="Noisy data", alpha=0.6, capsize=3)
    plt.plot(x, y_true, "k--", label="True function")
    plt.plot(x_new, y_new, "r-", label="Smoothed", lw=2)
    plt.fill_between(
        x_new,
        y_new - sigma_tot_new,
        y_new + sigma_tot_new,
        color="r",
        alpha=0.2,
        label="Total scatter band",
    )
    plt.xlabel("x")
    plt.ylabel("y")
    plt.legend()
    plt.grid(True, linestyle=":", alpha=0.6)
    plt.tight_layout()
    plt.show()

    print(f"Converged: {info_new['converged']}, niter = {info_new['niter']}")