Refactor dimension handling by introducing hilbert_dimensions_to_size and liouville_dimensions_to_size functions

Author: albertomercurio

docs/src/getting_started/type_stability.md

@@ -183,7 +183,7 @@ The `dims` field contains the dimensions of the subsystems (in this case, three
 
 ```@example type-stability
 function reshape_operator_data(dims)
-    op = Qobj(randn(prod(dims), prod(dims)), type=Operator(), dims=dims)
+    op = Qobj(randn(hilbert_dimensions_to_size(dims)...), type=Operator(), dims=dims)
     op_dims = op.dims
     op_data = op.data
     return reshape(op_data, vcat(op_dims, op_dims)...)

docs/src/resources/api.md

@@ -326,6 +326,8 @@ convert_unit
 row_major_reshape
 meshgrid
 enr_state_dictionaries
+hilbert_dimensions_to_size
+liouville_dimensions_to_size
 ```
 
 ## [Visualization](@id doc-API:Visualization)

ext/QuantumToolboxMakieExt.jl

@@ -226,7 +226,7 @@ function _plot_fock_distribution(
     kwargs...,
 ) where {SType<:Union{Bra,Ket,Operator}}
     ρ = ket2dm(ρ)
-    D = prod(ρ.dims)
+    D = hilbert_dimensions_to_size(ρ.dims)[1]
     isapprox(tr(ρ), 1, atol = 1e-4) || (@warn "The input ρ should be normalized.")
 
     xvec = 0:(D-1)

src/correlations.jl

@@ -113,7 +113,7 @@ function correlation_2op_2t(
     reverse::Bool = false,
     kwargs...,
 ) where {HOpType<:Union{Operator,SuperOperator},StateOpType<:Union{Ket,Operator}}
-    C = eye(prod(H.dimensions), dims = H.dimensions)
+    C = eye(hilbert_dimensions_to_size(H.dimensions)[1], dims = H.dimensions)
     if reverse
         corr = correlation_3op_2t(H, ψ0, tlist, τlist, c_ops, A, B, C; kwargs...)
     else

src/qobj/arithmetic_and_attributes.jl

@@ -672,15 +672,15 @@ Get the coherence value ``\alpha`` by measuring the expectation value of the des
 It returns both ``\alpha`` and the corresponding state with the coherence removed: ``\ket{\delta_\alpha} = \exp ( \alpha^* \hat{a} - \alpha \hat{a}^\dagger ) \ket{\psi}`` for a pure state, and ``\hat{\rho_\alpha} = \exp ( \alpha^* \hat{a} - \alpha \hat{a}^\dagger ) \hat{\rho} \exp ( -\bar{\alpha} \hat{a} + \alpha \hat{a}^\dagger )`` for a density matrix. These states correspond to the quantum fluctuations around the coherent state ``\ket{\alpha}`` or ``|\alpha\rangle\langle\alpha|``.
 """
 function get_coherence(ψ::QuantumObject{Ket})
-    a = destroy(prod(ψ.dimensions))
+    a = destroy(hilbert_dimensions_to_size(ψ.dimensions)[1])
     α = expect(a, ψ)
     D = exp(α * a' - conj(α) * a)
 
     return α, D' * ψ
 end
 
 function get_coherence(ρ::QuantumObject{Operator})
-    a = destroy(prod(ρ.dimensions))
+    a = destroy(hilbert_dimensions_to_size(ρ.dimensions)[1])
     α = expect(a, ρ)
     D = exp(α * a' - conj(α) * a)
 

src/qobj/dimensions.jl

@@ -3,6 +3,7 @@ This file defines the ProductDimensions structures, which can describe composite
 =#
 
 export AbstractDimensions, ProductDimensions, GeneralProductDimensions
+export hilbert_dimensions_to_size, liouville_dimensions_to_size
 
 abstract type AbstractDimensions{M,N} end
 
@@ -80,11 +81,52 @@ dimensions_to_dims(::Nothing) = nothing # for EigsolveResult.dimensions = nothin
 
 Base.length(::AbstractDimensions{N}) where {N} = N
 
-# need to specify return type `Int` for `_get_space_size`
-# otherwise the type of `prod(::ProductDimensions)` will be unstable
-_get_space_size(s::AbstractSpace)::Int = s.size
-Base.prod(dims::ProductDimensions) = prod(dims.to)
-Base.prod(spaces::NTuple{N,AbstractSpace}) where {N} = prod(_get_space_size, spaces)
+"""
+    hilbert_dimensions_to_size(dimensions)
+
+Returns the matrix dimensions `(m, n)` of an [`Operator`](@ref) with the given `dimensions`.
+
+For [`ProductDimensions`](@ref), returns `(m, m)` where `m` is the product of all subsystem Hilbert space dimensions.
+For [`GeneralProductDimensions`](@ref), returns `(m, n)` where `m` is the product of the `to` dimensions
+and `n` is the product of the `from` dimensions.
+
+If `dimensions` is an `Integer` or a vector/tuple of `Integer`s, it is automatically treated as [`ProductDimensions`](@ref).
+"""
+function hilbert_dimensions_to_size(dimensions::ProductDimensions)
+    m = prod(hilbert_dimensions_to_size, dimensions.to)
+    return (m, m)
+end
+function hilbert_dimensions_to_size(dimensions::GeneralProductDimensions)
+    m = prod(hilbert_dimensions_to_size, dimensions.to)
+    n = prod(hilbert_dimensions_to_size, dimensions.from)
+    return (m, n)
+end
+hilbert_dimensions_to_size(dimensions::Union{<:Integer,AbstractVector{<:Integer},NTuple{N,Integer}}) where {N} =
+    hilbert_dimensions_to_size(ProductDimensions(dimensions))
+
+"""
+    liouville_dimensions_to_size(dimensions)
+
+Returns the matrix dimensions `(m, n)` of a [`SuperOperator`](@ref) with the given `dimensions`.
+
+For [`ProductDimensions`](@ref), returns `(m, m)` where `m` is the product of all subsystem Liouville space dimensions
+(each Hilbert dimension `d` contributes `d²` to the product).
+For [`GeneralProductDimensions`](@ref), returns `(m, n)` where `m` is the product of the `to` dimensions
+and `n` is the product of the `from` dimensions.
+
+If `dimensions` is an `Integer` or a vector/tuple of `Integer`s, it is automatically treated as [`ProductDimensions`](@ref).
+"""
+function liouville_dimensions_to_size(dimensions::ProductDimensions)
+    m = prod(liouville_dimensions_to_size, dimensions.to)
+    return (m, m)
+end
+function liouville_dimensions_to_size(dimensions::GeneralProductDimensions)
+    m = prod(liouville_dimensions_to_size, dimensions.to)
+    n = prod(liouville_dimensions_to_size, dimensions.from)
+    return (m, n)
+end
+liouville_dimensions_to_size(dimensions::Union{<:Integer,AbstractVector{<:Integer},NTuple{N,Integer}}) where {N} =
+    liouville_dimensions_to_size(ProductDimensions(dimensions))
 
 Base.transpose(dimensions::ProductDimensions) = dimensions
 Base.transpose(dimensions::GeneralProductDimensions) = GeneralProductDimensions(dimensions.from, dimensions.to) # switch `to` and `from`

src/qobj/eigsolve.jl

@@ -424,7 +424,7 @@ end
         c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
         alg::AbstractODEAlgorithm = DP5(),
         params::NamedTuple = NamedTuple(),
-        ρ0::AbstractMatrix = rand_dm(prod(H.dimensions)).data,
+        ρ0::AbstractMatrix = rand_dm(hilbert_dimensions_to_size(H.dimensions)[1]).data,
         eigvals::Int = 1,
         krylovdim::Int = min(10, size(H, 1)),
         maxiter::Int = 200,
@@ -467,7 +467,7 @@ function eigsolve_al(
     c_ops::Union{Nothing,AbstractVector,Tuple} = nothing;
     alg::AbstractODEAlgorithm = DP5(),
     params::NamedTuple = NamedTuple(),
-    ρ0::AbstractMatrix = rand_dm(prod(H.dimensions)).data,
+    ρ0::AbstractMatrix = rand_dm(hilbert_dimensions_to_size(H.dimensions)[1]).data,
     eigvals::Int = 1,
     krylovdim::Int = min(10, size(H, 1)),
     maxiter::Int = 200,

src/qobj/energy_restricted.jl

@@ -76,6 +76,9 @@ Base.:(==)(s_enr1::EnrSpace, s_enr2::EnrSpace) = (s_enr1.size == s_enr2.size) &&
 
 dimensions_to_dims(s_enr::EnrSpace) = s_enr.dims
 
+hilbert_dimensions_to_size(s_enr::EnrSpace) = s_enr.size
+liouville_dimensions_to_size(s_enr::EnrSpace) = s_enr.size^2
+
 @doc raw"""
     enr_state_dictionaries(dims, n_excitations)
 

src/qobj/operators.jl

@@ -38,7 +38,7 @@ rand_unitary(
     distribution::Union{Symbol,Val} = Val(:haar),
 ) = rand_unitary(dimensions, makeVal(distribution))
 function rand_unitary(dimensions::Union{ProductDimensions,AbstractVector{Int},Tuple}, ::Val{:haar})
-    N = prod(dimensions)
+    N = hilbert_dimensions_to_size(dimensions)[1]
 
     # generate N x N matrix Z of complex standard normal random variates
     Z = randn(ComplexF64, N, N)
@@ -53,7 +53,7 @@ function rand_unitary(dimensions::Union{ProductDimensions,AbstractVector{Int},Tu
     return QuantumObject(to_dense(Q * Diagonal(Λ)); type = Operator(), dims = dimensions)
 end
 function rand_unitary(dimensions::Union{ProductDimensions,AbstractVector{Int},Tuple}, ::Val{:exp})
-    N = prod(dimensions)
+    N = hilbert_dimensions_to_size(dimensions)[1]
 
     # generate N x N matrix Z of complex standard normal random variates
     Z = randn(ComplexF64, N, N)
@@ -554,7 +554,7 @@ where ``\omega = \exp(\frac{2 \pi i}{N})``.
 """
 qft(dimensions::Int) = QuantumObject(_qft_op(dimensions), Operator(), dimensions)
 qft(dimensions::Union{ProductDimensions,AbstractVector{Int},Tuple}) =
-    QuantumObject(_qft_op(prod(dimensions)), Operator(), dimensions)
+    QuantumObject(_qft_op(hilbert_dimensions_to_size(dimensions)[1]), Operator(), dimensions)
 function _qft_op(N::Int)
     ω = exp(2.0im * π / N)
     arr = 0:(N-1)

src/qobj/quantum_object_base.jl

@@ -142,23 +142,23 @@ _check_QuantumObject(
 
 function _check_QuantumObject(type::Ket, dimensions::ProductDimensions, m::Int, n::Int)
     (n != 1) && throw(DomainError((m, n), "The size of the array is not compatible with Ket"))
-    (prod(dimensions) != m) && throw(
+    (hilbert_dimensions_to_size(dimensions)[1] != m) && throw(
         DimensionMismatch("Ket with dims = $(_get_dims_string(dimensions)) does not fit the array size = $((m, n))."),
     )
     return nothing
 end
 
 function _check_QuantumObject(type::Bra, dimensions::ProductDimensions, m::Int, n::Int)
     (m != 1) && throw(DomainError((m, n), "The size of the array is not compatible with Bra"))
-    (prod(dimensions) != n) && throw(
+    (hilbert_dimensions_to_size(dimensions)[1] != n) && throw(
         DimensionMismatch("Bra with dims = $(_get_dims_string(dimensions)) does not fit the array size = $((m, n))."),
     )
     return nothing
 end
 
 function _check_QuantumObject(type::Operator, dimensions::ProductDimensions, m::Int, n::Int)
-    L = prod(dimensions)
-    (L == m == n) || throw(
+    obj_size = hilbert_dimensions_to_size(dimensions)[1]
+    (obj_size == m == n) || throw(
         DimensionMismatch(
             "Operator with dims = $(_get_dims_string(dimensions)) does not fit the array size = $((m, n)).",
         ),
@@ -168,7 +168,8 @@ end
 
 function _check_QuantumObject(type::Operator, dimensions::GeneralProductDimensions, m::Int, n::Int)
     ((m == 1) || (n == 1)) && throw(DomainError((m, n), "The size of the array is not compatible with Operator"))
-    ((prod(dimensions.to) != m) || (prod(dimensions.from) != n)) && throw(
+    obj_size = hilbert_dimensions_to_size(dimensions)
+    ((obj_size[1] != m) || (obj_size[2] != n)) && throw(
         DimensionMismatch(
             "Operator with dims = $(_get_dims_string(dimensions)) does not fit the array size = $((m, n)).",
         ),
@@ -178,7 +179,7 @@ end
 
 function _check_QuantumObject(type::SuperOperator, dimensions::ProductDimensions, m::Int, n::Int)
     (m != n) && throw(DomainError((m, n), "The size of the array is not compatible with SuperOperator"))
-    (prod(dimensions) != sqrt(m)) && throw(
+    (liouville_dimensions_to_size(dimensions)[1] != m) && throw(
         DimensionMismatch(
             "SuperOperator with dims = $(_get_dims_string(dimensions)) does not fit the array size = $((m, n)).",
         ),
@@ -188,7 +189,7 @@ end
 
 function _check_QuantumObject(type::OperatorKet, dimensions::ProductDimensions, m::Int, n::Int)
     (n != 1) && throw(DomainError((m, n), "The size of the array is not compatible with OperatorKet"))
-    (prod(dimensions) != sqrt(m)) && throw(
+    (liouville_dimensions_to_size(dimensions)[1] != m) && throw(
         DimensionMismatch(
             "OperatorKet with dims = $(_get_dims_string(dimensions)) does not fit the array size = $((m, n)).",
         ),
@@ -198,7 +199,7 @@ end
 
 function _check_QuantumObject(type::OperatorBra, dimensions::ProductDimensions, m::Int, n::Int)
     (m != 1) && throw(DomainError((m, n), "The size of the array is not compatible with OperatorBra"))
-    (prod(dimensions) != sqrt(n)) && throw(
+    (liouville_dimensions_to_size(dimensions)[1] != n) && throw(
         DimensionMismatch(
             "OperatorBra with dims = $(_get_dims_string(dimensions)) does not fit the array size = $((m, n)).",
         ),