models.rs

4.6 kB · rust · 150 lines

1use crate::core::error::{value_error, Result};2use serde::{Deserialize, Serialize};3use std::collections::HashMap;45/// A point of the network.6#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]7pub struct Node {8    /// The coordinates, one per dimension.9    pub position: Vec<f64>,10    /// The node's place in the network's list.11    pub index: usize,12}1314/// A link between two nodes.15#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]16pub struct Branch {17    /// The index of the node the branch leaves.18    pub parent: usize,19    /// The index of the node the branch reaches.20    pub child: usize,21    /// The thickness of the branch.22    pub radius: f64,23}2425/// A spatial graph of nodes and branches.26#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]27pub struct Network {28    /// The dimension every position must match.29    pub dim: usize,30    /// The nodes in insertion order.31    pub nodes: Vec<Node>,32    /// The branches in insertion order.33    pub branches: Vec<Branch>,34}3536impl Network {37    /// Builds an empty network of the given dimension.38    pub fn new(dim: usize) -> Network {39        Network {40            dim,41            nodes: Vec::new(),42            branches: Vec::new(),43        }44    }45    /// Appends a node at the position and returns its index.46    ///47    /// # Errors48    ///49    /// Errors on a dimension mismatch.50    pub fn add_node(&mut self, position: Vec<f64>) -> Result<usize> {51        if position.len() != self.dim {52            return value_error(format!(53                "Expected {}D position, got {}D",54                self.dim,55                position.len()56            ));57        }58        let index = self.nodes.len();59        self.nodes.push(Node { position, index });60        Ok(index)61    }62    /// Appends a branch between two node indices.63    ///64    /// # Errors65    ///66    /// Errors when either index is out of range.67    pub fn add_branch(&mut self, parent: usize, child: usize, radius: f64) -> Result<()> {68        let n = self.nodes.len();69        if parent >= n || child >= n {70            return value_error(format!("Branch endpoints out of range: {parent}, {child}"));71        }72        self.branches.push(Branch {73            parent,74            child,75            radius,76        });77        Ok(())78    }79    /// Returns each node's branch count, indexed like the node list.80    ///81    /// # Errors82    ///83    /// Errors when a branch names a node the network does not hold.84    pub fn degree(&self) -> Result<Vec<usize>> {85        let n = self.nodes.len();86        let mut deg = vec![0; n];87        for b in &self.branches {88            if b.parent >= n || b.child >= n {89                return value_error(format!(90                    "Branch endpoints out of range: {}, {}",91                    b.parent, b.child92                ));93            }94            deg[b.parent] += 1;95            deg[b.child] += 1;96        }97        Ok(deg)98    }99    /// Returns the undirected neighbor lists of every node.100    ///101    /// # Errors102    ///103    /// Errors when a branch names a node the network does not hold.104    pub fn adjacency(&self) -> Result<HashMap<usize, Vec<usize>>> {105        let n = self.nodes.len();106        let mut adj: HashMap<usize, Vec<usize>> = (0..n).map(|i| (i, Vec::new())).collect();107        for b in &self.branches {108            if b.parent >= n || b.child >= n {109                return value_error(format!(110                    "Branch endpoints out of range: {}, {}",111                    b.parent, b.child112                ));113            }114            adj.entry(b.parent).or_default().push(b.child);115            adj.entry(b.child).or_default().push(b.parent);116        }117        Ok(adj)118    }119}120121#[cfg(test)]122mod tests {123    use super::*;124    #[test]125    fn build_and_measure() {126        let mut n = Network::new(2);127        let a = n.add_node(vec![0.0, 0.0]).unwrap();128        let b = n.add_node(vec![1.0, 0.0]).unwrap();129        n.add_branch(a, b, 1.0).unwrap();130        assert_eq!(n.degree().unwrap(), vec![1, 1]);131        assert!(n.add_node(vec![0.0]).is_err());132        assert!(n.add_branch(0, 5, 1.0).is_err());133    }134135    #[test]136    fn refuses_a_stray_dimension_and_a_branch_past_the_nodes() {137        let mut net = Network::new(2);138        net.add_node(vec![0.0, 0.0]).unwrap();139        net.add_node(vec![1.0, 0.0]).unwrap();140        assert!(net.add_node(vec![0.0]).is_err());141        assert!(net.add_branch(0, 5, 1.0).is_err());142        net.branches.push(Branch {143            parent: 0,144            child: 9,145            radius: 1.0,146        });147        assert!(net.degree().is_err());148        assert!(net.adjacency().is_err());149    }150}