[Numpy-discussion] What is the pythonic way to write a function that handles arrays and scalars?

Eric Wieser wieser.eric+numpy at gmail.com
Tue Dec 12 15:52:52 EST 2017


Using np.iscalar is a bad idea, as it fails for 0d arrays. x.ndim is the
better option there.

I’d maybe suggest not special-casing 0d arrays though, and using:

def func_for_scalars_or_vectors(x):
    x = np.asanyarray(x) # convert scalars to 0d arrays

    # The magic happens here

    return ret[()]  # convert 0d arrays to scalars

Eric

On Tue, 12 Dec 2017 at 11:58 Robert Kern <robert.kern at gmail.com> wrote:

On Wed, Dec 13, 2017 at 4:50 AM, Joe <solarjoe at posteo.org> wrote:
> >
> > Hi,
> >
> > the best example I found was this one:
> >
> > https://stackoverflow.com/a/29319864/7919597
> >
> > def func_for_scalars_or_vectors(x):
> >     x = np.asarray(x)
> >     scalar_input = False
> >     if x.ndim == 0:
> >         x = x[None]  # Makes x 1D
> >         scalar_input = True
> >
> >     # The magic happens here
> >
> >     if scalar_input:
> >         return np.squeeze(ret)
> >     return ret
> >
> > Is this as good as it gets or do you have other suggestions?
>
> In general structure, yeah. I'd probably use `np.isscalar(x)` for the test
> and `x = np.atleast_1d(x)` for the coercion for readability, but otherwise,
> that's it.
>
> --
> Robert Kern
> _______________________________________________
> NumPy-Discussion mailing list
> NumPy-Discussion at python.org
> https://mail.python.org/mailman/listinfo/numpy-discussion
>
​
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/numpy-discussion/attachments/20171212/16249a51/attachment.html>


More information about the NumPy-Discussion mailing list