allpy: 597c0d0e9124 allpy/graph.py

allpy

view allpy/graph.py @ 529:597c0d0e9124

Chaned sequence name format in extract_pfam.py to ">PDBID_PDBCHAIN ORIGINAL_NAME ORIGINAL_DESCRIPTION LIST_OF_ASSOCIATED_PDB_FRAGMENTS"

author	Daniil Alexeyevsky <dendik@kodomo.fbb.msu.ru>
date	Sat, 26 Feb 2011 18:59:48 +0300
parents	c755fb38896c
children	1fb830d7517e

line source

1 from datetime import datetime, timedelta

2 from copy import copy

4 class TimeoutError(Exception):

5 pass

7 class Graph(dict):

8 """ Undirected weighted graph

10 graph[vertex1][vertex2] = weight

11 * vertex1 != vertex2 (no such edges)

12 * weight is float in (0, 1] or 1 (if all edges are equal)

13 * symmetrical

15 Features:

16 * this graph cannot contain vertices without any edges

18 >>> g = Graph()

19 >>> g.set_edge(1, 2)

20 >>> g.fast_cliques()

21 Fast algorithm started

22 [frozenset([1, 2])]

23 >>> g.set_edge(2, 3)

24 >>> g.set_edge(1, 3)

25 >>> g.fast_cliques()

26 Fast algorithm started

27 [frozenset([1, 2, 3])]

28 >>> g.bron_kerbosh()

29 Bron and Kerbosh algorithm started

30 [frozenset([1, 2, 3])]

31 >>> g.cliques()

32 Bron and Kerbosh algorithm started

33 [frozenset([1, 2, 3])]

35 """

37 def copy(self):

38 cls = self.__class__

39 graph = cls()

40 for x, neighbours_dict in self.items():

41 for y, weight in neighbours_dict.items():

42 graph.set_edge(x, y, weight)

43 return graph

45 def set_edge(self, v1, v2, weight=1):

46 assert v1 is not v2

47 self.setdefault(v1, {})

48 self[v1][v2] = weight

49 self.setdefault(v2, {})

50 self[v2][v1] = weight

52 def drop_vertex(self, vertex):

53 """ Remove vertex and all involved edges """

54 for v in self[vertex]:

55 del self[v][vertex]

56 if not self[v]:

57 del self[v]

58 del self[vertex]

60 def vertex_weight(self, vertex):

61 """ Returns sum of weights of all connections of this vertex """

62 return sum(self[vertex].values())

64 def vertex_index(self, vertex):

65 return len(self[vertex])

67 def add_vertex(self, vertex, parent_graph):

68 """ Add vertex and corresponding edges from parent_graph

70 Added edges should be contained in self graph

71 (takes care of hanging edges)

72 """

73 for v in parent_graph[vertex]:

74 if v in self:

75 self.set_edge(vertex, v, parent_graph[vertex][v])

77 def bron_kerbosh(self, timeout=-1, minsize=1):

78 """ Bron and Kerboch algorithm implementation

80 returns list of cliques

81 clique is frozenset

82 if timeout=-1, it means infinity

83 if timeout has happened, raises TimeoutError

84 """

85 if timeout == 0:

86 raise TimeoutError()

87 print 'Bron and Kerbosh algorithm started'

88 cliques = []

90 class ExtendCall(object):

91 def __init__(call, candidates, used, compsub):

92 call.candidates = candidates

93 call.used = used

94 call.compsub = compsub

96 def get_v_from_candidates(call):

97 """ return some v from candidates """

98 return iter(call.candidates).next()

100 def used_candidates(call):

101 """ if v from used connected with all from candidates """

102 for u in call.used:

103 for c in call.candidates:

104 if not self.get(u,{}).get(c,None):

105 break

106 else:

107 return True

108 return False

109

110 def __call__(call):

111 """ process call and return list of recursive calls """

112 while call.candidates and not call.used_candidates():

113 v = call.get_v_from_candidates()

114 call.compsub.add(v)

115 new_candidates = set([i for i in call.candidates

116 if i is not v and self.get(i,{}).get(v,None)])

