Fix mutation behavior of position

src/treetensornetworks/projttns/abstractprojttn.jl

@@ -131,13 +131,43 @@ end
 function position(
   P::AbstractProjTTN{V}, psi::TTN{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)
+  P = invalidate_environments(P)
+  P = make_environments(P, psi)
+  return P
+end
+
+function position!(
+  P::AbstractProjTTN{V}, psi::TTN{V}, pos::Union{Vector{<:V},NamedEdge{V}}
+) where {V}
+  P = shift_position(P, pos)
+  P = invalidate_environments(P)
+  make_environments!(P, psi, e)
+  return P
+end
+
+function invalidate_environments(P::AbstractProjTTN)
+  ie = internal_edges(P)
+  newenvskeys = filter(!in(ie), keys(P.environments))
+  P = ProjTTN(P.pos, P.H, getindices(P.environments, newenvskeys))
+  return P
+end
+
+function invalidate_environment(P::AbstractProjTTN{V}, e::NamedEdge{V}) where {V}
+  newenvskeys = filter(isequal(e), keys(P.environments))
+  P = ProjTTN(P.pos, P.H, getindices(P.environments, newenvskeys))
+  return P
+end
+
+function make_environments(P::AbstractProjTTN{V}, psi::TTN{V}) where {V}
+  P = copy(P)
+  for e in incident_edges(P)
+    set!(P.environments, e, make_environment(P, psi, e))
   end
-  # make all environments surrounding new position
+  return P
+end
+
+function make_environments!(P::AbstractProjTTN{V}, psi::TTN{V}) where {V}
   for e in incident_edges(P)
     make_environment!(P, psi, e)
   end

src/treetensornetworks/projttns/projttn.jl

@@ -11,7 +11,7 @@ function ProjTTN(H::TTN)
   return ProjTTN(vertices(H), H, Dictionary{edgetype(H),ITensor}())
 end
 
-copy(P::ProjTTN) = ProjTTN(P.pos, copy(P.H), copy(P.environments))
+copy(P::ProjTTN) = ProjTTN(P.pos, copy(P.H), NamedGraphs.copy_keys_values(P.environments))
 
 # trivial if we choose to specify position as above; only kept to allow using alongside
 # ProjMPO
@@ -30,23 +30,7 @@ function make_environment!(P::ProjTTN{V}, psi::TTN{V}, e::NamedEdge{V})::ITensor
   if haskey(P.environments, 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)]))
-    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))))
-      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)
-      # TODO: actually use optimal contraction sequence here
-      env = reduce(*, itensor_map)
-    end
-    # cache
+    env = _compute_environment(P, psi, e)
     set!(P.environments, e, env)
   end
   @assert(
@@ -55,3 +39,42 @@ function make_environment!(P::ProjTTN{V}, psi::TTN{V}, e::NamedEdge{V})::ITensor
   )
   return env
 end
+
+function make_environment(P::ProjTTN{V}, psi::TTN{V}, e::NamedEdge{V})::ITensor where {V}
+  # invalidate environment for opposite edge direction if necessary
+  #P=copy(P)
+  reverse(e) ∈ incident_edges(P) || (P = invalidate_environment(P, reverse(e)))
+  # do nothing if valid environment already present
+  if haskey(P.environments, e)
+    env = environment(P, e)
+  else
+    env = _compute_environment(P, psi, e)
+  end
+  @assert(
+    hascommoninds(env, psi[src(e)]),
+    "Something went wrong, probably re-orthogonalized this edge in the same direction twice!"
+  )
+  return env
+end
+
+function _compute_environment(
+  P::ProjTTN{V}, psi::TTN{V}, e::NamedEdge{V}
+)::ITensor where {V}
+  if is_leaf(underlying_graph(P), src(e))
+    # leaves are easy
+    env = psi[src(e)] * P.H[src(e)] * dag(prime(psi[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))))
+    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)
+    # TODO: actually use optimal contraction sequence here
+    env = reduce(*, itensor_map)
+  end
+  return env
+end

test/test_treetensornetworks/test_position.jl

@@ -0,0 +1,44 @@
+using ITensors
+using ITensorNetworks
+using ITensorNetworks: position
+using Test
+
+@testset "ProjTTN position copy-safe" begin
+  # make a nontrivial TTN state and TTN operator
+
+  auto_fermion_enabled = ITensors.using_auto_fermion()
+  use_qns = true
+  cutoff = 1e-12
+
+  tooth_lengths = fill(2, 3)
+  c = named_comb_tree(tooth_lengths)
+  if use_qns  # test whether autofermion breaks things when using non-fermionic QNs
+    ITensors.enable_auto_fermion()
+  else        # when using no QNs, autofermion breaks # ToDo reference Issue in ITensors
+    ITensors.disable_auto_fermion()
+  end
+  s = siteinds("S=1/2", c; conserve_qns=use_qns)
+
+  os = ITensorNetworks.heisenberg(c)
+
+  H = TTN(os, s)
+
+  d = Dict()
+  for (i, v) in enumerate(vertices(s))
+    d[v] = isodd(i) ? "Up" : "Dn"
+  end
+  states = v -> d[v]
+  psi = TTN(s, states)
+
+  # actual test, verifies that position is copy safe
+  vs = vertices(s)
+  PH0 = ProjTTN(H)
+  PH0 = position(PH0, psi, [vs[2]])
+  PH = copy(PH0)
+  PH = position(PH, psi, [vs[2], vs[5]])
+  @test keys(PH.environments) != keys(PH0.environments)
+  if !auto_fermion_enabled
+    ITensors.disable_auto_fermion()
+  end
+end
+nothing