Refactor TensorNetwork constructors

Author: mtfishman

Project.toml

@@ -32,7 +32,7 @@ ITensorNetworksNextTensorOperationsExt = "TensorOperations"
 AbstractTrees = "0.4.5"
 Adapt = "4.3"
 BackendSelection = "0.1.6"
-Combinatorics = "1.0.3"
+Combinatorics = "1"
 DataGraphs = "0.2.7"
 DiagonalArrays = "0.3.23"
 Dictionaries = "0.4.5"

src/LazyNamedDimsArrays/evaluation_order.jl

@@ -91,16 +91,15 @@ function optimize_evaluation_order(
     return optimize_evaluation_order(alg, a)
 end
 
-struct Eager end
-
-default_optimize_evaluation_order_alg(a) = Eager()
+using BackendSelection: @Algorithm_str, Algorithm
+default_optimize_evaluation_order_alg(a) = Algorithm"eager"()
 
 function optimize_contraction_order_flattened(alg, a)
-    return error("Alg $alg not supported.")
+    return error("`alg = $alg` not supported.")
 end
 
 using Combinatorics: combinations
-function optimize_contraction_order_flattened(alg::Eager, a)
+function optimize_contraction_order_flattened(alg::Algorithm"eager", a)
     @assert ismul(a)
     arity(a) in (1, 2) && return a
     a1, a2 = argmin(combinations(arguments(a), 2)) do (a1, a2)

src/LazyNamedDimsArrays/lazynameddimsarray.jl

@@ -6,13 +6,18 @@ using WrappedUnions: @wrapped
     } <: AbstractNamedDimsArray{T, Any}
     union::Union{A, Mul{LazyNamedDimsArray{T, A}}}
 end
+
+parenttype(::Type{LazyNamedDimsArray{<:Any, A}}) where {A} = A
+parenttype(::Type{LazyNamedDimsArray{T}}) where {T} = AbstractNamedDimsArray{T}
+parenttype(::Type{LazyNamedDimsArray}) = AbstractNamedDimsArray
+
 function LazyNamedDimsArray(a::AbstractNamedDimsArray)
     # Use `eltype(typeof(a))` for arrays that have different
     # runtime and compile time eltypes, like `ITensor`.
     return LazyNamedDimsArray{eltype(typeof(a)), typeof(a)}(a)
 end
-function LazyNamedDimsArray(a::Mul{LazyNamedDimsArray{T, A}}) where {T, A}
-    return LazyNamedDimsArray{T, A}(a)
+function LazyNamedDimsArray(a::Mul{L}) where {L <: LazyNamedDimsArray}
+    return LazyNamedDimsArray{eltype(L), parenttype(L)}(a)
 end
 lazy(a::LazyNamedDimsArray) = a
 lazy(a::AbstractNamedDimsArray) = LazyNamedDimsArray(a)

src/LazyNamedDimsArrays/symbolicarray.jl

@@ -9,12 +9,6 @@ end
 function SymbolicArray(name, ax::Tuple{Vararg{AbstractUnitRange{<:Integer}}})
     return SymbolicArray{Any}(name, ax)
 end
-function SymbolicArray{T}(name, ax::AbstractUnitRange...) where {T}
-    return SymbolicArray{T}(name, ax)
-end
-function SymbolicArray(name, ax::AbstractUnitRange...)
-    return SymbolicArray{Any}(name, ax)
-end
 symname(a::SymbolicArray) = getfield(a, :name)
 Base.axes(a::SymbolicArray) = getfield(a, :axes)
 Base.size(a::SymbolicArray) = length.(axes(a))

src/LazyNamedDimsArrays/symbolicnameddimsarray.jl

@@ -2,9 +2,10 @@ using NamedDimsArrays: NamedDimsArray, dename, inds
 
 const SymbolicNamedDimsArray{T, N, Parent <: SymbolicArray{T, N}, DimNames} =
     NamedDimsArray{T, N, Parent, DimNames}
