Geometry Nodes: Sample grid node

This simple node finds the values of a volume grid at
positions in the local space used in geometry nodes
evaluation. There are three interpolation modes to
choose how to mix values between neighboring voxels.

For the implementation, first the values are sampled
with the grid's type directly, then implicit type conversions
are used to get the final type. This makes gives us flexibility
in case there aren't exact matches in support between grid
types and Blender types.

Pull Request: https://projects.blender.org/blender/blender/pulls/118397

scripts/startup/bl_ui/node_add_menu_geometry.py

@@ -666,6 +666,7 @@ class NODE_MT_category_GEO_VOLUME(Menu):
         layout = self.layout
         if context.preferences.experimental.use_new_volume_nodes:
             layout.menu("NODE_MT_geometry_node_GEO_VOLUME_READ")
+            layout.menu("NODE_MT_geometry_node_volume_sample")
             layout.menu("NODE_MT_geometry_node_GEO_VOLUME_WRITE")
             layout.separator()
         layout.menu("NODE_MT_geometry_node_GEO_VOLUME_OPERATIONS")
@@ -693,6 +694,16 @@ class NODE_MT_geometry_node_GEO_VOLUME_WRITE(Menu):
         node_add_menu.draw_assets_for_catalog(layout, "Volume/Write")
 
 
+class NODE_MT_geometry_node_volume_sample(Menu):
+    bl_idname = "NODE_MT_geometry_node_volume_sample"
+    bl_label = "Sample"
+
+    def draw(self, context):
+        layout = self.layout
+        node_add_menu.add_node_type(layout, "GeometryNodeSampleGrid")
+        node_add_menu.draw_assets_for_catalog(layout, "Volume/Sample")
+
+
 class NODE_MT_geometry_node_GEO_VOLUME_OPERATIONS(Menu):
     bl_idname = "NODE_MT_geometry_node_GEO_VOLUME_OPERATIONS"
     bl_label = "Operations"
@@ -785,6 +796,7 @@ classes = (
     NODE_MT_category_simulation,
     NODE_MT_category_GEO_VOLUME,
     NODE_MT_geometry_node_GEO_VOLUME_READ,
+    NODE_MT_geometry_node_volume_sample,
     NODE_MT_geometry_node_GEO_VOLUME_WRITE,
     NODE_MT_geometry_node_GEO_VOLUME_OPERATIONS,
     NODE_MT_geometry_node_GEO_VOLUME_PRIMITIVES,

source/blender/blenkernel/BKE_customdata.hh

@@ -812,4 +812,5 @@ void CustomData_debug_info_from_layers(const CustomData *data, const char *inden
 
 namespace blender::bke {
 std::optional<VolumeGridType> custom_data_type_to_volume_grid_type(eCustomDataType type);
+std::optional<eCustomDataType> volume_grid_type_to_custom_data_type(VolumeGridType type);
 }  // namespace blender::bke

source/blender/blenkernel/BKE_node.hh

@@ -1329,6 +1329,7 @@ void BKE_nodetree_remove_layer_n(struct bNodeTree *ntree, struct Scene *scene, i
 #define GEO_NODE_STORE_NAMED_GRID 2122
 #define GEO_NODE_SORT_ELEMENTS 2123
 #define GEO_NODE_MENU_SWITCH 2124
+#define GEO_NODE_SAMPLE_GRID 2125
 
 /** \} */
 

source/blender/blenkernel/intern/customdata.cc

@@ -5542,6 +5542,22 @@ std::optional<VolumeGridType> custom_data_type_to_volume_grid_type(const eCustom
   }
 }
 
+std::optional<eCustomDataType> volume_grid_type_to_custom_data_type(const VolumeGridType type)
+{
+  switch (type) {
+    case VOLUME_GRID_FLOAT:
+      return CD_PROP_FLOAT;
+    case VOLUME_GRID_VECTOR_FLOAT:
+      return CD_PROP_FLOAT3;
+    case VOLUME_GRID_INT:
+      return CD_PROP_INT32;
+    case VOLUME_GRID_BOOLEAN:
+      return CD_PROP_BOOL;
+    default:
+      return std::nullopt;
+  }
+}
+
 /** \} */
 
 }  // namespace blender::bke

source/blender/nodes/NOD_static_types.h

@@ -421,6 +421,7 @@ DefNode(GeometryNode, GEO_NODE_REPLACE_MATERIAL, 0, "REPLACE_MATERIAL", ReplaceM
 DefNode(GeometryNode, GEO_NODE_RESAMPLE_CURVE, 0, "RESAMPLE_CURVE", ResampleCurve, "Resample Curve", "Generate a poly spline for each input spline")
 DefNode(GeometryNode, GEO_NODE_REVERSE_CURVE, 0, "REVERSE_CURVE", ReverseCurve, "Reverse Curve", "Change the direction of curves by swapping their start and end data")
 DefNode(GeometryNode, GEO_NODE_ROTATE_INSTANCES, 0, "ROTATE_INSTANCES", RotateInstances, "Rotate Instances", "Rotate geometry instances in local or global space")
+DefNode(GeometryNode, GEO_NODE_SAMPLE_GRID, 0, "SAMPLE_GRID", SampleGrid, "Sample Grid", "")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_CURVE, def_geo_curve_sample, "SAMPLE_CURVE", SampleCurve, "Sample Curve", "Retrieve data from a point on a curve at a certain distance from its start")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_INDEX, def_geo_sample_index, "SAMPLE_INDEX", SampleIndex, "Sample Index", "Retrieve values from specific geometry elements")
 DefNode(GeometryNode, GEO_NODE_SAMPLE_NEAREST_SURFACE, 0, "SAMPLE_NEAREST_SURFACE", SampleNearestSurface, "Sample Nearest Surface", "Calculate the interpolated value of a mesh attribute on the closest point of its surface")

