ssvkernel.py

# SPDX-License-Identifier: BSD-3-Clause
# Copyright (c) 2026 The Regents of the University of California, Mazin Lab

"""
Locally adaptive kernel-density estimate (Shimazaki & Shinomoto 2010).

Vectorized version of the original ssvkernel from mkidcalculator.  All
mathematical operations are identical; Python-level loops over the grid
have been replaced with numpy broadcasting and batched FFTs.

Reference:
    H. Shimazaki and S. Shinomoto, "Kernel Bandwidth Optimization in
    Spike Rate Estimation," Journal of Computational Neuroscience
    29(1-2): 171-182, 2010.  doi:10.1007/s10827-009-0180-4
"""

import numpy as np


def ssvkernel(x, tin=None, M=80, nbs=0, WinFunc='Gauss', scale=1):
    """
    Locally adaptive kernel-density estimate for 1-D data.

    Parameters
    ----------
    x : array_like
        Samples drawn from the underlying density.
    tin : array_like, optional
        Evaluation points for the output density.
    M : int
        Number of candidate window sizes.  Default 80.
    nbs : int
        Bootstrap replicates for confidence intervals.  0 = skip.
    WinFunc : str
        Window function: 'Gauss', 'Boxcar', 'Laplace', or 'Cauchy'.
    scale : float
        Multiplicative scale on the optimal bandwidth.

    Returns
    -------
    y, t, optw, gs, C, confb95, yb
    """
    # --- set up evaluation grid -----------------------------------------------
    if tin is None:
        T = np.max(x) - np.min(x)
        dx = np.sort(np.diff(np.sort(x)))
        dt_samp = dx[np.nonzero(dx)][0]
        tin = np.linspace(np.min(x), np.max(x),
                          int(min(np.ceil(T / dt_samp), 1e3)))
        t = tin
        x_ab = x[(x >= min(tin)) & (x <= max(tin))]
    else:
        T = np.max(x) - np.min(x)
        x_ab = x[(x >= min(tin)) & (x <= max(tin))]
        dx = np.sort(np.diff(np.sort(x)))
        dt_samp = dx[np.nonzero(dx)][0]
        if dt_samp > min(np.diff(tin)):
            t = np.linspace(min(tin), max(tin),
                            int(min(np.ceil(T / dt_samp), 1e3)))
        else:
            t = tin

    dt = min(np.diff(t))

    # finest histogram
    thist = np.concatenate((t, (t[-1] + dt)[np.newaxis]))
    y_hist = np.histogram(x_ab, thist - dt / 2)[0] / dt
    L = y_hist.size
    N = float(np.sum(y_hist * dt))

    # candidate window sizes
    W = logexp(np.linspace(ilogexp(5 * dt), ilogexp(T), M))

    # --- local cost functions (M individual FFT convolutions) -----------------
    c = np.zeros((M, L))
    for j in range(M):
        w = W[j]
        yh = _fftkernel(y_hist, w / dt)
        c[j, :] = yh ** 2 - 2 * yh * y_hist + 2 / (2 * np.pi) ** 0.5 / w * y_hist

    # --- optimal bandwidths (batched FFT: M batches instead of M² calls) ------
    optws = np.zeros((M, L))
    for i in range(M):
        C_local = _batch_fftkernelWin(c, W[i] / dt, WinFunc)  # (M, L)
        optws[i, :] = W[np.argmin(C_local, axis=0)]

    # --- golden section search for stiffness parameter ------------------------
    k = 0
    gs = np.zeros((30, 1))
    C_arr = np.zeros((30, 1))
    tol = 1e-5
    a = 1e-12
    b = 1
    phi = (5 ** 0.5 + 1) / 2
    c1 = (phi - 1) * a + (2 - phi) * b
    c2 = (2 - phi) * a + (phi - 1) * b
    f1 = _cost_function(y_hist, N, t, dt, optws, W, WinFunc, c1, scale=scale)[0]
    f2 = _cost_function(y_hist, N, t, dt, optws, W, WinFunc, c2, scale=scale)[0]

    yopt = None
    optw = None
    while (np.abs(b - a) > tol * (abs(c1) + abs(c2))) and (k < 30):
        if f1 < f2:
            b = c2
            c2 = c1
            c1 = (phi - 1) * a + (2 - phi) * b
            f2 = f1
            f1, yv1, optwp1 = _cost_function(
                y_hist, N, t, dt, optws, W, WinFunc, c1, scale=scale)
            yopt = yv1 / np.sum(yv1 * dt)
            optw = optwp1
        else:
            a = c1
            c1 = c2
            c2 = (2 - phi) * a + (phi - 1) * b
            f1 = f2
            f2, yv2, optwp2 = _cost_function(
                y_hist, N, t, dt, optws, W, WinFunc, c2, scale=scale)
            yopt = yv2 / np.sum(yv2 * dt)
            optw = optwp2

        gs[k] = c1
        C_arr[k] = f1
        k += 1

    gs = gs[:k]
    C_arr = C_arr[:k]

    # --- bootstrap confidence intervals (skip if nbs == 0) --------------------
    nbs = int(nbs)
    if nbs > 0:
        yb = np.zeros((nbs, tin.size))
        for i in range(nbs):
            Nb = np.random.poisson(lam=N)
            idx = np.random.randint(0, int(N), int(Nb))
            xb = x_ab[idx]
            y_histb = np.histogram(xb, thist - dt / 2)[0]
            nz = y_histb.nonzero()
            y_histb_nz = y_histb[nz]
            t_nz = t[nz]
            # vectorized balloon estimator
            t_diff = t[:, None] - t_nz[None, :]     # (L, nnz)
            G = _gauss_2d(t_diff, optw[:, None])     # (L, nnz)
            yb_buf = np.sum(y_histb_nz[None, :] * G, axis=1) / Nb
            yb_buf = yb_buf / np.sum(yb_buf * dt)
            yb[i, :] = np.interp(tin, t, yb_buf)
        ybsort = np.sort(yb, axis=0)
        y95b = ybsort[int(np.floor(0.05 * nbs)), :]
        y95u = ybsort[int(np.floor(0.95 * nbs)), :]
        confb95 = np.stack([y95b, y95u], axis=0)
    else:
        yb = np.empty((0, tin.size))
        confb95 = np.full((2, tin.size), np.nan)

    y = np.interp(tin, t, yopt)
    optw = np.interp(tin, t, optw)
    t = tin

    return y, t, optw, gs, C_arr, confb95, yb


