forked from blender/blender
index-of-nearest-104619 #2
@ -1,7 +1,5 @@
|
|||||||
/* SPDX-License-Identifier: GPL-2.0-or-later */
|
/* SPDX-License-Identifier: GPL-2.0-or-later */
|
||||||
|
|
||||||
#include "BKE_attribute_math.hh"
|
|
||||||
|
|
||||||
#include "BLI_kdtree.h"
|
#include "BLI_kdtree.h"
|
||||||
#include "BLI_multi_value_map.hh"
|
#include "BLI_multi_value_map.hh"
|
||||||
#include "BLI_task.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();
|
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 {
|
class IndexOfNearestFieldInput final : public bke::GeometryFieldInput {
|
||||||
private:
|
private:
|
||||||
const Field<float3> positions_field_;
|
const Field<float3> positions_field_;
|
||||||
@ -26,70 +73,70 @@ class IndexOfNearestFieldInput final : public bke::GeometryFieldInput {
|
|||||||
|
|
||||||
public:
|
public:
|
||||||
IndexOfNearestFieldInput(Field<float3> positions_field, Field<int> group_field)
|
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)),
|
positions_field_(std::move(positions_field)),
|
||||||
group_field_(std::move(group_field))
|
group_field_(std::move(group_field))
|
||||||
{
|
{
|
||||||
}
|
}
|
||||||
|
|
||||||
GVArray get_varray_for_context(const bke::GeometryFieldContext &context,
|
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(positions_field_);
|
||||||
evaluator.add(group_field_);
|
evaluator.add(group_field_);
|
||||||
evaluator.evaluate();
|
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);
|
Array<int> result(mask.min_array_size());
|
||||||
const VArray<int> &group = evaluator.get_evaluated<int>(1);
|
|
||||||
|
|
||||||
MultiValueMap<int, int64_t> group_masks;
|
if (group.is_single()) {
|
||||||
mask.foreach_index([&](const int index) { group_masks.add(group[index], index); });
|
const IndexMask full_mask = positions.index_range();
|
||||||
|
KDTree_3d *tree = build_kdtree(positions, full_mask);
|
||||||
Array<int> indices(mask.min_array_size());
|
find_neighbors(*tree, positions, mask, result);
|
||||||
|
BLI_kdtree_3d_free(tree);
|
||||||
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);
|
|
||||||
}
|
}
|
||||||
|
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:
|
/* The grain size should be larger as each tree gets smaller. */
|
||||||
static int kdtree_find_neighboard(KDTree_3d *tree, const float3 &position, const int &index)
|
const int avg_tree_size = group_ids.size() / group_mask_storage.size();
|
||||||
{
|
const int grain_size = std::max(8192 / avg_tree_size, 1);
|
||||||
return BLI_kdtree_3d_find_nearest_cb_cpp(
|
threading::parallel_for(tree_masks.index_range(), grain_size, [&](const IndexRange range) {
|
||||||
tree,
|
for (const int i : range) {
|
||||||
position,
|
const IndexMask tree_mask = tree_masks[i];
|
||||||
0,
|
const IndexMask evaluate_mask = evaluate_masks.is_empty() ? tree_mask :
|
||||||
[index](const int other_new_i, const float * /*co*/, const float /*dist_sq*/) {
|
evaluate_masks[i];
|
||||||
if (index == other_new_i) {
|
if (tree_masks[i].size() < 2) {
|
||||||
return 0;
|
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:
|
public:
|
||||||
@ -99,53 +146,105 @@ class IndexOfNearestFieldInput final : public bke::GeometryFieldInput {
|
|||||||
group_field_.node().for_each_field_input_recursive(fn);
|
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_);
|
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 =
|
if (const auto *other_field = dynamic_cast<const IndexOfNearestFieldInput *>(&other)) {
|
||||||
dynamic_cast<const IndexOfNearestFieldInput *>(&other)) {
|
|
||||||
return positions_field_ == other_field->positions_field_ &&
|
return positions_field_ == other_field->positions_field_ &&
|
||||||
group_field_ == other_field->group_field_;
|
group_field_ == other_field->group_field_;
|
||||||
}
|
}
|
||||||
return false;
|
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_);
|
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)
|
static void node_geo_exec(GeoNodeExecParams params)
|
||||||
{
|
{
|
||||||
Field<float3> position_field = params.extract_input<Field<float3>>("Position");
|
Field<float3> position_field = params.extract_input<Field<float3>>("Position");
|
||||||
Field<int> group_field = params.extract_input<Field<int>>("Group ID");
|
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")) {
|
if (params.output_is_required("Index")) {
|
||||||
static auto clamp_fn = mf::build::SI1_SO<int, int>(
|
params.set_output("Index",
|
||||||
"Index Clamping",
|
Field<int>(std::make_shared<IndexOfNearestFieldInput>(
|
||||||
[](const int index) { return math::max(0, index); },
|
std::move(position_field), group_field)));
|
||||||
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));
|
|
||||||
}
|
}
|
||||||
|
|
||||||
if (params.output_is_required("Has Neighbor")) {
|
if (params.output_is_required("Has Neighbor")) {
|
||||||
static auto valid_fn = mf::build::SI1_SO<int, bool>(
|
params.set_output(
|
||||||
"Index Validating",
|
"Has Neighbor",
|
||||||
[](const int index) { return index != -1; },
|
Field<bool>(std::make_shared<HasNeighborFieldInput>(std::move(group_field))));
|
||||||
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));
|
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
Loading…
Reference in New Issue
Block a user