feat: add dot file creation from tensor tree

HepLean/Tensors/Tree/Basic.lean

@@ -54,7 +54,7 @@ inductive TensorTree (S : TensorStruct) : ∀ {n : ℕ}, (Fin n → S.C) → Typ
       (σ : (OverColor.mk c) ⟶ (OverColor.mk c1)) (t : TensorTree S c) : TensorTree S c1
   | prod {n m : ℕ} {c : Fin n → S.C} {c1 : Fin m → S.C}
     (t : TensorTree S c) (t1 : TensorTree S c1) : TensorTree S (Sum.elim c c1 ∘ finSumFinEquiv.symm)
-  | scale {n : ℕ} {c : Fin n → S.C} : S.k → TensorTree S c → TensorTree S c
+  | smul {n : ℕ} {c : Fin n → S.C} : S.k → TensorTree S c → TensorTree S c
   | mult {n m : ℕ} {c : Fin n.succ → S.C} {c1 : Fin m.succ → S.C} :
     (i : Fin n.succ) → (j : Fin m.succ) → TensorTree S c → TensorTree S c1 →
     TensorTree S (Sum.elim (c ∘ Fin.succAbove i) (c1 ∘ Fin.succAbove j) ∘ finSumFinEquiv.symm)
@@ -78,7 +78,7 @@ def size : ∀ {n : ℕ} {c : Fin n → S.C}, TensorTree S c → ℕ := fun
   | tensorNode _ => 1
   | add t1 t2 => t1.size + t2.size + 1
   | perm _ t => t.size + 1
-  | scale _ t => t.size + 1
+  | smul _ t => t.size + 1
   | prod t1 t2 => t1.size + t2.size + 1
   | mult _ _ t1 t2 => t1.size + t2.size + 1
   | contr _ _ t => t.size + 1
