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