3
\$\begingroup\$

This code calculates pi via collisions; it asks for a user input of N which determines the mass of the second block. It is fully working, it just takes forever to run when N >= 2. I want to be able have at least N=5 and have a reasonable runtime. The issue is in the velocity and x1, x2 calculation I believe. There are just so many that need to be appended that it takes forever to run and then just as long to append.

My code shows that I have tried to use Numba to speed up the runtime, but that doesn't seem to be helping. I am currently using the RK-4 method to update position, and I previously tried the Verlet method which did not seem to have an affect on the runtime. Any help on this is greatly appreciated.

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.animation import FuncAnimation
import matplotlib.animation as animation
from numba import jit
import time

start = time.time()

@jit(nopython=True)
def RK4_1(x, vx):
    return vx

@jit(nopython=True)
def RK4_2(x, vx):
    return 0

@jit(nopython=True)
def iterate(x1, x2, vx1, vx2, col):
    k1 = dt*RK4_1(x1, vx1)
    k2 = dt*RK4_1(x1 + k1/2, vx1)
    k3 = dt*RK4_1(x1 + k2/2, vx1)
    k4 = dt*RK4_1(x1 + k3, vx1)
    x1 += (1/6)*(k1 + 2*k2 + 2*k3 + k4)

    k1 = dt*RK4_1(x2, vx2)
    k2 = dt*RK4_1(x2 + k1/2, vx2)
    k3 = dt*RK4_1(x2 + k2/2, vx2)
    k4 = dt*RK4_1(x2 + k3, vx2)
    x2 += (1/6)*(k1 + 2*k2 + 2*k3 + k4)

    k1 = dt*RK4_2(x1, vx1)
    k2 = dt*RK4_2(x1, vx1 + k1/2)
    k3 = dt*RK4_2(x1, vx1 + k2/2)
    k4 = dt*RK4_2(x1, vx1 + k3)
    vx1 += (1/6)*(k1 + 2*k2 + 2*k3 + k4)

    k1 = dt*RK4_2(x2, vx2)
    k2 = dt*RK4_2(x2, vx2 + k1/2)
    k3 = dt*RK4_2(x2, vx2 + k2/2)
    k4 = dt*RK4_2(x2, vx2 + k3)
    vx2 += (1/6)*(k1 + 2*k2 + 2*k3 + k4)

    if x1 < 0:
        x1 = 0
        vx1 = -vx1
        col += 1

    if x2 < x1:
        x2 = x1
        vx1_i = vx1
        vx2_i = vx2
        vx1 = (2*m2*vx2_i + m1*vx1_i - m2*vx1_i)/(m1+m2)
        vx2 = (2*m1*vx1_i + m2*vx2_i - m1*vx2_i)/(m1+m2)
        col += 1

    return x1, x2, vx1, vx2, col

dt = 0.01

m1 = 1
N = int(input("Enter an integer N that will determine the mass of the second block: "))
m2 = 100**N

w1 = 1
w2 = w1*(100**N)**(1/3)
x1 = 1
x2 = 1.15
y1 = 1
y2 = 1
vx1 = 0
vx2 = -1
col = 0

x1arr = np.array([])
x2arr = np.array([])
y1arr = np.array([])
y2arr = np.array([])
vx1arr = np.array([])
vx2arr = np.array([])
colarr = np.array([])

t = 0

while (vx2 < 0) or (abs(vx1) > abs(vx2)):
    x1, x2, vx1, vx2, col = iterate(x1, x2, vx1, vx2, col)
    #print(vx1, vx2)
    t += dt

    x1arr = np.append(x1arr, x1)
    x2arr = np.append(x2arr, x2)
    y1arr = np.append(y1arr, y1)
    y2arr = np.append(y2arr, y2)
    vx1arr = np.append(vx1arr, vx1)
    vx2arr = np.append(vx2arr, vx2)
    colarr = np.append(colarr, col)

print("Number of collisions: %f" % (col))

speed = 1000

def update_plot(i, fig, scat1, scat2, txt):
    last = 0
    if(i>int(len(x1arr)/speed)-2):
        last=1  
    s = int(speed*i)
    scat1.set_data(x1arr[s],y1arr[s])
    scat2.set_data(x2arr[s],y2arr[s])
    #scat.set_sizes(lx1=5, lx2=5*N)
    txt.set_text('x1= %.3f   m1=%.0f\nx2= %.3f   m1=%.0f\nCollisions=%.0f\n t=%.3fs' % (x1arr[s],m1,x2arr[s],m2,colarr[s]+last,(s*dt))) #update of legend
    #print("Frame %d Rendered" % (s))
    return scat1, scat2, txt,

size = 5*N
fig =  plt.figure() 
ax = fig.add_subplot(111)
ax.set_xlim([0, 3]) #animation scale
ax.set_ylim([0,5])
ax.grid()
txt = ax.text(0.05, 0.8, '', transform=ax.transAxes) 
scat1, = ax.plot([], [],'s', c='r', markersize=5) 
scat2, = ax.plot([], [],'s', c='r', markersize=5*(N+1))
anim = FuncAnimation(fig, update_plot, fargs = (fig, scat1, scat2, txt), frames = int(len(x1arr)/speed), interval = 1, blit=True, repeat=False)
anim.save("originalpi.mp4", fps=30, bitrate=-1)

end = time.time()
print("Total runtime in seconds:  ", end-start)