-function symnameddims(name)
-    return lazy(NamedDimsArray(SymbolicArray(name), ()))
+function symnameddims(name, dims)
+    return lazy(NamedDimsArray(SymbolicArray(name, dename.(dims)), dims))
 end
+symnameddims(name) = symnameddims(name, ())
 using AbstractTrees: AbstractTrees
 function AbstractTrees.printnode(io::IO, a::SymbolicNamedDimsArray)
     print(io, symname(dename(a)))

src/abstracttensornetwork.jl

@@ -40,6 +40,7 @@ Base.copy(tn::AbstractTensorNetwork) = error("Not implemented")
 
 # Iteration
 Base.iterate(tn::AbstractTensorNetwork, args...) = iterate(vertex_data(tn), args...)
+Base.keys(tn::AbstractTensorNetwork) = vertices(tn)
 
 # TODO: This contrasts with the `DataGraphs.AbstractDataGraph` definition,
 # where it is defined as the `vertextype`. Does that cause problems or should it be changed?

src/contract_network.jl

@@ -1,46 +1,50 @@
 using BackendSelection: @Algorithm_str, Algorithm
 using Base.Broadcast: materialize
-using ITensorNetworksNext.LazyNamedDimsArrays: lazy, substitute, symnameddims
+using ITensorNetworksNext.LazyNamedDimsArrays: lazy, optimize_evaluation_order, substitute,
+    symnameddims
 
-# This is based on `MatrixAlgebraKit.select_algorithm`.
+# This is related to `MatrixAlgebraKit.select_algorithm`.
 # TODO: Define this in BackendSelection.jl.
-function select_algorithm(alg; kwargs...)
-    if alg isa Algorithm
-        @assert isempty(kwargs) "Cannot pass keyword arguments when `alg` is an `Algorithm`."
-        return alg
-    else
-        return Algorithm(alg; kwargs...)
-    end
+backend_value(::Algorithm{alg}) where {alg} = alg
+using BackendSelection: parameters
+function merge_parameters(alg::Algorithm; kwargs...)
+    return Algorithm(backend_value(alg); merge(parameters(alg), kwargs)...)
 end
+to_algorithm(alg::Algorithm; kwargs...) = merge_parameters(alg; kwargs...)
+to_algorithm(alg; kwargs...) = Algorithm(alg; kwargs...)
 
 # `contract_network`
 contract_network(alg::Algorithm, tn) = error("Not implemented.")
 function default_kwargs(::typeof(contract_network), tn)
     return (; alg = Algorithm"exact"(; order_alg = Algorithm"eager"()))
 end
 function contract_network(tn; alg = default_kwargs(contract_network, tn).alg, kwargs...)
-    return contract_network(select_algorithm(alg; kwargs...), tn)
+    return contract_network(to_algorithm(alg; kwargs...), tn)
 end
 
 # `contract_network(::Algorithm"exact", ...)`
 function contract_network(alg::Algorithm"exact", tn)
     order = @something begin
         get(alg, :order, nothing)
-        contraction_order(tn; alg = get(alg, :order_alg, default_kwargs(contraction_order, tn).alg))
+        contraction_order(
+            tn; alg = get(alg, :order_alg, default_kwargs(contraction_order, tn).alg)
+        )
     end
-    syms_to_ts = Dict(symnameddims(i) => lazy(tn[i]) for i in eachindex(tn))
+    syms_to_ts = Dict(symnameddims(i, Tuple(inds(tn[i]))) => lazy(tn[i]) for i in eachindex(tn))
     tn_expression = substitute(order, syms_to_ts)
     return materialize(tn_expression)
 end
 
 # `contraction_order`
-contraction_order(alg::Algorithm, tn) = error("Not implemented.")
+function contraction_order end
 default_kwargs(::typeof(contraction_order), tn) = (; alg = Algorithm"eager"())
 function contraction_order(tn; alg = default_kwargs(contraction_order, tn).alg, kwargs...)
-    return contraction_order(select_algorithm(alg; kwargs...), tn)
+    return contraction_order(to_algorithm(alg; kwargs...), tn)
 end
