Simulating a two-body collision problem to find digits of Pi

I came across a nice video on 3Blue1Brown's channel that highlights a very indirect way to find the digits in Pi. I'd suggest watching the whole video, but briefly: The setup is as above. A "small" body of unit mass is at rest on the left and a "large" body on the right is moving towards the left (the initial positions and velocities are irrelevant). Assuming perfectly elastic collisions and that the large body's mass to be the n-th power of 100, where n is a natural number, the total number of collisions is always floor(pi*(10n)).

The analytical reason for this is discussed here, but I've started learning a bit of Python and wanted to simulate this to improve my programming (i.e. I'm a beginner). Here's the code:

from fractions import Fraction

class Block:
# Creating a class since that helps keep track of attributes
def __init__(self, mass, velocity, position):
self.mass = mass
self.velocity = velocity
self.position = position

# Set initial conditions: the object on the left is at rest at x = 5 and has
# unit mass. The one on the right starts off at x = 10, with velocity =
# -5 units/s and has mass equal to 100^n, where n is user-specified.
# The number of collisions should be floor(Pi*10^n). e.g. n = 2 => 314,
# n = 3 => 3141, and so on

small = Block(1, 0, Fraction(32/10))
large = Block(100**int(input("Which power of 100 is the second mass? ")),
-7, Fraction(75,10))

# By "collision", we mean that either the position of both the objects is the
# same (both collide against each other) or the position of the small block is
# 0 (collision against wall)

def updateVelocities(collisions):
if(small.position == large.position != 0):
# Both blocks collide against each other
collisions += 1
temp = small.velocity

small.velocity = Fraction(((2*large.mass*large.velocity)+
(small.mass*small.velocity)-(large.mass*small.velocity)),
(small.mass + large.mass))

large.velocity = Fraction(((2*small.mass*temp)+(large.mass*large.velocity)
-(small.mass*large.velocity)),(small.mass + large.mass))

elif(small.position == 0 != large.position):
# The small block gets "reflected" off the wall
collisions += 1
small.velocity = -small.velocity

elif(small.position == large.position == 0):
# The rare instance in which both blocks move towards the wall and
# collide with the wall and against each other simultaneously
collisions += 2
small.velocity, large.velocity = -small.velocity, -large.velocity
else:
pass

return collisions

# Given the current positions and velocities, find the time to next collision
# This takes care of all different scenarios
def timeToNextCollision():
if(large.velocity >= small.velocity >= 0):
# Both blocks move towards right, but the large block is faster and the
# small block can't catch up
return float("inf")

elif(small.velocity >= 0 >= large.velocity):
# Both blocks are either moving towards each other, or one of the is at
# rest and the other is moving towards it. The wall is obviously ignored
# The condition small.velocity == 0 == large.velocity will also be ignored
# since if that's true, only the first if statement would be executed.
return Fraction(large.position - small.position,
small.velocity - large.velocity)

elif((large.velocity >= 0 and small.velocity < 0) or
(small.velocity <= large.velocity < 0)):
# Both blocks move towards left, but the large block can't catch up with
# the small block before the latter runs into the wall
return Fraction(-small.position, small.velocity)

elif(small.position == 0):
# Special case for when the small block is currently at the wall
if(large.velocity >= abs(small.velocity)):
# Small block can no longer catch up with large block
return float("inf")
else:
# Large block is either moving towards left or too slow moving towards
# the right. In either case, they will collide
return large.position/(abs(small.velocity) - large.velocity)
else:
# Both blocks move towards left, but large block is faster. If the
# distance between blocks is small enough compared to that between the wall
# and the small block, they will collide. Otherwise the small block will
# reach the wall before the large block has a chance to catch up
return min(Fraction(-small.position, small.velocity),
Fraction(large.position - small.position),
(small.velocity - large.velocity))

collisionCount = 0

while True:
t = timeToNextCollision()
if(t == float("inf")):
# No more collisions
break
# Update the distances to what they'll be during the next collision
small.position += small.velocity*t
large.position += large.velocity*t
# Update collision count AND velocities to post-collision values
collisionCount = updateVelocities(collisionCount)

print(collisionCount)

The biggest headache was dealing with float's precision issues. For example, for a collision to register, the updated positions of the two blocks should exactly be the same. But with float's rounding issues, there would be a slight difference in the positions and the program would go haywire.

Even though this was corrected by using the Fraction data type, the run time of the program is really slow. If n=2, the program finishes within milliseconds, whereas for n=3, it takes a whopping 115 seconds. I'm not really aware of all the nuances of Python nor do I have any computer science knowledge, so I'd be really grateful if I can get some guidance on how I can improve the code.

Off the top of my head, perhaps using Fraction affects the run time, or I wrote the if conditions in a clumsy way, or wrote redundant code. I'm confused because there are many possibilities. In the first video I linked to, at around the 2:00 min mark, there's a simulation of collisions between 1 kg and 1003 kg, and it's so smooth and quick!

