diffcb 0.1.0__py3-none-any.whl

dcb/kde.py

@@ -0,0 +1,394 @@
+"""
+dcb.kde — Differentiable Gaussian KDE and Derivative Utilities
+
+Two computation paths share the same formulas but differ in memory layout:
+
+  Dense path  (n × G matrix, small n)
+    gaussian_kde_grid, kde_derivatives, soft_mode_count_cross
+    — simple, autograd-friendly, O(n × G) memory.
+
+  Chunked path  (processes X in blocks, large n / GPU)
+    kde_derivatives_chunked, soft_mode_count_cross_from_derivs
+    — O(chunk × G) peak memory, same arithmetic.
+    kde_derivatives_chunked returns (f, f', f'') on the grid by
+    accumulating chunk contributions; soft_mode_count_cross_from_derivs
+    evaluates M̃_cross given pre-computed derivatives, enabling the
+    analytical IFT gradient without materialising the full n × G graph.
+
+All functions operate on 1D input tensors (univariate DCB).
+"""
+
+from __future__ import annotations
+
+import math
+
+import torch
+from torch import Tensor
+
+
+def gaussian_kde_grid(X: Tensor, h: float, grid: Tensor) -> Tensor:
+    """Evaluate the Gaussian KDE f_h on a grid of points.
+
+    Parameters
+    ----------
+    X : Tensor, shape (n,)
+        Observed data points.
+    h : float
+        Bandwidth (positive scalar).
+    grid : Tensor, shape (G,)
+        Evaluation points.
+
+    Returns
+    -------
+    Tensor, shape (G,)
+        f_h evaluated at each grid point.
+    """
+    # diff[i, j] = grid[j] - X[i], shape (n, G)
+    diff = grid.unsqueeze(0) - X.unsqueeze(1)
+    K = torch.exp(-0.5 * (diff / h) ** 2) / (math.sqrt(2 * math.pi) * h)
+    return K.mean(dim=0)  # (G,)
+
+
+def kde_derivatives(
+    X: Tensor, h: float, grid: Tensor
+) -> tuple[Tensor, Tensor, Tensor]:
+    """Evaluate the Gaussian KDE and its first two analytical derivatives on a grid.
+
+    Parameters
+    ----------
+    X : Tensor, shape (n,)
+        Observed data points.
+    h : float
+        Bandwidth (positive scalar).
+    grid : Tensor, shape (G,)
+        Evaluation points.
+
+    Returns
+    -------
+    (f, f_prime, f_double_prime) : tuple of Tensor, each shape (G,)
+        KDE value, first derivative, and second derivative at each grid point.
+        All outputs support autograd.
+    """
+    # diff[i, j] = grid[j] - X[i], shape (n, G)
+    diff = grid.unsqueeze(0) - X.unsqueeze(1)
+    K = torch.exp(-0.5 * (diff / h) ** 2) / (math.sqrt(2 * math.pi) * h)
+
+    f = K.mean(0)
+    f_prime = (-diff / h ** 2 * K).mean(0)
+    f_double_prime = ((diff ** 2 / h ** 4) - (1.0 / h ** 2)) * K
+    f_double_prime = f_double_prime.mean(0)
+
+    return f, f_prime, f_double_prime
+
+
+def kde_derivatives_chunked(
+    X: Tensor, h: float, grid: Tensor, chunk_size: int = 50_000
+) -> tuple[Tensor, Tensor, Tensor]:
+    """Chunked KDE derivatives — O(chunk_size × G) peak memory.
+
+    Identical arithmetic to `kde_derivatives` but accumulates contributions
+    from X in blocks of `chunk_size` so the peak GPU allocation is
+    O(chunk_size × G) instead of O(n × G).  Fully differentiable via autograd
+    (the chunk loop produces a valid computation graph).
+
+    Parameters
+    ----------
+    X : Tensor, shape (n,)
+    h : float
+    grid : Tensor, shape (G,)
+    chunk_size : int
+        Number of data points per block.  Default 50 000 uses ≈100 MB at
+        G = 512, float32.  Increase for better GPU utilisation; decrease if
+        GPU OOM.
+
+    Returns
+    -------
+    (f, f_prime, f_double_prime) : tuple of Tensor, each shape (G,)
+    """
+    n, G = X.shape[0], grid.shape[0]
+    f   = torch.zeros(G, dtype=X.dtype, device=X.device)
+    fp  = torch.zeros_like(f)
+    fpp = torch.zeros_like(f)
+    for start in range(0, n, chunk_size):
+        Xc   = X[start : start + chunk_size]
+        diff = grid.unsqueeze(0) - Xc.unsqueeze(1)          # (c, G)
+        K    = torch.exp(-0.5 * (diff / h) ** 2) / (math.sqrt(2 * math.pi) * h)
+        f.add_(K.sum(0))
+        fp.add_((-diff / h ** 2 * K).sum(0))
+        fpp.add_((((diff ** 2 / h ** 4) - (1.0 / h ** 2)) * K).sum(0))
+    f.div_(n); fp.div_(n); fpp.div_(n)
+    return f, fp, fpp
+
+
+def soft_mode_count_cross_from_derivs(
+    f: Tensor,
+    f_prime: Tensor,
+    f_double_prime: Tensor,
+    grid: Tensor,
+    eps: float,
+    tau: float,
+) -> Tensor:
+    """Evaluate M̃_cross from pre-computed KDE derivatives.
+
+    Same formula as `soft_mode_count_cross` but accepts (f, f', f'') directly
+    instead of recomputing them from X.  Used by the large-n analytical IFT
+    path: call `kde_derivatives_chunked` once to get (f, f', f''), then pass
+    them here with `requires_grad=True` to obtain ∂M̃/∂f' and ∂M̃/∂f'' via
+    autograd — only G-dimensional tensors need to stay in the graph.
+
+    Parameters
+    ----------
+    f, f_prime, f_double_prime : Tensor, each shape (G,)
+    grid : Tensor, shape (G,)
+    eps, tau : float
+
+    Returns
+    -------
+    Tensor, shape ()
+    """
+    dx = ((grid[-1] - grid[0]) / (grid.shape[0] - 1)).detach()
+    fpp_scale = f_double_prime.abs().median().clamp(min=1e-10)
+    eps_eff   = eps * fpp_scale * dx
+    tau_eff   = tau * fpp_scale
+
+    fp_j  = f_prime[:-1]
+    fp_j1 = f_prime[1:]
+    fpp_mid = (f_double_prime[:-1] + f_double_prime[1:]) / 2
+
+    pos_to_neg = torch.sigmoid( fp_j / eps_eff) * torch.sigmoid(-fp_j1 / eps_eff)
+    neg_to_pos = torch.sigmoid(-fp_j / eps_eff) * torch.sigmoid( fp_j1 / eps_eff)
+    crossing   = pos_to_neg + neg_to_pos
+    local_max  = torch.sigmoid(-fpp_mid / tau_eff)
+
+    f_peak = f.max().clamp(min=1e-10)
+    f_mid  = (f[:-1] + f[1:]) / 2
+    density_mask = torch.sigmoid((f_mid - 0.01 * f_peak) / (0.001 * f_peak))
+
+    return (crossing * local_max * density_mask).sum()
+
+
+def soft_mode_count_cross(
+    X: Tensor, h: float, grid: Tensor, eps: float, tau: float
+) -> Tensor:
+    """Compute the crossing-count soft mode count M̃_cross(h).
+
+    Counts positive-to-negative sign changes of f'_h that coincide with a local
+    maximum (f'' < 0), using sigmoid products to smoothly approximate the
+    discrete sign-change indicator.  Unlike the integral formula, M̃_cross is
+    bounded and converges to M(h) as ε → 0 on a dense grid, without the
+    divergence pathology of the integral formula in the unimodal regime.
+
+    Formula:
+        M̃_cross = Σ_{j=0}^{G-2}
+            [σ(f'_j/ε)·σ(−f'_{j+1}/ε) + σ(−f'_j/ε)·σ(f'_{j+1}/ε)]
+            · σ(−f''_mid,j / τ)
+
+    where f''_mid,j = (f''_j + f''_{j+1}) / 2 and σ is the logistic sigmoid.
+    The first bracket captures both + → − and − → + crossings; the second
+    bracket selects only local maxima (f'' < 0).  The result equals the number
+    of downward zero-crossings of f', i.e., the number of modes of f_h.
+
+    Parameters
+    ----------
+    X : Tensor, shape (n,)
+        Observed data points.
+    h : float
+        Bandwidth (positive scalar).
+    grid : Tensor, shape (G,)
+        Uniform evaluation grid.
+    eps : float
+        Sigmoid temperature controlling the sharpness of the zero-crossing
+        detector on f' (smaller = sharper, converges to true M(h) faster).
+    tau : float
+        Sigmoid temperature for the local-max selector on f''.
+
+    Returns
+    -------
+    Tensor, shape ()
+        Scalar soft mode count in [0, G−1], differentiable w.r.t. X.
+    """
+    f, f_prime, f_double_prime = kde_derivatives(X, h, grid)
+
+    # Calibrate eps to the local crossing scale, not the global f' amplitude.
+    # A single pos→neg crossing of f'_h happens over ~Δx in x-space; the f'
+    # value just outside the crossing pair is O(|f''_h| · Δx).  We need
+    # eps_eff << |f''_h| · Δx so non-crossing pairs saturate while the
+    # crossing pair contributes ~1.  Using eps_eff = eps · |f''_h| · Δx gives
+    # σ(|f'_crossing| / eps_eff) ≈ σ(1/(2·eps)) — well-calibrated for any h.
+    # Stop-gradient prevents these scale factors from entering the IFT backward.
+    dx = ((grid[-1] - grid[0]) / (grid.shape[0] - 1)).detach()
+    # fpp_scale retains gradient so the full ∂M/∂X and ∂M/∂h paths flow through
+    # the self-normalizing calibration — required for gradcheck to pass.
+    fpp_scale = f_double_prime.abs().median().clamp(min=1e-10)
+    eps_eff = eps * fpp_scale * dx   # dimensionless × curvature × grid-step
+    tau_eff = tau * fpp_scale        # dimensionless × curvature
+
+    fp_j  = f_prime[:-1]
+    fp_j1 = f_prime[1:]
+    fpp_mid = (f_double_prime[:-1] + f_double_prime[1:]) / 2
+
+    pos_to_neg = torch.sigmoid(fp_j / eps_eff) * torch.sigmoid(-fp_j1 / eps_eff)
+    neg_to_pos = torch.sigmoid(-fp_j / eps_eff) * torch.sigmoid(fp_j1 / eps_eff)
+    crossing   = pos_to_neg + neg_to_pos
+
+    local_max  = torch.sigmoid(-fpp_mid / tau_eff)
+
+    # Suppress spurious tail crossings: in low-density regions f'_h ≈ 0 but
+    # fluctuates, giving fractional sigmoid values on every grid pair.  Mask
+    # out intervals where the average density is below 1% of the peak.
+    # The threshold and sharpness are both relative to f_peak so this scales
+    # automatically with the KDE amplitude at any bandwidth.
+    f_peak = f.max().clamp(min=1e-10)
+    f_mid  = (f[:-1] + f[1:]) / 2
+    density_mask = torch.sigmoid((f_mid - 0.01 * f_peak) / (0.001 * f_peak))
+
+    return (crossing * local_max * density_mask).sum()
+
+
+def soft_mode_count(
+    X: Tensor, h: float, grid: Tensor, eps: float, tau: float
+) -> Tensor:
+    """Compute the soft (differentiable) mode count M̃(h).
+
+    Approximates the number of modes of f_h via the Riemann-sum integral
+
+        M̃(h) = Δx · Σ_j δ_ε(f'_h(grid_j)) · σ_τ(-f''_h(grid_j)) · |f''_h(grid_j)|
+
+    where δ_ε is a Gaussian approximation to the Dirac delta (peaking where
+    f' = 0) and σ_τ is a sigmoid that selects local maxima (f'' < 0).
+
+    Parameters
+    ----------
+    X : Tensor, shape (n,)
+        Observed data points.
+    h : float
+        Bandwidth (positive scalar).
+    grid : Tensor, shape (G,)
+        Uniform evaluation grid.
+    eps : float
+        Width of the Gaussian delta approximation (controls sharpness of the
+        zero-crossing detector on f').
+    tau : float
+        Temperature of the sigmoid (controls sharpness of the local-max selector
+        on f'').
+
+    Returns
+    -------
+    Tensor, shape ()
+        Scalar soft mode count, differentiable w.r.t. X.
+    """
+    _, f_prime, f_double_prime = kde_derivatives(X, h, grid)
+
+    # Gaussian approximation to Dirac delta centred at f' = 0
+    delta_eps = torch.exp(-0.5 * (f_prime / eps) ** 2) / (
+        math.sqrt(2 * math.pi) * eps
+    )
+
+    # Sigmoid selects points where f'' < 0, i.e. local maxima
+    sigma_tau = torch.sigmoid(-f_double_prime / tau)
+
+    integrand = delta_eps * sigma_tau * f_double_prime.abs()
+
+    dx = (grid[-1] - grid[0]) / (grid.shape[0] - 1)
+    return (integrand * dx).sum()
+
+
+def kde_derivatives_batched(
+    X: Tensor,
+    h_batch: Tensor,
+    grid: Tensor,
+    chunk_size: int = 10_000,
+) -> tuple[Tensor, Tensor, Tensor]:
+    """Evaluate Gaussian KDE derivatives for B bandwidths simultaneously.
+
+    Accumulates contributions from X in chunks to keep peak memory at
+    O(chunk_size × B × G) instead of O(n × B × G).
+
+    Parameters
+    ----------
+    X : Tensor, shape (n,)
+    h_batch : Tensor, shape (B,)
+        Candidate bandwidths.
+    grid : Tensor, shape (G,)
+    chunk_size : int
+        Number of data points per chunk. Default 10 000.
+
+    Returns
+    -------
+    (f, f_prime, f_double_prime) : tuple of Tensor, each shape (B, G)
+        KDE value and its first two derivatives, one row per bandwidth.
+    """
+    n, G, B = X.shape[0], grid.shape[0], h_batch.shape[0]
+    f   = torch.zeros(B, G, dtype=X.dtype, device=X.device)
+    fp  = torch.zeros_like(f)
+    fpp = torch.zeros_like(f)
+
+    # h_batch: (B,) → (1, B, 1) for broadcasting against (c, 1, G)
+    h = h_batch.view(1, B, 1)
+    h2 = h * h
+    norm = 1.0 / (math.sqrt(2 * math.pi) * h)  # (1, B, 1)
+
+    for start in range(0, n, chunk_size):
+        Xc = X[start : start + chunk_size]              # (c,)
+        # diff[i, j] = grid[j] - Xc[i] → (c, 1, G) broadcasts to (c, B, G)
+        diff = (grid.unsqueeze(0) - Xc.unsqueeze(1)).unsqueeze(1)  # (c, 1, G)
+        u = diff / h                                     # (c, B, G)
+        K = norm * torch.exp(-0.5 * u * u)              # (c, B, G)
+        f.add_(K.sum(0))                                 # (B, G)
+        fp.add_((-u / h * K).sum(0))                    # (B, G)
+        fpp.add_((((u * u - 1.0) / h2) * K).sum(0))    # (B, G)
+
+    f.div_(n); fp.div_(n); fpp.div_(n)
+    return f, fp, fpp
+
+
+def soft_mode_count_cross_batched(
+    f: Tensor,
+    fp: Tensor,
+    fpp: Tensor,
+    grid: Tensor,
+    eps: float,
+    tau: float,
+) -> Tensor:
+    """Vectorised M̃_cross over B bandwidths simultaneously.
+
+    Applies the same formula as `soft_mode_count_cross_from_derivs` but
+    operates on (B, G) tensors, returning one scalar per bandwidth.
+
+    Parameters
+    ----------
+    f, fp, fpp : Tensor, each shape (B, G)
+    grid : Tensor, shape (G,)
+    eps, tau : float
+
+    Returns
+    -------
+    Tensor, shape (B,)
+        One soft mode count per bandwidth.
+    """
+    dx = ((grid[-1] - grid[0]) / (grid.shape[0] - 1)).detach()
+
+    # Per-bandwidth calibration: median of |fpp| along grid axis
+    fpp_scale = fpp.abs().median(dim=1).values.clamp(min=1e-10)  # (B,)
+    eps_eff = eps * fpp_scale * dx       # (B,)
+    tau_eff = tau * fpp_scale            # (B,)
+
+    fp_j   = fp[:, :-1]                 # (B, G-1)
+    fp_j1  = fp[:, 1:]                  # (B, G-1)
+    fpp_mid = (fpp[:, :-1] + fpp[:, 1:]) / 2  # (B, G-1)
+
+    ee = eps_eff.unsqueeze(1)           # (B, 1)
+    te = tau_eff.unsqueeze(1)           # (B, 1)
+
+    pos_to_neg = torch.sigmoid( fp_j / ee) * torch.sigmoid(-fp_j1 / ee)
+    neg_to_pos = torch.sigmoid(-fp_j / ee) * torch.sigmoid( fp_j1 / ee)
+    crossing   = pos_to_neg + neg_to_pos
+    local_max  = torch.sigmoid(-fpp_mid / te)
+
+    f_peak = f.max(dim=1).values.clamp(min=1e-10)  # (B,)
+    f_mid  = (f[:, :-1] + f[:, 1:]) / 2            # (B, G-1)
+    fp_ref = (0.01 * f_peak).unsqueeze(1)           # (B, 1)
+    fw_ref = (0.001 * f_peak).unsqueeze(1)          # (B, 1)
+    density_mask = torch.sigmoid((f_mid - fp_ref) / fw_ref)
+
+    return (crossing * local_max * density_mask).sum(dim=1)  # (B,)

