Fix mutation behavior of position

src/ITensorNetworks.jl

@@ -20,6 +20,7 @@ using IterTools
 using KrylovKit: KrylovKit
 using LinearAlgebra
 using NamedGraphs
+
 using Observers
 using Observers.DataFrames: select!
 using Printf
@@ -33,6 +34,7 @@ using TimerOutputs
 
 using DataGraphs: IsUnderlyingGraph, edge_data_type, vertex_data_type
 using Graphs: AbstractEdge, AbstractGraph, Graph, add_edge!
+using NamedGraphs: copy_keys_values
 using ITensors:
   @Algorithm_str,
   @debug_check,

src/solvers/contract.jl

@@ -12,7 +12,7 @@ function contract_updater(
   v = ITensor(true)
   projected_operator = projected_operator![]
   for j in sites(projected_operator)
-    v *= projected_operator.psi0[j]
+    v *= init_state(projected_operator)[j]
   end
   vp = contract(projected_operator, v)
   return vp, (;)

src/treetensornetworks/projttns/abstractprojttn.jl

@@ -1,17 +1,22 @@
 abstract type AbstractProjTTN{V} end
 
+#accessors for fields common to all concrete subtypes
+environments(p::AbstractProjTTN) = p.environments
+operator(p::AbstractProjTTN) = p.operator
+underlying_graph(P::AbstractProjTTN) = underlying_graph(operator(P))
+pos(P::AbstractProjTTN) = P.pos
+
 copy(::AbstractProjTTN) = error("Not implemented")
 
 set_nsite(::AbstractProjTTN, nsite) = error("Not implemented")
 
 # silly constructor wrapper
 shift_position(::AbstractProjTTN, pos) = error("Not implemented")
 
-make_environment!(::AbstractProjTTN, psi, e) = error("Not implemented")
-
-underlying_graph(P::AbstractProjTTN) = underlying_graph(P.H)
-
-pos(P::AbstractProjTTN) = P.pos
+set_environments(p::AbstractProjTTN, environments) = error("Not implemented")
+set_environment(p::AbstractProjTTN, edge, environment) = error("Not implemented")
+make_environment!(P::AbstractProjTTN, psi, e) = error("Not implemented")
+make_environment(P::AbstractProjTTN, psi, e) = error("Not implemented")
 
 Graphs.edgetype(P::AbstractProjTTN) = edgetype(underlying_graph(P))
 
@@ -44,7 +49,7 @@ end
 
 environment(P::AbstractProjTTN, edge::Pair) = environment(P, edgetype(P)(edge))
 function environment(P::AbstractProjTTN{V}, edge::NamedEdge{V})::ITensor where {V}
-  return P.environments[edge]
+  return environments(P)[edge]
 end
 
 # there has to be a better way to do this...
@@ -69,9 +74,9 @@ function contract(P::AbstractProjTTN, v::ITensor)::ITensor
   else
     itensor_map = Union{ITensor,OneITensor}[] # TODO: will a Hamiltonian TTN tensor ever be a OneITensor?
     for s in sites(P)
-      site_envs = filter(hascommoninds(P.H[s]), environments)
+      site_envs = filter(hascommoninds(operator(P)[s]), environments)
       frst, scnd, rst = _separate_first_two(site_envs)
-      site_tensors = vcat(frst, scnd, P.H[s], rst)
+      site_tensors = vcat(frst, scnd, operator(P)[s], rst)
       append!(itensor_map, site_tensors)
     end
   end
@@ -103,9 +108,9 @@ end
 (P::AbstractProjTTN)(v::ITensor) = product(P, v)
 
 function Base.eltype(P::AbstractProjTTN)::Type
-  ElType = eltype(P.H(first(sites(P))))
+  ElType = eltype(operator(P)(first(sites(P))))
   for v in sites(P)
-    ElType = promote_type(ElType, eltype(P.H[v]))
+    ElType = promote_type(ElType, eltype(operator(P)[v]))
   end
   for e in incident_edges(P)
     ElType = promote_type(ElType, eltype(environments(P, e)))
@@ -121,25 +126,55 @@ function Base.size(P::AbstractProjTTN)::Tuple{Int,Int}
     end
   end
   for j in sites(P)
-    for i in inds(P.H[j])
+    for i in inds(operator(P)[j])
       plev(i) > 0 && (d *= dim(i))
     end
   end
   return (d, d)
 end
 
 function position(
-  P::AbstractProjTTN{V}, psi::TTN{V}, pos::Union{Vector{<:V},NamedEdge{V}}
+  P::AbstractProjTTN{V}, psi::AbstractTTN{V}, pos::Union{Vector{<:V},NamedEdge{V}}
+) where {V}
+  P = shift_position(P, pos)
+  P = invalidate_environments(P)
+  P = make_environments(P, psi)
+  return P
+end
+
+function position!(
+  P::AbstractProjTTN{V}, psi::AbstractTTN{V}, pos::Union{Vector{<:V},NamedEdge{V}}
 ) where {V}
-  # shift position
   P = shift_position(P, pos)
-  # invalidate environments corresponding to internal edges
   for e in internal_edges(P)
     unset!(P.environments, e)
   end
-  # make all environments surrounding new position
+  make_environments!(P, psi)
+  return P
+end
+
+function invalidate_environments(P::AbstractProjTTN)
+  ie = internal_edges(P)
+  newenvskeys = filter(!in(ie), keys(P.environments))
+  P = set_environments(P, getindices(P.environments, newenvskeys))
+  return P
+end
+
+function invalidate_environment(P::AbstractProjTTN, e::NamedEdge)
+  newenvskeys = filter(!isequal(e), keys(environments(P)))
+  P = set_environments(P, getindices(environments(P), newenvskeys))
+  return P
+end
+
+function make_environments(P::AbstractProjTTN, psi::AbstractTTN)
   for e in incident_edges(P)
-    make_environment!(P, psi, e)
+    P = make_environment(P, psi, e)
   end
   return P
 end
+
+function make_environments!(P::AbstractProjTTN, psi::AbstractTTN)
+  for e in incident_edges(P)
+    make_environment!(P, psi, e)
+  end
+end

src/treetensornetworks/projttns/projttn.jl

@@ -3,15 +3,16 @@ ProjTTN
 """
 struct ProjTTN{V} <: AbstractProjTTN{V}
   pos::Union{Vector{<:V},NamedEdge{V}} # TODO: cleanest way to specify effective Hamiltonian position?
-  H::TTN{V}
+  operator::AbstractTTN{V}
   environments::Dictionary{NamedEdge{V},ITensor}
 end
 
-function ProjTTN(H::TTN)
-  return ProjTTN(vertices(H), H, Dictionary{edgetype(H),ITensor}())
+function ProjTTN(operator::TTN)
+  return ProjTTN(vertices(operator), operator, Dictionary{edgetype(operator),ITensor}())
 end
 
-copy(P::ProjTTN) = ProjTTN(P.pos, copy(P.H), copy(P.environments))
+copy(P::ProjTTN) = ProjTTN(pos(P), copy(operator(P)), copy_keys_values(environments(P)))
+unsafe_copy(P::ProjTTN) = ProjTTN(pos(P), copy(operator(P)), copy(environments(P)))
 
 # trivial if we choose to specify position as above; only kept to allow using alongside
 # ProjMPO
@@ -20,38 +21,83 @@ function set_nsite(P::ProjTTN, nsite)
 end
 
 function shift_position(P::ProjTTN, pos)
-  return ProjTTN(pos, P.H, P.environments)
+  return ProjTTN(pos, operator(P), environments(P))
 end
 
-function make_environment!(P::ProjTTN{V}, psi::TTN{V}, e::NamedEdge{V})::ITensor where {V}
+set_environments(p::ProjTTN, environments) = ProjTTN(pos(p), operator(p), environments)
+function set_environment(p::ProjTTN, edge, env)
+  newenvs = merge(environments(p), Dictionary((edge,), (env,)))
+  return ProjTTN(pos(p), operator(p), newenvs)
+end
+set_environment!(p::ProjTTN, edge, env) = set!(environments(p), edge, env)
+
+function make_environment!(
+  P::ProjTTN{V}, state::AbstractTTN{V}, e::NamedEdge{V}
+)::ITensor where {V}
   # invalidate environment for opposite edge direction if necessary
-  reverse(e) ∈ incident_edges(P) || unset!(P.environments, reverse(e))
+  reverse(e) ∈ incident_edges(P) || unset!(environments(P), reverse(e))
   # do nothing if valid environment already present
-  if haskey(P.environments, e)
+  if haskey(environments(P), e)
     env = environment(P, e)
   else
     if is_leaf(underlying_graph(P), src(e))
       # leaves are easy
-      env = psi[src(e)] * P.H[src(e)] * dag(prime(psi[src(e)]))
+      env = state[src(e)] * operator(H)[src(e)] * dag(prime(state[src(e)]))
     else
       # construct by contracting neighbors
       neighbor_envs = ITensor[]
       for n in setdiff(neighbors(underlying_graph(P), src(e)), [dst(e)])
-        push!(neighbor_envs, make_environment!(P, psi, edgetype(P)(n, src(e))))
+        push!(neighbor_envs, make_environment!(P, state, edgetype(P)(n, src(e))))
       end
       # manually heuristic for contraction order: two environments, site tensors, then
       # other environments
       frst, scnd, rst = _separate_first_two(neighbor_envs)
-      itensor_map = vcat(psi[src(e)], frst, scnd, P.H[src(e)], dag(prime(psi[src(e)])), rst)
+      itensor_map = vcat(
+        state[src(e)], frst, scnd, operator(H)[src(e)], dag(prime(state[src(e)])), rst
+      )
       # TODO: actually use optimal contraction sequence here
       env = reduce(*, itensor_map)
     end
-    # cache
-    set!(P.environments, e, env)
+    set!(environments(P), e, env)
   end
   @assert(
-    hascommoninds(environment(P, e), psi[src(e)]),
+    hascommoninds(environment(P, e), state[src(e)]),
     "Something went wrong, probably re-orthogonalized this edge in the same direction twice!"
   )
   return env
 end
+
+function make_environment(
+  P::ProjTTN{V}, state::AbstractTTN{V}, e::NamedEdge{V}
+)::ProjTTN{V} where {V}
+  # invalidate environment for opposite edge direction if necessary
+  reverse(e) ∈ incident_edges(P) || (P = invalidate_environment(P, reverse(e)))
+  # do nothing if valid environment already present
+  if !haskey(environments(P), e)
+    if is_leaf(underlying_graph(P), src(e))
+      # leaves are easy
+      env = state[src(e)] * operator(P)[src(e)] * dag(prime(state[src(e)]))
+    else
+      # construct by contracting neighbors
+      neighbor_envs = ITensor[]
+      for n in setdiff(neighbors(underlying_graph(P), src(e)), [dst(e)])
+        P = make_environment(P, state, edgetype(P)(n, src(e)))
+        push!(neighbor_envs, environment(P, edgetype(P)(n, src(e))))
+      end
+      # manually heuristic for contraction order: two environments, site tensors, then
+      # other environments
+      frst, scnd, rst = _separate_first_two(neighbor_envs)
+      itensor_map = vcat(
+        state[src(e)], frst, scnd, operator(P)[src(e)], dag(prime(state[src(e)])), rst
+      )
+      # TODO: actually use optimal contraction sequence here
+      env = reduce(*, itensor_map)
+    end
+    P = set_environment(P, e, env)
+  end
+  @assert(
+    hascommoninds(environment(P, e), state[src(e)]),
+    "Something went wrong, probably re-orthogonalized this edge in the same direction twice!"
+  )
+  return P
+end

src/treetensornetworks/projttns/projttn_apply.jl

@@ -1,57 +1,114 @@
 struct ProjTTNApply{V} <: AbstractProjTTN{V}
   pos::Union{Vector{<:V},NamedEdge{V}}
-  psi0::TTN{V}
-  H::TTN{V}
+  init_state::AbstractTTN{V}
+  operator::AbstractTTN{V}
   environments::Dictionary{NamedEdge{V},ITensor}
 end
 
-function ProjTTNApply(psi0::TTN, H::TTN)
-  return ProjTTNApply(vertextype(H)[], psi0, H, Dictionary{edgetype(H),ITensor}())
+init_state(p::ProjTTNApply) = p.init_state
+
+function ProjTTNApply(init_state::AbstractTTN, operator::AbstractTTN)
+  return ProjTTNApply(
+    vertextype(operator)[], init_state, operator, Dictionary{edgetype(operator),ITensor}()
+  )
 end
 
 function copy(P::ProjTTNApply)
-  return ProjTTNApply(P.pos, copy(P.psi0), copy(P.H), copy(P.environments))
+  return ProjTTNApply(
+    pos(P), copy(init_state(P)), copy(operator(P)), copy_keys_values(environments(P))
+  )
 end
 
+function unsafe_copy(P::ProjTTNApply)
+  return ProjTTNApply(P.pos, copy(init_state(P)), copy(operator(P)), copy(environments(P)))
+end
 function set_nsite(P::ProjTTNApply, nsite)
   return P
 end
 
 function shift_position(P::ProjTTNApply, pos)
-  return ProjTTNApply(pos, P.psi0, P.H, P.environments)
+  return ProjTTNApply(pos, init_state(P), operator(P), environments(P))
+end
+
+function set_environments(p::ProjTTNApply, environments)
+  return ProjTTNApply(pos(p), init_state(p), operator(p), environments)
+end
+
+function set_environment(p::ProjTTNApply, edge, env)
+  newenv = merge(p.environments, Dictionary((edge,), (env,)))
+  return ProjTTNApply(pos(p), init_state(p), operator(p), newenv)
 end
 
+set_environment!(p::ProjTTNApply, edge, env) = set!(environments(P), edge, env)
+
 function make_environment!(
-  P::ProjTTNApply{V}, psi::TTN{V}, e::NamedEdge{V}
+  P::ProjTTNApply{V}, state::AbstractTTN{V}, e::NamedEdge{V}
 )::ITensor where {V}
   # invalidate environment for opposite edge direction if necessary
-  reverse(e) ∈ incident_edges(P) || unset!(P.environments, reverse(e))
+  reverse(e) ∈ incident_edges(P) || unset!(environments(P), reverse(e))
   # do nothing if valid environment already present
-  if haskey(P.environments, e)
+  if haskey(environments(P), e)
     env = environment(P, e)
   else
     if is_leaf(underlying_graph(P), src(e))
       # leaves are easy
-      env = P.psi0[src(e)] * P.H[src(e)] * dag(psi[src(e)])
+      env = init_state(P)[src(e)] * operator(P)[src(e)] * dag(state[src(e)])
     else
       # construct by contracting neighbors
       neighbor_envs = ITensor[]
       for n in setdiff(neighbors(underlying_graph(P), src(e)), [dst(e)])
-        push!(neighbor_envs, make_environment!(P, psi, edgetype(P)(n, src(e))))
+        push!(neighbor_envs, make_environment!(P, state, edgetype(P)(n, src(e))))
       end
       # manually heuristic for contraction order: two environments, site tensors, then
       # other environments
       frst, scnd, rst = _separate_first_two(neighbor_envs)
-      itensor_map = vcat(P.psi0[src(e)], frst, scnd, P.H[src(e)], dag(psi[src(e)]), rst)
+      itensor_map = vcat(
+        init_state(P)[src(e)], frst, scnd, operator(P)[src(e)], dag(state[src(e)]), rst
+      ) # no prime here in comparison to the same routine for Projttn
       # TODO: actually use optimal contraction sequence here
       env = reduce(*, itensor_map)
     end
     # cache
-    set!(P.environments, e, env)
+    set_environment!(P, e, env)
   end
   @assert(
-    hascommoninds(environment(P, e), psi[src(e)]),
+    hascommoninds(environment(P, e), state[src(e)]),
     "Something went wrong, probably re-orthogonalized this edge in the same direction twice!"
   )
   return env
 end
+
+function make_environment(
+  P::ProjTTNApply{V}, state::AbstractTTN{V}, e::NamedEdge{V}
+)::ProjTTNApply{V} where {V}
+  # invalidate environment for opposite edge direction if necessary
+  reverse(e) ∈ incident_edges(P) || (P = invalidate_environment(P, reverse(e)))
+  # do nothing if valid environment already present
+  if !haskey(environments(P), e)
+    if is_leaf(underlying_graph(P), src(e))
+      # leaves are easy
+      env = init_state(P)[src(e)] * operator(P)[src(e)] * dag(state[src(e)])
+    else
+      # construct by contracting neighbors
+      neighbor_envs = ITensor[]
+      for n in setdiff(neighbors(underlying_graph(P), src(e)), [dst(e)])
+        P = make_environment(P, state, edgetype(P)(n, src(e)))
+        push!(neighbor_envs, environment(P, edgetype(P)(n, src(e))))
+      end
+      # manually heuristic for contraction order: two environments, site tensors, then
+      # other environments
+      frst, scnd, rst = _separate_first_two(neighbor_envs)
+      itensor_map = vcat(
+        init_state(P)[src(e)], frst, scnd, operator(P)[src(e)], dag(state[src(e)]), rst
+      ) # no prime here in comparison to the same routine for Projttn
+      # TODO: actually use optimal contraction sequence here
+      env = reduce(*, itensor_map)
+    end
+    P = set_environment(P, e, env)
+  end
+  @assert(
+    hascommoninds(environment(P, e), state[src(e)]),
+    "Something went wrong, probably re-orthogonalized this edge in the same direction twice!"
+  )
+  return P
+end