How to efficiently find correspondences between two point sets without nested for loop in Pytorch?

username: ZhengdiYu
cooked: <p>Hello,</p>
<p>I now have two point sets (tensor) A and B that shape like</p>
<p>A.size() &gt;&gt;(50, 3) , example: [ [0, 0, 0], [0, 1, 2], …, [1, 1, 1]]</p>
<p>B.size() &gt;&gt;(10, 3)</p>
<p>where the first dimension stands for number of points and the second dim stands for coordinates (x,y,z)</p>
<p>To some extent, the question could also be simplified into " Finding common elements between two tensors ". Is there a quick way to do this without nested loop like :</p>
<pre><code class="lang-auto">def is_corr(a, b):
    corr = []
    index = []
    for idx, ele_b in enumerate(b):
        cur_b = ele_b
        for ids, ele_a in enumerate(a):
            if torch.equal(ele_a, cur_b):
                corr.append(ele_a)
                index.append(ids)

    return corr, index
</code></pre>
<p>Thanks</p>

username: Nikronic
cooked: <p>Hi,</p>
<p>I do not know there is any builtin function for this, but the best way I can think of to loop over elements in <code>b</code> and use it as index in <code>a</code>.<br>
I know this solution sounds weird but may be we can optimize it later!</p>
<pre><code class="lang-auto">indices = [torch.sum(a==b[idx], dim=1) == len(b[idx]) for idx in range(len(b))]
</code></pre>
<p><code>a==b[idx]</code> will return <code>[x, x, x]</code> in your shape where x is boolean. We looking for equality so sum should be equal to length of index tensor <code>b</code>.</p>
<p>Bests</p>

username: ZhengdiYu
cooked: <p>Hi,</p>
<p>Thanks for your quick reply, but it seems a bit tricky for me to understand, and I write a solution like:</p>
<pre><code class="lang-auto">space = torch.zeros((100, 100, 100)) # total coordinates space
pointset1 = torch.randint(0,10,(10240, 3)) # random sample
pointset2 = torch.randint(0,10,(256, 3))   # random sample

space[pointset1[:, 0], pointset1[:, 1], pointset1[:, 2]] += 1 
space[pointset2[:, 0], pointset2[:, 1], pointset2[:, 2]] += 1

corrs = space &gt; 1  ## Get intersaction
del space
print(torch.nonzero(corrs))
</code></pre>
<p>What do you think of it ?</p>

username: Nikronic
cooked: <p>Your idea is interesting!<br>
The coordinates you calculate are correct in <code>space</code>. But there is a few problems.</p>
<ol>
<li>How can you remap <code>corrs</code> obtained from <code>space</code> to values in <code>pointset1</code>?</li>
<li>It does not count occurrences more than once.<br>
About second problem,</li>
</ol>
<aside class="quote no-group" data-username="ZhengdiYu" data-post="3" data-topic="83544">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/z/3da27b/48.png" class="avatar"> ZhengdiYu:</div>
<blockquote>
<p><code>space[pointset1[:, 0], pointset1[:, 1], pointset1[:, 2]] += 1 </code></p>
</blockquote>
</aside>
<p>This line will not +1 for duplicate indices. For instance, if you define <code>pointset1</code> all ones, then only index <code>[1, 1, 1]</code> in <code>space</code> will be 1. But we expect it to be <code>len(pointset1)</code>.</p>
<p>Try this test:</p>
<pre><code class="lang-auto">pointset1 = torch.tensor([[1, 0, 1],
        [0, 0, 0],
        [1, 0, 0], 
        [1, 0, 0], 
        [0, 1, 1], 
        [0, 1, 1], 
        [0, 1, 0], 
        [1, 1, 1],
        [1, 0, 1],
        [1, 1, 1]])
pointset2 = torch.tensor([[1, 1, 1], [1, 0, 0]])
</code></pre>
<p>A way like computing histogram would help but still it is not possible to ocmpute it without for loop.</p>

username: ZhengdiYu
cooked: <p>Thank you! And I found my solution inapplicable  for my real problem. The thought of building a space map only works when the coordinates are all integer. However, in real problem, coordinates are float type like [0.5678, 1.2345, 1.3678] …</p>

username: Nikronic
cooked: <p>Actually, I was going to ask how to determine <code>space</code> size as it looked arbitrary. So, have you solved the issue?</p>

username: ZhengdiYu
cooked: <p>Yes, I think the code below could solve the problem:</p>
<p><code>((pointset1[:, None, :] == pointset2[None, ...]).all(dim=2)).nonzero()</code></p>
<p>this would return a tensor with shape of [ [index1_in_set1, index1_in_set2], [], [], …, []]</p>
<p>where each [index_in_set1 and index_in_set2] stands for a correspondence. And the two elements <code>index_in_set1</code> and <code>index_in_set2</code> stand for point’s index in the set1 and set 2 respectively.</p>