@@ -2,57 +2,57 @@ module ITensorNetworksEinExprsExt
22
33using ITensors: Index, ITensor, @Algorithm_str , inds, noncommoninds
44using ITensorNetworks:
5- ITensorNetworks,
6- ITensorList,
7- ITensorNetwork,
8- vertextype,
9- vertex_data,
10- contraction_sequence
5+ ITensorNetworks,
6+ ITensorList,
7+ ITensorNetwork,
8+ vertextype,
9+ vertex_data,
10+ contraction_sequence
1111using EinExprs: EinExprs, EinExpr, einexpr, SizedEinExpr
1212
1313function to_einexpr (ts:: ITensorList )
14- IndexType = Any
14+ IndexType = Any
1515
16- tensor_exprs = EinExpr{IndexType}[]
17- inds_dims = Dict {IndexType,Int} ()
16+ tensor_exprs = EinExpr{IndexType}[]
17+ inds_dims = Dict {IndexType, Int} ()
1818
19- for tensor_v in ts
20- inds_v = collect (inds (tensor_v))
21- push! (tensor_exprs, EinExpr {IndexType} (; head= inds_v))
22- merge! (inds_dims, Dict (inds_v .=> size (tensor_v)))
23- end
19+ for tensor_v in ts
20+ inds_v = collect (inds (tensor_v))
21+ push! (tensor_exprs, EinExpr {IndexType} (; head = inds_v))
22+ merge! (inds_dims, Dict (inds_v .=> size (tensor_v)))
23+ end
2424
25- externalinds_tn = reduce (noncommoninds, ts)
26- return SizedEinExpr (sum (tensor_exprs; skip= externalinds_tn), inds_dims)
25+ externalinds_tn = reduce (noncommoninds, ts)
26+ return SizedEinExpr (sum (tensor_exprs; skip = externalinds_tn), inds_dims)
2727end
2828
2929function tensor_inds_to_vertex (ts:: ITensorList )
30- IndexType = Any
31- VertexType = Int
30+ IndexType = Any
31+ VertexType = Int
3232
33- mapping = Dict {Set{IndexType},VertexType} ()
33+ mapping = Dict {Set{IndexType}, VertexType} ()
3434
35- for (v, tensor_v) in enumerate (ts)
36- inds_v = collect (inds (tensor_v))
37- mapping[Set (inds_v)] = v
38- end
35+ for (v, tensor_v) in enumerate (ts)
36+ inds_v = collect (inds (tensor_v))
37+ mapping[Set (inds_v)] = v
38+ end
3939
40- return mapping
40+ return mapping
4141end
4242
4343function ITensorNetworks. contraction_sequence (
44- :: Algorithm"einexpr" , tn:: ITensorList ; optimizer= EinExprs. Exhaustive ()
45- )
46- expr = to_einexpr (tn)
47- path = einexpr (optimizer, expr)
48- return to_contraction_sequence (path, tensor_inds_to_vertex (tn))
44+ :: Algorithm"einexpr" , tn:: ITensorList ; optimizer = EinExprs. Exhaustive ()
45+ )
46+ expr = to_einexpr (tn)
47+ path = einexpr (optimizer, expr)
48+ return to_contraction_sequence (path, tensor_inds_to_vertex (tn))
4949end
5050
5151function to_contraction_sequence (expr, tensor_inds_to_vertex)
52- EinExprs. nargs (expr) == 0 && return tensor_inds_to_vertex[Set (expr. head)]
53- return map (
54- expr -> to_contraction_sequence (expr, tensor_inds_to_vertex), EinExprs. args (expr)
55- )
52+ EinExprs. nargs (expr) == 0 && return tensor_inds_to_vertex[Set (expr. head)]
53+ return map (
54+ expr -> to_contraction_sequence (expr, tensor_inds_to_vertex), EinExprs. args (expr)
55+ )
5656end
5757
5858end
0 commit comments