Hide keyboard shortcuts

Hot-keys on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

1# 

2# (2018) Modified by Kwant Authors 

3# 

4# Modifications 

5# ============= 

6# Merged and modified from scipy/sparse/csgraph/_shortest_path.pyx 

7# 

8# All shortest path algorithms except for Dijkstra's removed. 

9# Implementation of Dijkstra's algorithm modified to allow for specific 

10# use-cases required by Flux. The changes are documented in the docstring. 

11# 

12# 

13# Copyright (c) 2001, 2002 Enthought, Inc. 

14# All rights reserved. 

15# 

16# Copyright (c) 2003-2017 SciPy Developers. 

17# All rights reserved. 

18# 

19# Copyright (c) 2011 Jake Vanderplas <vanderplas@astro.washington.edu> 

20# All rights reserved. 

21# 

22# Redistribution and use in source and binary forms, with or without 

23# modification, are permitted provided that the following conditions are met: 

24# 

25# a. Redistributions of source code must retain the above copyright notice, 

26# this list of conditions and the following disclaimer. 

27# b. Redistributions in binary form must reproduce the above copyright 

28# notice, this list of conditions and the following disclaimer in the 

29# documentation and/or other materials provided with the distribution. 

30# c. Neither the name of Enthought nor the names of the SciPy Developers 

31# may be used to endorse or promote products derived from this software 

32# without specific prior written permission. 

33# 

34# 

35# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" 

36# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE 

37# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE 

38# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDERS OR CONTRIBUTORS 

39# BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, 

40# OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF 

41# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS 

42# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN 

43# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) 

44# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF 

45# THE POSSIBILITY OF SUCH DAMAGE. 

46  

47import numpy as np 

48cimport numpy as np 

49from numpy.math cimport INFINITY as inf 

50  

51from libc.stdlib cimport malloc, free 

52from libc.string cimport memset 

53  

54ctypedef np.float64_t DTYPE_t 

55ctypedef np.int32_t ITYPE_t 

56ITYPE = np.int32 

57  

58  

59def dijkstra_directed( 

60 object graph, 

61 ITYPE_t[:] sources, 

62 ITYPE_t[:] targets, 

63 bint return_paths=True, 

64 bint return_predecessors=False): 

65 """Modified directed Dijkstra algorithm. 

66  

67 Edges with infinite weight are treated as if they do not exist. 

68  

69 The shortest paths between edge pairs 'zip(sources, targets)' 

70 are found. 

71  

72 If 'len(sources) == 1' then this routine can be used to 

73 compute the one-to-all distances by passing an integer 

74 greater than any node in 'graph' as the 'target'. 

75 In this case 'return_predecessors' may be specified, and 

76 the predecessor matrix will also be returned. 

77  

78 'return_predecessors' and 'return_paths' are mutually exclusive. 

79  

80 Returns 

81 ------- 

82 if return_paths: 

83 (paths, path_lengths) 

84 elif return_predecessors: 

85 (path_lengths, predecessors) 

86 else: 

87 path_lengths 

88  

89 """ 

90# Implementation of Dijkstra's algorithm modified to allow for 

91# early exit (when target is found) and return the path from source 

92# to target, rather than the whole predecessor matrix. In addition 

93# graph edges with infinite weight are treated as if they do not exist. 

94  

95 cdef ITYPE_t[:] csr_indices = graph.indices, csr_indptr = graph.indptr 

96 cdef DTYPE_t[:] csr_weights = graph.data 

97  

98 # This implementation of Dijkstra's algorithm is very tightly coupled 

99 # to our use-case in 'flux', so we allow ourselves to assert 

100 assert sources.shape[0] == targets.shape[0] 

101 assert graph.shape[0] == graph.shape[1] 

102 assert not (return_predecessors and return_paths) 

103 assert not (return_predecessors and (sources.shape[0] > 1)) 

104  

105 cdef unsigned int num_links = sources.shape[0], num_nodes = graph.shape[0] 

106  

107 cdef unsigned int i, k, j_source, j_target, j_current 

108 cdef ITYPE_t j 

109  

110 cdef DTYPE_t next_val 

111  

112 cdef FibonacciHeap heap 

113 cdef FibonacciNode *v 

114 cdef FibonacciNode *current_node 

115 cdef FibonacciNode* nodes = <FibonacciNode*> malloc(num_nodes * sizeof(FibonacciNode)) 

116  

117 cdef ITYPE_t[:] pred = np.empty((num_nodes,), dtype=ITYPE) 

118  

119  

120 # outputs 

121 cdef DTYPE_t[:] path_lengths = np.zeros((num_links,), float) 

122 cdef list paths 

