diff --git a/scripts/startup/bl_ui/node_add_menu_geometry.py b/scripts/startup/bl_ui/node_add_menu_geometry.py index 7700cd25397..d673f6941b2 100644 --- a/scripts/startup/bl_ui/node_add_menu_geometry.py +++ b/scripts/startup/bl_ui/node_add_menu_geometry.py @@ -220,6 +220,7 @@ class NODE_MT_geometry_node_GEO_GEOMETRY_SAMPLE(Menu): def draw(self, _context): layout = self.layout node_add_menu.add_node_type(layout, "GeometryNodeProximity") + node_add_menu.add_node_type(layout, "GeometryNodeIndexOfNearest") node_add_menu.add_node_type(layout, "GeometryNodeRaycast") node_add_menu.add_node_type(layout, "GeometryNodeSampleIndex") node_add_menu.add_node_type(layout, "GeometryNodeSampleNearest") diff --git a/source/blender/blenkernel/BKE_node.h b/source/blender/blenkernel/BKE_node.h index 03cec560682..8553f2fd0fe 100644 --- a/source/blender/blenkernel/BKE_node.h +++ b/source/blender/blenkernel/BKE_node.h @@ -1581,6 +1581,7 @@ void BKE_nodetree_remove_layer_n(struct bNodeTree *ntree, struct Scene *scene, i #define GEO_NODE_SDF_VOLUME_SPHERE 1196 #define GEO_NODE_MEAN_FILTER_SDF_VOLUME 1197 #define GEO_NODE_OFFSET_SDF_VOLUME 1198 +#define GEO_NODE_INDEX_OF_NEAREST 1199 /** \} */ diff --git a/source/blender/blenlib/BLI_kdtree_impl.h b/source/blender/blenlib/BLI_kdtree_impl.h index 4187724fbda..e42032e6402 100644 --- a/source/blender/blenlib/BLI_kdtree_impl.h +++ b/source/blender/blenlib/BLI_kdtree_impl.h @@ -105,6 +105,23 @@ inline void BLI_kdtree_nd_(range_search_cb_cpp)(const KDTree *tree, }, const_cast(&fn)); } + +template +inline int BLI_kdtree_nd_(find_nearest_cb_cpp)(const KDTree *tree, + const float co[KD_DIMS], + KDTreeNearest *r_nearest, + Fn &&fn) +{ + return BLI_kdtree_nd_(find_nearest_cb)( + tree, + co, + [](void *user_data, const int index, const float *co, const float dist_sq) { + Fn &fn = *static_cast(user_data); + return fn(index, co, dist_sq); + }, + &fn, + r_nearest); +} #endif #undef _BLI_CONCAT_AUX diff --git a/source/blender/nodes/NOD_static_types.h b/source/blender/nodes/NOD_static_types.h index a43f767bda2..1bd9896261e 100644 --- a/source/blender/nodes/NOD_static_types.h +++ b/source/blender/nodes/NOD_static_types.h @@ -324,6 +324,7 @@ DefNode(GeometryNode, GEO_NODE_FLIP_FACES, 0, "FLIP_FACES", FlipFaces, "Flip Fac DefNode(GeometryNode, GEO_NODE_GEOMETRY_TO_INSTANCE, 0, "GEOMETRY_TO_INSTANCE", GeometryToInstance, "Geometry to Instance", "Convert each input geometry into an instance, which can be much faster than the Join Geometry node when the inputs are large") DefNode(GeometryNode, GEO_NODE_IMAGE_INFO, 0, "IMAGE_INFO", ImageInfo, "Image Info", "Retrieve information about an image") DefNode(GeometryNode, GEO_NODE_IMAGE_TEXTURE, def_geo_image_texture, "IMAGE_TEXTURE", ImageTexture, "Image Texture", "Sample values from an image texture") +DefNode(GeometryNode, GEO_NODE_INDEX_OF_NEAREST, 0, "INDEX_OF_NEAREST", IndexOfNearest, "Index of Nearest", "Find the nearest element in the a group. Similar to the \"Sample Nearest\" node") DefNode(GeometryNode, GEO_NODE_IMAGE, def_geo_image, "IMAGE", InputImage, "Image", "Input image") DefNode(GeometryNode, GEO_NODE_INPUT_CURVE_HANDLES, 0, "INPUT_CURVE_HANDLES", InputCurveHandlePositions,"Curve Handle Positions", "Retrieve the position of each Bézier control point's handles") DefNode(GeometryNode, GEO_NODE_INPUT_CURVE_TILT, 0, "INPUT_CURVE_TILT", InputCurveTilt, "Curve Tilt", "Retrieve the angle at each control point used to twist the curve's normal around its tangent") @@ -398,7 +399,7 @@ DefNode(GeometryNode, GEO_NODE_ROTATE_INSTANCES, 0, "ROTATE_INSTANCES", RotateIn DefNode(GeometryNode, GEO_NODE_SAMPLE_CURVE, def_geo_curve_sample, "SAMPLE_CURVE", SampleCurve, "Sample Curve", "Retrieve data from a point on a curve at a certain distance from its start") DefNode(GeometryNode, GEO_NODE_SAMPLE_INDEX, def_geo_sample_index, "SAMPLE_INDEX", SampleIndex, "Sample Index", "Retrieve values from specific geometry elements") DefNode(GeometryNode, GEO_NODE_SAMPLE_NEAREST_SURFACE, def_geo_sample_nearest_surface, "SAMPLE_NEAREST_SURFACE", SampleNearestSurface, "Sample Nearest Surface", "Calculate the interpolated value of a mesh attribute on the closest point of its surface") -DefNode(GeometryNode, GEO_NODE_SAMPLE_NEAREST, def_geo_sample_nearest, "SAMPLE_NEAREST", SampleNearest, "Sample Nearest", "Find the element of a geometry closest to a position") +DefNode(GeometryNode, GEO_NODE_SAMPLE_NEAREST, def_geo_sample_nearest, "SAMPLE_NEAREST", SampleNearest, "Sample Nearest", "Find the element of a geometry closest to a position. Similar to the \"Index of Nearest\" node") DefNode(GeometryNode, GEO_NODE_SAMPLE_UV_SURFACE, def_geo_sample_uv_surface, "SAMPLE_UV_SURFACE", SampleUVSurface, "Sample UV Surface", "Calculate the interpolated values of a mesh attribute at a UV coordinate") DefNode(GeometryNode, GEO_NODE_SCALE_ELEMENTS, def_geo_scale_elements, "SCALE_ELEMENTS", ScaleElements, "Scale Elements", "Scale groups of connected edges and faces") DefNode(GeometryNode, GEO_NODE_SCALE_INSTANCES, 0, "SCALE_INSTANCES", ScaleInstances, "Scale Instances", "Scale geometry instances in local or global space") diff --git a/source/blender/nodes/geometry/CMakeLists.txt b/source/blender/nodes/geometry/CMakeLists.txt index 84515677805..93c7165e839 100644 --- a/source/blender/nodes/geometry/CMakeLists.txt +++ b/source/blender/nodes/geometry/CMakeLists.txt @@ -78,6 +78,7 @@ set(SRC nodes/node_geo_image.cc nodes/node_geo_image_info.cc nodes/node_geo_image_texture.cc + nodes/node_geo_index_of_nearest.cc nodes/node_geo_input_curve_handles.cc nodes/node_geo_input_curve_tilt.cc nodes/node_geo_input_id.cc diff --git a/source/blender/nodes/geometry/node_geometry_register.cc b/source/blender/nodes/geometry/node_geometry_register.cc index 4df52203529..ef78585a132 100644 --- a/source/blender/nodes/geometry/node_geometry_register.cc +++ b/source/blender/nodes/geometry/node_geometry_register.cc @@ -62,6 +62,7 @@ void register_geometry_nodes() register_node_type_geo_image_info(); register_node_type_geo_image_texture(); register_node_type_geo_image(); + register_node_type_geo_index_of_nearest(); register_node_type_geo_input_curve_handles(); register_node_type_geo_input_curve_tilt(); register_node_type_geo_input_id(); diff --git a/source/blender/nodes/geometry/node_geometry_register.hh b/source/blender/nodes/geometry/node_geometry_register.hh index 003515f9d63..d3a77f9bb93 100644 --- a/source/blender/nodes/geometry/node_geometry_register.hh +++ b/source/blender/nodes/geometry/node_geometry_register.hh @@ -59,6 +59,7 @@ void register_node_type_geo_geometry_to_instance(); void register_node_type_geo_image_info(); void register_node_type_geo_image_texture(); void register_node_type_geo_image(); +void register_node_type_geo_index_of_nearest(); void register_node_type_geo_input_curve_handles(); void register_node_type_geo_input_curve_tilt(); void register_node_type_geo_input_id(); diff --git a/source/blender/nodes/geometry/nodes/node_geo_index_of_nearest.cc b/source/blender/nodes/geometry/nodes/node_geo_index_of_nearest.cc new file mode 100644 index 00000000000..c03d8de174b --- /dev/null +++ b/source/blender/nodes/geometry/nodes/node_geo_index_of_nearest.cc @@ -0,0 +1,270 @@ +/* SPDX-License-Identifier: GPL-2.0-or-later */ + +#include "BLI_array.hh" +#include "BLI_kdtree.h" +#include "BLI_map.hh" +#include "BLI_task.hh" + +#include "node_geometry_util.hh" + +namespace blender::nodes::node_geo_index_of_nearest_cc { + +static void node_declare(NodeDeclarationBuilder &b) +{ + b.add_input(N_("Position")).implicit_field(implicit_field_inputs::position); + b.add_input(N_("Group ID")).supports_field().hide_value(); + + b.add_output(N_("Index")).field_source().description(N_("Index of nearest element")); + b.add_output(N_("Has Neighbor")).field_source(); +} + +static KDTree_3d *build_kdtree(const Span positions, const IndexMask mask) +{ + KDTree_3d *tree = BLI_kdtree_3d_new(mask.size()); + for (const int index : mask) { + BLI_kdtree_3d_insert(tree, index, positions[index]); + } + BLI_kdtree_3d_balance(tree); + return tree; +} + +static int find_nearest_non_self(const KDTree_3d &tree, const float3 &position, const int index) +{ + return BLI_kdtree_3d_find_nearest_cb_cpp( + &tree, position, 0, [index](const int other, const float * /*co*/, const float /*dist_sq*/) { + return index == other ? 0 : 1; + }); +} + +static void find_neighbors(const KDTree_3d &tree, + const Span positions, + const IndexMask mask, + MutableSpan r_indices) +{ + threading::parallel_for(mask.index_range(), 512, [&](const IndexRange range) { + for (const int index : mask.slice(range)) { + r_indices[index] = find_nearest_non_self(tree, positions[index], index); + } + }); +} + +class IndexOfNearestFieldInput final : public bke::GeometryFieldInput { + private: + const Field positions_field_; + const Field group_field_; + + public: + IndexOfNearestFieldInput(Field positions_field, Field group_field) + : bke::GeometryFieldInput(CPPType::get(), "Index of Nearest"), + positions_field_(std::move(positions_field)), + group_field_(std::move(group_field)) + { + } + + GVArray get_varray_for_context(const bke::GeometryFieldContext &context, + const IndexMask mask) const final + { + if (!context.attributes()) { + return {}; + } + const int domain_size = context.attributes()->domain_size(context.domain()); + fn::FieldEvaluator evaluator{context, domain_size}; + evaluator.add(positions_field_); + evaluator.add(group_field_); + evaluator.evaluate(); + const VArraySpan positions = evaluator.get_evaluated(0); + const VArray group = evaluator.get_evaluated(1); + + Array result; + + if (group.is_single()) { + result.reinitialize(mask.min_array_size()); + KDTree_3d *tree = build_kdtree(positions, IndexRange(domain_size)); + find_neighbors(*tree, positions, mask, result); + BLI_kdtree_3d_free(tree); + return VArray::ForContainer(std::move(result)); + } + + VectorSet group_indexing; + for (const int index : mask) { + const int group_id = group[index]; + group_indexing.add(group_id); + } + + /* Each group id has two corresponding index masks. One that contains all the points in the + * group, one that contains all the points in the group that should be looked up (this is the + * intersection of the points in the group and `mask`). In many cases, both of these masks are + * the same or very similar, so there is no benefit two separate masks. */ + const bool use_separate_lookup_indices = mask.size() < domain_size / 2; + + Array> all_indices_by_group_id(group_indexing.size()); + Array> lookup_indices_by_group_id; + + if (use_separate_lookup_indices) { + result.reinitialize(mask.min_array_size()); + lookup_indices_by_group_id.reinitialize(group_indexing.size()); + } + else { + result.reinitialize(domain_size); + } + + const auto build_group_masks = [&](const IndexMask mask, + MutableSpan> r_groups) { + for (const int index : mask) { + const int group_id = group[index]; + const int index_of_group = group_indexing.index_of_try(group_id); + if (index_of_group != -1) { + r_groups[index_of_group].append(index); + } + } + }; + + threading::parallel_invoke( + domain_size > 1024 && use_separate_lookup_indices, + [&]() { + if (use_separate_lookup_indices) { + build_group_masks(mask, lookup_indices_by_group_id); + } + }, + [&]() { build_group_masks(IndexMask(domain_size), all_indices_by_group_id); }); + + threading::parallel_for(group_indexing.index_range(), 256, [&](const IndexRange range) { + for (const int index : range) { + const IndexMask tree_mask = all_indices_by_group_id[index].as_span(); + const IndexMask lookup_mask = use_separate_lookup_indices ? + IndexMask(lookup_indices_by_group_id[index]) : + tree_mask; + KDTree_3d *tree = build_kdtree(positions, tree_mask); + find_neighbors(*tree, positions, lookup_mask, result); + BLI_kdtree_3d_free(tree); + } + }); + + return VArray::ForContainer(std::move(result)); + } + + public: + void for_each_field_input_recursive(FunctionRef fn) const + { + positions_field_.node().for_each_field_input_recursive(fn); + group_field_.node().for_each_field_input_recursive(fn); + } + + uint64_t hash() const final + { + return get_default_hash_2(positions_field_, group_field_); + } + + bool is_equal_to(const fn::FieldNode &other) const final + { + if (const auto *other_field = dynamic_cast(&other)) { + return positions_field_ == other_field->positions_field_ && + group_field_ == other_field->group_field_; + } + return false; + } + + std::optional preferred_domain(const GeometryComponent &component) const final + { + return bke::try_detect_field_domain(component, positions_field_); + } +}; + +class HasNeighborFieldInput final : public bke::GeometryFieldInput { + private: + const Field group_field_; + + public: + HasNeighborFieldInput(Field group_field) + : bke::GeometryFieldInput(CPPType::get(), "Has Neighbor"), + group_field_(std::move(group_field)) + { + } + + GVArray get_varray_for_context(const bke::GeometryFieldContext &context, + const IndexMask mask) const final + { + if (!context.attributes()) { + return {}; + } + const int domain_size = context.attributes()->domain_size(context.domain()); + if (domain_size == 1) { + return VArray::ForSingle(false, mask.min_array_size()); + } + + fn::FieldEvaluator evaluator{context, domain_size}; + evaluator.add(group_field_); + evaluator.evaluate(); + const VArray group = evaluator.get_evaluated(0); + + if (group.is_single()) { + return VArray::ForSingle(true, mask.min_array_size()); + } + + Map counts; + const VArraySpan group_span(group); + mask.foreach_index([&](const int i) { + counts.add_or_modify( + group_span[i], [](int *count) { *count = 0; }, [](int *count) { (*count)++; }); + }); + Array result(mask.min_array_size()); + mask.foreach_index([&](const int i) { result[i] = counts.lookup(group_span[i]) > 1; }); + return VArray::ForContainer(std::move(result)); + } + + public: + void for_each_field_input_recursive(FunctionRef fn) const + { + group_field_.node().for_each_field_input_recursive(fn); + } + + uint64_t hash() const final + { + return get_default_hash_2(3984756934876, group_field_); + } + + bool is_equal_to(const fn::FieldNode &other) const final + { + if (const auto *other_field = dynamic_cast(&other)) { + return group_field_ == other_field->group_field_; + } + return false; + } + + std::optional preferred_domain(const GeometryComponent &component) const final + { + return bke::try_detect_field_domain(component, group_field_); + } +}; + +static void node_geo_exec(GeoNodeExecParams params) +{ + Field position_field = params.extract_input>("Position"); + Field group_field = params.extract_input>("Group ID"); + + if (params.output_is_required("Index")) { + params.set_output("Index", + Field(std::make_shared( + std::move(position_field), group_field))); + } + + if (params.output_is_required("Has Neighbor")) { + params.set_output( + "Has Neighbor", + Field(std::make_shared(std::move(group_field)))); + } +} + +} // namespace blender::nodes::node_geo_index_of_nearest_cc + +void register_node_type_geo_index_of_nearest() +{ + namespace file_ns = blender::nodes::node_geo_index_of_nearest_cc; + + static bNodeType ntype; + + geo_node_type_base(&ntype, GEO_NODE_INDEX_OF_NEAREST, "Index of Nearest", NODE_CLASS_CONVERTER); + ntype.geometry_node_execute = file_ns::node_geo_exec; + ntype.declare = file_ns::node_declare; + nodeRegisterType(&ntype); +}