How to update a tensor element-wise in ITensors ?

I am trying to write an iterative solver for a non-linear equation that requires computing a large tensor for residuals R2 and using R2 to update amplitudes T2 , then those new amplitudes are further used to calculate new residuals and the process is repeated till the norm R2 is below a selected threshold. Till now, I had been using TensorOperations.jl where the update of the T2 tensor looked as follows :

function update_T2(T2::Array{Float64,4},R2::Array{Float64,4},Scaled_R2::Array{Float64,4})
    fvv = deserialize("fvv.jlbin")
    foo = deserialize("foo.jlbin")
    nv = deserialize("nv.jlbin")
    nocc = deserialize("nocc.jlbin")
    shiftp = 0.20
    for a in 1:nv, b in 1:nv, i in 1:nocc, j in 1:nocc
        Scaled_R2[a,b,i,j] = R2[a,b,i,j]/(fvv[a,a] + fvv[b,b]-foo[i,i] - foo[j,j]+shiftp)
    end
    for a in 1:nv, b in 1:nv, i in 1:nocc, j in 1:nocc
        T2[a,b,i,j] = T2[a,b,i,j] - Scaled_R2[a,b,i,j]
    end
    return T2
end

I want to update T2 the same way but now T2 has the data type ::ITensor. So I tried to write the following code

function update_T2(T2::ITensor,R2::ITensor,Scaled_R2::ITensor)
    fvv = ITensor(deserialize("fvv.jlbin"),a_1,a_2)
    foo = ITensor(deserialize("foo.jlbin"),i_1,i_2)
    nv = deserialize("nv.jlbin")::Int64
    nocc = deserialize("nocc.jlbin")::Int64
    shiftp = 0.20
    for a in 1:nv, b in 1:nv, i in 1:nocc, j in 1:nocc
        Scaled_R2[a,b,i,j] = R2[a,b,i,j]/(fvv[a,a] + fvv[b,b]-foo[i,i] - foo[j,j]+shiftp)
    end
    for a in 1:nv, b in 1:nv, i in 1:nocc, j in 1:nocc
        T2[a,b,i,j] = T2[a,b,i,j] - Scaled_R2[a,b,i,j]
    end
    return T2
end

Where a_1,a_2,i_1,i_2 are the indices of the tensors T2 Scaled_R2 and R2. The ITensor version seems to return a different value of T2 even when given the same initial guess of T2 and other parameters. Is there any problem in the way I am updating the tensor element-wise ?

Are you sure the indices of T2, R2, and Scaled_R2 are lined up correctly? If they are not ordered in the same way, your code will be incorrect. As suggested in your other post, I would try to rewrite your code taking advantage of the index labels of ITensors rather than naively translating your code.

As an example, the second for-loop you wrote:

for a in 1:nv, b in 1:nv, i in 1:nocc, j in 1:nocc
  T2[a,b,i,j] = T2[a,b,i,j] - Scaled_R2[a,b,i,j]
end

could just be written with broadcasting as T2 .= T2 .- Scaled_R2. By design of the ITensor type, writing for-loops over the elements of an ITensor is slower than the equivalent for-loop over Julia arrays so it is best to avoid that and use operations like broadcasting when possible.

I don’t see an obvious way to rewrite the first for-loop you wrote, it may be best to keep that in terms of arrays anyway and then convert to ITensors. However, again note that you need to make sure the indices of the ITensors area aligned properly, otherwise your indexing would be incorrect. I would recommend choosing a simple example and printing out the indices and elements of the tensors in the operations you are performing to try to debug any issues you are having.

Generally when indexing into ITensors you should also specify the Index as well to make sure the code is robust against permutations of the dimensions:

julia> using ITensors: Index, ITensor, permute

julia> i, j = Index.((2, 2))
((dim=2|id=953), (dim=2|id=176))

julia> a = ITensor([0.0 1.0; 2.0 0.0], i, j);

julia> @show a;
a = ITensor ord=2
Dim 1: (dim=2|id=953)
Dim 2: (dim=2|id=176)
NDTensors.Dense{Float64, Vector{Float64}}
 2×2
 0.0  1.0
 2.0  0.0

julia> a[i => 1, j => 2]
1.0

julia> a[1, 2]
1.0

julia> a = permute(a, j, i)
ITensor ord=2 (dim=2|id=176) (dim=2|id=953)
NDTensors.Dense{Float64, Vector{Float64}}

julia> a[i => 1, j => 2]
1.0

julia> a[1, 2]
2.0

I can’t explain the reason behind:

since your code does appear to be updating T2 in-place, though as I said above maybe in a way that is both slow and bug-prone. But also you haven’t shown the full code so we don’t know how T2 or the other input tensors are defined. You are clearly showing code snippets from a larger code and trying to debug them in isolation, which can be challenging for anyone (in particular us since we don’t see the rest of the code). A suggestion I would have in order to debug your code is to write a minimal example that appears to reproduce the same kind of issue you are seeing, say with small tensors with contrived values, and see if you can reproduce your guess for what might be going on. That is a technique I use when debugging code, it could be fruitful for you and also help you learn more about how ITensor works in the process.

