Geometry Nodes: improve Point Distribute node

This greatly simplifies the Point Distribute node. For a poisson disk
distribution, it now uses a simpler dart throwing variant. This results
in a slightly lower quality poisson disk distribution, but it still
fulfills our requirements: have a max density, minimum distance input
and stability while painting the density attribute.

This new implementation has a number of benefits over the old one:
* Much less and more readable code.
* Easier to extend with other distribution algorithms.
* Easier to transfer more attributes to the generated points later on.
* More predictable output when changing the max density and min distance.
* Works in 3d, so no projection on the xy plane is necessary.

This is related to T84640.

Differential Revision: https://developer.blender.org/D10104

1 files changed, 175 insertions, 180 deletions

diff --git a/source/blender/nodes/geometry/nodes/node_geo_point_distribute.cc b/source/blender/nodes/geometry/nodes/node_geo_point_distribute.cc
index f5f9c9f8830..1370f45877d 100644
--- a/source/blender/nodes/geometry/nodes/node_geo_point_distribute.cc
+++ b/source/blender/nodes/geometry/nodes/node_geo_point_distribute.cc
@@ -16,9 +16,11 @@
 
 #include "BLI_float3.hh"
 #include "BLI_hash.h"
+#include "BLI_kdtree.h"
 #include "BLI_math_vector.h"
 #include "BLI_rand.hh"
 #include "BLI_span.hh"
+#include "BLI_timeit.hh"
 
 #include "DNA_mesh_types.h"
 #include "DNA_meshdata_types.h"
@@ -34,7 +36,7 @@
 
 static bNodeSocketTemplate geo_node_point_distribute_in[] = {
     {SOCK_GEOMETRY, N_("Geometry")},
-    {SOCK_FLOAT, N_("Distance Min"), 0.1f, 0.0f, 0.0f, 0.0f, 0.0f, 100000.0f, PROP_NONE},
+    {SOCK_FLOAT, N_("Distance Min"), 0.0f, 0.0f, 0.0f, 0.0f, 0.0f, 100000.0f, PROP_NONE},
     {SOCK_FLOAT, N_("Density Max"), 1.0f, 0.0f, 0.0f, 0.0f, 0.0f, 100000.0f, PROP_NONE},
     {SOCK_STRING, N_("Density Attribute")},
     {SOCK_INT, N_("Seed"), 0, 0, 0, 0, -10000, 10000},
@@ -67,213 +69,196 @@ static float3 normal_to_euler_rotation(const float3 normal)
   return rotation;
 }
 
-static Vector<float3> random_scatter_points_from_mesh(const Mesh *mesh,
-                                                      const float density,
-                                                      const FloatReadAttribute &density_factors,
-                                                      Vector<float3> &r_normals,
-                                                      Vector<int> &r_ids,
-                                                      const int seed)
+static Span<MLoopTri> get_mesh_looptris(const Mesh &mesh)
 {
   /* This only updates a cache and can be considered to be logically const. */
-  const MLoopTri *looptris = BKE_mesh_runtime_looptri_ensure(const_cast<Mesh *>(mesh));
-  const int looptris_len = BKE_mesh_runtime_looptri_len(mesh);
+  const MLoopTri *looptris = BKE_mesh_runtime_looptri_ensure(const_cast<Mesh *>(&mesh));
+  const int looptris_len = BKE_mesh_runtime_looptri_len(&mesh);
+  return {looptris, looptris_len};
+}
 