dcb/layer.py

@@ -0,0 +1,231 @@
+"""
+dcb.layer — DifferentiableCriticalBandwidth PyTorch Module
+
+This module defines `DCBLayer`, the central contribution of the DCB package:
+a `torch.nn.Module` that accepts a 1D tensor of samples X and returns the
+critical bandwidth h_crit as a differentiable scalar. The forward pass
+evaluates the soft mode-counting integral M̃(h; X) over a dense evaluation
+grid Ω using Gaussian KDE derivatives from `dcb.kde`, then invokes the
+bisection solver in `dcb.solver` to locate h_crit as the root of
+M̃(h) - target_modes = 0. The backward pass is handled entirely by the
+custom `DCBFunction` (a `torch.autograd.Function`) which applies the IFT
+formula: ∂h_crit/∂X = -(∂M̃/∂h)^{-1} · (∂M̃/∂X)|_{h=h_crit}, bypassing
+the computational graph of the iterative solver and maintaining O(1) memory
+cost relative to the number of solver iterations. Hyperparameters ε and τ
+may be supplied explicitly or computed adaptively via `dcb.utils`.
+"""
+
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+from dcb.solver import find_h_crit, ift_gradient
+from dcb.utils import make_grid, silverman_bandwidth, adaptive_eps_tau, anneal_eps_tau
+
+
+class DCBFunction(torch.autograd.Function):
+    """Custom autograd Function for differentiable critical bandwidth.
+
+    Both the forward root-finding and the IFT backward use the same (eps, tau)
+    and formula ('cross' by default — the crossing-count M̃_cross that fixes the
+    m=1 divergence of the legacy integral formula).
+    """
+
+    @staticmethod
+    def forward(ctx, X, grid, eps, tau, target_modes, delta, formula, chunk_size,
+                brentq_n_max, g_brentq, use_hard_bisection, safe_backward, use_fft):
+        """Locate h_crit and save state for the backward pass."""
+        h_crit, cond_num = find_h_crit(
+            X, grid, eps, tau, target_modes,
+            formula=formula, brentq_n_max=brentq_n_max, chunk_size=chunk_size,
+            g_brentq=g_brentq, use_hard_bisection=use_hard_bisection,
+            use_fft=use_fft,
+        )
+        ctx.save_for_backward(X, grid)
+        ctx.h_crit = h_crit
+        ctx.eps = eps
+        ctx.tau = tau
+        ctx.delta = delta
+        ctx.formula = formula
+        ctx.chunk_size = chunk_size
+        ctx.cond_num = cond_num
+        ctx.safe_backward = safe_backward
+        return torch.tensor(h_crit, dtype=X.dtype, device=X.device)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        """Apply the IFT formula to compute ∂h_crit/∂X."""
+        X, grid = ctx.saved_tensors
+        grad_X = ift_gradient(
+            X, ctx.h_crit, grid, ctx.eps, ctx.tau, grad_output,
+            delta=ctx.delta, formula=ctx.formula, chunk_size=ctx.chunk_size,
+            safe_backward=ctx.safe_backward,
+        )
+        ctx.guard_triggered = ift_gradient.last_guard_triggered
+        ctx.denom_abs       = ift_gradient.last_denom_abs
+        # Gradients for: X, grid, eps, tau, target_modes, delta, formula,
+        #                chunk_size, brentq_n_max, g_brentq, use_hard_bisection,
+        #                safe_backward, use_fft
+        return grad_X, None, None, None, None, None, None, None, None, None, None, None, None
+
+
+class DCBLayer(nn.Module):
+    """Differentiable Critical Bandwidth as a PyTorch Module.
+
+    Computes the critical bandwidth h_crit for a 1D sample X — the smallest
+    bandwidth at which a Gaussian KDE has exactly `target_modes` modes — and
+    returns it as a differentiable scalar. Gradients flow back to X via the
+    Implicit Function Theorem.
+
+    The `anneal_factor` parameter scales eps and tau together, tightening the
+    soft mode-count approximation M̃ toward the true CBW.  Use anneal_factor=1.0
+    (default) during gradient-based training for stable IFT gradients; reduce
+    toward 0.05–0.10 for eval-time accuracy (in torch.no_grad()).  When
+    anneal_factor < 0.1 increase G proportionally (e.g. G=2048) so the grid
+    resolves the sharpened Gaussian delta approximation.
+
+    Parameters
+    ----------
+    target_modes : int
+        Number of modes to target (default 1).
+    G : int
+        Number of evaluation grid points. Default 512.
+    delta : float
+        Stabilisation floor for the IFT denominator. Default 1e-4.
+    eps : float or None
+        Gaussian-delta width. If None, computed adaptively each call.
+    tau : float or None
+        Sigmoid temperature. If None, computed adaptively each call.
+    anneal_factor : float
+        Scales eps and tau; use 1.0 for training, <1 for eval accuracy.
+    formula : str
+        'cross' (default) — M̃_cross crossing-count formula, fixes m=1 bias.
+        'integral' — legacy integral formula (kept for ablation comparison).
+    g_brentq : int
+        Grid resolution for the brentq objective (default 128, Round 15a).
+        Ignored when use_hard_bisection=True.
+    use_hard_bisection : bool
+        If True (default, Round 15b), bisect on the hard mode count — provably
+        monotone, no false roots. If False, use legacy brentq on M̃_cross.
+    adaptive_G : bool
+        If True (Round 17a), compute grid resolution G dynamically each forward
+        call as max(G, min(32768, int(G * max(1.0, (n/1000)**0.2)))). This
+        scales G with sample size: G=512 at n≤1K, ~706 at n=5K, ~1119 at n=50K,
+        ~2038 at n=1M. If False (default), use the fixed G value.
+    safe_backward : bool
+        If True (Round 17b), clamp the IFT denominator |∂M̃/∂h| to at least 0.1
+        (10× larger than the default guard of 0.01) to limit gradient magnitude
+        near mode-merging bifurcations. Default False. A warning is always emitted
+        when the denominator falls below 0.01, regardless of this flag.
+    use_fft : bool
+        Default True. Uses FFT-based mode counting (O(n + G log G)) for n > 50K,
+        eliminating subsampling bias. Falls back to direct KDE for n ≤ 50K (no
+        bias at small n). Set False only for legacy/ablation comparison.
+
+    Examples
+    --------
+    >>> layer = DCBLayer(target_modes=1)
+    >>> X = torch.cat([torch.randn(50) - 1, torch.randn(50) + 1]).requires_grad_(True)
+    >>> h = layer(X)
+    >>> h.backward()
+    >>> X.grad.shape
+    torch.Size([100])
+    """
+
+    def __init__(
+        self,
+        target_modes: int = 1,
+        G: int = 512,
+        delta: float = 1e-4,
+        eps: float = None,
+        tau: float = None,
+        anneal_factor: float = 1.0,
+        formula: str = 'cross',
+        chunk_size: int = 50_000,
+        brentq_n_max: int = 50_000,
+        g_brentq: int = 128,
+        use_hard_bisection: bool = True,
+        adaptive_G: bool = False,
+        safe_backward: bool = False,
+        use_fft: bool = True,
+    ):
+        super().__init__()
+        self.target_modes = target_modes
+        self.G = G
+        self.delta = delta
+        self._eps = eps
+        self._tau = tau
+        self.anneal_factor = anneal_factor
+        self.formula = formula
+        self.chunk_size = chunk_size
+        self.brentq_n_max = brentq_n_max
+        self.g_brentq = g_brentq
+        self.use_hard_bisection = use_hard_bisection
+        self.adaptive_G = adaptive_G
+        self.safe_backward = safe_backward
+        self.use_fft = use_fft
+        if use_fft and brentq_n_max != 50_000:
+            raise TypeError(
+                f"brentq_n_max={brentq_n_max} is meaningless when use_fft=True: the FFT path "
+                "does not subsample and ignores this parameter. Remove brentq_n_max from your "
+                "DCBLayer constructor, or set use_fft=False to use the legacy subsampling path."
+            )
+        if not use_fft and brentq_n_max != 50_000:
+            import warnings
+            warnings.warn(
+                "brentq_n_max is deprecated. When use_fft=True (default), this parameter is "
+                "ignored. Set use_fft=True (default) instead of tuning brentq_n_max.",
+                DeprecationWarning, stacklevel=2,
+            )
+
+    def set_anneal_factor(self, factor: float) -> None:
+        """Update the sharpening factor in-place (call before eval)."""
+        self.anneal_factor = float(factor)
+
+    def forward(self, X: Tensor) -> Tensor:
+        """Compute h_crit for sample X.
+
+        Parameters
+        ----------
+        X : Tensor, shape (n,)
+            1D sample tensor. Must be on a single device; may require grad.
+
+        Returns
+        -------
+        Tensor, shape ()
+            Scalar h_crit, differentiable w.r.t. X.
+        """
+        n = X.shape[0]
+        G_eff = (
+            max(self.G, min(32768, int(self.G * max(1.0, (n / 1000) ** 0.2))))
+            if self.adaptive_G else self.G
+        )
+        grid = make_grid(X.detach(), G_eff)
+
+        if self.formula == 'cross':
+            # M̃_cross self-normalises by std(f'_h) internally; eps/tau are
+            # dimensionless fractions.  anneal_factor still sharpens them.
+            eps = self._eps if self._eps is not None else 0.1
+            tau = self._tau if self._tau is not None else 0.2
+        else:
+            if self._eps is None or self._tau is None:
+                h0 = silverman_bandwidth(X.detach())
+                eps_a, tau_a = adaptive_eps_tau(X.detach(), h0, grid)
+                eps = self._eps if self._eps is not None else eps_a
+                tau = self._tau if self._tau is not None else tau_a
+            else:
+                eps, tau = self._eps, self._tau
+
+        eps_eff, tau_eff = anneal_eps_tau(eps, tau, self.anneal_factor)
+
+        return DCBFunction.apply(
+            X, grid, eps_eff, tau_eff, self.target_modes, self.delta, self.formula,
+            self.chunk_size, self.brentq_n_max, self.g_brentq, self.use_hard_bisection,
+            self.safe_backward, self.use_fft,
+        )
+
+
+# Public alias
+DifferentiableCriticalBandwidth = DCBLayer