Skip to content
This repository has been archived by the owner on Jul 7, 2024. It is now read-only.

Add order keyword argument in Chain constructor functions #47

Merged
merged 10 commits into from
Jun 19, 2024
Merged

Add order keyword argument in Chain constructor functions #47

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
10 commits
Select commit Hold shift + click to select a range
a34381c
Add order kwarg in Chain constructor
jofrevalles Jun 18, 2024
fc42668
Add tests
jofrevalles Jun 18, 2024
73449c3
Fix minor typos
jofrevalles Jun 18, 2024
f305b00
Add minor fixes in Chain constructor
jofrevalles Jun 19, 2024
1467f80
Enhance tests
jofrevalles Jun 19, 2024
5f50707
Format code
jofrevalles Jun 19, 2024
ff981fa
Apply @mofeing suggestions from code review
jofrevalles Jun 19, 2024
7375cb5
Remove unnecessary helper function
jofrevalles Jun 19, 2024
e52302b
Apply suggestions from code review
jofrevalles Jun 19, 2024
025a693
Minor aesthetic updates in code
jofrevalles Jun 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
95 changes: 79 additions & 16 deletions src/Ansatz/Chain.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -30,52 +30,98 @@ function Chain(tn::TensorNetwork, sites, args...; kwargs...)
Chain(Quantum(tn, sites), args...; kwargs...)
end

function Chain(::State, boundary::Periodic, arrays::Vector{<:AbstractArray})
defaultorder(::Type{Chain}, ::State) = (:o, :l, :r)
defaultorder(::Type{Chain}, ::Operator) = (:o, :i, :l, :r)

function Chain(::State, boundary::Periodic, arrays::Vector{<:AbstractArray}; order = defaultorder(Chain, State()))
@assert all(==(3) ∘ ndims, arrays) "All arrays must have 3 dimensions"
issetequal(order, defaultorder(Chain, State())) ||
throw(ArgumentError("order must be a permutation of $(String.(defaultorder(State())))"))

n = length(arrays)
symbols = [nextindex() for _ in 1:2n]

_tensors = map(enumerate(arrays)) do (i, array)
Tensor(array, [symbols[i], symbols[n+mod1(i - 1, n)], symbols[n+mod1(i, n)]])
inds = map(order) do dir
if dir == :o
symbols[i]
elseif dir == :r
symbols[n+mod1(i, n)]
elseif dir == :l
symbols[n+mod1(i - 1, n)]
else
throw(ArgumentError("Invalid direction: $dir"))
end
end
Tensor(array, inds)
end

sitemap = Dict(Site(i) => symbols[i] for i in 1:n)

Chain(Quantum(TensorNetwork(_tensors), sitemap), boundary)
end

function Chain(::State, boundary::Open, arrays::Vector{<:AbstractArray})
function Chain(::State, boundary::Open, arrays::Vector{<:AbstractArray}; order = defaultorder(Chain, State()))
@assert ndims(arrays[1]) == 2 "First array must have 2 dimensions"
@assert all(==(3) ∘ ndims, arrays[2:end-1]) "All arrays must have 3 dimensions"
@assert ndims(arrays[end]) == 2 "Last array must have 2 dimensions"
issetequal(order, defaultorder(Chain, State())) ||
throw(ArgumentError("order must be a permutation of $(String.(defaultorder(Chain, State())))"))

n = length(arrays)
symbols = [nextindex() for _ in 1:2n-1]
symbols = [nextindex() for _ in 1:2n]

_tensors = map(enumerate(arrays)) do (i, array)
if i == 1
Tensor(array, [symbols[1], symbols[1+n]])
_order = if i == 1
filter(x -> x != :l, order)
elseif i == n
Tensor(array, [symbols[n], symbols[n+mod1(n - 1, n)]])
filter(x -> x != :r, order)
else
Tensor(array, [symbols[i], symbols[n+mod1(i - 1, n)], symbols[n+mod1(i, n)]])
order
end

inds = map(_order) do dir
if dir == :o
symbols[i]
elseif dir == :r
symbols[n+mod1(i, n)]
elseif dir == :l
symbols[n+mod1(i - 1, n)]
else
throw(ArgumentError("Invalid direction: $dir"))
end
end
Tensor(array, inds)
end

sitemap = Dict(Site(i) => symbols[i] for i in 1:n)

Chain(Quantum(TensorNetwork(_tensors), sitemap), boundary)
end

function Chain(::Operator, boundary::Periodic, arrays::Vector{<:AbstractArray})
@assert all(==(4) ∘ ndims, arrays) "All arrays must have 3 dimensions"
function Chain(::Operator, boundary::Periodic, arrays::Vector{<:AbstractArray}; order = defaultorder(Chain, Operator()))
@assert all(==(4) ∘ ndims, arrays) "All arrays must have 4 dimensions"
issetequal(order, defaultorder(Chain, Operator())) ||
throw(ArgumentError("order must be a permutation of $(String.(defaultorder(Chain, Operator())))"))

