How to fix broadcasting issues with numpy.vectorize() - python

I am writing a custom function that I want to have behaving as if it where a numpy-function, having the ability to take in an array, and perform the same operation on every element of the input list, and returning a list of same shape with all the results.
Luckily, there is a solution for this: numpy.vectorize()
So I used that: I have a function that creates a model in the form of a sine wave, and it takes in two variables: one numpy list X containing some input values for the sine function, and one numpy list param that contains the four possible parameters that a sine curve can have.
import numpy as np
def sine(X, param):
#Unpacking param
A = param[0]
P = param[1]
Phi = param[2]
B = param[3]
#translating variables
#Phi = t0/P
f = X/P
Y = A*np.sin(2*np.pi*(f + Phi)) + B
return Y
Because only the input values X need the broadcasting while the all the parameters are necessary all the time, so, according to the documentation, the way to vecorise the function is as follows:
np_sine = np.vectorize(sine, excluded=['param']) #makes sine() behave like a numpy function
...so that param is properly excluded from vectorisation.
This method is useful, since I will be fitting this model to a dataset, which requires occasionally tweaking the parameters, meanwhile, with this method the code where I need it is only one line long:
CHIsqrt = np.sum(((ydata - np_sine(xdata, param))/yerr)**2)
where ydata, xdata and yerr are equally long lists of datapoints and where param is the list of four parameters.
Yet, the result was a broadcasting error:
File "C:\Users\Anonymous\AppData\Local\Programs\Python\Python36\lib\site-packages\numpy\lib\function_base.py", line 2831, in _vectorize_call outputs = ufunc(*inputs)
ValueError: operands could not be broadcast together with shapes (500,) (4,)
Since the list param is 4 elements long, I get that the function ignored my commands to exclude it from vectorisation. That is a problem.
I tried specifying that the end result should be a ndArray, which did not change the error.
np_sine = np.vectorize(sine, excluded=['param'], otypes=[np.ndarray])
What would be the correct way to use this function?

You've specified excluded wrong.
In [270]: def foo(x, param):
...: a,b,c = param
...: return a*x
...:
In [271]: f = np.vectorize(foo, excluded=[1]) # specify by position
In [272]: f(np.arange(4),[1,3,2])
Out[272]: array([0, 1, 2, 3])
For a keyword arg:
In [277]: def foo(x, param=[0,0,0]):
...: a,b,c = param
...: return a*x
...:
In [278]: f = np.vectorize(foo, excluded=['param'])
In [279]: f(np.arange(4),param=[1,3,2])
Out[279]: array([0, 1, 2, 3])

Related

Jax vectorization: vmap and/or numpy.vectorize?

what are the differences between jax.numpy.vectorizeand jax.vmap?
Here is a small snipset
import jax
import jax.numpy as jnp
def f(x):
return jnp.exp(-x)*jnp.sin(x)
gf = jax.grad(f)
x = jnp.arange(0,1,0.1)
jax.vmap(gf)(x)
jnp.vectorize(gf)(x)
Both computations give the same results:
DeviceArray([ 1. , 0.80998397, 0.63975394, 0.4888039 ,
0.35637075, 0.24149445, 0.14307144, 0.05990037,
-0.00927836, -0.06574923], dtype=float32)
How to decide which one to use, and if there is a difference in terms of performance?
jax.vmap and jax.numpy.vectorize have quite different semantics, and only happen to be similar in the case of a single 1D input as in your example.
The purpose of jax.vmap is to map a function over one or more inputs along a single explicit axis, as specified by the in_axes parameter. On the other hand, jax.numpy.vectorize maps a function over one or more inputs along zero or more implicit axes according to numpy broadcasting rules.
To see the difference, let's pass two 2-dimensional inputs and print the shape within the function:
import jax
import jax.numpy as jnp
def print_shape(x, y):
print(f"x.shape = {x.shape}")
print(f"y.shape = {y.shape}")
return x + y
x = jnp.zeros((20, 10))
y = jnp.zeros((20, 10))
_ = jax.vmap(print_shape)(x, y)
# x.shape = (10,)
# y.shape = (10,)
_ = jnp.vectorize(print_shape)(x, y)
# x.shape = ()
# y.shape = ()
Notice that vmap only maps along the first axis, while vectorize maps along both input axes.
And notice also that the implicit mapping of vectorize means it can be used much more flexibly; for example:
x2 = jnp.arange(10)
y2 = jnp.arange(20).reshape(20, 1)
def add(x, y):
# vectorize always maps over all axes, such that the function is applied elementwise
assert x.shape == y.shape == ()
return x + y
jnp.vectorize(add)(x2, y2).shape
# (20, 10)
vectorize will iterate over all axes of the inputs according to numpy broadcasting rules. On the other hand, vmap cannot handle this by default:
jax.vmap(add)(x2, y2)
# ValueError: vmap got inconsistent sizes for array axes to be mapped:
# arg 0 has shape (10,) and axis 0 is to be mapped
# arg 1 has shape (20, 1) and axis 0 is to be mapped
# so
# arg 0 has an axis to be mapped of size 10
# arg 1 has an axis to be mapped of size 20
To accomplish this same operation with vmap requires more thought, because there are two separate mapped axes, and some of the axes are broadcast. But you can accomplish the same thing this way:
jax.vmap(jax.vmap(add, in_axes=(None, 0)), in_axes=(0, None))(x2, y2[:, 0]).shape
# (20, 10)
This latter nested vmap is essentially what is happening under the hood when you use jax.numpy.vectorize.
As for which to use in any given situation:
if you want to map a function across a single, explicitly specified axis of the inputs, use jax.vmap
if you want a function's inputs to be mapped across zero or more axes according to numpy's broadcasting rules as applied to the input, use jax.numpy.vectorize.
in situations where the transforms are identical (for example when mapping a function of 1D inputs) lean toward using vmap, because it more directly does what you want to do.

