repo_description: Official PyTorch implementation of Superpoint Transformer [ICCV'23], SuperCluster [3DV'24 Oral], and EZ-SP [ICRA'26]
repo_topics: 3d 3dv2024 deep-learning efficient fast graph-clustering hierarchical iccv2023 lightweight panoptic-segmentation partition partitioning point-cloud pytorch semantic-segmentation superpoint transformer
code: import torch
  
  
  __all__ = ['fourier_position_encoder']
  
  
  def fourier_position_encoder(pos, dim, f_min=1e-1, f_max=1e1):
      """
      Heuristic: keeping ```f_min = 1 / f_max``` ensures that roughly 50%
      of the encoding dimensions are untouched and free to use. This is
      important when the positional encoding is added to learned feature
      embeddings. If the positional encoding uses too much of the encoding
      dimensions, it may be detrimental for the embeddings.
  
      The default `f_min` and `f_max` values are set so as to ensure
      a '~50% use of the encoding dimensions' and a '~1e-3 precision in
      the position encoding if pos is 1D'.
  
      :param pos: [M, M] Tensor
          Positions are expected to be in [-1, 1]
      :param dim: int
          Number of encoding dimensions, size of the encoding space. Note
          that increasing this is NOT the most direct way of improving
          spatial encoding precision or compactness. See `f_min` and
          `f_max` instead
      :param f_min: float
          Lower bound for the frequency range. Rules how much 'room' the
          positional encodings leave in the encoding space for additive
          embeddings
      :param f_max: float
          Upper bound for the frequency range. Rules how precise the
          encoding can be. Increase this if you need to capture finer
          spatial details
      :return:
      """
      assert pos.abs().max() <= 1, "Positions must be in [-1, 1]"
      assert 1 <= pos.dim() <= 2, "Positions must be a 1D or 2D tensor"
  
      # We preferably operate 2D tensors
      if pos.dim() == 1:
          pos = pos.view(-1, 1)
  
      # Make sure M divides dim
      N, M = pos.shape
      D = dim // M
      # assert dim % M == 0, "`dim` must be a multiple of the number of input spatial dimensions"
      # assert D % 2 == 0, "`dim / M` must be a even number"
  
      # To avoid uncomfortable border effects with -1 and +1 coordinates
      # having the same (or very close) encodings, we convert [-1, 1]
      # coordinates to [-π/2, π/2] for safety
      pos = pos * torch.pi / 2
  
      # Compute frequencies on a logarithmic range from f_min to f_max
      device = pos.device
      f_min = torch.tensor([f_min], device=device)
      f_max = torch.tensor([f_max], device=device)
      w = torch.logspace(f_max.log(), f_min.log(), D, device=device)
  
      # Compute sine and cosine encodings
      pos_enc = pos.view(N, M, 1) * w.view(1, -1)
      pos_enc[:, :, ::2] = pos_enc[:, :, ::2].cos()
      pos_enc[:, :, 1::2] = pos_enc[:, :, 1::2].sin()
      pos_enc = pos_enc.view(N, -1)
  
      # In case dim is not a multiple of 2 * M, we pad missing dimensions
      # with zeros
      if pos_enc.shape[1] < dim:
          zeros = torch.zeros(N, dim - pos_enc.shape[1], device=device)
          pos_enc = torch.hstack((pos_enc, zeros))
  
      return pos_enc
  