@@ -93,7 +93,7 @@ def tensor : ∀ {n : ℕ} {c : Fin n → S.C}, TensorTree S c → S.F.obj (Over
   | tensorNode t => t
   | add t1 t2 => t1.tensor + t2.tensor
   | perm σ t => (S.F.map σ).hom t.tensor
-  | scale a t => a • t.tensor
+  | smul a t => a • t.tensor
   | prod t1 t2 => (S.F.map (OverColor.equivToIso finSumFinEquiv).hom).hom
     ((S.F.μ _ _).hom (t1.tensor ⊗ₜ t2.tensor))
   | _ => 0

HepLean/Tensors/Tree/Dot.lean

@@ -0,0 +1,92 @@
+/-
+Copyright (c) 2024 Joseph Tooby-Smith. All rights reserved.
+Released under Apache 2.0 license as described in the file LICENSE.
+Authors: Joseph Tooby-Smith
+-/
+import HepLean.Tensors.Tree.Basic
+import Lean.Elab.Term
+/-!
+
+## Tensor trees to dot files
+
+Dot files are used by graphviz to create visualizations of graphs. Here we define a function
+that converts a tensor tree into a dot file.
+
+-/
+
+namespace TensorTree
+
+/-- Turns a nodes of tensor trees into nodes and edges of a dot file. -/
+def dotString (m  : ℕ) (nt : ℕ): ∀ {n : ℕ} {c : Fin n → S.C}, TensorTree S c → String := fun
+  | tensorNode _ =>
+    "node" ++ toString m ++ " [label=\"T" ++ toString nt ++ "\"];\n"
+  | add t1 t2 =>
+    let addNode := "node" ++ toString m ++ " [label=\"+\", shape=box];\n"
+    let edge1 := "node" ++ toString m ++ " -> node" ++ toString (m + 1) ++ ";\n"
+    let edge2 := "node" ++ toString m ++ " -> node" ++ toString (m + 2) ++ ";\n"
+    addNode ++ dotString (m + 1) nt t1 ++ dotString (m + 2) (nt + 1) t2 ++ edge1 ++ edge2
+  | perm σ t =>
+    let permNode := "node" ++ toString m ++ " [label=\"perm\", shape=box];\n"
+    let edge1 := "node" ++ toString m ++ " -> node" ++ toString (m + 1) ++ ";\n"
+    permNode ++ dotString (m + 1) nt t ++ edge1
+  | prod t1 t2 =>
+    let prodNode := "node" ++ toString m ++ " [label=\"⊗\", shape=box];\n"
+    let edge1 := "node" ++ toString m ++ " -> node" ++ toString (m + 1) ++ ";\n"
+    let edge2 := "node" ++ toString m ++ " -> node" ++ toString (m + 2) ++ ";\n"
+    prodNode ++ dotString (m + 1) nt t1 ++ dotString (m + 2) (nt + 1) t2 ++ edge1 ++ edge2
+  | smul k t =>
+    let smulNode := "node" ++ toString m ++ " [label=\"smul\", shape=box];\n"
+    let edge1 := "node" ++ toString m ++ " -> node" ++ toString (m + 1) ++ ";\n"
+    smulNode ++ dotString (m + 1) nt t ++ edge1
+  | mult _ _ t1 t2 =>
+    let multNode := "node" ++ toString m ++ " [label=\"mult\", shape=box];\n"
+    let edge1 := "node" ++ toString m ++ " -> node" ++ toString (m + 1) ++ ";\n"
+    let edge2 := "node" ++ toString m ++ " -> node" ++ toString (m + 2) ++ ";\n"
+    multNode ++ dotString (m + 1) nt t1 ++ dotString (m + 2) (nt + 1) t2 ++ edge1 ++ edge2
+  | eval _ _ t1 =>
+    let evalNode := "node" ++ toString m ++ " [label=\"eval\", shape=box];\n"
+    let edge1 := "node" ++ toString m ++ " -> node" ++ toString (m + 1) ++ ";\n"
+    evalNode ++ dotString (m + 1) nt t1 ++ edge1
+  | jiggle i t1 =>
+    let jiggleNode := "node" ++ toString m ++ " [label=\"jiggle\", shape=box];\n"
+    let edge1 := "node" ++ toString m ++ " -> node" ++ toString (m + 1) ++ ";\n"
+    jiggleNode ++ dotString (m + 1) nt t1 ++ edge1
+  | contr i j t1 =>
+    let contrNode := "node" ++ toString m ++ " [label=\"contr " ++ toString i ++ " "
+      ++ toString j ++ "\", shape=box];\n"
+    let edge1 := "node" ++ toString m ++ " -> node" ++ toString (m + 1) ++ ";\n"
+    contrNode ++ dotString (m + 1) nt t1 ++ edge1
+
+/-- Used to form a dot graph from a tensor tree. use e.g.
+  `#tensor_dot {T4 | i j l d ⊗ T5 | i j k a b}ᵀ.dot`.
+  Dot files can be viewed by copying and pasting into: https://dreampuf.github.io/GraphvizOnline. -/
+def dot {n : ℕ} {c : Fin n → S.C} (t : TensorTree S c) : String :=
+  "digraph G {\n" ++ dotString 0 0 t ++ "}"
+
+section dotElab
+
+open Lean.Elab.Command Lean Meta Lean.Elab
+
+syntax (name := tensorDot) "#tensor_dot " term : command
+
+/-- Adapted from `Lean.Elab.Command.elabReduce` in file copyright Microsoft Corporation. -/
+unsafe def dotElab  (term : Syntax)  : CommandElabM Unit :=
+  withoutModifyingEnv <| runTermElabM fun _ => Term.withDeclName `_reduce do
+    let e ← Term.elabTerm term none
+    Term.synthesizeSyntheticMVarsNoPostponing
+    let e ← Term.levelMVarToParam (← instantiateMVars e)
+    withTheReader Core.Context (fun ctx => { ctx with options := ctx.options.setBool `smartUnfolding false }) do
+      let e ← withTransparency (mode := TransparencyMode.all) <|
+        reduce e (skipProofs := true) (skipTypes := true)
+      let str ← Lean.Meta.evalExpr' String ``String e
+      println! str
+
+/-- Elaborator for the syntax tensorDot. -/
+@[command_elab tensorDot]
+unsafe def elabTensorDot: CommandElab
+  | `(#tensor_dot $term) => dotElab term
+  | _ => throwUnsupportedSyntax
+
+end dotElab
+
+end TensorTree

HepLean/Tensors/Tree/Elab.lean

@@ -5,6 +5,7 @@ Authors: Joseph Tooby-Smith
 -/
 import HepLean.Tensors.Tree.Basic
 import Lean.Elab.Term
+import HepLean.Tensors.Tree.Dot
 /-!
 
 ## Elaboration of tensor trees
@@ -22,6 +23,8 @@ open Lean.Elab.Term
 open Lean Meta Elab Tactic
 open IndexNotation
 
+namespace TensorTree
+
 /-!
 
 ## Indexies
@@ -100,6 +103,8 @@ syntax tensorExpr "+" tensorExpr : tensorExpr
 /-- Allowing brackets to be used in a tensor expression. -/
 syntax "(" tensorExpr ")" : tensorExpr
 
+syntax term "•" tensorExpr : tensorExpr
+
 namespace TensorNode
 
 /-!
@@ -283,18 +288,24 @@ elab_rules (kind:=tensorExprSyntax) : term
     return tensorTree
 
 variable {S : TensorStruct} {c4 : Fin 4 → S.C} (T4 : S.F.obj (OverColor.mk c4))
-  {c5 : Fin 5 → S.C} (T5 : S.F.obj (OverColor.mk c5))
+  {c5 : Fin 5 → S.C} (T5 : S.F.obj (OverColor.mk c5)) (a : S.k)
+
+variable (𝓣 : TensorTree S c4)
 
 /-!
 
 # Checks
 
 -/
+
 /-
-example : {T4 | i j k m}ᵀ = TensorTree.tensorNode T4 := by rfl
+#tensor_dot {T4 | i j l d ⊗ T5 | i j k m m}ᵀ.dot
 
 #check {T4 | i j l d ⊗ T5 | i j k a b}ᵀ
 
 #check {(T4 | i j l a ⊗ T5 | i j k c d) ⊗ T5 | i1 i2 i3 e f}ᵀ
 -/
 end ProdNode
+
+
+end TensorTree