Conditional operation with numpy arrays

Question

I am writing a program I use Euler's method to understand the behavior of a pendulum. I can give it a starting condition and run it. But now, I have to do it using a matrix with many starting conditions and during the iteration I have to normalize the value of 'teta' because it must not surpass |2pi| which happens for some conditions that accept the pendulum going all the way through the circle.

def pendulo(teta_zero, omega_zero, c_1, c_2, t_total, h_passo):
# inicializações
teta_omega  = np.empty((int(t_total/h_passo)+1, 2))
teta_omega[0][0] = teta_zero
teta_omega[0][1] = omega_zero


#recorrência

for i in range(0, int(t_total/h_passo), 1):
    teta_omega[i+1][0] = ( teta_omega[i][0] ) + ( h_passo*teta_omega[i][1] )
    teta_omega[i+1][1]= teta_omega[i][1] + (h_passo/28)*(forca_teta(teta_omega[i][0], c_1, c_2) - 10*teta_omega[i][1])

    #normalization
    if teta_omega[i+1][0] <= -np.pi:
        teta_omega[i+1][0] = teta_omega[i+1][0] + 2*np.pi

    elif teta_omega[i+1][0] > np.pi:
        teta_omega[i+1][0] = -2*np.pi + teta_omega[i+1][0]

My problem is when I try to do it over the matrices I have created with the two starting values using arrays (each dimension is a line of the matrix and each element is a value, one array with 'teta' values and the other with 'omega' values) because I can't use a similar code because,as it shows, the value of a matrix seems to be ambigous. I'm just a begginer, I've only had a introductory course in python and I'm learning numpy by myself, can you guys help me do it over a matrix?

def bacias(TH_IT, OM_IT, c_1, c_2, t_total, h_passo):

"""    # inicializações
teta_omega  = np.empty((int(t_total/h_passo)+1, 2))
teta_omega[0][0] = teta_zero
teta_omega[0][1] = omega_zero"""

#recorrência

for i in range(0, int(t_total/h_passo), 1):
    TH_IT = TH_IT + ( h_passo*OM_IT )
    OM_IT = OM_IT + (h_passo/28)*(forca_teta(TH_IT, c_1, c_2) - 10*OM_IT)

    #normalization
    if TH_IT <= -np.pi:
        TH_IT = teta_omega[i+1][0] + 2*np.pi

    elif TH_IT > np.pi:
        TH_IT = -2*np.pi + teta_omega[i+1][0]




return TH_IT, OM_IT

This returns me `ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all() I hope I've been clear enough, thanks in advance !

Answer 1

I solved it with 'where'

for i in range(0, int(t_total/h_passo), 1):
        TH_IT = TH_IT + ( h_passo*OM_IT )
        OM_IT = OM_IT + (h_passo/28)*(forca_teta(TH_IT, c_1, c_2) - 10*OM_IT)

        TH_IT[np.where(TH_IT <= -np.pi)] += 2*np.pi
        TH_IT[np.where(TH_IT > np.pi)] -= 2*np.pi