DeepCodeCraft/test_spatial_scatter.py at master · cswinter/DeepCodeCraft · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
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())