numpy function with float or ndarray as argument

Question

Consider the following code.

import numpy as np

def f(x, y):
    return (x[:, np.newaxis] - y[:]).sum(axis = 1)

x1 = np.linspace(1, 2, 5)
y1 = np.linspace(3, 4, 7)

print(f(x1, y1))

# x2 = 2. # (won't work)
# x2 = [2.] # (won't work either)
x2 = np.asarray([2.])
print(f(x2, y1))

# x3, _ = np.meshgrid(x1, x1) # won't work
# print(f(x3, y1))

The first call f(x1, y1) works as expected.

The second call, the float needs to be "converted" before it works (I'd like the conversion to be transparent for users, so the conversion should be included in the definition of the function).

The third call does not work at all and I don't know how to tune the definition of the function so that it does the same as the first call but for some points on a grid.

Any ideas? Thanks.

xnx · Accepted Answer · 2015-12-01 09:30:58Z

1

You could try the subtraction along the other axis and force x to be an array of a compatible dimension:

import numpy as np

def g(x,y):
    x = np.asarray(x)
    if len(x.shape)==2:
        return (x - y[:,np.newaxis, np.newaxis]).sum(axis=0)
    return (x - y[:,np.newaxis]).sum(axis=0)

x1 = np.linspace(1, 2, 5)
y1 = np.linspace(3, 4, 7)

print('x1 =', x1)
print('y1 =', y1)
print('g(x1,y1) =', g(x1, y1))

x2 = 2.
print('x2 =', x2)
print('y1 =', y1)
print('g(x2, y1) =', g(x2, y1))

x3, _ = np.meshgrid(x1, x1)

print('x3 =', x3)
print('g(x3,y1) =', g(x3, y1))

Output:

x1 = [ 1.    1.25  1.5   1.75  2.  ]
y1 = [ 3.          3.16666667  3.33333333  3.5         3.66666667  3.83333333
  4.        ]
g(x1,y1) = [-17.5  -15.75 -14.   -12.25 -10.5 ]
x2 = 2.0
y1 = [ 3.          3.16666667  3.33333333  3.5         3.66666667  3.83333333
  4.        ]
g(x2, y1) = [-10.5]
x3 = [[ 1.    1.25  1.5   1.75  2.  ]
 [ 1.    1.25  1.5   1.75  2.  ]
 [ 1.    1.25  1.5   1.75  2.  ]
 [ 1.    1.25  1.5   1.75  2.  ]
 [ 1.    1.25  1.5   1.75  2.  ]]
g(x3,y1) = [[-17.5  -15.75 -14.   -12.25 -10.5 ]
 [-17.5  -15.75 -14.   -12.25 -10.5 ]
 [-17.5  -15.75 -14.   -12.25 -10.5 ]
 [-17.5  -15.75 -14.   -12.25 -10.5 ]
 [-17.5  -15.75 -14.   -12.25 -10.5 ]]

answered Dec 1, 2015 at 9:30

xnx

25.7k11 gold badges75 silver badges117 bronze badges

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

1 Comment

cjorssen Over a year ago

Thanks (+1). It works. I was hoping not to have to do some if else but if there is no other options.

Leo · Accepted Answer · 2015-12-01 11:57:20Z

1

To answer the problem of transparent conversion of floats, at least, you could use the np.atleast_1d() method inside your definition:

def f(x, y):
    x_ = np.atleast_1d(x)
    y_ = np.atleast_1d(y)
    return (x_[:, np.newaxis] - y_[:]).sum(axis = 1)

edited Dec 1, 2015 at 11:57

answered Dec 1, 2015 at 9:23

Leo

1,11711 silver badges25 bronze badges

3 Comments

cjorssen Over a year ago

Thanks (+1). Good to know about that method.

Leo Over a year ago

Sorry about the typo. I take it from your edit that you don't want to ensure y is also at least a 1d array?

Leo Over a year ago

Ok, it might answer the question for others though. Would probably be incomplete otherwise.

Collectives™ on Stack Overflow

numpy function with float or ndarray as argument

2 Answers 2

1 Comment

3 Comments

Your Answer

Hot Network Questions

Collectives™ on Stack Overflow

2 Answers 2

1 Comment

3 Comments

Your Answer

Sign up or log in

Post as a guest

Related