content: import torch
  import torch.nn as nn
  import torch.nn.functional as F
  
  
  DEFAULT_TRIVEC_CONFIG = {
      'dim': 8,
      'rank': 8,
  }
  
  DEFAULT_VOXEL_CONFIG = {
      'solid': False,
  }
  
  DEFAULT_DECOPOLY_CONFIG = {
      'degree': 8,
      'rank': 16,
  }
  
  
  class DfsOctree:
      """
      Sparse Voxel Octree (SVO) implementation for PyTorch.
      Using Depth-First Search (DFS) order to store the octree.
      DFS order suits rendering and ray tracing.
  
      The structure and data are separatedly stored.
      Structure is stored as a continuous array, each element is a 3*32 bits descriptor.
      |-----------------------------------------|
      |      0:3 bits      |      4:31 bits     |
      |      leaf num      |       unused       |
      |-----------------------------------------|
      |               0:31  bits                |
      |                child ptr                |
      |-----------------------------------------|
      |               0:31  bits                |
      |                data ptr                 |
      |-----------------------------------------|
      Each element represents a non-leaf node in the octree.
      The valid mask is used to indicate whether the children are valid.
      The leaf mask is used to indicate whether the children are leaf nodes.
      The child ptr is used to point to the first non-leaf child. Non-leaf children descriptors are stored continuously from the child ptr.
      The data ptr is used to point to the data of leaf children. Leaf children data are stored continuously from the data ptr.
  
      There are also auxiliary arrays to store the additional structural information to facilitate parallel processing.
        - Position: the position of the octree nodes.
        - Depth: the depth of the octree nodes.
  
      Args:
          depth (int): the depth of the octree.
      """
  
      def __init__(
              self,
              depth,
              aabb=[0,0,0,1,1,1],
              sh_degree=2,
              primitive='voxel',
              primitive_config={},
              device='cuda',
          ):
          self.max_depth = depth
          self.aabb = torch.tensor(aabb, dtype=torch.float32, device=device)
          self.device = device
          self.sh_degree = sh_degree
          self.active_sh_degree = sh_degree
          self.primitive = primitive
          self.primitive_config = primitive_config
  
          self.structure = torch.tensor([[8, 1, 0]], dtype=torch.int32, device=self.device)
          self.position = torch.zeros((8, 3), dtype=torch.float32, device=self.device)
          self.depth = torch.zeros((8, 1), dtype=torch.uint8, device=self.device)
          self.position[:, 0] = torch.tensor([0.25, 0.75, 0.25, 0.75, 0.25, 0.75, 0.25, 0.75], device=self.device)
          self.position[:, 1] = torch.tensor([0.25, 0.25, 0.75, 0.75, 0.25, 0.25, 0.75, 0.75], device=self.device)
          self.position[:, 2] = torch.tensor([0.25, 0.25, 0.25, 0.25, 0.75, 0.75, 0.75, 0.75], device=self.device)
          self.depth[:, 0] = 1
  
          self.data = ['position', 'depth']
          self.param_names = []
  
          if primitive == 'voxel':
              self.features_dc = torch.zeros((8, 1, 3), dtype=torch.float32, device=self.device)
              self.features_ac = torch.zeros((8, (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device)
              self.data += ['features_dc', 'features_ac']
              self.param_names += ['features_dc', 'features_ac']
              if not primitive_config.get('solid', False):
                  self.density = torch.zeros((8, 1), dtype=torch.float32, device=self.device)
                  self.data.append('density')
                  self.param_names.append('density')
          elif primitive == 'gaussian':
              self.features_dc = torch.zeros((8, 1, 3), dtype=torch.float32, device=self.device)
              self.features_ac = torch.zeros((8, (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device)
              self.opacity = torch.zeros((8, 1), dtype=torch.float32, device=self.device)
              self.data += ['features_dc', 'features_ac', 'opacity']
              self.param_names += ['features_dc', 'features_ac', 'opacity']
          elif primitive == 'trivec':
              self.trivec = torch.zeros((8, primitive_config['rank'], 3, primitive_config['dim']), dtype=torch.float32, device=self.device)
              self.density = torch.zeros((8, primitive_config['rank']), dtype=torch.float32, device=self.device)
              self.features_dc = torch.zeros((8, primitive_config['rank'], 1, 3), dtype=torch.float32, device=self.device)
              self.features_ac = torch.zeros((8, primitive_config['rank'], (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device)
              self.density_shift = 0
              self.data += ['trivec', 'density', 'features_dc', 'features_ac']
              self.param_names += ['trivec', 'density', 'features_dc', 'features_ac']
          elif primitive == 'decoupoly':
              self.decoupoly_V = torch.zeros((8, primitive_config['rank'], 3), dtype=torch.float32, device=self.device)
              self.decoupoly_g = torch.zeros((8, primitive_config['rank'], primitive_config['degree']), dtype=torch.float32, device=self.device)
              self.density = torch.zeros((8, primitive_config['rank']), dtype=torch.float32, device=self.device)
              self.features_dc = torch.zeros((8, primitive_config['rank'], 1, 3), dtype=torch.float32, device=self.device)
              self.features_ac = torch.zeros((8, primitive_config['rank'], (sh_degree+1)**2-1, 3), dtype=torch.float32, device=self.device)
              self.density_shift = 0
              self.data += ['decoupoly_V', 'decoupoly_g', 'density', 'features_dc', 'features_ac']
              self.param_names += ['decoupoly_V', 'decoupoly_g', 'density', 'features_dc', 'features_ac']
  
          self.setup_functions()
  
      def setup_functions(self):
          self.density_activation = (lambda x: torch.exp(x - 2)) if self.primitive != 'trivec' else (lambda x: x)
          self.opacity_activation = lambda x: torch.sigmoid(x - 6)
          self.inverse_opacity_activation = lambda x: torch.log(x / (1 - x)) + 6
          self.color_activation = lambda x: torch.sigmoid(x)
  
      @property
      def num_non_leaf_nodes(self):
          return self.structure.shape[0]
      
      @property
      def num_leaf_nodes(self):
          return self.depth.shape[0]
  
      @property
      def cur_depth(self):
          return self.depth.max().item()
      
      @property
      def occupancy(self):
          return self.num_leaf_nodes / 8 ** self.cur_depth
      
      @property
      def get_xyz(self):
          return self.position
  
      @property
      def get_depth(self):
          return self.depth
  
      @property
      def get_density(self):
          if self.primitive == 'voxel' and self.voxel_config['solid']:
              return torch.full((self.position.shape[0], 1), 1000, dtype=torch.float32, device=self.device)
          return self.density_activation(self.density)
      
      @property
      def get_opacity(self):
          return self.opacity_activation(self.density)
  
      @property
      def get_trivec(self):
          return self.trivec
  
      @property
      def get_decoupoly(self):
          return F.normalize(self.decoupoly_V, dim=-1), self.decoupoly_g
  
      @property
      def get_color(self):
          return self.color_activation(self.colors)
  
      @property
      def get_features(self):
          if self.sh_degree == 0:
              return self.features_dc
          return torch.cat([self.features_dc, self.features_ac], dim=-2)
  
      def state_dict(self):
          ret = {'structure': self.structure, 'position': self.position, 'depth': self.depth, 'sh_degree': self.sh_degree, 'active_sh_degree': self.active_sh_degree, 'trivec_config': self.trivec_config, 'voxel_config': self.voxel_config, 'primitive': self.primitive}
          if hasattr(self, 'density_shift'):
              ret['density_shift'] = self.density_shift
          for data in set(self.data + self.param_names):
              if not isinstance(getattr(self, data), nn.Module):
                  ret[data] = getattr(self, data)
              else:
                  ret[data] = getattr(self, data).state_dict()
          return ret
  
      def load_state_dict(self, state_dict):
          keys = list(set(self.data + self.param_names + list(state_dict.keys()) + ['structure', 'position', 'depth']))
          for key in keys:
              if key not in state_dict:
                  print(f"Warning: key {key} not found in the state_dict.")
                  continue
              try:
                  if not isinstance(getattr(self, key), nn.Module):
                      setattr(self, key, state_dict[key])
                  else:
                      getattr(self, key).load_state_dict(state_dict[key])
              except Exception as e:
                  print(e)
                  raise ValueError(f"Error loading key {key}.")
  
      def gather_from_leaf_children(self, data):
          """
          Gather the data from the leaf children.
  
          Args:
              data (torch.Tensor): the data to gather. The first dimension should be the number of leaf nodes.
          """
          leaf_cnt = self.structure[:, 0]
          leaf_cnt_masks = [leaf_cnt == i for i in range(1, 9)]
          ret = torch.zeros((self.num_non_leaf_nodes,), dtype=data.dtype, device=self.device)
          for i in range(8):
              if leaf_cnt_masks[i].sum() == 0:
                  continue
              start = self.structure[leaf_cnt_masks[i], 2]
              for j in range(i+1):
                  ret[leaf_cnt_masks[i]] += data[start + j]
          return ret
  
      def gather_from_non_leaf_children(self, data):
          """
          Gather the data from the non-leaf children.
  
          Args:
              data (torch.Tensor): the data to gather. The first dimension should be the number of leaf nodes.
          """
          non_leaf_cnt = 8 - self.structure[:, 0]
          non_leaf_cnt_masks = [non_leaf_cnt == i for i in range(1, 9)]
          ret = torch.zeros_like(data, device=self.device)
          for i in range(8):
              if non_leaf_cnt_masks[i].sum() == 0:
                  continue
              start = self.structure[non_leaf_cnt_masks[i], 1]
              for j in range(i+1):
                  ret[non_leaf_cnt_masks[i]] += data[start + j]
          return ret
  
      def structure_control(self, mask):
          """
          Control the structure of the octree.
  
          Args:
              mask (torch.Tensor): the mask to control the structure. 1 for subdivide, -1 for merge, 0 for keep.
          """
          # Dont subdivide when the depth is the maximum.
          mask[self.depth.squeeze() == self.max_depth] = torch.clamp_max(mask[self.depth.squeeze() == self.max_depth], 0)
          # Dont merge when the depth is the minimum.
          mask[self.depth.squeeze() == 1] = torch.clamp_min(mask[self.depth.squeeze() == 1], 0)
  
          # Gather control mask
          structre_ctrl = self.gather_from_leaf_children(mask)
          structre_ctrl[structre_ctrl==-8] = -1
  
          new_leaf_num = self.structure[:, 0].clone()
          # Modify the leaf num.
          structre_valid = structre_ctrl >= 0
          new_leaf_num[structre_valid] -= structre_ctrl[structre_valid]                               # Add the new nodes.
          structre_delete = structre_ctrl < 0
          merged_nodes = self.gather_from_non_leaf_children(structre_delete.int())
          new_leaf_num += merged_nodes                                                                # Delete the merged nodes.
  
          # Update the structure array to allocate new nodes.
          mem_offset = torch.zeros((self.num_non_leaf_nodes + 1,), dtype=torch.int32, device=self.device)
          mem_offset.index_add_(0, self.structure[structre_valid, 1], structre_ctrl[structre_valid])  # Add the new nodes.
          mem_offset[:-1] -= structre_delete.int()                                                    # Delete the merged nodes.
          new_structre_idx = torch.arange(0, self.num_non_leaf_nodes + 1, dtype=torch.int32, device=self.device) + mem_offset.cumsum(0)
          new_structure_length = new_structre_idx[-1].item()
          new_structre_idx = new_structre_idx[:-1]
          new_structure = torch.empty((new_structure_length, 3), dtype=torch.int32, device=self.device)
          new_structure[new_structre_idx[structre_valid], 0] = new_leaf_num[structre_valid]
  
          # Initialize the new nodes.
          new_node_mask = torch.ones((new_structure_length,), dtype=torch.bool, device=self.device)
          new_node_mask[new_structre_idx[structre_valid]] = False
          new_structure[new_node_mask, 0] = 8                                                         # Initialize to all leaf nodes.
          new_node_num = new_node_mask.sum().item()
  
          # Rebuild child ptr.
          non_leaf_cnt = 8 - new_structure[:, 0]
          new_child_ptr = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), non_leaf_cnt.cumsum(0)[:-1]])
          new_structure[:, 1] = new_child_ptr + 1
  
          # Rebuild data ptr with old data.
          leaf_cnt = torch.zeros((new_structure_length,), dtype=torch.int32, device=self.device)
          leaf_cnt.index_add_(0, new_structre_idx, self.structure[:, 0])
          old_data_ptr = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), leaf_cnt.cumsum(0)[:-1]])
  
          # Update the data array
          subdivide_mask = mask == 1
          merge_mask = mask == -1
          data_valid = ~(subdivide_mask | merge_mask)
          mem_offset = torch.zeros((self.num_leaf_nodes + 1,), dtype=torch.int32, device=self.device)
          mem_offset.index_add_(0, old_data_ptr[new_node_mask], torch.full((new_node_num,), 8, dtype=torch.int32, device=self.device))    # Add data array for new nodes
          mem_offset[:-1] -= subdivide_mask.int()                                                                                         # Delete data elements for subdivide nodes
          mem_offset[:-1] -= merge_mask.int()                                                                                             # Delete data elements for merge nodes
          mem_offset.index_add_(0, self.structure[structre_valid, 2], merged_nodes[structre_valid])                                       # Add data elements for merge nodes
          new_data_idx = torch.arange(0, self.num_leaf_nodes + 1, dtype=torch.int32, device=self.device) + mem_offset.cumsum(0)
          new_data_length = new_data_idx[-1].item()
          new_data_idx = new_data_idx[:-1]
          new_data = {data: torch.empty((new_data_length,) + getattr(self, data).shape[1:], dtype=getattr(self, data).dtype, device=self.device) for data in self.data}
          for data in self.data:
              new_data[data][new_data_idx[data_valid]] = getattr(self, data)[data_valid]
  
          # Rebuild data ptr
          leaf_cnt = new_structure[:, 0]
          new_data_ptr = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), leaf_cnt.cumsum(0)[:-1]])
          new_structure[:, 2] = new_data_ptr
  
          # Initialize the new data array
          ## For subdivide nodes
          if subdivide_mask.sum() > 0:
              subdivide_data_ptr = new_structure[new_node_mask, 2]
              for data in self.data:
                  for i in range(8):
                      if data == 'position':
                          offset = torch.tensor([i // 4, (i // 2) % 2, i % 2], dtype=torch.float32, device=self.device) - 0.5
                          scale = 2 ** (-1.0 - self.depth[subdivide_mask])
                          new_data['position'][subdivide_data_ptr + i] = self.position[subdivide_mask] + offset * scale
                      elif data == 'depth':
                          new_data['depth'][subdivide_data_ptr + i] = self.depth[subdivide_mask] + 1
                      elif data == 'opacity':
                          new_data['opacity'][subdivide_data_ptr + i] = self.inverse_opacity_activation(torch.sqrt(self.opacity_activation(self.opacity[subdivide_mask])))
                      elif data == 'trivec':
                          offset = torch.tensor([i // 4, (i // 2) % 2, i % 2], dtype=torch.float32, device=self.device) * 0.5
                          coord = (torch.linspace(0, 0.5, self.trivec.shape[-1], dtype=torch.float32, device=self.device)[None] + offset[:, None]).reshape(1, 3, self.trivec.shape[-1], 1)
                          axis = torch.linspace(0, 1, 3, dtype=torch.float32, device=self.device).reshape(1, 3, 1, 1).repeat(1, 1, self.trivec.shape[-1], 1)
                          coord = torch.stack([coord, axis], dim=3).reshape(1, 3, self.trivec.shape[-1], 2).expand(self.trivec[subdivide_mask].shape[0], -1, -1, -1) * 2 - 1
                          new_data['trivec'][subdivide_data_ptr + i] = F.grid_sample(self.trivec[subdivide_mask], coord, align_corners=True)
                      else:
                          new_data[data][subdivide_data_ptr + i] = getattr(self, data)[subdivide_mask]
          ## For merge nodes
          if merge_mask.sum() > 0:
              merge_data_ptr = torch.empty((merged_nodes.sum().item(),), dtype=torch.int32, device=self.device)
              merge_nodes_cumsum = torch.cat([torch.zeros((1,), dtype=torch.int32, device=self.device), merged_nodes.cumsum(0)[:-1]])
              for i in range(8):
                  merge_data_ptr[merge_nodes_cumsum[merged_nodes > i] + i] = new_structure[new_structre_idx[merged_nodes > i], 2] + i
              old_merge_data_ptr = self.structure[structre_delete, 2]
              for data in self.data:
                  if data == 'position':
                      scale = 2 ** (1.0 - self.depth[old_merge_data_ptr])
                      new_data['position'][merge_data_ptr] = ((self.position[old_merge_data_ptr] + 0.5) / scale).floor() * scale + 0.5 * scale - 0.5
                  elif data == 'depth':
                      new_data['depth'][merge_data_ptr] = self.depth[old_merge_data_ptr] - 1
                  elif data == 'opacity':
                      new_data['opacity'][subdivide_data_ptr + i] = self.inverse_opacity_activation(self.opacity_activation(self.opacity[subdivide_mask])**2)
                  elif data == 'trivec':
                      new_data['trivec'][merge_data_ptr] = self.trivec[old_merge_data_ptr]
                  else:
                      new_data[data][merge_data_ptr] = getattr(self, data)[old_merge_data_ptr]
  
          # Update the structure and data array
          self.structure = new_structure
          for data in self.data:
              setattr(self, data, new_data[data])
  
          # Save data array control temp variables
          self.data_rearrange_buffer = {
              'subdivide_mask': subdivide_mask,
              'merge_mask': merge_mask,
              'data_valid': data_valid,
              'new_data_idx': new_data_idx,
              'new_data_length': new_data_length,
              'new_data': new_data
          } 
  