test_spatial_scatter.py

from spatial import spatial_scatter, relative_positions
import torch

items = torch.tensor([
    [[1.0, -1.0], [5.0, -5.0], [3., 3.]],
    [[5., 5.], [7., -7.], [2., 2.]],
])
positions = torch.tensor([
    [[0., 0.1], [1., 0.], [10., 10.]],
    [[0., 0.1], [1., 0.], [99.5, 9.5]],
])

origin = torch.tensor([
    [[0., 0.], [0., 1.]],
    [[0., 0.], [100., 10.]],
])
direction = torch.tensor([
    [[0., 1.], [0., -1.]],
    [[0., 1.], [0., -1.]],
])

relpos = relative_positions(origin, direction, positions)
map = spatial_scatter(
    items,
    relpos,
    nray=8,
    nring=5,
    inner_radius=1,
)
print(map.size())
print(map)

assert((map == torch.tensor([
        [[[[ 0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.],
           [ 0.,  0.,  5.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  3.,  0.,  0.,  0.,  0.]],

          [[ 0.,  0.,  0.,  0., -1.,  0.,  0.,  0.],
           [ 0.,  0., -5.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  3.,  0.,  0.,  0.,  0.]]],


         [[[ 0.,  0.,  0.,  0.,  1.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  5.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  3.,  0.]],

          [[ 0.,  0.,  0.,  0., -1.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0., -5.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  3.,  0.]]]],



        [[[[ 0.,  0.,  0.,  0.,  5.,  0.,  0.,  0.],
           [ 0.,  0.,  7.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  2.,  0.,  0.,  0.,  0.,  0.]],

          [[ 0.,  0.,  0.,  0.,  5.,  0.,  0.,  0.],
           [ 0.,  0., -7.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  2.,  0.,  0.,  0.,  0.,  0.]]],


         [[[ 0.,  0.,  0.,  2.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0., 12.,  0.,  0.,  0.,  0.,  0.]],

          [[ 0.,  0.,  0.,  2.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0.,  0.,  0.,  0.,  0.,  0.,  0.],
           [ 0.,  0., -2.,  0.,  0.,  0.,  0.,  0.]]]]])
).all())