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)
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 CompiledFunction
s:
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.
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
CompiledFunction
s such as cf
and cg
instead of MapThread
. Anyways, one would probably need a thourough case analysis for that.
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)
Outer[Plus, a, b, 1]
works (and should be much faster).
$\endgroup$
Commented
Aug 27, 2018 at 12:36
Outer[]
$\endgroup$
Outer
. $\endgroup$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$