Numpy how to use np.cumprod to rewrite python for i in range function-CodePudding

I have two python functions. The first one:

mt = np.array([1, 2, 3, 4, 5, 6, 7])
age, interest = 3, 0.5

def getnpx(mt, age, interest):
    val = 1
    initval = 1
    for i in range(age, 6):
        val = val * mt[i]
        intval = val / (1   interest) ** (i   1 - age)
        initval = initval   intval
    return initval

The output is:

48.111111111111114

In order to make it faster, I used numpy to vectorize it:

def getnpx_(mt, age, interest):
    print(np.cumprod(mt[age:6]) / (1   interest)**np.arange(1, 7 - age))
    return 1   (np.cumprod(mt[age:6]) / (1   interest)**np.arange(1, 7 - age)).sum()

getnpx_(mt, age, interest)

It works and the output is still:

48.111111111111114

However I have no idea how to rewrite my second function by numpy:

pt1 = np.array([1, 2, 3, 4, 5, 6, 7])
pt2 = np.array([2, 4, 3, 4, 7, 4, 8])
pvaltable = np.array([0, 0, 0, 0, 0, 0, 0])

def jointpval(pt1, pt2, age1, age2):
    j = age1
    for i in range(age2, 6):
        k = min(j, 135)
        pvaltable[i] = pt1[k] * pt2[i]
        j = j   1
    return pvaltable

jointpval(pt1, pt2, 3, 4)

Output:

array([ 0,  0,  0,  0, 28, 20,  0])

I expect to be able to convert the loop

for i in range(age2, 6):

To something like:

np.cumprod(pt1[age:6])

The final output should be the same as:

array([ 0,  0,  0,  0, 28, 20,  0])

CodePudding user response：

I found this solution:

import numpy as np
pt1 = np.array([1, 2, 3, 4, 5, 6, 7])
pt2 = np.array([2, 4, 3, 4, 7, 4, 8])

def jointpval(pt1, pt2, age1, age2):
    pvaltable = np.zeros(len(pt1))
    idx2 = np.arange(age2, 6)
    idx1 = np.arange(len(idx2))   age1
    idx1 = np.where(idx1 > 135, 135, idx1) 
    pvaltable[idx2] = pt1[idx1] * pt2[idx2]
    return pvaltable

Where jointpval(pt1, pt2, 3, 4) returns

array([ 0.,  0.,  0.,  0., 28., 20.,  0.])

CodePudding user response：

I would recommend not hard coding your array sizes. For example:

def getnpx_(mt, age, interest):
    return 1   (np.cumprod(mt[age:-1]) / (1   interest)**np.arange(1, len(mt) - age)).sum()

Notice that the index age:-1 agnostic to the size of mt, and mt.size - age saves you the trouble of recoding the function every time you have a different mt.

Your solution for the second case is pretty much on-point. I might recommend using np.clip rather than where to also set the lower bound to zero:

def jointpval(pt1, pt2, age1, age2):
    pvaltable = np.zeros(len(pt1))
    idx = np.arange(len(pt2) - age2 - 1)
    idx1 = np.clip(idx   age1, 0, len(pt1) - 1)
    idx2 = idx   age2
    pvaltable[idx2] = pt1[idx1] * pt2[idx2]
    return pvaltable

You could also use np.minimum to implement vectorized min directly:

idx1 = np.minimum(idx   age1, len(pt1) - 1)

Keep in mind that ranges are exclusive on the upper bound. Both of your functions are truncating the computation of the last element for this reason. It is likely that idx = np.arange(len(pt2) - age2 - 1) should really be idx = np.arange(len(pt2) - age2) in jointpval, while mt[age:-1] should be mt[age:], and np.arange(1, len(mt) - age) should be np.arange(len(mt) - age) in genpx_.