Ytdyklly

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*(10ⁿ)).

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 100⁶ 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 instead, 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.