# Copyright Contributors to the OpenVDB Project
# SPDX-License-Identifier: Apache-2.0
#
"""Functional API for loading and saving grid batches in NanoVDB format."""
from __future__ import annotations
from typing import TYPE_CHECKING, overload
import torch
from ..jagged_tensor import JaggedTensor
from ..types import DeviceIdentifier, resolve_device
if TYPE_CHECKING:
from .._fvdb_cpp import NanoVDBGridMetadata
from ..grid import Grid
from ..grid_batch import GridBatch
def _wrap_grid(cpp_impl):
from ..grid_batch import GridBatch
return GridBatch(data=cpp_impl)
# ---------------------------------------------------------------------------
# Load
# ---------------------------------------------------------------------------
@overload
def load_nanovdb(
path: str,
*,
device: DeviceIdentifier = "cpu",
verbose: bool = False,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load_nanovdb(
path: str,
*,
indices: list[int],
device: DeviceIdentifier = "cpu",
verbose: bool = False,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load_nanovdb(
path: str,
*,
index: int,
device: DeviceIdentifier = "cpu",
verbose: bool = False,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load_nanovdb(
path: str,
*,
names: list[str],
device: DeviceIdentifier = "cpu",
verbose: bool = False,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
@overload
def load_nanovdb(
path: str,
*,
name: str,
device: DeviceIdentifier = "cpu",
verbose: bool = False,
) -> tuple[GridBatch, JaggedTensor, list[str]]: ...
[docs]
def load_nanovdb(
path: str,
*,
indices: list[int] | None = None,
index: int | None = None,
names: list[str] | None = None,
name: str | None = None,
device: DeviceIdentifier = "cpu",
verbose: bool = False,
) -> tuple[GridBatch, JaggedTensor, list[str]]:
"""Load a grid batch from a ``.nvdb`` file.
Args:
path (str): Path to the ``.nvdb`` file.
indices (list[int] | None): Optional list of grid indices to load.
index (int | None): Optional single grid index to load.
names (list[str] | None): Optional list of grid names to load.
name (str | None): Optional single grid name to load.
device (DeviceIdentifier): Device to load onto. Defaults to ``"cpu"``.
verbose (bool): Print information about loaded grids.
Returns:
grid_batch (GridBatch): The loaded grid batch.
data (JaggedTensor): Per-voxel data.
names (list[str]): Grid names.
.. seealso:: :func:`load_nanovdb_single`
"""
from .._fvdb_cpp import load as _load
device = resolve_device(device)
selectors = [indices is not None, index is not None, names is not None, name is not None]
if sum(selectors) > 1:
raise ValueError("Only one of indices, index, names, or name can be specified")
if indices is not None:
grid_impl, data_impl, names_out = _load(path, indices, device, verbose)
elif index is not None:
grid_impl, data_impl, names_out = _load(path, index, device, verbose)
elif names is not None:
grid_impl, data_impl, names_out = _load(path, names, device, verbose)
elif name is not None:
grid_impl, data_impl, names_out = _load(path, name, device, verbose)
else:
grid_impl, data_impl, names_out = _load(path, device, verbose)
return _wrap_grid(grid_impl), JaggedTensor(impl=data_impl), names_out
# ---------------------------------------------------------------------------
# Metadata
# ---------------------------------------------------------------------------
def read_nanovdb_metadata(path: str) -> list[NanoVDBGridMetadata]:
"""Read per-grid metadata from a ``.nvdb`` file without loading voxel data.
This is a lightweight way to enumerate which grids are in a file (their names, value types,
voxel counts, etc.) before deciding which ones to load. It only reads file headers, so it is
much cheaper than :func:`load_nanovdb` and works on files containing a mix of grid value types
(e.g. a float density grid alongside a Vec3f velocity grid) where the full batch load is not
representable as a single :class:`JaggedTensor`.
Args:
path (str): Path to the ``.nvdb`` file.
Returns:
list[NanoVDBGridMetadata]: One entry per grid, in file order. Each entry exposes
``name``, ``type`` (e.g. ``"float"``, ``"Vec3f"``, ``"OnIndex"``), ``grid_class``
(e.g. ``"SDF"``, ``"FOG"``, ``"TENSOR"``, ``"?"``), ``voxel_count``, ``voxel_size``,
``index_bbox_min`` and ``index_bbox_max``.
Example:
>>> for m in read_nanovdb_metadata("explosion.nvdb"):
... print(m.name, m.type, m.voxel_count)
"""
from .._fvdb_cpp import read_metadata as _read_metadata
return _read_metadata(path)
def grid_names_in_nanovdb(path: str) -> list[str]:
"""Return the list of grid names stored in a ``.nvdb`` file, in file order.
This is a convenience wrapper around :func:`read_nanovdb_metadata` that returns just the names.
Args:
path (str): Path to the ``.nvdb`` file.
Returns:
list[str]: Grid names in the order they appear in the file. Names are not guaranteed to be
unique; NanoVDB allows duplicate grid names.
"""
return [m.name for m in read_nanovdb_metadata(path)]
# ---------------------------------------------------------------------------
# Save
# ---------------------------------------------------------------------------
[docs]
def save_nanovdb(
grid: GridBatch,
path: str,
data: JaggedTensor | None = None,
names: list[str] | str | None = None,
name: str | None = None,
compressed: bool = False,
verbose: bool = False,
) -> None:
"""Save a grid batch and optional voxel data to a ``.nvdb`` file.
Args:
grid (GridBatch): The grid batch to save.
path (str): File path (should have ``.nvdb`` extension).
data (JaggedTensor | None): Optional voxel data to save with the grids.
names (list[str] | str | None): Names for each grid, or a single name for all.
name (str | None): Single name for all grids (takes precedence over ``names``).
compressed (bool): Use Blosc compression. Default ``False``.
verbose (bool): Print information about saved grids. Default ``False``.
.. seealso:: :func:`save_nanovdb_single`
"""
from .._fvdb_cpp import save as _save
grid_data = grid.data
data_impl = data._impl if data else None
if name is not None:
_save(path, grid_data, data_impl, name, compressed, verbose)
elif names is not None:
if isinstance(names, str):
_save(path, grid_data, data_impl, names, compressed, verbose)
else:
_save(path, grid_data, data_impl, names, compressed, verbose)
else:
_save(path, grid_data, data_impl, [], compressed, verbose)
# ---------------------------------------------------------------------------
# Single variants (Grid + torch.Tensor)
# ---------------------------------------------------------------------------
def _wrap_single_grid(cpp_impl):
from ..grid import Grid
return Grid(data=cpp_impl)
[docs]
def load_nanovdb_single(
path: str,
*,
index: int = 0,
name: str | None = None,
device: DeviceIdentifier = "cpu",
verbose: bool = False,
) -> tuple[Grid, torch.Tensor, str]:
"""Load a single grid from a ``.nvdb`` file.
Args:
path (str): Path to the ``.nvdb`` file.
index (int): Grid index to load. Default ``0``.
name (str | None): Optional grid name to load (overrides ``index``).
device (DeviceIdentifier): Device to load onto. Defaults to ``"cpu"``.
verbose (bool): Print information about loaded grids.
Returns:
grid (Grid): The loaded single grid.
data (torch.Tensor): Per-voxel data.
name (str): Grid name.
.. seealso:: :func:`load_nanovdb`
"""
if name is not None:
gb, jt_data, names_out = load_nanovdb(path, name=name, device=device, verbose=verbose)
else:
gb, jt_data, names_out = load_nanovdb(path, index=index, device=device, verbose=verbose)
return _wrap_single_grid(gb.data), jt_data.jdata, names_out[0] if names_out else ""
[docs]
def save_nanovdb_single(
grid: Grid,
path: str,
data: torch.Tensor | None = None,
name: str | None = None,
compressed: bool = False,
verbose: bool = False,
) -> None:
"""Save a single grid and optional voxel data to a ``.nvdb`` file.
Args:
grid (Grid): The single grid to save.
path (str): File path (should have ``.nvdb`` extension).
data (torch.Tensor | None): Optional voxel data as a plain tensor.
name (str | None): Optional name for the grid.
compressed (bool): Use Blosc compression. Default ``False``.
verbose (bool): Print information about saved grids. Default ``False``.
.. seealso:: :func:`save_nanovdb`
"""
from .._fvdb_cpp import save as _save
grid_data = grid.data
if data is not None:
data_impl = JaggedTensor(data)._impl
else:
data_impl = None
if name is not None:
_save(path, grid_data, data_impl, name, compressed, verbose)
else:
_save(path, grid_data, data_impl, [], compressed, verbose)