Electric force between particles using numpy arrays

Question

I am trying to simulate a particle flying at another particle while undergoing electrical repulsion (or attraction), called Rutherford-scattering. I have succeeded in simulating (a few) particles using for loops and python lists. However, now I want to use numpy arrays instead. The model will use the following steps:

1. For all particles:
   1. Calculate radial distance with all other particles
   2. Calculate the angle with all other particles
   3. Calculate netto force in x-direction and y-direction
2. Create matrix with netto xForce and yForce for each particle
3. Create accelaration (also x and y component) matrix by a = F/mass
4. Update speed matrix
5. Update position matrix

My problem is that I do not know how I can use numpy arrays in calculating the force components. Here follows my code which is not runnable.

import numpy as np
# I used this function to calculate the force while using for-loops.
def force(x1, y1, x2, x2):
    angle =  math.atan((y2 - y1)/(x2 - x1))
    dr = ((x1-x2)**2 + (y1-y2)**2)**0.5
    force = charge2 * charge2 / dr**2 
    xforce = math.cos(angle) * force 
    yforce = math.sin(angle) * force

    # The direction of force depends on relative location
    if x1 > x2 and y1<y2:
        xforce = xforce
        yforce = yforce
    elif x1< x2 and y1< y2:
        xforce = -1 * xforce
        yforce = -1 * yforce
    elif x1 > x2 and y1 > y2:
        xforce = xforce
        yforce = yforce
    else:
        xforce = -1 * xforce
        yforce = -1* yforce

    return xforce, yforce

def update(array):

    # this for loop defeats the entire use of numpy arrays
    for particle in range(len(array[0])):
        # find distance of all particles pov from 1 particle

        # find all x-forces and y-forces on that particle

        xforce = # sum of all x-forces from all particles
        yforce = # sum of all y-forces from all particles
        force_arr[0, particle] = xforce
        force_arr[1, particle] = yforce

    return force

# begin parameters
t = 0
N = 3
masses = np.ones(N)
charges = np.ones(N)
loc_arr = np.random.rand(2, N)
speed_arr = np.random.rand(2, N)
acc_arr = np.random.rand(2, N)
force = np.random.rand(2, N)


while t < 0.5:
    force_arr = update(loc_arry)
    acc_arr = force_arr / masses
    speed_arr += acc_array
    loc_arr += speed_arr
    t += dt

    # plot animation

Answer 1

One approach to model this problem with arrays may be:

• define the point coordinates as a Nx2 array. (This will help with extensibility if you advance to 3-D points later)
• define the intermediate variables distance, angle, force as NxN arrays to represent the pairwise interactions

Numpy things to know about:

• You can call most numeric functions on arrays if the arrays have the same shape (or conforming shapes, which is a nontrivial topic...)
• meshgrid helps you generate the array indices necessary to shapeshift your Nx2 arrays to compute NxN results
• and a tangential note (ha ha) arctan2() computes a signed angle, so you can bypass the complex "which quadrant" logic

For example you can do something like this. Note in get_dist and get_angle the arithmetic operations between points take place in the bottom-most dimension:

import numpy as np

# 2-D locations of particles
points = np.array([[1,0],[2,1],[2,2]])
N = len(points)  # 3

def get_dist(p1, p2):
    r = p2 - p1
    return np.sqrt(np.sum(r*r, axis=2))

def get_angle(p1, p2):
    r = p2 - p1
    return np.arctan2(r[:,:,1], r[:,:,0])

ii = np.arange(N)
ix, iy = np.meshgrid(ii, ii)

dist = get_dist(points[ix], points[iy])
angle = get_angle(points[ix], points[iy])
# ... compute force
# ... apply the force, etc.

For the sample 3-point vector shown above:

In [246]: dist
Out[246]: 
array([[0.        , 1.41421356, 2.23606798],
       [1.41421356, 0.        , 1.        ],
       [2.23606798, 1.        , 0.        ]])

In [247]: angle / np.pi     # divide by Pi to make the numbers recognizable
Out[247]: 
array([[ 0.        , -0.75      , -0.64758362],
       [ 0.25      ,  0.        , -0.5       ],
       [ 0.35241638,  0.5       ,  0.        ]])

Answer 2

Here is one go with only a loop for each time step, and it should work for any number of dimensions, I have tested with 3 too:

from matplotlib import pyplot as plt
import numpy as np

fig, ax = plt.subplots()

N = 4
ndim = 2
masses = np.ones(N)
charges = np.array([-1, 1, -1, 1]) * 2
# loc_arr = np.random.rand(N, ndim)
loc_arr = np.array(((-1,0), (1,0), (0,-1), (0,1)), dtype=float)
speed_arr = np.zeros((N, ndim))

# compute charge matrix, ie c1 * c2
charge_matrix = -1 * np.outer(charges, charges)

time = np.linspace(0, 0.5)
dt = np.ediff1d(time).mean()

for i, t in enumerate(time):
    # get (dx, dy) for every point
    delta = (loc_arr.T[..., np.newaxis] - loc_arr.T[:, np.newaxis]).T
    # calculate Euclidean distance
    distances = np.linalg.norm(delta, axis=-1)
    # and normalised unit vector
    unit_vector = (delta.T / distances).T
    unit_vector[np.isnan(unit_vector)] = 0 # replace NaN values with 0

    # calculate force
    force = charge_matrix / distances**2 # norm gives length of delta vector
    force[np.isinf(force)] = 0 # NaN forces are 0

    # calculate acceleration in all dimensions
    acc = (unit_vector.T * force / masses).T.sum(axis=1)
    # v = a * dt
    speed_arr += acc * dt

    # increment position, xyz = v * dt
    loc_arr += speed_arr * dt 

    # plotting
    if not i:
        color = 'k'
        zorder = 3
        ms = 3
        for i, pt in enumerate(loc_arr):
            ax.text(*pt + 0.1, s='{}q {}m'.format(charges[i], masses[i]))
    elif i == len(time)-1:
        color = 'b'
        zroder = 3
        ms = 3
    else:
        color = 'r'
        zorder = 1
        ms = 1
    ax.plot(loc_arr[:,0], loc_arr[:,1], '.', color=color, ms=ms, zorder=zorder)