P.S. I used the Block class just to improve readability. I've tried using a simple list instead of a Block class object, but that only shaves off around 8 seconds or so. Not much improvement. I also tried using the numpy double data type - again, same rounding issues as with the default float type.

• what version of python are you using? – arnau vivet pares Jan 22 at 18:42
• @arnauvivetpares: I'm using 3.7 – Shirish Kulhari Jan 24 at 11:31

First, this is an awesome video! Upvoted for that reason alone. :)

If n=2, the program finishes within milliseconds, whereas for n=3, it takes a whopping 115 seconds.

Do you know about big-O notation? Just off the top of my head after watching that video: for n=2 you're computing the number of collisions for a 1kg block and a 100**2 = 10,000kg block, which we know from the video should be 314 collisions. For n=3 you're simulating 3141 collisions. That's 10 times as many collisions, simulated one at a time, so it should take 10 times as long. In general your program is going to take O(10n) steps to compute its result. So you shouldn't be surprised if it gets real slow real fast.

However, you're saying the difference between n=2 and n=3 is a factor of more than 100. That's surprising. I think you're right to blame Fraction — that is, you're dealing with much bigger numerators and denominators in the n=3 case, and bigger numbers naturally take more time to manipulate.

if(large.velocity >= small.velocity >= 0):

In general I like your use of chained comparisons... but why did you almost invariably put the biggest number on the left and the smallest on the right? That's backwards from how we usually write number lines.

And then sometimes you don't chain the comparison at all, for no reason I can detect:

elif((large.velocity >= 0 and small.velocity < 0) or
(small.velocity <= large.velocity < 0)):

I'd write this as

elif (small.velocity < 0 <= large.velocity) or (small.velocity <= large.velocity < 0):

Notice that you don't need parens around the condition of an if or elif in Python, and so it's not usual to write them.

return large.position/(abs(small.velocity) - large.velocity)

You didn't use Fraction here. Was that an oversight? Also, if this is a floating-point division, I might want to blame "repeated conversion from Fraction to floating-point and back" for some of your performance problems.

large = Block(100**int(input("Which power of 100 is the second mass? ")),
-7, Fraction(75,10))

I strongly recommend moving the input onto its own source line. Mixing user input into the middle of an arithmetic expression is just asking for trouble. Plus, this lets you unambiguously name the n that you were trying to talk about in your question:

n = int(input("Which power of 100 is the second mass? "))
large = Block(100**n, -7, Fraction(75,10))

Actually, hang on; back up!

# Set initial conditions: the object on the left is at rest at x = 5 and has
# unit mass. The one on the right starts off at x = 10, with velocity =
# -5 units/s and has mass equal to 100^n, where n is user-specified.
# The number of collisions should be floor(Pi*10^n). e.g. n = 2 => 314,
# n = 3 => 3141, and so on

small = Block(1, 0, Fraction(32/10))
large = Block(100**int(input("Which power of 100 is the second mass? ")),
-7, Fraction(75,10))

That comment is ridiculously untrue! The object on the left is at rest at x = 3.2 (or x = 3 if you're in Python 2), and the object on the right starts off at x = 7.5 with velocity -7 units per second, not -5. So the comment is completely wrong. Besides, starting the big block with anything other than "velocity -1" is just wasting bits and CPU cycles. Who wants to multiply anything by 32/10 when you could be multiplying it by 1?

Also, all of that initial setup should be encapsulated into a __main__ block:

if __name__ == '__main__':
n = int(input("Which power of 100 is the second mass? "))
small = Block(1, 0, 1)
large = Block(100**n, -1, 2)
collisionCount = 0

while True:
t = timeToNextCollision(small, large)
if t == float("inf"):
# No more collisions
break
# Update the distances to what they'll be during the next collision
small.position += small.velocity * t
large.position += large.velocity * t
# Update collision count AND velocities to post-collision values
collisionCount = updateVelocities(small, large, collisionCount)
print(collisionCount)

I changed your timeToNextCollision() to timeToNextCollision(small, large), passing the blocks to it as parameters, since it needs to look at the blocks in order to know what to do.

# Both blocks move towards left, but large block is faster. If the
# distance between blocks is small enough compared to that between the wall
# and the small block, they will collide. Otherwise the small block will
# reach the wall before the large block has a chance to catch up
return min(Fraction(-small.position, small.velocity),
Fraction(large.position - small.position),
(small.velocity - large.velocity))

I strongly recommend running your whole program through pylint or pyflakes and fixing all the style issues. Here specifically, I think it would benefit from the usual Python indentation style, which looks like this:

return min(
Fraction(-small.position, small.velocity),
Fraction(large.position - small.position),
(small.velocity - large.velocity)
)

This makes it very clear that you're taking the min of three things, not the usual two — and also you're constructing a Fraction from one argument, not the usual two. If this is intentional behavior, then the indentation is important because it communicates to your reader, "Hey, I know what I'm typing, don't worry" — and if this is unintentional behavior, then hey, you just found one of your bugs!

Finally, let's fix your performance issue.

As I said, I assume that your program takes so long because you're manipulating gigantic numerators and denominators. Above about 2**64 (or maybe 2**32, I'm not sure), Python is going to switch from native integer representation to bignum representation, and get super slow. Reading the fractions docs tells me that there's a limit_denominator method that's used precisely to keep the numerator and denominator small. So let's use it!

