ratio-markov 0.5.1

//! # Clustering module
//!
//! ## License
//!
//! This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0.
//! If a copy of the MPL was not distributed with this file,
//! You can obtain one at <https://mozilla.org/MPL/2.0/>.
//!
//! **Code examples both in the docstrings and rendered documentation are free to use.**

use std::cmp::max;
use std::collections::HashMap;

use crate::{DMatrix, DVector, MarkovFlow, utils};

/// Clustering a graph based on Markov steady-state flow calculations on its adjacency matrix.
///
/// # Arguments
///
/// * `adj` - Square adjacency matrix with nodes in identical order on both axis.
///   Only positive adjacency values are used in the calculations.
/// * `alpha` - Expansion coefficient. Power to which the transfer probability matrix is
///   raised in each iteration.
/// * `beta` - Inflation coefficient. Element-wise power that is applied to all values
///   in each iteration.
/// * `mu` - Evaporation constant. Factor by which flow is reduced after passing
///   through a node.
/// * `max_iter` - Maximum number of iterations to go through after which the then
///   calculated transfer probability matrix will be assumed to be the attractor matrix.
///
/// # Returns
///
/// A hashmap of attractor node IDs to vectors that represent clusters of node indices
/// that correspond to the row/column indices in the original adjacency matrix.
///
/// # Notes
///
/// The Markov clustering algorithm works by calculating a transfer probability matrix
/// and using it to make "steps" (iterations) using a predefined procedure. Tightly
/// coupled nodes will attract more and more towards attractor nodes and converge
/// towards a steady state system. When this is approximated, the procedure will break
/// and the maximum value in each column determines the "cluster" it belongs to.
pub fn cluster(
    adj: &DMatrix<f64>,
    alpha: usize,
    beta: f64,
    mu: f64,
    max_iter: usize,
) -> HashMap<usize, Vec<usize>> {
    let flow = MarkovFlow::new(&(adj + adj.transpose()), mu, false);
    let max_degree = max(1, utils::get_max_degree(&flow.adjacency_matrix, true)) as f64;
    let threshold = (mu * max_degree).powi(-((alpha + 1) as i32));

    let attractor_matrix = if alpha > 1 {
        clustering_cycle(flow, alpha, beta, max_iter, threshold)
    } else {
        adj.clone_owned()
    };

    let cluster_ids = get_cluster_ids(&attractor_matrix);
    get_cluster_map(&cluster_ids)
}

/// get_tpm calculates the transfer probability matrix used in the Markov clustering algorithm.
///
/// # Arguments
///
/// * `flow` - A Markov flow object.
///
/// # Returns
///
/// A transfer probability matrix.
pub fn get_tpm(flow: MarkovFlow) -> DMatrix<f64> {
    // Add up, set diagonal to zero and apply column normalization.
    let mut tpm = flow.influence_matrix + flow.dependency_matrix;
    tpm.fill_diagonal(0.0);

    // Set diagonal to the maximum value in each column and normalize.
    for (j, mut col) in tpm.column_iter_mut().enumerate() {
        col[j] = col.max();
        let sum = col.sum();
        if sum.abs() < f64::EPSILON {
            col.fill(0.0);
            col[j] = 1.0;
            continue;
        }
        col *= 1.0 / sum;
    }

    tpm
}

/// Perform a Markov clustering cycle.
///
/// # Arguments
///
/// * `flow` - Markov flow object.
/// * `alpha` - Expansion coefficient. Power to which the transfer probability matrix is
///   raised in each iteration.
/// * `beta` - Inflation coefficient. Element-wise power that is applied to all values
///   in each iteration.
/// * `max_iter` - Maximum number of iterations to go through after which the then
///   calculated transfer probability matrix will be assumed to be the attractor matrix.
/// * `threshold` - Pruning threshold. Values below this threshold are pruned after each
///   cycle to boost convergence.
///
/// # Returns
///
/// Attractor matrix. Indices of the maximum value in each column correspond to the
/// attractor node (cluster) each node belongs to.
pub fn clustering_cycle(
    flow: MarkovFlow,
    alpha: usize,
    beta: f64,
    max_iter: usize,
    threshold: f64,
) -> DMatrix<f64> {
    let mut tpm = get_tpm(flow);
    let (rows, cols) = tpm.shape();
    utils::prune_matrix(&mut tpm, threshold);
    let mut last_tpm = DMatrix::<f64>::zeros(rows, cols);

    // Try to converge within max iterations.
    for _ in 0..max_iter {
        if tpm.relative_eq(&last_tpm, f64::EPSILON, f64::EPSILON) {
            break;
        }
        last_tpm = tpm.clone_owned();
        tpm.pow_mut(alpha as u32);
        utils::matrix_elementwise_power(&mut tpm, beta);
        utils::prune_matrix(&mut tpm, threshold);
        utils::normalize_matrix_columns(&mut tpm);
    }

    // Converged TPM (or last iteration) is considered the attractor matrix.
    tpm
}

/// Get a vector containing the cluster ID assignments (row indices of maximum values in each column).
///
/// # Arguments
///
/// * `attractor_matrix` - Attractor matrix. Indices of the maximum value in each column
///   correspond to the attractor node (cluster) each node belongs to.
///
/// # Returns
///
/// Vector containing the attractor node (cluster) ID for each node.
pub fn get_cluster_ids(attractor_matrix: &DMatrix<f64>) -> DVector<usize> {
    let (_, cols) = attractor_matrix.shape();
    DVector::from_fn(cols, |i, _| attractor_matrix.column(i).imax())
}

/// Get a cluster mapping of cluster indices to original node indices.
///
/// # Arguments
///
/// * `cluster_ids` - Vector containing the attractor node (cluster) ID for each node.
///
/// # Returns
///
/// A hashmap of attractor node IDs to vectors that represent clusters of node indices
/// that correspond to the row/column indices in the original adjacency matrix.
pub fn get_cluster_map(cluster_ids: &DVector<usize>) -> HashMap<usize, Vec<usize>> {
    let mut cluster_map = HashMap::<usize, Vec<usize>>::new();
    for (i, id) in cluster_ids.iter().enumerate() {
        if cluster_map.contains_key(id) {
            cluster_map.entry(*id).and_modify(|v| v.push(i));
        } else {
            cluster_map.insert(*id, vec![i]);
        }
    }
    cluster_map
}

#[cfg(test)]
mod tests {
    use super::*;
    #[test]
    fn test_trivial_example() {
        let adj = DMatrix::from_row_slice(
            4,
            4,
            &[
                0.0, 1.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0,
            ],
        );
        let cmap = cluster(&adj, 2, 2.0, 2.0, 10 ^ 100);
        let reference: HashMap<usize, Vec<usize>> =
            [(0, vec![0, 1]), (2, vec![2, 3])].into_iter().collect();
        assert_eq!(cmap, reference);
    }
}