Space and Dimensions structure to align with QuTiP

[ytdHuang]

Space

Define a new AbstractSpace

abstract type AbstractSpace end

# this show function is for printing AbstractDimensions
Base.show(io::IO, svec::SVector{N,AbstractSpace}) where {N} = print(io, "[", join(svec, ", "), "]")

The normal Hilbert space will be a structure called Space with a single field size:

struct Space <: AbstractSpace
    size::Int

    function Space(size::Int)
        (size < 1) && throw(DomainError(size, "The size of Space must be positive integer (≥ 1)."))
        return new(size)
    end
end
Base.show(io::IO, s::Space) = print(io, s.size)

The AbstractSpace opens a door for us to implement energy-restricted space in the future (just a brief pseudo-code here):

struct EnrSpace <: AbstractSpace
    size::Int
    dims
    state2idx
    idx2state
    n_excitations
end
Base.show(io::IO, s::EnrSpace) = print(io, join(s.dims, ", ")) # don't print the brackets []

Dimensions

Define a new AbstractDimensions

abstract type AbstractDimensions{N} end

The original dims can be implemented as a structure called Dimensions:

struct Dimensions{N} <: AbstractDimensions{N}
    to::SVector{N,AbstractSpace}
end
Base.show(io::IO, D::Dimensions) = print(io, D.to)

function Dimensions(dims::Union{AbstractVector{T},NTuple{N,T}}) where {T<:Integer,N}
    QuantumToolbox._non_static_array_warning("dims", dims)
    L = length(dims)
    (L > 0) || throw(DomainError(dims, "The argument dims must be of non-zero length"))

    return Dimensions(SVector{L,AbstractSpace}(Space.(dims)))
end
Dimensions(dims::Any) = throw(
    ArgumentError(
        "The argument dims must be a Tuple or a StaticVector of non-zero length and contain only positive integers.",
    ),
)

For example,

julia> d = Dimensions((2, 2, 2, 3, 4))
[2, 2, 2, 3, 4]

julia> typeof(d)
Dimensions{5}

The AbstractDimensions opens a door for us to implement CompoundDimensions:

struct CompoundDimensions{N} <: AbstractDimensions{N}
    # note that the number `N` should be the same for both `to` and `from`
    to::SVector{N,AbstractSpace}   # space acting on the left
    from::SVector{N,AbstractSpace} # space acting on the right
end
Base.show(io::IO, D::CompoundDimensions) = print(io, "[", D.to, ", ", D.from, "]")

By overloading some existing functions, I think we can implement it in our current code. For example,

Base.:(*)(i::Int, s::Space) = i * s.size
Base.:(*)(s1::Space, s2::Space) = s1.size * s2.size

Base.prod(dims::Dimensions) = prod(dims.to)

# not sure we should implement it like this or not for CompoundDimensions
Base.prod(dims::CompoundDimensions) = maximum(prod(dims.to), prod(dims.from))

Change definition of Qobj and QobjEvo

abstract type AbstractQuantumObject{DataType,ObjType,DimsType} end

struct Qobj{..., DimsType} <: AbstractQuantumObject{..., DimsType<:AbstractDimensions}
    ...
    dims::DimsType
end

struct QobjEvo{..., DimsType} <: AbstractQuantumObject{..., DimsType<:AbstractDimensions}
    ...
    dims::DimsType
end