Search code examples
pythonnumpyvectorizationnumpy-ndarray

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


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])

Solution

  • 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.])