ax.set_aspect('equal')

The above example produces, where the black and blue points signify the start and end positions, respectively:

And when charges are equal charges = np.ones(N) * 2 the system symmetry is preserved and the charges repel:

And finally with some random initial velocities speed_arr = np.random.rand(N, 2):

EDIT

Made a small change to the code above to make sure it was correct. (I was missing -1 on the resultant force, ie. force between +/+ should be negative, and I was summing down the wrong axis, apologies for that. Now in the cases where masses[0] = 5, the system evolves correctly:

Answer 3

The classic approach is to calculate electric field for all particles in the system. Say you have 3 charged particles all with positive charge:

particles = np.array([[1,0,0],[2,1,0],[2,2,0]]) # location of each particle
q = np.array([1,1,1]) # charge of each particle

The easiest way to compute the electric field at each particle`s location is for loop:

def for_method(pos,q):
    """Computes electric field vectors for all particles using for-loop."""
    Evect = np.zeros( (len(pos),len(pos[0])) ) # define output electric field vector
    k =  1 / (4 * np.pi * const.epsilon_0) * np.ones((len(pos),len(pos[0]))) * 1.602e-19 # make this into matrix as matrix addition is faster
    # alternatively you can get rid of np.ones and just define this as a number
    
    for i, v0 in enumerate(pos): # s_p - selected particle | iterate over all particles | v0 reference particle
        for v, qc in zip(pos,q): # loop over all particles and calculate electric force sum | v particle being calculated for
            if all((v0 == v)):   # do not compute for the same particle
                continue
            else:
                r = v0 - v       #
                Evect[i] += r / np.linalg.norm(r) ** 3 * qc #! multiply by charge
    return Evect * k
# to find electric field at each particle`s location call
for_method(particles, q)

This function returns array of vectors with the same shape as input particles array. To find force on each, you simply multiply this vector with q array of charges. From there on, you can easily find your acceleration and integrate the system using your favourite ODE solver.

Performance Optimization & Accuracy

For method is the slowest possible approach. The field can be computed using solely linear algebra granting significant speed boost. Following code is very efficient Numpy matrix "one-liner" (almost one-liner) to this problem:

def CPU_matrix_method(pos,q):
    """Classic vectorization of for Coulomb law using numpy arrays."""
    k = 1 / (4 * np.pi * const.epsilon_0) * np.ones((len(pos),3)) * 1.602e-19 # define electric constant
    dist = distance.cdist(pos,pos)  # compute distances
    return k * np.sum( (( np.tile(pos,len(pos)).reshape((len(pos),len(pos),3)) - np.tile(pos,(len(pos),1,1))) * q.reshape(len(q),1)).T * np.power(dist,-3, where = dist != 0),axis = 1).T

Note that this and following code also return electric field vector for each particle.

You can get even higher performance if you offload this onto the GPU using Cupy library. Following code is almost identical to the CPU_matrix_method, I only expanded the one-liner a little so that you could see better what is going on:

def GPU_matrix_method(pos,q):
    """GPU Coulomb law vectorization.
    Takes in numpy arrays, performs computations and returns cupy array"""
    # compute distance matrix between each particle
    k_cp = 1 / (4 * cp.pi * const.epsilon_0) * cp.ones((len(pos),3)) * 1.602e-19 # define electric constant, runs faster if this is matrix
    dist = cp.array(distance.cdist(pos,pos)) # could speed this up with cupy cdist function! use this: cupyx.scipy.spatial.distance.cdist
    pos, q = cp.array(pos), cp.array(q) # load inputs to GPU memory
    dist_mod = cp.power(dist,-3)        # compute inverse cube of distance
    dist_mod[dist_mod == cp.inf] = 0    # set all infinity entries to 0 (i.e. diagonal elements/ same particle-particle pairs)
    # compute by magic
    return k_cp * cp.sum((( cp.tile(pos,len(pos)).reshape((len(pos),len(pos),3)) - cp.tile(pos,(len(pos),1,1))) * q.reshape(len(q),1)).T * dist_mod, axis = 1).T

Regarding the accuracy of the mentioned algorithms, if you compute the 3 methods on the particles array you get identical results:

[[-6.37828367e-10 -7.66608512e-10  0.00000000e+00]
 [ 5.09048221e-10 -9.30757576e-10  0.00000000e+00]
 [ 1.28780145e-10  1.69736609e-09  0.00000000e+00]]

Regarding the performance, I computed each algorithm on systems ranging from 2 to 5000 charged particles. Additionally I also included Numba precompiled version of the for_method to make the for-loop approach competitive: We see that for-loop performs terribly needing over 400 seconds to compute for system with 5000 particles. Zooming in to the bottom part: This shows that matrix approach to this problem is orders of magnitude better. To be exact the 5000 particle evaluation took 18.5s for Numba for-loop, 4s for CPU matrix(5 times faster than Numba), and 0.8s for GPU matrix* (23 times faster than Numba). The significant difference shows for larger arrays.

* GPU used was Nvidia K100.