1 Like

Again, as a broader point I would suggest trying to rewrite your code more broadly in terms of ITensor, taking advantage of the ITensor Index system, rather than trying to rewrite your code in snippets, which can lead to these kinds of challenges.

Apologies for the extremely late reply. As you rightly pointed out, the second loop could be substituted with a broadcasting statement, whereas the first one cannot be (simply at least). I tried replacing the second loop by the above to give the resulting code

function update_T2(T2::ITensor,R2::ITensor,Scaled_R2::ITensor)
    fvv = ITensor(deserialize("fvv.jlbin"),a_1,a_2)
    foo = ITensor(deserialize("foo.jlbin"),i_1,i_2)
    nv = deserialize("nv.jlbin")::Int64
    nocc = deserialize("nocc.jlbin")::Int64
    shiftp = 0.20
    for a in 1:nv, b in 1:nv, i in 1:nocc, j in 1:nocc
        Scaled_R2[a,b,i,j] = R2[a,b,i,j]/(fvv[a,a] + fvv[b,b]-foo[i,i] - foo[j,j]+shiftp)
    end
    # for a in 1:nv, b in 1:nv, i in 1:nocc, j in 1:nocc
    #     T2[a,b,i,j] = T2[a,b,i,j] - Scaled_R2[a,b,i,j]
    # end
    T2 .= T2 .- Scaled_R2
    return T2
end

But the same problem still persists. I understand that looping over ITensors is slower than broadcasting. But is the way I have tried to modify the tensor element wise wrong ?

Here is a minimal rewrite of your code that demonstrates it should be updating T2:

using ITensors: ITensor, Index, random_itensor
using Random: Random
function main()
  Random.seed!(1234)
  a_1, a_2, i_1, i_2 = Index.((2, 2, 2, 2))
  R2 = random_itensor(a_1, a_2, i_1, i_2)
  T2 = random_itensor(a_1, a_2, i_1, i_2)
  Scaled_R2 = random_itensor(a_1, a_2, i_1, i_2)
  function update_T2!(T2::ITensor, R2::ITensor, Scaled_R2::ITensor)
    fvv = random_itensor(a_1, a_2)
    foo = random_itensor(i_1, i_2)
    nv = 2
    nocc = 2
    shiftp = 0.20
    for a in 1:nv, b in 1:nv, i in 1:nocc, j in 1:nocc
      numerator = R2[a_1 => a, a_2 => b, i_1 => i, i_2 => j]
      denominator = fvv[a_1 => a, a_2 => a] + fvv[a_1 => b, a_2 => b] - foo[i_1 => i, i_2 => i] - foo[i_1 => j, i_2 => j] + shiftp
      Scaled_R2[a_1 => a, a_2 => b, i_1 => i, i_2 => j] = numerator / denominator
    end
    T2 .-= Scaled_R2
    return T2
  end
  @show T2
  update_T2!(T2, R2, Scaled_R2)
  @show T2
  return nothing
end

When I run main() I get the output:

julia> main()
T2 = ITensor ord=4
Dim 1: (dim=2|id=592)
Dim 2: (dim=2|id=579)
Dim 3: (dim=2|id=308)
Dim 4: (dim=2|id=949)
NDTensors.Dense{Float64, Vector{Float64}}
 2×2×2×2
[:, :, 1, 1] =
  1.6218076699653694  -0.5347089043694231
 -0.7636889550217101  -0.8371161519624881

[:, :, 2, 1] =
 -0.09989094954741905  -0.240570071975397
 -0.726142325634263     0.9427328405545989

[:, :, 1, 2] =
 0.44588457807333104  -1.4451163776340619
 0.4604268914916924    0.3716847353731026

[:, :, 2, 2] =
 1.9522884266156713  -0.7578371325188665
 0.3237050517386015   1.1693426324343965
T2 = ITensor ord=4
Dim 1: (dim=2|id=592)
Dim 2: (dim=2|id=579)
Dim 3: (dim=2|id=308)
Dim 4: (dim=2|id=949)
NDTensors.Dense{Float64, Vector{Float64}}
 2×2×2×2
[:, :, 1, 1] =
 2.212330480674426   -2.2440191884987355
 1.0922381970413777  -0.82496215138259

[:, :, 2, 1] =
 -0.3416953884869767  8.401621007271215
 -5.33919777706203    0.12228179003341455

[:, :, 1, 2] =
  0.7516874457199619   0.8084381798095097
 -2.353150269542973   -0.21578545586345294

[:, :, 2, 2] =
 2.2143522987879845  0.07399335046884059
 0.3979494267800609  0.276824871640005

Note that I am indexing into the ITensors in the loop using the ITensor Indices, which is better practice since otherwise the code can become incorrect if ITensors are input with Indices that are ordered in a different way. I also shortened T2 .= T2 .- Scaled_R2 to T2 .-= Scaled_R2, those are equivalent to each other.

Also functions that change object in-place in Julia have a convention that their name ends in !.

So the issue you are seeing is probably coming from a different part of your code.

1 Like

This topic was automatically closed 10 days after the last reply. New replies are no longer allowed.