117 new_used = set([i for i in call.used

118 if i is not v and self.get(i,{}).get(v,None)])

119 if (not new_candidates) and (not new_used):

120 if len(call.compsub) >= minsize:

121 cliques.append(copy(frozenset(call.compsub)))

122 else:

123 yield ExtendCall(new_candidates,

124 new_used, copy(call.compsub))

125 call.compsub.remove(v)

126 call.candidates.remove(v)

127 call.used.add(v)

128

129 stop_time = datetime.now() + timedelta(0, timeout)

130 stack = [ExtendCall(copy(set(self)), set(), set())()]

131 while stack:

132 if timeout != -1 and datetime.now() > stop_time:

133 raise TimeoutError()

134 top = stack[-1]

135 try:

136 call = top.next()

137 stack.append(call())

138 except StopIteration:

139 stack.pop()

140 return cliques

141

142 def drop_bad_vertices(self):

143 """ drop vertices untill self becomes a clique """

144 while len(self) >= 1:

145 indexes = dict((v, self.vertex_index(v)) for v in self)

146 min_index = min(indexes.values())

147 bad_vertices = [v for (v, i) in indexes.items()

148 if i == min_index]

149 if len(bad_vertices) == len(self):

150 break # clique

151 weights = dict([(self.vertex_weight(v),v) for v in bad_vertices])

152 min_weight = min(weights.keys())

153 vertex = weights[min_weight]

154 self.drop_vertex(vertex)

155

156 def expand_to_max_clique(self, parent_graph):

157 """ add vertices untill self becomes a nonextendible clique """

158 weights = dict((v, parent_graph.vertex_index(v))

159 for v in parent_graph)

160 while True:

161 candidates = {}

162 for vertex in set(parent_graph) - set(self):

163 neighbours = set(parent_graph[vertex])

164 if set(self).issubset(neighbours):

165 weight = parent_graph.vertex_weight(vertex)

166 candidates[weight] = vertex

167 if not candidates:

168 break

169 max_weight = max(candidates.keys())

170 vertex = candidates[max_weight]

171 self.add_vertex(vertex, parent_graph)

172

173 def fast_cliques(self, minsize=1):

174 """ returns list of cliques

175

176 clique is frozenset

177 """

178 print 'Fast algorithm started'

179 cliques = []

180 while True:

181 graph = self.copy()

182 for clique in cliques:

183 for v in clique:

184 if v in graph:

185 graph.drop_vertex(v)

186 if not graph:

187 break

188 graph.drop_bad_vertices()

189 graph.expand_to_max_clique(self)

190 clique = copy(frozenset(graph))

191 if len(clique) >= minsize:

192 cliques.append(clique)

193 else:

194 break

195 return cliques

196

197 def connected_components(self):

198 """ return list of connected components, remove single vertices """

199 used = set()

200 graphs = []

201 candidates = set(self)

202 while candidates:

203 graph = Graph()

204 queue = set()

205 queue.add(candidates.pop())

206 while queue:

207 v = queue.pop()

208 candidates.discard(v)

209 for v1 in self.get(v,{}):

210 if v1 in candidates:

211 graph.set_edge(v, v1, self[v][v1])

212 queue.add(v1)

213 graphs.append(graph)

214 return graphs

215

216 def cliques(self, timeout=-1, minsize=1):

217 """ returns length-sorted list of cliques

218

219 clique is frozenset

220

221 try to execute bron_kerbosh

222 if it raises TimeoutError, executes fast_cliques

223 """

224 cliques = []

225 for graph in self.connected_components():

226 if len(graph) >= minsize:

227 for v, neighbours_dict in graph.items():

228 if len(neighbours_dict) < minsize and v in graph:

229 graph.drop_vertex(v)

230 try:

231 cliques += graph.bron_kerbosh(timeout, minsize)

232 except TimeoutError:

233 cliques += graph.fast_cliques(minsize)

234 cliques.sort(key=lambda clique: len(clique), reverse=True)

235 return cliques

236

237 if __name__ == "__main__":

238 import doctest

239 doctest.testmod()

240