diff --git a/HepLean.lean b/HepLean.lean index bc26329..a5eeb98 100644 --- a/HepLean.lean +++ b/HepLean.lean @@ -124,4 +124,5 @@ import HepLean.Tensors.IndexNotation.TensorIndex import HepLean.Tensors.MulActionTensor import HepLean.Tensors.RisingLowering import HepLean.Tensors.Tree.Basic +import HepLean.Tensors.Tree.Dot import HepLean.Tensors.Tree.Elab diff --git a/HepLean/Tensors/Tree/Basic.lean b/HepLean/Tensors/Tree/Basic.lean index 82acd69..3d1572d 100644 --- a/HepLean/Tensors/Tree/Basic.lean +++ b/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 diff --git a/HepLean/Tensors/Tree/Dot.lean b/HepLean/Tensors/Tree/Dot.lean new file mode 100644 index 0000000..4b7087c --- /dev/null +++ b/HepLean/Tensors/Tree/Dot.lean @@ -0,0 +1,95 @@ +/- +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 for showing the dot file associated with a tensor tree. + Use as `#tensor_dot {T4 | i j l d ⊗ T5 | i j k a b}ᵀ.dot`. -/ +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 diff --git a/HepLean/Tensors/Tree/Elab.lean b/HepLean/Tensors/Tree/Elab.lean index 7ad4771..eb5fc70 100644 --- a/HepLean/Tensors/Tree/Elab.lean +++ b/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,9 @@ syntax tensorExpr "+" tensorExpr : tensorExpr /-- Allowing brackets to be used in a tensor expression. -/ syntax "(" tensorExpr ")" : tensorExpr +/-- Scalar multiplication for tensors. -/ +syntax term "•" tensorExpr : tensorExpr + namespace TensorNode /-! @@ -283,18 +289,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