Merge branch 'index_grid' into 'main'

NanoVDB Index grid support

See merge request omniverse/warp!341

docs/modules/functions.rst

@@ -1563,6 +1563,32 @@ Geometry
 
 Volumes
 ---------------
+.. py:function:: volume_sample(id: uint64, uvw: vec3f, sampling_mode: int32, dtype: Any)
+
+    Sample the volume of type `dtype` given by ``id`` at the volume local-space point ``uvw``.
+
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`
+
+
+.. py:function:: volume_sample_grad(id: uint64, uvw: vec3f, sampling_mode: int32, grad: Any, dtype: Any)
+
+    Sample the volume given by ``id`` and its gradient at the volume local-space point ``uvw``.
+
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`
+
+
+.. py:function:: volume_lookup(id: uint64, i: int32, j: int32, k: int32, dtype: Any)
+
+    Returns the value of voxel with coordinates ``i``, ``j``, ``k`` for a volume of type type `dtype`.
+
+    If the voxel at this index does not exist, this function returns the background value.
+
+
+.. py:function:: volume_store(id: uint64, i: int32, j: int32, k: int32, value: Any)
+
+    Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.
+
+
 .. py:function:: volume_sample_f(id: uint64, uvw: vec3f, sampling_mode: int32) -> float
 
     Sample the volume given by ``id`` at the volume local-space point ``uvw``.
@@ -1625,6 +1651,35 @@ Volumes
     Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.
 
 
+.. py:function:: volume_sample_index(id: uint64, uvw: vec3f, sampling_mode: int32, voxel_data: Array[Any], background: Any)
+
+    Sample the volume given by ``id`` at the volume local-space point ``uvw``.
+
+    Values for allocated voxels are read from the ``voxel_data`` array, and `background` is used as the value of non-existing voxels.
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR`.
+    This function is available for both index grids and classical volumes.
+
+
+
+.. py:function:: volume_sample_grad_index(id: uint64, uvw: vec3f, sampling_mode: int32, voxel_data: Array[Any], background: Any, grad: Any)
+
+    Sample the volume given by ``id`` and its gradient at the volume local-space point ``uvw``.
+
+    Values for allocated voxels are read from the ``voxel_data`` array, and `background` is used as the value of non-existing voxels.
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR`.
+    This function is available for both index grids and classical volumes.
+
+
+
+.. py:function:: volume_lookup_index(id: uint64, i: int32, j: int32, k: int32) -> int32
+
+    Returns the index associated to the voxel with coordinates ``i``, ``j``, ``k``.
+
+    If the voxel at this index does not exist, this function returns -1.
+    This function is available for both index grids and classical volumes.
+
+
+
 .. py:function:: volume_index_to_world(id: uint64, uvw: vec3f) -> vec3f
 
     Transform a point ``uvw`` defined in volume index space to world space given the volume's intrinsic affine transformation.

docs/modules/runtime.rst

@@ -975,12 +975,12 @@ or use built-in closest-point or trilinear interpolation to sample grid data fro
 
 Volume objects can be created directly from Warp arrays containing a NanoVDB grid, from the contents of a
 standard ``.nvdb`` file using :func:`load_from_nvdb() <warp.Volume.load_from_nvdb>`,
+from an uncompressed in-memory buffer using :func:`load_from_address() <warp.Volume.load_from_address>`,
 or from a dense 3D NumPy array using :func:`load_from_numpy() <warp.Volume.load_from_numpy>`.
 
-Volumes can also be created using :func:`allocate() <warp.Volume.allocate>` or
-:func:`allocate_by_tiles() <warp.Volume.allocate_by_tiles>`. The values for a Volume object can be modified in a Warp
-kernel using :func:`wp.volume_store_f() <warp.volume_store_f>`, :func:`wp.volume_store_v() <warp.volume_store_v>`, and
-:func:`wp.volume_store_i() <warp.volume_store_i>`.
+Volumes can also be created using :func:`allocate() <warp.Volume.allocate>`, 
+:func:`allocate_by_tiles() <warp.Volume.allocate_by_tiles>` or :func:`allocate_by_voxels() <warp.Volume.allocate_by_voxels>`. 
+The values for a Volume object can be modified in a Warp kernel using :func:`wp.volume_store() <warp.volume_store>`.
 
 .. note::
     Warp does not currently support modifying the topology of sparse volumes at runtime.
@@ -995,8 +995,11 @@ Below we give an example of creating a Volume object from an existing NanoVDB fi
 
 .. note::
     Files written by the NanoVDB library, commonly marked by the ``.nvdb`` extension, can contain multiple grids with
