r"""Encodes :attr:`x`, :attr:`y`, :attr:`z` coordinates to the shuffled keys based on pre-computed look up tables. The speed of this function is much faster than the method based on for-loop. Args: x (torch.Tensor): The x coordinate. y (torch.Tensor): The y coordinate.
(
x: torch.Tensor,
y: torch.Tensor,
z: torch.Tensor,
b: Optional[Union[torch.Tensor, int]] = None,
depth: int = 16,
)
| 64 | |
| 65 | |
| 66 | def xyz2key( |
| 67 | x: torch.Tensor, |
| 68 | y: torch.Tensor, |
| 69 | z: torch.Tensor, |
| 70 | b: Optional[Union[torch.Tensor, int]] = None, |
| 71 | depth: int = 16, |
| 72 | ): |
| 73 | r"""Encodes :attr:`x`, :attr:`y`, :attr:`z` coordinates to the shuffled keys |
| 74 | based on pre-computed look up tables. The speed of this function is much |
| 75 | faster than the method based on for-loop. |
| 76 | |
| 77 | Args: |
| 78 | x (torch.Tensor): The x coordinate. |
| 79 | y (torch.Tensor): The y coordinate. |
| 80 | z (torch.Tensor): The z coordinate. |
| 81 | b (torch.Tensor or int): The batch index of the coordinates, and should be |
| 82 | smaller than 32768. If :attr:`b` is :obj:`torch.Tensor`, the size of |
| 83 | :attr:`b` must be the same as :attr:`x`, :attr:`y`, and :attr:`z`. |
| 84 | depth (int): The depth of the shuffled key, and must be smaller than 17 (< 17). |
| 85 | """ |
| 86 | |
| 87 | EX, EY, EZ = _key_lut.encode_lut(x.device) |
| 88 | x, y, z = x.long(), y.long(), z.long() |
| 89 | |
| 90 | mask = 255 if depth > 8 else (1 << depth) - 1 |
| 91 | key = EX[x & mask] | EY[y & mask] | EZ[z & mask] |
| 92 | if depth > 8: |
| 93 | mask = (1 << (depth - 8)) - 1 |
| 94 | key16 = EX[(x >> 8) & mask] | EY[(y >> 8) & mask] | EZ[(z >> 8) & mask] |
| 95 | key = key16 << 24 | key |
| 96 | |
| 97 | if b is not None: |
| 98 | b = b.long() |
| 99 | key = b << 48 | key |
| 100 | |
| 101 | return key |
| 102 | |
| 103 | |
| 104 | def key2xyz(key: torch.Tensor, depth: int = 16): |
nothing calls this directly
no test coverage detected