Support MPO-(MPS with dangling site indices) contraction

[shinaoka]

I made a fresh start after #44 was merged into the main branch.

For QTT, we may want to contract an MPO (with two site indices at each vertex) and an MPS with dangling site indices. I found several issues, but let me sout them out one by one. The first issue is the following.

A minimum example for naive algoritm works:

nbit = 3
sites = siteinds("Qubit", nbit)

H = randomMPO(sites) + randomMPO(sites)
H = replaceprime(H, 1=>2, 0=>1)

psi = randomMPO(sites) + randomMPO(sites)

# MPO-MPO
Hpsi_ref = contract(H, psi; alg="naive")

# MPO-MPS
Hpsi_ref2 = contract(H, MPS([A for A in psi]); alg="naive")

@show H
@show psi
@show Hpsi_ref
@show Hpsi_ref2
@assert MPS([A for A in Hpsi_ref]) ≈ Hpsi_ref2

H = MPO
[1] ((dim=2|id=58|"Qubit,Site,n=1")'', (dim=2|id=58|"Qubit,Site,n=1")', (dim=2|id=68|"Link,l=1")')
[2] ((dim=2|id=656|"Qubit,Site,n=2")'', (dim=2|id=656|"Qubit,Site,n=2")', (dim=2|id=121|"Link,l=2")', (dim=2|id=68|"Link,l=1")')
[3] ((dim=2|id=270|"Qubit,Site,n=3")'', (dim=2|id=270|"Qubit,Site,n=3")', (dim=2|id=121|"Link,l=2")')

psi = MPO
[1] ((dim=2|id=58|"Qubit,Site,n=1")', (dim=2|id=58|"Qubit,Site,n=1"), (dim=2|id=787|"Link,l=1"))
[2] ((dim=2|id=656|"Qubit,Site,n=2")', (dim=2|id=656|"Qubit,Site,n=2"), (dim=2|id=46|"Link,l=2"), (dim=2|id=787|"Link,l=1"))
[3] ((dim=2|id=270|"Qubit,Site,n=3")', (dim=2|id=270|"Qubit,Site,n=3"), (dim=2|id=46|"Link,l=2"))

Hpsi_ref = MPO
[1] ((dim=2|id=58|"Qubit,Site,n=1")'', (dim=2|id=58|"Qubit,Site,n=1"), (dim=4|id=716|"CMB,Link"))
[2] ((dim=2|id=656|"Qubit,Site,n=2")'', (dim=2|id=656|"Qubit,Site,n=2"), (dim=4|id=206|"CMB,Link"), (dim=4|id=716|"CMB,Link"))
[3] ((dim=2|id=270|"Qubit,Site,n=3")'', (dim=2|id=270|"Qubit,Site,n=3"), (dim=4|id=206|"CMB,Link"))

Hpsi_ref2 = MPS
[1] ((dim=2|id=58|"Qubit,Site,n=1")'', (dim=2|id=58|"Qubit,Site,n=1"), (dim=4|id=596|"CMB,Link"))
[2] ((dim=2|id=656|"Qubit,Site,n=2")'', (dim=2|id=656|"Qubit,Site,n=2"), (dim=4|id=816|"CMB,Link"), (dim=4|id=596|"CMB,Link"))
[3] ((dim=2|id=270|"Qubit,Site,n=3")'', (dim=2|id=270|"Qubit,Site,n=3"), (dim=4|id=816|"CMB,Link"))

fitting algorithm does not work:

# MPO-MPS
contract(H, MPS([A for A in psi]); alg="fit")

BoundsError: attempt to access 1-element Vector{Index{Int64}} at index [2]

Stacktrace:
  [1] getindex
    @ ./array.jl:924 [inlined]
  [2] replaceinds(is::Tuple{Index{Int64}, Index{Int64}, Index{Int64}}, inds1::Vector{Index{Int64}}, inds2::Vector{Index{Int64}})
    @ ITensors /sharehome/shinaoka/.julia/packages/ITensors/fpBnt/src/indexset.jl:570
  [3] replaceinds(::NDTensors.DenseTensor{Float64, 3, Tuple{Index{Int64}, Index{Int64}, Index{Int64}}, NDTensors.Dense{Float64, Vector{Float64}}}, ::Vector{Index{Int64}}, ::Vararg{Vector{Index{Int64}}}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ ITensors /sharehome/shinaoka/.julia/packages/ITensors/fpBnt/src/itensor.jl:1480
  [4] replaceinds(::NDTensors.DenseTensor{Float64, 3, Tuple{Index{Int64}, Index{Int64}, Index{Int64}}, NDTensors.Dense{Float64, Vector{Float64}}}, ::Vector{Index{Int64}}, ::Vector{Index{Int64}})
    @ ITensors /sharehome/shinaoka/.julia/packages/ITensors/fpBnt/src/itensor.jl:1479
  [5] replaceinds!(::ITensor, ::Vector{Index{Int64}}, ::Vararg{Vector{Index{Int64}}}; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ ITensors /sharehome/shinaoka/.julia/packages/ITensors/fpBnt/src/itensor.jl:1484
  [6] replaceinds!(::ITensor, ::Vector{Index{Int64}}, ::Vector{Index{Int64}})
    @ ITensors /sharehome/shinaoka/.julia/packages/ITensors/fpBnt/src/itensor.jl:1483
  [7] contract(::ITensors.Algorithm{:fit}, A::MPO, psi0::MPS; init_state::MPS, nsweeps::Int64, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ ITensorNetworks /sharehome/shinaoka/git/ITensorNetworks.jl/src/treetensornetworks/solvers/contract_operator_state.jl:39
  [8] contract
    @ /sharehome/shinaoka/git/ITensorNetworks.jl/src/treetensornetworks/solvers/contract_operator_state.jl:13 [inlined]
  [9] contract(A::MPO, ψ::MPS; alg::String, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}})
    @ ITensors /sharehome/shinaoka/.julia/packages/ITensors/fpBnt/src/mps/mpo.jl:567
 [10] top-level scope
    @ In[14]:2

The replacement of site indices does not work with an MPS with dangling site indices. Shall we just remove this part?
https://github.com/mtfishman/ITensorNetworks.jl/blob/main/src/treetensornetworks/solvers/contract_operator_state.jl#L38-L42

If we move those lines, it will be user's responsibility to provide an appropriate init_state with the same site indices as H * psi0. The following default value assumes that init_state lives in the same space as psi0, which is the case if H is a time-evoluation operator. For a more general MPO, this default value is not appropriate. What do you think about removing these lines for the replacement and adding a check on the site indces on init_state?

https://github.com/mtfishman/ITensorNetworks.jl/blob/main/src/treetensornetworks/solvers/contract_operator_state.jl#L17

[mtfishman]

Yes, I think we should remove any automated index replacement and leave it up to the user to provide a proper init_state.

For example, it doesn't make sense to try to replace the indices if the matrix A in Ax = b is rectangular, in which case x and b fundamentally have different spaces. Additionally, if the MPO and MPS have multiple sites per index, that index replacement isn't even well defined.

EDIT: Sorry, I lost track and thought we were talking about linsolve functionality. However, my previous comment still stands, that if we want users to provide a guess C0 for A * B ≈ C we should require them to provide a C with the correct indices, since it won't always be possible to map incorrect inputs to the space of A * B (for example if A is a rectangular, in the sense that the spaces of the bra and ket are not even the same size).

[mtfishman]

Finally, in the signature of the function I don't think we should specialize A::IsTreeOperator and psi0::IsTreeState.

Instead I think it should more generally be defined as:

function ITensors.contract(
  ::ITensors.Algorithm"fit",
  A::IsTree,
  B::IsTree;
  init_guess=[...], # Determine an initial guess by analyzing the external indices of `A * B`
  kwargs...,
)
  [...]
end

where:

const IsTree = Union{MPS,MPO,TTNS,TTNO}

to make it as general as possible. Going forward, we might make this into a trait, and also just get ride of TTNS and TTNO in favor of a single type TTN, since I think there is generally not much reason to have different types, at least based on how they are designed right now.

EDIT: In #47 I went ahead and removed the TTNS and TTNO types in favor of just TTN, and also defined IsTree = Union{MPS,MPO,TTN} here.

[shinaoka]

Agreed with your comments & suggestions. Should we set a default value to init_guess by automatically analyzing the site indices (and creating a random TNN object similar to A * B with some bond dimensions)? This may be possible but it sounds a bit complicated (especially how to determine appropriate bond dimensions).

Candidate 1 is backward compatible. Candidate 2 is explicit. In addition, we could write a separate function for generating a random initial guess with a given global bond dimension from A and B.

Candidate 1:

function ITensors.contract(
  ::ITensors.Algorithm"fit",
  A::IsTree,
  B::IsTree;
  init_guess=B, # Error is thrown if B does not have the same idices as `A * B`.
  kwargs...,
)
  [...]
end

Candidate 2:

function ITensors.contract(
  ::ITensors.Algorithm"fit",
  A::IsTree,
  B::IsTree;
  init_guess, # The user must specify explicitly
  kwargs...,
)
  [...]
end

[mtfishman]

@shinaoka I don't think using init_guess=B makes sense since it won't have the same external indices as A * B.

I don't think there is any particular challenge to analyzing A and B and determining the proper external indices for the resulting network (related to #40). However, it is tricky creating random initial guesses, particularly for determining the QNs of the link indices if there are symmetric tensors. For DMRG, we use randomMPS for that, which internally uses a random circuit to create a random state of some finite bond dimension.

I think like for DMRG we should have the users specify an initial guess unless we come up with something good that generally works. I know this algorithm has a tendency to get stuck in local minima so the initial guess is particular important. Something to consider would be to use a cheaper algorithm for contracting the TTNs approximately, such as the zip-up algorithm. But we would have to test that out to decide if it is a good default in general.

[mtfishman]

Also, we've decided that for now we won't explicitly support accepting MPS/MPO types in the functions in ITensorNetworks.jl (I'll change this throughout the package), so the signature should be:

function ITensors.contract(
  ::ITensors.Algorithm"fit",
  A::AbstractTTN,
  B::AbstractTTN;
  init_guess,
  kwargs...,
)
  # [...]
end

Users can convert their MPS/MPO to TTN if they want to use the new functions available in this package.

[shinaoka]

Thank you for letting me know about the decision! I will test the code once you make changes throughout the package.

[mtfishman]

@shinaoka in #50 I changed this contract function to have the signature:

function contract(
  ::Algorithm"fit",
  tn1::AbstractTTN,
  tn2::AbstractTTN;
  init,
  kwargs...,
)
  # [...]
end

which helps us move towards not having a distinction between states and operators.

Unfortunately that function now seems to be broken, which I think happened when we originally generalized the solver functions from MPO/MPS to TTNs in #44.

If you have time, it would be very helpful if you could investigate why it isn't working!

[shinaoka]

Thank you!

It seems to be working on my side. Do you have any example code?

nbit = 3
sites = siteinds("Qubit", nbit)

M1 = replaceprime(randomMPO(sites) + randomMPO(sites), 1=>2, 0=>1)
M2 = randomMPO(sites) + randomMPO(sites)
M12_ref = contract(M1, M2; alg="naive")
t12_ref = TreeTensorNetwork([M12_ref[v] for v in eachindex(M12_ref)])

t1 = TreeTensorNetwork([M1[v] for v in eachindex(M1)])
t2 = TreeTensorNetwork([M2[v] for v in eachindex(M2)])
t12 = contract(t1, t2; alg="fit", init=t12_ref)
#t12 = contract(t1, t2; alg="fit")

function toarr(t)
    arr = reduce(*, [t[v] for v in vertices(t)])
    return Array(arr, vcat(sites, sites''))
end

@show t12 ≈ t12_ref
@show toarr(t12) ≈ toarr(t12_ref)

[leburgel]

Turns out there were some issues in the tests due to a bug in inner and changes in the default initial state for contract (#52), but the function itself should be fine I think.

[shinaoka]

I see. So, shall we add the above code to the test?

[mtfishman]

@shinaoka that seems like a reasonable test to add. Does it seem like the current contract(::TTN, ::TTN) code addresses this original issue (i.e. can we close this issue)?