New MPO/MPO index manipulation interface (#575)

ITensor · Feb 5, 2021 · 913e832 · 913e832
1 parent 783f8b4
commit 913e832
Show file tree

Hide file tree

Showing 11 changed files with 344 additions and 210 deletions.
diff --git a/docs/src/MPSandMPO.md b/docs/src/MPSandMPO.md
@@ -42,8 +42,8 @@ hasqns(::ITensors.AbstractMPS)
 ## Obtaining and finding indices
 
 ```@docs
-common_siteind(::ITensors.AbstractMPS, ::ITensors.AbstractMPS, ::Int)
-common_siteinds(::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+siteinds(::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS, ::Int)
+siteinds(::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS, ::Int)
 findsite
 findsites
 firstsiteinds
@@ -54,27 +54,67 @@ siteinds(::MPS)
 siteind(::MPO, ::Int)
 siteinds(::MPO)
 siteinds(::ITensors.AbstractMPS, ::Int)
-unique_siteind(::ITensors.AbstractMPS, ::ITensors.AbstractMPS, ::Int)
-unique_siteinds(::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
 ```
 
 ## Priming and tagging
 
 ```@docs
 prime(::ITensors.AbstractMPS)
 prime!(::ITensors.AbstractMPS)
+prime(::typeof(linkinds), ::ITensors.AbstractMPS)
+prime!(::typeof(linkinds), ::ITensors.AbstractMPS)
+prime(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+prime!(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+prime(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+prime!(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
 setprime(::ITensors.AbstractMPS)
 setprime!(::ITensors.AbstractMPS)
+setprime(::typeof(linkinds), ::ITensors.AbstractMPS)
+setprime!(::typeof(linkinds), ::ITensors.AbstractMPS)
+setprime(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+setprime!(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+setprime(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+setprime!(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
 noprime(::ITensors.AbstractMPS)
 noprime!(::ITensors.AbstractMPS)
+noprime(::typeof(linkinds), ::ITensors.AbstractMPS)
+noprime!(::typeof(linkinds), ::ITensors.AbstractMPS)
+noprime(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+noprime!(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+noprime(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+noprime!(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
 addtags(::ITensors.AbstractMPS)
 addtags!(::ITensors.AbstractMPS)
+addtags(::typeof(linkinds), ::ITensors.AbstractMPS)
+addtags!(::typeof(linkinds), ::ITensors.AbstractMPS)
+addtags(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+addtags!(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+addtags(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+addtags!(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
 removetags(::ITensors.AbstractMPS)
 removetags!(::ITensors.AbstractMPS)
+removetags(::typeof(linkinds), ::ITensors.AbstractMPS)
+removetags!(::typeof(linkinds), ::ITensors.AbstractMPS)
+removetags(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+removetags!(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+removetags(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+removetags!(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
 replacetags(::ITensors.AbstractMPS)
 replacetags!(::ITensors.AbstractMPS)
+replacetags(::typeof(linkinds), ::ITensors.AbstractMPS)
+replacetags!(::typeof(linkinds), ::ITensors.AbstractMPS)
+replacetags(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+replacetags!(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+replacetags(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+replacetags!(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
 settags(::ITensors.AbstractMPS)
 settags!(::ITensors.AbstractMPS)
+settags(::typeof(linkinds), ::ITensors.AbstractMPS)
+settags!(::typeof(linkinds), ::ITensors.AbstractMPS)
+settags(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+settags!(::typeof(siteinds), ::typeof(commoninds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+settags(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
+settags!(::typeof(siteinds), ::typeof(uniqueinds), ::ITensors.AbstractMPS, ::ITensors.AbstractMPS)
 ```
 
 ## Operations

diff --git a/src/ITensors.jl b/src/ITensors.jl
@@ -84,6 +84,7 @@ include("qn/qnitensor.jl")
 # MPS/MPO
 #
 include("mps/abstractmps.jl")
+include("mps/deprecated.jl")
 include("mps/mps.jl")
 include("mps/mpo.jl")
 include("mps/sweeps.jl")

diff --git a/src/broadcast.jl b/src/broadcast.jl
@@ -5,17 +5,26 @@
 
 BroadcastStyle(::Type{<:IndexSet}) = Style{IndexSet}()
 
-broadcastable(t::IndexSet) = t
+BroadcastStyle(::AbstractArrayStyle{0}, b::Style{IndexSet}) = b
+BroadcastStyle(a::AbstractArrayStyle, ::Style{IndexSet}) = a
+BroadcastStyle(a::Style{Tuple}, ::Style{IndexSet}) = a
 
-# This is required for speed and type stability
-# (not sure why, but I guess the default is slow)
 instantiate(bc::Broadcasted{Style{IndexSet}, Nothing}) = bc
+function instantiate(bc::Broadcasted{Style{IndexSet}})
+  check_broadcast_axes(bc.axes, bc.args...)
+  return bc
+end
+
+broadcastable(t::IndexSet) = t
+
+indexset_or_tuple(t::Tuple{Vararg{<:Index}}) = IndexSet(t)
+indexset_or_tuple(t) = t
 
-function Base.copy(bc::Broadcasted{Style{IndexSet}})
-  dim = axes(bc)
-  length(dim) == 1 || throw(DimensionMismatch("IndexSet only supports one dimension"))
-  N = length(dim[1])
-  return IndexSet(ntuple(k -> bc[k], Val(N)))
+@inline function copy(bc::Broadcasted{Style{IndexSet}})
+    dim = axes(bc)
+    length(dim) == 1 || throw(DimensionMismatch("tuple only supports one dimension"))
+    N = length(dim[1])
+    return indexset_or_tuple(ntuple(k -> @inbounds(_broadcast_getindex(bc, k)), Val(N)))
 end
 
 #

diff --git a/src/imports.jl b/src/imports.jl
@@ -32,6 +32,7 @@ import Base:
   keys,
   lastindex,
   length,
+  map,
   map!,
   ndims,
   push!,
@@ -45,11 +46,13 @@ import Base:
 
 import Base.Broadcast:
   # types
+  AbstractArrayStyle,
   Broadcasted,
   BroadcastStyle,
   DefaultArrayStyle,
   Style,
   # functions
+  _broadcast_getindex,
   broadcasted,
   broadcastable,
   instantiate

diff --git a/src/indexset.jl b/src/indexset.jl
@@ -969,9 +969,18 @@ end
 swaptags(is::IndexSet, tags1, tags2, args...; kwargs...) =
   swaptags(fmatch(args...; kwargs...), is, tags1, tags2)
 
-replaceinds(is::IndexSet, rep_inds::Pair{ <: Index, <: Index}...) =
+replaceinds(is::IndexSet, rep_inds::Pair{<: Index, <: Index}...) =
   replaceinds(is, zip(rep_inds...)...)
 
+replaceinds(is::IndexSet, rep_inds::Vector{Pair{<: Index, <: Index}}) =
+  replaceinds(is, rep_inds...)
+
+replaceinds(is::IndexSet, rep_inds::Tuple{Vararg{Pair{<: Index, <: Index}}}) =
+  replaceinds(is, rep_inds...)
+
+replaceinds(is::IndexSet, rep_inds::Pair) =
+  replaceinds(is, Tuple(first(rep_inds)) .=> Tuple(last(rep_inds)))
+
 # Check that the QNs are all the same
 hassameflux(i1::Index, i2::Index) = (dim(i1) == dim(i2))