forked from blender/blender
index-of-nearest-104619 #2
@ -1,7 +1,5 @@
|
||||
/* SPDX-License-Identifier: GPL-2.0-or-later */
|
||||
|
||||
#include "BKE_attribute_math.hh"
|
||||
|
||||
#include "BLI_kdtree.h"
|
||||
#include "BLI_multi_value_map.hh"
|
||||
#include "BLI_task.hh"
|
||||
@ -19,6 +17,55 @@ static void node_declare(NodeDeclarationBuilder &b)
|
||||
b.add_output<decl::Bool>(N_("Has Neighbor")).field_source();
|
||||
}
|
||||
|
||||
static KDTree_3d *build_kdtree(const Span<float3> &positions, const IndexMask mask)
|
||||
{
|
||||
KDTree_3d *tree = BLI_kdtree_3d_new(mask.size());
|
||||
for (const int i : mask) {
|
||||
BLI_kdtree_3d_insert(tree, i, positions[i]);
|
||||
}
|
||||
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<float3> positions,
|
||||
const IndexMask mask,
|
||||
MutableSpan<int> indices)
|
||||
{
|
||||
threading::parallel_for(mask.index_range(), 1024, [&](const IndexRange range) {
|
||||
for (const int i : mask.slice(range)) {
|
||||
indices[i] = find_nearest_non_self(tree, positions[i], i);
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
static Vector<IndexMask> masks_from_group_ids(const Span<int> group_ids,
|
||||
const IndexMask mask,
|
||||
MultiValueMap<int, int64_t> &storage)
|
||||
{
|
||||
mask.foreach_index([&](const int i) { storage.add(group_ids[i], i); });
|
||||
Vector<IndexMask> masks;
|
||||
masks.reserve(storage.size());
|
||||
for (const Span<int64_t> indices : storage.values()) {
|
||||
masks.append(indices);
|
||||
}
|
||||
return masks;
|
||||
}
|
||||
|
||||
static Vector<IndexMask> masks_from_group_ids(const Span<int> group_ids,
|
||||
MultiValueMap<int, int64_t> &storage)
|
||||
{
|
||||
return masks_from_group_ids(group_ids, group_ids.index_range(), storage);
|
||||
}
|
||||
|
||||
class IndexOfNearestFieldInput final : public bke::GeometryFieldInput {
|
||||
private:
|
||||
const Field<float3> positions_field_;
|
||||
@ -26,70 +73,70 @@ class IndexOfNearestFieldInput final : public bke::GeometryFieldInput {
|
||||
|
||||
public:
|
||||
IndexOfNearestFieldInput(Field<float3> positions_field, Field<int> group_field)
|
||||
: bke::GeometryFieldInput(CPPType::get<int>(), "Nearest to"),
|
||||
: bke::GeometryFieldInput(CPPType::get<int>(), "Index of Nearest"),
|
||||
positions_field_(std::move(positions_field)),
|
||||
group_field_(std::move(group_field))
|
||||
{
|
||||
}
|
||||
|
||||
GVArray get_varray_for_context(const bke::GeometryFieldContext &context,
|
||||
IndexMask mask) const final
|
||||
const IndexMask mask) const final
|
||||
{
|
||||
fn::FieldEvaluator evaluator{context, &mask};
|
||||
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<float3> positions = evaluator.get_evaluated<float3>(0);
|
||||
const VArray<int> group = evaluator.get_evaluated<int>(1);
|
||||
|
||||
const VArray<float3> &positions = evaluator.get_evaluated<float3>(0);
|
||||
const VArray<int> &group = evaluator.get_evaluated<int>(1);
|
||||
Array<int> result(mask.min_array_size());
|
||||
|
||||
MultiValueMap<int, int64_t> group_masks;
|
||||
mask.foreach_index([&](const int index) { group_masks.add(group[index], index); });
|
||||
|
||||
Array<int> indices(mask.min_array_size());
|
||||
|
||||
const auto nearest_for = [this, &positions](const IndexMask mask, MutableSpan<int> r_indices) {
|
||||
devirtualize_varray(positions, [mask, r_indices, this](const auto positions) {
|
||||
KDTree_3d *tree = BLI_kdtree_3d_new(mask.size());
|
||||
mask.foreach_index([tree, positions](const int index) {
|
||||
BLI_kdtree_3d_insert(tree, index, positions[index]);
|
||||
});
|
||||
|
||||
BLI_kdtree_3d_balance(tree);
|
||||
|
||||
threading::parallel_for(mask.index_range(), 512, [&](const IndexRange range) {
|
||||
mask.slice(range).foreach_index([&](const auto index) {
|
||||
r_indices[index] = this->kdtree_find_neighboard(tree, positions[index], index);
|
||||
});
|
||||
});
|
||||
|
||||
BLI_kdtree_3d_free(tree);
|
||||
});
|
||||
};
|
||||
|
||||
for (const Span<int64_t> mask_span : group_masks.values()) {
|
||||
if (mask_span.size() == 1) {
|
||||
indices[mask_span.first()] = -1;
|
||||
}
|
||||
nearest_for(mask_span, indices);
|
||||
if (group.is_single()) {
|
||||
const IndexMask full_mask = positions.index_range();
|
||||
KDTree_3d *tree = build_kdtree(positions, full_mask);
|
||||
find_neighbors(*tree, positions, mask, result);
|
||||
BLI_kdtree_3d_free(tree);
|
||||
}
|
||||
else {
|
||||
/* The goal is to build each tree and use it immediately, rather than building all trees and
|
||||
* sampling them later. That should help to keep the tree in caches before balancing and when
|
||||
* sampling many points. */
|
||||
const VArraySpan<int> group_ids(group);
|
||||
MultiValueMap<int, int64_t> group_mask_storage;
|
||||
const Vector<IndexMask> tree_masks = masks_from_group_ids(group_ids, group_mask_storage);
|
||||
|
||||
return VArray<int>::ForContainer(std::move(indices));
|
||||
}
|
||||
MultiValueMap<int, int64_t> evaluate_masks_storage;
|
||||
Vector<IndexMask> evaluate_masks;
|
||||
if (mask.size() < domain_size) {
|
||||
/* Separate masks for evaluation are only necessary if the mask mask
|
||||
* for field input evaluation doesn't have every element selected. */
|
||||
evaluate_masks = masks_from_group_ids(group_ids, mask, evaluate_masks_storage);
|
||||
}
|
||||
|
||||
protected:
|
||||
static int kdtree_find_neighboard(KDTree_3d *tree, const float3 &position, const int &index)
|
||||
{
|
||||
return BLI_kdtree_3d_find_nearest_cb_cpp(
|
||||
tree,
|
||||
position,
|
||||
0,
|
||||
[index](const int other_new_i, const float * /*co*/, const float /*dist_sq*/) {
|
||||
if (index == other_new_i) {
|
||||
return 0;
|
||||
/* The grain size should be larger as each tree gets smaller. */
|
||||
const int avg_tree_size = group_ids.size() / group_mask_storage.size();
|
||||
const int grain_size = std::max(8192 / avg_tree_size, 1);
|
||||
threading::parallel_for(tree_masks.index_range(), grain_size, [&](const IndexRange range) {
|
||||
for (const int i : range) {
|
||||
const IndexMask tree_mask = tree_masks[i];
|
||||
const IndexMask evaluate_mask = evaluate_masks.is_empty() ? tree_mask :
|
||||
evaluate_masks[i];
|
||||
if (tree_masks[i].size() < 2) {
|
||||
result.as_mutable_span().fill_indices(evaluate_mask.indices(), 0);
|
||||
}
|
||||
return 1;
|
||||
});
|
||||
else {
|
||||
KDTree_3d *tree = build_kdtree(positions, tree_mask);
|
||||
find_neighbors(*tree, positions, evaluate_mask, result);
|
||||
BLI_kdtree_3d_free(tree);
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
return VArray<int>::ForContainer(std::move(result));
|
||||
}
|
||||
|
||||
public:
|
||||
@ -99,53 +146,105 @@ class IndexOfNearestFieldInput final : public bke::GeometryFieldInput {
|
||||
group_field_.node().for_each_field_input_recursive(fn);
|
||||
}
|
||||
|
||||
uint64_t hash() const override
|
||||
uint64_t hash() const final
|
||||
{
|
||||
return get_default_hash_2(positions_field_, group_field_);
|
||||
}
|
||||
|
||||
bool is_equal_to(const fn::FieldNode &other) const override
|
||||
bool is_equal_to(const fn::FieldNode &other) const final
|
||||
{
|
||||
if (const IndexOfNearestFieldInput *other_field =
|
||||
dynamic_cast<const IndexOfNearestFieldInput *>(&other)) {
|
||||
if (const auto *other_field = dynamic_cast<const IndexOfNearestFieldInput *>(&other)) {
|
||||
return positions_field_ == other_field->positions_field_ &&
|
||||
group_field_ == other_field->group_field_;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
std::optional<eAttrDomain> preferred_domain(const GeometryComponent &component) const override
|
||||
std::optional<eAttrDomain> 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<int> group_field_;
|
||||
|
||||
public:
|
||||
HasNeighborFieldInput(Field<int> group_field)
|
||||
: bke::GeometryFieldInput(CPPType::get<bool>(), "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());
|
||||
fn::FieldEvaluator evaluator{context, domain_size};
|
||||
evaluator.add(group_field_);
|
||||
evaluator.evaluate();
|
||||
const VArray<int> group = evaluator.get_evaluated<int>(0);
|
||||
|
||||
if (group.is_single()) {
|
||||
return VArray<bool>::ForSingle(true, mask.min_array_size());
|
||||
}
|
||||
|
||||
/* When a group ID is contained in the set, it means there is only one element with that ID. */
|
||||
Map<int, int> counts;
|
||||
const VArraySpan<int> group_span(group);
|
||||
mask.foreach_index([&](const int i) {
|
||||
counts.add_or_modify(
|
||||
group_span[i], [](int *count) { *count = 0; }, [](int *count) { (*count)++; });
|
||||
});
|
||||
Array<bool> result(mask.min_array_size());
|
||||
mask.foreach_index([&](const int i) { result[i] = counts.lookup(group_span[i]) > 1; });
|
||||
return VArray<bool>::ForContainer(std::move(result));
|
||||
}
|
||||
|
||||
public:
|
||||
void for_each_field_input_recursive(FunctionRef<void(const FieldInput &)> 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<const HasNeighborFieldInput *>(&other)) {
|
||||
return group_field_ == other_field->group_field_;
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
std::optional<eAttrDomain> preferred_domain(const GeometryComponent &component) const final
|
||||
{
|
||||
return bke::try_detect_field_domain(component, group_field_);
|
||||
}
|
||||
};
|
||||
|
||||
static void node_geo_exec(GeoNodeExecParams params)
|
||||
{
|
||||
Field<float3> position_field = params.extract_input<Field<float3>>("Position");
|
||||
Field<int> group_field = params.extract_input<Field<int>>("Group ID");
|
||||
|
||||
Field<int> index_of_nearest_field(std::make_shared<IndexOfNearestFieldInput>(
|
||||
std::move(position_field), std::move(group_field)));
|
||||
|
||||
if (params.output_is_required("Index")) {
|
||||
static auto clamp_fn = mf::build::SI1_SO<int, int>(
|
||||
"Index Clamping",
|
||||
[](const int index) { return math::max(0, index); },
|
||||
mf::build::exec_presets::Materialized());
|
||||
auto clamp_op = std::make_shared<FieldOperation>(
|
||||
FieldOperation(std::move(clamp_fn), {index_of_nearest_field}));
|
||||
params.set_output("Index", Field<int>(clamp_op, 0));
|
||||
params.set_output("Index",
|
||||
Field<int>(std::make_shared<IndexOfNearestFieldInput>(
|
||||
std::move(position_field), group_field)));
|
||||
}
|
||||
|
||||
if (params.output_is_required("Has Neighbor")) {
|
||||
static auto valid_fn = mf::build::SI1_SO<int, bool>(
|
||||
"Index Validating",
|
||||
[](const int index) { return index != -1; },
|
||||
mf::build::exec_presets::Materialized());
|
||||
auto valid_op = std::make_shared<FieldOperation>(
|
||||
FieldOperation(std::move(valid_fn), {std::move(index_of_nearest_field)}));
|
||||
params.set_output("Has Neighbor", Field<bool>(valid_op, 0));
|
||||
params.set_output(
|
||||
"Has Neighbor",
|
||||
Field<bool>(std::make_shared<HasNeighborFieldInput>(std::move(group_field))));
|
||||
}
|
||||
}
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user