I often write equations using numpy that take scalars as inputs and returns another scalar or a vector. Then later I find that I would like to do the same thing, but with one or more vectors as inputs. I'm trying to figure out a way to make one function work in both cases without sprinkling in all sorts of if tests calling np.isscalar() or np.atleast1d() (unless that's the only way).
Is it possible to handle scalar and vector inputs with only one function or am I stuck with multiple implementations?
As an example, here are some functions to convert x, y angles into a north, east, down unit vector (assume the angles are in radians already). I've included a scalar version, vectorized version, and one using np.meshgrid() with calls to the scalar version. I'm trying to avoid doing explicit for looping or calls to np.vectorize().
Scalar Example
import numpy as np
def xy_to_nez_scalar(x, y):
n = np.sin(y)
e = np.sin(x) * np.cos(y)
z = np.cos(x) * np.cos(y)
return np.array([n, e, z])
x = 1
y = 1
nez1 = xy_to_nez_scalar(x, y)
print(f'{nez1=}')
print(f'{nez1.shape=}\n')
producing
nez1=array([0.84147098, 0.45464871, 0.29192658])
nez1.shape=(3,)
Vectorized Example
# X and Y are arrays, so we'll follow the convention that X
# is a column of N rows and Y is a row of M columns. Then we would
# return an array of shape (N, M, 3).
def xy_to_nez_vector(x, y):
n = np.sin(y)
e = np.sin(x[:, np.newaxis]) * np.cos(y)
z = np.cos(x[:, np.newaxis]) * np.cos(y)
# make sure n has the same shape as e and z
nn = np.broadcast_to(n, e.shape)
nez = np.stack([nn, e, z], axis=-1)
return nez
x_array = np.arange(4)
y_array = np.arange(2)
nez2 = xy_to_nez_vector(x_array, y_array)
print(f'{nez2=}')
print(f'{nez2.shape=}\n')
producing
nez2=array([[[ 0. , 0. , 1. ],
[ 0.84147098, 0. , 0.54030231]],
[[ 0. , 0.84147098, 0.54030231],
[ 0.84147098, 0.45464871, 0.29192658]],
[[ 0. , 0.90929743, -0.41614684],
[ 0.84147098, 0.4912955 , -0.2248451 ]],
[[ 0. , 0.14112001, -0.9899925 ],
[ 0.84147098, 0.07624747, -0.53489523]]])
nez2.shape=(4, 2, 3)
Meshgrid Example
# note this produces a (3, M, N) result.
xv, yv = np.meshgrid(x_array, y_array)
nez3 = xy_to_nez_scalar(xv, yv)
print(f'{nez3=}')
print(f'{nez3.shape=}\n')
# fix things by transposing.
nez4 = nez3.T
print(f'{nez4=}')
print(f'{nez4.shape=}\n')
print(np.allclose(nez2, nez4))
producing
nez3=array([[[ 0. , 0. , 0. , 0. ],
[ 0.84147098, 0.84147098, 0.84147098, 0.84147098]],
[[ 0. , 0.84147098, 0.90929743, 0.14112001],
[ 0. , 0.45464871, 0.4912955 , 0.07624747]],
[[ 1. , 0.54030231, -0.41614684, -0.9899925 ],
[ 0.54030231, 0.29192658, -0.2248451 , -0.53489523]]])
nez3.shape=(3, 2, 4)
nez4=array([[[ 0. , 0. , 1. ],
[ 0.84147098, 0. , 0.54030231]],
[[ 0. , 0.84147098, 0.54030231],
[ 0.84147098, 0.45464871, 0.29192658]],
[[ 0. , 0.90929743, -0.41614684],
[ 0.84147098, 0.4912955 , -0.2248451 ]],
[[ 0. , 0.14112001, -0.9899925 ],
[ 0.84147098, 0.07624747, -0.53489523]]])
nez4.shape=(4, 2, 3)
True
CodePudding user response：
This may be bordering on a frame challenge, but I would suggest changing your implementation philosophy slightly to conform to what most numpy functions do already. This has two advantages: (1) experienced numpy users will know what to expect from your functions, and (2) the scalar-vector problems go away.
Normally if faced with a function like xy_to_nez(x, y), I would expect it to take arrays x and y, and return something that has the broadcasted shape of the two, with 3 as either the first or last dimension. The choice of putting 3 in the last dimension is totally fine here. However, magically meshing the arrays together instead of broadcasting them is a rather un-num-pythonic thing to do.
Have the user tell you what you want explicitly (a core tenet of python, item #2 in import this). For example, given a broadcasting interface as suggested above, your scalar function is nearly complete. The only change you would need to make is to stack along the last axis instead of the first, as np.array does:
def xy_to_nez(x, y):
n = np.sin(y)
e = np.sin(x) * np.cos(y)
z = np.cos(x) * np.cos(y)
return np.stack(np.broadcast_arrays(n, e, z), -1)
I would expect the following three examples to work as nez1, nez2 and nez4:
>>> xy_to_nez(1, 1) # shape: 3
array([0.84147098, 0.45464871, 0.29192658])
>>> xy_to_nez(np.arange(4)[:, None], np.arange(2)) # shape: 4, 2
array([[[ 0. , 0. , 1. ],
[ 0.84147098, 0. , 0.54030231]],
[[ 0. , 0.84147098, 0.54030231],
[ 0.84147098, 0.45464871, 0.29192658]],
[[ 0. , 0.90929743, -0.41614684],
[ 0.84147098, 0.4912955 , -0.2248451 ]],
[[ 0. , 0.14112001, -0.9899925 ],
[ 0.84147098, 0.07624747, -0.53489523]]])
>>> xy_to_nez(*np.meshgrid(np.arange(4), np.arange(2), indexing='ij'))
array([[[ 0. , 0. , 1. ],
[ 0.84147098, 0. , 0.54030231]],
[[ 0. , 0.84147098, 0.54030231],
[ 0.84147098, 0.45464871, 0.29192658]],
[[ 0. , 0.90929743, -0.41614684],
[ 0.84147098, 0.4912955 , -0.2248451 ]],
[[ 0. , 0.14112001, -0.9899925 ],
[ 0.84147098, 0.07624747, -0.53489523]]])
In the third example, I corrected the problem with nez3 by passing indexing='ij' to np.meshgrid. The issue there was not with vectorization, but with the shape of the meshgrid you were passing in.
I would not expect xy_to_nez(np.arange(4), np.arange(2)) to work at all: the arrays don't broadcast together and we shouldn't try to figure out how to combine them. To see why, pretend that they had three random dimensions each. Do you interleave? Do you place one set after the other? What if some of the dimensions broadcast but others don't? Leave it up to the user to eliminate those questions from consideration.
At the same time:
>>> xy_to_nez(np.arange(4), np.arange(4)) # Shape: 4, 3
array([[ 0. , 0. , 1. ],
[ 0.84147098, 0.45464871, 0.29192658],
[ 0.90929743, -0.37840125, 0.17317819],
[ 0.14112001, -0.13970775, 0.98008514]])