n = length(arrays)
symbols = [nextindex() for _ in 1:3n]

_tensors = map(enumerate(arrays)) do (i, array)
Tensor(array, [symbols[i], symbols[i+n], symbols[2n+mod1(i - 1, n)], symbols[2n+mod1(i, n)]])
inds = map(order) do dir
if dir == :o
symbols[i]
elseif dir == :i
symbols[i+n]
elseif dir == :l
symbols[2n+mod1(i - 1, n)]
elseif dir == :r
symbols[2n+mod1(i, n)]
else
throw(ArgumentError("Invalid direction: $dir"))
end
end
Tensor(array, inds)
end

sitemap = Dict(Site(i) => symbols[i] for i in 1:n)
Expand All @@ -84,22 +130,39 @@ function Chain(::Operator, boundary::Periodic, arrays::Vector{<:AbstractArray})
Chain(Quantum(TensorNetwork(_tensors), sitemap), boundary)
end

function Chain(::Operator, boundary::Open, arrays::Vector{<:AbstractArray})
function Chain(::Operator, boundary::Open, arrays::Vector{<:AbstractArray}; order = defaultorder(Chain, Operator()))
@assert ndims(arrays[1]) == 3 "First array must have 3 dimensions"
@assert all(==(4) ∘ ndims, arrays[2:end-1]) "All arrays must have 4 dimensions"
@assert ndims(arrays[end]) == 3 "Last array must have 3 dimensions"
issetequal(order, defaultorder(Chain, Operator())) ||
throw(ArgumentError("order must be a permutation of $(String.(defaultorder(Chain, Operator())))"))

n = length(arrays)
symbols = [nextindex() for _ in 1:3n-1]

_tensors = map(enumerate(arrays)) do (i, array)
if i == 1
Tensor(array, [symbols[1], symbols[n+1], symbols[1+2n]])
_order = if i == 1
filter(x -> x != :l, order)
elseif i == n
Tensor(array, [symbols[n], symbols[2n], symbols[2n+mod1(n - 1, n)]])
filter(x -> x != :r, order)
else
Tensor(array, [symbols[i], symbols[i+n], symbols[2n+mod1(i - 1, n)], symbols[2n+mod1(i, n)]])
order
end
jofrevalles marked this conversation as resolved.
Show resolved Hide resolved

inds = map(_order) do dir
if dir == :o
symbols[i]
elseif dir == :i
symbols[i+n]
elseif dir == :l
symbols[2n+mod1(i - 1, n)]
elseif dir == :r
symbols[2n+mod1(i, n)]
else
throw(ArgumentError("Invalid direction: $dir"))
end
end
Tensor(array, inds)
end

sitemap = Dict(Site(i) => symbols[i] for i in 1:n)
Expand Down
193 changes: 162 additions & 31 deletions test/Ansatz/Chain_test.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,35 +1,166 @@
@testset "Chain ansatz" begin
qtn = Chain(State(), Periodic(), [rand(2, 4, 4) for _ in 1:3])
@test socket(qtn) == State()
@test ninputs(qtn) == 0
@test noutputs(qtn) == 3
@test issetequal(sites(qtn), [site"1", site"2", site"3"])
@test boundary(qtn) == Periodic()
@test leftindex(qtn, site"1") == rightindex(qtn, site"3") != nothing

qtn = Chain(State(), Open(), [rand(2, 2), rand(2, 2, 2), rand(2, 2)])
@test socket(qtn) == State()
@test ninputs(qtn) == 0
@test noutputs(qtn) == 3
@test issetequal(sites(qtn), [site"1", site"2", site"3"])
@test boundary(qtn) == Open()
@test leftindex(qtn, site"1") == rightindex(qtn, site"3") == nothing

qtn = Chain(Operator(), Periodic(), [rand(2, 2, 4, 4) for _ in 1:3])
@test socket(qtn) == Operator()
@test ninputs(qtn) == 3
@test noutputs(qtn) == 3
@test issetequal(sites(qtn), [site"1", site"2", site"3", site"1'", site"2'", site"3'"])
@test boundary(qtn) == Periodic()
@test leftindex(qtn, site"1") == rightindex(qtn, site"3") != nothing

