Let's say I have a conditional calculation like this
import numpy as np
def foo(a, b, c):
if a <= 0 or c == 0:
return foo_special(a, b, c)
return foo_normal(a, b, c)
def foo_normal(a, b, c):
return np.log(a) * b / c
def foo_special(a, b, c):
return a + b + c
and I want to make foo work with arrays. A naive approach would be
def foo_naive(a, b, c):
out = np.zeros_like(a)
special = (a <= 0) | (c == 0)
out[special] = foo_special(a[special], b[special], c[special])
out[~special] = foo_normal(a[~special], b[~special], c[~special])
return out
but this assumes that a, b and c have the same shape. However, normally, numpy allows me to mix scalars with arrays as well as arrays of different shapes, e.g.
foo_normal(
4,
np.array([1, 2, 3]),
np.array([[5, 6, 7]]).T,
)
works just fine (returning a 3 x 3 array) and I want to keep that behavior.
Thus, what I need to do (I think), is to bring all input arrays into the same, final shape before giving them to foo_naive. How can I do this?