-  Vector<float3> points;
+static void sample_mesh_surface(const Mesh &mesh,
+                                const float base_density,
+                                const FloatReadAttribute *density_factors,
+                                const int seed,
+                                Vector<float3> &r_positions,
+                                Vector<float3> &r_bary_coords,
+                                Vector<int> &r_looptri_indices)
+{
+  Span<MLoopTri> looptris = get_mesh_looptris(mesh);
 
-  for (const int looptri_index : IndexRange(looptris_len)) {
+  for (const int looptri_index : looptris.index_range()) {
     const MLoopTri &looptri = looptris[looptri_index];
-    const int v0_index = mesh->mloop[looptri.tri[0]].v;
-    const int v1_index = mesh->mloop[looptri.tri[1]].v;
-    const int v2_index = mesh->mloop[looptri.tri[2]].v;
-    const float3 v0_pos = mesh->mvert[v0_index].co;
-    const float3 v1_pos = mesh->mvert[v1_index].co;
-    const float3 v2_pos = mesh->mvert[v2_index].co;
-    const float v0_density_factor = std::max(0.0f, density_factors[v0_index]);
-    const float v1_density_factor = std::max(0.0f, density_factors[v1_index]);
-    const float v2_density_factor = std::max(0.0f, density_factors[v2_index]);
-    const float looptri_density_factor = (v0_density_factor + v1_density_factor +
-                                          v2_density_factor) /
-                                         3.0f;
+    const int v0_index = mesh.mloop[looptri.tri[0]].v;
+    const int v1_index = mesh.mloop[looptri.tri[1]].v;
+    const int v2_index = mesh.mloop[looptri.tri[2]].v;
+    const float3 v0_pos = mesh.mvert[v0_index].co;
+    const float3 v1_pos = mesh.mvert[v1_index].co;
+    const float3 v2_pos = mesh.mvert[v2_index].co;
+
+    float looptri_density_factor = 1.0f;
+    if (density_factors != nullptr) {
+      const float v0_density_factor = std::max(0.0f, (*density_factors)[v0_index]);
+      const float v1_density_factor = std::max(0.0f, (*density_factors)[v1_index]);
+      const float v2_density_factor = std::max(0.0f, (*density_factors)[v2_index]);
+      looptri_density_factor = (v0_density_factor + v1_density_factor + v2_density_factor) / 3.0f;
+    }
     const float area = area_tri_v3(v0_pos, v1_pos, v2_pos);
 
     const int looptri_seed = BLI_hash_int(looptri_index + seed);
     RandomNumberGenerator looptri_rng(looptri_seed);
 
-    const float points_amount_fl = area * density * looptri_density_factor;
+    const float points_amount_fl = area * base_density * looptri_density_factor;
     const float add_point_probability = fractf(points_amount_fl);
     const bool add_point = add_point_probability > looptri_rng.get_float();
     const int point_amount = (int)points_amount_fl + (int)add_point;
 
     for (int i = 0; i < point_amount; i++) {
-      const float3 bary_coords = looptri_rng.get_barycentric_coordinates();
+      const float3 bary_coord = looptri_rng.get_barycentric_coordinates();
       float3 point_pos;
-      interp_v3_v3v3v3(point_pos, v0_pos, v1_pos, v2_pos, bary_coords);
-      points.append(point_pos);
-
-      /* Build a hash stable even when the mesh is deformed. */
-      r_ids.append(((int)(bary_coords.hash()) + looptri_index));
-
-      float3 tri_normal;
-      normal_tri_v3(tri_normal, v0_pos, v1_pos, v2_pos);
-      r_normals.append(tri_normal);
+      interp_v3_v3v3v3(point_pos, v0_pos, v1_pos, v2_pos, bary_coord);
+      r_positions.append(point_pos);
+      r_bary_coords.append(bary_coord);
+      r_looptri_indices.append(looptri_index);
     }
   }
+}
 
-  return points;
+BLI_NOINLINE static KDTree_3d *build_kdtree(Span<float3> positions)
+{
+  KDTree_3d *kdtree = BLI_kdtree_3d_new(positions.size());
+  for (const int i : positions.index_range()) {
+    BLI_kdtree_3d_insert(kdtree, i, positions[i]);
+  }
+  BLI_kdtree_3d_balance(kdtree);
+  return kdtree;
 }
 
-struct RayCastAll_Data {
-  void *bvhdata;
+BLI_NOINLINE static void update_elimination_mask_for_close_points(
+    Span<float3> positions, const float minimum_distance, MutableSpan<bool> elimination_mask)
+{
+  if (minimum_distance <= 0.0f) {
+    return;
+  }
 
-  BVHTree_RayCastCallback raycast_callback;
+  KDTree_3d *kdtree = build_kdtree(positions);
 
-  /** The original coordinate the result point was projected from. */
-  float2 raystart;
+  for (const int i : positions.index_range()) {
+    if (elimination_mask[i]) {
+      continue;
+    }
 
-  const Mesh *mesh;
-  float base_weight;
-  FloatReadAttribute *density_factors;
-  Vector<float3> *projected_points;
-  Vector<float3> *normals;
-  Vector<int> *stable_ids;
-  float cur_point_weight;
-};
+    struct CallbackData {
+      int index;
+      MutableSpan<bool> elimination_mask;
+    } callback_data = {i, elimination_mask};
+
+    BLI_kdtree_3d_range_search_cb(
+        kdtree,
+        positions[i],
+        minimum_distance,
+        [](void *user_data, int index, const float *UNUSED(co), float UNUSED(dist_sq)) {
+          CallbackData &callback_data = *static_cast<CallbackData *>(user_data);
+          if (index != callback_data.index) {
+            callback_data.elimination_mask[index] = true;
+          }
+          return true;
+        },
+        &callback_data);
+  }
+  BLI_kdtree_3d_free(kdtree);
+}
 
-static void project_2d_bvh_callback(void *userdata,
-                                    int index,
-                                    const BVHTreeRay *ray,
-                                    BVHTreeRayHit *hit)
+BLI_NOINLINE static void update_elimination_mask_based_on_density_factors(
+    const Mesh &mesh,
+    const FloatReadAttribute &density_factors,
+    Span<float3> bary_coords,
+    Span<int> looptri_indices,
+    MutableSpan<bool> elimination_mask)
 {
-  struct RayCastAll_Data *data = (RayCastAll_Data *)userdata;
-  data->raycast_callback(data->bvhdata, index, ray, hit);
-  if (hit->index != -1) {
-    /* This only updates a cache and can be considered to be logically const. */
-    const MLoopTri *looptris = BKE_mesh_runtime_looptri_ensure(const_cast<Mesh *>(data->mesh));
-    const MVert *mvert = data->mesh->mvert;
+  Span<MLoopTri> looptris = get_mesh_looptris(mesh);
+  for (const int i : bary_coords.index_range()) {
+    if (elimination_mask[i]) {
+      continue;
+    }
 
-    const MLoopTri &looptri = looptris[index];
-    const FloatReadAttribute &density_factors = data->density_factors[0];
+    const MLoopTri &looptri = looptris[looptri_indices[i]];
+    const float3 bary_coord = bary_coords[i];
 
-    const int v0_index = data->mesh->mloop[looptri.tri[0]].v;
-    const int v1_index = data->mesh->mloop[looptri.tri[1]].v;
-    const int v2_index = data->mesh->mloop[looptri.tri[2]].v;
+    const int v0_index = mesh.mloop[looptri.tri[0]].v;
+    const int v1_index = mesh.mloop[looptri.tri[1]].v;
+    const int v2_index = mesh.mloop[looptri.tri[2]].v;
 
     const float v0_density_factor = std::max(0.0f, density_factors[v0_index]);
     const float v1_density_factor = std::max(0.0f, density_factors[v1_index]);
     const float v2_density_factor = std::max(0.0f, density_factors[v2_index]);
 
-    /* Calculate barycentric weights for hit point. */
-    float3 weights;
-    interp_weights_tri_v3(
-        weights, mvert[v0_index].co, mvert[v1_index].co, mvert[v2_index].co, hit->co);
+    const float probablity = v0_density_factor * bary_coord.x + v1_density_factor * bary_coord.y +
+                             v2_density_factor * bary_coord.z;
 
-    float point_weight = weights[0] * v0_density_factor + weights[1] * v1_density_factor +
-                         weights[2] * v2_density_factor;
-
-    point_weight *= data->base_weight;
-
-    if (point_weight >= FLT_EPSILON && data->cur_point_weight <= point_weight) {
-      data->projected_points->append(hit->co);
-
-      /* Build a hash stable even when the mesh is deformed. */
-      data->stable_ids->append((int)data->raystart.hash());
-
-      data->normals->append(hit->no);
+    const float hash = BLI_hash_int_01(bary_coord.hash());
+    if (hash > probablity) {
+      elimination_mask[i] = true;
     }
   }
 }
 
-static Vector<float3> poisson_scatter_points_from_mesh(const Mesh *mesh,
-                                                       const float density,
-                                                       const float minimum_distance,
-                                                       const FloatReadAttribute &density_factors,
-                                                       Vector<float3> &r_normals,
-                                                       Vector<int> &r_ids,
-                                                       const int seed)
+BLI_NOINLINE static void eliminate_points_based_on_mask(Span<bool> elimination_mask,
+                                                        Vector<float3> &positions,
+                                                        Vector<float3> &bary_coords,
+                                                        Vector<int> &looptri_indices)
 {
-  Vector<float3> points;
-
-  if (minimum_distance <= FLT_EPSILON || density <= FLT_EPSILON) {
-    return points;
-  }
-
-  /* Scatter points randomly on the mesh with higher density (5-7) times higher than desired for
-   * good quality possion disk distributions. */
-  int quality = 5;
-  const int output_points_target = 1000;
-  points.resize(output_points_target * quality);
-
-  const float required_area = output_points_target *
-                              (2.0f * sqrtf(3.0f) * minimum_distance * minimum_distance);
-  const float point_scale_multiplier = sqrtf(required_area);
-
-  {
-    const int rnd_seed = BLI_hash_int(seed);
-    RandomNumberGenerator point_rng(rnd_seed);
-
-    for (int i = 0; i < points.size(); i++) {
-      points[i].x = point_rng.get_float() * point_scale_multiplier;
-      points[i].y = point_rng.get_float() * point_scale_multiplier;
-      points[i].z = 0.0f;
+  for (int i = positions.size() - 1; i >= 0; i--) {
+    if (elimination_mask[i]) {
+      positions.remove_and_reorder(i);
+      bary_coords.remove_and_reorder(i);
+      looptri_indices.remove_and_reorder(i);
     }
   }
+}
 
-  /* Eliminate the scattered points until we get a possion distribution. */
-  Vector<float3> output_points(output_points_target);
-
-  const float3 bounds_max = float3(point_scale_multiplier, point_scale_multiplier, 0);
-  poisson_disk_point_elimination(&points, &output_points, 2.0f * minimum_distance, bounds_max);
-  Vector<float3> final_points;
-  r_ids.reserve(output_points_target);
-  final_points.reserve(output_points_target);
-
-  /* Check if we have any points we should remove from the final possion distribition. */
-  BVHTreeFromMesh treedata;
-  BKE_bvhtree_from_mesh_get(&treedata, const_cast<Mesh *>(mesh), BVHTREE_FROM_LOOPTRI, 2);
-
-  float3 bb_min, bb_max;
-  BLI_bvhtree_get_bounding_box(treedata.tree, bb_min, bb_max);
-
-  struct RayCastAll_Data data;
-  data.bvhdata = &treedata;
-  data.raycast_callback = treedata.raycast_callback;
-  data.mesh = mesh;
-  data.projected_points = &final_points;
-  data.stable_ids = &r_ids;
-  data.normals = &r_normals;
-  data.density_factors = const_cast<FloatReadAttribute *>(&density_factors);
-  data.base_weight = std::min(
-      1.0f, density / (output_points.size() / (point_scale_multiplier * point_scale_multiplier)));
-
-  const float max_dist = bb_max[2] - bb_min[2] + 2.0f;
-  const float3 dir = float3(0, 0, -1);
-  float3 raystart;
-  raystart.z = bb_max[2] + 1.0f;
-
-  float tile_start_x_coord = bb_min[0];
-  int tile_repeat_x = ceilf((bb_max[0] - bb_min[0]) / point_scale_multiplier);
-
-  float tile_start_y_coord = bb_min[1];
-  int tile_repeat_y = ceilf((bb_max[1] - bb_min[1]) / point_scale_multiplier);
-
-  for (int x = 0; x < tile_repeat_x; x++) {
-    float tile_curr_x_coord = x * point_scale_multiplier + tile_start_x_coord;
-    for (int y = 0; y < tile_repeat_y; y++) {
-      float tile_curr_y_coord = y * point_scale_multiplier + tile_start_y_coord;
-      for (int idx = 0; idx < output_points.size(); idx++) {
-        raystart.x = output_points[idx].x + tile_curr_x_coord;
-        raystart.y = output_points[idx].y + tile_curr_y_coord;
-
-        data.cur_point_weight = (float)idx / (float)output_points.size();
-        data.raystart = raystart;
-
-        BLI_bvhtree_ray_cast_all(
-            treedata.tree, raystart, dir, 0.0f, max_dist, project_2d_bvh_callback, &data);
-      }
-    }
+BLI_NOINLINE static void compute_remaining_point_data(const Mesh &mesh,
+                                                      Span<float3> bary_coords,
+                                                      Span<int> looptri_indices,
+                                                      MutableSpan<float3> r_normals,
+                                                      MutableSpan<int> r_ids,
+                                                      MutableSpan<float3> r_rotations)
+{
+  Span<MLoopTri> looptris = get_mesh_looptris(mesh);
+  for (const int i : bary_coords.index_range()) {
+    const int looptri_index = looptri_indices[i];
+    const MLoopTri &looptri = looptris[looptri_index];
+    const float3 &bary_coord = bary_coords[i];
+
+    const int v0_index = mesh.mloop[looptri.tri[0]].v;
+    const int v1_index = mesh.mloop[looptri.tri[1]].v;
+    const int v2_index = mesh.mloop[looptri.tri[2]].v;
+    const float3 v0_pos = mesh.mvert[v0_index].co;
+    const float3 v1_pos = mesh.mvert[v1_index].co;
+    const float3 v2_pos = mesh.mvert[v2_index].co;
+
+    r_ids[i] = (int)(bary_coord.hash()) + looptri_index;
+    normal_tri_v3(r_normals[i], v0_pos, v1_pos, v2_pos);
+    r_rotations[i] = normal_to_euler_rotation(r_normals[i]);
   }
+}
 
-  return final_points;
+static void sample_mesh_surface_with_minimum_distance(const Mesh &mesh,
+                                                      const float max_density,
+                                                      const float minimum_distance,
+                                                      const FloatReadAttribute &density_factors,
+                                                      const int seed,
+                                                      Vector<float3> &r_positions,
+                                                      Vector<float3> &r_bary_coords,
+                                                      Vector<int> &r_looptri_indices)
+{
+  sample_mesh_surface(
+      mesh, max_density, nullptr, seed, r_positions, r_bary_coords, r_looptri_indices);
+  Array<bool> elimination_mask(r_positions.size(), false);
+  update_elimination_mask_for_close_points(r_positions, minimum_distance, elimination_mask);
+  update_elimination_mask_based_on_density_factors(
+      mesh, density_factors, r_bary_coords, r_looptri_indices, elimination_mask);
+  eliminate_points_based_on_mask(elimination_mask, r_positions, r_bary_coords, r_looptri_indices);
 }
 
 static void geo_node_point_distribute_exec(GeoNodeExecParams params)
@@ -309,25 +294,37 @@ static void geo_node_point_distribute_exec(GeoNodeExecParams params)
       density_attribute, ATTR_DOMAIN_POINT, 1.0f);
   const int seed = params.get_input<int>("Seed");
 
-  Vector<int> stable_ids;
-  Vector<float3> normals;
-  Vector<float3> points;
+  Vector<float3> positions;
+  Vector<float3> bary_coords;
+  Vector<int> looptri_indices;
   switch (distribute_method) {
     case GEO_NODE_POINT_DISTRIBUTE_RANDOM:
-      points = random_scatter_points_from_mesh(
-          mesh_in, density, density_factors, normals, stable_ids, seed);
+      sample_mesh_surface(
+          *mesh_in, density, &density_factors, seed, positions, bary_coords, looptri_indices);
       break;
     case GEO_NODE_POINT_DISTRIBUTE_POISSON:
-      const float min_dist = params.extract_input<float>("Distance Min");
-      points = poisson_scatter_points_from_mesh(
-          mesh_in, density, min_dist, density_factors, normals, stable_ids, seed);
+      const float minimum_distance = params.extract_input<float>("Distance Min");
+      sample_mesh_surface_with_minimum_distance(*mesh_in,
+                                                density,
+                                                minimum_distance,
+                                                density_factors,
+                                                seed,
+                                                positions,
+                                                bary_coords,
+                                                looptri_indices);
       break;
   }
-
-  PointCloud *pointcloud = BKE_pointcloud_new_nomain(points.size());
-  memcpy(pointcloud->co, points.data(), sizeof(float3) * points.size());
-  for (const int i : points.index_range()) {
-    *(float3 *)(pointcloud->co + i) = points[i];
+  const int tot_points = positions.size();
+  Array<float3> normals(tot_points);
+  Array<int> stable_ids(tot_points);
+  Array<float3> rotations(tot_points);
+  compute_remaining_point_data(
+      *mesh_in, bary_coords, looptri_indices, normals, stable_ids, rotations);
+
+  PointCloud *pointcloud = BKE_pointcloud_new_nomain(tot_points);
+  memcpy(pointcloud->co, positions.data(), sizeof(float3) * tot_points);
+  for (const int i : positions.index_range()) {
+    *(float3 *)(pointcloud->co + i) = positions[i];
     pointcloud->radius[i] = 0.05f;
   }
 
@@ -355,9 +352,7 @@ static void geo_node_point_distribute_exec(GeoNodeExecParams params)
     Float3WriteAttribute rotations_attribute = point_component.attribute_try_ensure_for_write(
         "rotation", ATTR_DOMAIN_POINT, CD_PROP_FLOAT3);
     MutableSpan<float3> rotations_span = rotations_attribute.get_span();
-    for (const int i : rotations_span.index_range()) {
-      rotations_span[i] = normal_to_euler_rotation(normals[i]);
-    }
+    rotations_span.copy_from(rotations);
     rotations_attribute.apply_span();
   }