# Help with type instabilities

**URL:** <https://discourse.julialang.org/t/help-with-type-instabilities/3897>\
**Category:** General Usage\
**Created:** [May 24, 2017, 3:21am UTC](https://discourse.julialang.org/t/help-with-type-instabilities/3897 "2017-05-24T03:21:53Z")\
**Posts on this page:** 10\
**Page:** 1

<div class="post-metadata">

**Author:** ![jebej](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jebej/32/1784_2.png) [@jebej](https://discourse.julialang.org/u/jebej)\
**Post date:** [May 24, 2017, 3:21am UTC](https://discourse.julialang.org/t/help-with-type-instabilities/3897/1 "2017-05-24T03:21:53Z")

</div>

I’m having issues removing type instabilities from the function below, and would appreciate some tips.

```julia
A = rand(36,36);
out = [2]; sysdims = [2,9,2];
@code_warntype ptrace2(A,out,sysdims)

```

```julia
function ptrace2{T<:Number}(A::Matrix{T},out,sysdims)
    nsys = length(sysdims)
    keep = setdiff(1:nsys,out)
    keepdims = sysdims[keep]
    outdims = sysdims[out]::Vector{Int}
    N = prod(keepdims)
    R = Base.product(range.(1,keepdims)...)
    S = Base.product(range.(1,outdims)...)
    B = zeros(T,N,N)
    inds_ii = Vector{Int}(nsys)
    inds_jj = Vector{Int}(nsys)
    for r in R, q in R
        i, j = tindex(q,keepdims), tindex(r,keepdims)
        for l=1:length(keep)
            inds_ii[keep[l]] = q[l]
            inds_jj[keep[l]] = r[l]
        end
        for s in S
            for l=1:length(out)
                inds_ii[out[l]] = s[l]
                inds_jj[out[l]] = s[l]
            end
            ii = tindex(inds_ii,sysdims)
            jj = tindex(inds_jj,sysdims)
            B[i,j] += A[ii,jj]
        end
    end
    return B
end

function tindex(inds,sysdims)
    N = length(inds)
    i = inds[N]
    d = 1
    for n = N-1:-1:1
        d *= sysdims[n+1]
        i += d * (inds[n]-1)
    end
    return i
end

```

---

<div class="post-metadata">

**Author:** ![traktofon](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/traktofon/32/591_2.png) [@traktofon](https://discourse.julialang.org/u/traktofon)\
**Post date:** [May 24, 2017, 4:25am UTC](https://discourse.julialang.org/t/help-with-type-instabilities/3897/2 "2017-05-24T04:25:05Z")

</div>

In this case a function barrier can help. Here is my attempt (no change of the code except splitting the function into two parts):

```julia
function ptrace2{T<:Number}(A::Matrix{T},out,sysdims)
    nsys = length(sysdims)
    keep = setdiff(1:nsys,out)
    keepdims = sysdims[keep]
    outdims = sysdims[out]::Vector{Int}
    N = prod(keepdims)
    R = Base.product(range.(1,keepdims)...)
    S = Base.product(range.(1,outdims)...)
    B = zeros(T,N,N)
    _ptrace2!(B,A,R,S,nsys,sysdims,keep,keepdims,out)
end

function _ptrace2!(B,A,R,S,nsys,sysdims,keep,keepdims,out)
    inds_ii = Vector{Int}(nsys)
    inds_jj = Vector{Int}(nsys)
    for r in R, q in R
        i, j = tindex(q,keepdims), tindex(r,keepdims)
        for l=1:length(keep)
            inds_ii[keep[l]] = q[l]
            inds_jj[keep[l]] = r[l]
        end
        for s in S
            for l=1:length(out)
                inds_ii[out[l]] = s[l]
                inds_jj[out[l]] = s[l]
            end
            ii = tindex(inds_ii,sysdims)
            jj = tindex(inds_jj,sysdims)
            B[i,j] += A[ii,jj]
        end
    end
    return B
end

```

Using your example data, I get (Julia 0.6.0-rc2) with the barrier:

> julia\> @btime ptrace2($A, $out, $sysdims);  
> 5.268 μs (25 allocations: 2.09 KiB)

while without the barrier (original code):

> julia\> @btime ptrace2($A, $out, $sysdims);  
> 57.671 μs (810 allocations: 18.64 KiB)

---

<div class="post-metadata">

**Author:** ![smldis](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/smldis/32/1735_2.png) [@smldis](https://discourse.julialang.org/u/smldis)\
**Post date:** [May 24, 2017, 6:31am UTC](https://discourse.julialang.org/t/help-with-type-instabilities/3897/3 "2017-05-24T06:31:03Z")

</div>

It seems to me that `R = Base.product(range.(1,keepdims)...)` is not inferred since the type of `Base.product(range.(1,keepdims)...)` depends on the expansion, in your casethe type inferred is `Base.Iterators.Prod2{UnitRange{Int64},UnitRange{Int64}}`.

I think splitting the function is ok so you have a dispatch on the type of R. Maybe someone can give a better solution with a more generic iterator?

---

<div class="post-metadata">

**Author:** ![jebej](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jebej/32/1784_2.png) [@jebej](https://discourse.julialang.org/u/jebej)\
**Post date:** [May 24, 2017, 4:29pm UTC](https://discourse.julialang.org/t/help-with-type-instabilities/3897/4 "2017-05-24T16:29:57Z")

</div>

Thanks!

---

<div class="post-metadata">

**Author:** ![rdeits](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/rdeits/32/286_2.png) [@rdeits](https://discourse.julialang.org/u/rdeits)\
**Post date:** [May 24, 2017, 6:03pm UTC](https://discourse.julialang.org/t/help-with-type-instabilities/3897/5 "2017-05-24T18:03:13Z")

</div>

Just as an aside, instead of `Base.product(range.(1, keepdims)...)` you might consider using a `CartesianRange(keepdims...)` which is designed for the kind of multidimensional indexing that you’re doing.

Unfortunately, `CartesianRange` also relies on knowing the number of dimensions for type-inference, so it doesn’t solve your immediate issue. But it might at least make your code more understandable and perhaps more performant.

---

<div class="post-metadata">

**Author:** ![jebej](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jebej/32/1784_2.png) [@jebej](https://discourse.julialang.org/u/jebej)\
**Post date:** [May 24, 2017, 8:31pm UTC](https://discourse.julialang.org/t/help-with-type-instabilities/3897/6 "2017-05-24T20:31:23Z")

</div>

I tried using `CartesianRange`, but it actually made the function slower…

In any case, this is what I have now, after some more clever permutation of indices, if anyone wants to take a stab at making it faster:

```julia
julia> A = rand(36,36); out = (2,); sysdims = (2,9,2);
julia> @btime ptrace3($A,$out,$sysdims)
  4.221 μs (35 allocations: 2.63 KiB)

```

```julia
function ptrace3{T,M,nsub}(A::Matrix{T}, out::NTuple{M,Int}, sysdims::NTuple{nsub,Int})
    rkeep = (nsub + 1) .- reverse(setdiff(1:nsub,out))
    rout = (nsub + 1) .- reverse([out...])
    rsysdims = reverse(sysdims)
    rkeepdims = rsysdims[rkeep]::NTuple{nsub-M,Int}
    routdims = rsysdims[rout]::NTuple{M,Int}
    N = prod(rkeepdims)::Int
    R = product(range.(1,rkeepdims)...)
    S = product(range.(1,routdims)...)
    B = zeros(T,N,N)
    _ptrace3!(B,A,R,S,nsub,rsysdims,rkeepdims,rkeep,rout)
end

function _ptrace3!(B,A,R,S,nsub,rsysdims,rkeepdims,rkeep,rout)
    K = length(rkeepdims)
    O = nsub - K
    rinds_ii = Vector{Int}(nsub)
    rinds_jj = Vector{Int}(nsub)
    for r in R, q in R
        i, j = tindex2(q,rkeepdims), tindex2(r,rkeepdims)
        #println("i: $i, j: $j")
        for k=1:K
            rinds_ii[rkeep[k]] = q[k]
            rinds_jj[rkeep[k]] = r[k]
        end
        for s in S
            for o=1:O
                rinds_ii[rout[o]] = s[o]
                rinds_jj[rout[o]] = s[o]
            end
            ii, jj = tindex2(rinds_ii,rsysdims), tindex2(rinds_jj,rsysdims)
            #println("ii: $ii, jj: $jj")
            B[i,j] += A[ii,jj]
        end
    end
    return B
end

function tindex2{N}(rinds,rsysdims::NTuple{N,Int})
    i = rinds[1]
    d = 1
    @inbounds for n = 2:N
        d *= rsysdims[n-1]
        i += d * (rinds[n]-1)
    end
    return i
end

```

---

<div class="post-metadata">

**Author:** ![Dan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dan/32/42581_2.png) [@Dan](https://discourse.julialang.org/u/Dan)\
**Post date:** [May 27, 2017, 1:45pm UTC](https://discourse.julialang.org/t/help-with-type-instabilities/3897/7 "2017-05-27T13:45:13Z")

</div>

There is a package called `TensorOperations` which does this kind of calculation and many more. It is optimized and given the simplicity of code, it would probably save a lot of time down the line and allow someone else to worry about further optimizations.

The code for `ptrace` would look like this:

```julia
using BenchmarkTools
using TensorOperations

A = rand(2,9,2,2,9,2);

julia> @btime tensortrace(A,[:x1,:t,:z1,:x2,:t,:z2])
  5.668 μs (65 allocations: 5.20 KiB)
2×2×2×2 Array{Float64,4}:
[:, :, 1, 1] =
 5.1795 5.58526
 4.18046 6.27735

[:, :, 2, 1] =
 5.70867 5.40797
 6.11036 4.07155

[:, :, 1, 2] =
 5.26413 5.49429
 4.90246 3.99036

[:, :, 2, 2] =
 5.26287 3.62941
 6.27622 4.9152 

```

If a 2D matrix is really necessary (not really), then a helper functions would do the job:

```julia
function mat(A::Array)
    s,d = size(A),ndims(A)
    s[1:d>>>1] != s[(1:d>>>1) + (d>>>1)] && error("tensor size cannot be matrixfied: $s")
    n = prod(s[1:d>>>1])
    return reshape(A,n,n)
end

julia> @btime mat(tensortrace(A,[:x1,:t,:z1,:x2,:t,:z2]))
  5.924 μs (70 allocations: 5.39 KiB)
4×4 Array{Float64,2}:
 5.1795 5.70867 5.26413 5.26287
 4.18046 6.11036 4.90246 6.27622
 5.58526 5.40797 5.49429 3.62941
 6.27735 4.07155 3.99036 4.9152 

```

Note the timings/allocations are very similar.

---

<div class="post-metadata">

**Author:** ![jebej](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/jebej/32/1784_2.png) [@jebej](https://discourse.julialang.org/u/jebej)\
**Post date:** [May 27, 2017, 1:54pm UTC](https://discourse.julialang.org/t/help-with-type-instabilities/3897/8 "2017-05-27T13:54:15Z")

</div>

Hey thanks for the reply. I had indeed looked at TensorOperations, but I do need everything in 2D matrix form. I improved the code further (with some [help](https://discourse.julialang.org/t/type-stable-difference-of-tuples/3933/6)) and managed to get the partial trace even faster:

```julia
julia> using Schrodinger, BenchmarkTools
julia> A = rand(36,36); out = (2,); sysdims = (2,9,2);
julia> @benchmark ptrace($A,$out,$sysdims)
BenchmarkTools.Trial:
  memory estimate: 480 bytes
  allocs estimate: 5
  --------------
  minimum time: 1.275 μs (0.00% GC)
  median time: 1.431 μs (0.00% GC)
  mean time: 1.520 μs (2.17% GC)
  maximum time: 181.442 μs (97.48% GC)
  --------------
  samples: 10000
  evals/sample: 10

```

The code is now available [here](https://github.com/jebej/Schrodinger.jl/blob/master/src/math/ptrace.jl) (the package itself is not fully ready for prime time) for those interested.

---

<div class="post-metadata">

**Author:** ![Dan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dan/32/42581_2.png) [@Dan](https://discourse.julialang.org/u/Dan)\
**Post date:** [May 27, 2017, 3:05pm UTC](https://discourse.julialang.org/t/help-with-type-instabilities/3897/9 "2017-05-27T15:05:09Z")

</div>

Depending on the size of the matrices involved, it might be better to go with TensorOperations. Small tests can obscure the coefficients of the running-time surface. In particular, we have the following for TensorOperations vs. Schrodinger implementation:

A comparable TensorOperations implementation with wrapping functions to make interface like `ptrace`’s:

```julia
using TensorOperations
using BenchmarkTools

A = rand(4000,4000); out = (2,); sysdims = (10,40,10);

function tens(A::Matrix,sysdims)
    size(A)==(prod(sysdims),prod(sysdims)) || error("sysdims product must equal row and col number")
    return reshape(A,sysdims...,sysdims...)
end

function mat(A::Array)
    s,d = size(A),ndims(A)
    s[1:d>>>1] != s[(1:d>>>1) + (d>>>1)] && error("tensor size cannot be matrixfied: $s")
    n = prod(s[1:d>>>1])
    return reshape(A,n,n)
end

function traceout(A,out)
    inds = collect(1:ndims(A))
    inds[out+(ndims(A)>>>1)] = inds[out]
    return mat(tensortrace(A,inds))
end

ptrace(A,out,sysdims) = traceout(reshape(A,sysdims...,sysdims...),collect(out))

```

With benchmark:

```julia
julia> @benchmark ptrace($A,$out,$sysdims)
BenchmarkTools.Trial: 
  memory estimate: 83.69 KiB
  allocs estimate: 72
  --------------
  minimum time: 2.559 ms (0.00% GC)
  median time: 2.712 ms (0.00% GC)
  mean time: 2.756 ms (0.11% GC)
  maximum time: 5.178 ms (17.54% GC)
  --------------
  samples: 1810
  evals/sample: 1

```

Compared to ptrace from Schrodinger.jl:

```julia
A = rand(4000,4000); out = (2,); sysdims = (10,40,10);

julia> @benchmark ptrace(A,out,sysdims)
BenchmarkTools.Trial: 
  memory estimate: 78.47 KiB
  allocs estimate: 6
  --------------
  minimum time: 4.812 ms (0.00% GC)
  median time: 5.410 ms (0.00% GC)
  mean time: 5.461 ms (0.03% GC)
  maximum time: 10.425 ms (0.00% GC)
  --------------
  samples: 914
  evals/sample: 1

```

---

<div class="post-metadata">

**Author:** ![Dan](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/dan/32/42581_2.png) [@Dan](https://discourse.julialang.org/u/Dan)\
**Post date:** [May 27, 2017, 3:18pm UTC](https://discourse.julialang.org/t/help-with-type-instabilities/3897/10 "2017-05-27T15:18:57Z")

</div>

If the helper functions are confusing, here is a more compact version (with less parameter checking):

```julia
using TensorOperations
using BenchmarkTools

A = rand(4000,4000); out = (2,); sysdims = (10,40,10);

mat(A::Array) = ( n = prod(size(A)[1:ndims(A)>>>1]) ; reshape(A,n,n) )

tens(M::Matrix,sysdims) = reshape(M,sysdims...,sysdims...)

traceout(A,out) = ( n = ndims(A)>>>1 ; mat(tensortrace(A,[i-n in out ? i-n : i for i=1:2n])) )

ptrace(M,out,sysdims) = traceout(tens(M,sysdims),out)

@benchmark ptrace($A,$out,$sysdims)

```