This is my first question here on the forum, so I hope I'm doing evrything right. My problem is, I am supposed to solve an ODE with python. The ODE describes the movement of an electron inside of a "bubble" (a bubble is something that occurs in Plasmaphysics when you shoot with a laser at plasma but thats not that important for the question im asking). So my task is to build a programm that can plot the trajectory of such an electron. For that I have to basically solve 2 different ODE's since the electron inside the bubble is moving according to an ODE and if it exits the bubble it is in a field-free-area where it should move according to the classical equations of motion.

In the following I will show you the code if only 1 ODE is solved:

import numpy as np
from scipy.integrate import odeint
import matplotlib.pyplot as plt
from math import sqrt,pi

def f(r,t):
    px = r[0]
    py = r[1]
    x = r[2]
    y = r[3]
    gamma = sqrt(1+px**2+py**2)
    
    dpx_dt = -(1+V)*(x-V*t)/4 + py*y/(4*gamma) 
    dpy_dt = -(1+px/gamma)*y/4
    dx_dt = px/gamma 
    dy_dt = py/gamma
    
    return [dpx_dt, dpy_dt, dx_dt, dy_dt]

def Plot_Aufgabe5(t,p_x,p_y,xi,y):
    #Plotten der Trajektorie des Elektrons
    plt.plot(xi,y)
    
    return
      
#Konstanten
r_b     = 10
m       = 9.109*10**-31                          
c       = 2.99*10**8                              
roh     = 9   

#Anfangsbedingungen
gamma_0 = 4
V       = np.sqrt(((gamma_0**2) - 1)/(gamma_0**2))

t_end   = 300                                     
N       = 1000                                    
t       = np.linspace(0,t_end,N)                  
y       = [0, 0, np.sqrt(r_b**2 - roh**2), roh]   

#Array mit äquidistanten Werten erstellen
Äqui_array = np.linspace(0,2*np.pi,10)

#Array mit verschd. Anfangsimpulsen
Impuls_array = [-1,-0.5,0.5,1]

for j in Impuls_array:
    plt.figure()
    #Plotten eines Kreises
    theta = np.linspace(0,2*np.pi,t_end)
    Punkt_x = r_b*np.cos(theta)
    Punkt_y = r_b*np.sin(theta)
    plt.plot(Punkt_x,Punkt_y)

    for i in range(10):
        #Neue Anfangsbedingungen aufstellen
        y   = [j,j,r_b*np.cos(Äqui_array[i]),r_b*np.sin(Äqui_array[i])]
    
        #Ausführen des Solvers
        sol = odeint(f,y,t)

        p_x = sol[:,0]
        p_y = sol[:,1]
        x   = sol[:,2]
        xi  = x - V*t
        y   = sol[:,3]

        #Plot
        Plot_Aufgabe5(t,p_x,p_y,xi,y) 
    
    plt.axis([-15,15,-10,10])
    plt.xlabel(r'$xi_{num}$ = $k_p$$\xi$')
    plt.ylabel(r'$y_{num}$ = $k_p$y')
    plt.title('Trajektorie des Elekrons in einer Bubble mit Impuls p = ' + str(j) +  '*mc')
    plt.show()

Plots for the code for 1 ODE

As you can see in the plot, the electron goes outside the bubble but it comes back in which should not happen. (Also the plot is for 10 different electrons which start on the edge of the bubble)

Therefore I made an approach to change between 2 ODE's according to the position of the electron however that did not quite work:

def f(r,t):
    px = r[0]                                                    #boundary conditions for the electron
    py = r[1]
    x = r[2]
    y = r[3]
    gamma = sqrt(1+px**2+py**2)
    
    if ((x**2 + y**2) <= r_b**2):                                #if electron is inside bubble
        dpx_dt = -(1+V)*(x-V*t)/4 + py*y/(4*gamma) 
        dpy_dt = -(1+px/gamma)*y/4
        dx_dt = px/gamma
        dy_dt = py/gamma
    else:                                                        #if it is outside of the bubble 
        dpx_dt = 0
        dpy_dt = 0
        dx_dt = px/gamma
        dy_dt = py/gamma

    result = [dpx_dt, dpy_dt, dx_dt, dy_dt]
    
    return result

My Idea here was that with the if-statement i check, whether the electron is inside the bubble or not and change the ODE according to that. However it does not really work out and im not sure why, underneath you can find the plots that come with that function (also the rest of the code stayed the same as above):

Plots for the Function to change between 2 ODE's

I hope my problem gets clear and that maybe someone has an idea that could help me! Good day to evryone :)