barenes-hut.py

一个用python编写的N体模拟问题

#! /usr/bin/python

# A simple Barnes-Hut N-Body problem solver

import sys

from random import *
from math import sqrt
from Tkinter import *
import string

# each body will be a tuple (mass, dx, dy, x, y) where
#   x and y are offsets from (0,0) at lower left and
#   dx and dy are the velocities in the x and y dimensions
BodyList = []

MaxX = 200       # assume meters
MaxY = 200
MaxMass = 1.0e9  # kilograms
TimeStep = 1  # assume timestep in seconds
G = (6.673e-11) # Gravitational constant in Nm^2/kg^2 
Theta = 1.0      # opening angle
StepCount = 500

Scale = 4  # one grid point is this many pixels

global ForceCalculations
ForceCalculations = 0

class TreeNode:
    """ This is the basic quadtree node. It contains the boundary of the square
        represented by this node of the tree. It also has the list of contained
        bodies and the aggregate mass """
    #recursionDepth = 0

    def __init__(self, x1, y1, x2, y2, bodies, canvas):
        #TreeNode.recursionDepth = TreeNode.recursionDepth + 1
        #print 'recursion depth is %d' % (TreeNode.recursionDepth,)
        self.canvas = canvas #temporary
        self.minX = x1
        self.minY = y1
        self.maxX = x2
        self.maxY = y2
        self.edgeLength = x2 - x1     # square areas
        self.containedBodies = []
        self.containedMass = 0.0
        self.addBodies(bodies)
        #print 'Built node from (%d,%d) to (%d,%d) with %d bodies' % (
        #    self.minX, self.minY, self.maxX, self.maxY, len(self.containedBodies))
        self.children = self.buildQuadTree()
        if self.notEmpty():
            self.computeCenterOfMass()
            #comment the next two operations out to avoid drawing the bounding boxes
            canvas.create_rectangle(x1*Scale, y1*Scale,x2*Scale, y2*Scale, outline='red')
            canvas.create_rectangle(self.centerOfMass[0] * Scale, self.centerOfMass[1] * Scale,
                                    self.centerOfMass[0] * Scale + 1,
                                    self.centerOfMass[1] * Scale + 1, outline='green')
        #TreeNode.recursionDepth = TreeNode.recursionDepth - 1
        
    def inBounds(self, x, y):
        return self.minX <= x < self.maxX and self.minY <= y < self.maxY

    def addBodies(self, bodies):
        for b in bodies:
            if self.inBounds(b[3], b[4]):
                self.containedBodies.append(b)
                #self.containedMass = self.containedMass + b[0] # done later

    def centerPoint(self):
        return self.minX + (self.maxX-self.minX) / 2.0, \
               self.minY + (self.maxY-self.minY) / 2.0

    def notEmpty(self):
        """ returns 0 if no bodies are contained, otherwise the length (true) """
        #print "Treenode contains",  len(self.containedBodies), "bodies"
        return len(self.containedBodies)

    def buildQuadTree(self):
        """ Build a quad tree containing the nodes of this tree node in
            four subtrees of equal area """
        if len(self.containedBodies) > 1: # done if only 1 (or 0)
            cx, cy = self.centerPoint()
            childList = [
                TreeNode(self.minX, self.minY, cx, cy, self.containedBodies, canvas), # bottom left
                TreeNode(cx, self.minY, self.maxX, cy, self.containedBodies, canvas), # bottom right
                TreeNode(self.minX, cy, cx, self.maxY, self.containedBodies, canvas), # top left
                TreeNode(cx, cy, self.maxX, self.maxY, self.containedBodies, canvas)] # top right
            return tuple(filter(TreeNode.notEmpty, childList))
        else:
            return ()  # empty set

    def computeCenterOfMass(self):
        if len(self.containedBodies) == 1:  #only one mass, so that's it
            b = self.containedBodies[0]
            self.containedMass = b[0]
            self.centerOfMass = (b[3], b[4])
        else:
            xSum = 0.0
            ySum = 0.0
            for child in self.children:
                m = child.containedMass
                cm = child.centerOfMass
                xSum += m * cm[0]
                ySum += m * cm[1]
                self.containedMass += m
            self.centerOfMass = (xSum / self.containedMass,
                                 ySum / self.containedMass)



