Geometry Nodes: Add Selection to Attribute Statistics

This adds a bool field selection input to the Attribute Statistics node.
This is useful for running calculations on a subset of the input field data
rather that then whole set.

Differential Revision: https://developer.blender.org/D13520
This commit is contained in:
Johnny Matthews 2021-12-15 08:04:49 -06:00
parent e85d7d5a28
commit 1a833dbdb9
1 changed files with 39 additions and 31 deletions

View File

@ -29,6 +29,7 @@ namespace blender::nodes::node_geo_attribute_statistic_cc {
static void node_declare(NodeDeclarationBuilder &b)
{
b.add_input<decl::Geometry>(N_("Geometry"));
b.add_input<decl::Bool>(N_("Selection")).default_value(true).supports_field().hide_value();
b.add_input<decl::Float>(N_("Attribute")).hide_value().supports_field();
b.add_input<decl::Vector>(N_("Attribute"), "Attribute_001").hide_value().supports_field();
@ -66,7 +67,8 @@ static void node_init(bNodeTree *UNUSED(tree), bNode *node)
static void node_update(bNodeTree *ntree, bNode *node)
{
bNodeSocket *socket_geo = (bNodeSocket *)node->inputs.first;
bNodeSocket *socket_float_attr = socket_geo->next;
bNodeSocket *socket_selection = socket_geo->next;
bNodeSocket *socket_float_attr = socket_selection->next;
bNodeSocket *socket_float3_attr = socket_float_attr->next;
bNodeSocket *socket_float_mean = (bNodeSocket *)node->outputs.first;
@ -148,38 +150,36 @@ static float median_of_sorted_span(const Span<float> data)
static void node_geo_exec(GeoNodeExecParams params)
{
GeometrySet geometry_set = params.get_input<GeometrySet>("Geometry");
const bNode &node = params.node();
const CustomDataType data_type = static_cast<CustomDataType>(node.custom1);
const AttributeDomain domain = static_cast<AttributeDomain>(node.custom2);
int64_t total_size = 0;
Vector<const GeometryComponent *> components = geometry_set.get_components_for_read();
for (const GeometryComponent *component : components) {
if (component->attribute_domain_supported(domain)) {
total_size += component->attribute_domain_size(domain);
}
}
if (total_size == 0) {
params.set_default_remaining_outputs();
return;
}
const Field<bool> selection_field = params.get_input<Field<bool>>("Selection");
switch (data_type) {
case CD_PROP_FLOAT: {
const Field<float> input_field = params.get_input<Field<float>>("Attribute");
Array<float> data = Array<float>(total_size);
int offset = 0;
Vector<float> data;
for (const GeometryComponent *component : components) {
if (component->attribute_domain_supported(domain)) {
GeometryComponentFieldContext field_context{*component, domain};
const int domain_size = component->attribute_domain_size(domain);
fn::FieldEvaluator data_evaluator{field_context, domain_size};
MutableSpan<float> component_result = data.as_mutable_span().slice(offset, domain_size);
data_evaluator.add_with_destination(input_field, component_result);
data_evaluator.add(input_field);
data_evaluator.set_selection(selection_field);
data_evaluator.evaluate();
offset += domain_size;
const VArray<float> &component_data = data_evaluator.get_evaluated<float>(0);
const IndexMask selection = data_evaluator.get_evaluated_selection_as_mask();
const int next_data_index = data.size();
data.resize(next_data_index + selection.size());
MutableSpan<float> selected_data = data.as_mutable_span().slice(next_data_index,
selection.size());
for (const int i : selection.index_range()) {
selected_data[i] = component_data[selection[i]];
}
}
}
@ -200,7 +200,7 @@ static void node_geo_exec(GeoNodeExecParams params)
const bool variance_required = params.output_is_required("Standard Deviation") ||
params.output_is_required("Variance");
if (total_size != 0) {
if (data.size() != 0) {
if (sort_required) {
std::sort(data.begin(), data.end());
median = median_of_sorted_span(data);
@ -211,7 +211,7 @@ static void node_geo_exec(GeoNodeExecParams params)
}
if (sum_required || variance_required) {
sum = compute_sum<float>(data);
mean = sum / total_size;
mean = sum / data.size();
if (variance_required) {
variance = compute_variance(data, mean);
@ -238,18 +238,26 @@ static void node_geo_exec(GeoNodeExecParams params)
}
case CD_PROP_FLOAT3: {
const Field<float3> input_field = params.get_input<Field<float3>>("Attribute_001");
Array<float3> data = Array<float3>(total_size);
int offset = 0;
Vector<float3> data;
for (const GeometryComponent *component : components) {
if (component->attribute_domain_supported(domain)) {
GeometryComponentFieldContext field_context{*component, domain};
const int domain_size = component->attribute_domain_size(domain);
fn::FieldEvaluator data_evaluator{field_context, domain_size};
MutableSpan<float3> component_result = data.as_mutable_span().slice(offset, domain_size);
data_evaluator.add_with_destination(input_field, component_result);
data_evaluator.add(input_field);
data_evaluator.set_selection(selection_field);
data_evaluator.evaluate();
offset += domain_size;
const VArray<float3> &component_data = data_evaluator.get_evaluated<float3>(0);
const IndexMask selection = data_evaluator.get_evaluated_selection_as_mask();
const int next_data_index = data.size();
data.resize(data.size() + selection.size());
MutableSpan<float3> selected_data = data.as_mutable_span().slice(next_data_index,
selection.size());
for (const int i : selection.index_range()) {
selected_data[i] = component_data[selection[i]];
}
}
}
@ -274,9 +282,9 @@ static void node_geo_exec(GeoNodeExecParams params)
Array<float> data_y;
Array<float> data_z;
if (sort_required || variance_required) {
data_x.reinitialize(total_size);
data_y.reinitialize(total_size);
data_z.reinitialize(total_size);
data_x.reinitialize(data.size());
data_y.reinitialize(data.size());
data_z.reinitialize(data.size());
for (const int i : data.index_range()) {
data_x[i] = data[i].x;
data_y[i] = data[i].y;
@ -284,7 +292,7 @@ static void node_geo_exec(GeoNodeExecParams params)
}
}
if (total_size != 0) {
if (data.size() != 0) {
if (sort_required) {
std::sort(data_x.begin(), data_x.end());
std::sort(data_y.begin(), data_y.end());
@ -301,7 +309,7 @@ static void node_geo_exec(GeoNodeExecParams params)
}
if (sum_required || variance_required) {
sum = compute_sum(data.as_span());
mean = sum / total_size;
mean = sum / data.size();
if (variance_required) {
const float x_variance = compute_variance(data_x, mean.x);