Systematically turn (numpy) 1d-array of size 1 to scalar

Question

I have been scratching my head over this for the last hour.

Imagine I have a function that takes as argument a number, or an array up to dimension 1. I'd like it to return a scalar (not a 0d array) in the former case, and an array of the same shape in the latter case, like ufuncs do.

The current implementation of my function does something along the lines of

def func(x: Real | np.ndarray, arr: np.ndarray):
    """Illustrate an actually more complicated  function"""
    return arr @ np.sin(arr[:, None] * x)

and I know arr is a 1d array. It is promoted to 2d so that there is no broadcast issue with the element-wise multiplication. The issue being that a 1d array is systematically returned. The nicety being that the cases

x scalar and len(arr) == 1;
x scalar and len(arr) > 1 ;
x array and len(arr) == 1 ;
x array and len(arr) >= 1 and (len(x) != len(arr) or len(x) == len(arr))

are covered in a one-liner (the tautological last statement is here for emphasis).

I tried

@np.vectorize
def func(x, arr):
    return arr @ np.sin(arr * x)

which consistently returns a 1d array as well, and is probably not the best performance-wise. I looked at functools.singledispatch, which leads to much duplication, and I'd probably forget about some corner cases. A solution would be

def func(x, arr):
    res = arr @ np.sin(arr[:, None] * x)
    if len(res) == 1:
        return res.item()
    return res

but I have many such functions, and it doesn't feel very pythonic? I can write a decorator to write this check,

def give_me_a_scalar(f):
    @functools.wraps(f)
    def wrapper(*args, **kwargs):
        res = f(*args, **kwargs)
        if len(res) == 1:
            return res.item()
        return res
    return wrapper

seems to be doing what I want, but I struggle to believe nothing like that already exists. Am I missing something simpler?

You don't want return np.squeeze(arr @ np.sin(arr[:, None] * x))? — mozway
– mozway, Commented Jul 2, 2024 at 14:43
@mozway A squeezed 1d array of length 1 unfortunately becomes a 0d array, not a scalar. But thank you, it made me realise that another option would be arr @ np.atleast_1d((arr[:, None] * x).squeeze()); which is a bit convoluted, maybe? — Aubergine
– Aubergine, Commented Jul 2, 2024 at 15:12
After struggling with this problem a number of times I acquired a taste for 0d arrays. — aerobiomat
– aerobiomat, Commented Jul 2, 2024 at 15:22

simon · Accepted Answer · 2024-07-03 16:51:04Z

2

The following seems to work; it is mainly based on this answer to a related question, combining it with @mozway's comment:

def func(x, arr):
    res = arr @ np.sin(arr[:, None] * x)
    return res.squeeze()[()]

That is, after squeezing, you can index into the result with the empty tuple (). What is essential here is that squeeze() will turn a single-valued array into a 0-d array, and the empty index will turn the latter into a scalar¹ (just like item()), but neither squeeze() nor the empty index will alter the shape of the result if it is more than single-valued (or raise an exception, as item() does).

It is still a bit ugly though. Another potential downside: for the scalar case, unlike item(), the result still has a Numpy data type (so e.g. it is a np.float64 scalar rather than a float scalar) – but it is a scalar rather than a 0-d array nevertheless.

If you want to hide the ugliness, you could still use it as a decorator. I have no solution for the data type though (although this might be a non-issue).

Update: I just realized, this approach is even used in an example in the documentation of np.squeeze(), so I guess it could be considered "canonical" (bottom-most example there):

>>> x = np.array([[1234]])
>>> x.shape
(1, 1)
>>> np.squeeze(x)
array(1234)  # 0d array
>>> np.squeeze(x).shape
()
>>> np.squeeze(x)[()]
1234

The last line will actually print np.int64(1234) rather than 1234 in Numpy 2.0, owing to the altered representation of Numpy scalars.

_{¹) Also see the Numpy doc: An empty (tuple) index is a full scalar index into a zero-dimensional array. x[()] returns a scalar if x is zero-dimensional and a view otherwise.}

edited Jul 3, 2024 at 16:51

answered Jul 2, 2024 at 15:38

simon

6,8642 gold badges18 silver badges33 bronze badges

Sign up to request clarification or add additional context in comments.

2 Comments

Aubergine Over a year ago

Good catch! I saw that thread but missed this particular answer, and I had stopped reading the documentation of squeeze at the Returns section. And I didn't know that indexing with an empty tuple was a thing. I ran some tests with the size of x ranging from 1 to 1e5, and it is consistently slower than the test on length + item solution by a few percents, but it is more compact.

simon Over a year ago

@Aubergine I guess I would probably go with the decorator then and use the length check + item solution. Or you could use item in a try-except block. Unlike you wrote in an earlier comment, personally I would not consider that an abuse of exception handling in Python. See for example here: Easier to ask for forgiveness than permission. … This clean and fast style is characterized by the presence of many try and except statements. It might still be slower than length check + item though, which, I guess, is your main criterion.

Collectives™ on Stack Overflow

Systematically turn (numpy) 1d-array of size 1 to scalar

1 Answer 1

2 Comments

Your Answer

Linked

Hot Network Questions

Collectives™ on Stack Overflow

1 Answer 1

2 Comments

Your Answer

Sign up or log in

Post as a guest

Linked

Related