Add broadcasting support for obj .= scalar (#202)

* add `copyto!` methods for `DefaultArrayStyle{0}`

* add simple `.= 1.0` tests for ops and states

* have fallback error for obscure broadcasting

Author: apkille

src/operators_dense.jl

@@ -455,6 +455,17 @@ Base.iterate(a::DataOperator, idx) = iterate(a.data, idx)
     end
     return dest
 end
+# same implementation as Base.copyto!(dest::AbstractArray, bc::Broadcasted{<:AbstractArrayStyle{0}}) in https://github.com/JuliaLang/julia/blob/master/base/broadcast.jl
+@inline function Base.copyto!(dest::DataOperator{BL,BR}, bc::Broadcast.Broadcasted{Style,Axes,F,Args}) where {BL,BR,Style<:Broadcast.DefaultArrayStyle{0},Axes,F,Args}
+    # Typically, we must independently execute bc for every storage location in `dest`, but:
+    # IF we're in the common no-op identity case with no nested args (like `dest .= val`),
+    if bc.f === identity && bc.args isa Tuple{Any} && Broadcast.isflat(bc)
+        # THEN we can just extract the argument and `fill!` the destination with it
+        return fill!(dest, bc.args[1][])
+    else
+        throw(ArgumentError("no fallback implementation has been defined outside of dest .= val."))
+    end
+end
 @inline Base.copyto!(A::DataOperator{BL,BR},B::DataOperator{BL,BR}) where {BL,BR} = (copyto!(A.data,B.data); A)
 @inline Base.copyto!(dest::DataOperator{BL,BR}, bc::Broadcast.Broadcasted{Style,Axes,F,Args}) where {BL,BR,Style<:DataOperatorStyle,Axes,F,Args} =
     throw(IncompatibleBases())

src/states.jl

@@ -227,6 +227,17 @@ Base.@propagate_inbounds Base.Broadcast._broadcast_getindex(x::T, i) where {T<:U
     end
     return dest
 end
+# same implementation as Base.copyto!(dest::AbstractArray, bc::Broadcasted{<:AbstractArrayStyle{0}}) in https://github.com/JuliaLang/julia/blob/master/base/broadcast.jl
+@inline function Base.copyto!(dest::Ket{B}, bc::Broadcast.Broadcasted{Style,Axes,F,Args}) where {B,Style<:Broadcast.DefaultArrayStyle{0},Axes,F,Args}
+    # Typically, we must independently execute bc for every storage location in `dest`, but:
+    # IF we're in the common no-op identity case with no nested args (like `dest .= val`),
+    if bc.f === identity && bc.args isa Tuple{Any} && Broadcast.isflat(bc)
+        # THEN we can just extract the argument and `fill!` the destination with it
+        return fill!(dest, bc.args[1][])
+    else
+        throw(ArgumentError("no fallback implementation has been defined outside of dest .= val."))
+    end
+end
 @inline Base.copyto!(dest::Ket{B1}, bc::Broadcast.Broadcasted{Style,Axes,F,Args}) where {B1,B2,Style<:KetStyle{B2},Axes,F,Args} =
     throw(IncompatibleBases())
 
@@ -240,6 +251,17 @@ end
     end
     return dest
 end
+# same implementation as Base.copyto!(dest::AbstractArray, bc::Broadcasted{<:AbstractArrayStyle{0}}) in https://github.com/JuliaLang/julia/blob/master/base/broadcast.jl
+@inline function Base.copyto!(dest::Bra{B}, bc::Broadcast.Broadcasted{Style,Axes,F,Args}) where {B,Style<:Broadcast.DefaultArrayStyle{0},Axes,F,Args}
+    # Typically, we must independently execute bc for every storage location in `dest`, but:
+    # IF we're in the common no-op identity case with no nested args (like `dest .= val`),
+    if bc.f === identity && bc.args isa Tuple{Any} && Broadcast.isflat(bc)
+        # THEN we can just extract the argument and `fill!` the destination with it
+        return fill!(dest, bc.args[1][])
+    else
+        throw(ArgumentError("no fallback implementation has been defined outside of dest .= val."))
+    end
+end
 @inline Base.copyto!(dest::Bra{B1}, bc::Broadcast.Broadcasted{Style,Axes,F,Args}) where {B1,B2,Style<:BraStyle{B2},Axes,F,Args} =
     throw(IncompatibleBases())
 

test/test_operators_dense.jl

@@ -382,6 +382,9 @@ op1 .= op1_ .+ 3 * op1_
 bf = FockBasis(3)
 op3 = randoperator(bf)
 @test_throws QuantumOpticsBase.IncompatibleBases op1 .+ op3
+op .= 1.0
+@test op == fill!(zero(op), 1.0)
+@test_throws ArgumentError op .= 1.0 .+ 1.0
 
 # Dimension mismatches
 b1, b2, b3 = NLevelBasis.((2,3,4))  # N is not a type parameter

test/test_states.jl

@@ -166,6 +166,12 @@ psi_ .+= psi123
 bra_ = copy(bra123)
 bra_ .= 3*bra123
 @test bra_ == 3*dagger(psi123)
+bra .= 1.0
+ket .= 1.0
+@test bra == Bra(bf, [1.0 + 0.0im, 1.0 + 0.0im])
+@test ket == Ket(bf, [1.0 + 0.0im, 1.0 + 0.0im])
+@test_throws ArgumentError bra .= 1.0 .+ 1.0
+@test_throws ArgumentError ket .= 1.0 .+ 1.0
 
 end # testset