Automatically handle arrays and scalars for item assignment

Question

I have a Python function that usually handles numpy arrays, and it does something like:

result = a*x
result[a<=0] = float('NaN')

However, sometimes x and/or a are just float numbers and when this happens this code fails. To deal with this situation I am currently doing this:

x = np.asarray(x).astype(float)
a = np.asarray(a).astype(float)
result = a*x
if a.shape == ():
    if a <= 0:
        result *= float('NaN')
else:
    result[a<=0] = float('NaN')

This does exactly what I want to do, but it requires code pseudo-replication and when we have many parameters a, b, c,... or more complicated conditions it is prone to errors.

I have tried

x = np.asarray(x).astype(float)
a = np.asarray(a).astype(float)
result = a*x
result[a<=0] = float('NaN')

but it produces

TypeError: 'numpy.float64' object does not support item assignment

Is there a better way of doing this? I would like to keep the notation result[a<=0]=float('NaN') and if a is a scalar apply it to all elements of result without having to manually check if it is dimension 0.

Answer 1

It sounds like you want numpy scalars to behave like one-element 1d arrays, and everything else behave normally (the problem is that numpy arrays are mutable, but scalars are not). Semantically speaking the correct tool is np.atleast_1d():

# dummy input
a = np.array(1.0)
x = np.array(42.0)

result = np.atleast_1d(a*x)
result[a <= 0] = np.nan

This gives you

>>> result
array([42.])

which sounds like something you want to end up with. If for some reason you want the result to be a scalar instead (which I do not recommend, as this would set you up for similar kind of loops to jump through downstream) you can convert the result with multiple approaches, but exactly how you do that would depend on the expected behaviour you want (considering edge cases such as 1d-vs-0d single-element cases).