visgraph 0.1.0

//! Module containing functionality for the force-directed layout.
//!
//! The main function is [`force_directed_layout`], which returns a position map function that
//! arranges nodes in a force-directed layout.

use petgraph::visit::{EdgeRef, IntoEdgeReferences, IntoNodeReferences, NodeIndexable, NodeRef};

/// Default number of iterations for the [`force_directed_layout`] function.
pub const DEFAULT_ITERATIONS: u32 = 1000;
/// Default initial temperature for the [`force_directed_layout`] function.
pub const DEFAULT_INITIAL_TEMPERATURE: f32 = 0.1;
const CLIPPING_VALUE: f32 = 0.01;

/// Returns a position map function that arranges nodes using a force-directed layout. The specific
/// algorithm used is the Fruchterman-Reingold algorithm, see [Reference](#reference).
///
/// The returned position map is normalized to [0.0, 1.0].
///
/// # Reference
///
/// This is an implementation of the Fruchterman-Reingold force-directed algorithm as presented in
/// the original paper:
///
/// Fruchterman, T. M. J., Reingold, E. M. (1991). Graph drawing by force-directed placement
/// <https://doi.org/10.1002/spe.4380211102>.
pub fn force_directed_layout<G>(
    graph: &G,
    iterations: u32,
    inital_temperature: f32,
) -> impl Fn(G::NodeId) -> (f32, f32) + '_
where
    G: IntoNodeReferences + IntoEdgeReferences + NodeIndexable,
{
    let node_count = graph.node_references().count();
    let mut positions = vec![(0.0f32, 0.0f32); graph.node_bound()];

    if node_count > 0 {
        // Initialize positions randomly
        let mut rng = fastrand::Rng::new();

        for node_ref in graph.node_references() {
            let x = rng.f32();
            let y = rng.f32();
            let idx = graph.to_index(node_ref.id());
            positions[idx] = (x, y);
        }

        // Simulation parameters
        let k = (1.0 / (node_count as f32)).sqrt();

        let edges: Vec<(usize, usize)> = graph
            .edge_references()
            .map(|edge| (graph.to_index(edge.source()), graph.to_index(edge.target())))
            .collect();

        let node_indices: Vec<usize> = graph
            .node_references()
            .map(|node_ref| graph.to_index(node_ref.id()))
            .collect();

        for iteration in 0..iterations {
            let mut displacements = vec![(0.0f32, 0.0f32); graph.node_bound()];

            // Calculate repulsive forces between all pairs of nodes
            for i in 0..node_indices.len() {
                for j in (i + 1)..node_indices.len() {
                    let idx_i = node_indices[i];
                    let idx_j = node_indices[j];

                    let delta_x = positions[idx_i].0 - positions[idx_j].0;
                    let delta_y = positions[idx_i].1 - positions[idx_j].1;
                    let distance = (delta_x * delta_x + delta_y * delta_y)
                        .sqrt()
                        .max(CLIPPING_VALUE);

                    // Repulsive force: f_r = k^2 / d
                    let repulsion = k * k / distance;
                    let force_x = (delta_x / distance) * repulsion;
                    let force_y = (delta_y / distance) * repulsion;

                    displacements[idx_i].0 += force_x;
                    displacements[idx_i].1 += force_y;
                    displacements[idx_j].0 -= force_x;
                    displacements[idx_j].1 -= force_y;
                }
            }

            // Calculate attractive forces along edges
            for &(source_idx, target_idx) in &edges {
                let delta_x = positions[source_idx].0 - positions[target_idx].0;
                let delta_y = positions[source_idx].1 - positions[target_idx].1;
                let distance = (delta_x * delta_x + delta_y * delta_y)
                    .sqrt()
                    .max(CLIPPING_VALUE);

                let attraction = distance * distance / k;
                let force_x = (delta_x / distance) * attraction;
                let force_y = (delta_y / distance) * attraction;

                displacements[source_idx].0 -= force_x;
                displacements[source_idx].1 -= force_y;
                displacements[target_idx].0 += force_x;
                displacements[target_idx].1 += force_y;
            }

            // Apply displacements with cooling
            let curr_temp =
                inital_temperature - (0.1 * iteration as f32) / ((iterations + 1) as f32);
            for &idx in &node_indices {
                let disp_len = (displacements[idx].0 * displacements[idx].0
                    + displacements[idx].1 * displacements[idx].1)
                    .sqrt();

                if disp_len > 0.0 {
                    let limited_disp_len = disp_len.min(curr_temp);
                    positions[idx].0 += (displacements[idx].0 / disp_len) * limited_disp_len;
                    positions[idx].1 += (displacements[idx].1 / disp_len) * limited_disp_len;
                }
            }
        }

        // Normalize positions to [0.0, 1.0]
        if !positions.is_empty() {
            let mut min_x = f32::INFINITY;
            let mut max_x = f32::NEG_INFINITY;
            let mut min_y = f32::INFINITY;
            let mut max_y = f32::NEG_INFINITY;

            for &idx in &node_indices {
                min_x = min_x.min(positions[idx].0);
                max_x = max_x.max(positions[idx].0);
                min_y = min_y.min(positions[idx].1);
                max_y = max_y.max(positions[idx].1);
            }

            let range_x = max_x - min_x;
            let range_y = max_y - min_y;

            if range_x > 0.0 && range_y > 0.0 {
                for &idx in &node_indices {
                    positions[idx].0 = (positions[idx].0 - min_x) / range_x;
                    positions[idx].1 = (positions[idx].1 - min_y) / range_y;
                }
            } else if range_x > 0.0 {
                for &idx in &node_indices {
                    positions[idx].0 = (positions[idx].0 - min_x) / range_x;
                    positions[idx].1 = 0.5;
                }
            } else if range_y > 0.0 {
                for &idx in &node_indices {
                    positions[idx].0 = 0.5;
                    positions[idx].1 = (positions[idx].1 - min_y) / range_y;
                }
            } else {
                for &idx in &node_indices {
                    positions[idx] = (0.5, 0.5);
                }
            }
        }
    }

    move |node_id| positions[NodeIndexable::to_index(&graph, node_id)]
}