ecsgame/Util/kdtree.h

207 lines
7.9 KiB
C++

#include "geometry.h"
#include <algorithm>
#include <glm/glm.hpp>
#include <iostream>
#include <string>
#include <vector>
class KDTree {
public:
KDTree(std::vector<Point *> points) { root = build(points, 0); }
~KDTree() = default; // TODO: Delete all allocated Nodes
Triangle *intersect_ray(Ray ray, Vector &result) {
return intersect_ray_recurse(result, ray, root, 1000.0, 0);
}
std::string to_string() {
std::string str = "";
to_string_recurse(str, root, 0);
return str;
}
std::vector<Vector> get_lines() {
std::vector<Vector> lines;
float mins[3] = {-1000, -1000, -1000};
float maxs[3] = {1000, 1000, 1000};
get_lines_recursive(lines, root, mins, maxs);
return lines;
}
private:
Node *root;
int MAX_DEPTH = 500;
// Returns a comparator lambda for assessing which of the two points has a
// greater coordinate in the given axis.
auto get_point_comparator(int axis) {
return [axis](Point *p1, Point *p2) {
return p1->pos[axis] < p2->pos[axis];
};
}
Node *build(std::vector<Point *> points, int depth) {
// Exit conditions
if (points.empty() || depth > MAX_DEPTH) { return nullptr; }
// Select axis by choosing the one with maximal extent
float max_extent = 0;
int axis = 0;
for (int it_axis = 0; it_axis < 3; it_axis++) {
// Get extent along this axis
auto comparator = get_point_comparator(it_axis);
Point *min = *std::min_element(points.begin(), points.end(), comparator);
Point *max = *std::max_element(points.begin(), points.end(), comparator);
float extent = max->pos[it_axis] - min->pos[it_axis];
// Is it greater than max_extent?
if (extent > max_extent) {
// If so, make this the splitting axis
max_extent = extent;
axis = it_axis;
}
}
// Choose the median as the pivot and sort the points into
// left-of-median and right-of-median using nth_element
int middle = points.size() / 2;
std::nth_element(points.begin(), points.begin() + middle, points.end(),
get_point_comparator(axis));
Point *median = points[middle];
// TODO: This copies. Can we split the vector into two without copying?
std::vector<Point *> left_of_median(points.begin(), points.begin() + middle);
std::vector<Point *> right_of_median(points.begin() + middle + 1, points.end());
// Create node, recursively call to construct subtree
return new Node(axis, median, build(left_of_median, depth + 1),
build(right_of_median, depth + 1));
}
Triangle *intersect_ray_recurse(Vector &result, Ray ray, Node *node, float max_distance,
int depth) {
// Exit condition: There was no collision
if (node == nullptr) { return nullptr; }
// Is the left or right child node closer to this point?
Node *near =
ray.origin[node->axis] > node->point->pos[node->axis] ? node->right : node->left;
Node *far = near == node->right ? node->left : node->right;
// Check for collisions in this order (stopping if an intersection is found):
// 1. In the nearer section
// 2. With the point in this current node
// 3. In the further section
// If the axes are not parallel, our max_distance decreases, since we've already covered
// some area. `t` represents the distance from this node to the splitting plane.
float t = ray.direction[node->axis] != 0.0
? (node->point->pos[node->axis] - ray.origin[node->axis]) /
ray.direction[node->axis]
: max_distance;
Triangle *near_result = intersect_ray_recurse(result, ray, near, t, depth + 1);
// If the nearer segment had a collision, we're done! We're only interested in the closest
// collision.
if (near_result != nullptr) { return near_result; }
// No collision in the nearer side, so check for a collision directly here
if (ray.intersects_triangle(node->point->triangle, result)) {
// We do have a collision here, so we're done and can return this point!
return node->point->triangle;
}
// No collision here either. Does it make sense to also check the far node?
// Only if the axes are not parallel and if that area is not behind us
if (ray.direction[node->axis] != 0.0 && t >= 0.0) {
// It does make sense to check the far node.
// For this, calculate a new ray origin and continue towards that direction, but with
// the new origin (we can leave behind what we already checked)
return intersect_ray_recurse(result, Ray(ray.origin + ray.direction * t, ray.direction),
far, max_distance - t, depth + 1);
}
// If nothing worked, return a nullptr
return nullptr;
}
void to_string_recurse(std::string &str, Node *node, int depth) {
if (node == nullptr) { return; }
Point *point = node->point;
str += std::string(depth * 2, ' ') + std::to_string(point->pos[0]) + ", " +
std::to_string(point->pos[1]) + ", " + std::to_string(point->pos[2]) +
" with axis " + std::to_string(node->axis) + "\n";
to_string_recurse(str, node->left, depth + 1);
to_string_recurse(str, node->right, depth + 1);
}
void get_lines_recursive(std::vector<Vector> &lines, Node *node, float mins[3], float maxs[3]) {
if (node == nullptr) return;
int axis = node->axis;
// Add lines for this node
// We add a whole bounding box, so we need 8 points -> 12 lines
lines.emplace_back(Vector(mins[0], mins[1], mins[2]));
lines.emplace_back(Vector(maxs[0], mins[1], mins[2]));
lines.emplace_back(Vector(maxs[0], mins[1], mins[2]));
lines.emplace_back(Vector(maxs[0], maxs[1], mins[2]));
lines.emplace_back(Vector(maxs[0], maxs[1], mins[2]));
lines.emplace_back(Vector(mins[0], maxs[1], mins[2]));
lines.emplace_back(Vector(mins[0], maxs[1], mins[2]));
lines.emplace_back(Vector(mins[0], mins[1], mins[2]));
lines.emplace_back(Vector(mins[0], mins[1], maxs[2]));
lines.emplace_back(Vector(maxs[0], mins[1], maxs[2]));
lines.emplace_back(Vector(maxs[0], mins[1], maxs[2]));
lines.emplace_back(Vector(maxs[0], maxs[1], maxs[2]));
lines.emplace_back(Vector(maxs[0], maxs[1], maxs[2]));
lines.emplace_back(Vector(mins[0], maxs[1], maxs[2]));
lines.emplace_back(Vector(mins[0], maxs[1], maxs[2]));
lines.emplace_back(Vector(mins[0], mins[1], maxs[2]));
lines.emplace_back(Vector(mins[0], mins[1], mins[2]));
lines.emplace_back(Vector(mins[0], mins[1], maxs[2]));
lines.emplace_back(Vector(maxs[0], mins[1], mins[2]));
lines.emplace_back(Vector(maxs[0], mins[1], maxs[2]));
lines.emplace_back(Vector(mins[0], maxs[1], mins[2]));
lines.emplace_back(Vector(mins[0], maxs[1], maxs[2]));
lines.emplace_back(Vector(maxs[0], maxs[1], mins[2]));
lines.emplace_back(Vector(maxs[0], maxs[1], maxs[2]));
// TODO: We don't actually have to add all as some are always going to overlap with the
// previous lines...
// Recurse on the left
float new_maxs[3] = {maxs[0], maxs[1], maxs[2]}; // Copy initially
new_maxs[axis] = node->point->pos[axis];
get_lines_recursive(lines, node->left, mins, new_maxs);
// Recurse on the right
float new_mins[3] = {mins[0], mins[1], mins[2]}; // Copy initially
new_mins[axis] = node->point->pos[axis];
get_lines_recursive(lines, node->right, new_mins, maxs);
}
};