fix: cleaner dispatching onto StaticArrays

docs/src/tutorials.md

@@ -60,13 +60,13 @@ can specify an array type in its first (and second) arguments. Let's see an exam
 ```jldoctest
 julia> using StaticArrays
 
-julia> state = coherentstate(SVector{2}, SMatrix{2,2}, QuadPairBasis(1), 1.0-im)
+julia> state = coherentstate(SVector, SMatrix, QuadPairBasis(1), 1.0-im)
 GaussianState for 1 mode.
   symplectic basis: QuadPairBasis
-mean: 2-element SVector{2, Float64} with indices SOneTo(2):
+mean: 2-element SVector{2, Any} with indices SOneTo(2):
   2.0
  -2.0
-covariance: 2×2 SMatrix{2, 2, Float64, 4} with indices SOneTo(2)×SOneTo(2):
+covariance: 2×2 SMatrix{2, 2, Any, 4} with indices SOneTo(2)×SOneTo(2):
  1.0  0.0
  0.0  1.0
 

ext/StaticArraysExt/StaticArraysExt.jl

@@ -3,10 +3,12 @@ module StaticArraysExt
 using StaticArrays: SVector, SMatrix, SArray
 
 using Gabs
-import Gabs: ptrace, tensor, ⊗, _promote_output_matrix, _promote_output_vector, 
-            _generaldyne_map
+using Gabs: SymplecticBasis, QuadPairBasis
+
+import Gabs: ptrace, tensor, ⊗, _promote_output_matrix, _promote_output_vector,
+            _generaldyne_map, _infer_types
 
 include("utils.jl")
 include("measurements.jl")
 
-end
+end

ext/StaticArraysExt/utils.jl

@@ -14,4 +14,44 @@
 end
 Base.@propagate_inbounds function _promote_output_matrix(::Type{<:SMatrix}, mat_out, out_dim::Tuple)
     return SMatrix{out_dim[1],out_dim[2]}(mat_out)
-end
+end
+Base.@propagate_inbounds function _promote_output_matrix(::Type{T1}, ::Type{T2}, mat_out) where {T1<:SMatrix, T2<:AbstractMatrix}
+    return SMatrix{size(mat_out,1), size(mat_out,2)}(mat_out)
+end
+Base.@propagate_inbounds function _promote_output_vector(::Type{T1}, ::Type{T2}, vec_out) where {T1<:SVector, T2<:AbstractVector}
+    return SVector{length(vec_out)}(vec_out)
+end
+
+abstract type ArrayTrait end
+
+struct DenseArrayTrait <: ArrayTrait end
+struct StaticArrayTrait <: ArrayTrait end
+
+array_trait(::Type{<:Array}) = DenseArrayTrait()
+array_trait(::Type{<:SArray}) = StaticArrayTrait()
+array_trait(::Type{<:UnionAll}) = StaticArrayTrait()
+
+function _infer_types(::DenseArrayTrait, nmodes, T = Float64)
+    disp_type = Vector{T}
+    transform_type = Matrix{T}
+    return disp_type, transform_type
+end
+function _infer_types(::StaticArrayTrait, nmodes, T = Float64)
+    disp_type = SArray{Tuple{2*nmodes}, T}
+    transform_type = SArray{Tuple{2*nmodes, 2*nmodes}, T}
+    return disp_type, transform_type
+end
+function _infer_types(T1, T2, basis)
+    nmodes = basis.nmodes
+    elT1 = eltype(T1)
+    elT2 = eltype(T2)
+    disp_type1, _ = _infer_types(array_trait(T1), nmodes, elT1)
+    _, transform_type2 = _infer_types(array_trait(T2), nmodes, elT2)
+    return disp_type1, transform_type2
+end
+function _infer_types(T, basis)
+    nmodes = basis.nmodes
+    elT = eltype(T)
+    disp_type, transform_type = _infer_types(array_trait(T), nmodes, elT)
+    return disp_type, transform_type
+end

src/channels.jl

@@ -3,50 +3,75 @@
 ##
 
 function displace(::Type{Td}, ::Type{Tt}, basis::SymplecticBasis{N}, alpha::A, noise::M) where {Td,Tt,N<:Int,A,M}
+    disp_type, transform_type = _infer_types(Td, Tt, basis)
     disp, transform = _displace(basis, alpha)
-    return GaussianChannel(basis, Td(disp), Tt(transform), Tt(noise))
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise))
+end
+function displace(::Type{T}, basis::SymplecticBasis{N}, alpha::A, noise::M) where {T,N<:Int,A,M}
+    disp_type, transform_type = _infer_types(T, basis)
+    disp, transform = _displace(basis, alpha)
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise))
 end
-displace(::Type{T}, basis::SymplecticBasis{N}, alpha::A, noise::M) where {T,N<:Int,A,M} = displace(T, T, basis, alpha, noise)
 function displace(basis::SymplecticBasis{N}, alpha::A, noise::M) where {N<:Int,A,M}
     disp, transform = _displace(basis, alpha)
     return GaussianChannel(basis, disp, transform, noise)
 end
 
 function squeeze(::Type{Td}, ::Type{Tt}, basis::SymplecticBasis{N}, r::R, theta::R, noise::M) where {Td,Tt,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(Td, Tt, basis)
+    disp, transform = _squeeze(basis, r, theta)
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise))
+end
+function squeeze(::Type{T}, basis::SymplecticBasis{N}, r::R, theta::R, noise::M) where {T,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(T, basis)
     disp, transform = _squeeze(basis, r, theta)
-    return GaussianChannel(basis, Td(disp), Tt(transform), Tt(noise))
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise))
 end
-squeeze(::Type{T}, basis::SymplecticBasis{N}, r::R, theta::R, noise::M) where {T,N<:Int,R,M} = squeeze(T, T, basis, r, theta, noise)
 function squeeze(basis::SymplecticBasis{N}, r::R, theta::R, noise::M) where {N<:Int,R,M}
     disp, transform = _squeeze(basis, r, theta)
     return GaussianChannel(basis, disp, transform, noise)
 end
 
 function twosqueeze(::Type{Td}, ::Type{Tt}, basis::SymplecticBasis{N}, r::R, theta::R, noise::M) where {Td,Tt,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(Td, Tt, basis)
     disp, transform = _twosqueeze(basis, r, theta)
-    return GaussianChannel(basis, Td(disp), Tt(transform), Tt(noise))
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise))
+end
+function twosqueeze(::Type{T}, basis::SymplecticBasis{N}, r::R, theta::R, noise::M) where {T,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(T, basis)
+    disp, transform = _twosqueeze(basis, r, theta)
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise))
 end
-twosqueeze(::Type{T}, basis::SymplecticBasis{N}, r::R, theta::R, noise::M) where {T,N<:Int,R,M} = twosqueeze(T, T, basis, r, theta, noise)
 function twosqueeze(basis::SymplecticBasis{N}, r::R, theta::R, noise::M) where {N<:Int,R,M}
     disp, transform = _twosqueeze(basis, r, theta)
     return GaussianChannel(basis, disp, transform, noise)
 end
 
 function phaseshift(::Type{Td}, ::Type{Tt}, basis::SymplecticBasis{N}, theta::R, noise::M) where {Td,Tt,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(Td, Tt, basis)
+    disp, transform = _phaseshift(basis, theta)
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise))
+end
+function phaseshift(::Type{T}, basis::SymplecticBasis{N}, theta::R, noise::M) where {T,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(T, T, basis)
     disp, transform = _phaseshift(basis, theta)
-    return GaussianChannel(basis, Td(disp), Tt(transform), Tt(noise))
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise))
 end
-phaseshift(::Type{T}, basis::SymplecticBasis{N}, theta::R, noise::M) where {T,N<:Int,R,M} = phaseshift(T, T, basis, theta, noise)
 function phaseshift(basis::SymplecticBasis{N}, theta::R, noise::M) where {N<:Int,R,M}
     disp, transform = _phaseshift(basis, theta)
     return GaussianChannel(basis, disp, transform, noise)
 end
 
 function beamsplitter(::Type{Td}, ::Type{Tt}, basis::SymplecticBasis{N}, transmit::R, noise::M) where {Td,Tt,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(Td, Tt, basis)
     disp, transform = _beamsplitter(basis, transmit)
-    return GaussianChannel(basis, Td(disp), Tt(transform), Tt(noise))
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise))
+end
+function beamsplitter(::Type{T}, basis::SymplecticBasis{N}, transmit::R, noise::M) where {T,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(T, basis)
+    disp, transform = _beamsplitter(basis, transmit)
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise))
 end
-beamsplitter(::Type{T}, basis::SymplecticBasis{N}, transmit::R, noise::M) where {T,N<:Int,R,M} = beamsplitter(T, T, basis, transmit, noise)
 function beamsplitter(basis::SymplecticBasis{N}, transmit::R, noise::M) where {N<:Int,R,M}
     disp, transform = _beamsplitter(basis, transmit)
     return GaussianChannel(basis, disp, transform, noise)
@@ -84,25 +109,30 @@ noise: 2×2 Matrix{Float64}:
 ```
 """
 function attenuator(::Type{Td}, ::Type{Tt}, basis::SymplecticBasis{N}, theta::R, n::M; ħ = 2) where {Td,Tt,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(Td, Tt, basis)
+    disp, transform, noise = _attenuator(basis, theta, n)
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise); ħ = ħ)
+end
+function attenuator(::Type{T}, basis::SymplecticBasis{N}, theta::R, n::M; ħ = 2) where {T,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(T, basis)
     disp, transform, noise = _attenuator(basis, theta, n)
-    return GaussianChannel(basis, Td(disp), Tt(transform), Tt(noise); ħ = ħ)
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise); ħ = ħ)
 end
-attenuator(::Type{T}, basis::SymplecticBasis{N}, theta::R, n::M; ħ = 2) where {T,N<:Int,R,M} = attenuator(T, T, basis, theta, n; ħ = ħ)
 function attenuator(basis::SymplecticBasis{N}, theta::R, n::M; ħ = 2) where {N<:Int,R,M}
     disp, transform, noise = _attenuator(basis, theta, n)
     return GaussianChannel(basis, disp, transform, noise; ħ = ħ)
 end
 function _attenuator(basis::Union{QuadPairBasis{N},QuadBlockBasis{N}}, theta::R, n::M) where {N<:Int,R<:Real,M<:Int}
     nmodes = basis.nmodes
-    disp = zeros(R, 2*nmodes) 
+    disp = zeros(R, 2*nmodes)
     transform = Matrix{R}(cos(theta) * I, 2*nmodes, 2*nmodes)
     noise = Matrix{R}((sin(theta))^2 * n * I, 2*nmodes, 2*nmodes)
     return disp, transform, noise
 end
 function _attenuator(basis::QuadPairBasis{N}, theta::R, n::M) where {N<:Int,R<:Vector,M<:Vector}
     nmodes = basis.nmodes
     Rt = eltype(R)
-    disp = zeros(Rt, 2*nmodes) 
+    disp = zeros(Rt, 2*nmodes)
     transform = zeros(Rt, 2*nmodes, 2*nmodes)
     noise = zeros(Rt, 2*nmodes, 2*nmodes)
     @inbounds for i in Base.OneTo(nmodes)
@@ -168,17 +198,22 @@ noise: 2×2 Matrix{Float64}:
 ```
 """
 function amplifier(::Type{Td}, ::Type{Tt}, basis::SymplecticBasis{N}, r::R, n::M; ħ = 2) where {Td,Tt,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(Td, Tt, basis)
+    disp, transform, noise = _amplifier(basis, r, n)
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise); ħ = ħ)
+end
+function amplifier(::Type{T}, basis::SymplecticBasis{N}, r::R, n::M; ħ = 2) where {T,N<:Int,R,M}
+    disp_type, transform_type = _infer_types(T, basis)
     disp, transform, noise = _amplifier(basis, r, n)
