diff --git a/src/treetensornetworks/projttns/abstractprojttn.jl b/src/treetensornetworks/projttns/abstractprojttn.jl index 0a9c162c..b63546cb 100644 --- a/src/treetensornetworks/projttns/abstractprojttn.jl +++ b/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) - # remove internal edges (out of place) + 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(pos, P.H, getindices(P.environments, newenvskeys)) - # make all environments surrounding new position + 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 + 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 diff --git a/src/treetensornetworks/projttns/projttn.jl b/src/treetensornetworks/projttns/projttn.jl index 18baa3e2..faaa5541 100644 --- a/src/treetensornetworks/projttns/projttn.jl +++ b/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 diff --git a/test/test_treetensornetworks/test_position.jl b/test/test_treetensornetworks/test_position.jl index d5d0475d..a53a882d 100644 --- a/test/test_treetensornetworks/test_position.jl +++ b/test/test_treetensornetworks/test_position.jl @@ -41,3 +41,4 @@ using Test ITensors.disable_auto_fermion() end end +nothing