while True:
t = timeToNextCollision(small, large)
if t == float("inf"):
# No more collisions
break
# Update the distances to what they'll be during the next collision
small.position += small.velocity * t
large.position += large.velocity * t
collisionCount = updateVelocities(collisionCount, small, large)

# Here's the new code!
small.position = small.position.limit_denominator(2**32)
large.position = large.position.limit_denominator(2**32)
small.velocity = small.velocity.limit_denominator(2**32)
large.velocity = large.velocity.limit_denominator(2**32)

With just this change (and the cleanups mentioned in this review, including fixing that bug that pyflakes would have found), I see your program taking the O(10n) that we expect it to:

$time echo '2' | python x.py Which power of 100 is the second mass? 314 real 0m0.240s user 0m0.196s sys 0m0.018s$ time echo '3' | python x.py
Which power of 100 is the second mass? 3141

real 0m1.721s
user 0m1.697s
sys  0m0.015s

\$ time echo '4' | python x.py
Which power of 100 is the second mass? 31415

real 0m22.497s
user 0m20.226s
sys  0m0.160s

Problem solved!

• Thanks so much for the detailed feedback! Regarding the min formula with nested fractions, it's indeed technically incorrect but it runs fine and gives the correct answer, which is strange. I corrected that though. Regarding the comment, I'd actually updated the program with new values just for more testing and forgot to update the comment. Also, I want to upload this thing to Github (though it may seem like a trivial toy program to experienced programmers). Are you okay with me crediting you (because you made a big improvement to the performance)? Let me know! – Shirish Kulhari Jan 21 at 6:43
• @ShirishKulhari: Sure, I don't mind. Re the typo not mattering: Consider whether that branch of the code is ever even reached. (You could put an assert False in there to see what happens.) – Quuxplusone Jan 21 at 7:12
• Okay, I'll credit you by your username in the program comment - hopefully that's the right protocol. One more thing, you restricted the numerator and denominator of the positions to 32 bits. Couldn't I do the same for the velocities as well? Or would that start causing inaccurate results? I noticed when I changed the 2^32 in your newly added lines to 2^16, it gives the result 3.14169 for n=100^5. 2**32 is totally okay, but should the same limiting be done for all the rest of fraction calculations too? – Shirish Kulhari Jan 21 at 7:17
• @Graipher: Yeah I wrote all that in the comments, providing a link to this question and his profile and all. Hopefully that's enough? github.com/Shirish-Kulhari/numerical-methods/blob/master/… – Shirish Kulhari Jan 21 at 17:52

I would like to explain why my code has a better implementation of the problem. one thing to understand is that all possible collisions will be over when the right side block is moving right and the left block is also moving right with a lower velocity. this situation will continue till infinity.

Also, we dont need the position parameter at all as we only need to check velocity change to know a collision. The case when the left block is moving towards the wall has also been considered in the last if statement where i reverse its velocity. I am new in this forum and if any changes are needed, just leave a note and will do the same.

I am a beginner to programming and can code only in MATLAB. I found this problem interesting and hence tried solving it.

%Written by Shubham Wani, Mechanical Undergrad at NIT Trichy.
%Set the mb(stands for Mass B) variable to powers of 100
%COLLIDING BLOCKS COMPUTE PI
%Takes about 0.043737 seconds for calculations for 100000000 kg Mass B
tic
clearvars;
clear ;
clear clc;
ma = 1;
%set mb to various powers of 100
mb = 10000000000;
va(1)= double(0);
vb(1)= double(-10);
n=1;
while (true)
**%check if it is the last possible collision**
if (vb(n)>0 && va(n)<vb(n))
break;
end
**%Calculate new velocities after the collision**
va(n+1)=(ma*va(n)+mb*vb(n)+mb*(vb(n)-va(n)))/(ma+mb);
va(n+1);
vb(n+1)=(ma*va(n)+mb*vb(n)+ma*(va(n)-vb(n)))/(ma+mb);
vb(n+1);
n=n+1;
**%if mass A is moving towards left, invert its speed.**
if (va(n)<0)
va(n+1)=-va(n);
vb(n+1)=vb(n);
n=n+1;
end
end
**%Number of collisions=n-1 as indexing starts from 1**
disp(n-1);
toc

the comments are self explanatory.

• @shirish-kulhari this code is also available on my github handle shubhamwani376 – Shubham Wani Jan 27 at 15:10
• That's very nice! +1 for the faster code, though I think I may have seen another similar implementation in a different programming language. Either would work much faster than my code. I guess vectorizing the velocity update would also speed things up a bit - not entirely sure though.. – Shirish Kulhari Jan 29 at 9:51