9
$\begingroup$

One should expect that the implementation of the built-in function Mean was quick. However, as I recently observed, this is not true at all, in particular for matrices. Just compare Mean to the following simple function:

mean[A_?TensorQ] := ConstantArray[1./Length[A], Length[A]].A

A test with a decently sized matrix:

A = RandomReal[{-1, 1}, {4000, 5000}];
a = Mean[A]; // AbsoluteTiming//First
b = mean[A]; // AbsoluteTiming//First
Max[Abs[a - b]]

(* 0.220566 *)
(* 0.008163 *)
(* 1.71738*10^-16 *)

Wow, for a built-in function, this is really, really bad. For higher rank matrices, this gap is not as pronounced but it is still there:

A = RandomReal[{-1, 1}, {400, 500, 300}];
a = Mean[A]; // AbsoluteTiming // First
b = mean[A]; // AbsoluteTiming // First
Max[Abs[a - b]]

(* 0.110258 *)
(* 0.024915 *)
(* 2.22045*10^-16 *)

Okay, for the first (and also for the last slot), there is this workaround by using Dot. What about a slot in between? E.g. how to compute Total[A,{2}]/Dimensions[A][[2]] as fast as possible?

PS.: The same mess appears with Total, a function that I usually value quite high, in particular since Total allows for level specs.

$\endgroup$
  • 1
    $\begingroup$ Using exact values in the targets and in the function reverses the performance dramatically. Hypothesis: I'd venture the dot with machine precision takes a path using Intel MKL, while Mean/Total, being built to handle general forms (e.g. symbolic, mixed) dont. $\endgroup$ – ciao Oct 22 '17 at 8:43
  • $\begingroup$ @ciao You are totally right about exact arithmetic. I really aim for PackedArrays. What I wanted to point out is that it might be a good idea to overload Mean for PackedArrays to something that actually uses Intel MKL. $\endgroup$ – Henrik Schumacher Oct 22 '17 at 9:15
7
$\begingroup$

Here is a compiled version of Total that is faster for intermediate levels:

total[d_, l_] := With[
    {o = ConstantArray[1., Dimensions[d][[l]]], n = Length@Dimensions@d-l+1},
    Switch[l,
        1, o.d,

        Length@Dimensions@d, d.o,

        _,
        With[{fc = Compile[{{data, _Real, n}}, o.data, RuntimeAttributes->{Listable}]},
            fc[d]
        ]
    ]
]

For your higher rank example:

d = RandomReal[{-1, 1}, {400, 500, 300}];

r1 = total[d, 2]; //RepeatedTiming
r2 = Total[d, {2}]; //RepeatedTiming

r1 == r2

{0.0632, Null}

{0.11, Null}

True

I will leave the modification to compute the mean instead to you.

$\endgroup$
  • $\begingroup$ Shiny! I wasn't aware that we can control the threading level of listable compiled functions this way. Thank you, this answers also several further questions! =o) $\endgroup$ – Henrik Schumacher Oct 23 '17 at 7:56
3
$\begingroup$

First, don't use TensorQ. It is undocumented and may change in the future. The documented function to use is ArrayQ. Second, here an arbitrary-rank version of your function,

mean[A_?ArrayQ, slot_] :=  With[
   {d = Part[Dimensions[A], slot]}, 
   ConstantArray[1./d, d].Transpose[A, 1 <-> slot]
]

Unfortunately, it is slower than Total for higher-rank, because the cost of rearranging memory is too high. I suspect that Carl's approach is probably the best you can do, but I can't immediately prove that.

$\endgroup$
  • $\begingroup$ Thanks for the hint on ArrayQ. And you are right, transposing is not an option... $\endgroup$ – Henrik Schumacher Oct 24 '17 at 5:58

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.