# ---------------------------------------------------------------------------
# Cost function (vectorized — no Python loops over L)
# ---------------------------------------------------------------------------

def _cost_function(y_hist, N, t, dt, optws, WIN, WinFunc, g, scale=1.0):
    L = y_hist.size
    M = WIN.size

    # --- optwv: optimal variable bandwidth per grid point ---
    gs_all = optws / WIN[:, None]          # (M, L)
    gs_max = gs_all.max(axis=0)            # (L,)
    gs_min = gs_all.min(axis=0)            # (L,)

    # last index (highest row) where gs >= g, per column
    mask_ge = gs_all >= g                  # (M, L)
    # reversed argmax gives index of last True
    last_idx = M - 1 - np.argmax(mask_ge[::-1], axis=0)

    optwv = np.where(
        g > gs_max, WIN.min(),
        np.where(g < gs_min, WIN.max(), g * WIN[last_idx]))

    # --- Nadaraya-Watson kernel regression (vectorized) ---
    w_nw = optwv / g                       # (L,)
    t_diff = t[:, None] - t[None, :]       # (L, L)

    if WinFunc == 'Boxcar':
        Z = _boxcar_2d(t_diff, w_nw[None, :])
    elif WinFunc == 'Laplace':
        Z = _laplace_2d(t_diff, w_nw[None, :])
    elif WinFunc == 'Cauchy':
        Z = _cauchy_2d(t_diff, w_nw[None, :])
    else:
        Z = _gauss_2d(t_diff, w_nw[None, :])

    optwp = np.sum(optwv[None, :] * Z, axis=1) / np.sum(Z, axis=1) * scale

    # --- balloon estimator (vectorized) ---
    nz = y_hist.nonzero()
    y_hist_nz = y_hist[nz]
    t_nz = t[nz]
    t_diff_nz = t[:, None] - t_nz[None, :]   # (L, nnz)
    G = _gauss_2d(t_diff_nz, optwp[:, None])  # (L, nnz)
    yv = np.sum(y_hist_nz[None, :] * dt * G, axis=1)
    yv = yv * N / np.sum(yv * dt)

    # cost
    cg = yv ** 2 - 2 * yv * y_hist + 2 / (2 * np.pi) ** 0.5 / optwp * y_hist
    Cg = np.sum(cg * dt)

    return Cg, yv, optwp


