ipfrs_tensorlogic/
visualization.rs

1//! Visualization utilities for computation graphs and proofs.
2//!
3//! This module provides tools for exporting graphs and proofs to DOT format
4//! (Graphviz) for visualization and debugging.
5//!
6//! # Examples
7//!
8//! ## Visualizing a Computation Graph
9//!
10//! ```
11//! use ipfrs_tensorlogic::{ComputationGraph, GraphNode, TensorOp, GraphVisualizer};
12//!
13//! let mut graph = ComputationGraph::new();
14//!
15//! // Create nodes
16//! let input = GraphNode::new("input".to_string(), TensorOp::Input {
17//!     name: "x".to_string(),
18//! });
19//! graph.add_node(input).expect("write to String is infallible");
20//! graph.mark_input("input".to_string());
21//!
22//! let relu = GraphNode::new("relu".to_string(), TensorOp::ReLU)
23//!     .add_input("input".to_string());
24//! graph.add_node(relu).expect("write to String is infallible");
25//! graph.mark_output("relu".to_string());
26//!
27//! // Export to DOT format
28//! let dot = GraphVisualizer::to_dot(&graph);
29//! println!("{}", dot);
30//! // Save to file: std::fs::write("graph.dot", dot).expect("write to String is infallible");
31//! // Render: dot -Tpng graph.dot -o graph.png
32//! ```
33
34use crate::computation_graph::{ComputationGraph, TensorOp};
35use crate::proof_storage::ProofFragment;
36use std::fmt::Write as FmtWrite;
37
38/// Visualizer for computation graphs
39pub struct GraphVisualizer;
40
41impl GraphVisualizer {
42    /// Export a computation graph to DOT format
43    ///
44    /// The output can be rendered using Graphviz:
45    /// ```bash
46    /// dot -Tpng graph.dot -o graph.png
47    /// dot -Tsvg graph.dot -o graph.svg
48    /// ```
49    pub fn to_dot(graph: &ComputationGraph) -> String {
50        let mut dot = String::new();
51        writeln!(dot, "digraph ComputationGraph {{").expect("write to String is infallible");
52        writeln!(dot, "  rankdir=TB;").expect("write to String is infallible");
53        writeln!(dot, "  node [shape=box, style=filled];").expect("write to String is infallible");
54        writeln!(dot).expect("write to String is infallible");
55
56        // Write nodes
57        for (node_id, node) in &graph.nodes {
58            let color = Self::node_color(&node.op);
59            let shape = if graph.inputs.contains(node_id) {
60                "ellipse"
61            } else if graph.outputs.contains(node_id) {
62                "doubleoctagon"
63            } else {
64                "box"
65            };
66
67            let label = Self::format_operation(&node.op);
68            writeln!(
69                dot,
70                "  \"{}\" [label=\"{}\\n{}\", fillcolor=\"{}\", shape={}];",
71                Self::escape(node_id),
72                Self::escape(node_id),
73                label,
74                color,
75                shape
76            )
77            .expect("write to String is infallible");
78        }
79
80        writeln!(dot).expect("write to String is infallible");
81
82        // Write edges
83        for (node_id, node) in &graph.nodes {
84            for input in &node.inputs {
85                writeln!(
86                    dot,
87                    "  \"{}\" -> \"{}\";",
88                    Self::escape(input),
89                    Self::escape(node_id)
90                )
91                .expect("write to String is infallible");
92            }
93        }
94
95        // Add legend
96        writeln!(dot).expect("write to String is infallible");
97        writeln!(dot, "  subgraph cluster_legend {{").expect("write to String is infallible");
98        writeln!(dot, "    label=\"Legend\";").expect("write to String is infallible");
99        writeln!(dot, "    style=filled;").expect("write to String is infallible");
100        writeln!(dot, "    fillcolor=lightgrey;").expect("write to String is infallible");
101        writeln!(
102            dot,
103            "    legend_input [label=\"Input\", shape=ellipse, fillcolor=lightblue];"
104        )
105        .expect("write to String is infallible");
106        writeln!(
107            dot,
108            "    legend_output [label=\"Output\", shape=doubleoctagon, fillcolor=lightgreen];"
109        )
110        .expect("write to String is infallible");
111        writeln!(
112            dot,
113            "    legend_compute [label=\"Compute\", shape=box, fillcolor=lightyellow];"
114        )
115        .expect("write to String is infallible");
116        writeln!(dot, "  }}").expect("write to String is infallible");
117
118        writeln!(dot, "}}").expect("write to String is infallible");
119        dot
120    }
121
122    /// Get color for a node based on operation type
123    fn node_color(op: &TensorOp) -> &'static str {
124        match op {
125            TensorOp::Input { .. } | TensorOp::Constant { .. } => "lightblue",
126            TensorOp::MatMul | TensorOp::Einsum { .. } => "orange",
127            TensorOp::Add | TensorOp::Mul | TensorOp::Sub | TensorOp::Div => "yellow",
128            TensorOp::ReLU
129            | TensorOp::Tanh
130            | TensorOp::Sigmoid
131            | TensorOp::GELU
132            | TensorOp::Softmax { .. } => "lightgreen",
133            TensorOp::LayerNorm { .. } | TensorOp::BatchNorm { .. } => "lightcoral",
134            TensorOp::Dropout { .. } => "plum",
135            TensorOp::Reshape { .. } | TensorOp::Transpose { .. } | TensorOp::Slice { .. } => {
136                "lightyellow"
137            }
138            _ => "white",
139        }
140    }
141
142    /// Format operation for display
143    fn format_operation(op: &TensorOp) -> String {
144        match op {
145            TensorOp::Input { name } => format!("Input({})", name),
146            TensorOp::Constant { value_cid } => format!("Const(cid:{})", &value_cid[..8]),
147            TensorOp::MatMul => "MatMul".to_string(),
148            TensorOp::Einsum { subscripts } => format!("Einsum({})", subscripts),
149            TensorOp::Add => "Add".to_string(),
150            TensorOp::Mul => "Multiply".to_string(),
151            TensorOp::Sub => "Subtract".to_string(),
152            TensorOp::Div => "Divide".to_string(),
153            TensorOp::ReLU => "ReLU".to_string(),
154            TensorOp::Tanh => "Tanh".to_string(),
155            TensorOp::Sigmoid => "Sigmoid".to_string(),
156            TensorOp::GELU => "GELU".to_string(),
157            TensorOp::Softmax { axis } => format!("Softmax(axis={})", axis),
158            TensorOp::LayerNorm {
159                normalized_shape: _,
160                eps,
161            } => format!("LayerNorm(ε={:.1e})", eps),
162            TensorOp::BatchNorm { eps, momentum } => {
163                format!("BatchNorm(ε={:.1e}, μ={:.2})", eps, momentum)
164            }
165            TensorOp::Dropout { p } => format!("Dropout({:.2})", p),
166            TensorOp::Reshape { shape } => format!("Reshape({:?})", shape),
167            TensorOp::Transpose { axes } => format!("Transpose({:?})", axes),
168            TensorOp::ReduceSum { axes, keepdims: _ } => format!("ReduceSum({:?})", axes),
169            TensorOp::ReduceMean { axes, keepdims: _ } => format!("ReduceMean({:?})", axes),
170            TensorOp::Concat { axis } => format!("Concat(axis={})", axis),
171            TensorOp::Split { axis, sections } => {
172                format!("Split(axis={}, n={})", axis, sections.len())
173            }
174            TensorOp::Gather { axis } => format!("Gather(axis={})", axis),
175            TensorOp::Scatter { axis } => format!("Scatter(axis={})", axis),
176            TensorOp::Slice {
177                start,
178                end,
179                strides,
180            } => format!("Slice({:?}:{:?}:{:?})", start, end, strides),
181            TensorOp::Pad { padding, mode: _ } => format!("Pad({:?})", padding),
182            TensorOp::Exp => "Exp".to_string(),
183            TensorOp::Log => "Log".to_string(),
184            TensorOp::Pow { exponent } => format!("Pow({})", exponent),
185            TensorOp::Sqrt => "Sqrt".to_string(),
186            TensorOp::FusedLinear => "FusedLinear".to_string(),
187            TensorOp::FusedAddReLU => "FusedAdd+ReLU".to_string(),
188            TensorOp::FusedBatchNormReLU { eps, momentum } => {
189                format!("FusedBN+ReLU(ε={:.1e}, μ={:.2})", eps, momentum)
190            }
191            TensorOp::FusedLayerNormDropout {
192                normalized_shape: _,
193                eps,
194                dropout_p,
195            } => format!("FusedLN+Dropout(ε={:.1e}, p={:.2})", eps, dropout_p),
196        }
197    }
198
199    /// Escape special characters for DOT format
200    fn escape(s: &str) -> String {
201        s.replace('\"', "\\\"")
202            .replace('\n', "\\n")
203            .replace('\t', "\\t")
204    }
205
206    /// Export graph statistics
207    pub fn graph_stats(graph: &ComputationGraph) -> String {
208        let mut stats = String::new();
209        writeln!(stats, "Graph Statistics:").expect("write to String is infallible");
210        writeln!(stats, "  Total nodes: {}", graph.nodes.len())
211            .expect("write to String is infallible");
212        writeln!(stats, "  Input nodes: {}", graph.inputs.len())
213            .expect("write to String is infallible");
214        writeln!(stats, "  Output nodes: {}", graph.outputs.len())
215            .expect("write to String is infallible");
216
217        // Count operation types
218        let mut op_counts = std::collections::HashMap::new();
219        for node in graph.nodes.values() {
220            let op_name = Self::operation_name(&node.op);
221            *op_counts.entry(op_name).or_insert(0) += 1;
222        }
223
224        writeln!(stats, "  Operation counts:").expect("write to String is infallible");
225        let mut ops: Vec<_> = op_counts.into_iter().collect();
226        ops.sort_by_key(|a| std::cmp::Reverse(a.1));
227        for (op, count) in ops {
228            writeln!(stats, "    {}: {}", op, count).expect("write to String is infallible");
229        }
230
231        stats
232    }
233
234    fn operation_name(op: &TensorOp) -> &'static str {
235        match op {
236            TensorOp::Input { .. } => "Input",
237            TensorOp::Constant { .. } => "Constant",
238            TensorOp::MatMul => "MatMul",
239            TensorOp::Einsum { .. } => "Einsum",
240            TensorOp::Add => "Add",
241            TensorOp::Mul => "Mul",
242            TensorOp::Sub => "Sub",
243            TensorOp::Div => "Div",
244            TensorOp::ReLU => "ReLU",
245            TensorOp::Tanh => "Tanh",
246            TensorOp::Sigmoid => "Sigmoid",
247            TensorOp::GELU => "GELU",
248            TensorOp::Softmax { .. } => "Softmax",
249            TensorOp::LayerNorm { .. } => "LayerNorm",
250            TensorOp::BatchNorm { .. } => "BatchNorm",
251            TensorOp::Dropout { .. } => "Dropout",
252            TensorOp::Reshape { .. } => "Reshape",
253            TensorOp::Transpose { .. } => "Transpose",
254            TensorOp::ReduceSum { .. } => "ReduceSum",
255            TensorOp::ReduceMean { .. } => "ReduceMean",
256            TensorOp::Concat { .. } => "Concat",
257            TensorOp::Split { .. } => "Split",
258            TensorOp::Gather { .. } => "Gather",
259            TensorOp::Scatter { .. } => "Scatter",
260            TensorOp::Slice { .. } => "Slice",
261            TensorOp::Pad { .. } => "Pad",
262            TensorOp::Exp => "Exp",
263            TensorOp::Log => "Log",
264            TensorOp::Pow { .. } => "Pow",
265            TensorOp::Sqrt => "Sqrt",
266            TensorOp::FusedLinear => "FusedLinear",
267            TensorOp::FusedAddReLU => "FusedAddReLU",
268            TensorOp::FusedBatchNormReLU { .. } => "FusedBatchNormReLU",
269            TensorOp::FusedLayerNormDropout { .. } => "FusedLayerNormDropout",
270        }
271    }
272}
273
274/// Visualizer for proof trees
275pub struct ProofVisualizer;
276
277impl ProofVisualizer {
278    /// Export a proof tree to DOT format
279    ///
280    /// The proof is rendered as a tree with the conclusion at the top
281    /// and premises as child nodes.
282    pub fn to_dot(proof: &ProofFragment, id: usize) -> String {
283        let mut dot = String::new();
284        writeln!(dot, "digraph ProofTree {{").expect("write to String is infallible");
285        writeln!(dot, "  rankdir=TB;").expect("write to String is infallible");
286        writeln!(dot, "  node [shape=box, style=\"filled,rounded\"];")
287            .expect("write to String is infallible");
288        writeln!(dot).expect("write to String is infallible");
289
290        let mut node_counter = 0;
291        Self::write_proof_node(&mut dot, proof, id, &mut node_counter);
292
293        writeln!(dot, "}}").expect("write to String is infallible");
294        dot
295    }
296
297    fn write_proof_node(
298        dot: &mut String,
299        proof: &ProofFragment,
300        node_id: usize,
301        counter: &mut usize,
302    ) {
303        let color = if proof.premise_refs.is_empty() {
304            "lightblue" // Fact (no premises)
305        } else {
306            "lightyellow" // Rule application
307        };
308
309        let conclusion_str = format!("{:?}", proof.conclusion);
310        writeln!(
311            dot,
312            "  node_{} [label=\"{}\", fillcolor=\"{}\"];",
313            node_id,
314            GraphVisualizer::escape(&conclusion_str),
315            color
316        )
317        .expect("write to String is infallible");
318
319        // Write premise references as child nodes
320        for premise_ref in &proof.premise_refs {
321            *counter += 1;
322            let premise_id = *counter;
323            let premise_str = if let Some(ref hint) = premise_ref.conclusion_hint {
324                hint.clone()
325            } else {
326                format!("CID: {}", premise_ref.cid)
327            };
328            writeln!(
329                dot,
330                "  node_{} [label=\"{}\", fillcolor=\"lightgray\"];",
331                premise_id,
332                GraphVisualizer::escape(&premise_str)
333            )
334            .expect("write to String is infallible");
335            writeln!(dot, "  node_{} -> node_{};", node_id, premise_id)
336                .expect("write to String is infallible");
337        }
338
339        // Add rule information
340        if let Some(ref rule_ref) = proof.rule_applied {
341            writeln!(
342                dot,
343                "  node_{}_rule [label=\"Rule: {}\", shape=note, fillcolor=\"lightyellow\"];",
344                node_id,
345                GraphVisualizer::escape(&rule_ref.rule_id)
346            )
347            .expect("write to String is infallible");
348            writeln!(
349                dot,
350                "  node_{}_rule -> node_{} [style=dashed];",
351                node_id, node_id
352            )
353            .expect("write to String is infallible");
354        }
355    }
356
357    /// Generate a textual explanation of a proof
358    pub fn explain(proof: &ProofFragment, depth: usize) -> String {
359        let mut explanation = String::new();
360        let indent = "  ".repeat(depth);
361
362        writeln!(explanation, "{}Prove: {:?}", indent, proof.conclusion)
363            .expect("write to String is infallible");
364
365        if proof.premise_refs.is_empty() {
366            writeln!(explanation, "{}  ✓ This is a known fact", indent)
367                .expect("write to String is infallible");
368        } else {
369            if let Some(ref rule_ref) = proof.rule_applied {
370                writeln!(explanation, "{}  Using rule: {}", indent, rule_ref.rule_id)
371                    .expect("write to String is infallible");
372            }
373            writeln!(
374                explanation,
375                "{}  Requires proving {} premise(s):",
376                indent,
377                proof.premise_refs.len()
378            )
379            .expect("write to String is infallible");
380            for (i, premise_ref) in proof.premise_refs.iter().enumerate() {
381                let hint = premise_ref
382                    .conclusion_hint
383                    .as_deref()
384                    .unwrap_or("(premise)");
385                writeln!(explanation, "{}    {}. {}", indent, i + 1, hint)
386                    .expect("write to String is infallible");
387            }
388        }
389
390        if let Some(complexity) = proof.metadata.complexity {
391            writeln!(explanation, "{}  Complexity: {} steps", indent, complexity)
392                .expect("write to String is infallible");
393        }
394        writeln!(explanation, "{}  Depth: {}", indent, proof.metadata.depth)
395            .expect("write to String is infallible");
396
397        explanation
398    }
399
400    /// Generate a summary of proof statistics
401    pub fn proof_stats(proof: &ProofFragment) -> String {
402        let mut stats = String::new();
403        writeln!(stats, "Proof Statistics:").expect("write to String is infallible");
404        writeln!(stats, "  ID: {}", proof.id).expect("write to String is infallible");
405        writeln!(stats, "  Direct premises: {}", proof.premise_refs.len())
406            .expect("write to String is infallible");
407
408        writeln!(
409            stats,
410            "  Complexity: {} steps",
411            proof.metadata.complexity.unwrap_or(0)
412        )
413        .expect("write to String is infallible");
414        writeln!(stats, "  Depth: {}", proof.metadata.depth)
415            .expect("write to String is infallible");
416        if let Some(ref created_by) = proof.metadata.created_by {
417            writeln!(stats, "  Created by: {}", created_by).expect("write to String is infallible");
418        }
419
420        if proof.premise_refs.is_empty() {
421            writeln!(stats, "  Type: Fact (axiom)").expect("write to String is infallible");
422        } else {
423            writeln!(stats, "  Type: Rule application").expect("write to String is infallible");
424            if let Some(ref rule_ref) = proof.rule_applied {
425                writeln!(stats, "  Rule: {}", rule_ref.rule_id)
426                    .expect("write to String is infallible");
427            }
428        }
429
430        if !proof.substitution.is_empty() {
431            writeln!(stats, "  Substitutions: {}", proof.substitution.len())
432                .expect("write to String is infallible");
433        }
434
435        stats
436    }
437}
438
439#[cfg(test)]
440mod tests {
441    use super::*;
442    use crate::{ComputationGraph, GraphNode, Predicate, TensorOp, Term};
443
444    #[test]
445    fn test_graph_to_dot() {
446        let mut graph = ComputationGraph::new();
447
448        let input = GraphNode::new(
449            "input".to_string(),
450            TensorOp::Input {
451                name: "x".to_string(),
452            },
453        );
454        graph.add_node(input).expect("test: should succeed");
455        graph.mark_input("input".to_string());
456
457        let relu =
458            GraphNode::new("relu".to_string(), TensorOp::ReLU).add_input("input".to_string());
459        graph.add_node(relu).expect("test: should succeed");
460        graph.mark_output("relu".to_string());
461
462        let dot = GraphVisualizer::to_dot(&graph);
463
464        assert!(dot.contains("digraph ComputationGraph"));
465        assert!(dot.contains("\"input\""));
466        assert!(dot.contains("\"relu\""));
467        assert!(dot.contains("\"input\" -> \"relu\""));
468    }
469
470    #[test]
471    fn test_graph_stats() {
472        let mut graph = ComputationGraph::new();
473
474        let input = GraphNode::new(
475            "input".to_string(),
476            TensorOp::Input {
477                name: "x".to_string(),
478            },
479        );
480        graph.add_node(input).expect("test: should succeed");
481
482        let relu =
483            GraphNode::new("relu".to_string(), TensorOp::ReLU).add_input("input".to_string());
484        graph.add_node(relu).expect("test: should succeed");
485
486        let stats = GraphVisualizer::graph_stats(&graph);
487
488        assert!(stats.contains("Total nodes: 2"));
489        assert!(stats.contains("Input: 1"));
490        assert!(stats.contains("ReLU: 1"));
491    }
492
493    #[test]
494    fn test_proof_to_dot() {
495        use crate::proof_storage::{ProofFragmentRef, ProofMetadata, RuleRef};
496
497        let conclusion = Predicate::new(
498            "ancestor".to_string(),
499            vec![
500                Term::Const(crate::Constant::String("Alice".to_string())),
501                Term::Const(crate::Constant::String("Bob".to_string())),
502            ],
503        );
504
505        let proof = ProofFragment {
506            id: "proof_1".to_string(),
507            conclusion,
508            rule_applied: Some(RuleRef {
509                rule_id: "ancestor_rule".to_string(),
510                rule_cid: None,
511                rule: None,
512            }),
513            premise_refs: vec![ProofFragmentRef {
514                cid: ipfrs_core::Cid::default(),
515                conclusion_hint: Some("parent(Alice, Bob)".to_string()),
516            }],
517            substitution: vec![],
518            metadata: ProofMetadata {
519                created_at: None,
520                created_by: None,
521                complexity: Some(2),
522                depth: 1,
523                custom: std::collections::HashMap::new(),
524            },
525        };
526
527        let dot = ProofVisualizer::to_dot(&proof, 0);
528
529        assert!(dot.contains("digraph ProofTree"));
530        assert!(dot.contains("ancestor"));
531        assert!(dot.contains("parent"));
532    }
533
534    #[test]
535    fn test_proof_explain() {
536        use crate::proof_storage::ProofMetadata;
537
538        let conclusion = Predicate::new(
539            "test".to_string(),
540            vec![Term::Const(crate::Constant::String("A".to_string()))],
541        );
542
543        let proof = ProofFragment {
544            id: "proof_2".to_string(),
545            conclusion,
546            rule_applied: None,
547            premise_refs: vec![],
548            substitution: vec![],
549            metadata: ProofMetadata {
550                created_at: None,
551                created_by: None,
552                complexity: None,
553                depth: 0,
554                custom: std::collections::HashMap::new(),
555            },
556        };
557
558        let explanation = ProofVisualizer::explain(&proof, 0);
559
560        assert!(explanation.contains("Prove"));
561        assert!(explanation.contains("known fact"));
562    }
563
564    #[test]
565    fn test_proof_stats() {
566        use crate::proof_storage::ProofMetadata;
567
568        let conclusion = Predicate::new(
569            "test".to_string(),
570            vec![Term::Const(crate::Constant::String("A".to_string()))],
571        );
572
573        let proof = ProofFragment {
574            id: "proof_3".to_string(),
575            conclusion,
576            rule_applied: None,
577            premise_refs: vec![],
578            substitution: vec![],
579            metadata: ProofMetadata {
580                created_at: None,
581                created_by: None,
582                complexity: None,
583                depth: 0,
584                custom: std::collections::HashMap::new(),
585            },
586        };
587
588        let stats = ProofVisualizer::proof_stats(&proof);
589
590        assert!(stats.contains("Proof Statistics"));
591        assert!(stats.contains("Type: Fact"));
592    }
593}
ipfrs_tensorlogic/visualization.rs

ipfrs_tensorlogic/
visualization.rs
