Add dijkstra shortest_path_lengths functions (#96)

* Add dijkstra shortest_path_lengths functions

This commit adds a new function for finding the shortest path lengths
from a node. This is built using the dijkstra's algorithm function from
the upstream petgraph lib, however because the data for the graphs being
traversed are PyObjects we needed to be able to handle python exceptions
gracefully (instead of panicking) which was not possible with the upstream
dijkstra function. So this function has been forked into retworkx and
modified to hand PyResult<> wrapped returns from the callables.

* Add tests for edge cases

* Improve input node detection for output dict

docs/source/api.rst

@@ -48,6 +48,8 @@ Algorithm Functions
    retworkx.digraph_all_simple_paths
    retworkx.graph_astar_shortest_path
    retworkx.digraph_astar_shortest_path
+   retworkx.graph_dijkstra_shortest_path_lengths
+   retworkx.digraph_dijkstra_shortest_path_lengths
    retworkx.graph_greedy_color
 
 Exceptions

src/dijkstra.rs

@@ -0,0 +1,143 @@
+// Licensed under the Apache License, Version 2.0 (the "License"); you may
+// not use this file except in compliance with the License. You may obtain
+// a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations
+// under the License.
+
+// This module is copied and forked from the upstream petgraph repository,
+// specifically:
+// https://github.com/petgraph/petgraph/blob/0.5.1/src/dijkstra.rs
+// this was necessary to modify the error handling to allow python callables
+// to be use for the input functions for edge_cost and return any exceptions
+// raised in Python instead of panicking
+
+use std::collections::BinaryHeap;
+use std::hash::Hash;
+
+use hashbrown::hash_map::Entry::{Occupied, Vacant};
+use hashbrown::HashMap;
+
+use petgraph::algo::Measure;
+use petgraph::visit::{EdgeRef, IntoEdges, VisitMap, Visitable};
+
+use pyo3::prelude::*;
+
+use crate::astar::MinScored;
+
+/// \[Generic\] Dijkstra's shortest path algorithm.
+///
+/// Compute the length of the shortest path from `start` to every reachable
+/// node.
+///
+/// The graph should be `Visitable` and implement `IntoEdges`. The function
+/// `edge_cost` should return the cost for a particular edge, which is used
+/// to compute path costs. Edge costs must be non-negative.
+///
+/// If `goal` is not `None`, then the algorithm terminates once the `goal` node's
+/// cost is calculated.
+///
+/// Returns a `HashMap` that maps `NodeId` to path cost.
+/// # Example
+/// ```rust
+/// use petgraph::Graph;
+/// use petgraph::algo::dijkstra;
+/// use petgraph::prelude::*;
+/// use std::collections::HashMap;
+///
+/// let mut graph : Graph<(),(),Directed>= Graph::new();
+/// let a = graph.add_node(()); // node with no weight
+/// let b = graph.add_node(());
+/// let c = graph.add_node(());
+/// let d = graph.add_node(());
+/// let e = graph.add_node(());
+/// let f = graph.add_node(());
+/// let g = graph.add_node(());
+/// let h = graph.add_node(());
+/// // z will be in another connected component
+/// let z = graph.add_node(());
+///
+/// graph.extend_with_edges(&[
+///     (a, b),
+///     (b, c),
+///     (c, d),
+///     (d, a),
+///     (e, f),
+///     (b, e),
+///     (f, g),
+///     (g, h),
+///     (h, e)
+/// ]);
+/// // a ----> b ----> e ----> f
+/// // ^       |       ^       |
+/// // |       v       |       v
+/// // d <---- c       h <---- g
+///
+/// let expected_res: HashMap<NodeIndex, usize> = [
+///      (a, 3),
+///      (b, 0),
+///      (c, 1),
+///      (d, 2),
+///      (e, 1),
+///      (f, 2),
+///      (g, 3),
+///      (h, 4)
+///     ].iter().cloned().collect();
+/// let res = dijkstra(&graph,b,None, |_| 1);
+/// assert_eq!(res, expected_res);
+/// // z is not inside res because there is not path from b to z.
+/// ```
+pub fn dijkstra<G, F, K>(
+    graph: G,
+    start: G::NodeId,
+    goal: Option<G::NodeId>,
+    mut edge_cost: F,
+) -> PyResult<HashMap<G::NodeId, K>>
+where
+    G: IntoEdges + Visitable,
+    G::NodeId: Eq + Hash,
+    F: FnMut(G::EdgeRef) -> PyResult<K>,
+    K: Measure + Copy,
+{
+    let mut visited = graph.visit_map();
+    let mut scores = HashMap::new();
+    let mut visit_next = BinaryHeap::new();
+    let zero_score = K::default();
+    scores.insert(start, zero_score);
+    visit_next.push(MinScored(zero_score, start));
+    while let Some(MinScored(node_score, node)) = visit_next.pop() {
+        if visited.is_visited(&node) {
+            continue;
+        }
+        if goal.as_ref() == Some(&node) {
+            break;
+        }
+        for edge in graph.edges(node) {
+            let next = edge.target();
+            if visited.is_visited(&next) {
+                continue;
+            }
+            let cost = edge_cost(edge)?;
+            let next_score = node_score + cost;
+            match scores.entry(next) {
+                Occupied(ent) => {
+                    if next_score < *ent.get() {
+                        *ent.into_mut() = next_score;
+                        visit_next.push(MinScored(next_score, next));
+                    }
+                }
+                Vacant(ent) => {
+                    ent.insert(next_score);
+                    visit_next.push(MinScored(next_score, next));
+                }
+            }
+        }
+        visited.visit(node);
+    }
+    Ok(scores)
+}

src/lib.rs

@@ -22,6 +22,7 @@ extern crate rand_pcg;
 mod astar;
 mod dag_isomorphism;
 mod digraph;
+mod dijkstra;
 mod graph;
 
 use std::cmp::{Ordering, Reverse};
@@ -901,6 +902,116 @@ fn digraph_all_simple_paths(
     Ok(result)
 }
 
+/// Compute the lengths of the shortest paths for a PyGraph object using
+/// Dijkstra's algorithm
+///
+/// :param PyGraph graph: The input graph to use
+/// :param int node: The node index to use as the source for finding the
+///     shortest paths from
+/// :param edge_cost_fn: A python callable that will take in 1 parameter, an
+///     edge's data object and will return a float that represents the
+///     cost/weight of that edge. It must be non-negative
+/// :param int goal: An optional node index to use as the end of the path.
+///     When specified the traversal will stop when the goal is reached and
+///     the output dictionary will only have a single entry with the length
+///     of the shortest path to the goal node.
+///
+/// :returns: A dictionary of the shortest paths from the provided node where
+///     the key is the node index of the end of the path and the value is the
+///     cost/sum of the weights of path
+/// :rtype: dict
+#[pyfunction]
+#[text_signature = "(graph, node, edge_cost_fn, /, goal=None)"]
+fn graph_dijkstra_shortest_path_lengths(
+    py: Python,
+    graph: &graph::PyGraph,
+    node: usize,
+    edge_cost_fn: PyObject,
+    goal: Option<usize>,
+) -> PyResult<PyObject> {
+    let edge_cost_callable = |a: &PyObject| -> PyResult<f64> {
+        let res = edge_cost_fn.call1(py, (a,))?;
+        let raw = res.to_object(py);
+        Ok(raw.extract(py)?)
+    };
+
+    let start = NodeIndex::new(node);
+    let goal_index: Option<NodeIndex> = match goal {
+        Some(node) => Some(NodeIndex::new(node)),
+        None => None,
+    };
+
+    let res = dijkstra::dijkstra(graph, start, goal_index, |e| {
+        edge_cost_callable(e.weight())
+    })?;
+    let out_dict = PyDict::new(py);
+    for (index, value) in res {
+        let int_index = index.index();
+        if int_index == node {
+            continue;
+        }
+        if (goal.is_some() && goal.unwrap() == int_index) || goal.is_none() {
+            out_dict.set_item(int_index, value)?;
+        }
+    }
+    Ok(out_dict.into())
+}
+
+/// Compute the lengths of the shortest paths for a PyDiGraph object using
+/// Dijkstra's algorithm
+///
+/// :param PyDiGraph graph: The input graph to use
+/// :param int node: The node index to use as the source for finding the
+///     shortest paths from
+/// :param edge_cost_fn: A python callable that will take in 1 parameter, an
+///     edge's data object and will return a float that represents the
+///     cost/weight of that edge. It must be non-negative
+/// :param int goal: An optional node index to use as the end of the path.
+///     When specified the traversal will stop when the goal is reached and
+///     the output dictionary will only have a single entry with the length
+///     of the shortest path to the goal node.
+///
+/// :returns: A dictionary of the shortest paths from the provided node where
+///     the key is the node index of the end of the path and the value is the
+///     cost/sum of the weights of path
+/// :rtype: dict
+#[pyfunction]
+#[text_signature = "(graph, node, edge_cost_fn, /, goal=None)"]
+fn digraph_dijkstra_shortest_path_lengths(
+    py: Python,
+    graph: &digraph::PyDiGraph,
+    node: usize,
+    edge_cost_fn: PyObject,
+    goal: Option<usize>,
+) -> PyResult<PyObject> {
+    let edge_cost_callable = |a: &PyObject| -> PyResult<f64> {
+        let res = edge_cost_fn.call1(py, (a,))?;
+        let raw = res.to_object(py);
+        Ok(raw.extract(py)?)
+    };
+
+    let start = NodeIndex::new(node);
+    let goal_index: Option<NodeIndex> = match goal {
+        Some(node) => Some(NodeIndex::new(node)),
+        None => None,
+    };
+
+    let res = dijkstra::dijkstra(graph, start, goal_index, |e| {
+        edge_cost_callable(e.weight())
+    })?;
+    let out_dict = PyDict::new(py);
+    for (index, value) in res {
+        let int_index = index.index();
+        if int_index == node {
+            continue;
+        }
+        if (goal.is_some() && goal.unwrap() == int_index) || goal.is_none() {
+            out_dict.set_item(int_index, value)?;
+        }
+    }
+    Ok(out_dict.into())
+}
+
 /// Compute the A* shortest path for a PyGraph
 ///
 /// :param PyGraph graph: The input graph to use
@@ -1240,6 +1351,8 @@ fn retworkx(py: Python<'_>, m: &PyModule) -> PyResult<()> {
     m.add_wrapped(wrap_pyfunction!(graph_adjacency_matrix))?;
     m.add_wrapped(wrap_pyfunction!(graph_all_simple_paths))?;
     m.add_wrapped(wrap_pyfunction!(digraph_all_simple_paths))?;
+    m.add_wrapped(wrap_pyfunction!(graph_dijkstra_shortest_path_lengths))?;
+    m.add_wrapped(wrap_pyfunction!(digraph_dijkstra_shortest_path_lengths))?;
     m.add_wrapped(wrap_pyfunction!(graph_astar_shortest_path))?;
     m.add_wrapped(wrap_pyfunction!(digraph_astar_shortest_path))?;
     m.add_wrapped(wrap_pyfunction!(graph_greedy_color))?;