# ---------------------------------------------------------------------------
# Batched FFT convolution
# ---------------------------------------------------------------------------

def _batch_fftkernelWin(rows, w, WinFunc):
    """Convolve each row of *rows* with the same kernel of width *w*.

    Replaces M individual fftkernelWin calls with one batched FFT.
    """
    L = rows.shape[1]
    Lmax = L + 3 * w
    n = int(2 ** np.ceil(np.log2(Lmax)))

    X = np.fft.fft(rows, n, axis=1)       # (M, n)

    f = np.linspace(0, n - 1, n) / n
    f = np.concatenate((-f[:n // 2 + 1], f[1:n // 2][::-1]))
    t = 2 * np.pi * f

    if WinFunc == 'Boxcar':
        a = 12 ** 0.5 * w
        K = 2 * np.sin(a * t / 2) / (a * t)
        K[0] = 1
    elif WinFunc == 'Laplace':
        K = 1 / (1 + (w * 2 * np.pi * f) ** 2 / 2)
    elif WinFunc == 'Cauchy':
        K = np.exp(-w * np.abs(2 * np.pi * f))
    else:
        K = np.exp(-0.5 * (w * 2 * np.pi * f) ** 2)

    y = np.real(np.fft.ifft(X * K[None, :], n, axis=1))
    return y[:, :L]


def _fftkernel(x, w):
    """FFT-based Gaussian convolution (single vector)."""
    L = x.size
    Lmax = L + 3 * w
    n = int(2 ** np.ceil(np.log2(Lmax)))
    X = np.fft.fft(x, n)
    f = np.linspace(0, n - 1, n) / n
    f = np.concatenate((-f[:n // 2 + 1], f[1:n // 2][::-1]))
    K = np.exp(-0.5 * (w * 2 * np.pi * f) ** 2)
    y = np.real(np.fft.ifft(X * K, n))
    return y[:L]


# ---------------------------------------------------------------------------
# Kernel functions (2-D broadcasting versions for vectorized loops)
# ---------------------------------------------------------------------------

def _gauss_2d(x, w):
    return 1 / (2 * np.pi) ** 2 / w * np.exp(-x ** 2 / 2 / w ** 2)


def _laplace_2d(x, w):
    return 1 / 2 ** 0.5 / w * np.exp(-2 ** 0.5 / w / np.abs(x))


def _cauchy_2d(x, w):
    return 1 / (np.pi * w * (1 + (x / w) ** 2))


def _boxcar_2d(x, w):
    a = 12 ** 0.5 * w
    y = np.where(np.abs(x) <= a / 2, 1 / a, 0.0)
    return y


# ---------------------------------------------------------------------------
# Utility functions
# ---------------------------------------------------------------------------

def logexp(x):
    y = np.zeros(x.shape)
    y[x < 1e2] = np.log(1 + np.exp(x[x < 1e2]))
    y[x >= 1e2] = x[x >= 1e2]
    return y


def ilogexp(x):
    y = np.zeros(x.shape)
    y[x < 1e2] = np.log(np.exp(x[x < 1e2]) - 1)
    y[x >= 1e2] = x[x >= 1e2]
    return y