Support for arrays in addition to tuples in LinearCombination

[oschulz]

We use LinearMaps extensively in MGVI.jl. We're adding support for Newton-CG-optimization now which needs means of linear maps. The number of maps in the mean/sum is not fixed though, so we'd like to use an array of maps instead of a tuple in LinearCombination. Would a pull request to extend LinearCombination that way be welcome?

[dkarrasch]

I'm glad to hear LinearMaps.jl is used. @jagot took a similar approach in #84, but then didn't finish because it didn't do the job he needed, but you may find some inspiration there. On the other hand, mean (almost) already works.

using LinearMaps, Statistics
A = LinearMap(rand(3,3))
B = A + A + A
julia> mean(B.maps)
3×3 LinearMaps.ScaledMap{Float64} with scale: 0.3333333333333333 of
  3×3 LinearMaps.LinearCombination{Float64} with 3 maps:
    3×3 LinearMaps.WrappedMap{Float64} of
      3×3 Matrix{Float64}
    3×3 LinearMaps.ScaledMap{Float64} with scale: 1.0 of
      3×3 LinearMaps.WrappedMap{Float64} of
        3×3 Matrix{Float64}
    3×3 LinearMaps.WrappedMap{Float64} of
      3×3 Matrix{Float64}

The fact that the maps form a tuple shouldn't be a big restriction, because these tuples are all lazy, i.e., adding another map to an existing LinearCombination forms a new tuple (of pointers to the maps) of lenght one bigger than before. That should be really cheap and fast. If you have loads of maps, then tuples are no longer efficient, which was @jagot's motivation to generalize to vectors IIRC. Otherwise, we could add a method like

Statistics.mean(A::LinearCombination) = mean(A.maps)

which would add another dependency just for the sake of this one method, but that's not too bad. Oh, but wait, what would change if the maps field was a vector instead of a tuple? Another option is to just use the above definition in your code, i.e., if you know that you will have a B::LinearCombination at some point, then just hard-code mean(B.maps) instead of some generic mean(B), which wouldn't be generic after all because it doesn't make sense for any other map type. Let me know what you think.

[oschulz]

because these tuples are all lazy, i.e., adding another map to an existing LinearCombination forms a new tuple (of pointers to the maps) of lenght one bigger than before

Yes, the problem is the load on the compiler, with new data types created all the time for tuples of different length (and n can be hundreds of even thousands in our use case). Mean itself works fine indeed already.

Statistics.mean(A::LinearCombination) = mean(A.maps)

Yes, forgot to mention, I would like to specialize mean and sum. And from what I can see LinearCombination should work just fine if we allow As<:Union{LinearMapTuple,AbstractArray{<:LinearMap} or similar.

Just wanted to check with you before preparing a PR.

[oschulz]

I'm glad to hear LinearMaps.jl is used.

It's awesome! We found it to be a very powerful abstraction. We use it for lazy Jacobians too, with both left- and right-multiplication (backed by ForwardDiff and Zygote, by default).

[oschulz]

CC MGVI devs @apmypb, @jknollm and @vindex10

[dkarrasch]

Before you start with your own PR, check out #84. It does a few more things (to block concatenation), but it already has the LinearCombination extension that you want, including some adoption of the mul! machinery. So you could play with it and see how things work.

[oschulz]

Thanks, will do! Should I just base a PR on top of #84?

[dkarrasch]

That was just a suggestion for a quick start. It's not important which way you go. #84 has more things that we might keep for when there is actual interest in it.

[oschulz]

Ok, I'll start with LinearCombination and specialized sum then.