source/blender/nodes/geometry/CMakeLists.txt

@@ -155,6 +155,7 @@ set(SRC
   nodes/node_geo_remove_attribute.cc
   nodes/node_geo_repeat.cc
   nodes/node_geo_rotate_instances.cc
+  nodes/node_geo_sample_grid.cc
   nodes/node_geo_sample_index.cc
   nodes/node_geo_sample_nearest.cc
   nodes/node_geo_sample_nearest_surface.cc

source/blender/nodes/geometry/nodes/node_geo_sample_grid.cc

@@ -0,0 +1,303 @@
+/* SPDX-FileCopyrightText: 2024 Blender Authors
+ *
+ * SPDX-License-Identifier: GPL-2.0-or-later */
+
+#include "BKE_customdata.hh"
+#include "BKE_type_conversions.hh"
+#include "BKE_volume_grid.hh"
+#include "BKE_volume_openvdb.hh"
+
+#include "BLI_index_mask.hh"
+#include "BLI_virtual_array.hh"
+
+#include "NOD_rna_define.hh"
+#include "NOD_socket_search_link.hh"
+
+#include "UI_interface.hh"
+#include "UI_resources.hh"
+
+#include "RNA_enum_types.hh"
+
+#ifdef WITH_OPENVDB
+#  include <openvdb/tools/Interpolation.h>
+#endif
+
+#include "node_geometry_util.hh"
+
+namespace blender::nodes::node_geo_sample_grid_cc {
+
+enum class InterpolationMode {
+  Nearest = 0,
+  TriLinear = 1,
+  TriQuadratic = 2,
+};
+
+static void node_declare(NodeDeclarationBuilder &b)
+{
+  const bNode *node = b.node_or_null();
+  if (!node) {
+    return;
+  }
+  const eNodeSocketDatatype data_type = eNodeSocketDatatype(node->custom1);
+
+  b.add_input(data_type, "Grid").hide_value();
+  b.add_input<decl::Vector>("Position").implicit_field(implicit_field_inputs::position);
+
+  b.add_output(data_type, "Value").dependent_field({1});
+}
+
+static std::optional<eNodeSocketDatatype> node_type_for_socket_type(const bNodeSocket &socket)
+{
+  switch (socket.type) {
+    case SOCK_FLOAT:
+      return SOCK_FLOAT;
+    case SOCK_BOOLEAN:
+      return SOCK_BOOLEAN;
+    case SOCK_INT:
+      return SOCK_INT;
+    case SOCK_VECTOR:
+    case SOCK_RGBA:
+      return SOCK_VECTOR;
+    default:
+      return std::nullopt;
+  }
+}
+
+static void node_gather_link_search_ops(GatherLinkSearchOpParams &params)
+{
+  if (!U.experimental.use_new_volume_nodes) {
+    return;
+  }
+  const std::optional<eNodeSocketDatatype> node_type = node_type_for_socket_type(
+      params.other_socket());
+  if (!node_type) {
+    return;
+  }
+  if (params.in_out() == SOCK_IN) {
+    params.add_item(IFACE_("Grid"), [node_type](LinkSearchOpParams &params) {
+      bNode &node = params.add_node("GeometryNodeSampleGrid");
+      node.custom1 = *node_type;
+      params.update_and_connect_available_socket(node, "Grid");
+    });
+    const eNodeSocketDatatype other_type = eNodeSocketDatatype(params.other_socket().type);
+    if (params.node_tree().typeinfo->validate_link(other_type, SOCK_VECTOR)) {
+      params.add_item(IFACE_("Position"), [node_type](LinkSearchOpParams &params) {
+        bNode &node = params.add_node("GeometryNodeSampleGrid");
+        params.update_and_connect_available_socket(node, "Position");
+      });
+    }
+  }
+  else {
+    params.add_item(IFACE_("Value"), [node_type](LinkSearchOpParams &params) {
+      bNode &node = params.add_node("GeometryNodeSampleGrid");
+      node.custom1 = *node_type;
+      params.update_and_connect_available_socket(node, "Value");
+    });
+  }
+}
+
+static void node_layout(uiLayout *layout, bContext * /*C*/, PointerRNA *ptr)
+{
+  uiItemR(layout, ptr, "data_type", UI_ITEM_NONE, "", ICON_NONE);
+  uiItemR(layout, ptr, "interpolation_mode", UI_ITEM_NONE, "", ICON_NONE);
+}
+
+static void node_init(bNodeTree * /*tree*/, bNode *node)
+{
+  node->custom1 = SOCK_FLOAT;
+  node->custom2 = int16_t(InterpolationMode::TriLinear);
+}
+
+#ifdef WITH_OPENVDB
+
+template<typename T>
+void sample_grid(const bke::OpenvdbGridType<T> &grid,
+                 const InterpolationMode interpolation,
+                 const Span<float3> positions,
+                 const IndexMask &mask,
+                 MutableSpan<T> dst)
+{
+  using GridType = bke::OpenvdbGridType<T>;
+  using GridValueT = typename GridType::ValueType;
+  using AccessorT = typename GridType::ConstAccessor;
+  using TraitsT = typename bke::VolumeGridTraits<T>;
+  AccessorT accessor = grid.getConstAccessor();
+
+  auto sample_data = [&](auto sampler) {
+    mask.foreach_index([&](const int64_t i) {
+      const float3 &pos = positions[i];
+      GridValueT value = sampler.wsSample(openvdb::Vec3R(pos.x, pos.y, pos.z));
+      dst[i] = TraitsT::to_blender(value);
+    });
+  };
+
+  /* Use to the Nearest Neighbor sampler for Bool grids (no interpolation). */
+  InterpolationMode real_interpolation = interpolation;
+  if constexpr (std::is_same_v<T, bool>) {
+    real_interpolation = InterpolationMode::Nearest;
+  }
+  switch (real_interpolation) {
+    case InterpolationMode::TriLinear: {
+      openvdb::tools::GridSampler<AccessorT, openvdb::tools::BoxSampler> sampler(accessor,
+                                                                                 grid.transform());
+      sample_data(sampler);
+      break;
+    }
+    case InterpolationMode::TriQuadratic: {
+      openvdb::tools::GridSampler<AccessorT, openvdb::tools::QuadraticSampler> sampler(
+          accessor, grid.transform());
+      sample_data(sampler);
+      break;
+    }
+    case InterpolationMode::Nearest: {
+      openvdb::tools::GridSampler<AccessorT, openvdb::tools::PointSampler> sampler(
+          accessor, grid.transform());
+      sample_data(sampler);
+      break;
+    }
+  }
+}
+
+template<typename Fn> void convert_to_static_type(const VolumeGridType type, const Fn &fn)
+{
+  switch (type) {
+    case VOLUME_GRID_BOOLEAN:
+      fn(bool());
+      break;
+    case VOLUME_GRID_FLOAT:
+      fn(float());
+      break;
+    case VOLUME_GRID_INT:
+      fn(int());
+      break;
+    case VOLUME_GRID_MASK:
+      fn(bool());
+      break;
+    case VOLUME_GRID_VECTOR_FLOAT:
+      fn(float3());
+      break;
+    default:
+      break;
+  }
+}
+
+class SampleGridFunction : public mf::MultiFunction {
+  bke::GVolumeGrid grid_;
+  InterpolationMode interpolation_;
+  mf::Signature signature_;
+
+ public:
+  SampleGridFunction(bke::GVolumeGrid grid, InterpolationMode interpolation)
+      : grid_(std::move(grid)), interpolation_(interpolation)
+  {
+    BLI_assert(grid_);
+
+    const std::optional<eNodeSocketDatatype> data_type = bke::grid_type_to_socket_type(
+        grid_->grid_type());
+    const CPPType *cpp_type = bke::socket_type_to_geo_nodes_base_cpp_type(*data_type);
+    mf::SignatureBuilder builder{"Sample Volume", signature_};
+    builder.single_input<float3>("Position");
+    builder.single_output("Value", *cpp_type);
+    this->set_signature(&signature_);
+  }
+
+  void call(const IndexMask &mask, mf::Params params, mf::Context /*context*/) const override
+  {
+    const VArraySpan<float3> positions = params.readonly_single_input<float3>(0, "Position");
+    GMutableSpan dst = params.uninitialized_single_output(1, "Value");
+
+    bke::VolumeTreeAccessToken tree_token;
+    convert_to_static_type(grid_->grid_type(), [&](auto dummy) {
+      using T = decltype(dummy);
+      sample_grid<T>(
+          grid_.typed<T>().grid(tree_token), interpolation_, positions, mask, dst.typed<T>());
+    });
+  }
+};
+
+#endif /* WITH_OPENVDB */
+
+static void node_geo_exec(GeoNodeExecParams params)
+{
+#ifdef WITH_OPENVDB
+  const bNode &node = params.node();
+  const eNodeSocketDatatype data_type = eNodeSocketDatatype(node.custom1);
+  const InterpolationMode interpolation = InterpolationMode(node.custom2);
+
+  Field<float3> position = params.extract_input<Field<float3>>("Position");
+  bke::GVolumeGrid grid = params.extract_input<bke::GVolumeGrid>("Grid");
+  if (!grid) {
+    params.set_default_remaining_outputs();
+    return;
+  }
+
+  auto fn = std::make_shared<SampleGridFunction>(std::move(grid), interpolation);
+  auto op = FieldOperation::Create(std::move(fn), {std::move(position)});
+
+  const bke::DataTypeConversions &conversions = bke::get_implicit_type_conversions();
+  const CPPType &output_type = *bke::socket_type_to_geo_nodes_base_cpp_type(data_type);
+  const GField output_field = conversions.try_convert(fn::GField(std::move(op)), output_type);
+  params.set_output("Value", std::move(output_field));
+
+#else
+  node_geo_exec_with_missing_openvdb(params);
+#endif
+}
+
+static const EnumPropertyItem *data_type_filter_fn(bContext * /*C*/,
+                                                   PointerRNA * /*ptr*/,
+                                                   PropertyRNA * /*prop*/,
+                                                   bool *r_free)
+{
+  *r_free = true;
+  return enum_items_filter(
+      rna_enum_node_socket_type_items, [](const EnumPropertyItem &item) -> bool {
+        return ELEM(item.value, SOCK_FLOAT, SOCK_INT, SOCK_BOOLEAN, SOCK_VECTOR);
+      });
+}
+
+static void node_rna(StructRNA *srna)
+{
+  RNA_def_node_enum(srna,
+                    "data_type",
+                    "Data Type",
+                    "Node socket data type",
+                    rna_enum_node_socket_type_items,
+                    NOD_inline_enum_accessors(custom1),
+                    CD_PROP_FLOAT,
+                    data_type_filter_fn);
+
+  static const EnumPropertyItem interpolation_mode_items[] = {
+      {int(InterpolationMode::Nearest), "NEAREST", 0, "Nearest Neighbor", ""},
+      {int(InterpolationMode::TriLinear), "TRILINEAR", 0, "Trilinear", ""},
+      {int(InterpolationMode::TriQuadratic), "TRIQUADRATIC", 0, "Triquadratic", ""},
+      {0, nullptr, 0, nullptr, nullptr},
+  };
+
+  RNA_def_node_enum(srna,
+                    "interpolation_mode",
+                    "Interpolation Mode",
+                    "How to interpolate the values between neighboring voxels",
+                    interpolation_mode_items,
+                    NOD_inline_enum_accessors(custom2),
+                    int(InterpolationMode::TriLinear));
+}
+
+static void node_register()
+{
+  static bNodeType ntype;
+
+  geo_node_type_base(&ntype, GEO_NODE_SAMPLE_GRID, "Sample Grid", NODE_CLASS_CONVERTER);
+  ntype.initfunc = node_init;
+  ntype.declare = node_declare;
+  ntype.gather_link_search_ops = node_gather_link_search_ops;
+  ntype.geometry_node_execute = node_geo_exec;
+  ntype.draw_buttons = node_layout;
+  ntype.geometry_node_execute = node_geo_exec;
+  nodeRegisterType(&ntype);
+
+  node_rna(ntype.rna_ext.srna);
+}
+NOD_REGISTER_NODE(node_register)
+
+}  // namespace blender::nodes::node_geo_sample_grid_cc