Numpy Apply Custom Function Along Axis (or Map)

Question

Related to this question:

I have the following code with custom functions:

import numpy as np
from scipy.integrate import quad
from scipy.special import binom
from scipy.stats import norm
import math as mt
ranks = np.array([[8, 5, 3, 7, 6, 2, 1, 4],
       [8, 5, 3, 7, 6, 2, 1, 4]])

def E(r, n):
    print(r,n)
    def f(x):
        F = Phi(x)
        return x*(1-F)**(r-1)*F**(n-r)*phi(x)
    return -r*binom(n, r)*quad(f, -inf, inf)[0]

I want to apply the function E over each row of the array ranks .

I tried np.apply_along_axis but the problem I'm having is that I have 2 arguments for the function: the individual rank of an element in a row of an array ( r ) and the length of that row ( n ).

Also, I'm wondering if it would be faster and better anyway to use map as with the accepted answer to the question referenced above (because speed is important with this).

Here's what I've tried:

result = np.apply_along_axis(E, 1, ranks)

Issue: I need to provide the n argument. Also, this may be slower than mapping.

#not sure if the "*" is needed or what it does.
def E2(i):
    return [*E(ranks[i], len(ranks[i]))]

it = 2 #2 rows (iterations)
indx = np.arange(it)
result = np.array(list(map(E2, indx))) 

Issue: Needs 1 rank at a time, per row. Not sure how to achieve this.

Thanks in advance!

Answer 1

Since you haven't specified the imports for Phi , quad etc, I won't try to provide a working example.

def E(r, n):
    print(r,n)
    def f(x):
        F = Phi(x)
        return x*(1-F)**(r-1)*F**(n-r)*phi(x)
    return -r*binom(n, r)*quad(f, -inf, inf)[0]

A simple list comprehension:

res = [E(r,n) for r in ranks]

I'd prefer to move the f(x) definition out of E . That may be more of a matter of style than performance.

def f(x,r):
    F = Phi(x)
    return x*(1-F)**(r-1)*F**(n-r)*phi(x)

def E(r,n):
    temp = quad(f, -inf, inf, args=(r,))[0]
    return -r*binom(n, r)*temp