非均匀B场中粒子的路径

Question

我已经完成了我的项目，它模拟了一个粒子被困在一个称为磁镜。的设备中的路径，我希望你对我的实现有一些想法和改进。

以下是代码：

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

• 而不是4.0 * pow(10, -7)，而是更喜欢科学符号，即4e-7
• 几乎在所有情况下，都不需要.0后缀。浮点提升将使用整数值执行正确的操作。
• 更喜欢使用内部元组而不是内部列表进行数组初始化，因为它们是不可变的；比如np.array((0, 0, 1))。
• 添加PEP484类型提示。
• 将全局代码移动到函数中。
• 相对于matplotlib，它更喜欢面向对象(ax.)接口而不是隐式状态(plt.)接口。
• 算出一个函数来计算偶极子。
• 正确使用np.linalg.norm，而不是手动平方根。
• 不要使用循环初始化Bf；将此矢量化
• 不需要时间数组--你不用它。您所需要的只是一个样本计数，它等于您的结束时间(1秒)除以您的时间增量。
• 也不需要速度阵列。你所需要的只是当前的速度。
• 你做了一堆完全覆盖的绘图工作，包括你的streamplot。不要这样做两次，只要做一次。
• 你的质子常数是不像应该的那样准确。这个错误会在你的时间迭代过程中累积，最终会使你的路径明显偏离它应该去的地方。
• 一些排版，如Foward -> Forward，positon -> position
• 在rp[it]这样的地方使用隐式外轴索引。这比显式地显示正在发生的事情更令人困惑--索引第一个轴，在第二个轴上取一个切片，比如rp[it, :]。
• 实际上，您并不关心rp从零开始--它可以从未初始化(“空”)开始。

建议

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