Advertisement
JontePonte

Recursive quadtree

Nov 11th, 2024 (edited)
119
0
Never
Not a member of Pastebin yet? Sign Up, it unlocks many cool features!
C# 6.96 KB | None | 0 0
  1. using System.Collections.Generic;
  2. using UnityEngine;
  3. using Vector2 = UnityEngine.Vector2;
  4. using Particle = Simulation.Particle;
  5. public class QuadTree
  6. {
  7.     // Arbitrary constant to indicate how many elements can be stored in this quad
  8.     public static int Capacity;
  9.     //How many times tree can split
  10.     public static float MaxDepth;
  11.     public static float OriginalSize;
  12.     //Handles edge cases
  13.     public static float OverlapTolerance;
  14.     //How much the algorithm should group together particles
  15.     public static float ApproximationThreshold;
  16.    
  17.     //The average center of mass and mass in a node
  18.     Vector2 centerOfMass;
  19.     float mass;
  20.    
  21.     // Axis-aligned bounding box to represent the
  22.     //boundaries of this quad tree
  23.     AABB boundary;
  24.    
  25.     //Objects in this quad tree node
  26.     List<Particle> particles;
  27.     int particleCount;
  28.    
  29.     // Children
  30.     QuadTree northWest;
  31.     QuadTree northEast;
  32.     QuadTree southWest;
  33.     QuadTree southEast;
  34.  
  35.     Particle avgParticle;
  36.    
  37.     //True if the quad subdivides
  38.     bool divided;
  39.  
  40.     public QuadTree(AABB boundary)
  41.     {
  42.         this.boundary = boundary;
  43.         avgParticle = new();
  44.         particles = new();
  45.     }
  46.  
  47.     public void DrawTree()
  48.     {
  49.         DrawSquare(boundary.position, boundary.size);
  50.         Debug.DrawLine(centerOfMass, centerOfMass + new Vector2(0.1f, 0), Color.blue);
  51.        
  52.         if (divided)
  53.         {
  54.             northWest.DrawTree();
  55.             northEast.DrawTree();
  56.             southWest.DrawTree();
  57.             southEast.DrawTree();
  58.         }
  59.     }
  60.     static void DrawSquare(Vector2 position, float size)
  61.     {
  62.         Debug.DrawLine(position, new Vector2(position.x, position.y + size), Color.red);
  63.         Debug.DrawLine(new Vector2(position.x, position.y + size), new Vector2(position.x + size, position.y + size), Color.red);
  64.         Debug.DrawLine(new Vector2(position.x + size, position.y + size), new Vector2(position.x + size, position.y), Color.red);
  65.         Debug.DrawLine(new Vector2(position.x + size, position.y), new Vector2(position.x, position.y), Color.red);
  66.     }
  67.    
  68.     //TODO: Update velocity in Simulation instead
  69.     public Particle UpdateVelocity(Particle particle, float timeStep)
  70.     {
  71.         //If the node hasn't been divided, calculate the force exerted by all particles in this node
  72.         if (!divided)
  73.         {
  74.             for (int i = 0; i < particleCount; i++)
  75.             {
  76.                 if(particles[i].position == particle.position) continue;
  77.                 particle.velocity += Simulation.CalculateAcceleration(particle, particles[i]) * timeStep;
  78.             }
  79.  
  80.             return particle;
  81.         }
  82.         //Otherwise, calculate the ratio s/d. If node is sufficiently far away, treat it as a single particle
  83.         float ratio = boundary.size / (particle.position - centerOfMass).magnitude;
  84.         if (ratio < ApproximationThreshold)
  85.         {
  86.             particle.velocity += Simulation.CalculateAcceleration(particle, avgParticle) * timeStep;
  87.             return particle;
  88.         }
  89.         //Otherwise, check children if the node has been divided
  90.         particle = northWest.UpdateVelocity(particle, timeStep);
  91.         particle = northEast.UpdateVelocity(particle, timeStep);
  92.         particle = southWest.UpdateVelocity(particle, timeStep);
  93.         particle = southEast.UpdateVelocity(particle, timeStep);
  94.         return particle;
  95.     }
  96.  
  97.     public void CalculateMassAndCenterOfMass()
  98.     {
  99.         for (int i = 0; i < particleCount; i++)
  100.         {
  101.             mass += particles[i].mass;
  102.         }
  103.  
  104.         float x = 0;
  105.         float y = 0;
  106.         for (int i = 0; i < particleCount; i++)
  107.         {
  108.             x += particles[i].position.x * particles[i].mass;
  109.             y += particles[i].position.y * particles[i].mass;
  110.         }
  111.  
  112.         x /= mass;
  113.         y /= mass;
  114.         centerOfMass = new Vector2(x, y);
  115.  
  116.         if (divided)
  117.         {
  118.             northWest.CalculateMassAndCenterOfMass();
  119.             northEast.CalculateMassAndCenterOfMass();
  120.             southWest.CalculateMassAndCenterOfMass();
  121.             southEast.CalculateMassAndCenterOfMass();
  122.         }
  123.  
  124.         avgParticle.position = centerOfMass;
  125.         avgParticle.mass = mass;
  126.     }
  127.    
  128.     public void Insert(Particle particle)
  129.     {
  130.         //If the quad has already divided, insert it into on of the children
  131.         if (divided)
  132.         {
  133.             InsertIntoChildren(particle);
  134.         }
  135.         else
  136.         {
  137.             // Ignore objects that do not belong in this quad tree
  138.             if (!boundary.Intersects(particle, OverlapTolerance)) return;
  139.            
  140.             // If there is space in this quad tree and if doesn't have subdivisions, add the object here
  141.             if (particleCount < Capacity)
  142.             {
  143.                 particles.Add(particle);
  144.                 particleCount++;
  145.             }
  146.             // Otherwise, subdivide and then add the point to whichever node will accept it
  147.             else
  148.             {
  149.                 Subdivide(particle);
  150.             }
  151.  
  152.         }
  153.     }
  154.  
  155.     void Subdivide(Particle particle)
  156.     {
  157.         float size = boundary.size / 2;
  158.         //If the node is big enough
  159.         if (OriginalSize / size < MaxDepth * 4)
  160.         {
  161.             //Create 4 children
  162.             northWest = new(new AABB
  163.             {
  164.                 position = new Vector2(boundary.position.x, boundary.position.y + boundary.size / 2f),
  165.                 size = size
  166.             });
  167.             northEast = new(new AABB
  168.             {
  169.                 position = new Vector2(boundary.position.x + boundary.size / 2, boundary.position.y + boundary.size / 2f),
  170.                 size = size
  171.             });
  172.             southWest = new(new AABB
  173.             {
  174.                 position = new Vector2(boundary.position.x, boundary.position.y),
  175.                 size = size
  176.             });
  177.             southEast = new(new AABB
  178.             {  
  179.                 position = new Vector2(boundary.position.x + boundary.size / 2, boundary.position.y),
  180.                 size = size
  181.             });
  182.        
  183.             //Move all objects from this quad into children
  184.             InsertIntoChildren(particle);
  185.             for (int i = 0; i < particleCount; i++)
  186.             {
  187.                 InsertIntoChildren(particles[i]);
  188.             }
  189.  
  190.             divided = true;
  191.         }
  192.     }
  193.  
  194.     void InsertIntoChildren(Particle particle)
  195.     {
  196.         northWest.Insert(particle);
  197.         northEast.Insert(particle);
  198.         southWest.Insert(particle);
  199.         southEast.Insert(particle);
  200.     }
  201. }
  202.  
  203. public struct AABB
  204. {
  205.     //Bottom left of rectangle
  206.     public Vector2 position;
  207.     public float size;
  208.  
  209.     public bool Intersects(Particle p, float margin)
  210.     {
  211.         return p.position.x > position.x - margin && p.position.x < position.x + size + margin &&
  212.                p.position.y > position.y - margin && p.position.y < position.y + size + margin;
  213.     }
  214. }
  215.  
Advertisement
Add Comment
Please, Sign In to add comment
Advertisement