mod_moder/blender
mod_moder/blender
0
1
Fork 0
forked from blender/blender
Code Pull Requests Activity

index-of-nearest-104619 #2

Merged
Iliya Katushenock merged 62 commits from HooglyBoogly/blender:index-of-nearest-104619 into index_of_nearest 2023-04-20 21:19:53 +02:00
Conversation 0 Commits 62 Files Changed 226 +171 -72
1 changed files with 171 additions and 72 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files
Showing only changes of commit f8d47be900 - Show all commits

245
source/blender/nodes/geometry/nodes/node_geo_index_of_nearest.cc
View File

@ -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);
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);
});
});
Array<int> result(mask.min_array_size());
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);
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);
}
/* 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);
}
else {
KDTree_3d *tree = build_kdtree(positions, tree_mask);
find_neighbors(*tree, positions, evaluate_mask, result);
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);
}
return VArray<int>::ForContainer(std::move(indices));
}
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;
}
return 1;
});
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))));
}
}