-    return GaussianChannel(basis, Td(disp), Tt(transform), Tt(noise); ħ = ħ)
+    return GaussianChannel(basis, disp_type(disp), transform_type(transform), transform_type(noise); ħ = ħ)
 end
-amplifier(::Type{T}, basis::SymplecticBasis{N}, r::R, n::M; ħ = 2) where {T,N<:Int,R,M} = amplifier(T, T, basis, r, n; ħ = ħ)
 function amplifier(basis::SymplecticBasis{N}, r::R, n::M; ħ = 2) where {N<:Int,R,M}
     disp, transform, noise = _amplifier(basis, r, n)
     return GaussianChannel(basis, disp, transform, noise; ħ = ħ)
 end
 function _amplifier(basis::Union{QuadPairBasis{N},QuadBlockBasis{N}}, r::R, n::M) where {N<:Int,R<:Real,M<:Int}
     nmodes = basis.nmodes
-    disp = zeros(R, 2*nmodes) 
+    disp = zeros(R, 2*nmodes)
     transform = Matrix{R}(cosh(r) * I, 2*nmodes, 2*nmodes)
     noise = Matrix{R}((sinh(r))^2 * n * I, 2*nmodes, 2*nmodes)
     return disp, transform, noise
@@ -244,7 +279,10 @@ function _tensor(op1::GaussianChannel{B1,D1,T1}, op2::GaussianChannel{B2,D2,T2})
     block1, block2 = Base.OneTo(2*nmodes1), Base.OneTo(2*nmodes2)
     # initialize direct sum of displacement vectors
     disp1, disp2 = op1.disp, op2.disp