123 if return_paths: 

124 paths = [] 

125  

126 for i in range(num_links): 

127 j_source = sources[i] 

128 j_target = targets[i] 

129  

130 for k in range(num_nodes): 

131 initialize_node(&nodes[k], k) 

132 pred[:] = -1 # only useful for debugging 

133  

134 heap.min_node = NULL 

135 insert_node(&heap, &nodes[j_source]) 

136  

137 while heap.min_node: 

138 v = remove_min(&heap) 

139 v.state = SCANNED 

140  

141 if v.index == j_target: 

142 path_lengths[i] = v.val 

143 if return_paths: 

144 paths.append(_calculate_path(pred, j_source, j_target)) 

145 break # next iteration of outer 'for' loop 

146  

147 for j in range(csr_indptr[v.index], csr_indptr[v.index + 1]): 

148 if csr_weights[j] == inf: 

149 # Treat infinite weight links as missing 

150 continue 

151 j_current = csr_indices[j] 

152 current_node = &nodes[j_current] 

153 if current_node.state != SCANNED: 

154 next_val = v.val + csr_weights[j] 

155 if current_node.state == NOT_IN_HEAP: 

156 current_node.state = IN_HEAP 

157 current_node.val = next_val 

158 insert_node(&heap, current_node) 

159 pred[j_current] = v.index 

160 elif current_node.val > next_val: 

161 decrease_val(&heap, current_node, 

162 next_val) 

163 pred[j_current] = v.index 

164  

165 free(nodes) 

166 if return_paths: 

167 return paths, path_lengths 

168 elif return_predecessors: 

169 return path_lengths, pred 

170 else: 

171 return path_lengths 

172  

173  

174cdef list _calculate_path(ITYPE_t[:] pred, int j_source, int j_target): 

175 visited = [] 

176 cdef int node = j_target 

177 while node != j_source: 

178 visited.append(node) 

179 node = pred[node] 

180 visited.append(j_source) 

181 return visited 

182  

183  

184###################################################################### 

185# FibonacciNode structure 

186# This structure and the operations on it are the nodes of the 

187# Fibonacci heap. 

188# 

189cdef enum FibonacciState: 

190 SCANNED 

191 NOT_IN_HEAP 

192 IN_HEAP 

193  

194  

195cdef struct FibonacciNode: 

196 unsigned int index 

197 unsigned int rank 

198 FibonacciState state 

199 DTYPE_t val 

200 FibonacciNode* parent 

201 FibonacciNode* left_sibling 

202 FibonacciNode* right_sibling 

203 FibonacciNode* children 

204  

205  

206cdef void initialize_node(FibonacciNode* node, 

207 unsigned int index, 

208 DTYPE_t val=0): 

209 # Assumptions: - node is a valid pointer 

210 # - node is not currently part of a heap 

211 node.index = index 

212 node.val = val 

213 node.rank = 0 

214 node.state = NOT_IN_HEAP 

215  

216 node.parent = NULL 

217 node.left_sibling = NULL 

218 node.right_sibling = NULL 

219 node.children = NULL 

220  

221  

222cdef FibonacciNode* rightmost_sibling(FibonacciNode* node): 

223 # Assumptions: - node is a valid pointer 

224 cdef FibonacciNode* temp = node 

225 while(temp.right_sibling): 

226 temp = temp.right_sibling 

227 return temp 

228  

229  

230cdef FibonacciNode* leftmost_sibling(FibonacciNode* node): 

231 # Assumptions: - node is a valid pointer 

232 cdef FibonacciNode* temp = node 

233 while(temp.left_sibling): 

234 temp = temp.left_sibling 

235 return temp 

236  

237  

238cdef void add_child(FibonacciNode* node, FibonacciNode* new_child): 

239 # Assumptions: - node is a valid pointer 

240 # - new_child is a valid pointer 

241 # - new_child is not the sibling or child of another node 

242 new_child.parent = node 

243  

244 if node.children: 

245 add_sibling(node.children, new_child) 

246 else: 

247 node.children = new_child 

248 new_child.right_sibling = NULL 

249 new_child.left_sibling = NULL 

250 node.rank = 1 

251  

252  

253cdef void add_sibling(FibonacciNode* node, FibonacciNode* new_sibling): 

254 # Assumptions: - node is a valid pointer 

255 # - new_sibling is a valid pointer 

256 # - new_sibling is not the child or sibling of another node 

257 cdef FibonacciNode* temp = rightmost_sibling(node) 

258 temp.right_sibling = new_sibling 

259 new_sibling.left_sibling = temp 

260 new_sibling.right_sibling = NULL 

261 new_sibling.parent = node.parent 

