layout.rs

1use super::models::Network;2use crate::core::error::{value_error, Result};3use crate::core::Rng;45const SPAN: usize = 300;6const FLOOR: f64 = 0.02;7const GAIN: f64 = 0.1;8const JITTER: f64 = 0.05;9const APART: f64 = 1e-9;1011/// A force-directed layout: every node repels every other, every branch pulls its ends together, and a cooling cap on the move per tick lets the lattice settle.12///13/// The forces are Fruchterman and Reingold's: repulsion `k^2 / d` between every pair, attraction14/// `d^2 / k` along every branch, with `k` the ideal length `extent / n^(1/dim)` read off the15/// starting box. The temperature caps the move of any node in one tick and cools linearly from16/// a tenth of the extent to a floor over the first ticks, then holds, so the layout keeps17/// creeping toward rest. The seed jitters the start so a symmetric lattice can fold.18///19/// ```20/// let square = [0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0];21/// let ring = [(0, 1), (1, 2), (2, 3), (3, 0)];22/// let mut layout = mrlyrs::math::graph::Layout::new(&square, &ring, 2, 1).unwrap();23/// assert!(layout.step(500) < 1e-3);24/// ```25#[derive(Clone)]26pub struct Layout {27    dim: usize,28    positions: Vec<f64>,29    branches: Vec<(usize, usize)>,30    ideal: f64,31    heat: f64,32    tick: usize,33    energy: f64,34    moved: f64,35    rng: Rng,36}3738impl Layout {39    /// Starts a layout from flat positions, `dim` floats per node, and the branch pairs.40    ///41    /// # Errors42    ///43    /// Errors when a branch points off the list, or the positions are not dim floats a node.44    pub fn new(45        positions: &[f64],46        branches: &[(usize, usize)],47        dim: usize,48        seed: u64,49    ) -> Result<Layout> {50        if dim == 0 || !positions.len().is_multiple_of(dim) {51            return value_error("positions must hold dim floats per node.");52        }53        let count = positions.len() / dim;54        if branches.iter().any(|&(a, b)| a >= count || b >= count) {55            return value_error("a branch points past the last node.");56        }57        let mut extent: f64 = 0.0;58        for axis in 0..dim {59            let column = positions.iter().skip(axis).step_by(dim).copied();60            let low = column.clone().fold(f64::INFINITY, f64::min);61            let high = column.fold(f64::NEG_INFINITY, f64::max);62            extent = extent.max(high - low);63        }64        if extent <= 0.0 || extent.is_nan() {65            extent = 1.0;66        }67        let ideal = extent / (count.max(1) as f64).powf(1.0 / dim as f64);68        let mut rng = Rng::new(seed);69        let positions = positions70            .iter()71            .map(|&p| p + (rng.unit() - 0.5) * 2.0 * JITTER * ideal)72            .collect();73        Ok(Layout {74            dim,75            positions,76            branches: branches.to_vec(),77            ideal,78            heat: extent / 10.0,79            tick: 0,80            energy: 0.0,81            moved: 0.0,82            rng,83        })84    }85    /// Starts a layout from a network's own positions and branches.86    ///87    /// # Errors88    ///89    /// Errors when a branch names a node the network does not hold.90    pub fn from_network(network: &Network, seed: u64) -> Result<Layout> {91        let positions: Vec<f64> = network92            .nodes93            .iter()94            .flat_map(|node| node.position.iter().copied())95            .collect();96        let branches: Vec<(usize, usize)> = network97            .branches98            .iter()99            .map(|b| (b.parent, b.child))100            .collect();101        Layout::new(&positions, &branches, network.dim, seed)102    }103    /// Returns the node count.104    pub fn nodes(&self) -> usize {105        self.positions.len() / self.dim106    }107    /// Returns the ideal branch length `k`.108    pub fn ideal(&self) -> f64 {109        self.ideal110    }111    /// Returns the cap on one node's move in the next tick.112    pub fn temperature(&self) -> f64 {113        let cooled = 1.0 - self.tick.min(SPAN) as f64 / SPAN as f64;114        self.heat * cooled.max(FLOOR)115    }116    /// Returns the ticks stepped so far.117    pub fn ticks(&self) -> usize {118        self.tick119    }120    /// Runs the ticks and returns the energy left: the mean net force per node in units of `k`.121    pub fn step(&mut self, ticks: usize) -> f64 {122        for _ in 0..ticks {123            self.tick_once();124        }125        self.energy126    }127    fn tick_once(&mut self) {128        let (dim, n, k) = (self.dim, self.nodes(), self.ideal);129        let mut push = vec![0.0; n * dim];130        let mut delta = vec![0.0; dim];131        for i in 0..n {132            for j in (i + 1)..n {133                let d = gap(&self.positions, dim, i, j, &mut delta, &mut self.rng);134                let f = k * k / (d * d);135                for a in 0..dim {136                    push[i * dim + a] += delta[a] * f;137                    push[j * dim + a] -= delta[a] * f;138                }139            }140        }141        for &(i, j) in &self.branches {142            let d = gap(&self.positions, dim, i, j, &mut delta, &mut self.rng);143            let f = d / k;144            for a in 0..dim {145                push[i * dim + a] -= delta[a] * f;146                push[j * dim + a] += delta[a] * f;147            }148        }149        let cap = self.temperature();150        let (mut force, mut moved) = (0.0, 0.0);151        for i in 0..n {152            let len = (0..dim)153                .map(|a| push[i * dim + a].powi(2))154                .sum::<f64>()155                .sqrt();156            force += len;157            if len == 0.0 {158                continue;159            }160            let reach = (len * GAIN).min(cap);161            moved += reach;162            for a in 0..dim {163                self.positions[i * dim + a] += push[i * dim + a] / len * reach;164            }165        }166        self.energy = force / (n.max(1) as f64 * k);167        self.moved = moved / n.max(1) as f64;168        self.tick += 1;169    }170    /// Returns the positions, `dim` floats per node.171    pub fn positions(&self) -> &[f64] {172        &self.positions173    }174    /// Returns the mean net force per node in units of `k` after the last tick.175    pub fn energy(&self) -> f64 {176        self.energy177    }178    /// Returns the mean distance a node moved in the last tick.179    pub fn moved(&self) -> f64 {180        self.moved181    }182}183184fn gap(positions: &[f64], dim: usize, i: usize, j: usize, delta: &mut [f64], rng: &mut Rng) -> f64 {185    for a in 0..dim {186        delta[a] = positions[i * dim + a] - positions[j * dim + a];187    }188    let d = delta.iter().map(|x| x * x).sum::<f64>().sqrt();189    if d >= APART {190        return d;191    }192    for x in delta.iter_mut() {193        *x = (rng.unit() - 0.5) * APART;194    }195    delta196        .iter()197        .map(|x| x * x)198        .sum::<f64>()199        .sqrt()200        .max(APART / 2.0)201}202203#[cfg(test)]204mod tests {205    use super::*;206    fn ring() -> Layout {207        let square = [0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0];208        let ring = [(0, 1), (1, 2), (2, 3), (3, 0)];209        Layout::new(&square, &ring, 2, 1).unwrap()210    }211    #[test]212    fn the_four_cycle_relaxes_to_a_square() {213        let mut layout = ring();214        assert!((layout.ideal() - 0.5).abs() < 1e-12);215        let rest = layout.step(500);216        assert!(rest < 1e-3, "energy {rest}");217        let p = layout.positions();218        let gaps: Vec<f64> = [(0, 1), (1, 2), (2, 3), (3, 0)]219            .iter()220            .map(|&(a, b)| {221                ((p[2 * a] - p[2 * b]).powi(2) + (p[2 * a + 1] - p[2 * b + 1]).powi(2)).sqrt()222            })223            .collect();224        let spread = gaps.iter().cloned().fold(0.0, f64::max)225            - gaps.iter().cloned().fold(f64::INFINITY, f64::min);226        assert!(spread < 1e-3, "gaps {gaps:?}");227        let side = 0.5 * 1.5f64.cbrt();228        assert!((gaps[0] - side).abs() < 1e-2, "side {}", gaps[0]);229        assert_eq!(layout.ticks(), 500);230    }231    #[test]232    fn the_seed_replays_and_the_cap_cools() {233        let mut a = ring();234        let mut b = ring();235        a.step(20);236        b.step(20);237        assert_eq!(a.positions(), b.positions());238        assert!(239            a.temperature()240                < Layout::new(&[0.0, 0.0, 1.0, 1.0], &[], 2, 1)241                    .unwrap()242                    .temperature()243        );244        assert!(a.moved() > 0.0);245    }246    #[test]247    fn the_faults_are_named() {248        assert!(Layout::new(&[0.0, 0.0, 1.0], &[], 2, 1).is_err());249        assert!(Layout::new(&[0.0, 0.0], &[(0, 1)], 2, 1).is_err());250        assert!(Layout::new(&[0.0, 0.0, 0.0, 0.0], &[(0, 1)], 2, 1)251            .unwrap()252            .step(3)253            .is_finite());254    }255}