diff --git a/scripts/startup/bl_ui/node_add_menu_geometry.py b/scripts/startup/bl_ui/node_add_menu_geometry.py index d615be54da4..714d775befe 100644 --- a/scripts/startup/bl_ui/node_add_menu_geometry.py +++ b/scripts/startup/bl_ui/node_add_menu_geometry.py @@ -677,6 +677,9 @@ class NODE_MT_category_GEO_UTILITIES_FIELD(Menu): node_add_menu.add_node_type(layout, "GeometryNodeAccumulateField") node_add_menu.add_node_type(layout, "GeometryNodeFieldAtIndex") node_add_menu.add_node_type(layout, "GeometryNodeFieldOnDomain") + node_add_menu.add_node_type(layout, "GeometryNodeFieldAverage") + node_add_menu.add_node_type(layout, "GeometryNodeFieldMinAndMax") + node_add_menu.add_node_type(layout, "GeometryNodeFieldVariance") node_add_menu.draw_assets_for_catalog(layout, "Utilities/Field") diff --git a/source/blender/makesrna/intern/rna_nodetree.cc b/source/blender/makesrna/intern/rna_nodetree.cc index 4f560c54206..e4b12a27a48 100644 --- a/source/blender/makesrna/intern/rna_nodetree.cc +++ b/source/blender/makesrna/intern/rna_nodetree.cc @@ -12786,7 +12786,10 @@ static void rna_def_nodes(BlenderRNA *brna) define("GeometryNode", "GeometryNodeExtrudeMesh"); define("GeometryNode", "GeometryNodeFaceOfCorner"); define("GeometryNode", "GeometryNodeFieldAtIndex"); + define("GeometryNode", "GeometryNodeFieldAverage"); + define("GeometryNode", "GeometryNodeFieldMinAndMax"); define("GeometryNode", "GeometryNodeFieldOnDomain"); + define("GeometryNode", "GeometryNodeFieldVariance"); define("GeometryNode", "GeometryNodeFillCurve"); define("GeometryNode", "GeometryNodeFilletCurve"); define("GeometryNode", "GeometryNodeFlipFaces"); diff --git a/source/blender/nodes/geometry/CMakeLists.txt b/source/blender/nodes/geometry/CMakeLists.txt index d37e245b0b5..198fab2f6eb 100644 --- a/source/blender/nodes/geometry/CMakeLists.txt +++ b/source/blender/nodes/geometry/CMakeLists.txt @@ -78,6 +78,9 @@ set(SRC nodes/node_geo_evaluate_closure.cc nodes/node_geo_evaluate_on_domain.cc nodes/node_geo_extrude_mesh.cc + nodes/node_geo_field_average.cc + nodes/node_geo_field_min_and_max.cc + nodes/node_geo_field_variance.cc nodes/node_geo_flip_faces.cc nodes/node_geo_foreach_geometry_element.cc nodes/node_geo_geometry_to_instance.cc diff --git a/source/blender/nodes/geometry/nodes/node_geo_field_average.cc b/source/blender/nodes/geometry/nodes/node_geo_field_average.cc new file mode 100644 index 00000000000..ee6067f4593 --- /dev/null +++ b/source/blender/nodes/geometry/nodes/node_geo_field_average.cc @@ -0,0 +1,332 @@ +/* SPDX-FileCopyrightText: 2025 Blender Authors + * + * SPDX-License-Identifier: GPL-2.0-or-later */ + +#include "BKE_attribute_math.hh" + +#include "BLI_array.hh" +#include "BLI_generic_virtual_array.hh" +#include "BLI_math_vector.h" +#include "BLI_vector.hh" +#include "BLI_virtual_array.hh" + +#include "NOD_rna_define.hh" +#include "NOD_socket_search_link.hh" + +#include "RNA_enum_types.hh" + +#include "node_geometry_util.hh" + +#include "UI_interface.hh" +#include "UI_resources.hh" + +#include + +namespace blender::nodes::node_geo_field_average_cc { + +static void node_declare(NodeDeclarationBuilder &b) +{ + const bNode *node = b.node_or_null(); + + if (node != nullptr) { + const eCustomDataType data_type = eCustomDataType(node->custom1); + b.add_input(data_type, "Value") + .supports_field() + .description("The values the mean and median will be calculated from"); + } + + b.add_input("Group ID", "Group Index") + .supports_field() + .hide_value() + .description("An index used to group values together for multiple separate operations"); + + if (node != nullptr) { + const eCustomDataType data_type = eCustomDataType(node->custom1); + b.add_output(data_type, "Mean") + .field_source_reference_all() + .description("The sum of all values in each group divided by the size of said group"); + b.add_output(data_type, "Median") + .field_source_reference_all() + .description( + "The middle value in each group when all values are sorted from lowest to highest"); + } +} + +static void node_layout(uiLayout *layout, bContext * /*C*/, PointerRNA *ptr) +{ + uiItemR(layout, ptr, "data_type", UI_ITEM_NONE, "", ICON_NONE); + uiItemR(layout, ptr, "domain", UI_ITEM_NONE, "", ICON_NONE); +} + +static void node_init(bNodeTree * /*tree*/, bNode *node) +{ + node->custom1 = CD_PROP_FLOAT; + node->custom2 = int16_t(AttrDomain::Point); +} + +enum class Operation { Mean = 0, Median = 1 }; + +static std::optional node_type_from_other_socket(const bNodeSocket &socket) +{ + switch (socket.type) { + case SOCK_FLOAT: + case SOCK_BOOLEAN: + case SOCK_INT: + return CD_PROP_FLOAT; + case SOCK_VECTOR: + case SOCK_RGBA: + return CD_PROP_FLOAT3; + default: + return std::nullopt; + } +} + +static void node_gather_link_searches(GatherLinkSearchOpParams ¶ms) +{ + const NodeDeclaration &declaration = *params.node_type().static_declaration; + search_link_ops_for_declarations(params, declaration.inputs); + + const std::optional type = node_type_from_other_socket(params.other_socket()); + if (!type) { + return; + } + if (params.in_out() == SOCK_OUT) { + params.add_item( + IFACE_("Mean"), + [type](LinkSearchOpParams ¶ms) { + bNode &node = params.add_node("GeometryNodeFieldAverage"); + node.custom1 = *type; + params.update_and_connect_available_socket(node, "Mean"); + }, + 0); + params.add_item( + IFACE_("Median"), + [type](LinkSearchOpParams ¶ms) { + bNode &node = params.add_node("GeometryNodeFieldAverage"); + node.custom1 = *type; + params.update_and_connect_available_socket(node, "Median"); + }, + -1); + } + else { + params.add_item( + IFACE_("Value"), + [type](LinkSearchOpParams ¶ms) { + bNode &node = params.add_node("GeometryNodeFieldAverage"); + node.custom1 = *type; + params.update_and_connect_available_socket(node, "Value"); + }, + 0); + } +} + +template static T calculate_median(MutableSpan values) +{ + if constexpr (std::is_same::value) { + Array x_vals(values.size()), y_vals(values.size()), z_vals(values.size()); + + for (const int i : values.index_range()) { + float3 value = values[i]; + x_vals[i] = value.x; + y_vals[i] = value.y; + z_vals[i] = value.z; + } + + return float3(calculate_median(x_vals), + calculate_median(y_vals), + calculate_median(z_vals)); + } + else { + const auto middle_itr = values.begin() + values.size() / 2; + std::nth_element(values.begin(), middle_itr, values.end()); + if (values.size() % 2 == 0) { + const auto left_middle_itr = std::max_element(values.begin(), middle_itr); + return math::midpoint(*left_middle_itr, *middle_itr); + } + return *middle_itr; + } +} + +class FieldAverageInput final : public bke::GeometryFieldInput { + private: + GField input_; + Field group_index_; + AttrDomain source_domain_; + Operation operation_; + + public: + FieldAverageInput(const AttrDomain source_domain, + GField input, + Field group_index, + Operation operation) + : bke::GeometryFieldInput(input.cpp_type(), "Calculation"), + input_(std::move(input)), + group_index_(std::move(group_index)), + source_domain_(source_domain), + operation_(operation) + { + } + + GVArray get_varray_for_context(const bke::GeometryFieldContext &context, + const IndexMask & /*mask*/) const final + { + const AttributeAccessor attributes = *context.attributes(); + const int64_t domain_size = attributes.domain_size(source_domain_); + if (domain_size == 0) { + return {}; + } + + const bke::GeometryFieldContext source_context{context, source_domain_}; + fn::FieldEvaluator evaluator{source_context, domain_size}; + evaluator.add(input_); + evaluator.add(group_index_); + evaluator.evaluate(); + const GVArray g_values = evaluator.get_evaluated(0); + const VArray group_indices = evaluator.get_evaluated(1); + + GVArray g_outputs; + + bke::attribute_math::convert_to_static_type(g_values.type(), [&](auto dummy) { + using T = decltype(dummy); + if constexpr (is_same_any_v) { + const VArraySpan values = g_values.typed(); + + if (operation_ == Operation::Mean) { + if (group_indices.is_single()) { + const T mean = std::reduce(values.begin(), values.end(), T()) / domain_size; + g_outputs = VArray::ForSingle(mean, domain_size); + } + else { + Map> sum_and_counts; + for (const int i : values.index_range()) { + auto &pair = sum_and_counts.lookup_or_add(group_indices[i], std::make_pair(T(), 0)); + pair.first = pair.first + values[i]; + pair.second = pair.second + 1; + } + + Array outputs(domain_size); + for (const int i : values.index_range()) { + const auto &pair = sum_and_counts.lookup(group_indices[i]); + outputs[i] = pair.first / pair.second; + } + g_outputs = VArray::ForContainer(std::move(outputs)); + } + } + else { + if (group_indices.is_single()) { + Array sorted_values(values); + T median = calculate_median(sorted_values); + g_outputs = VArray::ForSingle(median, domain_size); + } + else { + Map> groups; + for (const int i : values.index_range()) { + groups.lookup_or_add(group_indices[i], Vector()).append(values[i]); + } + + Map medians; + for (MutableMapItem> group : groups.items()) { + medians.add(group.key, calculate_median(group.value)); + } + + Array outputs(domain_size); + for (const int i : values.index_range()) { + outputs[i] = medians.lookup(group_indices[i]); + } + g_outputs = VArray::ForContainer(std::move(outputs)); + } + } + } + }); + + return attributes.adapt_domain(std::move(g_outputs), source_domain_, context.domain()); + } + + uint64_t hash() const override + { + return get_default_hash(input_, group_index_, source_domain_, operation_); + } + + bool is_equal_to(const fn::FieldNode &other) const override + { + if (const FieldAverageInput *other_field = dynamic_cast(&other)) { + return input_ == other_field->input_ && group_index_ == other_field->group_index_ && + source_domain_ == other_field->source_domain_ && + operation_ == other_field->operation_; + } + return false; + } + + std::optional preferred_domain( + const GeometryComponent & /*component*/) const override + { + return source_domain_; + } +}; + +static void node_geo_exec(GeoNodeExecParams params) +{ + const AttrDomain source_domain = AttrDomain(params.node().custom2); + + const Field group_index_field = params.extract_input>("Group Index"); + const GField input_field = params.extract_input("Value"); + if (params.output_is_required("Mean")) { + params.set_output( + "Mean", + GField{std::make_shared( + source_domain, input_field, group_index_field, Operation::Mean)}); + } + if (params.output_is_required("Median")) { + params.set_output( + "Median", + GField{std::make_shared( + source_domain, input_field, group_index_field, Operation::Median)}); + } +} + +static void node_rna(StructRNA *srna) +{ + static EnumPropertyItem items[] = { + {CD_PROP_FLOAT, "FLOAT", 0, "Float", "Floating-point value"}, + {CD_PROP_FLOAT3, "FLOAT_VECTOR", 0, "Vector", "3D vector with floating-point values"}, + {0, nullptr, 0, nullptr, nullptr}, + }; + + RNA_def_node_enum(srna, + "data_type", + "Data Type", + "Type of data the outputs are calculated from", + items, + NOD_inline_enum_accessors(custom1), + CD_PROP_FLOAT); + + RNA_def_node_enum(srna, + "domain", + "Domain", + "", + rna_enum_attribute_domain_items, + NOD_inline_enum_accessors(custom2), + int(AttrDomain::Point), + nullptr, + true); +} + +static void node_register() +{ + static blender::bke::bNodeType ntype; + + geo_node_type_base(&ntype, "GeometryNodeFieldAverage"); + ntype.ui_name = "Field Average"; + ntype.ui_description = "Calculate the mean and median of a given field"; + ntype.nclass = NODE_CLASS_CONVERTER; + ntype.geometry_node_execute = node_geo_exec; + ntype.initfunc = node_init; + ntype.draw_buttons = node_layout; + ntype.declare = node_declare; + ntype.gather_link_search_ops = node_gather_link_searches; + blender::bke::node_register_type(ntype); + node_rna(ntype.rna_ext.srna); +} +NOD_REGISTER_NODE(node_register) + +} // namespace blender::nodes::node_geo_field_average_cc diff --git a/source/blender/nodes/geometry/nodes/node_geo_field_min_and_max.cc b/source/blender/nodes/geometry/nodes/node_geo_field_min_and_max.cc new file mode 100644 index 00000000000..7ff2895a3e4 --- /dev/null +++ b/source/blender/nodes/geometry/nodes/node_geo_field_min_and_max.cc @@ -0,0 +1,318 @@ +/* SPDX-FileCopyrightText: 2025 Blender Authors + * + * SPDX-License-Identifier: GPL-2.0-or-later */ + +#include "BKE_attribute_math.hh" + +#include "BLI_array.hh" +#include "BLI_generic_virtual_array.hh" +#include "BLI_virtual_array.hh" + +#include "NOD_rna_define.hh" +#include "NOD_socket_search_link.hh" + +#include "RNA_enum_types.hh" + +#include "node_geometry_util.hh" + +#include "UI_interface.hh" +#include "UI_resources.hh" + +namespace blender::nodes::node_geo_field_min_and_max_cc { + +static void node_declare(NodeDeclarationBuilder &b) +{ + const bNode *node = b.node_or_null(); + + if (node != nullptr) { + const eCustomDataType data_type = eCustomDataType(node->custom1); + b.add_input(data_type, "Value") + .supports_field() + .description("The values the mininum and maximum will be calculated from"); + } + + b.add_input("Group ID", "Group Index") + .supports_field() + .hide_value() + .description("An index used to group values together for multiple separate operations"); + + if (node != nullptr) { + const eCustomDataType data_type = eCustomDataType(node->custom1); + b.add_output(data_type, "Min") + .field_source_reference_all() + .description("The lowest value in each group"); + b.add_output(data_type, "Max") + .field_source_reference_all() + .description("The highest value in each group"); + } +} + +static void node_layout(uiLayout *layout, bContext * /*C*/, PointerRNA *ptr) +{ + uiItemR(layout, ptr, "data_type", UI_ITEM_NONE, "", ICON_NONE); + uiItemR(layout, ptr, "domain", UI_ITEM_NONE, "", ICON_NONE); +} + +static void node_init(bNodeTree * /*tree*/, bNode *node) +{ + node->custom1 = CD_PROP_FLOAT; + node->custom2 = int16_t(AttrDomain::Point); +} + +enum class Operation { Min = 0, Max = 1 }; + +static std::optional node_type_from_other_socket(const bNodeSocket &socket) +{ + switch (socket.type) { + case SOCK_FLOAT: + return CD_PROP_FLOAT; + case SOCK_BOOLEAN: + case SOCK_INT: + return CD_PROP_INT32; + case SOCK_VECTOR: + case SOCK_RGBA: + case SOCK_ROTATION: + return CD_PROP_FLOAT3; + default: + return std::nullopt; + } +} + +static void node_gather_link_searches(GatherLinkSearchOpParams ¶ms) +{ + const NodeDeclaration &declaration = *params.node_type().static_declaration; + search_link_ops_for_declarations(params, declaration.inputs); + + const std::optional type = node_type_from_other_socket(params.other_socket()); + if (!type) { + return; + } + if (params.in_out() == SOCK_OUT) { + params.add_item( + IFACE_("Min"), + [type](LinkSearchOpParams ¶ms) { + bNode &node = params.add_node("GeometryNodeFieldMinAndMax"); + node.custom1 = *type; + params.update_and_connect_available_socket(node, "Min"); + }, + 0); + params.add_item( + IFACE_("Max"), + [type](LinkSearchOpParams ¶ms) { + bNode &node = params.add_node("GeometryNodeFieldMinAndMax"); + node.custom1 = *type; + params.update_and_connect_available_socket(node, "Max"); + }, + -1); + } + else { + params.add_item( + IFACE_("Value"), + [type](LinkSearchOpParams ¶ms) { + bNode &node = params.add_node("GeometryNodeFieldMinAndMax"); + node.custom1 = *type; + params.update_and_connect_available_socket(node, "Value"); + }, + 0); + } +} + +template struct MinMaxInfo { + static inline const T min_initial_value = []() { + if constexpr (std::is_same_v) { + return float3(std::numeric_limits::max()); + } + else { + return std::numeric_limits::max(); + } + }(); + + static inline const T max_initial_value = []() { + if constexpr (std::is_same_v) { + return float3(std::numeric_limits::lowest()); + } + else { + return std::numeric_limits::lowest(); + } + }(); +}; + +class FieldMinMaxInput final : public bke::GeometryFieldInput { + private: + GField input_; + Field group_index_; + AttrDomain source_domain_; + Operation operation_; + + public: + FieldMinMaxInput(const AttrDomain source_domain, + GField input, + Field group_index, + Operation operation) + : bke::GeometryFieldInput(input.cpp_type(), "Calculation"), + input_(std::move(input)), + group_index_(std::move(group_index)), + source_domain_(source_domain), + operation_(operation) + { + } + + GVArray get_varray_for_context(const bke::GeometryFieldContext &context, + const IndexMask & /*mask*/) const final + { + const AttributeAccessor attributes = *context.attributes(); + const int64_t domain_size = attributes.domain_size(source_domain_); + if (domain_size == 0) { + return {}; + } + + const bke::GeometryFieldContext source_context{context, source_domain_}; + fn::FieldEvaluator evaluator{source_context, domain_size}; + evaluator.add(input_); + evaluator.add(group_index_); + evaluator.evaluate(); + const GVArray g_values = evaluator.get_evaluated(0); + const VArray group_indices = evaluator.get_evaluated(1); + + GVArray g_outputs; + + bke::attribute_math::convert_to_static_type(g_values.type(), [&](auto dummy) { + using T = decltype(dummy); + if constexpr (is_same_any_v) { + const VArray values = g_values.typed(); + + if (operation_ == Operation::Min) { + if (group_indices.is_single()) { + T result = MinMaxInfo::min_initial_value; + for (const int i : values.index_range()) { + result = math::min(result, values[i]); + } + g_outputs = VArray::ForSingle(result, domain_size); + } + else { + Map results; + for (const int i : values.index_range()) { + T &value = results.lookup_or_add(group_indices[i], MinMaxInfo::min_initial_value); + value = math::min(value, values[i]); + } + Array outputs(domain_size); + for (const int i : values.index_range()) { + outputs[i] = results.lookup(group_indices[i]); + } + g_outputs = VArray::ForContainer(std::move(outputs)); + } + } + else { + if (group_indices.is_single()) { + T result = MinMaxInfo::max_initial_value; + for (const int i : values.index_range()) { + result = math::max(result, values[i]); + } + g_outputs = VArray::ForSingle(result, domain_size); + } + else { + Map results; + for (const int i : values.index_range()) { + T &value = results.lookup_or_add(group_indices[i], MinMaxInfo::max_initial_value); + value = math::max(value, values[i]); + } + Array outputs(domain_size); + for (const int i : values.index_range()) { + outputs[i] = results.lookup(group_indices[i]); + } + g_outputs = VArray::ForContainer(std::move(outputs)); + } + } + } + }); + + return attributes.adapt_domain(std::move(g_outputs), source_domain_, context.domain()); + } + + uint64_t hash() const override + { + return get_default_hash(input_, group_index_, source_domain_, operation_); + } + + bool is_equal_to(const fn::FieldNode &other) const override + { + if (const FieldMinMaxInput *other_field = dynamic_cast(&other)) { + return input_ == other_field->input_ && group_index_ == other_field->group_index_ && + source_domain_ == other_field->source_domain_ && + operation_ == other_field->operation_; + } + return false; + } + + std::optional preferred_domain( + const GeometryComponent & /*component*/) const override + { + return source_domain_; + } +}; + +static void node_geo_exec(GeoNodeExecParams params) +{ + const AttrDomain source_domain = AttrDomain(params.node().custom2); + + const Field group_index_field = params.extract_input>("Group Index"); + const GField input_field = params.extract_input("Value"); + if (params.output_is_required("Min")) { + params.set_output("Min", + GField{std::make_shared( + source_domain, input_field, group_index_field, Operation::Min)}); + } + if (params.output_is_required("Max")) { + params.set_output("Max", + GField{std::make_shared( + source_domain, input_field, group_index_field, Operation::Max)}); + } +} + +static void node_rna(StructRNA *srna) +{ + static EnumPropertyItem items[] = { + {CD_PROP_FLOAT, "FLOAT", 0, "Float", "Floating-point value"}, + {CD_PROP_INT32, "INT", 0, "Integer", "32-bit integer"}, + {CD_PROP_FLOAT3, "FLOAT_VECTOR", 0, "Vector", "3D vector with floating-point values"}, + {0, nullptr, 0, nullptr, nullptr}, + }; + + RNA_def_node_enum(srna, + "data_type", + "Data Type", + "Type of data the outputs are calculated from", + items, + NOD_inline_enum_accessors(custom1), + CD_PROP_FLOAT); + + RNA_def_node_enum(srna, + "domain", + "Domain", + "", + rna_enum_attribute_domain_items, + NOD_inline_enum_accessors(custom2), + int(AttrDomain::Point), + nullptr, + true); +} + +static void node_register() +{ + static blender::bke::bNodeType ntype; + + geo_node_type_base(&ntype, "GeometryNodeFieldMinAndMax"); + ntype.ui_name = "Field Min & Max"; + ntype.ui_description = "Calculate the minimum and maximum of a given field"; + ntype.nclass = NODE_CLASS_CONVERTER; + ntype.geometry_node_execute = node_geo_exec; + ntype.initfunc = node_init; + ntype.draw_buttons = node_layout; + ntype.declare = node_declare; + ntype.gather_link_search_ops = node_gather_link_searches; + blender::bke::node_register_type(ntype); + node_rna(ntype.rna_ext.srna); +} +NOD_REGISTER_NODE(node_register) + +} // namespace blender::nodes::node_geo_field_min_and_max_cc diff --git a/source/blender/nodes/geometry/nodes/node_geo_field_variance.cc b/source/blender/nodes/geometry/nodes/node_geo_field_variance.cc new file mode 100644 index 00000000000..a7683b01451 --- /dev/null +++ b/source/blender/nodes/geometry/nodes/node_geo_field_variance.cc @@ -0,0 +1,334 @@ +/* SPDX-FileCopyrightText: 2025 Blender Authors + * + * SPDX-License-Identifier: GPL-2.0-or-later */ + +#include "BKE_attribute_math.hh" + +#include "BLI_array.hh" +#include "BLI_generic_virtual_array.hh" +#include "BLI_virtual_array.hh" + +#include "NOD_rna_define.hh" +#include "NOD_socket_search_link.hh" + +#include "RNA_enum_types.hh" + +#include "node_geometry_util.hh" + +#include "UI_interface.hh" +#include "UI_resources.hh" + +#include + +namespace blender::nodes::node_geo_field_variance_cc { + +static void node_declare(NodeDeclarationBuilder &b) +{ + const bNode *node = b.node_or_null(); + + if (node != nullptr) { + const eCustomDataType data_type = eCustomDataType(node->custom1); + b.add_input(data_type, "Value") + .supports_field() + .description("The values the standard deviation and variance will be calculated from"); + } + + b.add_input("Group ID", "Group Index") + .supports_field() + .hide_value() + .description("An index used to group values together for multiple separate operations"); + + if (node != nullptr) { + const eCustomDataType data_type = eCustomDataType(node->custom1); + b.add_output(data_type, "Standard Deviation") + .field_source_reference_all() + .description("The square root of the variance for each group"); + b.add_output(data_type, "Variance") + .field_source_reference_all() + .description("The expected squared deviation from the mean for each group"); + } +} + +static void node_layout(uiLayout *layout, bContext * /*C*/, PointerRNA *ptr) +{ + uiItemR(layout, ptr, "data_type", UI_ITEM_NONE, "", ICON_NONE); + uiItemR(layout, ptr, "domain", UI_ITEM_NONE, "", ICON_NONE); +} + +static void node_init(bNodeTree * /*tree*/, bNode *node) +{ + node->custom1 = CD_PROP_FLOAT; + node->custom2 = int16_t(AttrDomain::Point); +} + +enum class Operation { StdDev = 0, Variance = 1 }; + +static std::optional node_type_from_other_socket(const bNodeSocket &socket) +{ + switch (socket.type) { + case SOCK_FLOAT: + case SOCK_BOOLEAN: + case SOCK_INT: + return CD_PROP_FLOAT; + case SOCK_VECTOR: + case SOCK_RGBA: + return CD_PROP_FLOAT3; + default: + return std::nullopt; + } +} + +static void node_gather_link_searches(GatherLinkSearchOpParams ¶ms) +{ + const NodeDeclaration &declaration = *params.node_type().static_declaration; + search_link_ops_for_declarations(params, declaration.inputs); + + const std::optional type = node_type_from_other_socket(params.other_socket()); + if (!type) { + return; + } + if (params.in_out() == SOCK_OUT) { + params.add_item( + IFACE_("Standard Deviation"), + [type](LinkSearchOpParams ¶ms) { + bNode &node = params.add_node("GeometryNodeFieldVariance"); + node.custom1 = *type; + params.update_and_connect_available_socket(node, "Standard Deviation"); + }, + 0); + params.add_item( + IFACE_("Variance"), + [type](LinkSearchOpParams ¶ms) { + bNode &node = params.add_node("GeometryNodeFieldVariance"); + node.custom1 = *type; + params.update_and_connect_available_socket(node, "Variance"); + }, + -1); + } + else { + params.add_item( + IFACE_("Value"), + [type](LinkSearchOpParams ¶ms) { + bNode &node = params.add_node("GeometryNodeFieldVariance"); + node.custom1 = *type; + params.update_and_connect_available_socket(node, "Value"); + }, + 0); + } +} + +class FieldVarianceInput final : public bke::GeometryFieldInput { + private: + GField input_; + Field group_index_; + AttrDomain source_domain_; + Operation operation_; + + public: + FieldVarianceInput(const AttrDomain source_domain, + GField input, + Field group_index, + Operation operation) + : bke::GeometryFieldInput(input.cpp_type(), "Calculation"), + input_(std::move(input)), + group_index_(std::move(group_index)), + source_domain_(source_domain), + operation_(operation) + { + } + + GVArray get_varray_for_context(const bke::GeometryFieldContext &context, + const IndexMask & /*mask*/) const final + { + const AttributeAccessor attributes = *context.attributes(); + const int64_t domain_size = attributes.domain_size(source_domain_); + if (domain_size == 0) { + return {}; + } + + const bke::GeometryFieldContext source_context{context, source_domain_}; + fn::FieldEvaluator evaluator{source_context, domain_size}; + evaluator.add(input_); + evaluator.add(group_index_); + evaluator.evaluate(); + const GVArray g_values = evaluator.get_evaluated(0); + const VArray group_indices = evaluator.get_evaluated(1); + + GVArray g_outputs; + + bke::attribute_math::convert_to_static_type(g_values.type(), [&](auto dummy) { + using T = decltype(dummy); + if constexpr (is_same_any_v) { + const VArraySpan values = g_values.typed(); + + if (operation_ == Operation::StdDev) { + if (group_indices.is_single()) { + const T mean = std::reduce(values.begin(), values.end(), T()) / domain_size; + const T sum_of_squared_diffs = std::reduce( + values.begin(), values.end(), T(), [mean](T accumulator, const T &value) { + T difference = mean - value; + return accumulator + difference * difference; + }); + g_outputs = VArray::ForSingle(math::sqrt(sum_of_squared_diffs / domain_size), + domain_size); + } + else { + Map> sum_and_counts; + Map deviations; + + for (const int i : values.index_range()) { + auto &pair = sum_and_counts.lookup_or_add(group_indices[i], std::make_pair(T(), 0)); + pair.first = pair.first + values[i]; + pair.second = pair.second + 1; + } + + for (const int i : values.index_range()) { + const auto &pair = sum_and_counts.lookup(group_indices[i]); + T mean = pair.first / pair.second; + T deviation = (mean - values[i]); + deviation = deviation * deviation; + + T &dev_sum = deviations.lookup_or_add(group_indices[i], T()); + dev_sum = dev_sum + deviation; + } + + Array outputs(domain_size); + for (const int i : values.index_range()) { + const auto &pair = sum_and_counts.lookup(group_indices[i]); + outputs[i] = math::sqrt(deviations.lookup(group_indices[i]) / pair.second); + } + g_outputs = VArray::ForContainer(std::move(outputs)); + } + } + else { + if (group_indices.is_single()) { + const T mean = std::reduce(values.begin(), values.end(), T()) / domain_size; + const T sum_of_squared_diffs = std::reduce( + values.begin(), values.end(), T(), [mean](T accumulator, const T &value) { + T difference = mean - value; + return accumulator + difference * difference; + }); + g_outputs = VArray::ForSingle(sum_of_squared_diffs / domain_size, domain_size); + } + else { + Map> sum_and_counts; + Map deviations; + + for (const int i : values.index_range()) { + auto &pair = sum_and_counts.lookup_or_add(group_indices[i], std::make_pair(T(), 0)); + pair.first = pair.first + values[i]; + pair.second = pair.second + 1; + } + + for (const int i : values.index_range()) { + const auto &pair = sum_and_counts.lookup(group_indices[i]); + T mean = pair.first / pair.second; + T deviation = (mean - values[i]); + deviation = deviation * deviation; + + T &dev_sum = deviations.lookup_or_add(group_indices[i], T()); + dev_sum = dev_sum + deviation; + } + + Array outputs(domain_size); + for (const int i : values.index_range()) { + const auto &pair = sum_and_counts.lookup(group_indices[i]); + outputs[i] = deviations.lookup(group_indices[i]) / pair.second; + } + g_outputs = VArray::ForContainer(std::move(outputs)); + } + } + } + }); + + return attributes.adapt_domain(std::move(g_outputs), source_domain_, context.domain()); + } + + uint64_t hash() const override + { + return get_default_hash(input_, group_index_, source_domain_, operation_); + } + + bool is_equal_to(const fn::FieldNode &other) const override + { + if (const FieldVarianceInput *other_field = dynamic_cast(&other)) { + return input_ == other_field->input_ && group_index_ == other_field->group_index_ && + source_domain_ == other_field->source_domain_ && + operation_ == other_field->operation_; + } + return false; + } + + std::optional preferred_domain( + const GeometryComponent & /*component*/) const override + { + return source_domain_; + } +}; + +static void node_geo_exec(GeoNodeExecParams params) +{ + const AttrDomain source_domain = AttrDomain(params.node().custom2); + + const Field group_index_field = params.extract_input>("Group Index"); + const GField input_field = params.extract_input("Value"); + if (params.output_is_required("Standard Deviation")) { + params.set_output( + "Standard Deviation", + GField{std::make_shared( + source_domain, input_field, group_index_field, Operation::StdDev)}); + } + if (params.output_is_required("Variance")) { + params.set_output( + "Variance", + GField{std::make_shared( + source_domain, input_field, group_index_field, Operation::Variance)}); + } +} + +static void node_rna(StructRNA *srna) +{ + static EnumPropertyItem items[] = { + {CD_PROP_FLOAT, "FLOAT", 0, "Float", "Floating-point value"}, + {CD_PROP_FLOAT3, "FLOAT_VECTOR", 0, "Vector", "3D vector with floating-point values"}, + {0, nullptr, 0, nullptr, nullptr}, + }; + + RNA_def_node_enum(srna, + "data_type", + "Data Type", + "Type of data the outputs are calculated from", + items, + NOD_inline_enum_accessors(custom1), + CD_PROP_FLOAT); + + RNA_def_node_enum(srna, + "domain", + "Domain", + "", + rna_enum_attribute_domain_items, + NOD_inline_enum_accessors(custom2), + int(AttrDomain::Point), + nullptr, + true); +} + +static void node_register() +{ + static blender::bke::bNodeType ntype; + + geo_node_type_base(&ntype, "GeometryNodeFieldVariance"); + ntype.ui_name = "Field Variance"; + ntype.ui_description = "Calculate the standard deviation and variance of a given field"; + ntype.nclass = NODE_CLASS_CONVERTER; + ntype.geometry_node_execute = node_geo_exec; + ntype.initfunc = node_init; + ntype.draw_buttons = node_layout; + ntype.declare = node_declare; + ntype.gather_link_search_ops = node_gather_link_searches; + blender::bke::node_register_type(ntype); + node_rna(ntype.rna_ext.srna); +} +NOD_REGISTER_NODE(node_register) + +} // namespace blender::nodes::node_geo_field_variance_cc