qtn = Chain(Operator(), Open(), [rand(2, 2, 4), rand(2, 2, 4, 4), rand(2, 2, 4)])
@test socket(qtn) == Operator()
@test ninputs(qtn) == 3
@test noutputs(qtn) == 3
@test issetequal(sites(qtn), [site"1", site"2", site"3", site"1'", site"2'", site"3'"])
@test boundary(qtn) == Open()
@test leftindex(qtn, site"1") == rightindex(qtn, site"3") == nothing
@testset "Periodic boundary" begin
@testset "State" begin
qtn = Chain(State(), Periodic(), [rand(2, 4, 4) for _ in 1:3])
@test socket(qtn) == State()
@test ninputs(qtn) == 0
@test noutputs(qtn) == 3
@test issetequal(sites(qtn), [site"1", site"2", site"3"])
@test boundary(qtn) == Periodic()
@test leftindex(qtn, site"1") == rightindex(qtn, site"3") != nothing

arrays = [rand(2, 1, 4), rand(2, 4, 3), rand(2, 3, 1)]
qtn = Chain(State(), Periodic(), arrays) # Default order (:o, :l, :r)

@test size(tensors(qtn; at = Site(1))) == (2, 1, 4)
@test size(tensors(qtn; at = Site(2))) == (2, 4, 3)
@test size(tensors(qtn; at = Site(3))) == (2, 3, 1)

@test leftindex(qtn, Site(1)) == rightindex(qtn, Site(3))
@test leftindex(qtn, Site(2)) == rightindex(qtn, Site(1))
@test leftindex(qtn, Site(3)) == rightindex(qtn, Site(2))

arrays = [permutedims(array, (3, 1, 2)) for array in arrays] # now we have (:r, :o, :l)
qtn = Chain(State(), Periodic(), arrays, order = [:r, :o, :l])

@test size(tensors(qtn; at = Site(1))) == (4, 2, 1)
@test size(tensors(qtn; at = Site(2))) == (3, 2, 4)
@test size(tensors(qtn; at = Site(3))) == (1, 2, 3)

@test leftindex(qtn, Site(1)) == rightindex(qtn, Site(3))
@test leftindex(qtn, Site(2)) == rightindex(qtn, Site(1))
@test leftindex(qtn, Site(3)) == rightindex(qtn, Site(2))

for i in 1:nsites(qtn)
@test size(TensorNetwork(qtn), inds(qtn; at = Site(i))) == 2
end
end
jofrevalles marked this conversation as resolved.
Show resolved Hide resolved

@testset "Operator" begin
qtn = Chain(Operator(), Periodic(), [rand(2, 2, 4, 4) for _ in 1:3])
@test socket(qtn) == Operator()
@test ninputs(qtn) == 3
@test noutputs(qtn) == 3
@test issetequal(sites(qtn), [site"1", site"2", site"3", site"1'", site"2'", site"3'"])
@test boundary(qtn) == Periodic()
@test leftindex(qtn, site"1") == rightindex(qtn, site"3") != nothing

arrays = [rand(2, 4, 1, 3), rand(2, 4, 3, 6), rand(2, 4, 6, 1)] # Default order (:o, :i, :l, :r)
qtn = Chain(Operator(), Periodic(), arrays)

@test size(tensors(qtn; at = Site(1))) == (2, 4, 1, 3)
@test size(tensors(qtn; at = Site(2))) == (2, 4, 3, 6)
@test size(tensors(qtn; at = Site(3))) == (2, 4, 6, 1)

@test leftindex(qtn, Site(1)) == rightindex(qtn, Site(3))
@test leftindex(qtn, Site(2)) == rightindex(qtn, Site(1))
@test leftindex(qtn, Site(3)) == rightindex(qtn, Site(2))

for i in 1:length(arrays)
@test size(TensorNetwork(qtn), inds(qtn; at = Site(i))) == 2
@test size(TensorNetwork(qtn), inds(qtn; at = Site(i; dual = true))) == 4
end

arrays = [permutedims(array, (4, 1, 3, 2)) for array in arrays] # now we have (:r, :o, :l, :i)
qtn = Chain(Operator(), Periodic(), arrays, order = [:r, :o, :l, :i])

@test size(tensors(qtn; at = Site(1))) == (3, 2, 1, 4)
@test size(tensors(qtn; at = Site(2))) == (6, 2, 3, 4)
@test size(tensors(qtn; at = Site(3))) == (1, 2, 6, 4)

@test leftindex(qtn, Site(1)) == rightindex(qtn, Site(3)) !== nothing
@test leftindex(qtn, Site(2)) == rightindex(qtn, Site(1)) !== nothing
@test leftindex(qtn, Site(3)) == rightindex(qtn, Site(2)) !== nothing

for i in 1:length(arrays)
@test size(TensorNetwork(qtn), inds(qtn; at = Site(i))) == 2
@test size(TensorNetwork(qtn), inds(qtn; at = Site(i; dual = true))) == 4
end
end
end

