Broadcast matrix valued function, output into existing array

General Usage
,
1

How do I broadcast the function matop below such that the output car1 resembles car, i.e., car1 is an array which can be accessed as car1[:,1,2]?
Presently, car1 looks like a vector, with size (N,), instead of (N,2,2).

using MKL
using LinearAlgebra

function matop(a, b)
    return a .* [b a; a b];
end

N = 100;
aar = range(0,0.1,N);
bar = range(0,0.05,N);

car = zeros(Float64,N,2,2);
for hi = 1:N
    car[hi,:,:] = matop(aar[hi],bar[hi]);
end

# car1 = matop.(aar,bar); #??
2

You could just use stack:

julia> car1 = stack(matop.(aar, bar), dims=1);

julia> car == car1
true

If you want to output into an existing array (cf. title)

julia> car2 = similar(car);

julia> broadcast!(matop, eachslice(car2, dims=1), aar, bar);

julia> car2 == car
true

is an option.

3
  1. This is the wrong order for locality
  2. Working with lots of 2x2 arrays is an ideal case for StaticArrays — instead of a 3d array, use a 1d array of 2x2 SMatrix’s. This will also fix (1).

Using an array of SMatrix also makes it easier to use broadcasting. For example:

using StaticArrays

function matop(a, b)
    return a * @SMatrix[b a; a b];
end

N = 100;
aar = range(0,0.1,N);
bar = range(0,0.05,N);
car = matop.(aar, bar)

and then you can access elements as e.g. car[i][1,2]. This is much more efficient than your previous implementation (especially if you put it into a function to avoid globals), because it no longer heap-allocates an Array on every call to matop.

4 Likes
4

That being said, especially if you are coming from Matlab or Python or R, you should re-think the urge to break every computation into a sequence of discrete, simple, “vectorized” operations. As much as possible you want to combine operations into a single loop and avoid creating temporary arrays for intermediate values.

Here, why are you allocating an array of simple 2x2 matrices like this? Why not construct them, use them, then discard them without ever storing them in a big car array? (Using StaticArrays for things like 2x2 matrix computations is still a good idea, however.)

Loops are fast in Julia. You don’t have to “vectorize” everything.

6 Likes
5

Thanks, makes sense.
I needed the same batch of matrices multiple times, hence I decided to store and reuse them. Otherwise, for one-time use, constructing them on the fly would be ideal.
Also, StaticArrays looks like the way to go.

using MKL
using LinearAlgebra
using StaticArrays
using BenchmarkTools

N = 10000000;
aar = range(0,0.1,N);
bar = range(0,0.05,N);

function matop(a, b)
    return a .* Array([b a; a b]);
end


function basic(a, b)
    c = zeros(Float64,2,2,length(a));
    for hi = 1:length(a)
        c[:,:,hi] = matop(a[hi],b[hi]);
    end
    return c
end
@btime car = basic($aar, $bar);

function broadcast1(a, b)
    c = zeros(Float64,2,2,length(a));
    broadcast!(matop, eachslice(c, dims=3), a, b);
    return c
end
@btime car1 = broadcast1($aar, $bar);

function broadcast2(a, b)
    c = stack(matop.(a, b), dims=3);
    return c
end
@btime car2 = broadcast2($aar, $bar);

function matops(a, b)
    return a * @SMatrix[b a; a b];
end
function broadcast3(a, b)
    c = matops.(a, b);
    return c
end
@btime car3s = broadcast3($aar, $bar);

car = basic(aar, bar);
car1 = broadcast1(aar, bar);
car2 = broadcast2(aar, bar);
car3s = broadcast3(aar, bar);
car3 = zeros(Float64,2,2,N);
for hi = 1:N
    car3[:,:,hi] = car3s[hi][1:2,1:2];
end

car≈car1
car≈car2
car≈car3

 3.048 s (100000003 allocations: 4.02 GiB)
 2.923 s (100000003 allocations: 4.02 GiB)
 4.251 s (100000006 allocations: 4.10 GiB)
 75.003 ms (3 allocations: 305.18 MiB)
6

Have you tried:

function matop(a, b)
    return a * @SMatrix [b a; a b];
end
car4 = Array{Float64}(undef, N, 2, 2)
map!(matop, eachslice(car4; dims=1), aar, bar)
7

Here you risk creating two arrays, first [a b a; a b], and then a second one when you pass it to the Array constructor. Also, it’s unidiomatic. Just write [b a; a b].

8

And if you want to optimize further, you can write the above as

[a*b a^2; a^2 a*b]

to further reduce allocations.

BTW, trailing semicolons inside functions etc are not needed in Julia. The do no particular harm either, but you may find that style guides recommend against them.