AABB tree for atoms

scattering.js

@@ -44,20 +44,37 @@ function veryClose(x, y) {
 }
 
 /*****************************************
- * Parameters
+ * Object Classes
  ****************************************/
 
-let width = 800;
-let height = 600;
+class Rect {
+  constructor(args) {
+    if ('r1' in args) {
+      this.x_0 = Math.min(args.r1.x_0, args.r2.x_0);
+      this.y_0 = Math.min(args.r1.y_0, args.r2.y_0);
+      this.x_1 = Math.max(args.r1.x_1, args.r2.x_1);
+      this.y_1 = Math.max(args.r1.y_1, args.r2.y_1);
+    } else {
+      this.x_0 = args.x_0;
+      this.x_1 = args.x_1;
+      this.y_0 = args.y_0;
+      this.y_1 = args.y_1;
+    }
+  }
+
+  area() {
+    return (this.x_1 - this.x_0)*(this.y_1 - this.y_0);
+  }
+}
+
+let windowRect = new Rect({x_0: 0, x_1: 800, y_0: 0, y_1: 600})
+let width = windowRect.x_1;
+let height = windowRect.y_1;
 
 /* Infinite lines are drawn with length = 2*longestDiag which will always
  * cross the entire screen. */
 let longestDiag = distance(width, height);
 
-/*****************************************
- * Object Classes
- ****************************************/
-
 /* Basic infinite line class, useful for physics simulation */
 class InfiniteLine {
   constructor(x_0, y_0, angle) {
@@ -66,8 +83,7 @@ class InfiniteLine {
     this.angle = angle;
   }
 
-  /* If the line is outside of the canvas */
-  cullable() {
+  insideOfRect(r) {
     // implicit equation of a line:
     // (y - y_0)cos(theta) - (x - x_0)*sin(theta) = 0
     // Evaluate the LHS for two points of a line segment. If the
@@ -79,15 +95,20 @@ class InfiniteLine {
     const upper_left = f(0,0);
     const upper_right = f(width,0);
     if (sign(upper_left) != sign(upper_right))
-      return false;
+      return true;
 
     const lower_left = f(0,height);
     if (sign(lower_left) != sign(upper_left))
-      return false;
+      return true;
 
     const lower_right = f(width,height);
-    return sign(lower_right) == sign(lower_left)
-        && sign(lower_right) == sign(upper_right);
+    return sign(lower_right) != sign(lower_left)
+        || sign(lower_right) != sign(upper_right);
+  }
+
+  /* If the line is outside of the canvas */
+  cullable() {
+    return !this.insideOfRect(windowRect);
   }
 
   // Returns true if drawn.
@@ -95,6 +116,7 @@ class InfiniteLine {
     if (this.cullable())
       return false;
 
+    fill(color(0,0,0));
     translate(this.x_0, this.y_0);
     rotate(this.angle);
     strokeWeight(2);
@@ -116,6 +138,7 @@ class PlaneWave extends InfiniteLine {
   }
 
   drawArrow() {
+    fill(color(0,0,0));
     translate(this.x_0, this.y_0);
     rotate(this.angle);
     triangle(-5, 0, 0, 15, 5, 0);
@@ -128,14 +151,91 @@ class Atom {
   constructor(x,y) {
     this.x = x;
     this.y = y;
+    this.boundRect = new Rect({
+      x_0: x - Atom.radius,
+      x_1: x + Atom.radius,
+      y_0: y - Atom.radius,
+      y_1: y + Atom.radius
+    });
+  }
+
+  updatePos(x,y) {
+    this.x = x;
+    this.y = y;
+    this.boundRect.x_0 = x - Atom.radius;
+    this.boundRect.x_1 = x + Atom.radius;
+    this.boundRect.y_0 = y - Atom.radius;
+    this.boundRect.y_1 = y + Atom.radius;
   }
 
   draw() {
-    color(255,0,0);
+    fill(color(255,0,0));
     circle(this.x, this.y, Atom.radius);
   }
 }
 
+class AABBNode {
+  constructor(args) {
+    this.leftNode = null;
+    this.rightNode = null;
+    if (args !== undefined) {
+      this.obj = args.obj;
+      this.rect = args.rect;
+    } else {
+      this.obj = null;
+      this.rect = null;
+    }
+  }
+
+  add(obj, rect) {
+    // Empty node (the root). Add object.
+    if (this.obj === null && this.leftNode === null && this.rightNode === null) {
+      this.obj = obj;
+      this.rect = rect;
+    } else if (this.obj !== null) {
+      // Leaf node. Split node.
+      this.leftNode = new AABBNode({obj: this.obj, rect: this.rect});
+      this.rightNode = new AABBNode({obj: obj, rect: rect});
+      this.rect = new Rect({r1: this.rect, r2: rect});
+      this.obj = null;
+    } else {
+      // General case with two children.
+      const newLeftRect = new Rect({r1: this.leftNode.rect, r2: rect});
+      const newRightRect = new Rect({r1: this.rightNode.rect, r2: rect});
+
+      if (newLeftRect.area() < newRightRect.area()) {
+        this.leftNode.add(obj, rect);
+      } else {
+        this.rightNode.add(obj, rect);
+      }
+
+      this.rect = new Rect({r1 : this.leftNode.rect, r2 : this.rightNode.rect})
+    }
+  }
+
+  draw(n) {
+    if (this.rect !== null) {
+      fill(color(0, n, 0, 100));
+      quad(this.rect.x_0, this.rect.y_0,
+           this.rect.x_0, this.rect.y_1,
+           this.rect.x_1, this.rect.y_1,
+           this.rect.x_1, this.rect.y_0);
+    }
+
+    if (this.obj !== null) {
+      this.obj.draw();
+    }
+
+    if (this.leftNode) {
+      this.leftNode.draw(n+10);
+    }
+
+    if (this.rightNode) {
+      this.rightNode.draw(n+10);
+    }
+  }
+}
+
 /********************************
  * Mouse handler classes.
  *
@@ -204,7 +304,6 @@ class AtomMouseHandler extends MouseHandler {
     } else {
       const oldAtom = this.currentAtom;
       this.currentAtom = null;
-      console.log('abcd');
       return this.stateTo('base', oldAtom);
     }
   }
@@ -218,8 +317,7 @@ class AtomMouseHandler extends MouseHandler {
   }
 
   handleMouseMove() {
-    this.currentAtom.x = mouseX;
-    this.currentAtom.y = mouseY;
+    this.currentAtom.updatePos(mouseX, mouseY);
     this.currentAtom.draw();
     return null;
   }
@@ -277,7 +375,7 @@ class SimulationMouseHandler extends MouseHandler {
 class Simulation {
   constructor() {
     this.planeWaves = [];
-    this.atoms = [];
+    this.atomTree = new AABBNode();
     this.mouseHandler = new SimulationMouseHandler();
   }
 
@@ -286,7 +384,7 @@ class Simulation {
     if (handleRet instanceof PlaneWave) {
       this.planeWaves.push(handleRet);
     } else if (handleRet instanceof Atom) {
-      this.atoms.push(handleRet);
+      this.atomTree.add(handleRet, handleRet.boundRect);
     }
 
     for (let ind = 0; ind < this.planeWaves.length; ind++) {
@@ -296,9 +394,7 @@ class Simulation {
       }
     }
 
-    for (let ind = 0; ind < this.atoms.length; ind++) {
-      this.atoms[ind].draw();
-    }
+    this.atomTree.draw(0);
   }
 }