plt.show()
\$\endgroup\$
  • \$\begingroup\$ I'm confused, can RK4_1(x, vx) just be replaced with vx? so k1 = dt*RK4_1(x1, vx1) is k1 = dt*vx1. And k1 = dt*RK4_2(x1, vx1) is k1 = 0. \$\endgroup\$ – Peilonrayz May 3 at 16:44
  • \$\begingroup\$ Yes I believe I could do that, I wrote it this way so that the RK 4 process was visible. Do you think changing it would decrease runtime? \$\endgroup\$ – Maria May 3 at 16:57
  • \$\begingroup\$ What exactly do you mean by "calculates pi"? It certainly doesn't print it. \$\endgroup\$ – Peter Taylor May 4 at 13:56
  • \$\begingroup\$ If you run this code for, say, N=1, an animation will appear that shows the blocks colliding. The number of collisions count out pi. For N=1 the number of collisions is 31, for N=2 it is 314, etc. \$\endgroup\$ – Maria May 5 at 14:37
4
\$\begingroup\$

I am currently using the RK-4 method to update position

Firstly, the naming is unhelpful here. A name like RK4_1 implies that the method does the quadrature. I think that function is really dx_dt, RK4_2 is really dvx_dt, and there would be room to pull out a function like

def RK4_step(t, x, dx_dt, dt):
    k1 = dt * dx_dt(x, t)
    k2 = dt * dx_dt(x + k1 / 2, t + dt / 2)
    k3 = dt * dx_dt(x + k2 / 2, t + dt / 2)
    k4 = dt * dx_dt(x + k3, t + dt)
    return x + (1/6) * (k1 + 2 * k2 + 2 * k3 + k4)

But secondly, RK4 is completely overkill here. If we inline

@jit(nopython=True)
def RK4_1(x, vx):
    return vx

@jit(nopython=True)
def RK4_2(x, vx):
    return 0

@jit(nopython=True)
def iterate(x1, x2, vx1, vx2, col):
    k1 = dt*RK4_1(x1, vx1)
    k2 = dt*RK4_1(x1 + k1/2, vx1)
    k3 = dt*RK4_1(x1 + k2/2, vx1)
    k4 = dt*RK4_1(x1 + k3, vx1)
    x1 += (1/6)*(k1 + 2*k2 + 2*k3 + k4)

    ...

    k1 = dt*RK4_2(x1, vx1)
    k2 = dt*RK4_2(x1, vx1 + k1/2)
    k3 = dt*RK4_2(x1, vx1 + k2/2)
    k4 = dt*RK4_2(x1, vx1 + k3)
    vx1 += (1/6)*(k1 + 2*k2 + 2*k3 + k4)

we get

def iterate(x1, x2, vx1, vx2, col):
    x1 += dt*vx1

    ...

    vx1 += 0

It seems that it would be much quicker to solve simultaneous linear equations to work out when the collision will happen rather than to use quadrature. In fact, I had a go, and the actual simulation takes less than a second for N = 5.


My code shows that I have tried to use Numba to speed up the runtime, but that doesn't seem to be helping.

The jit is on the wrong parts. I added some debug printing and observed that it slowed down over time, at which point I very quickly spotted the real problem:

    x1arr = np.append(x1arr, x1)
    x2arr = np.append(x2arr, x2)
    y1arr = np.append(y1arr, y1)
    y2arr = np.append(y2arr, y2)
    vx1arr = np.append(vx1arr, vx1)
    vx2arr = np.append(vx2arr, vx2)
    colarr = np.append(colarr, col)

This is copying the entire array every time, so the length of a step is proportional to the number of steps taken.

Solution: use normal lists rather than numpy arrays. numpy is not a panacea: it's good for parallel processing, but that's not what you want here.

However, the one advantage that numpy arrays give you here is that they don't have length bounds. In testing with N = 3 using normal Python lists I got a memory error after slightly more than 245 million steps and 3140 collisions. It might need a 2D structure to handle the large quantities of data. Of course, that doesn't prevent the extremely likely problem of the plot running into memory problems too. A complete structural rethink may be necessary.


If you do decide to implement RK4, it's more correct to do a single multi-variable quadrature rather than multiple single-variable ones. Here is where there is a good reason to use numpy arrays. What I mean by multi-variable quadrature is something like

k1_x, k1_vx = dt * RK4(x, vx), dt * RK4(vx, ax)
k2_x, k2_vx = dt * RK4(x + k1_x / 2, vx + k1_vx / 2), dt * RK4(vx + k1_vx / 2, ax)

etc. (NB For clarity I separated out the variables. For speed you'd put x and vx in a single numpy array).


It makes no sense to start timing before taking the input: the time I take to type 3 and press Enter is not really runtime.


In terms of general readability, the inline mixing of functions and top-level code is not helpful. It's good to put constants like dt at the top, so that when the reader comes to the usage of dt in iterate they've already seen it. The main loop should probably be a function which returns the arrays (or perhaps a single 2D array) with intermediate values; the plot creation should be a function, and then Python best practice is for the "main" functionality to go at the end with a guard:

if __name__ == "__main__":
    N = int(input("Enter an integer N that will determine the mass of the second block: "))
    start = time.time()
    data = generate_collisions(N)
    plot = graph_collisions(data)
    end = time.time()
    print("Total runtime in seconds:  ", end - start)
    plot.show()
\$\endgroup\$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.