10
$\begingroup$

What is the Mathematica equivalent of the following Python code with the vectors' broadcast addition?

import numpy as np

a = np.random.rand(5000, 1, 5);
b = np.random.rand(1, 500, 5);

result = a + b #shape: (5000, 500, 5) 
$\endgroup$
5
  • 1
    $\begingroup$ What is the output of this addition? $\endgroup$ Commented Aug 27, 2018 at 12:16
  • $\begingroup$ The output shape is (3,8,5) $\endgroup$ Commented Aug 27, 2018 at 12:17
  • $\begingroup$ I don't know of such a built-in method. $\endgroup$ Commented Aug 27, 2018 at 12:28
  • 2
    $\begingroup$ Maybe one can play around with Outer. $\endgroup$ Commented Aug 27, 2018 at 12:37
  • 1
    $\begingroup$ Outer[Plus, a[[All, 1]], b[[1]], 1] should be fast. broadcastedJIT[Plus, a, b], from my answer, is two times faster on my computer. $\endgroup$
    – jkuczm
    Commented Aug 28, 2018 at 10:08

2 Answers 2

8
$\begingroup$

For your specific case (dimension 1 only in the first two slots), this might work:

a = RandomReal[{0, 1}, {5000, 1, 5}];
b = RandomReal[{0, 1}, {1, 500, 5}];

c1 = Flatten[
 Outer[Plus, a, b, 2],
 {{1, 3}, {2, 4}}
 ]// RepeatedTiming // First

0.255

It is a bit more tedious to use Compile but also a bit faster:

Creating the CompiledFunctions:

cf = Compile[{{a, _Real, 2}},
   Table[Flatten[a, 1], {500}],
   CompilationTarget -> "WVM",
   RuntimeAttributes -> {Listable},
   Parallelization -> True
   ];
cg = Compile[{{b, _Real, 3}},
   Table[Flatten[b, 1], {5000}],
   CompilationTarget -> "WVM",
   RuntimeAttributes -> {Listable},
   Parallelization -> True
   ];

Running the actual code:

c2 = Plus[cf[a], cg[b]]; // RepeatedTiming // First
Max[Abs[c1 - c2]]

0.19

0.

Final remarks

The general case may be treated by a suitable combination of ArrayReshape, MapThread, Outer, and Flatten. Or, maybe even better, by ad-hoc compliled, Listable CompiledFunctions such as cf and cg instead of MapThread. Anyways, one would probably need a thourough case analysis for that.

$\endgroup$
3
$\begingroup$

You can use:

a = RandomReal[{0,1}, {3,5}]
b = RandomReal[{0,1}, {8,5}]
c = Table[a[[i]] + b[[j]], {i, Length[a]}, {j, Length[b]}]

Dimension[c] will be {3,8,5}, similar to result.shape in Python of (3,8,5)

$\endgroup$
3
  • 5
    $\begingroup$ In this simplified case, also Outer[Plus, a, b, 1] works (and should be much faster). $\endgroup$ Commented Aug 27, 2018 at 12:36
  • $\begingroup$ @HenrikSchumacher that is pretty cool! I did not know about Outer[] $\endgroup$
    – Lee
    Commented Aug 27, 2018 at 12:41
  • $\begingroup$ Yes, I've tried that but it's not fast for large arrays $\endgroup$ Commented Aug 27, 2018 at 14:30

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