Given linear index, return (i,j,k) Args: ind: # (*, n) long, x + y * sx + z * sx * sy (*, n) long, z + y * sz + x * sy * sz size: (*, 3) long Returns: (*, n, 3) long
(
ind: torch.Tensor, # (*, n)
size: torch.Tensor, # (*, 3)
)
| 39 | |
| 40 | |
| 41 | def ind2sub( |
| 42 | ind: torch.Tensor, # (*, n) |
| 43 | size: torch.Tensor, # (*, 3) |
| 44 | ) -> torch.Tensor: |
| 45 | """ |
| 46 | Given linear index, return (i,j,k) |
| 47 | |
| 48 | Args: |
| 49 | ind: |
| 50 | # (*, n) long, x + y * sx + z * sx * sy |
| 51 | (*, n) long, z + y * sz + x * sy * sz |
| 52 | size: |
| 53 | (*, 3) long |
| 54 | |
| 55 | Returns: |
| 56 | (*, n, 3) long |
| 57 | """ |
| 58 | |
| 59 | # xy_size = size[..., 0:1] * size[..., 1:2] # (*, 1) |
| 60 | # zs = torch.div(ind, xy_size, rounding_mode='floor') # (*, n) |
| 61 | # ind = ind - zs * xy_size # (*, n) |
| 62 | # ys = torch.div(ind, size[..., 0:1], rounding_mode='floor') # (*, n) |
| 63 | # xs = ind - ys * size[..., 0:1] # (*, n) |
| 64 | |
| 65 | yz_size = size[..., 1:2] * size[..., 2:3] # (*, 1) |
| 66 | xs = torch.div(ind, yz_size, rounding_mode='floor') # (*, n) |
| 67 | ind = ind - xs * yz_size # (*, n) |
| 68 | ys = torch.div(ind, size[..., 2:3], rounding_mode='floor') # (*, n) |
| 69 | zs = ind - ys * size[..., 2:3] # (*, n) |
| 70 | |
| 71 | idx = torch.stack((xs, ys, zs), dim=-1) # (*, n, 3) |
| 72 | return idx |
| 73 | |
| 74 | |
| 75 | def get_grid_idx( |
nothing calls this directly
no outgoing calls
no test coverage detected