Simplifying ODE with sympy and then solving numerically with Scipy

91 views Asked by At

I want to solve numerically the equations of motion for a relativistic particle in a E/M field:

$$ \frac{d\vec{p}}{dt} \ = \ q (\vec{E}+ \vec{u} \times \vec{B}) $$

enter image description here

where, \ \vec{p} \ = \ m \gamma\vec{v}, and,
enter image description here.

I have started by simplifying with Sympy to solve the first equation for enter image description here. Using

""" Define symbols"""
t,m,q, c = smp.symbols('t m q c', Real =True)
x, y, z, vx, vy, vz, Ex, Ey, Ez, Bx, By, Bz, gamma = smp.symbols('x y z vx vy vz  E_x E_y E_z B_x B_y B_z gamma ',cls=smp.Function, real=True)

""" Define functions"""

x    = x(t);                y   = y(t);               z    = z(t);                 # Position   
vx   = vx(t);              vy   = vy(t);              vz   = vz(t);                # Velocity  
Ex   = Ex(x, y, z, t);     Ey   = Ey(x, y, z, t);     Ez   = Ez(x, y, z, t);       # Components of E field 
Bx   = Bx(x, y, z, t);     By   = By(x, y, z, t);     Bz   = Bz(x, y, z, t);       # Components of B field
x_d  = smp.diff(x,t);      y_d  = smp.diff(y,t);      z_d  = smp.diff(z,t);        # dx_{i}/dt ,    x_{i} = x, y, z
vx   = x_d;                vy   = y_d;                vz   = z_d;                  # v_{i} ,           i  = x, y, z 
vx_d = smp.diff(x,t,2);    vy_d = smp.diff(y,t,2);    vz_d = smp.diff(z,t,2);      # d^{2}x_{i}/dt^2 , i  = x, y, z 

""" Define Matrices"""
r      = smp.Matrix([x, y, z])
E      = smp.Matrix([Ex, Ey, Ez])
B      = smp.Matrix([Bx, By, Bz])
v      = smp.Matrix([vx, vy, vz])

""" Define gamma """
v_norm = v.norm()
gamma  = gamma(v_norm)
gamma  = 1/smp.sqrt(1 - (v.norm()/c)**2)

""" Define momentum """
p    = m*gamma*v

""" 2nd order ODE of motion """
de1 = smp.diff(p,t)-q*(E +v.cross(B)/c)

""" Solve for dv_{i}/dt, i = x, y, z """ 
sols   = smp.solve([de1[0],de1[1], de1[2]], (vx_d, vy_d, vz_d), Simplify =False, rational=False)

"""Create numpy functions that we can use with numerical methods"""
list_p = (c, m, q, x, y, z, vx, vy, vz, Ex, Ey, Ez, Bx, By, Bz)


dvx_dt_f  = smp.lambdify(list_p,sols[vx_d], modules=smp)
dx_dt_f   = smp.lambdify(vx,vx, modules=smp)

dvy_dt_f  = smp.lambdify(list_p,sols[vy_d], modules=smp)
dy_dt_f   = smp.lambdify(vy, vy, modules=smp)

dvz_dt_f  = smp.lambdify(list_p,sols[vz_d], modules=smp)
dz_dt_f   = smp.lambdify(vz, vz, modules=smp)

Then I want to use the lamdified expressions and integrate them using Scipy's ivp in the following form.

def dSdt(S,t):
    vx, x, vy, y, vz, z = S
    return[dvx_dt_f(c, m, q, x, y, z, vx, vy, vz, Ex, Ey, Ez, Bx, By, Bz),
           dx_dt_f(vx),
           dvy_dt_f(c, m, q, x, y, z, vx, vy, vz, Ex, Ey, Ez, Bx, By, Bz),
           dy_dt_f(vy),
           dvz_dt_f(c, m, q, x, y, z, vx, vy, vz, Ex, Ey, Ez, Bx, By, Bz),
           dz_dt_f(vz)
          ]

for initial conditions, v0 = (x, y, z, vx, vy, vz, Ex, Ey, Ez, Bx, By, Bz)

However, instead of getting the Sympy solution in terms of Vx, Vy, Vz, needed for the numerical integration I get it in terms of the time derivatives $\frac{dx}{dt}.

enter image description here

What can I do to change the derivatives with $v_x, v_y, v_z$?

0

There are 0 answers