Scalar and Arrays: indexing the same way?

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.

Any other ways of dealing with this?

1 Like

There is a package called FillArrays.jl:

That was fast! Thanks!

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.

3 Likes

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))

this is also possibly a good use of broadcasting

Unfortunately, the more complicated stuff here doesn’t lend itself easily to broadcasting.

1 Like