Skip to content

embedding — Embedding Loop

SG-t-SNE-Pi gradient descent optimization loop.

Translated from ref/sgtsnepi/src/gradient_descend.cpp.

sgtsne_embedding(P, d=2, max_iter=1000, early_exag=250, alpha=12.0, eta=200.0, h=1.0, Y0=None, random_state=None)

Run SG-t-SNE-Pi gradient descent.

Parameters:

Name Type Description Default
P csc_matrix

Symmetrized, normalized probability matrix (n, n).

 required
d int

Embedding dimensions (1, 2, or 3).

 2
max_iter int

Total iterations.

 1000
early_exag int

Number of early exaggeration iterations.

 250
alpha float

Exaggeration multiplier (drops to 1 after early_exag).

 12.0
eta float

Learning rate.

 200.0
h float

Grid side length for FFT.

 1.0
Y0 ndarray or None

Initial embedding of shape (n, d). If None, random init with scale 0.01 (matching Julia tutorial convention).

 None
random_state int or None

Random seed for reproducibility.

 None

Returns:

Name Type Description
Y ndarray of shape (n, d)

Final embedding coordinates.
Source code in src/pysgtsnepi/embedding.py 
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
def sgtsne_embedding(
    P: csc_matrix,
    d: int = 2,
    max_iter: int = 1000,
    early_exag: int = 250,
    alpha: float = 12.0,
    eta: float = 200.0,
    h: float = 1.0,
    Y0: np.ndarray | None = None,
    random_state: int | None = None,
) -> np.ndarray:
    """Run SG-t-SNE-Pi gradient descent.

    Parameters
    ----------
    P : csc_matrix
        Symmetrized, normalized probability matrix (n, n).
    d : int
        Embedding dimensions (1, 2, or 3).
    max_iter : int
        Total iterations.
    early_exag : int
        Number of early exaggeration iterations.
    alpha : float
        Exaggeration multiplier (drops to 1 after early_exag).
    eta : float
        Learning rate.
    h : float
        Grid side length for FFT.
    Y0 : ndarray or None
        Initial embedding of shape (n, d). If None, random init
        with scale 0.01 (matching Julia tutorial convention).
    random_state : int or None
        Random seed for reproducibility.

    Returns
    -------
    Y : ndarray of shape (n, d)
        Final embedding coordinates.
    """
    n = P.shape[0]
    rng = np.random.default_rng(random_state)

    # Initialize embedding
    if Y0 is not None:
        Y = Y0.copy().astype(np.float64)
    else:
        Y = rng.standard_normal((n, d)) * 0.01

    # Momentum parameters (same as C++ reference)
    momentum = 0.5
    final_momentum = 0.8
    mom_switch_iter = 250

    # Allocate state
    uY = np.zeros((n, d), dtype=np.float64)
    gains = np.ones((n, d), dtype=np.float64)

    current_alpha = alpha

    for it in range(max_iter):
        # Compute attractive forces
        Fattr = attractive_forces(P, Y)

        # Compute repulsive forces
        Frep, _zeta = compute_repulsive_forces(Y, h)

        # Gradient: dY = alpha * Fattr - Frep
        dY = current_alpha * Fattr - Frep

        # Gain adaptation
        gains = np.where(np.sign(dY) != np.sign(uY), gains + 0.2, gains * 0.8)
        np.clip(gains, 0.01, None, out=gains)

        # Velocity update
        uY = momentum * uY - eta * gains * dY

        # Position update
        Y += uY

        # Zero-mean center
        Y -= Y.mean(axis=0)

        # Stop early exaggeration
        if it == early_exag:
            current_alpha = 1.0

        # Switch momentum
        if it == mom_switch_iter:
            momentum = final_momentum

    return Y