models.rs

3.2 kB · rust · 104 lines

1use mrlycore::errors::{value_error, Result};2use std::collections::HashMap;34/// A point of the network.5#[derive(Clone, Debug, PartialEq)]6pub struct Node {7    /// The coordinates, one per dimension.8    pub position: Vec<f64>,9    /// The node's place in the network's list.10    pub index: usize,11}1213/// A link between two nodes.14#[derive(Clone, Debug, PartialEq)]15pub struct Branch {16    /// The index of the node the branch leaves.17    pub parent: usize,18    /// The index of the node the branch reaches.19    pub child: usize,20    /// The thickness of the branch.21    pub radius: f64,22}2324/// A spatial graph of nodes and branches.25#[derive(Clone, Debug, PartialEq)]26pub struct Network {27    /// The dimension every position must match.28    pub dim: usize,29    /// The nodes in insertion order.30    pub nodes: Vec<Node>,31    /// The branches in insertion order.32    pub branches: Vec<Branch>,33}3435impl Network {36    /// Builds an empty network of the given dimension.37    pub fn new(dim: usize) -> Network {38        Network {39            dim,40            nodes: Vec::new(),41            branches: Vec::new(),42        }43    }44    /// Appends a node at the position and returns its index, or an error on a dimension mismatch.45    pub fn add_node(&mut self, position: Vec<f64>) -> Result<usize> {46        if position.len() != self.dim {47            return value_error(format!(48                "Expected {}D position, got {}D",49                self.dim,50                position.len()51            ));52        }53        let index = self.nodes.len();54        self.nodes.push(Node { position, index });55        Ok(index)56    }57    /// Appends a branch between two node indices, or an error when either is out of range.58    pub fn add_branch(&mut self, parent: usize, child: usize, radius: f64) -> Result<()> {59        let n = self.nodes.len();60        if parent >= n || child >= n {61            return value_error(format!("Branch endpoints out of range: {parent}, {child}"));62        }63        self.branches.push(Branch {64            parent,65            child,66            radius,67        });68        Ok(())69    }70    /// Returns each node's branch count, indexed like the node list.71    pub fn degree(&self) -> Vec<usize> {72        let mut deg = vec![0; self.nodes.len()];73        for b in &self.branches {74            deg[b.parent] += 1;75            deg[b.child] += 1;76        }77        deg78    }79    /// Returns the undirected neighbor lists of every node.80    pub fn adjacency(&self) -> HashMap<usize, Vec<usize>> {81        let mut adj: HashMap<usize, Vec<usize>> =82            (0..self.nodes.len()).map(|i| (i, Vec::new())).collect();83        for b in &self.branches {84            adj.get_mut(&b.parent).unwrap().push(b.child);85            adj.get_mut(&b.child).unwrap().push(b.parent);86        }87        adj88    }89}9091#[cfg(test)]92mod tests {93    use super::*;94    #[test]95    fn build_and_measure() {96        let mut n = Network::new(2);97        let a = n.add_node(vec![0.0, 0.0]).unwrap();98        let b = n.add_node(vec![1.0, 0.0]).unwrap();99        n.add_branch(a, b, 1.0).unwrap();100        assert_eq!(n.degree(), vec![1, 1]);101        assert!(n.add_node(vec![0.0]).is_err());102        assert!(n.add_branch(0, 5, 1.0).is_err());103    }104}