repulsive — Repulsive Forces

FFT-accelerated repulsive forces for SG-t-SNE-Pi.

Translated from ref/sgtsnepi/src/qq.cpp, gridding.cpp, nuconv.cpp, non_periodic_conv.cpp.

compute_repulsive_forces(Y, h)

Compute FFT-accelerated repulsive forces.

Parameters:

Name	Type	Description	Default
Y	ndarray of shape (n, d)	Current embedding coordinates.	required
h	float	Grid side length for FFT.	required

Returns:

Name	Type	Description
Frep	ndarray of shape (n, d)	Repulsive forces.
zeta	float	Normalization constant.

Source code in src/pysgtsnepi/repulsive.py

def compute_repulsive_forces(Y: np.ndarray, h: float) -> tuple[np.ndarray, float]:
    """Compute FFT-accelerated repulsive forces.

    Parameters
    ----------
    Y : ndarray of shape (n, d)
        Current embedding coordinates.
    h : float
        Grid side length for FFT.

    Returns
    -------
    Frep : ndarray of shape (n, d)
        Repulsive forces.
    zeta : float
        Normalization constant.
    """
    n, d = Y.shape
    eps = np.finfo(np.float64).eps

    # Guard against NaN/Inf input
    if not np.all(np.isfinite(Y)):
        return np.zeros((n, d), dtype=np.float64), 1.0

    # Work on copies — do NOT modify the embedding Y in-place
    Y_shifted = Y - Y.min(axis=0)

    # Global max across all dims
    maxy = Y_shifted.max()
    if maxy < eps:
        return np.zeros((n, d), dtype=np.float64), 1.0

    # Grid size
    n_grid = max(math.ceil(maxy / h), 14)
    n_grid = _best_grid_size(n_grid)

    # Grid positions in [0, n_grid-1]
    Y_grid = Y_shifted / maxy * (n_grid - 1)
    # Clamp to avoid edge issues
    Y_grid = np.clip(Y_grid, 0, n_grid - 1 - eps)

    # Exact grid spacing in data space
    h_exact = maxy / (n_grid - 1 - eps)

    ng = n_grid + 3  # padded grid: 4-point stencil at pos n_grid-1 needs idx n_grid+2
    n_vec = d + 1  # number of value channels

    # Setup scattered values: VScat[:, 0] = 1, VScat[:, 1:] = Y_shifted
    VScat = np.empty((n, n_vec), dtype=np.float64)
    VScat[:, 0] = 1.0
    VScat[:, 1:] = Y_shifted

    # Ensure grid coordinates are 2D for the numba kernels
    Y_grid_2d = np.ascontiguousarray(Y_grid.reshape(n, d))

    # S2G: scatter to grid
    grid_shape = tuple([ng] * d + [n_vec])
    VGrid = np.zeros(grid_shape, dtype=np.float64)
    _S2G[d](VGrid, Y_grid_2d, VScat, ng, n, n_vec)

    # G2G: FFT convolution
    PhiGrid = np.zeros(grid_shape, dtype=np.float64)
    _CONV[d](PhiGrid, VGrid, h_exact, ng, n_vec)

    # G2S: grid to scatter
    PhiScat = np.zeros((n, n_vec), dtype=np.float64)
    _G2S[d](PhiScat, PhiGrid, Y_grid_2d, ng, n, n_vec)

    # Compute Z and repulsive forces using original-scale coords
    Frep, zeta = _zeta_and_force(np.ascontiguousarray(Y_shifted), PhiScat, n, d)

    return Frep, zeta