-    Dt = promote_type(eltype(disp1), eltype(disp2))
+    elD1 = eltype(disp1) isa Type ? eltype(disp1) : Float64
+    elD2 = eltype(disp2) isa Type ? eltype(disp2) : Float64
+    Dt = promote_type(elD1, elD2)
+    Dt = Dt == Any ? Float64 : Dt
     disp′ = zeros(Dt, 2*nmodes)
     @inbounds for i in block1
         disp′[i] = disp1[i]
@@ -254,7 +292,10 @@ function _tensor(op1::GaussianChannel{B1,D1,T1}, op2::GaussianChannel{B2,D2,T2})
     end
     # initialize direct sum of transform and noise matrices
     trans1, trans2 = op1.transform, op2.transform
-    Tt = promote_type(eltype(trans1), eltype(trans2))
+    elT1 = eltype(trans1) isa Type ? eltype(trans1) : Float64
+    elT2 = eltype(trans2) isa Type ? eltype(trans2) : Float64
+    Tt = promote_type(elT1, elT2)
+    Tt = Tt == Any ? Float64 : Tt
     transform′ = zeros(Tt, 2*nmodes, 2*nmodes)
     noise1, noise2 = op1.noise, op2.noise
     noise′ = zeros(Tt, 2*nmodes, 2*nmodes)
