rune-node2vec 0.1.0

use crate::rng::Rng;

/// Parameters passed to the skip-gram trainer.
pub(crate) struct TrainParams<'a> {
    pub(crate) n_nodes: usize,
    pub(crate) dim: usize,
    pub(crate) walks: &'a [Vec<usize>],
    pub(crate) window_size: usize,
    pub(crate) neg_samples: usize,
    pub(crate) n_epochs: usize,
    pub(crate) initial_lr: f64,
    pub(crate) noise_table: &'a [usize],
}

/// Trains embeddings with skip-gram negative sampling over all walks.
///
/// `embeddings` is a flat row-major array of shape `[n_nodes × dim]`.
/// After each epoch the learning rate decays linearly towards `min_lr`.
pub(crate) fn train(embeddings: &mut [f64], params: &TrainParams<'_>, rng: &mut Rng) {
    let min_lr = params.initial_lr * 0.0001;

    for epoch in 0..params.n_epochs {
        let progress = epoch as f64 / params.n_epochs.max(1) as f64;
        let lr = params.initial_lr - (params.initial_lr - min_lr) * progress;

        for walk in params.walks {
            process_walk(embeddings, params, walk, lr, rng);
        }
    }

    let _ = params.n_nodes;
}

/// Processes one walk: for each centre node, updates embeddings with all context pairs.
fn process_walk(
    embeddings: &mut [f64],
    params: &TrainParams<'_>,
    walk: &[usize],
    lr: f64,
    rng: &mut Rng,
) {
    let walk_len = walk.len();

    for (center_pos, &center_node) in walk.iter().enumerate() {
        let context_start = center_pos.saturating_sub(params.window_size);
        let context_end = (center_pos + params.window_size + 1).min(walk_len);

        for (offset, &context_node) in walk[context_start..context_end].iter().enumerate() {
            if context_start + offset == center_pos {
                continue;
            }
            run_pair_update(
                embeddings,
                &PairParams {
                    dim: params.dim,
                    center: center_node,
                    context: context_node,
                    neg_samples: params.neg_samples,
                    lr,
                    noise_table: params.noise_table,
                },
                rng,
            );
        }
    }
}

/// Parameters for updating a single (centre, context) pair.
struct PairParams<'a> {
    dim: usize,
    center: usize,
    context: usize,
    neg_samples: usize,
    lr: f64,
    noise_table: &'a [usize],
}

fn run_pair_update(embeddings: &mut [f64], p: &PairParams<'_>, rng: &mut Rng) {
    let dim = p.dim;
    let center = p.center;
    let context = p.context;

    let mut grad_center = vec![0.0f64; dim];

    // Positive pair: label = 1.
    {
        let dot = dot_product(embeddings, dim, center, context);
        let sigma = sigmoid(dot);
        let positive_factor = (1.0 - sigma) * p.lr;

        accumulate_gradient(&mut grad_center, embeddings, dim, context, positive_factor);
        let center_row = embeddings_row(embeddings, dim, center);
        update_vector(embeddings, dim, context, &center_row, positive_factor);
    }

    // Negative samples: label = 0.
    for _ in 0..p.neg_samples {
        let noise_node = p.noise_table[rng.next_usize(p.noise_table.len())];
        let dot = dot_product(embeddings, dim, center, noise_node);
        let sigma = sigmoid(dot);
        let negative_factor = -sigma * p.lr;

        accumulate_gradient(&mut grad_center, embeddings, dim, noise_node, negative_factor);
        let center_row = embeddings_row(embeddings, dim, center);
        update_vector(embeddings, dim, noise_node, &center_row, negative_factor);
    }

    // Apply accumulated gradient to centre.
    let start = center * dim;
    for d in 0..dim {
        embeddings[start + d] += grad_center[d];
    }
}

fn sigmoid(x: f64) -> f64 {
    1.0 / (1.0 + (-x).exp())
}

fn dot_product(embeddings: &[f64], dim: usize, row_a: usize, row_b: usize) -> f64 {
    let a = row_a * dim;
    let b = row_b * dim;
    (0..dim).map(|d| embeddings[a + d] * embeddings[b + d]).sum()
}

/// Returns a copy of one embedding row (needed before borrowing mutably).
fn embeddings_row(embeddings: &[f64], dim: usize, row: usize) -> Vec<f64> {
    embeddings[row * dim..(row + 1) * dim].to_vec()
}

/// Accumulates `factor × embeddings[row]` into `grad`.
fn accumulate_gradient(grad: &mut [f64], embeddings: &[f64], dim: usize, row: usize, factor: f64) {
    let start = row * dim;
    for d in 0..dim {
        grad[d] += factor * embeddings[start + d];
    }
}

/// Adds `factor × center_row` to `embeddings[row]`.
fn update_vector(embeddings: &mut [f64], dim: usize, row: usize, center_row: &[f64], factor: f64) {
    let start = row * dim;
    for d in 0..dim {
        embeddings[start + d] += factor * center_row[d];
    }
}

/// Builds the noise (negative sampling) table.
///
/// Each node appears proportional to `degree^(3/4)`. The table has `table_size`
/// entries. Nodes with no edges are excluded; if all nodes are isolated the table
/// is filled with zeros.
pub(crate) fn build_noise_table(degrees: &[usize], table_size: usize) -> Vec<usize> {
    let weights: Vec<f64> = degrees.iter().map(|&d| (d as f64).powf(0.75)).collect();
    let total: f64 = weights.iter().sum();

    if total == 0.0 {
        return vec![0usize; table_size];
    }

    let mut table = Vec::with_capacity(table_size);
    let mut cumulative = 0.0;
    let mut node_idx = 0;

    for i in 0..table_size {
        let threshold = (i as f64 + 1.0) / table_size as f64;
        cumulative += weights[node_idx] / total;
        while cumulative < threshold && node_idx + 1 < weights.len() {
            node_idx += 1;
            cumulative += weights[node_idx] / total;
        }
        table.push(node_idx);
    }

    table
}