Is there an equivalent to R apply function in Python?

I am trying to find the Python equivalent to R's apply function but with multidimensional arrays.
For example, when called the following code:
z <- array(1, dim = 2:4)
apply(z, 1, sum)
The result is:
[1] 12 12
and when called with two values for margin:
apply(z, c(1,2), sum)
The result is:
[,1] [,2] [,3]
[1,] 4 4 4
[2,] 4 4 4
I found that the sum function in numpy can be used, but not in the same consistent way:
For example:
import numpy as np
xx= np.ones((2,3,4))
np.sum(xx,axis=(1,2))
The result is:
array([12., 12.])
but I can't find a function that equivalent to apply in its manner specifically when dealing with margin=c(1,2). Could anyone help?
The equivalent in NumPy is:
xx.sum(axis=2)
That is, you are summing over axis 2 (the last dimension), which as its length is 4, leaves the other two dimensions (2,3) as the shape of the result:
array([[4., 4., 4.],
[4., 4., 4.]])
Perhaps a more literal translation of your R code would be:
np.apply_over_axes(np.sum, xx, 2)
Which gives a similar result but transposed. This is likely to be slower, however, and is not idiomatic unless the actual operation you're performing is something more complicated than sum.
np.apply_over_axes is different from R's apply in several ways.
First, np.apply_over_axes needs collapsing axes to be specified,
whereas R's apply needs remaining axes to be specified.
Secondly, np.apply_over_axes applies function iteratively as the documentation stated below. The result is the same for np.sum but it could be different for other functions.
func is called as res = func(a, axis), where axis is the first element of axes. The result res of the function call must have either the same dimensions as a or one less dimension. If res has one less dimension than a, a dimension is inserted before axis. The call to func is then repeated for each axis in axes, with res as the first argument.
And the func for np.apply_over_axes needs to be in particular format and the return of func needs to be in particular shape for np.apply_over_axes to perform correctly.
Here's an example how np.apply_over_axes fails
>>> arr.shape
(5, 4, 3, 2)
>>> np.apply_over_axes(np.mean, arr, (0,1))
array([[[[ 0.05856732, -0.14844212],
[ 0.34214183, 0.24319846],
[-0.04807454, 0.04752829]]]])
>>> np_mean = lambda x: np.mean(x)
>>> np.apply_over_axes(np_mean, arr, (0,1))
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "<__array_function__ internals>", line 5, in apply_over_axes
File "/Users/kwhkim/opt/miniconda3/envs/rtopython2-pip/lib/python3.8/site-packages/numpy/lib/shape_base.py", line 495, in apply_over_axes
res = func(*args)
TypeError: <lambda>() takes 1 positional argument but 2 were given
Since there seems to be no equivalent function in Python,
I made a function that is similar to R's apply
def np_apply(arr, axes_remain, fun, *args, **kwargs):
axes_remain = tuple(set(axes_remain))
arr_shape = arr.shape
axes_to_move = set(range(len(arr.shape)))
for axis in axes_remain:
axes_to_move.remove(axis)
axes_to_move = tuple(axes_to_move)
arr, axes_to_move
arr2 = np.moveaxis(arr, axes_to_move, [-x for x in list(range(1,len(axes_to_move)+1))]).copy()
#if arr2.flags.c_contiguous:
arr2 = arr2.reshape([arr_shape[x] for x in axes_remain]+[-1])
return np.apply_along_axis(fun, -1, arr2, *args, **kwargs)
It works fine at least for the sample example as above(not exactly the same as the result above but math.close() returns True for nearly all elements)
>>> np_apply(arr, (2,3), np.mean)
array([[ 0.05856732, -0.14844212],
[ 0.34214183, 0.24319846],
[-0.04807454, 0.04752829]])
>>> np_apply(arr, (2,3), np_mean)
array([[ 0.05856732, -0.14844212],
[ 0.34214183, 0.24319846],
[-0.04807454, 0.04752829]])
For the function to work smoothly for large multidimensional array,
it needs to be optimized. For instance,
array should be prevented from copying.
Anyway it seems to work as a proof-of-concept and I hope it helps.
PS)
arr is generated by arr = np.random.normal(0,1,(5,4,3,2))

