# Path of a particle in a non-uniform B field

I have completed my project which simulates the path of a particle trapped in a device known as a magnetic mirror. I would like your thoughts and improvements on my implementation.

Here is the code:

from numba import jit

import numpy as np
import matplotlib.pyplot as plt

# Plot attribute settings
plt.rc("xtick", labelsize="large")
plt.rc("ytick", labelsize="large")
plt.rc("axes", labelsize="xx-large")
plt.rc("axes", titlesize="xx-large")
plt.rc("figure", figsize=(8,8))

# constants
mu_0 = np.pi * 4.0 * pow(10, -7) # permeability of free space [kg*m*s^-2*A^-2]
q_p = 1.6022 * pow(10, -19)         # proton charge [coulombs]
m_p = 1.6726 * pow(10, -27)            # proton mass [kg]

mu = 10000.0 * np.array([0.0, 0.0, 1.0]) # set magnetic moment to point in z direction

''' Calculates magnetic bottle field '''
''' This function will come in handy when we illustrate the field along and ultimately illustrate the
particle motion in field '''
def bot_field(x,y,z):

z_disp = 10.0 # displacement of the two magnetic dipoles with respect to zero (one at z = -z_disp, the other at +z_disp)

# point dipole A
pos_A = np.array([0.0, 0.0, z_disp])        # set the position of the first dipole
r_A = np.array([x,y,z]) - pos_A             # find the difference between this point and point of the observer
rmag_A = np.sqrt(sum(r_A**2))
B1_A = 3.0*r_A*np.dot(mu,r_A) / (rmag_A**5)   # calculate the first term to the magnetic field
B2_A = -1.0 * mu / (rmag_A**3)                # calculate the second term

# point dipole B
pos_B = np.array([0.0, 0.0, -z_disp])  # set the position of the first dipole
r_B = np.array([x,y,z]) - pos_B        # find the difference between this position and the observation position
rmag_B = np.sqrt(sum(r_B**2))
B1_B = 3.0*r_B*np.dot(mu,r_B) / (rmag_B**5) # calculate the first term to the magnetic field
B2_B = -1.0 * mu / (rmag_B**3)              # calculate the second term

return ((mu_0/(4.0*np.pi)) * (B1_A + B2_A + B1_B + B2_B)) # return field due to magnetic bottle

'''Setting up graph for dipole magnetic field'''
y = np.arange(-10.0, 10.0, .1) # create a grid of points from y = -10 to 10
z = np.arange(-10.0, 10.0, .1) # create a grid of points from z = -10 to 10
Y, Z = np.meshgrid(y,z)        # create a rectangular grid out of y and z
len_i, len_j = np.shape(Y)       # define dimensions, for use in iteration
Bf = np.zeros((len_i,len_j,3))   # initialize all points with zero magnetic field

''' iterate through the grid and set magnetic field values at each point '''
for i in range(0, len_i):
for j in range(0, len_j):
Bf[i,j] = bot_field(0.0, Y[i,j], Z[i,j])

plt.streamplot(Y,Z, Bf[:,:,1], Bf[:,:,2], color='orange') # plot the magnetic field
plt.xlim(-10.0,10.0)
plt.ylim(-10.0,10.0)
plt.xlabel("$y$ (m)")
plt.ylabel("$z$ (m)")
plt.title("Magnetic Field in a Magnetic Bottle")

q = 2.0*q_p # charge of helium-4
m = 4.0*m_p # mass of helium-4
QoverM = q/m

dt = pow(10, -5) # timestep

t = np.arange(0.0, 1.0, dt) # array for times
rp = np.zeros((len(t), 3)) # array for position values
vp = np.zeros((len(t), 3)) # array for velocity values

v_o = 100 # set the initial velocity to 100 m/s
rp[0,:] = np.array([0.0, -5.0, 0.0]) # initialize the position to y=-5, 5m above the lower dipole
vp[0,:] = np.array([0.0, 0.0, v_o]) # initialize the velocity to be in the z-direction

