Vectorized operations

vectorization is important to write fast and compact code

Here we will often be working with $N$ particles, which gives $N(N-1)/2$ possible pairs of particles. When $N=4$ then there are $6$ possible pairs. We will use np.where() to efficiently apply logical operations to all elements of the arrays. Read the following code, and understand the role of unravel_index to target the particle and index which corresponds to a given logical operation.

import numpy as np

# Create an array of 4 particles in 2 dimensions
position = np.array([[0.25, 0.25], [0.75, 0.25], [0.25, 0.75], [0.75, 0.75]])
print("shape of position", position.shape)

velocity = np.random.normal(size=position.shape)  # velocities with the same dimensions as the position array

# there are N(N-1)/2=6 pairs, which are labeled i,j
i, j = np.triu_indices(position.shape[0], k=1)
print("indexing i, j", i,j)
rij = position[i]-position[j]  # set of all 6 separation vectors
print("rij\n",rij)

# vectorized calculation as a function of each component of the velocity
v0 = np.where(velocity >0 , velocity, velocity*velocity  )
print("v0\n", v0) 

#find the disk and direction corresponding to smallest v0
disk, direction = np.unravel_index(v0.argmin(), v0.shape)
print("smallest value of v0: disk direction", disk, direction)

shape of position (4, 2)
indexing i, j [0 0 0 1 1 2] [1 2 3 2 3 3]
rij
 [[-0.5  0. ]
 [ 0.  -0.5]
 [-0.5 -0.5]
 [ 0.5 -0.5]
 [ 0.  -0.5]
 [-0.5  0. ]]
v0
 [[0.70798478 1.26364151]
 [1.66493774 1.55396407]
 [0.23827062 0.75921693]
 [0.6386523  1.02382915]]
smallest value of v0: disk direction 2 0