I have a situation where an array of objects, say galaxies, can have either the same value for a property, or each has a different value. I would like to use the same function in both cases. For example, if the property is weights,
function do_something(galaxies::Array; weights=1)
output_array = zeros(size(galaxies))
for i=1:length(galaxies)
w = weights[i] # doesn't work if weights is a scalar
# ... more complicated stuff here
output_array[i] += w
end
return output_array
end
galaxies = rand(5)
do_something(galaxies, weights=1)
do_something(galaxies, weights=rand(length(galaxies)))
So every galaxy by default has the same weight 1, but it will also be common to give an array of weights, one for each galaxy.
What is a good way to make weights work with both a scalar and an array?
As a potential solution, I was thinking of creating a special ArrayScalar struct that returns a constant no matter what indices are given to it. But it seems like overkill.
Another option would be to turn the scalar into weights = ones(length(weights)). But that seems wasteful.
This is also a classic use of Julia’s multiple dispatch - you want a function that behaves slightly different depending on the type of a parameter.
One way would be to factor out the weight access:
get_weight(weights::Integer, index) = weights
get_weight(weights::Vector, index) = weights[index]
function do_something(galaxies::Array; weights=1)
output_array = zeros(size(galaxies))
for i=1:length(galaxies)
w = get_weight(weights, i)
# ... more complicated stuff here
output_array[i] += w
end
return output_array
end
The nice thing about this is there is no performance cost, the compiler will produce two versions of this depending on the type of weights and call the appropriate one. And no data overhead.
Actually, on reflection this is also possibly a good use of broadcasting. Julia’s broadcasting already handles the case of when one argument is either a vector or a scalar so you could rewrite as:
function do_something(galaxy, weight)
# ... more complicated stuff here
return calculated_value
end
# supply default value for weight
do_something(galaxy) = do_something(galaxy, 1)
galaxies = rand(5)
single_weight = do_something.(galaxies, 1)
multiple_weights = do_something.(galaxies, rand(5))