''' Model the particle motion in the field at each time step (Foward Euler Method) '''
for it in np.arange(0, len(t)-1,1):
Bp = bot_field(rp[it,0], rp[it, 1], rp[it,2]) # input the current particle position into to get the magnetic field
Ap = QoverM * np.cross(vp[it,:], Bp)          # calculate the magnetic force on the particle
vp[it+1] = vp[it] + dt*Ap                     # update the velocity of the particle based on this force
rp[it+1] = rp[it] + dt*vp[it]                 # update the positon of the particle based on this velocity
if (np.sqrt(np.sum(rp[it+1]**2)) > 20.0): # If the particle escapes (i.e. exceeds 20 m from origin), cease calculations
break

''' Plot the particle motion in the bottle '''
plt.streamplot(Y,Z, Bf[:,:,1], Bf[:,:,2], color="orange")
plt.plot(rp[:,1], rp[:,2], color='navy')
plt.xlim(-10.0,10.0)
plt.ylim(-10.0,10.0)
plt.xlabel("$y$ (m)")
plt.ylabel("$z$ (m)")
plt.title("Motion of Charged Particle in a Magnetic Bottle")


• Nov 28, 2021 at 23:25

• Rather than 4.0 * pow(10, -7), prefer scientific notation literals, i.e. 4e-7
• In nearly all cases there's no need for a .0 suffix. Float promotion will do the right thing with an integer value.
• Prefer using inner tuples rather than inner lists for array initialization, since they're immutable; like np.array((0, 0, 1)).
• Move your global code into functions.
• Prefer the object-oriented (ax.) rather than implicit-state (plt.) interface to matplotlib
• Factor out a function to do your dipole calculation
• Make proper use of np.linalg.norm instead of a manual square-and-square-root
• Do not use a loop to initialize Bf; vectorise this
• There's no need for a time array - you don't use it. All you need is a sample count, which is equal to your end time (one second) divided by your time delta.
• There's also no need for a velocity array. All you need is the current velocity.
• You do a bunch of plotting work that's entirely overwritten, including your streamplot. Don't do this twice - just do it once.
• Your constants for the proton are not as accurate as they should be. This error accumulates over your time iteration, and eventually makes your path deviate perceptibly from where it's supposed to go.
• Some typos, such as Foward -> Forward, positon -> position
• You use implicit outer axis indexing in places like rp[it]. This is more confusing than explicitly showing what's going on - indexing the first axis, taking a slice over the second axis, like rp[it, :].
• You don't actually care that rp starts off as zeros - it can start off as uninitialized ("empty").

## Suggested

from typing import Tuple

import numpy as np
import matplotlib.pyplot as plt

# constants
mu_0 = np.pi * 4e-7            # permeability of free space [kg*m*s^-2*A^-2]
q_p = 1.602_176_634e-19        # proton charge [coulombs]
m_p = 1.672_621_923_695_1e-27  # proton mass [kg]
mu = np.array((0, 0, 1e4))     # set magnetic moment to point in z direction

