How should I modify this function in order to use it with arrays as arguments and to give arrays as return?

Viewed 34

I have a Python function that uses normally scalar values as arguments, transforms them as arrays during the computation and returns scalar values at the end.

def referentialConversion(thetaI, thetaO, deltaPhi):
    omegaI = np.array([-np.sin(thetaI), 0, np.cos(thetaI)])
    omegaO = np.array([np.sin(thetaO) * np.cos(deltaPhi), np.sin(thetaO) * np.sin(deltaPhi), np.cos(thetaO)])

    h = (omegaI + omegaO) / np.linalg.norm(omegaI + omegaO)
    b = np.array([1, 0, 0])
    t = np.array([0, 1, 0])
    n = np.array([0, 0, 1])

    b_prime = np.cross(n, h) / np.linalg.norm(np.cross(n, h))
    t_prime = np.cross(b_prime, h)

    x_h, y_h, z_h = np.dot(h, b), np.dot(h, t), np.dot(h, n)
    x_d, y_d, z_d = np.dot(omegaI, b_prime), np.dot(omegaI, b_prime), np.dot(omegaI, h)

    phi_h   = np.arctan2(y_h, x_h)
    phi_d   = np.arctan2(y_d, x_d)
    theta_h = np.arccos(z_h)
    theta_d = np.arccos(z_d)

    return phi_d, theta_h, theta_d

If I am using the function with scalar arguments, it is perfectly working:

thetaI   = np.linspace(0, np.pi/2, 500)
thetaO   = thetaI
deltaPhi = np.linspace(0, 2*np.pi, 2000)

phiD   = np.array([])
thetaH = np.array([])
thetaD = np.array([])
for theta_i in thetaI:
    for theta_o in thetaO:
        for delta_phi in deltaPhi:
            phi_d, theta_h, theta_d = referentialConversion(0.5, 0.5, 0.5) # This works
            phiD   = np.append(phiD, phi_d)
            thetaH = np.append(thetaH, theta_h)
            thetaD = np.append(thetaD, theta_d)

But since I need to do the computation for a lot of values, I would prefer to do it using NumPy arrays rather than a for loop.

thetaI   = np.linspace(0, np.pi/2, 500)
thetaO   = thetaI
deltaPhi = np.linspace(0, 2*np.pi, 2000)

phi_d, theta_h, theta_d = referentialConversion(thetaI, thetaO, deltaPhi) # Doesn't work
# ValueError: operands could not be broadcast together with shapes...

I am not very used to NumPy, is there a way to enable my function to do the last example ?

0 Answers
Related