From 558078fe359b75bfa2eccfc10378bdd9c7e5a453 Mon Sep 17 00:00:00 2001 From: Ben Ruijl Date: Tue, 10 Sep 2024 17:25:30 +0200 Subject: [PATCH] Add graph support to Python API --- src/api/python.rs | 214 ++++++++++++++++++++++++++++++++++++++++++++++ src/graph.rs | 83 +++++++++++------- src/tensors.rs | 12 ++- symbolica.pyi | 109 +++++++++++++++++++++++ 4 files changed, 382 insertions(+), 36 deletions(-) diff --git a/src/api/python.rs b/src/api/python.rs index 4386ca54..d30e007b 100644 --- a/src/api/python.rs +++ b/src/api/python.rs @@ -45,6 +45,7 @@ use crate::{ CompileOptions, CompiledEvaluator, EvaluationFn, ExpressionEvaluator, FunctionMap, InlineASM, OptimizationSettings, }, + graph::Graph, id::{ Condition, Match, MatchSettings, MatchStack, Pattern, PatternAtomTreeIterator, PatternOrMap, PatternRestriction, ReplaceIterator, Replacement, WildcardAndRestriction, @@ -87,6 +88,7 @@ pub fn create_symbolica_module(m: &PyModule) -> PyResult<&PyModule> { m.add_class::()?; m.add_class::()?; m.add_class::()?; + m.add_class::()?; m.add_function(wrap_pyfunction!(symbol_shorthand, m)?)?; m.add_function(wrap_pyfunction!(number_shorthand, m)?)?; @@ -9959,3 +9961,215 @@ impl PythonNumericalIntegrator { Ok((stats.avg, stats.err, stats.chi_sq / stats.cur_iter as f64)) } } + +/// A graph that supported directional edges, parallel edges, self-edges and custom data on the nodes and edges. +#[pyclass(name = "Graph", module = "symbolica")] +#[derive(Clone, PartialEq, Eq, Hash)] +struct PythonGraph { + graph: Graph, +} + +#[pymethods] +impl PythonGraph { + /// Create an empty graph. + #[new] + fn new() -> Self { + Self { + graph: Graph::new(), + } + } + + /// Print the graph in a human-readable format. + fn __str__(&self) -> String { + format!("{}", self.graph) + } + + /// Hash the graph. + fn __hash__(&self) -> u64 { + let mut hasher = ahash::AHasher::default(); + self.graph.hash(&mut hasher); + hasher.finish() + } + + /// Copy the graph. + fn __copy__(&self) -> PythonGraph { + Self { + graph: self.graph.clone(), + } + } + + /// Get the number of nodes. + fn __len__(&self) -> usize { + self.graph.nodes().len() + } + + /// Compare two graphs. + fn __richcmp__(&self, other: &Self, op: CompareOp) -> PyResult { + match op { + CompareOp::Eq => Ok(self.graph == other.graph), + CompareOp::Ne => Ok(self.graph != other.graph), + _ => Err(exceptions::PyTypeError::new_err(format!( + "Inequalities between graphs are not allowed", + ))), + } + } + + /// Generate all connected graphs with `external_edges` half-edges and the given allowed list + /// of vertex connections. + /// + /// Returns the canonical form of the graph and the size of its automorphism group (including edge permutations). + #[classmethod] + fn generate( + _cls: &PyType, + external_edges: Vec<(ConvertibleToExpression, ConvertibleToExpression)>, + vertex_signatures: Vec>, + max_vertices: Option, + max_loops: Option, + max_bridges: Option, + allow_self_loops: Option, + ) -> PyResult> { + if max_vertices.is_none() && max_loops.is_none() { + return Err(exceptions::PyValueError::new_err( + "At least one of max_vertices or max_loop must be set", + )); + } + + let external_edges: Vec<_> = external_edges + .into_iter() + .map(|(a, b)| (a.to_expression().expr, b.to_expression().expr)) + .collect(); + let vertex_signatures: Vec<_> = vertex_signatures + .into_iter() + .map(|v| v.into_iter().map(|x| x.to_expression().expr).collect()) + .collect(); + + Ok(Graph::generate( + &external_edges, + &vertex_signatures, + max_vertices, + max_loops, + max_bridges, + allow_self_loops.unwrap_or(false), + ) + .into_iter() + .map(|(k, v)| (Self { graph: k }, Atom::new_num(v).into())) + .collect()) + } + + /// Convert the graph to a graphviz dot string. + fn to_dot(&self) -> String { + self.graph.to_dot() + } + + /// Convert the graph to a mermaid string. + fn to_mermaid(&self) -> String { + self.graph.to_mermaid() + } + + /// Add a node with data `data` to the graph, returning the index of the node. + /// The default data is the number 0. + fn add_node(&mut self, data: Option) -> usize { + self.graph + .add_node(data.map(|x| x.to_expression().expr).unwrap_or_default()) + } + + /// Add an edge between the `source` and `target` nodes, returning the index of the edge. + /// Optionally, the edge can be set as directed. The default data is the number 0. + #[pyo3(signature = (source, target, directed = false, data = None))] + fn add_edge( + &mut self, + source: usize, + target: usize, + directed: bool, + data: Option, + ) -> PyResult { + self.graph + .add_edge( + source, + target, + directed, + data.map(|x| x.to_expression().expr).unwrap_or_default(), + ) + .map_err(|e| exceptions::PyValueError::new_err(e)) + } + + /// Get the `idx`th node. + fn __getitem__(&self, idx: isize) -> PyResult<(Vec, PythonExpression)> { + self.node(idx) + } + + /// Get the number of nodes. + fn num_nodes(&self) -> usize { + self.graph.nodes().len() + } + + /// Get the number of edges. + fn num_edges(&self) -> usize { + self.graph.edges().len() + } + + /// Get the number of loops. + fn num_loops(&self) -> usize { + self.graph.num_loops() + } + + /// Get the `idx`th node, consisting of the edge indices and the data. + fn node(&self, idx: isize) -> PyResult<(Vec, PythonExpression)> { + if idx.unsigned_abs() < self.graph.nodes().len() { + let n = if idx < 0 { + self.graph + .node(self.graph.nodes().len() - idx.abs() as usize) + } else { + self.graph.node(idx as usize) + }; + Ok((n.edges.clone(), n.data.clone().into())) + } else { + Err(PyIndexError::new_err(format!( + "Index {} out of bounds: the graph only has {} nodes.", + idx, + self.graph.nodes().len(), + ))) + } + } + + /// Get the `idx`th edge, consisting of the the source vertex, target vertex, whether the edge is directed, and the data. + fn edge(&self, idx: isize) -> PyResult<(usize, usize, bool, PythonExpression)> { + if idx.unsigned_abs() < self.graph.edges().len() { + let e = if idx < 0 { + self.graph + .edge(self.graph.edges().len() - idx.abs() as usize) + } else { + self.graph.edge(idx as usize) + }; + Ok(( + e.vertices.0, + e.vertices.1, + e.directed, + e.data.clone().into(), + )) + } else { + Err(PyIndexError::new_err(format!( + "Index {} out of bounds: the graph only has {} edges.", + idx, + self.graph.edges().len(), + ))) + } + } + + /// Write the graph in a canonical form. + /// Returns the canonicalized graph, the vertex map, the automorphism group size, and the orbit. + fn canonize(&self) -> (PythonGraph, Vec, PythonExpression, Vec) { + let c = self.graph.canonize(); + ( + Self { graph: c.graph }, + c.vertex_map, + Atom::new_num(c.automorphism_group_size).into(), + c.orbit, + ) + } + + /// Return true `iff` the graph is isomorphic to `other`. + fn is_isomorphic(&self, other: &PythonGraph) -> bool { + self.graph.is_isomorphic(&other.graph) + } +} diff --git a/src/graph.rs b/src/graph.rs index d661c6a7..df8d818e 100644 --- a/src/graph.rs +++ b/src/graph.rs @@ -255,7 +255,17 @@ impl Graph { /// Add an edge between vertex indices `source` and `target` to the graph, with arbitrary data. /// If `directed` is true, the edge is directed from `source` to `target`. - pub fn add_edge(&mut self, source: usize, target: usize, directed: bool, data: E) { + pub fn add_edge( + &mut self, + source: usize, + target: usize, + directed: bool, + data: E, + ) -> Result { + if source >= self.nodes.len() || target >= self.nodes.len() { + return Err("Invalid node index"); + } + let index = self.edges.len(); self.edges.push(Edge { vertices: if !directed && source > target { @@ -268,6 +278,7 @@ impl Graph { }); self.nodes[source].edges.push(index); self.nodes[target].edges.push(index); + Ok(index) } /// Delete the last added edge. This operation is O(1). @@ -457,7 +468,8 @@ impl Graph Graph= consume_count) { *count -= consume_count; - self.add_edge(source, cur_target, false, e.clone()); + self.add_edge(source, cur_target, false, e.clone()).unwrap(); self.distribute_edges( source, cur_target, @@ -774,7 +786,7 @@ impl = stack.iter().map(|x| x.selected_vertex.unwrap()).collect(); @@ -1026,7 +1038,8 @@ impl = (0..self.nodes.len()) @@ -1097,7 +1110,13 @@ impl::default(); node.partition = vec![vec![0, 1]]; @@ -1293,8 +1312,8 @@ mod test { let mut g = Graph::new(); let n0 = g.add_node(0); let n1 = g.add_node(0); - g.add_edge(n0, n1, true, 0); - g.add_edge(n0, n1, false, 1); + g.add_edge(n0, n1, true, 0).unwrap(); + g.add_edge(n0, n1, false, 1).unwrap(); let mut node = SearchTreeNode::::default(); node.partition = vec![vec![0, 1]]; @@ -1306,20 +1325,20 @@ mod test { fn isomorphic() { let mut g = Graph::new(); let n0 = g.add_node(0); - g.add_edge(n0, n0, false, 0); - g.add_edge(n0, n0, false, 0); + g.add_edge(n0, n0, false, 0).unwrap(); + g.add_edge(n0, n0, false, 0).unwrap(); let mut g1 = Graph::new(); let n0 = g1.add_node(0); - g1.add_edge(n0, n0, false, 0); + g1.add_edge(n0, n0, false, 0).unwrap(); assert!(!g.is_isomorphic(&g1)); - g1.add_edge(n0, n0, true, 0); + g1.add_edge(n0, n0, true, 0).unwrap(); assert!(!g.is_isomorphic(&g1)); - g.add_edge(n0, n0, true, 0); - g1.add_edge(n0, n0, false, 0); + g.add_edge(n0, n0, true, 0).unwrap(); + g1.add_edge(n0, n0, false, 0).unwrap(); assert!(g.is_isomorphic(&g1)); let _ = g.add_node(1); @@ -1344,22 +1363,22 @@ mod test { let n7 = g.add_node(0); let n8 = g.add_node(1); - g.add_edge(n0, n1, false, 0); - g.add_edge(n0, n3, false, 0); - g.add_edge(n1, n2, false, 0); - g.add_edge(n1, n3, false, 0); - g.add_edge(n1, n4, false, 0); - g.add_edge(n1, n5, false, 0); - g.add_edge(n2, n5, false, 0); - g.add_edge(n3, n4, false, 0); - g.add_edge(n3, n6, false, 0); - g.add_edge(n3, n7, false, 0); - g.add_edge(n4, n5, false, 0); - g.add_edge(n4, n7, false, 0); - g.add_edge(n5, n7, false, 0); - g.add_edge(n5, n8, false, 0); - g.add_edge(n6, n7, false, 0); - g.add_edge(n7, n8, false, 0); + g.add_edge(n0, n1, false, 0).unwrap(); + g.add_edge(n0, n3, false, 0).unwrap(); + g.add_edge(n1, n2, false, 0).unwrap(); + g.add_edge(n1, n3, false, 0).unwrap(); + g.add_edge(n1, n4, false, 0).unwrap(); + g.add_edge(n1, n5, false, 0).unwrap(); + g.add_edge(n2, n5, false, 0).unwrap(); + g.add_edge(n3, n4, false, 0).unwrap(); + g.add_edge(n3, n6, false, 0).unwrap(); + g.add_edge(n3, n7, false, 0).unwrap(); + g.add_edge(n4, n5, false, 0).unwrap(); + g.add_edge(n4, n7, false, 0).unwrap(); + g.add_edge(n5, n7, false, 0).unwrap(); + g.add_edge(n5, n8, false, 0).unwrap(); + g.add_edge(n6, n7, false, 0).unwrap(); + g.add_edge(n7, n8, false, 0).unwrap(); let c = g.canonize(); diff --git a/src/tensors.rs b/src/tensors.rs index 90aa5328..8d0d5bbe 100644 --- a/src/tensors.rs +++ b/src/tensors.rs @@ -354,7 +354,8 @@ impl<'a> AtomView<'a> { .map(|c| c[p].into()) .unwrap_or(Atom::Zero.into()), ), - ); + ) + .unwrap(); connections[p] = None; } else { connections[p] = Some(n); @@ -387,7 +388,8 @@ impl<'a> AtomView<'a> { .map(|c| c[p].into()) .unwrap_or(Atom::Zero.into()), ), - ); + ) + .unwrap(); connections[p] = None; } else { connections[p] = Some(start + i); @@ -411,7 +413,8 @@ impl<'a> AtomView<'a> { }, Atom::Zero.into(), ), - ); + ) + .unwrap(); } } @@ -421,7 +424,8 @@ impl<'a> AtomView<'a> { start, true, (HiddenData::new(0, 0), Atom::Zero.into()), - ); + ) + .unwrap(); } Ok(()) diff --git a/symbolica.pyi b/symbolica.pyi index 74320713..8836923b 100644 --- a/symbolica.pyi +++ b/symbolica.pyi @@ -3316,3 +3316,112 @@ class RandomNumberGenerator: def __new__(_cls, seed: int, stream_id: int): """Create a new random number generator with a given `seed` and `stream_id`. For parallel runs, each thread or instance generating samples should use the same `seed` but a different `stream_id`.""" + + +class Graph: + """A graph that supported directional edges, parallel edges, self-edges and expression data on the nodes and edges.""" + + def __new__(_cls): + """Create a new empty graph.""" + + def __str__(self) -> str: + """Print the graph in a human-readable format.""" + + def __hash__(self) -> int: + """Hash the graph.""" + + def __copy__(self) -> Graph: + """Copy the graph.""" + + def __len__(self) -> int: + """Get the number of nodes in the graph.""" + + def __eq__(self, other: Graph) -> bool: + """Compare two graphs.""" + + def __neq__(self, other: Graph) -> bool: + """Compare two graphs.""" + + def __getitem__(self, idx: int) -> tuple[Sequence[int], Expression]: + """Get the `idx`th node, consisting of the edge indices and the data.""" + + @classmethod + def generate(_cls, + external_nodes: Sequence[tuple[Expression | int, Expression | int]], + vertex_signatures: Sequence[Sequence[Expression | int]], + max_vertices: Optional[int] = None, + max_loops: Optional[int] = None, + max_bridges: Optional[int] = None, + allow_self_edges: bool = False,) -> dict[Graph, Expression]: + """Generate all connected graphs with `external_edges` half-edges and the given allowed list + of vertex connections. + + Returns the canonical form of the graph and the size of its automorphism group (including edge permutations). + + Examples + -------- + >>> from symbolica import * + >>> g, q, qb, gh, ghb = S('g', 'q', 'qb', 'gh', 'ghb') + >>> Graph.generate([(1, g), (2, g)], + >>> [[g, g, g], [g, g, g, g], + >>> [q, qb, g], [gh, ghb, g]], max_loops=2) + >>> for (g, sym) in graphs.items(): + >>> print(f'Symmetry factor = 1/{sym}:') + >>> print(g.to_dot()) + + generates all connected graphs up to 2 loops with the specified vertices. + + Parameters + ---------- + external_nodes: Sequence[tuple[Expression | int, Expression | int]] + The external edges, consisting of a tuple of the node data and the edge data. + If the node data is the same, flip symmetries will be recongized. + vertex_signatures: Sequence[Sequence[Expression | int]] + The allowed connections for each vertex. + max_vertices: int, optional + The maximum number of vertices in the graph. + max_loops: int, optional + The maximum number of loops in the graph. + max_bridges: int, optional + The maximum number of bridges in the graph. + allow_self_edges: bool, optional + Whether self-edges are allowed. + """ + + def to_dot(self) -> str: + """Convert the graph to a graphviz dot string.""" + + def to_mermaid(self) -> str: + """Convert the graph to a mermaid string.""" + + def num_nodes(self) -> int: + """Get the number of nodes in the graph.""" + + def num_edges(self) -> int: + """Get the number of edges in the graph.""" + + def num_loops(self) -> int: + """Get the number of loops in the graph.""" + + def node(self, idx: int) -> tuple[Sequence[int], Expression]: + """Get the `idx`th node, consisting of the edge indices and the data.""" + + def edge(self, idx: int) -> tuple[int, int, bool, Expression]: + """Get the `idx`th edge, consisting of the the source vertex, target vertex, whether the edge is directed, and the data.""" + + def add_node(self, data: Optional[Expression | int] = None) -> int: + """Add a node with data `data` to the graph, returning the index of the node. + The default data is the number 0. + """ + + def add_edge(self, source: int, target: int, directed: bool = False, data: Optional[Expression | int] = None) -> int: + """Add an edge between the `source` and `target` nodes, returning the index of the edge. + + Optionally, the edge can be set as directed. The default data is the number 0. + """ + + def canonize(self) -> Tuple[Graph, Sequence[int], Expression, Sequence[int]]: + """Write the graph in a canonical form. Returns the canonicalized graph, the vertex map, the automorphism group size, and the orbit.""" + + def is_isomorphic(self, other: Graph) -> bool: + """Check if the graph is isomorphic to another graph."""