Fix type-instability for prod (#363)

Author: ytdHuang

src/qobj/dimensions.jl

@@ -72,7 +72,11 @@ Base.:(==)(dims::AbstractDimensions, vect::AbstractVector{T}) where {T} = vect =
 
 Base.length(::AbstractDimensions{N}) where {N} = N
 
+# need to specify return type `Int` for `_get_space_size`
+# otherwise the type of `prod(::Dimensions)` will be unstable
+_get_space_size(s::AbstractSpace)::Int = s.size
 Base.prod(dims::Dimensions) = prod(dims.to)
+Base.prod(spaces::SVector{N,AbstractSpace}) where {N} = prod(_get_space_size, spaces)
 
 LinearAlgebra.transpose(dims::Dimensions) = dims
 LinearAlgebra.transpose(dims::CompoundDimensions) = CompoundDimensions(dims.from, dims.to) # switch `to` and `from`

src/qobj/quantum_object_base.jl

@@ -176,7 +176,7 @@ function _check_QuantumObject(type::BraQuantumObject, dims::Dimensions, m::Int,
 end
 
 function _check_QuantumObject(type::OperatorQuantumObject, dims::Dimensions, m::Int, n::Int)
-    L = prod(dims.to)
+    L = prod(dims)
     (L == m == n) || throw(DimensionMismatch("Operator with dims = $(dims) does not fit the array size = $((m, n))."))
     return nothing
 end

src/qobj/space.jl

@@ -26,8 +26,3 @@ struct Space <: AbstractSpace
     end
 end
 Base.string(s::Space) = string(s.size) # this is only used when printing AbstractDimensions
-
-# for `prod(::Dimensions)`
-Base.:(*)(i::Int, s::AbstractSpace) = i * s.size
-Base.:(*)(s1::AbstractSpace, s2::AbstractSpace) = s1.size * s2.size
-Base.prod(spaces::SVector{1,AbstractSpace}) = spaces[1].size # for `Dimensions.to` has only a single Space

test/core-test/correlations_and_spectrum.jl

@@ -28,8 +28,7 @@
     @test all(spec2 .≈ spec3)
 
     @testset "Type Inference spectrum" begin
-        # TODO: fix this type inference issue (originally from `steadystate`)
-        # @inferred correlation_2op_1t(H, nothing, t_l, c_ops, a', a; progress_bar = Val(false))
+        @inferred correlation_2op_1t(H, nothing, t_l, c_ops, a', a; progress_bar = Val(false))
         @inferred spectrum_correlation_fft(t_l, corr1)
         @inferred spectrum(H, ω_l2, c_ops, a', a)
         @inferred spectrum(H, ω_l2, c_ops, a', a; solver = PseudoInverse())

test/core-test/steady_state.jl

@@ -38,8 +38,7 @@
         @inferred steadystate(H, psi0, t_l[end], c_ops, solver = solver)
         @inferred steadystate(L, psi0, t_l[end], solver = solver)
 
-        # TODO: fix the following type inference issues
-        #= solver = SteadyStateDirectSolver()
+        solver = SteadyStateDirectSolver()
         @inferred steadystate(H, c_ops, solver = solver)
         @inferred steadystate(L, solver = solver)
 
@@ -52,7 +51,7 @@
         @inferred steadystate(L, solver = solver)
 
         @inferred steadystate(H, c_ops)
-        @inferred steadystate(L) =#
+        @inferred steadystate(L)
     end
 
     H = a_d * a
@@ -74,15 +73,14 @@
     @test abs(sum(sol_me.expect[1, end-100:end]) / 101 - expect(e_ops[1], ρ_ss2)) < 1e-3
 
     @testset "Type Inference (steadystate_floquet)" begin
-        # TODO: fix the following type inference issues
-        #= @inferred steadystate_floquet(H, -1im * 0.5 * H_t, 1im * 0.5 * H_t, 1, c_ops, solver = SSFloquetLinearSystem())
+        @inferred steadystate_floquet(H, -1im * 0.5 * H_t, 1im * 0.5 * H_t, 1, c_ops, solver = SSFloquetLinearSystem())
         @inferred steadystate_floquet(
             H,
             -1im * 0.5 * H_t,
             1im * 0.5 * H_t,
             1,
             c_ops,
             solver = SSFloquetEffectiveLiouvillian(),
-        ) =#
+        )
     end
 end