How to call a function with parameters as matrix?

I am trying to call scipy.stats.multivariate_normal with four different parameters for mu and sigma. And then for each generated probability density function I need to call that pdf on an array of say, 10 values.
For simplicity let's say that above mentioned function is addXY:
def addXY(x, y):
return x+y
params=[[1,2],[1,3],[1,4],[1,5]] # mu and sigma, four versions
inputs=[1,2,3] # values, in this case 3 of them
matrix = []
for pdf_params in params:
row = []
for inp in inputs:
entry = addXY(*pdf_params)
row.append(entry*inp)
matrix.append(row)
print matrix
Is this pythonic?
Is there a way to pass params and inputs and get a matrix with all combinations in it that is more pythonic/vectorized/faster?
!Important notice: Inputs in the example are scalar values (I've set scalar values to simplify problem description, I am actually using array of n-dimensional vectors and thus multivariate_normal pdf).
Hints and tips about similar operations are welcome.
Based on your description of what you are trying to compute, you don't need multivariate_normal. You are calling the PDF method with a set of scalar values for a distribution with a scalar mu and sigma. So you can use the pdf() method of scipy.stats.norm. This method will broadcast its arguments, so by passing in arrays with the proper shape, you can compute the PDF for the different values of mu and sigma in one call. Here's an example.
Here are your x values (you called them inputs), and the parameters:
In [23]: x = np.array([1, 2, 3])
In [24]: params = np.array([[1, 2], [1, 3], [1, 4], [1, 5]])
For convenience, separate the parameters into arrays of mu and sigma values.
In [25]: mu = params[:, 0]
In [26]: sig = params[:, 1]
We'll use scipy.stats.norm to compute the PDF.
In [27]: from scipy.stats import norm
This call computes the PDF for the desired combinations of x and parameters. mu.reshape(-1, 1) and sig.reshape(-1, 1) are 2D arrays with shape (4, 1). x has shape (3,), so when these arguments are broadcast, the result has shape (4, 3). Each row is the PDF evaluated at x for one of the pairs of mu and sigma.
In [28]: p = norm.pdf(x, loc=mu.reshape(-1, 1), scale=sig.reshape(-1, 1))
In [29]: p
Out[29]:
array([[ 0.19947114, 0.17603266, 0.12098536],
[ 0.13298076, 0.12579441, 0.10648267],
[ 0.09973557, 0.09666703, 0.08801633],
[ 0.07978846, 0.07820854, 0.07365403]])
In other words, the rows of p are:
norm.pdf(x, loc=mu[0], scale=sig[0])
norm.pdf(x, loc=mu[1], scale=sig[1])
norm.pdf(x, loc=mu[2], scale=sig[2])
norm.pdf(x, loc=mu[3], scale=sig[3])
This is only my idea to shorten the code and utilize more library.
In your code, in fact, you do not use numpy, scipy. Question will be whether you would like to use numpy.array for further data processing.
Option 1: just use list to present array and list of list to present matrix:
from itertools import product
matrix_list = [sum(param)*input_x for param, input_x in product(params, inputs)]
matrix = zip(*[iter(matrix_list)]*len(inputs))
print matrix
Credit for using zip method should be given to
convert a flat list to list of list in python
Option 2: use numpy.array and numpy.matrix for further processing
from itertools import product
import numpy as np
matrix_array = np.array([sum(param)*input_x for param, input_x in product(params, inputs)])
matrix = matrix_array.reshape(len(params),len(inputs))
print matrix

Numpy function to get shape of added arrays

tl;dr: How do I predict the shape returned by numpy broadcasting across several arrays without having to actually add the arrays?
I have a lot of scripts that make use of numpy (Python) broadcasting rules so that essentially 1D inputs result in a multiple-dimension output. For a basic example, the ideal gas law (pressure = rho * R_d * temperature) might look like
def rhoIdeal(pressure,temperature):
rho = np.zeros_like(pressure + temperature)
rho += pressure / (287.05 * temperature)
return rho
It's not necessary here, but in more complicated functions it's very useful to initialize the array with the right shape. If pressure and temperature have the same shape, then rho also has that shape. If pressure has shape (n,) and temperature has shape (m,), I can call
rhoIdeal(pressure[:,np.newaxis], temperature[np.newaxis,:])
to get rho with shape (n,m). This lets me make plots with multiple values of temperature without having to loop over rhoIdeal, while still allowing the script to accept arrays of the same shape and compute the result element-by-element.
My question is: Is there a built-in function to return the shape compatible with several inputs? Something that behaves like
def returnShape(list_of_arrays):
return np.zeros_like(sum(list_of_arrays)).shape
without actually having to sum the arrays? If there's no built-in function, what would a good implementation look like?
You could use np.broadcast. This function returns an object encapsulating the result of broadcasting two or more arrays together. No actual operation (e.g. addition) is performed - the object simply has some of the same attributes that an array produced by means of other operations would have (shape, ndim, etc.).
For example:
x = np.array([1,2,3]) # shape (3,)
y = x.reshape(3,1) # shape (3, 1)
z = np.ones((5,1,1)) # shape (5, 1, 1)
Then you can check what the shape of the array returned by broadcasting x, y and z would be by inspecting the shape attribute:
>>> np.broadcast(x, y, z).shape
(5, 3, 3)
This means that you could implement your function simply as follows:
def returnShape(*args):
return np.broadcast(*args).shape

How can I turn this function into a unfunc without for loops?

I am trying to plot a function I created against a range of values (y-axis vs. x-axis).
The operation I would like to compute is common in "matrix multiplication" :
r^T * C * r
where r^T should be of shape (1,100), r of shape (100,1), and C is a matrix of shape (100,100) (or an ndarray shape 100,100) . Multiplied together using numpy.dot(), the output should be a single value.
The function only has one input, which can be an array of data.
import numpy as np
# The user first sets the values used by the function
# Not "true code", because input() too complex for the question at hand
r = data # an numpy ndarray of 100 values, (100,)
original_matrix = M # set matrix, such that M.shape = (100, 100)
param = array of data # EITHER an array of values, shape (50,),
# OR one value, i.e. a 32/64-bit float
# e.g. parameters = np.array of 50 values
def function(param):
# using broadcasting, "np.sum(param * original_matrix for i in r)"
new_matrix = np.sum(param[:, None, None] * original_matrix, axis=0)
# now perform r^T * C * r
return np.dot( r.transpose(), np.dot( new_matrix, r) )
Calling the function
function(param)
results in one value, in format = numpy.float64.
I would like to plot this function against a series of values, i.e. I need this function to input a np.array and output a np.cdarray, must like other ufuncs in NumPy. The function will evaluate each element in the ndarray, and plot this as a function.
For example,
import pylab
X = np.arange(100)
Y = sin(X)
pylab.plot(X, Y)
outputs
Given that my original function (which is solely a function of the array "parameters") results in np.float64 format, how can I turn this function into a ufunc? I would like to plot my function on the y-axis against parameters on the x-axis.
What if you change your function to take a single paramater rather than an array?
Then you could just do
X = range(50)
Y = [function(x) for x in X]
pylab.plot(X, Y)
I can offer two solutions
You can make (almost) any function a ufunc using np.vectorize which handles numbers as well as np.arrays like the np.sin function
def my_func_1(param):
# using broadcasting, "np.sum(param * original_matrix for i in r)"
new_matrix = np.sum(param * original_matrix[None,:,:], axis=0)
# now perform r^T * C * r
return np.dot( r.transpose(), np.dot( new_matrix, r) )
my_vec_func_1 = np.vectorize(my_func_1)
Note that np.vectorize does not really vectorize your code ... I just makes automatically a forloop if an array is passed as an argument. There is not gain in runtime by using it ... see the timings below.
You can define a truly vectorized function which takes (for the following code) only one-dimensional lists or np.arrays as an argument:
def my_vec_func_2(param):
param = np.asarray(param)
new_matrix = np.sum(param[:,None,None,None] * original_matrix[None,None,:,:],axis=1)
return np.dot(r, np.dot(new_matrix,r).transpose())
Truly vectorized codes are usually considerably faster than forloops. Why the gain is so small in this case I can not explain for this case ...
Timings
I used the following code to test the runtime
import numpy as np
from numpy.random import randint
r = randint(10,size=(100)) # an numpy ndarray of 100 values, (100,)
original_matrix = randint(30,size=(100,100))
timeit my_vec_func_1(np.arange(10000))
1 loops, best of 3: 508 ms per loop
timeit my_vec_func_2(np.arange(10000))
1 loops, best of 3: 488 ms per loop
timeit [my_func_1(x) for x in np.arange(10000)]
1 loops, best of 3: 505 ms per loop

Categories