I wish to write a function foo, which takes, among arguments, a Tuple. Each element in the Tuple can be a scalar or an Array of arbitrary dimension. The scalars and array elements are all of the same type Mytype{Nmy} (Nmy shall be common for all elements in the Tuple).
My function will return a scalar or an Array of scalars, where the scalars are of type Mytype{Nmy}
I need to know N, so that I can write typestable code
using StaticArrays
struct Mytype{N}
x :: Float64
dx :: SVector{N,Float64}
end
Mytype(x,dx::Vector{Float64}) =Mytype(x,SVector{size(dx,1),Float64}(dx))
Base.:(+)(a::Mytype{N},b::Mytype{N}) where {N} = Mytype{N}(a.x+b.x,a.dx+b.dx)
MytypeContainer{N} = Union{Mytype{N},Array{Mytype{N}}}
function foo(args::NTuple{nx,MytypeContainer{N}}) where {N,nx}
s = Mytype{N}(0.,zeros(N))
for arg in args
if isa(arg,Mytype{N})
s += arg
else
for x in arg
s += x
end
end
end
return s
end
v1 = Mytype{2}(3.,[1.,2.])
v2 = Mytype{2}(4.,[5.,6.])
v = [v1,v2]
@show foo((v,))
@show foo((v1,))
@show foo((v1,v2))
@show foo((v,v))
@show foo((v1,v))
In ways I do not understand,
::Tuple{Union{Mytype{N},Array{Mytype{N}}}}
is not expressing what I intend: a tuple of any type of containers of Mytype{N}, and give me N for typestable allocation.
P.S. The code above is now updated and seems to pass the test.