-    various compression methods, but a :class:`Volume` object represents a single NanoVDB grid therefore only files with
-    a single grid are supported. NanoVDB's uncompressed and zip-compressed file formats are supported.
+    various compression methods, but a :class:`Volume` object represents a single NanoVDB grid. 
+    The first grid is loaded by default, then  Warp volumes corresponding to the other grids in the file can be created
+    using repeated calls to :func:`load_next_grid() <warp.Volume.load_next_grid>`.
+    NanoVDB's uncompressed and zip-compressed file formats are supported out-of-the-box, blosc compressed files require
+    the `blosc` Python package to be installed.
 
 To sample the volume inside a kernel we pass a reference to it by ID, and use the built-in sampling modes::
 
@@ -1014,11 +1017,35 @@ To sample the volume inside a kernel we pass a reference to it by ID, and use th
         q = wp.volume_world_to_index(volume, p)
 
         # sample volume with trilinear interpolation
-        f = wp.volume_sample_f(volume, q, wp.Volume.LINEAR)
+        f = wp.volume_sample(volume, q, wp.Volume.LINEAR, dtype=float)
 
         # write result
         samples[tid] = f
 
+Warp also supports NanoVDB index grids, which provide a memory-efficient linearization of voxel indices that can refer 
+to values in arbitrarily shaped arrays::
+
+    @wp.kernel
+    def sample_index_grid(volume: wp.uint64,
+                         points: wp.array(dtype=wp.vec3),
+                         voxel_values: wp.array(dtype=Any)):
+
+        tid = wp.tid()
+
+        # load sample point in world-space
+        p = points[tid]
+
+        # transform position to the volume's local-space
+        q = wp.volume_world_to_index(volume, p)
+
+        # sample volume with trilinear interpolation
+        background_value = voxel_values.dtype(0.0)
+        f = wp.volume_sample_index(volume, q, wp.Volume.LINEAR, voxel_values, background_value)
+
+The coordinates of all indexable voxels can be recovered using :func:`get_voxels() <warp.Volume.get_voxels>`.
+NanoVDB grids may also contains embedded *blind* data arrays; those can be accessed with the 
+:func:`feature_array() <warp.Volume.feature_array>` function.
+
 .. autoclass:: Volume
     :members:
     :undoc-members:

warp/builtins.py

@@ -2103,6 +2103,112 @@ def spatial_vector_constructor_value_func(arg_types, kwds, templates):
 # ---------------------------------
 # Volumes
 
+_volume_supported_value_types = {
+    int32,
+    int64,
+    uint32,
+    float32,
+    float64,
+    vec3f,
+    vec3d,
+    vec4f,
+    vec4d,
+}
+
+
+def volume_value_func(arg_types, kwds, templates):
+    try:
+        dtype = kwds["dtype"]
+    except KeyError as err:
+        raise RuntimeError(
+            "'dtype' keyword argument must be specified when calling generic volume lookup or sampling functions"
+        ) from err
+
+    if dtype not in _volume_supported_value_types:
+        raise RuntimeError(f"Unsupported volume type '{type_repr(dtype)}'")
+
+    templates.append(dtype)
+
+    return dtype
+
+
+add_builtin(
+    "volume_sample",
+    input_types={"id": uint64, "uvw": vec3, "sampling_mode": int, "dtype": Any},
+    value_func=volume_value_func,
+    export=False,
+    group="Volumes",
+    doc="""Sample the volume of type `dtype` given by ``id`` at the volume local-space point ``uvw``.
+
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`""",
+)
+
+
+def check_volume_value_grad_compatibility(dtype, grad_dtype):
+    if type_is_vector(dtype):
+        expected = matrix(shape=(type_length(dtype), 3), dtype=type_scalar_type(dtype))
+    else:
+        expected = vector(length=3, dtype=dtype)
+
+    if not types_equal(grad_dtype, expected):
+        raise RuntimeError(f"Incompatible gradient type, expected {type_repr(expected)}, got {type_repr(grad_dtype)}")
+
+
+def volume_sample_grad_value_func(arg_types, kwds, templates):
+    dtype = volume_value_func(arg_types, kwds, templates)
+
+    if len(arg_types) < 4:
+        raise RuntimeError("'volume_sample_grad' requires 4 positional arguments")
+
+    grad_type = arg_types[3]
+    check_volume_value_grad_compatibility(dtype, grad_type)
+    return dtype
+
+
+add_builtin(
+    "volume_sample_grad",
+    input_types={"id": uint64, "uvw": vec3, "sampling_mode": int, "grad": Any, "dtype": Any},
+    value_func=volume_sample_grad_value_func,
+    export=False,
+    group="Volumes",
+    doc="""Sample the volume given by ``id`` and its gradient at the volume local-space point ``uvw``.
+
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR.`""",
+)
+
+add_builtin(
+    "volume_lookup",
+    input_types={"id": uint64, "i": int, "j": int, "k": int, "dtype": Any},
+    value_type=int,
+    value_func=volume_value_func,
+    export=False,
+    group="Volumes",
+    doc="""Returns the value of voxel with coordinates ``i``, ``j``, ``k`` for a volume of type type `dtype`.
+
+    If the voxel at this index does not exist, this function returns the background value.""",
+)
+
+
+def volume_store_value_func(arg_types, kwds, templates):
+    if len(arg_types) < 4:
+        raise RuntimeError("'volume_store' requires 5 positional arguments")
+
+    dtype = arg_types[4]
+    if dtype not in _volume_supported_value_types:
+        raise RuntimeError(f"Unsupported volume type '{type_repr(dtype)}'")
+
+    return None
+
+
+add_builtin(
+    "volume_store",
+    value_func=volume_store_value_func,
+    input_types={"id": uint64, "i": int, "j": int, "k": int, "value": Any},
+    export=False,
+    group="Volumes",
+    doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
+)
+
 add_builtin(
     "volume_sample_f",
     input_types={"id": uint64, "uvw": vec3, "sampling_mode": int},
@@ -2192,6 +2298,81 @@ def spatial_vector_constructor_value_func(arg_types, kwds, templates):
     doc="""Store ``value`` at the voxel with coordinates ``i``, ``j``, ``k``.""",
 )
 
