Refactor and generalize tensor network constructors (#155)

Author: mtfishman

Project.toml

@@ -1,7 +1,7 @@
 name = "ITensorNetworks"
 uuid = "2919e153-833c-4bdc-8836-1ea460a35fc7"
 authors = ["Matthew Fishman <mfishman@flatironinstitute.org> and contributors"]
-version = "0.6"
+version = "0.7"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -56,14 +56,14 @@ GraphsFlows = "0.1.1"
 ITensors = "0.3.58"
 IsApprox = "0.1"
 IterTools = "1.4.0"
-KrylovKit = "0.6.0"
+KrylovKit = "0.6, 0.7"
 NamedGraphs = "0.1.23"
 Observers = "0.2"
 PackageExtensionCompat = "1"
 Requires = "1.3"
 SerializedElementArrays = "0.1"
 SimpleTraits = "0.9"
-SparseArrayKit = "0.2.1"
+SparseArrayKit = "0.2.1, 0.3"
 SplitApplyCombine = "1.2"
 StaticArrays = "1.5.12"
 StructWalk = "0.2"

README.md

@@ -105,13 +105,13 @@ and 4 edge(s):
 
 with vertex data:
 4-element Dictionaries.Dictionary{Tuple{Int64, Int64}, Any}
- (1, 1) │ ((dim=2|id=712|"1×1,2×1"), (dim=2|id=598|"1×1,1×2"))
- (2, 1) │ ((dim=2|id=712|"1×1,2×1"), (dim=2|id=457|"2×1,2×2"))
- (1, 2) │ ((dim=2|id=598|"1×1,1×2"), (dim=2|id=683|"1×2,2×2"))
- (2, 2) │ ((dim=2|id=457|"2×1,2×2"), (dim=2|id=683|"1×2,2×2"))
+ (1, 1) │ ((dim=2|id=74|"1×1,2×1"), (dim=2|id=723|"1×1,1×2"))
+ (2, 1) │ ((dim=2|id=74|"1×1,2×1"), (dim=2|id=823|"2×1,2×2"))
+ (1, 2) │ ((dim=2|id=723|"1×1,1×2"), (dim=2|id=712|"1×2,2×2"))
+ (2, 2) │ ((dim=2|id=823|"2×1,2×2"), (dim=2|id=712|"1×2,2×2"))
 
 julia> tn[1, 1]
-ITensor ord=2 (dim=2|id=712|"1×1,2×1") (dim=2|id=598|"1×1,1×2")
+ITensor ord=2 (dim=2|id=74|"1×1,2×1") (dim=2|id=723|"1×1,1×2")
 NDTensors.EmptyStorage{NDTensors.EmptyNumber, NDTensors.Dense{NDTensors.EmptyNumber, Vector{NDTensors.EmptyNumber}}}
 
 julia> neighbors(tn, (1, 1))
@@ -135,8 +135,8 @@ and 1 edge(s):
 
 with vertex data:
 2-element Dictionaries.Dictionary{Tuple{Int64, Int64}, Any}
- (1, 1) │ ((dim=2|id=712|"1×1,2×1"), (dim=2|id=598|"1×1,1×2"))
- (1, 2) │ ((dim=2|id=598|"1×1,1×2"), (dim=2|id=683|"1×2,2×2"))
+ (1, 1) │ ((dim=2|id=74|"1×1,2×1"), (dim=2|id=723|"1×1,1×2"))
+ (1, 2) │ ((dim=2|id=723|"1×1,1×2"), (dim=2|id=712|"1×2,2×2"))
 
 julia> tn_2 = subgraph(v -> v[1] == 2, tn)
 ITensorNetworks.ITensorNetwork{Tuple{Int64, Int64}} with 2 vertices:
@@ -149,8 +149,8 @@ and 1 edge(s):
 
 with vertex data:
 2-element Dictionaries.Dictionary{Tuple{Int64, Int64}, Any}
- (2, 1) │ ((dim=2|id=712|"1×1,2×1"), (dim=2|id=457|"2×1,2×2"))
- (2, 2) │ ((dim=2|id=457|"2×1,2×2"), (dim=2|id=683|"1×2,2×2"))
+ (2, 1) │ ((dim=2|id=74|"1×1,2×1"), (dim=2|id=823|"2×1,2×2"))
+ (2, 2) │ ((dim=2|id=823|"2×1,2×2"), (dim=2|id=712|"1×2,2×2"))
 ```
 
 
@@ -176,9 +176,9 @@ and 2 edge(s):
 
 with vertex data:
 3-element Dictionaries.Dictionary{Int64, Vector{ITensors.Index}}
- 1 │ ITensors.Index[(dim=2|id=830|"S=1/2,Site,n=1")]
- 2 │ ITensors.Index[(dim=2|id=369|"S=1/2,Site,n=2")]
- 3 │ ITensors.Index[(dim=2|id=558|"S=1/2,Site,n=3")]
+ 1 │ ITensors.Index[(dim=2|id=683|"S=1/2,Site,n=1")]
+ 2 │ ITensors.Index[(dim=2|id=123|"S=1/2,Site,n=2")]
+ 3 │ ITensors.Index[(dim=2|id=656|"S=1/2,Site,n=3")]
 
 and edge data:
 0-element Dictionaries.Dictionary{NamedGraphs.NamedEdge{Int64}, Vector{ITensors.Index}}
@@ -196,9 +196,9 @@ and 2 edge(s):
 
 with vertex data:
 3-element Dictionaries.Dictionary{Int64, Any}
- 1 │ ((dim=2|id=830|"S=1/2,Site,n=1"), (dim=2|id=186|"1,2"))
- 2 │ ((dim=2|id=369|"S=1/2,Site,n=2"), (dim=2|id=186|"1,2"), (dim=2|id=430|"2,3…
- 3 │ ((dim=2|id=558|"S=1/2,Site,n=3"), (dim=2|id=430|"2,3"))
+ 1 │ ((dim=2|id=683|"S=1/2,Site,n=1"), (dim=2|id=382|"1,2"))
+ 2 │ ((dim=2|id=123|"S=1/2,Site,n=2"), (dim=2|id=382|"1,2"), (dim=2|id=190|"2,3…
+ 3 │ ((dim=2|id=656|"S=1/2,Site,n=3"), (dim=2|id=190|"2,3"))
 
 julia> tn2 = ITensorNetwork(s; link_space=2)
 ITensorNetworks.ITensorNetwork{Int64} with 3 vertices:
@@ -213,9 +213,9 @@ and 2 edge(s):
 
 with vertex data:
 3-element Dictionaries.Dictionary{Int64, Any}
- 1 │ ((dim=2|id=830|"S=1/2,Site,n=1"), (dim=2|id=994|"1,2"))
- 2 │ ((dim=2|id=369|"S=1/2,Site,n=2"), (dim=2|id=994|"1,2"), (dim=2|id=978|"2,3…
- 3 │ ((dim=2|id=558|"S=1/2,Site,n=3"), (dim=2|id=978|"2,3"))
+ 1 │ ((dim=2|id=683|"S=1/2,Site,n=1"), (dim=2|id=934|"1,2"))
+ 2 │ ((dim=2|id=123|"S=1/2,Site,n=2"), (dim=2|id=934|"1,2"), (dim=2|id=614|"2,3…
+ 3 │ ((dim=2|id=656|"S=1/2,Site,n=3"), (dim=2|id=614|"2,3"))
 
 julia> @visualize tn1;
     ⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀ 

src/ModelNetworks/ModelNetworks.jl

@@ -1,6 +1,6 @@
 module ModelNetworks
 using Graphs: degree, dst, edges, src
-using ..ITensorNetworks: IndsNetwork, delta_network, insert_missing_internal_inds, itensor
+using ..ITensorNetworks: IndsNetwork, delta_network, insert_linkinds, itensor
 using ITensors: commoninds, diagITensor, inds, noprime
 using LinearAlgebra: Diagonal, eigen
 using NamedGraphs: NamedGraph
@@ -17,7 +17,7 @@ OPTIONAL ARGUMENT:
 function ising_network(
   eltype::Type, s::IndsNetwork, beta::Number; h::Number=0.0, szverts=nothing
 )
-  s = insert_missing_internal_inds(s, edges(s); internal_inds_space=2)
+  s = insert_linkinds(s; link_space=2)
   tn = delta_network(eltype, s)
   if (szverts != nothing)
     for v in szverts

src/abstractindsnetwork.jl

@@ -23,13 +23,57 @@ end
 # TODO: Define a generic fallback for `AbstractDataGraph`?
 DataGraphs.edge_data_type(::Type{<:AbstractIndsNetwork{V,I}}) where {V,I} = Vector{I}
 
+## TODO: Bring these back.
+## function indsnetwork_getindex(is::AbstractIndsNetwork, index)
+##   return get(data_graph(is), index, indtype(is)[])
+## end
+## 
+## function Base.getindex(is::AbstractIndsNetwork, index)
+##   return indsnetwork_getindex(is, index)
+## end
+## 
+## function Base.getindex(is::AbstractIndsNetwork, index::Pair)
+##   return indsnetwork_getindex(is, index)
+## end
+## 
+## function Base.getindex(is::AbstractIndsNetwork, index::AbstractEdge)
+##   return indsnetwork_getindex(is, index)
+## end
+## 
+## function indsnetwork_setindex!(is::AbstractIndsNetwork, value, index)
+##   data_graph(is)[index] = value
+##   return is
+## end
+## 
+## function Base.setindex!(is::AbstractIndsNetwork, value, index)
+##   indsnetwork_setindex!(is, value, index)
+##   return is
+## end
+## 
+## function Base.setindex!(is::AbstractIndsNetwork, value, index::Pair)
+##   indsnetwork_setindex!(is, value, index)
+##   return is
+## end
+## 
+## function Base.setindex!(is::AbstractIndsNetwork, value, index::AbstractEdge)
+##   indsnetwork_setindex!(is, value, index)
+##   return is
+## end
+## 
+## function Base.setindex!(is::AbstractIndsNetwork, value::Index, index)
+##   indsnetwork_setindex!(is, value, index)
+##   return is
+## end
+
 # 
 # Index access
 # 
 
 function ITensors.uniqueinds(is::AbstractIndsNetwork, edge::AbstractEdge)
+  # TODO: Replace with `is[v]` once `getindex(::IndsNetwork, ...)` is smarter.
   inds = IndexSet(get(is, src(edge), Index[]))
   for ei in setdiff(incident_edges(is, src(edge)), [edge])
+    # TODO: Replace with `is[v]` once `getindex(::IndsNetwork, ...)` is smarter.
     inds = unioninds(inds, get(is, ei, Index[]))
   end
   return inds
@@ -39,8 +83,8 @@ function ITensors.uniqueinds(is::AbstractIndsNetwork, edge::Pair)
   return uniqueinds(is, edgetype(is)(edge))
 end
 
-function Base.union(tn1::AbstractIndsNetwork, tn2::AbstractIndsNetwork; kwargs...)
-  return IndsNetwork(union(data_graph(tn1), data_graph(tn2); kwargs...))
+function Base.union(is1::AbstractIndsNetwork, is2::AbstractIndsNetwork; kwargs...)
+  return IndsNetwork(union(data_graph(is1), data_graph(is2); kwargs...))
 end
 
 function NamedGraphs.rename_vertices(f::Function, tn::AbstractIndsNetwork)
@@ -51,31 +95,49 @@ end
 # Convenience functions
 # 
 
+function promote_indtypeof(is::AbstractIndsNetwork)
+  sitetype = mapreduce(promote_indtype, vertices(is); init=Index{Int}) do v
+    # TODO: Replace with `is[v]` once `getindex(::IndsNetwork, ...)` is smarter.
+    return mapreduce(typeof, promote_indtype, get(is, v, Index[]); init=Index{Int})
+  end
+  linktype = mapreduce(promote_indtype, edges(is); init=Index{Int}) do e
+    # TODO: Replace with `is[e]` once `getindex(::IndsNetwork, ...)` is smarter.
+    return mapreduce(typeof, promote_indtype, get(is, e, Index[]); init=Index{Int})
+  end
+  return promote_indtype(sitetype, linktype)
+end
+
 function union_all_inds(is_in::AbstractIndsNetwork...)
   @assert all(map(ug -> ug == underlying_graph(is_in[1]), underlying_graph.(is_in)))
   is_out = IndsNetwork(underlying_graph(is_in[1]))
   for v in vertices(is_out)
+    # TODO: Remove this check.
     if any(isassigned(is, v) for is in is_in)
+      # TODO: Change `get` to `getindex`.
       is_out[v] = unioninds([get(is, v, Index[]) for is in is_in]...)
     end
   end
   for e in edges(is_out)
+    # TODO: Remove this check.
     if any(isassigned(is, e) for is in is_in)
+      # TODO: Change `get` to `getindex`.
       is_out[e] = unioninds([get(is, e, Index[]) for is in is_in]...)
     end
   end
   return is_out
 end
 
-function insert_missing_internal_inds(
+function insert_linkinds(
   indsnetwork::AbstractIndsNetwork,
   edges=edges(indsnetwork);
-  internal_inds_space=trivial_space(indsnetwork),
+  link_space=trivial_space(indsnetwork),
 )
   indsnetwork = copy(indsnetwork)
   for e in edges
+    # TODO: Change to check if it is empty.
     if !isassigned(indsnetwork, e)
-      iₑ = Index(internal_inds_space, edge_tag(e))
+      iₑ = Index(link_space, edge_tag(e))
+      # TODO: Allow setting with just `Index`.
       indsnetwork[e] = [iₑ]
     end
   end

src/abstractitensornetwork.jl

@@ -116,11 +116,12 @@ end
 # TODO: broadcasting
 
 function Base.union(tn1::AbstractITensorNetwork, tn2::AbstractITensorNetwork; kwargs...)
-  tn = ITensorNetwork(union(data_graph(tn1), data_graph(tn2)); kwargs...)
+  # TODO: Use a different constructor call here?
+  tn = _ITensorNetwork(union(data_graph(tn1), data_graph(tn2)); kwargs...)
   # Add any new edges that are introduced during the union
   for v1 in vertices(tn1)
     for v2 in vertices(tn2)
-      if hascommoninds(tn[v1], tn[v2])
+      if hascommoninds(tn, v1 => v2)
         add_edge!(tn, v1 => v2)
       end
     end
@@ -129,7 +130,8 @@ function Base.union(tn1::AbstractITensorNetwork, tn2::AbstractITensorNetwork; kw
 end
 
 function NamedGraphs.rename_vertices(f::Function, tn::AbstractITensorNetwork)
-  return ITensorNetwork(rename_vertices(f, data_graph(tn)))
+  # TODO: Use a different constructor call here?
+  return _ITensorNetwork(rename_vertices(f, data_graph(tn)))
 end
 
 #
@@ -172,6 +174,8 @@ function Base.Vector{ITensor}(tn::AbstractITensorNetwork)
 end
 
 # Convenience wrapper
+# TODO: Delete this and just use `Vector{ITensor}`, or maybe
+# it should output a dictionary or be called `eachtensor`?
 itensors(tn::AbstractITensorNetwork) = Vector{ITensor}(tn)
 
 #
@@ -182,10 +186,13 @@ function LinearAlgebra.promote_leaf_eltypes(tn::AbstractITensorNetwork)
   return LinearAlgebra.promote_leaf_eltypes(itensors(tn))
 end
 
-function trivial_space(tn::AbstractITensorNetwork)
-  return trivial_space(tn[first(vertices(tn))])
+function promote_indtypeof(tn::AbstractITensorNetwork)
+  return mapreduce(promote_indtype, vertices(tn)) do v
+    return indtype(tn[v])
+  end
 end
 
+# TODO: Delete in favor of `scalartype`.
 function ITensors.promote_itensor_eltype(tn::AbstractITensorNetwork)
   return LinearAlgebra.promote_leaf_eltypes(tn)
 end
@@ -464,7 +471,6 @@ function NDTensors.contract(
   neighbors_src = setdiff(neighbors(tn, src(edge)), [dst(edge)])
   neighbors_dst = setdiff(neighbors(tn, dst(edge)), [src(edge)])
   new_itensor = tn[src(edge)] * tn[dst(edge)]
-
   # The following is equivalent to:
   #
   # tn[dst(edge)] = new_itensor
@@ -480,6 +486,7 @@ function NDTensors.contract(
   for n_dst in neighbors_dst
     add_edge!(tn, merged_vertex => n_dst)
   end
+
   setindex_preserve_graph!(tn, new_itensor, merged_vertex)
 
   return tn
@@ -736,7 +743,8 @@ function norm_network(tn::AbstractITensorNetwork)
     setindex_preserve_graph!(tndag, dag(tndag[v]), v)
   end
   tnket = rename_vertices(v -> (v, 2), data_graph(prime(tndag; sites=[])))
-  tntn = ITensorNetwork(union(tnbra, tnket))
+  # TODO: Use a different constructor here?
+  tntn = _ITensorNetwork(union(tnbra, tnket))
   for v in vertices(tn)
     if !isempty(commoninds(tntn[(v, 1)], tntn[(v, 2)]))
       add_edge!(tntn, (v, 1) => (v, 2))
@@ -809,6 +817,9 @@ end
 
 Base.show(io::IO, graph::AbstractITensorNetwork) = show(io, MIME"text/plain"(), graph)
 
+# TODO: Move to an `ITensorNetworksVisualizationInterfaceExt`
+# package extension (and define a `VisualizationInterface` package
+# based on `ITensorVisualizationCore`.).
 function ITensorVisualizationCore.visualize(
   tn::AbstractITensorNetwork,
   args...;
@@ -865,13 +876,13 @@ function site_combiners(tn::AbstractITensorNetwork{V}) where {V}
   return Cs
 end
 
-function insert_missing_internal_inds(
-  tn::AbstractITensorNetwork, edges; internal_inds_space=trivial_space(tn)
+function insert_linkinds(
+  tn::AbstractITensorNetwork, edges=edges(tn); link_space=trivial_space(tn)
 )
   tn = copy(tn)
   for e in edges
-    if !hascommoninds(tn[src(e)], tn[dst(e)])
-      iₑ = Index(internal_inds_space, edge_tag(e))
+    if !hascommoninds(tn, e)
+      iₑ = Index(link_space, edge_tag(e))
       X = onehot(iₑ => 1)
       tn[src(e)] *= X
       tn[dst(e)] *= dag(X)
@@ -880,12 +891,10 @@ function insert_missing_internal_inds(
   return tn
 end
 
-function insert_missing_internal_inds(
-  tn::AbstractITensorNetwork; internal_inds_space=trivial_space(tn)
-)
-  return insert_internal_inds(tn, edges(tn); internal_inds_space)
-end
-
+# TODO: What to output? Could be an `IndsNetwork`. Or maybe
+# that would be a different function `commonindsnetwork`.
+# Even in that case, this could output a `Dictionary`
+# from the edges to the common inds on that edge.
 function ITensors.commoninds(tn1::AbstractITensorNetwork, tn2::AbstractITensorNetwork)
   inds = Index[]
   for v1 in vertices(tn1)
@@ -911,8 +920,8 @@ function ITensorMPS.add(tn1::AbstractITensorNetwork, tn2::AbstractITensorNetwork
 
   if !issetequal(edges_tn1, edges_tn2)
     new_edges = union(edges_tn1, edges_tn2)
-    tn1 = insert_missing_internal_inds(tn1, new_edges)
-    tn2 = insert_missing_internal_inds(tn2, new_edges)
+    tn1 = insert_linkinds(tn1, new_edges)
+    tn2 = insert_linkinds(tn2, new_edges)
   end
 
   edges_tn1, edges_tn2 = edges(tn1), edges(tn2)

src/indsnetwork.jl

@@ -2,10 +2,10 @@ using DataGraphs: DataGraphs, DataGraph, IsUnderlyingGraph, map_data, vertex_dat
 using Dictionaries: AbstractDictionary, Indices
 using Graphs: Graphs
 using Graphs.SimpleGraphs: AbstractSimpleGraph
-# using LinearAlgebra: I # Not sure if this is needed
 using ITensors: Index, dag
 using ITensors.ITensorVisualizationCore: ITensorVisualizationCore, visualize
-using NamedGraphs: NamedGraphs, AbstractNamedGraph, NamedEdge, NamedGraph, vertextype
+using NamedGraphs:
+  NamedGraphs, AbstractNamedGraph, NamedEdge, NamedGraph, named_path_graph, vertextype
 
 struct IndsNetwork{V,I} <: AbstractIndsNetwork{V,I}
   data_graph::DataGraph{V,Vector{I},Vector{I},NamedGraph{V},NamedEdge{V}}