Path of a particle in a non-uniform B field

Question

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")

Answer 1

• 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)).
• Add PEP484 type hints.
• 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()