262 if new_sibling.parent: 

263 new_sibling.parent.rank += 1 

264  

265  

266cdef void remove(FibonacciNode* node): 

267 # Assumptions: - node is a valid pointer 

268 if node.parent: 

269 node.parent.rank -= 1 

270 if node.left_sibling: 

271 node.parent.children = node.left_sibling 

272 elif node.right_sibling: 

273 node.parent.children = node.right_sibling 

274 else: 

275 node.parent.children = NULL 

276  

277 if node.left_sibling: 

278 node.left_sibling.right_sibling = node.right_sibling 

279 if node.right_sibling: 

280 node.right_sibling.left_sibling = node.left_sibling 

281  

282 node.left_sibling = NULL 

283 node.right_sibling = NULL 

284 node.parent = NULL 

285  

286  

287###################################################################### 

288# FibonacciHeap structure 

289# This structure and operations on it use the FibonacciNode 

290# routines to implement a Fibonacci heap 

291  

292ctypedef FibonacciNode* pFibonacciNode 

293  

294  

295cdef struct FibonacciHeap: 

296 FibonacciNode* min_node 

297 pFibonacciNode[100] roots_by_rank # maximum number of nodes is ~2^100. 

298  

299  

300cdef void insert_node(FibonacciHeap* heap, 

301 FibonacciNode* node): 

302 # Assumptions: - heap is a valid pointer 

303 # - node is a valid pointer 

304 # - node is not the child or sibling of another node 

305 if heap.min_node: 

306 add_sibling(heap.min_node, node) 

307 if node.val < heap.min_node.val: 

308 heap.min_node = node 

309 else: 

310 heap.min_node = node 

311  

312  

313cdef void decrease_val(FibonacciHeap* heap, 

314 FibonacciNode* node, 

315 DTYPE_t newval): 

316 # Assumptions: - heap is a valid pointer 

317 # - newval <= node.val 

318 # - node is a valid pointer 

319 # - node is not the child or sibling of another node 

320 # - node is in the heap 

321 node.val = newval 

322 if node.parent and (node.parent.val >= newval): 

323 remove(node) 

324 insert_node(heap, node) 

325 elif heap.min_node.val > node.val: 

326 heap.min_node = node 

327  

328  

329cdef void link(FibonacciHeap* heap, FibonacciNode* node): 

330 # Assumptions: - heap is a valid pointer 

331 # - node is a valid pointer 

332 # - node is already within heap 

333  

334 cdef FibonacciNode *linknode 

335 cdef FibonacciNode *parent 

336 cdef FibonacciNode *child 

337  

338 if heap.roots_by_rank[node.rank] == NULL: 

339 heap.roots_by_rank[node.rank] = node 

340 else: 

341 linknode = heap.roots_by_rank[node.rank] 

342 heap.roots_by_rank[node.rank] = NULL 

343  

344 if node.val < linknode.val or node == heap.min_node: 

345 remove(linknode) 

346 add_child(node, linknode) 

347 link(heap, node) 

348 else: 

349 remove(node) 

350 add_child(linknode, node) 

351 link(heap, linknode) 

352  

353  

354cdef FibonacciNode* remove_min(FibonacciHeap* heap): 

355 # Assumptions: - heap is a valid pointer 

356 # - heap.min_node is a valid pointer 

357 cdef FibonacciNode *temp 

358 cdef FibonacciNode *temp_right 

359 cdef FibonacciNode *out 

360 cdef unsigned int i 

361  

362 # make all min_node children into root nodes 

363 if heap.min_node.children: 

364 temp = leftmost_sibling(heap.min_node.children) 

365 temp_right = NULL 

366  

367 while temp: 

368 temp_right = temp.right_sibling 

369 remove(temp) 

370 add_sibling(heap.min_node, temp) 

371 temp = temp_right 

372  

373 heap.min_node.children = NULL 

374  

375 # choose a root node other than min_node 

376 temp = leftmost_sibling(heap.min_node) 

377 if temp == heap.min_node: 

378 if heap.min_node.right_sibling: 

379 temp = heap.min_node.right_sibling 

380 else: 

381 out = heap.min_node 

382 heap.min_node = NULL 

383 return out 

384  

385 # remove min_node, and point heap to the new min 

386 out = heap.min_node 

387 remove(heap.min_node) 

388 heap.min_node = temp 

389  

390 # re-link the heap 

391 for i in range(100): 

392 heap.roots_by_rank[i] = NULL 

393  

394 while temp: 

395 if temp.val < heap.min_node.val: 

396 heap.min_node = temp 

397 temp_right = temp.right_sibling 

398 link(heap, temp) 

399 temp = temp_right 

400  

401 return out