+
+def volume_sample_index_value_func(arg_types, kwds, templates):
+    if len(arg_types) != 5:
+        raise RuntimeError("'volume_sample_index' requires 5 positional arguments")
+
+    dtype = arg_types[3].dtype
+
+    if not types_equal(dtype, arg_types[4]):
+        raise RuntimeError("The 'voxel_data' array and the 'background' value must have the same dtype")
+
+    return dtype
+
+
+add_builtin(
+    "volume_sample_index",
+    input_types={"id": uint64, "uvw": vec3, "sampling_mode": int, "voxel_data": array(dtype=Any), "background": Any},
+    value_func=volume_sample_index_value_func,
+    export=False,
+    group="Volumes",
+    doc="""Sample the volume given by ``id`` at the volume local-space point ``uvw``.
+
+    Values for allocated voxels are read from the ``voxel_data`` array, and `background` is used as the value of non-existing voxels.
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR`.
+    This function is available for both index grids and classical volumes.
+    """,
+)
+
+
+def volume_sample_grad_index_value_func(arg_types, kwds, templates):
+    if len(arg_types) != 6:
+        raise RuntimeError("'volume_sample_grad_index' requires 6 positional arguments")
+
+    dtype = arg_types[3].dtype
+
+    if not types_equal(dtype, arg_types[4]):
+        raise RuntimeError("The 'voxel_data' array and the 'background' value must have the same dtype")
+
+    grad_type = arg_types[5]
+    check_volume_value_grad_compatibility(dtype, grad_type)
+    return dtype
+
+
+add_builtin(
+    "volume_sample_grad_index",
+    input_types={
+        "id": uint64,
+        "uvw": vec3,
+        "sampling_mode": int,
+        "voxel_data": array(dtype=Any),
+        "background": Any,
+        "grad": Any,
+    },
+    value_func=volume_sample_grad_index_value_func,
+    export=False,
+    group="Volumes",
+    doc="""Sample the volume given by ``id`` and its gradient at the volume local-space point ``uvw``.
+
+    Values for allocated voxels are read from the ``voxel_data`` array, and `background` is used as the value of non-existing voxels.
+    Interpolation should be :attr:`warp.Volume.CLOSEST` or :attr:`wp.Volume.LINEAR`.
+    This function is available for both index grids and classical volumes.
+   """,
+)
+
+add_builtin(
+    "volume_lookup_index",
+    input_types={"id": uint64, "i": int, "j": int, "k": int},
+    value_type=int32,
+    group="Volumes",
+    doc="""Returns the index associated to the voxel with coordinates ``i``, ``j``, ``k``.
+
+    If the voxel at this index does not exist, this function returns -1.
+    This function is available for both index grids and classical volumes.
+    """,
+)
+
 add_builtin(
     "volume_index_to_world",
     input_types={"id": uint64, "uvw": vec3},