In [115]:

    

from sympy import sin, cos, tan, atan2, sqrt, Matrix, var, init_printing, simplify, pprint

init_printing(use_unicode=True)

var('tanLambda,d0,phi0,z0', real = True)
var('ptinv', real= True, positive = True)
var('x,y,z,p_x,p_y,p_z', real = True)
var('alpha, P', real = True)
var('T,V,D', real = True, positive = True)
_T = p_x**2 + p_y**2
_P = p_x * y - p_y * x
_V = x**2 + y**2 
_D = sqrt(_T + alpha *(2*_P + _V * alpha))

_ptinv = 1 / sqrt(_T)  
_tanLambda = p_z / sqrt(_T) 
_phi0 = atan2(p_y - x * alpha, p_x + y * alpha) 
_z0 = z - p_z * atan2(p_x * x * alpha + p_y * y * alpha, _T + alpha * _P) / alpha 
_d0 = (2 * _P + alpha * _V) / (_T + _D )

X = Matrix([
    _tanLambda, _ptinv, _d0, _phi0, _z0
])

Y = Matrix([x, y, z, p_x, p_y, p_z])
J = X.jacobian(Y)
J.simplify()

A = J[:, :3]
B = J[:, 3:]


    

In [116]:

    

import matplotlib as mpl 
from mpl_toolkits.mplot3d import Axes3D
import numpy as np
import matplotlib.pyplot as plt 
%matplotlib inline

# Position of the particle
pos_x = 3.0
pos_y = 2.0
pos_z = 1.0

# Momentum of the particle at this position
mom_x = 1.0
mom_y = 2.0
mom_z = 3.0

# Aplha containing units, magnetic field and charge
draw_alpha = 1.0

mom_tinv = 1.0 / np.sqrt(mom_x**2 + mom_y**2)

fig = plt.figure()
ax = fig.gca(projection='3d')

t = np.linspace(0, 10, 100)
draw_mom_z = mom_z * np.ones(len(t))
draw_mom_y = np.cos(draw_alpha * t) / mom_tinv
draw_mom_x = -np.sin(draw_alpha * t) / mom_tinv

draw_z = pos_z + draw_mom_z * t
draw_x = pos_x + np.cos(draw_alpha * t) / (draw_alpha * mom_tinv)
draw_y = pos_y + np.sin(draw_alpha * t) / (draw_alpha * mom_tinv)
ax.plot(draw_x, draw_y, draw_z, label='Helix')
ax.legend()

plt.show()

for i in range(0, 10):
    print(X.subs([(p_x, draw_mom_x[i]),(p_y, draw_mom_y[i]), (p_z, draw_mom_z[i]), (x, draw_x[i]), (y, draw_y[i]), (z, draw_z[i]), (alpha, 1.0)]))


    

    













    




Matrix([[1.34164078649987], [0.447213595499958], [0.929632483024007], [-0.982793723247329], [-3.0*pi + 2.7640078106427]])
Matrix([[1.34164078649987], [0.447213595499958], [0.929632483024007], [-0.982793723247329], [-3.0*pi + 2.7640078106427]])
Matrix([[1.34164078649987], [0.447213595499958], [0.929632483024007], [-0.982793723247329], [-3.0*pi + 2.7640078106427]])
Matrix([[1.34164078649987], [0.447213595499958], [0.929632483024007], [-0.982793723247329], [-3.0*pi + 2.7640078106427]])
Matrix([[1.34164078649987], [0.447213595499958], [0.929632483024007], [-0.982793723247329], [-3.0*pi + 2.7640078106427]])
Matrix([[1.34164078649987], [0.447213595499958], [0.929632483024008], [-0.982793723247329], [-3.0*pi + 2.7640078106427]])
Matrix([[1.34164078649987], [0.447213595499958], [0.929632483024008], [-0.982793723247329], [2.7640078106427 + 3.0*pi]])
Matrix([[1.34164078649987], [0.447213595499958], [0.929632483024007], [-0.982793723247329], [2.7640078106427 + 3.0*pi]])
Matrix([[1.34164078649987], [0.447213595499958], [0.929632483024007], [-0.982793723247329], [2.7640078106427 + 3.0*pi]])
Matrix([[1.34164078649987], [0.447213595499958], [0.929632483024007], [-0.982793723247329], [2.7640078106427 + 3.0*pi]])

In [122]:

    

# Without magnetic field or zero charge
fig = plt.figure()
ax = fig.gca(projection='3d')

t = np.linspace(0, 10, 100)
draw_mom_z = mom_z * np.ones(len(t))
draw_mom_y = mom_y * np.ones(len(t))
draw_mom_x = mom_x * np.ones(len(t))

draw_z = pos_z + draw_mom_z * t
draw_x = pos_x + draw_mom_x * t
draw_y = pos_y + draw_mom_y * t
ax.plot(draw_x, draw_y, draw_z, label='Helix')
ax.legend()