def make_dipole(pos: np.ndarray, xyz: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
r = xyz - pos  # difference between this point and point of the observer
rmag = np.linalg.norm(r, axis=-1)
if isinstance(rmag, np.ndarray):
rmag = rmag[:, :, np.newaxis]

'''
https://numpy.org/doc/stable/reference/generated/numpy.dot.html
If a is an N-D array and b is a 1-D array,
it is a sum product over the last axis of a and b.
'''

dot = np.dot(r, mu)
if isinstance(dot, np.ndarray):
dot = dot[:, :, np.newaxis]

B1 = 3*r*dot / rmag**5   # first term to the magnetic field
B2 = -mu / rmag**3       # second term
return B1, B2

def bot_field(xyz: np.ndarray) -> np.ndarray:
"""
Calculates magnetic bottle field.
This function will come in handy when we illustrate the field along and ultimately illustrate the
particle motion in field
"""
z_disp = 10  # displacement of the two magnetic dipoles with respect to zero (one at z = -z_disp, the other at +z_disp)
pos = np.array((0, 0, z_disp))  # position of the first dipole

# point dipoles
B1_A, B2_A = make_dipole(pos, xyz)
B1_B, B2_B = make_dipole(-pos, xyz)

return mu_0/4/np.pi * (B1_A + B2_A + B1_B + B2_B)  # field due to magnetic bottle

def calculate_field() -> Tuple[
np.ndarray,  # Y
np.ndarray,  # Z
np.ndarray,  # Bf
]:
"""Setting up graph for dipole magnetic field"""
y = np.arange(-10, 10, 0.1)  # create a grid of points from y = -10 to 10
z = np.arange(-10, 10, 0.1)  # create a grid of points from z = -10 to 10

# Each 200*200
Y, Z = np.meshgrid(y, z)     # create a rectangular grid out of y and z
X = np.zeros_like(Y)

# Each 200*200*3
XYZ = np.stack((X, Y, Z), axis=2)
Bf = bot_field(XYZ)

return Y, Z, Bf

def calculate_path() -> np.ndarray:
q = 2 * q_p  # charge of helium-4
m = 4 * m_p  # mass of helium-4

dt = 1e-5     # timestep
end_time = 1  # seconds
n_samples = round(end_time / dt)

rp = np.empty((n_samples, 3))       # position values
rp[0, :] = np.array((0, -5, 0))  # initial position (m), above the lower dipole

v_o = 100                                     # initial velocity (m/s)
vp = np.array((0, 0, v_o), dtype=np.float64)  # velocity is in the z-direction

# Model the particle motion in the field at each time step (Forward Euler Method)
for it in range(n_samples - 1):
Bp = bot_field(rp[it, :])        # input the current particle position to get the magnetic field
rp[it+1, :] = rp[it, :] + dt*vp  # update the position of the particle based on this velocity

distance = np.linalg.norm(rp[it+1, :])
if distance > 20:  # If the particle escapes (i.e. exceeds 20 m from origin), cease calculations
break

Ap = q / m * np.cross(vp, Bp)  # calculate the magnetic force on the particle
vp += dt*Ap                    # update the velocity of the particle based on this force

return rp

def make_figure() -> Tuple[plt.Figure, plt.Axes]:
fig, ax = plt.subplots()
ax.set_xlim(-10, 10)
ax.set_ylim(-10, 10)
ax.set_xlabel("$y$ (m)")
ax.set_ylabel("$z$ (m)")
ax.set_title("Motion of Charged Particle in a Magnetic Bottle")
return fig, ax

def plot_field(ax: plt.Axes, Y: np.ndarray, Z: np.ndarray, Bf: np.ndarray) -> None:
ax.streamplot(Y, Z, Bf[:, :, 1], Bf[:, :, 2], color='orange')  # plot the magnetic field

def plot_path(ax: plt.Axes, rp: np.ndarray) -> None:
"""Plot the particle motion in the bottle"""
ax.plot(rp[:, 1], rp[:, 2], color='navy')

def main() -> None:
fig, ax = make_figure()

Y, Z, Bf = calculate_field()
plot_field(ax, Y, Z, Bf)

rp = calculate_path()
plot_path(ax, rp)

plt.show()

if __name__ == '__main__':
main()


• I was wondering. I do not know what your background is with numerical analysis, but do you believe the results could also be improved if I replaced Eule Method with a 4th order Runge-Kutta method?
– Dila
Dec 7, 2021 at 19:07
• @Dila I'm sorry, but I don't consider myself qualified to make that judgement. I've used both RK and Euler before but can't meaningfully speak about their trade-offs for your scenario. Dec 8, 2021 at 3:32