layout.rs

1use super::models::Network;2use mrlycore::errors::{value_error, Result};3use mrlycore::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 = mrlynum::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, or an error when a branch points off the list.40    pub fn new(41        positions: &[f64],42        branches: &[(usize, usize)],43        dim: usize,44        seed: u64,45    ) -> Result<Layout> {46        if dim == 0 || !positions.len().is_multiple_of(dim) {47            return value_error("positions must hold dim floats per node.");48        }49        let count = positions.len() / dim;50        if branches.iter().any(|&(a, b)| a >= count || b >= count) {51            return value_error("a branch points past the last node.");52        }53        let mut extent: f64 = 0.0;54        for axis in 0..dim {55            let column = positions.iter().skip(axis).step_by(dim).copied();56            let low = column.clone().fold(f64::INFINITY, f64::min);57            let high = column.fold(f64::NEG_INFINITY, f64::max);58            extent = extent.max(high - low);59        }60        if extent <= 0.0 || extent.is_nan() {61            extent = 1.0;62        }63        let ideal = extent / (count.max(1) as f64).powf(1.0 / dim as f64);64        let mut rng = Rng::new(seed);65        let positions = positions66            .iter()67            .map(|&p| p + (rng.unit() - 0.5) * 2.0 * JITTER * ideal)68            .collect();69        Ok(Layout {70            dim,71            positions,72            branches: branches.to_vec(),73            ideal,74            heat: extent / 10.0,75            tick: 0,76            energy: 0.0,77            moved: 0.0,78            rng,79        })80    }81    /// Starts a layout from a network's own positions and branches.82    pub fn from_network(network: &Network, seed: u64) -> Result<Layout> {83        let positions: Vec<f64> = network84            .nodes85            .iter()86            .flat_map(|node| node.position.iter().copied())87            .collect();88        let branches: Vec<(usize, usize)> = network89            .branches90            .iter()91            .map(|b| (b.parent, b.child))92            .collect();93        Layout::new(&positions, &branches, network.dim, seed)94    }95    /// Returns the node count.96    pub fn nodes(&self) -> usize {97        self.positions.len() / self.dim98    }99    /// Returns the ideal branch length `k`.100    pub fn ideal(&self) -> f64 {101        self.ideal102    }103    /// Returns the cap on one node's move in the next tick.104    pub fn temperature(&self) -> f64 {105        let cooled = 1.0 - self.tick.min(SPAN) as f64 / SPAN as f64;106        self.heat * cooled.max(FLOOR)107    }108    /// Returns the ticks stepped so far.109    pub fn ticks(&self) -> usize {110        self.tick111    }112    /// Runs the ticks and returns the energy left: the mean net force per node in units of `k`.113    pub fn step(&mut self, ticks: usize) -> f64 {114        for _ in 0..ticks {115            self.tick_once();116        }117        self.energy118    }119    fn tick_once(&mut self) {120        let (dim, n, k) = (self.dim, self.nodes(), self.ideal);121        let mut push = vec![0.0; n * dim];122        let mut delta = vec![0.0; dim];123        for i in 0..n {124            for j in (i + 1)..n {125                let d = gap(&self.positions, dim, i, j, &mut delta, &mut self.rng);126                let f = k * k / (d * d);127                for a in 0..dim {128                    push[i * dim + a] += delta[a] * f;129                    push[j * dim + a] -= delta[a] * f;130                }131            }132        }133        for &(i, j) in &self.branches {134            let d = gap(&self.positions, dim, i, j, &mut delta, &mut self.rng);135            let f = d / k;136            for a in 0..dim {137                push[i * dim + a] -= delta[a] * f;138                push[j * dim + a] += delta[a] * f;139            }140        }141        let cap = self.temperature();142        let (mut force, mut moved) = (0.0, 0.0);143        for i in 0..n {144            let len = (0..dim)145                .map(|a| push[i * dim + a].powi(2))146                .sum::<f64>()147                .sqrt();148            force += len;149            if len == 0.0 {150                continue;151            }152            let reach = (len * GAIN).min(cap);153            moved += reach;154            for a in 0..dim {155                self.positions[i * dim + a] += push[i * dim + a] / len * reach;156            }157        }158        self.energy = force / (n.max(1) as f64 * k);159        self.moved = moved / n.max(1) as f64;160        self.tick += 1;161    }162    /// Returns the positions, `dim` floats per node.163    pub fn positions(&self) -> &[f64] {164        &self.positions165    }166    /// Returns the mean net force per node in units of `k` after the last tick.167    pub fn energy(&self) -> f64 {168        self.energy169    }170    /// Returns the mean distance a node moved in the last tick.171    pub fn moved(&self) -> f64 {172        self.moved173    }174}175176fn gap(positions: &[f64], dim: usize, i: usize, j: usize, delta: &mut [f64], rng: &mut Rng) -> f64 {177    for a in 0..dim {178        delta[a] = positions[i * dim + a] - positions[j * dim + a];179    }180    let d = delta.iter().map(|x| x * x).sum::<f64>().sqrt();181    if d >= APART {182        return d;183    }184    for x in delta.iter_mut() {185        *x = (rng.unit() - 0.5) * APART;186    }187    delta188        .iter()189        .map(|x| x * x)190        .sum::<f64>()191        .sqrt()192        .max(APART / 2.0)193}194195#[cfg(test)]196mod tests {197    use super::*;198    fn ring() -> Layout {199        let square = [0.0, 0.0, 1.0, 0.0, 1.0, 1.0, 0.0, 1.0];200        let ring = [(0, 1), (1, 2), (2, 3), (3, 0)];201        Layout::new(&square, &ring, 2, 1).unwrap()202    }203    #[test]204    fn the_four_cycle_relaxes_to_a_square() {205        let mut layout = ring();206        assert!((layout.ideal() - 0.5).abs() < 1e-12);207        let rest = layout.step(500);208        assert!(rest < 1e-3, "energy {rest}");209        let p = layout.positions();210        let gaps: Vec<f64> = [(0, 1), (1, 2), (2, 3), (3, 0)]211            .iter()212            .map(|&(a, b)| {213                ((p[2 * a] - p[2 * b]).powi(2) + (p[2 * a + 1] - p[2 * b + 1]).powi(2)).sqrt()214            })215            .collect();216        let spread = gaps.iter().cloned().fold(0.0, f64::max)217            - gaps.iter().cloned().fold(f64::INFINITY, f64::min);218        assert!(spread < 1e-3, "gaps {gaps:?}");219        let side = 0.5 * 1.5f64.cbrt();220        assert!((gaps[0] - side).abs() < 1e-2, "side {}", gaps[0]);221        assert_eq!(layout.ticks(), 500);222    }223    #[test]224    fn the_seed_replays_and_the_cap_cools() {225        let mut a = ring();226        let mut b = ring();227        a.step(20);228        b.step(20);229        assert_eq!(a.positions(), b.positions());230        assert!(231            a.temperature()232                < Layout::new(&[0.0, 0.0, 1.0, 1.0], &[], 2, 1)233                    .unwrap()234                    .temperature()235        );236        assert!(a.moved() > 0.0);237    }238    #[test]239    fn the_faults_are_named() {240        assert!(Layout::new(&[0.0, 0.0, 1.0], &[], 2, 1).is_err());241        assert!(Layout::new(&[0.0, 0.0], &[(0, 1)], 2, 1).is_err());242        assert!(Layout::new(&[0.0, 0.0, 0.0, 0.0], &[(0, 1)], 2, 1)243            .unwrap()244            .step(3)245            .is_finite());246    }247}