mert
mert

Reputation: 2002

Numpy arange over numpy arrays

I have a function that basically returns generalized harmonic number.

def harmonic(limit, z):
   return numpy.sum(1.0/numpy.arange(1, limit+1)**z)

Here is two examples for the current function definition:

>>> harmonic(1, 1)
1.0
>>> harmonic(2, 1)
1.5

As you might guess this works fine when limit is scalar, but how can I make this function work with 1D and 2D arrays as well?

The following demonstrates an example output of the function I want to achieve

>>> limit = np.array([[1, 2], [3, 4]])
>>> harmonic(limit, 1)
array([[1.0, 1.5], [1.833, 2.083]])

Upvotes: 1

Views: 1605

Answers (2)

DSM
DSM

Reputation: 352979

If you're only interested in vectorizing over limit and not z, as in the example you showed, then I think you can use np.vectorize:

>>> h = np.vectorize(harmonic)
>>> h(1, 1)
array(1.0)
>>> h(2, 1)
array(1.5)
>>> h([[1,2], [3,4]], 1)
array([[ 1.        ,  1.5       ],
       [ 1.83333333,  2.08333333]])
>>> h([[1,2], [3,4]], 2)
array([[ 1.        ,  1.25      ],
       [ 1.36111111,  1.42361111]])

Note that this will return 0-dimensional arrays for the scalar case.

Actually, on second thought, it should work for the z case too:

>>> h([[2,2], [2,2]], [[1,2],[3,4]])
array([[ 1.5   ,  1.25  ],
       [ 1.125 ,  1.0625]])

Upvotes: 5

kiriloff
kiriloff

Reputation: 26335

arange generates evenly spaced 1D ndarray in range [1,limit+1] in your example.

Now say you want an multi-dim ndarray of evenly spaced arrays. Then you may use arange to generate each component of your 2D ndarray. You convert result of arange to a python list with list(), to make it the right format to be an argument of ndarray constructor.

It all depends on your purpose. As you deal with math. analysis, what you look for may be a grid:

>>> np.mgrid[0:5,0:5]
array([[[0, 0, 0, 0, 0],
        [1, 1, 1, 1, 1],
        [2, 2, 2, 2, 2],
        [3, 3, 3, 3, 3],
        [4, 4, 4, 4, 4]],
       [[0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4],
        [0, 1, 2, 3, 4]]])

More here.

EDIT: After you posted the code : as DSM mentions, np.vectorize is a good way to do. From doc,

class numpy.vectorize(pyfunc, otypes='', doc=None, excluded=None, 
   cache=False)

Generalized function class.

Define a vectorized function which takes a nested sequence of objects or numpy arrays as inputs and returns a numpy array as output. The vectorized function evaluates pyfunc over successive tuples of the input arrays like the python map function, except it uses the broadcasting rules of numpy.

Upvotes: 0

Related Questions