Reference

N^3-Tree

class N3Tree(N=2, data_dim=None, depth_limit=10, init_reserve=1, init_refine=0, geom_resize_fact=1.0, radius=0.5, center=[0.5, 0.5, 0.5], data_format='RGBA', extra_data=None, device='cpu', dtype=torch.float32, map_location=None)

Bases: torch.nn.modules.module.Module

PyTorch \(N^3\)-tree library with CUDA acceleration. By \(N^3\)-tree we mean a 3D tree with branching factor N at each interior node, where \(N=2\) is the familiar octree.

Warning

nn.Parameters can change size, which makes current optimizers invalid. If any refine(): or :code:`shrink_to_fit() call returns True, or expand(), shrink() is used, please re-make any optimizers

__init__(N=2, data_dim=None, depth_limit=10, init_reserve=1, init_refine=0, geom_resize_fact=1.0, radius=0.5, center=[0.5, 0.5, 0.5], data_format='RGBA', extra_data=None, device='cpu', dtype=torch.float32, map_location=None)

Construct N^3 Tree

Parameters

  • N – int branching factor N

  • data_dim – int size of data stored at each leaf (NEW in 0.2.28: optional if data_format other than RGBA is given). If data_format = “RGBA” or empty, this defaults to 4.

  • depth_limit – int maximum depth of tree to stop branching/refining Note that the root is at depth -1. Size N^[-10] leaves (1/1024 for octree) for example are depth 9. max_depth applies to the same depth values.

  • init_reserve – int amount of nodes to reserve initially

  • init_refine – int number of times to refine entire tree initially inital resolution will be [N^(init_refine + 1)]^3. initial max_depth will be init_refine.

  • geom_resize_fact – float geometric resizing factor

  • radius – float or list, 1/2 side length of cube (possibly in each dim)

  • center – list center of space

  • data_format – a string to indicate the data format. RGBA | SH# | SG# | ASG#

  • extra_data – extra data to include with tree

  • device – str device to put data

  • dtype – str tree data type, torch.float32 (default) | torch.float64

  • map_location – str DEPRECATED old name for device (will override device and warn)

set(indices, values, cuda=True)

Set tree values,

Parameters

  • indices – torch.Tensor (Q, 3)

  • values – torch.Tensor (Q, K)

  • cuda – whether to use CUDA kernel if available. If false, uses only PyTorch version.

forward(indices, cuda=True, want_node_ids=False, world=True)

Get tree values. Differentiable.

Parameters

  • indices – (Q, 3) the points

  • cuda – whether to use CUDA kernel if available. If false, uses only PyTorch version.

  • want_node_ids – if true, returns node ID for each query.

  • world – use world space instead of [0,1]^3, default True

Returns

(Q, data_dim), [(Q)]

snap(indices)

Snap indices to lowest corner of corresponding leaf voxel

Parameters

indices – (B, 3) indices to snap

Returns

(B, 3)

partial(data_sel=None, data_format=None, dtype=None, device=None)

Get partial tree with some of the data dimensions (channels) E.g. tree.partial(-1) to get tree with data_dim 1 of last channel only

Parameters

  • data_sel – data channel selector, default is all channels

  • data_format – data format for new tree, default is current format

  • dtype – new data type, torch.float32 | torch.float64

  • device – where to put result tree

Returns

partial N3Tree (copy)

expand(data_format, data_dim=None, remap=None)

Modify the size of the data stored at the octree leaves.

Parameters

  • data_format – new data format, RGBA | SH# | SG# | ASG#

  • data_dim – data dimension; inferred from data_format by default only needed if data_format is RGBA.

  • remap – mapping of old data to new data. For each leaf, we will do new_data[remap] = old_data if the new data_dim is equal or larger, or new_data = old_data[remap] else. By default, this will be inferred automatically (maps basis functions in the correct way).

Warning

Will change the nn.Parameter size (data), breaking optimizer! Please re-create the optimizer

shrink(data_format, data_dim=None, remap=None)

Modify the size of the data stored at the octree leaves. (Alias of expand, because it can actually both expand and shrink)

Parameters

  • data_format – new data format, RGBA | SH# | SG# | ASG#

  • data_dim – data dimension; inferred from data_format by default only needed if data_format is RGBA.

  • remap – mapping of old data to new data. For each leaf, we will do new_data = old_data[remap] if the new data_dim gets smaller, or new_data[remap] = old_data else. By default, this will be inferred automatically (maps basis functions in the correct way).

Warning

Will change the nn.Parameter size (data), breaking optimizer! Please re-create the optimizer

clone(device=None)

Deep copy the tree

Parameters

device – device of output tree (could e.g. copy cuda tree to cpu)

merge(frontier_sel=None, op=<built-in method mean of type object>)

Merge leaves into selected ‘frontier’ nodes (i.e., nodes for which all children are leaves). Use shrink_to_fit() to recover memory freed.

Parameters

  • frontier_sel – selector (int, mask, list of indices etc) for frontier nodes. In same order as reduce_frontier(). Default all nodes. Typical use: use reduce_frontier(...) to determine conditions for merge, then pass bool mask (of length n_frontier) or indices to merge().

  • op – reduction to combine child leaves into node. E.g. torch.max, torch.mean. Should take a positional argument x (B, N, data_dim) and a named parameter dim (always 1), and return a matrix of (B, data_dim). If a tuple is returned, uses first result.

property n_frontier

Number of frontier nodes

property frontier_depth

Depth of frontier nodes, size (n_frontier)

reduce_frontier(op=<built-in method mean of type object>, dim=None, grad=False)

Reduce child leaf values for each ‘frontier’ node (i.e., nodes for which all children are leaves).

Parameters

  • op – reduction to combine child leaves into node. E.g. torch.max, torch.mean. Should take a positional argument x (B, N, in_dim <= data_dim) and a named parameter dim (always 1), and return a matrix of (B, your_out_dim).

  • dim – dimension(s) of data to return, e.g. -1 returns last data dimension for all ‘frontier’ nodes

  • grad – if True, returns a tensor differentiable wrt tree data. Default False.

Example

def mean_and_var_func(x, dim=1):
    # (n_frontier, tree.N^3, data_dim) -> (n_frontier, 2 * data_dim)
    # Outputs mean and variance over children of each frontier node
    return torch.cat([torch.mean(x, dim=1),
                      torch.var(x, dim=1)], dim=-1)
# returns: (n_frontier, 2 * data_dim)
tree.reduce_frontier(mean_and_var_func)
Returns

reduced tensor

max_frontier(dim=None, grad=False)

Takes max over child leaf values for each ‘frontier’ node (i.e., nodes for which all children are leaves). This is simply reduce_frontier with torch.max operation, taking the returned values and discarding the argmax part.

Parameters

  • dim – dimension(s) of data to return, e.g. -1 returns last data dimension for all ‘frontier’ nodes

  • grad – if True, returns a tensor differentiable wrt tree data. Default False.

Returns

reduced tensor

diam_frontier(dim=None, grad=False, scale=1.0)

Takes diameter over child leaf values for each ‘frontier’ node (i.e., nodes for which all children are leaves).

Parameters

  • dim – dimension(s) of data to return, e.g. -1 returns last data dimension for all ‘frontier’ nodes

  • grad – if True, returns a tensor differentiable wrt tree data. Default False.

Returns

reduced tensor

check_integrity()

Do some checks to verify the tree’s structural integrity, mostly for debugging. Errors with message if check fails; does nothing else.

refine(repeats=1, sel=None)

Refine each selected leaf node, respecting depth_limit.

Parameters

  • repeats – int number of times to repeat refinement

  • sel – (N, 4) node selector. Default selects all leaves.

Returns

True iff N3Tree.data parameter was resized, requiring optimizer reinitialization if you’re using an optimizer

Warning

The parameter tree.data can change due to refinement. If any refine() call returns True, please re-make any optimizers using tree.params().

Warning

The selector sel is assumed to contain unique leaf indices. If there are duplicates memory will be wasted. We do not dedup here for efficiency reasons.

shrink_to_fit()

Shrink data & buffers to tightly needed fit tree data, possibly dealing with fragmentation caused by merging. This is called by the save() function by default, unless shrink=False is specified there.

Warning

Will change the nn.Parameter size (data), breaking optimizer!

property max_depth

Maximum tree depth - 1

accumulate_weights(op: str = 'sum')

Begin weight accumulation.

Parameters

op – reduction to apply weight in each voxel, sum | max

Warning

Weight accumulator has not been validated and may have bugs

Example

with tree.accumulate_weights() as accum:
    ...

# (n_leaves) in same order as values etc.
accum = accum()
save(path, shrink=True, compress=True)

Save to from npz file

Parameters

  • path – npz path

  • shrink – if True (default), applies shrink_to_fit before saving

  • compress – whether to compress the npz; may be slow

classmethod load(path, device='cpu', dtype=torch.float32, map_location=None)

Load from npz file

Parameters

  • path – npz path

  • device – str device to put data

  • dtype – str torch.float32 (default) | torch.float64

  • map_location – str DEPRECATED old name for device

world2tree(indices)

Scale world points to tree (\([0,1]^3\))

tree2world(indices)

Scale tree points (\([0,1]^3\)) to world accoording to center/radius

N^3-Tree View

class N3TreeView(tree, key)

Bases: object

__init__(tree, key)

Initialize self. See help(type(self)) for accurate signature.

refine(repeats=1)

Refine selected leaves using tree.refine

property values

Values of the selected leaves (autograd enabled)

Returns

(n_leaves, data_dim) float32 note this is 2D even if key is int

property values_nograd

Values of the selected leaves (no autograd)

Returns

(n_leaves, data_dim) float32 note this is 2D even if key is int

property depths

Get a list of selected leaf depths in tree, in same order as values, corners. Root is at depth -1. Any children of root will have depth 0, etc.

Returns

(n_leaves) int32

property lengths

Get a list of selected leaf side lengths in tree (world dimensions), in same order as values, corners, depths

Returns

(n_leaves, 3) float

property lengths_local

Get a list of selected leaf side lengths in tree (local index \([0, 1]^3\)), in same order as values, corners, depths

Returns

(n_leaves) float

property corners

Get a list of selected leaf lower corners in tree (world coordinates), in same order as values, lengths, depths

Returns

(n_leaves, 3) float

property corners_local

Get a list of selected leaf lower corners in tree (local index \([0, 1]^3\)), in same order as values, lengths, depths

Returns

(n_leaves, 3) float

sample(n_samples, device=None)

Sample n_samples uniform points in each selected leaf (world coordinates)

Parameters

  • n_samples – samples for each leaf

  • device – device to output random samples in

Returns

(n_leaves, n_samples, 3) float

sample_local(n_samples)

Sample n_samples uniform points in each selected leaf (local index \([0, 1]^3\))

Returns

(n_leaves, n_samples, 3) float

aux(arr)

Index an auxiliary tree data array of size (capacity, N, N, N, Any) using this view

normal_(mean=0.0, std=1.0)

Set all values to random normal FIXME: inefficient

Parameters

  • mean – normal mean

  • std – normal std

uniform_(min=0.0, max=1.0)

Set all values to random uniform FIXME: inefficient

Parameters

  • min – interval min

  • max – interval max

clamp_(min=None, max=None)

Clamp. FIXME: inefficient

Parameters

  • min – clamp min value, None=disable

  • max – clamp max value, None=disable

relu_()

Apply relu to all elements. FIXME: inefficient

sigmoid_()

Apply sigmoid to all elements. FIXME: inefficient

nan_to_num_(inf_val=20000.0)

Convert nans to 0.0 and infs to inf_val FIXME: inefficient

Differentiable Volume Renderer

class VolumeRenderer(tree, step_size: float = 0.001, background_brightness: float = 1.0, ndc: Optional[svox.renderer.NDCConfig] = None, min_comp: int = 0, max_comp: int = - 1, density_softplus: bool = False, rgb_padding: float = 0.0)

Bases: torch.nn.modules.module.Module

Volume renderer

__init__(tree, step_size: float = 0.001, background_brightness: float = 1.0, ndc: Optional[svox.renderer.NDCConfig] = None, min_comp: int = 0, max_comp: int = - 1, density_softplus: bool = False, rgb_padding: float = 0.0)

Construct volume renderer associated with given N^3 tree.

The renderer traces rays with origins/dirs within the octree boundaries, detection ray-voxel intersections. The color and density within each voxel is assumed constant, and no interpolation is performed.

For each intersection point, it queries the tree, assuming the last data dimension is density (sigma) and the rest of the dimensions are color, formatted according to tree.data_format. It then applies SH/SG/ASG basis functions, if any, according to viewdirs. Sigmoid will be applied to these colors to normalize them, and optionally a shifted softplus is applied to the density.

Parameters

  • tree – N3Tree instance for rendering

  • step_size – float step size eps, added to each voxel aabb intersection step

  • background_brightness – float background brightness, 1.0 = white

  • ndc – NDCConfig, NDC coordinate configuration, namedtuple(width, height, focal). None = no NDC, use usual coordinates

  • min_comp – minimum SH/SG component to render.

  • max_comp – maximum SH/SG component to render, -1=last. Set min_comp = max_comp to render a particular component. Default means all.

  • density_softplus – if true, applies \(\log(1 + \exp(sigma - 1))\). Mind the shift -1! (from mip-NeRF). Please note softplus will NOT be compatible with volrend, please pre-apply it .

  • rgb_padding – to avoid oversaturating the sigmoid, applies * (1 + 2 * rgb_padding) - rgb_padding to colors after sigmoid (from mip-NeRF). Please note the padding will NOT be compatible with volrend, although most likely the effect is very small. 0.001 is a reasonable value to try.

forward(rays: svox.renderer.Rays, cuda=True, fast=False)

Render a batch of rays. Differentiable.

Parameters

  • rays – namedtuple Rays of origins (B, 3), dirs (B, 3), viewdirs (B, 3)

  • cuda – whether to use CUDA kernel if available. If false, uses only PyTorch version. Only True supported right now

  • fast – if True, enables faster evaluation, potentially leading to some loss of accuracy.

Returns

(B, rgb_dim). Where rgb_dim is tree.data_dim - 1 if data_format.format == DataFormat.RGBA or (tree.data_dim - 1) / tree.data_format.basis_dim else.

render_persp(c2w, width=800, height=800, fx=1111.111, fy=None, cuda=True, fast=False)

Render a perspective image. Differentiable.

Parameters

  • c2w – torch.Tensor (3, 4) or (4, 4) camera pose matrix (c2w)

  • width – int output image width

  • height – int output image height

  • fx – float output image focal length (x)

  • fy – float output image focal length (y), if not specified uses fx

  • cuda – whether to use CUDA kernel if available. If false, uses only PyTorch version. Only True supported right now

  • fast – if True, enables faster evaluation, potentially leading to some loss of accuracy.

Returns

(height, width, rgb_dim) where rgb_dim is tree.data_dim - 1 if data_format.format == DataFormat.RGBA or (tree.data_dim - 1) / tree.data_format.basis_dim else.