92 lines
4.2 KiB
Text
92 lines
4.2 KiB
Text
/-
|
||
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
|