Skip to content

Commit

Permalink
Refactor and generalize tensor network constructors (#155)
Browse files Browse the repository at this point in the history
  • Loading branch information
@mtfishman
mtfishman authored Apr 12, 2024
1 parent ae4ad2c commit 906f184
Show file tree
Hide file tree
Showing 30 changed files with 657 additions and 691 deletions.
6 changes: 3 additions & 3 deletions Project.toml
View file
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "ITensorNetworks"
uuid = "2919e153-833c-4bdc-8836-1ea460a35fc7"
authors = ["Matthew Fishman <[email protected]> and contributors"]
version = "0.6"
version = "0.7"

[deps]
AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
Expand Down Expand Up @@ -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"
Expand Down
36 changes: 18 additions & 18 deletions README.md
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand All @@ -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:
Expand All @@ -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"))
```


Expand All @@ -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}}
Expand All @@ -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:
Expand All @@ -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;
⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀⠀
Expand Down
4 changes: 2 additions & 2 deletions src/ModelNetworks/ModelNetworks.jl
View file
Original file line number Diff line number Diff line change
@@ -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
Expand All @@ -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
Expand Down
72 changes: 67 additions & 5 deletions src/abstractindsnetwork.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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)
Expand All @@ -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
Expand Down
47 changes: 28 additions & 19 deletions src/abstractitensornetwork.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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

#
Expand Down Expand Up @@ -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)

#
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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))
Expand Down Expand Up @@ -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...;
Expand Down Expand Up @@ -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)
Expand All @@ -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)
Expand All @@ -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)
Expand Down
4 changes: 2 additions & 2 deletions src/indsnetwork.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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}}
Expand Down
Loading

2 comments on commit 906f184

@mtfishman
Copy link
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator register

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/104741

Tip: Release Notes

Did you know you can add release notes too? Just add markdown formatted text underneath the comment after the text
"Release notes:" and it will be added to the registry PR, and if TagBot is installed it will also be added to the
release that TagBot creates. i.e.

@JuliaRegistrator register

Release notes:

## Breaking changes

- blah

To add them here just re-invoke and the PR will be updated.

Tagging

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.7.0 -m "<description of version>" 906f184fa63915c3436d97bb2021554b1387fb31
git push origin v0.7.0

Please sign in to comment.