plt.show()

X_zero = X.subs([(alpha, 0.0)])
X_zero[4] = z - p_z * (p_x * x + p_y * y) / _T
X_zero

for i in range(0, 10):
    print(X_zero.subs([(p_x, draw_mom_x[i]),(p_y, draw_mom_y[i]), (p_z, draw_mom_z[i]), (x, draw_x[i]), (y, draw_y[i]), (z, draw_z[i])]))


    

    













    




Matrix([[1.34164078649987], [0.447213595499958], [-1.10557280900008], [1.10714871779409], [-3.20000000000000]])
Matrix([[1.34164078649987], [0.447213595499958], [-1.10557280900008], [1.10714871779409], [-3.20000000000000]])
Matrix([[1.34164078649987], [0.447213595499958], [-1.10557280900008], [1.10714871779409], [-3.20000000000000]])
Matrix([[1.34164078649987], [0.447213595499958], [-1.10557280900008], [1.10714871779409], [-3.20000000000000]])
Matrix([[1.34164078649987], [0.447213595499958], [-1.10557280900008], [1.10714871779409], [-3.20000000000000]])
Matrix([[1.34164078649987], [0.447213595499958], [-1.10557280900008], [1.10714871779409], [-3.20000000000000]])
Matrix([[1.34164078649987], [0.447213595499958], [-1.10557280900008], [1.10714871779409], [-3.20000000000000]])
Matrix([[1.34164078649987], [0.447213595499958], [-1.10557280900008], [1.10714871779409], [-3.20000000000000]])
Matrix([[1.34164078649987], [0.447213595499958], [-1.10557280900008], [1.10714871779409], [-3.20000000000000]])
Matrix([[1.34164078649987], [0.447213595499958], [-1.10557280900008], [1.10714871779409], [-3.20000000000000]])

In [114]:

    

_P2 = 2 * p_x * y - 2 * p_y * x
_P2A = 2 * alpha * p_x * y - 2 * alpha * p_y * x
_VAA = alpha **2 * x**2 + alpha **2 * y**2 
_DE = sqrt(T + alpha *(2*P + V * alpha))

subs = [(_D, D), (_P, P),(_V, V), (_T, T), (_P2, 2 * P),  (_P2A, 2 * alpha * P), (_VAA, alpha ** 2 * V), (_DE, D)]

A = A.subs(subs)
A


    

    
Out[114]:





$$\left[\begin{matrix}0 & 0 & 0\\0 & 0 & 0\\\frac{1}{D \left(D + T\right)^{2}} \left(2 D \left(D + T\right) - \alpha \left(2 P + V \alpha\right)\right) \left(\alpha x - p_{y}\right) & \frac{1}{D \left(D + T\right)^{2}} \left(2 D \left(D + T\right) - \alpha \left(2 P + V \alpha\right)\right) \left(\alpha y + p_{x}\right) & 0\\- \frac{\alpha \left(\alpha y + p_{x}\right)}{\left(\alpha x - p_{y}\right)^{2} + \left(\alpha y + p_{x}\right)^{2}} & \frac{\alpha \left(\alpha x - p_{y}\right)}{\left(\alpha x - p_{y}\right)^{2} + \left(\alpha y + p_{x}\right)^{2}} & 0\\- \frac{p_{z} \left(\alpha y + p_{x}\right)}{2 P \alpha + T + V \alpha^{2}} & \frac{p_{z} \left(\alpha x - p_{y}\right)}{2 P \alpha + T + V \alpha^{2}} & 1\end{matrix}\right]$$

In [95]:

    

B = B.subs(subs)
B


    

    
Out[95]:





$$\left[\begin{matrix}- \frac{p_{x}}{T^{\frac{3}{2}}} & - \frac{p_{y}}{T^{\frac{3}{2}}} & 0\\- \frac{p_{x} p_{z}}{T^{\frac{3}{2}}} & - \frac{p_{y} p_{z}}{T^{\frac{3}{2}}} & \frac{1}{\sqrt{T}}\\\frac{\alpha x - p_{y}}{\left(\alpha x - p_{y}\right)^{2} + \left(\alpha y + p_{x}\right)^{2}} & \frac{\alpha y + p_{x}}{\left(\alpha x - p_{y}\right)^{2} + \left(\alpha y + p_{x}\right)^{2}} & 0\\- \frac{p_{z} \left(x \left(P \alpha + T\right) - \left(\alpha y + 2 p_{x}\right) \left(p_{x} x + p_{y} y\right)\right)}{\alpha^{2} \left(p_{x} x + p_{y} y\right)^{2} + \left(P \alpha + T\right)^{2}} & - \frac{p_{z} \left(y \left(P \alpha + T\right) + \left(\alpha x - 2 p_{y}\right) \left(p_{x} x + p_{y} y\right)\right)}{\alpha^{2} \left(p_{x} x + p_{y} y\right)^{2} + \left(P \alpha + T\right)^{2}} & - \frac{1}{\alpha} \operatorname{atan_{2}}{\left (\alpha \left(p_{x} x + p_{y} y\right),P \alpha + T \right )}\\\frac{1}{D \left(D + T\right)^{2}} \left(2 D y \left(D + T\right) - \left(2 P + V \alpha\right) \left(2 D p_{x} + \alpha y + p_{x}\right)\right) & - \frac{1}{D \left(D + T\right)^{2}} \left(2 D x \left(D + T\right) + \left(2 P + V \alpha\right) \left(2 D p_{y} - \alpha x + p_{y}\right)\right) & 0\end{matrix}\right]$$

