How to preserve custom type after broadcasting?

General Usage
1

I’d like to preserve the custom type after it has been broadcasted. Currently I have the following:

struct SigmaPoints{T} <: AbstractMatrix{T}
    x0::Array{T, 1}
    xi::Array{T, 2}
end
Base.size(S::SigmaPoints) = (size(S.xi, 1), size(S.xi, 2) + 1)
Base.getindex(S::SigmaPoints{T}, inds::Vararg{Int,2}) where {T} = if inds[2] == 1 S.x0[inds[1]] else S.xi[inds[1], inds[2] - 1] end
Base.setindex!(S::SigmaPoints{T}, val, inds::Vararg{Int,2}) where {T} = if inds[2] == 1 S.x0[inds[1]] = val else S.xi[inds[1], inds[2] - 1] = val end

SigmaPoints is essentially a matrix where the first column is saved in x0 instead of xi

julia> S = SigmaPoints(ones(2), zeros(2,4))
4×5 SigmaPoints{Float64}:
 1.0  0.0  0.0  0.0  0.0
 1.0  0.0  0.0  0.0  0.0
julia> S .+ ones(2)
4×5 Array{Float64,2}:
 2.0  1.0  1.0  1.0  1.0
 2.0  1.0  1.0  1.0  1.0

Mathematically that’s correct. However, I’d like to preserve the SigmaPoints type for basic operations (e.g. multiplications, additions, etc) with scalars and vectors.
I have read the manual, but I am struggling to map the example to this case.

2

Define BroadcastStyle and similar for your type:

Base.BroadcastStyle(::Type{<:SigmaPoints}) = Broadcast.ArrayStyle{SigmaPoints}()

Base.similar(bc::Broadcast.Broadcasted{Broadcast.ArrayStyle{SigmaPoints}}, ::Type{T}) where T =
    SigmaPoints(similar(bc.args[1].x0), similar(bc.args[1].xi))

And then:

julia> SigmaPoints(ones(2), zeros(2,4)) .+ ones(2)
2×5 SigmaPoints{Float64}:
 2.0  1.0  1.0  1.0  1.0
 2.0  1.0  1.0  1.0  1.0
4 Likes