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: