Dijkstra algorithm in Pytorch

Ajinkya_Ambatwar: <p>Hello!<br>
I am working on 3D point clouds. I have the SPARSE MATRIX representation of the graph structure of the point cloud (like <code>csr_matrix</code> in <code>scipy.sparse</code>). I want to club together the points that are within certain threshold of the Geodesic distance (approximated by the path length in the graph) and process them together. TO FIND such points, I need to run some shortest path finding algorithm like <code>Dijkstra's</code>. In a nutshell, my idea is like this</p>
<pre><code class="lang-auto">1) Sample K points out of N points (that I could do using Furthest Point Sampling)
2) Find the nearest Geodesic neighbours (using BackProp supported algorithm) for each of K points
3) Process the neighbours for each point using some Neural Network
</code></pre>
<p>This will go in my <code>forward</code> function.<br>
Is there a way to implement Dijkstra’s(that support BackProp) in my functionality?</p>
<p>Or any other idea that I can implement?</p>
<p>Thank you!</p>

Ajinkya_Ambatwar: <p>I created my custom implementation for Dijkstra using priority queues as discussed <a href="https://cp-algorithms.com/graph/dijkstra_sparse.html" rel="noopener nofollow ugc">here</a><br>
For the same, I created a custom <code>PriorityQ</code> class using torch function as below</p>
<pre><code class="lang-auto">class priorityQ_torch(object):
    """Priority Q implelmentation in PyTorch

    Args:
        object ([torch.Tensor]): [The Queue to work on]
    """

    def __init__(self, val):
        self.q = torch.tensor([[val, 0]])
        # self.top = self.q[0]
        # self.isEmpty = self.q.shape[0] == 0

    def push(self, x):
        """Pushes x to q based on weightvalue in x. Maintains ascending order

        Args:
            q ([torch.Tensor]): [The tensor queue arranged in ascending order of weight value]
            x ([torch.Tensor]): [[index, weight] tensor to be inserted]

        Returns:
            [torch.Tensor]: [The queue tensor after correct insertion]
        """
        if type(x) == np.ndarray:
            x = torch.tensor(x)
        if self.isEmpty():
            self.q = x
            self.q = torch.unsqueeze(self.q, dim=0)
            return
        idx = torch.searchsorted(self.q.T[1], x[1])
        print(idx)
        self.q = torch.vstack([self.q[0:idx], x, self.q[idx:]]).contiguous()

    def top(self):
        """Returns the top element from the queue

        Returns:
            [torch.Tensor]: [top element]
        """
        return self.q[0]

    def pop(self):
        """pops(without return) the highest priority element with the minimum weight

        Args:
            q ([torch.Tensor]): [The tensor queue arranged in ascending order of weight value]

        Returns:
            [torch.Tensor]: [highest priority element]
        """
        if self.isEmpty():
            print("Can Not Pop")
        self.q = self.q[1:]

    def isEmpty(self):
        """Checks is the priority queue is empty

        Args:
            q ([torch.Tensor]): [The tensor queue arranged in ascending order of weight value]

        Returns:
            [Bool] : [Returns True is empty]
        """
        return self.q.shape[0] == 0

</code></pre>
<p>Now dijkstra, with adjacency matrix(with graph weights as input)</p>
<pre><code class="lang-auto">def dijkstra(adj):
    n = adj.shape[0]
    distance_matrix = torch.zeros([n, n])
    for i in range(n):
        u = torch.zeros(n, dtype=torch.bool)
        d = np.inf * torch.ones(n)
        d[i] = 0
        q = priorityQ_torch(i)
        while not q.isEmpty():
            v, d_v = q.top()  # point and distance
            v = v.int()
            q.pop()
            if d_v != d[v]:
                continue
            for j, py in enumerate(adj[v]):
                if py == 0 and j != v:
                    continue
                else:
                    to = j
                    weight = py
                    if d[v] + py &lt; d[to]:
                        d[to] = d[v] + py
                        q.push(torch.Tensor([to, d[to]]))
        distance_matrix[i] = d
    return distance_matrix
</code></pre>
<p>Returns shortest path distance matrix for the graph points!</p>
<p>Hope this helps someone! <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=9" title=":slight_smile:" class="emoji" alt=":slight_smile:"></p>
<p>Thank you!</p>