-
-# `contraction_order(::Algorithm"eager", ...)`
-function contraction_order(alg::Algorithm"eager", tn)
-    return error("Eager not implemented.")
+function contraction_order(alg::Algorithm"left_associative", tn)
+    return prod(i -> symnameddims(i, Tuple(inds(tn[i]))), eachindex(tn))
+end
+function contraction_order(alg::Algorithm, tn)
+    s = contraction_order(tn; alg = Algorithm"left_associative"())
+    return optimize_evaluation_order(s; alg)
 end

src/tensornetwork.jl

@@ -1,9 +1,10 @@
+using Combinatorics: combinations
 using DataGraphs: DataGraphs, AbstractDataGraph, DataGraph
 using Dictionaries: AbstractDictionary, Indices, dictionary
 using Graphs: AbstractSimpleGraph
 using NamedDimsArrays: AbstractNamedDimsArray, dimnames
 using NamedGraphs: NamedGraphs, NamedEdge, NamedGraph, vertextype
-using NamedGraphs.GraphsExtensions: arranged_edges, vertextype
+using NamedGraphs.GraphsExtensions: add_edges!, arrange_edge, arranged_edges, vertextype
 
 function _TensorNetwork end
 
@@ -18,45 +19,65 @@ struct TensorNetwork{V, VD, UG <: AbstractGraph{V}, Tensors <: AbstractDictionar
         return new{V, VD, UG, Tensors}(underlying_graph, tensors)
     end
 end
+# This assumes the tensor connectivity matches the graph structure.
+function _TensorNetwork(graph::AbstractGraph, tensors)
+    return _TensorNetwork(graph, Dictionary(keys(tensors), values(tensors)))
+end
 
 DataGraphs.underlying_graph(tn::TensorNetwork) = getfield(tn, :underlying_graph)
 DataGraphs.vertex_data(tn::TensorNetwork) = getfield(tn, :tensors)
 function DataGraphs.underlying_graph_type(type::Type{<:TensorNetwork})
     return fieldtype(type, :underlying_graph)
 end
 
-# Determine the graph structure from the tensors.
-function TensorNetwork(t::AbstractDictionary)
-    g = NamedGraph(eachindex(t))
-    for v1 in vertices(g)
-        for v2 in vertices(g)
-            if v1 ≠ v2
-                if !isdisjoint(dimnames(t[v1]), dimnames(t[v2]))
-                    add_edge!(g, v1 => v2)
-                end
+# For a collection of tensors, return the edges implied by shared indices
+# as a list of `edgetype` edges of keys/vertices.
+function tensornetwork_edges(edgetype::Type, tensors)
+    # We need to collect the keys since in the case of `tensors::AbstractDictionary`,
+    # `keys(tensors)::AbstractIndices`, which is indexed by `keys(tensors)` rather
+    # than `1:length(keys(tensors))`, which is assumed by `combinations`.
+    verts = collect(keys(tensors))
+    return filter(
+        !isnothing, map(combinations(verts, 2)) do (v1, v2)
+            if !isdisjoint(inds(tensors[v1]), inds(tensors[v2]))
+                return arrange_edge(edgetype(v1, v2))
             end
+            return nothing
         end
-    end
-    return _TensorNetwork(g, t)
+    )
 end
-function TensorNetwork(tensors::AbstractDict)
-    return TensorNetwork(Dictionary(tensors))
+tensornetwork_edges(tensors) = tensornetwork_edges(NamedEdge, tensors)
+
+function TensorNetwork(f::Base.Callable, graph::AbstractGraph)
+    tensors = Dictionary(vertices(graph), f.(vertices(graph)))
+    return TensorNetwork(graph, tensors)
+end
+function TensorNetwork(graph::AbstractGraph, tensors)
+    tn = _TensorNetwork(graph, tensors)
+    fix_links!(tn)
+    return tn
 end
 
-function TensorNetwork(graph::AbstractGraph, tensors::AbstractDictionary)
-    tn = TensorNetwork(tensors)
-    arranged_edges(tn) ⊆ arranged_edges(graph) ||
+# Insert trivial links for missing edges, and also check
+# the vertices and edges are consistent between the graph and tensors.
+function fix_links!(tn::AbstractTensorNetwork)
+    graph = underlying_graph(tn)
+    tensors = vertex_data(tn)
+    @assert issetequal(vertices(graph), keys(tensors)) "Graph vertices and tensor keys must match."
+    tn_edges = tensornetwork_edges(edgetype(graph), tensors)
+    tn_edges ⊆ arranged_edges(graph) ||
         error("The edges in the tensors do not match the graph structure.")
-    for e in setdiff(arranged_edges(graph), arranged_edges(tn))
+    for e in setdiff(arranged_edges(graph), tn_edges)
         insert_trivial_link!(tn, e)
     end
     return tn
 end
-function TensorNetwork(graph::AbstractGraph, tensors::AbstractDict)
-    return TensorNetwork(graph, Dictionary(tensors))
-end
-function TensorNetwork(f, graph::AbstractGraph)
-    return TensorNetwork(graph, Dict(v => f(v) for v in vertices(graph)))
+
+# Determine the graph structure from the tensors.
+function TensorNetwork(tensors)
+    graph = NamedGraph(keys(tensors))
+    add_edges!(graph, tensornetwork_edges(tensors))
+    return _TensorNetwork(graph, tensors)
 end
 
 function Base.copy(tn::TensorNetwork)
@@ -65,10 +86,9 @@ end
 TensorNetwork(tn::TensorNetwork) = copy(tn)
 TensorNetwork{V}(tn::TensorNetwork{V}) where {V} = copy(tn)
 function TensorNetwork{V}(tn::TensorNetwork) where {V}
-    g′ = convert_vertextype(V, underlying_graph(tn))
-    d = vertex_data(tn)
-    d′ = dictionary(V(k) => d[k] for k in eachindex(d))
-    return TensorNetwork(g′, d′)
+    g = convert_vertextype(V, underlying_graph(tn))
+    d = dictionary(V(k) => tn[k] for k in keys(d))
+    return TensorNetwork(g, d)
 end
 
 NamedGraphs.convert_vertextype(::Type{V}, tn::TensorNetwork{V}) where {V} = tn

test/test_contract_network.jl

@@ -2,8 +2,7 @@ using Graphs: edges
 using NamedGraphs.GraphsExtensions: arranged_edges, incident_edges
 using NamedGraphs.NamedGraphGenerators: named_grid
 using ITensorBase: Index, ITensor
-using ITensorNetworksNext:
-    TensorNetwork, linkinds, siteinds, contract_network
+using ITensorNetworksNext: TensorNetwork, linkinds, siteinds, contract_network
 using TensorOperations: TensorOperations
 using Test: @test, @testset
 
@@ -15,8 +14,8 @@ using Test: @test, @testset
         C = ITensor([5.0, 1.0], j)
         D = ITensor([-2.0, 3.0, 4.0, 5.0, 1.0], k)
 
-        ABCD_1 = contract_network([A, B, C, D]; alg = "exact", sequence_alg = "leftassociative")
-        ABCD_2 = contract_network([A, B, C, D]; alg = "exact", sequence_alg = "optimal")
+        ABCD_1 = contract_network([A, B, C, D]; alg = "exact", order_alg = "leftassociative")
+        ABCD_2 = contract_network([A, B, C, D]; alg = "exact", order_alg = "optimal")
 
         @test ABCD_1 == ABCD_2
     end

test/test_lazynameddimsarrays.jl

@@ -95,7 +95,7 @@ using WrappedUnions: unwrap
         @test unwrap(a1) isa NamedDimsArray
         @test dename(a1) isa SymbolicArray
         @test dename(unwrap(a1)) isa SymbolicArray
-        @test dename(unwrap(a1)) == SymbolicArray(:a1)
+        @test dename(unwrap(a1)) == SymbolicArray(:a1, ())
         @test inds(a1) == ()
         @test dimnames(a1) == ()