collision.py

inivels = (0,-1)
fp = open('col.dat','w')

def col(vels,a):
    (u,v) = vels # u : small ball || v : large ball
    uu = ((1-a)*u + 2*a*v)/(1+a)
    vv = (2*u + (a-1)*v)/(1+a)
    return (uu,vv)

def cols(vels,a):
    (u,v) = vels
    count = 0
    while True:
        uu,vv = col((u,v),a)
        if u <= v:
            return count
        elif (u > v) and (uu >= 0):
            count += 1
            return count
        else:
            count += 2
            u = -uu
            v = vv


def col_write(n):
    i = 0
    while i<=n:
        count = cols(inivels,i)
        fp.write(str(i)+' '+str(count)+'\n')
        i += 1

print(u'衝突回数を求める 1を入力')
print(u'指定した質量比までの衝突回数を求めてcol.datに出力 2を入力')
print(u'それ以外で終了します')
sel = int(input(u'入力 '))
while (sel==1 or sel==2):
    if sel == 1:
        n = int(input('M/n = '))
        print(cols(inivels,n))
        sel = int(input(u'入力 '))
    if sel == 2:
        n = int(input('Max M/m = '))
        col_write(n)
        sel = int(input(u'入力 '))