@testset "Open boundary" begin
@testset "State" begin
qtn = Chain(State(), Open(), [rand(2, 2), rand(2, 2, 2), rand(2, 2)])
@test socket(qtn) == State()
@test ninputs(qtn) == 0
@test noutputs(qtn) == 3
@test issetequal(sites(qtn), [site"1", site"2", site"3"])
@test boundary(qtn) == Open()
@test leftindex(qtn, site"1") == rightindex(qtn, site"3") == nothing

arrays = [rand(2, 1), rand(2, 1, 3), rand(2, 3)]
qtn = Chain(State(), Open(), arrays) # Default order (:o, :l, :r)

@test size(tensors(qtn; at = Site(1))) == (2, 1)
@test size(tensors(qtn; at = Site(2))) == (2, 1, 3)
@test size(tensors(qtn; at = Site(3))) == (2, 3)

@test leftindex(qtn, Site(1)) == rightindex(qtn, Site(3)) === nothing
@test leftindex(qtn, Site(2)) == rightindex(qtn, Site(1))
@test leftindex(qtn, Site(3)) == rightindex(qtn, Site(2))

arrays = [permutedims(arrays[1], (2, 1)), permutedims(arrays[2], (3, 1, 2)), permutedims(arrays[3], (1, 2))] # now we have (:r, :o, :l)
qtn = Chain(State(), Open(), arrays, order = [:r, :o, :l])

@test size(tensors(qtn; at = Site(1))) == (1, 2)
@test size(tensors(qtn; at = Site(2))) == (3, 2, 1)
@test size(tensors(qtn; at = Site(3))) == (2, 3)

@test leftindex(qtn, Site(1)) == rightindex(qtn, Site(3)) === nothing
@test leftindex(qtn, Site(2)) == rightindex(qtn, Site(1)) !== nothing
@test leftindex(qtn, Site(3)) == rightindex(qtn, Site(2)) !== nothing

for i in 1:nsites(qtn)
@test size(TensorNetwork(qtn), inds(qtn; at = Site(i))) == 2
end
end
@testset "Operator" begin
qtn = Chain(Operator(), Open(), [rand(2, 2, 4), rand(2, 2, 4, 4), rand(2, 2, 4)])
@test socket(qtn) == Operator()
@test ninputs(qtn) == 3
@test noutputs(qtn) == 3
@test issetequal(sites(qtn), [site"1", site"2", site"3", site"1'", site"2'", site"3'"])
@test boundary(qtn) == Open()
@test leftindex(qtn, site"1") == rightindex(qtn, site"3") == nothing

arrays = [rand(2, 4, 1), rand(2, 4, 1, 3), rand(2, 4, 3)] # Default order (:o :i, :l, :r)
qtn = Chain(Operator(), Open(), arrays)

@test size(tensors(qtn; at = Site(1))) == (2, 4, 1)
@test size(tensors(qtn; at = Site(2))) == (2, 4, 1, 3)
@test size(tensors(qtn; at = Site(3))) == (2, 4, 3)

@test leftindex(qtn, Site(1)) == rightindex(qtn, Site(3)) === nothing
@test leftindex(qtn, Site(2)) == rightindex(qtn, Site(1)) !== nothing
@test leftindex(qtn, Site(3)) == rightindex(qtn, Site(2)) !== nothing

for i in 1:length(arrays)
@test size(TensorNetwork(qtn), inds(qtn; at = Site(i))) == 2
@test size(TensorNetwork(qtn), inds(qtn; at = Site(i; dual = true))) == 4
end

arrays = [
permutedims(arrays[1], (3, 1, 2)),
permutedims(arrays[2], (4, 1, 3, 2)),
permutedims(arrays[3], (1, 3, 2)),
] # now we have (:r, :o, :l, :i)
qtn = Chain(Operator(), Open(), arrays, order = [:r, :o, :l, :i])

@test size(tensors(qtn; at = Site(1))) == (1, 2, 4)
@test size(tensors(qtn; at = Site(2))) == (3, 2, 1, 4)
@test size(tensors(qtn; at = Site(3))) == (2, 3, 4)

@test leftindex(qtn, Site(1)) == rightindex(qtn, Site(3)) === nothing
@test leftindex(qtn, Site(2)) == rightindex(qtn, Site(1)) !== nothing
@test leftindex(qtn, Site(3)) == rightindex(qtn, Site(2)) !== nothing

for i in 1:length(arrays)
@test size(TensorNetwork(qtn), inds(qtn; at = Site(i))) == 2
@test size(TensorNetwork(qtn), inds(qtn; at = Site(i; dual = true))) == 4
end
end
end

@testset "Site" begin
using Qrochet: leftsite, rightsite
Expand Down
Loading