def ComputeForcesAndUpdateBodies(list, tree):
    """ Traverse the list and compute the forces on each body using the Barnes-Hut
        model. The tree is a quad tree where each node has center of mass and mass
        already computed
        Returns: an updated list of bodies"""
    newList = []
    for b in list:
        force = ComputeForce(b, tree)   # force is a vector in a tuple
        # velocity(new) = velocity(old) + force * dt / mass
        vx = b[1] + force[0] * TimeStep / b[0]
        vy = b[2] + force[1] * TimeStep / b[0]
        # position(new) = position(old) + velocity(new) * dt
        x = b[3] + vx * TimeStep
        y = b[4] + vy * TimeStep
        # now make a new entry in the new list
        if 0 <= x <= MaxX and 0 <= y <= MaxY: #only if in bounds
            newList.append((b[0], vx, vy, x, y))
    return newList

def ComputeForce(body, tree):
    global ForceCalculations
    position = (body[3], body[4])
    range = ComputeRange(position, tree.centerOfMass)
    if range < 0.2:  # This is a hack to avoid collision issues and infinities
        return (0.0, 0.0)   # same location, so no force (avoid divide by 0)
    # if the range is greater than the grid length/theta or only one element in tree
    if range >= (tree.edgeLength / Theta) or len(tree.containedBodies) == 1:
        # compute the force from the tree's aggregate mass and center of mass
        # f = G * m1 * m2 / r^2, but must separate x and y components
        rcubed = range * range * range
        invariant = G * body[0] * tree.containedMass / rcubed
        fx = invariant * (tree.centerOfMass[0] - position[0])
        fy = invariant * (tree.centerOfMass[1] - position[1])
        ForceCalculations = ForceCalculations + 1
        return (fx, fy)
    else:  # must try the child nodes of tree
        fx = fy = 0.0
        for child in tree.children:
            childForce = ComputeForce(body, child)
            fx += childForce[0]
            fy += childForce[1]
        return (fx, fy)

def ComputeRange(a, b):
    xd = b[0] - a[0]
    yd = b[1] - a[1]
    return sqrt(xd*xd + yd*yd)
    
def PointInList(x, y, list):
    for b in list:
        if x == b[3] and y == b[4]: return 1
    return 0

def GenerateBodies(count):
    for i in range(count):
        x = float(randrange(0, MaxX))
        y = float(randrange(0, MaxY))
        while PointInList(x, y, BodyList):
            x = float(randrange(0, MaxX))
            y = float(randrange(0, MaxY))
        mass = 10000.0 + randrange(0, MaxMass)
        BodyList.append((mass, 0.0, 0.0, x, y))



def DrawBodies(canvas, bodies):
    for b in bodies:
        x = int(b[3] * Scale)
        y = int(b[4] * Scale)
        r = int(b[0]) / (MaxMass / 10)   # not exactly the radius, but it will do
        canvas.create_oval(x-r, y-r, x+r, y+r, fill='white')


if __name__ == "__main__":
    # the command line arguments are: filename step_count
    # Both are optional and have default values/operations, but the step_count
    # can't be specified without the filename being present
    if len(sys.argv) == 1:
        GenerateBodies(20)
    else:
        file = open(sys.argv[1], 'r')   #one body per line
        bodies = file.readlines()
        file.close()
        for b in bodies:
            body = string.split(b)      # split at whitespace
            BodyList.append((float(body[0]), float(body[1]), float(body[2]), float(body[3]), float(body[4])))
        if (len(sys.argv) > 2):         # step_count present
            StepCount = int(sys.argv[2]);

    canvas = Canvas(width = Scale * MaxX + Scale, height = Scale * MaxY + Scale,
                    bg='black')
    canvas.pack();  # show the canvas
    DrawBodies(canvas, BodyList)
    canvas.update()
    for i in range(StepCount):
        root = TreeNode(0, 0, MaxX, MaxY, BodyList, canvas)
        ForceCalculations = 0
        BodyList = ComputeForcesAndUpdateBodies(BodyList, root)
        print 'Iteration %4d - Force calculations: %d' % (i, ForceCalculations) 
        #canvas.delete('all')
        DrawBodies(canvas, BodyList)
        canvas.update()
    for b in BodyList:
        for item in b:                  # each element in the tuple
            print item,
        print                           # newline