In [96]:

    

X = X.subs(subs)
X


    

    
Out[96]:





$$\left[\begin{matrix}\frac{1}{\sqrt{T}}\\\frac{p_{z}}{\sqrt{T}}\\\operatorname{atan_{2}}{\left (- \alpha x + p_{y},\alpha y + p_{x} \right )}\\z - \frac{p_{z}}{\alpha} \operatorname{atan_{2}}{\left (\alpha p_{x} x + \alpha p_{y} y,P \alpha + T \right )}\\\frac{2 P + V \alpha}{D + T}\end{matrix}\right]$$

In [97]:

    

c0 = X - A * Matrix([x,y,z]) - B * Matrix([p_x, p_y, p_z])
c0


    

    
Out[97]:





$$\left[\begin{matrix}\frac{1}{\sqrt{T}} + \frac{p_{x}^{2}}{T^{\frac{3}{2}}} + \frac{p_{y}^{2}}{T^{\frac{3}{2}}}\\\frac{p_{x}^{2} p_{z}}{T^{\frac{3}{2}}} + \frac{p_{y}^{2} p_{z}}{T^{\frac{3}{2}}}\\\frac{\alpha x \left(\alpha y + p_{x}\right)}{\left(\alpha x - p_{y}\right)^{2} + \left(\alpha y + p_{x}\right)^{2}} - \frac{\alpha y \left(\alpha x - p_{y}\right)}{\left(\alpha x - p_{y}\right)^{2} + \left(\alpha y + p_{x}\right)^{2}} - \frac{p_{x} \left(\alpha x - p_{y}\right)}{\left(\alpha x - p_{y}\right)^{2} + \left(\alpha y + p_{x}\right)^{2}} - \frac{p_{y} \left(\alpha y + p_{x}\right)}{\left(\alpha x - p_{y}\right)^{2} + \left(\alpha y + p_{x}\right)^{2}} + \operatorname{atan_{2}}{\left (- \alpha x + p_{y},\alpha y + p_{x} \right )}\\\frac{p_{x} p_{z} \left(x \left(P \alpha + T\right) - \left(\alpha y + 2 p_{x}\right) \left(p_{x} x + p_{y} y\right)\right)}{\alpha^{2} \left(p_{x} x + p_{y} y\right)^{2} + \left(P \alpha + T\right)^{2}} + \frac{p_{y} p_{z} \left(y \left(P \alpha + T\right) + \left(\alpha x - 2 p_{y}\right) \left(p_{x} x + p_{y} y\right)\right)}{\alpha^{2} \left(p_{x} x + p_{y} y\right)^{2} + \left(P \alpha + T\right)^{2}} + \frac{p_{z} x \left(\alpha y + p_{x}\right)}{2 P \alpha + T + V \alpha^{2}} - \frac{p_{z} y \left(\alpha x - p_{y}\right)}{2 P \alpha + T + V \alpha^{2}} + \frac{p_{z}}{\alpha} \operatorname{atan_{2}}{\left (\alpha \left(p_{x} x + p_{y} y\right),P \alpha + T \right )} - \frac{p_{z}}{\alpha} \operatorname{atan_{2}}{\left (\alpha p_{x} x + \alpha p_{y} y,P \alpha + T \right )}\\\frac{2 P + V \alpha}{D + T} - \frac{p_{x}}{D \left(D + T\right)^{2}} \left(2 D y \left(D + T\right) - \left(2 P + V \alpha\right) \left(2 D p_{x} + \alpha y + p_{x}\right)\right) + \frac{p_{y}}{D \left(D + T\right)^{2}} \left(2 D x \left(D + T\right) + \left(2 P + V \alpha\right) \left(2 D p_{y} - \alpha x + p_{y}\right)\right) - \frac{x}{D \left(D + T\right)^{2}} \left(2 D \left(D + T\right) - \alpha \left(2 P + V \alpha\right)\right) \left(\alpha x - p_{y}\right) - \frac{y}{D \left(D + T\right)^{2}} \left(2 D \left(D + T\right) - \alpha \left(2 P + V \alpha\right)\right) \left(\alpha y + p_{x}\right)\end{matrix}\right]$$

In [ ]: