Skip to content

Commit

Permalink
Fix mutation behavior of position (#131)
Browse files Browse the repository at this point in the history
  • Loading branch information
b-kloss authored Jan 31, 2024
1 parent 3fc272e commit cf59738
Show file tree
Hide file tree
Showing 6 changed files with 154 additions and 62 deletions.
4 changes: 2 additions & 2 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -37,13 +37,13 @@ Combinatorics = "1"
Compat = "3, 4"
DataGraphs = "0.1.7"
DataStructures = "0.18"
Dictionaries = "0.3.15"
Dictionaries = "0.4"
Distributions = "0.25.86"
DocStringExtensions = "0.8, 0.9"
Graphs = "1.8"
GraphsFlows = "0.1.1"
ITensors = "0.3.23"
IsApprox = "0.1.7"
IsApprox = "0.1"
IterTools = "1.4.0"
KrylovKit = "0.6.0"
NamedGraphs = "0.1.11"
Expand Down
2 changes: 1 addition & 1 deletion src/solvers/contract.jl
Original file line number Diff line number Diff line change
Expand Up @@ -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, (;)
Expand Down
60 changes: 38 additions & 22 deletions src/treetensornetworks/projttns/abstractprojttn.jl
Original file line number Diff line number Diff line change
@@ -1,17 +1,22 @@
abstract type AbstractProjTTN{V} end

environments(::AbstractProjTTN) = error("Not implemented")
operator(::AbstractProjTTN) = error("Not implemented")
pos(::AbstractProjTTN) = error("Not implemented")

underlying_graph(P::AbstractProjTTN) = error("Not implemented")

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))

Expand Down Expand Up @@ -43,8 +48,8 @@ function internal_edges(P::AbstractProjTTN{V})::Vector{NamedEdge{V}} where {V}
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]
function environment(P::AbstractProjTTN, edge::AbstractEdge)
return environments(P)[edge]
end

# there has to be a better way to do this...
Expand All @@ -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
Expand Down Expand Up @@ -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)))
Expand All @@ -121,25 +126,36 @@ 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}}
) where {V}
# shift position
function position(P::AbstractProjTTN, psi::AbstractTTN, pos)
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
P = invalidate_environments(P)
P = make_environments(P, psi)
return P
end

function invalidate_environments(P::AbstractProjTTN)
ie = internal_edges(P)
newenvskeys = filter(!in(ie), keys(environments(P)))
P = set_environments(P, getindices(environments(P), newenvskeys))
return P
end

function invalidate_environment(P::AbstractProjTTN, e::AbstractEdge)
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
47 changes: 30 additions & 17 deletions src/treetensornetworks/projttns/projttn.jl
Original file line number Diff line number Diff line change
Expand Up @@ -3,15 +3,21 @@ ProjTTN
"""
struct ProjTTN{V} <: AbstractProjTTN{V}
pos::Union{Vector{<:V},NamedEdge{V}} # TODO: cleanest way to specify effective Hamiltonian position?
H::TTN{V}
operator::TTN{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(environments(P)))

#accessors for fields
environments(p::ProjTTN) = p.environments
operator(p::ProjTTN) = p.operator
underlying_graph(P::ProjTTN) = underlying_graph(operator(P))
pos(P::ProjTTN) = P.pos

# trivial if we choose to specify position as above; only kept to allow using alongside
# ProjMPO
Expand All @@ -20,38 +26,45 @@ 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

set_environments(p::ProjTTN, environments) = ProjTTN(pos(p), operator(p), environments)
set_environment(p::ProjTTN, edge, env) = set_environment!(copy(p), edge, env)
function set_environment!(p::ProjTTN, edge, env)
set!(environments(p), edge, env)
return p
end

function make_environment!(P::ProjTTN{V}, psi::TTN{V}, e::NamedEdge{V})::ITensor where {V}
function make_environment(P::ProjTTN, state::AbstractTTN, e::AbstractEdge)
# invalidate environment for opposite edge direction if necessary
reverse(e) incident_edges(P) || unset!(P.environments, reverse(e))
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
if !haskey(environments(P), e)
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(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)])
push!(neighbor_envs, make_environment!(P, psi, edgetype(P)(n, src(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(psi[src(e)], frst, scnd, P.H[src(e)], dag(prime(psi[src(e)])), rst)
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
# cache
set!(P.environments, e, env)
P = 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
return P
end
56 changes: 36 additions & 20 deletions src/treetensornetworks/projttns/projttn_apply.jl
Original file line number Diff line number Diff line change
@@ -1,57 +1,73 @@
struct ProjTTNApply{V} <: AbstractProjTTN{V}
pos::Union{Vector{<:V},NamedEdge{V}}
psi0::TTN{V}
H::TTN{V}
init_state::TTN{V}
operator::TTN{V}
environments::Dictionary{NamedEdge{V},ITensor}
end

function ProjTTNApply(psi0::TTN, H::TTN)
return ProjTTNApply(vertextype(H)[], psi0, H, Dictionary{edgetype(H),ITensor}())
environments(p::ProjTTNApply) = p.environments
operator(p::ProjTTNApply) = p.operator
underlying_graph(p::ProjTTNApply) = underlying_graph(operator(p))
pos(p::ProjTTNApply) = p.pos
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(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 make_environment!(
P::ProjTTNApply{V}, psi::TTN{V}, e::NamedEdge{V}
)::ITensor where {V}
set_environment(p::ProjTTNApply, edge, env) = set_environment!(copy(p), edge, env)
function set_environment!(p::ProjTTNApply, edge, env)
set!(environments(p), edge, env)
return p
end

function make_environment(P::ProjTTNApply, state::AbstractTTN, e::AbstractEdge)
# invalidate environment for opposite edge direction if necessary
reverse(e) incident_edges(P) || unset!(P.environments, reverse(e))
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
if !haskey(environments(P), e)
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))))
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(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)
P = 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
return P
end
47 changes: 47 additions & 0 deletions test/test_treetensornetworks/test_position.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
using ITensors
using ITensorNetworks
using ITensorNetworks: position, environments
using Test

@testset "ProjTTN position" 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 out of place
vs = vertices(s)
PH = ProjTTN(H)
PH = position(PH, psi, [vs[2]])
original_keys = deepcopy(keys(environments(PH)))
# test out-of-placeness of position
PHc = position(PH, psi, [vs[2], vs[5]])
@test keys(environments(PH)) == original_keys
@test keys(environments(PHc)) != original_keys

if !auto_fermion_enabled
ITensors.disable_auto_fermion()
end
end
nothing

0 comments on commit cf59738

Please sign in to comment.