sim...?

dokokaokashii.py

import random
#import pylab
import copy

G = 1.0
H = 3.5
M = 1.0

def ngp(data, n):
    m = [[0] * (n+1) for i in range(n+1)]
    print "length of data in ngp: %d" % len(data)
    for i in data:
        x = int(round(i[0]))
        y = int(round(i[1]))
        m[x][y] += 1
    return m

def montecarlo(num, func):
    count = 0
    res = []
    while count < num:
        r = random.uniform(-5, 5)
        s = random.uniform(-5, 5)
        if func(r, s):
            count += 1
            res.append([r, s])
    return res

def expanding(data, v, dt, f):
    l = len(data)
    res1 = []
    res2 = []
    for i in range(l):
        px = int(round(data[i][0]))
        py = int(round(data[i][1]))
        fpx = M * f[px][py][0]
        fpy = M * f[px][py][1]
        nvx = v[i][0] + (fpx / M) * dt
        nvy = v[i][1] + (fpy / M) * dt
        nx = data[i][0] + nvx * dt
        ny = data[i][1] + nvy * dt
        res1.append([nx, ny])
        res2.append([nvx, nvy])
    return res1, res2

def pde(start, rho):
    mat = copy.deepcopy(start)
    for i in range(50):
        s = 0
        for x in range(1,len(mat[0]) - 1):
            for y in range(1,len(mat) - 1):
                old = mat[x][y]
                new = (mat[x-1][y]+mat[x+1][y]+mat[x][y-1]+mat[x][y+1])
                f = G * rho[x][y]
                mat[x][y] = (new - f)/4.0
                s += abs(mat[x][y] - old)
        if s < 0.00001:
            break
    print 'count = %d' % i
    print 'error = %f' % s
    return mat

def gf(phi, n):
    res = [[0] * (n+1) for i in range(n+1)]
    for i in range(len(res)-1):
        for j in range(len(res)-1):
            fx = -(phi[i+1][j]-phi[i][j])
            fy = -(phi[i][j+1]-phi[i][j])
            res[i][j] = [fx, fy]
    return res

def rm(c, v):
    res1 = []
    res2 = []
    for i in range(len(c)):
        if 0 < round(c[i][0]) < 100 and 0 < round(c[i][1]) < 100:
            res1.append([c[i][0], c[i][1]])
            res2.append([v[i][0], v[i][1]])
    return res1, res2

def graph(data):
    x = []
    y = []
    for i in data:
        x.append(i[0])
        y.append(i[1])
    pylab.scatter(x, y, marker = 'o', c = [0,0,0])
    pylab.xlim([0,100])
    pylab.ylim([0,100])
    pylab.show()

def main():
    mesh = 100
    nk = 80
    dt = 0.1
    phi = [[0] * (mesh+1) for i in range(mesh+1)]
    c = montecarlo(500, lambda x,y: x*x + y*y <= 25)
    v = []
    for i in range(len(c)):
        vx = H * c[i][0]
        vy = H * c[i][1]
        c[i][0] += 50
        c[i][1] += 50
        v.append([vx,vy])
    #graph(c)
    for i in range(nk):
        rho = ngp(c, mesh)
        phi = pde(phi, rho)
        f = gf(phi, mesh)
        c, v = expanding(c, v, dt, f)
        c, v = rm(c, v)
        name = str(i) + '.txt'
        f = open(name, 'w')
        for i in c:
            f.write('%f %f\n' % (i[0], i[1]))
        f.close()
    #graph(c)
    print len(c)

if __name__ == '__main__':
    main()