@@ -279,7 +320,10 @@ function _tensor(op1::GaussianChannel{B1,D1,T1}, op2::GaussianChannel{B2,D2,T2})
     block1, block2 = Base.OneTo(nmodes1), Base.OneTo(nmodes2)
     # initialize direct sum of displacement vectors
     disp1, disp2 = op1.disp, op2.disp
-    Dt = promote_type(eltype(disp1), eltype(disp2))
+    elD1 = eltype(disp1) isa Type ? eltype(disp1) : Float64
+    elD2 = eltype(disp2) isa Type ? eltype(disp2) : Float64
+    Dt = promote_type(elD1, elD2)
+    Dt = Dt == Any ? Float64 : Dt
     disp′ = zeros(Dt, 2*nmodes)
     @inbounds for i in block1
         disp′[i] = disp1[i]
@@ -291,7 +335,10 @@ function _tensor(op1::GaussianChannel{B1,D1,T1}, op2::GaussianChannel{B2,D2,T2})
     end
     # initialize direct sum of transform and noise matrices
     trans1, trans2 = op1.transform, op2.transform
-    Tt = promote_type(eltype(trans1), eltype(trans2))
+    elT1 = eltype(trans1) isa Type ? eltype(trans1) : Float64
+    elT2 = eltype(trans2) isa Type ? eltype(trans2) : Float64
+    Tt = promote_type(elT1, elT2)
+    Tt = Tt == Any ? Float64 : Tt
     transform′ = zeros(Tt, 2*nmodes, 2*nmodes)
     noise1, noise2 = op1.noise, op2.noise
     noise′ = zeros(Tt, 2*nmodes, 2*nmodes)