Instantly share code, notes, and snippets.

Embed
What would you like to do?
N-body simulation using Barnes-Hut method and Processing.
import java.awt.geom.Rectangle2D;
import java.awt.geom.Point2D;
public static final int FRAME_X = 1024;
public static final int FRAME_Y = 1024;
public static final int GENERATION_NUM = 4;
public static final int WALL_POWER = 10;
public static final int COFF_COULUMB = 1;
private ArrayList<Body> bodies;
private ArrayList<Edge> edges;
private QuadTree tree;
void setup() {
size(FRAME_X,FRAME_Y);
colorMode(HSB, 360, 256, 256);
smooth();
frameRate(15);
bodies = new ArrayList<Body>();
edges = new ArrayList<Edge>();
Body root = new Body();
bodies.add(root);
build(bodies,edges,root,GENERATION_NUM);
tree = new QuadTree(bodies);
}
void build(ArrayList<Body> bodies,ArrayList<Edge> edges,Body self,int gen_num){
if(--gen_num<=0)return;
int children_num = int(random(1,5));
for(int i=0;i<children_num;i++){
Body child = new Body();
bodies.add(child);
edges.add(new Edge(self,child));
build(bodies,edges,child,gen_num);
}
}
void draw() {
background(360);
tree.reCalc();//rebuild quad tree
tree.calcCoulumb(bodies);//calc acceleration (Coulumb's Power)
for(Edge edge : edges)edge.calcSpring();//calc accelretion (spring power)
for(Body body : bodies){
body.calcWallPower();//calc accelretion (repeling Wall)
body.calcFriction();
}
for(Edge edge : edges)edge.draw();
for(Body body : bodies){
body.proceed();//move bodies
body.draw();
}
tree.drawRectAndCenter();//debug code
}
public class QuadTree{
private Node root;
public QuadTree(ArrayList<Body> bodies){
root = new Node(new Rectangle2D.Float(0,0,FRAME_X,FRAME_Y));
for(Body body : bodies)root.addBody(body);
root.calcCenter();//calc the total weight and the center of gravity
}
public void drawRectAndCenter(){
root.draw();
}
public void reCalc(){
ArrayList<Body> recalc_bodies = new ArrayList<Body>();
root.reCalc(recalc_bodies);
for(Body recalc_body : recalc_bodies)root.addBody(recalc_body);
root.calcCenter();
}
public void calcCoulumb(ArrayList<Body> bodies){
for(Body body : bodies)root.calcCoulumb(body);
}
public class Node{
private final float[] RECT_X = {0,0,0.5,0.5};
private final float[] RECT_Y = {0,0.5,0,0.5};
private final int INTERNAL = 0;
private final int EXTERNAL_NO_BODY = 1;
private final int EXTERNAL_BODY = 2;
private Node[] children = new Node[4];
private Rectangle2D.Float my_rect;
private Rectangle2D.Float[] rects = new Rectangle2D.Float[4];
private Point2D.Float center = new Point2D.Float();
private float weight = 0;
private Body body;
private int mode = EXTERNAL_NO_BODY;
public Node(Rectangle2D.Float rect){
my_rect = rect;
for(int i=0;i<4;i++)
rects[i] = new Rectangle2D.Float(rect.x + rect.width*RECT_X[i],rect.y + rect.height*RECT_Y[i],rect.width/2,rect.height/2);
}
public void addBody(Body body){
switch(mode){
case INTERNAL:
//this is an internal node, and hand over this body to its children.
for(int i=0;i<4;i++)
if(rects[i].contains(body.point))children[i].addBody(body);
return;
case EXTERNAL_NO_BODY:
//this is an external node but do not have a body, so add body to itself.
this.body = body;
mode = EXTERNAL_BODY;
return;
case EXTERNAL_BODY:
//this is an external node and have a body, so hand over this job to its children and become an internal node.
for(int i=0;i<4;i++){
children[i]=new Node(rects[i]);
if(rects[i].contains(this.body.point))children[i].addBody(this.body);
if(rects[i].contains(body.point))children[i].addBody(body);
}
this.body = null;
mode = INTERNAL;
return;
}
}
public void calcCoulumb(Body body){
PVector r = new PVector(body.point.x - this.center.x,body.point.y - this.center.y);
if(r.magSq()==0)return;
switch(mode){
case EXTERNAL_NO_BODY:
return;
case INTERNAL:
if(r.mag()/my_rect.width < 0.5){
//when the distance is not so far, this node hand over the calc to its children.
for(Node child : children)child.calcCoulumb(body);
return;
}
case EXTERNAL_BODY:
if(this.body==body)return;
else{
PVector ac = r.get();
ac.normalize();
ac.mult(COFF_COULUMB*(this.weight * body.weight) / r.magSq());
body.ac.add(ac);
}
return;
}
}
public void calcCenter(){
switch(mode){
case INTERNAL:
float wx=0,wy=0;//the center of gravity
weight = 0;
for(Node child : children){
child.calcCenter();
weight += child.weight;
wx += child.weight * child.center.x;
wy += child.weight * child.center.y;
}
center.setLocation(int(wx/weight),int(wy/weight));
return;
case EXTERNAL_NO_BODY:
weight = 0;
return;
case EXTERNAL_BODY:
weight = this.body.weight;
center = new Point2D.Float(this.body.point.x,this.body.point.y);
return;
}
}
public void reCalc(ArrayList<Body> recalc_bodies){
switch(mode){
case INTERNAL:
for(Node child : children)child.reCalc(recalc_bodies);
for(Node child : children)if(child.mode != EXTERNAL_NO_BODY)return;
children = new Node[4];
mode = EXTERNAL_NO_BODY;
return;
case EXTERNAL_NO_BODY:
return;
case EXTERNAL_BODY:
if(!my_rect.contains(body.point)){
recalc_bodies.add(body);
this.body = null;
mode = EXTERNAL_NO_BODY;
}
return;
}
}
public void draw(){
switch(mode){
case EXTERNAL_NO_BODY:
return;
case INTERNAL:
for(Node child : children)child.draw();
case EXTERNAL_BODY:
noFill();
strokeWeight(2);
stroke(0);
rect(my_rect.x,my_rect.y,my_rect.width,my_rect.height);
ellipse(center.x,center.y,10,10);
return;
}
}
}
}
public class Edge{
public float n_length = 1,k = 0.0001;
public Body src,dst;
public Edge(Body src,Body dst){
this.src = src;
this.dst = dst;
}
public void calcSpring(){
PVector ac = new PVector(dst.point.x - src.point.x,dst.point.y - src.point.y);
float d = ac.mag() - n_length;
ac.normalize();
ac.mult(k*d);
src.ac.add(ac);
ac.mult(-1);
dst.ac.add(ac);
}
public void draw(){
stroke(50);
strokeWeight(2);
line(int(src.point.x),int(src.point.y),int(dst.point.x),int(dst.point.y));
}
}
public class Body{
public int h,s,v,a,radius;
public Point2D.Float point;
public float weight = 30;
public PVector ac = new PVector(0,0);
public PVector vel = new PVector(0,0);
private float k = 0.01;
public Body(){
point = new Point2D.Float((int)random(0,FRAME_X),(int)random(0,FRAME_Y));
h = int(random(0,90));
s = 255;
v = 255;
a = 200;
radius = 20;
}
public void proceed(){
vel.add(ac);
ac.set(0,0);
point.x += vel.x;
point.y += vel.y;
if(point.x<0)
point.x=FRAME_X-10;
else if(point.x>=FRAME_X-1)
point.x=10;
else if(point.y<0)
point.y=FRAME_Y-10;
else if(point.y>=FRAME_Y-1)
point.y=10;
else return;
vel.mult(0.1);
}
public void calcWallPower(){
float[][] vectors = {{0,point.y},{0,-(FRAME_Y-point.y)},{point.x,0},{-(FRAME_X-point.x),0}};
for(float[] vector : vectors){
PVector r = new PVector(vector[0],vector[1]);
PVector ac = r.get();
ac.normalize();
ac.mult(WALL_POWER/r.mag());
this.ac.add(ac);
}
}
public void calcFriction(){
PVector fr = vel.get();
fr.normalize();
fr.mult((-1)*k*vel.magSq());
ac.add(fr);
}
public void draw(){
noStroke();
fill(h,s,v,a);
ellipse(point.x,point.y,radius*2,radius*2);
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment