ptrace is type instable

[Lightup1]

julia> @code_warntype ptrace(dm(ψ₀),1)
MethodInstance for QuantumOpticsBase.ptrace(::Operator{CompositeBasis{Vector{Int64}, Tuple{SubspaceBasis{Vector{Int64}, PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Ket{PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Vector{ComplexF64}}, 0x9711e39bc569856c, UInt64}, SpinBasis{1//2, Int64}}}, CompositeBasis{Vector{Int64}, Tuple{SubspaceBasis{Vector{Int64}, PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Ket{PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Vector{ComplexF64}}, 0x9711e39bc569856c, UInt64}, SpinBasis{1//2, Int64}}}, Matrix{ComplexF64}}, ::Int64)
  from ptrace(a::DataOperator, indices) in QuantumOpticsBase at C:\Users\dell\.julia\packages\QuantumOpticsBase\bHpyR\src\operators_dense.jl:169
Arguments
  #self#::Core.Const(QuantumOpticsBase.ptrace)
  a::Operator{CompositeBasis{Vector{Int64}, Tuple{SubspaceBasis{Vector{Int64}, PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Ket{PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Vector{ComplexF64}}, 0x9711e39bc569856c, UInt64}, SpinBasis{1//2, Int64}}}, CompositeBasis{Vector{Int64}, Tuple{SubspaceBasis{Vector{Int64}, PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Ket{PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Vector{ComplexF64}}, 0x9711e39bc569856c, UInt64}, SpinBasis{1//2, Int64}}}, Matrix{ComplexF64}}
  indices::Int64
Locals
  result::Any
  rank::Int64
Body::Operator
1 ─       QuantumOpticsBase.check_ptrace_arguments(a, indices)
│   %2  = Base.getproperty(a, :basis_l)::CompositeBasis{Vector{Int64}, Tuple{SubspaceBasis{Vector{Int64}, PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Ket{PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Vector{ComplexF64}}, 0x9711e39bc569856c, UInt64}, SpinBasis{1//2, Int64}}}
│   %3  = Base.getproperty(%2, :shape)::Vector{Int64}
│         (rank = QuantumOpticsBase.length(%3))
│   %5  = Core.apply_type(QuantumOpticsBase.Val, rank)::Type{Val{_A}} where _A
│   %6  = Base.getproperty(a, :data)::Matrix{ComplexF64}
│   %7  = Base.getproperty(a, :basis_l)::CompositeBasis{Vector{Int64}, Tuple{SubspaceBasis{Vector{Int64}, PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Ket{PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Vector{ComplexF64}}, 0x9711e39bc569856c, UInt64}, SpinBasis{1//2, Int64}}}
│   %8  = Base.getproperty(%7, :shape)::Vector{Int64}
│   %9  = Base.getproperty(a, :basis_r)::CompositeBasis{Vector{Int64}, Tuple{SubspaceBasis{Vector{Int64}, PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Ket{PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Vector{ComplexF64}}, 0x9711e39bc569856c, UInt64}, SpinBasis{1//2, Int64}}}
│   %10 = Base.getproperty(%9, :shape)::Vector{Int64}
│         (result = QuantumOpticsBase._ptrace(%5, %6, %8, %10, indices))
│   %12 = Base.getproperty(a, :basis_l)::CompositeBasis{Vector{Int64}, Tuple{SubspaceBasis{Vector{Int64}, PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Ket{PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Vector{ComplexF64}}, 0x9711e39bc569856c, UInt64}, SpinBasis{1//2, Int64}}}
│   %13 = QuantumOpticsBase.ptrace(%12, indices)::Any
│   %14 = Base.getproperty(a, :basis_r)::CompositeBasis{Vector{Int64}, Tuple{SubspaceBasis{Vector{Int64}, PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Ket{PositionBasis{-47.0, 26.008912220730075, Int64, Float64}, Vector{ComplexF64}}, 0x9711e39bc569856c, UInt64}, SpinBasis{1//2, Int64}}}
│   %15 = QuantumOpticsBase.ptrace(%14, indices)::Any
│   %16 = QuantumOpticsBase.Operator(%13, %15, result)::Operator
└──       return %16

When one want to access its field to do something like:

function fout(t,rho)
    rho_ptr=ptrace(rho,1)
    real(dot(ψ_proj.data,rho_ptr.data,ψ_proj.data)/tr(rho_ptr))
end
tout, projs = timeevolution.master_nh_dynamic(tspan, rho0, H_func; fout=fout,kwargs...)

typeof(projs) will be Vector{Any} since typeof(rho_ptr.data) is Any

And interesting thing is that, for ket_ptr=ptrace(Ket,index), insdeat of recognizing ket_ptr.data to be Any, it recognize it to be Matrix{ComplexF64}.

[amilsted]

This is a tricky one. The unstable part is the ptrace() of a CompositeBasis, e.g. ptrace(a.basis_l, indices). This is hard to stabilize, as the type of the output depends on the index, or indices, selected. Stabilization would have to rely on constant propagation, which might sometimes work...

[amilsted]

Oh, and we turn the rank into a Val type, which also can't generally be inferred at compile time.

I'm not sure how much we should worry, however, as ptrace() in any case needs to allocate memory for the output. I would hope the runtime dispatch does not add much overhead.