def IsRegularGraph(topo: nx.DiGraph) -> bool: """Dtermine a graph is regular or not, i.e. all nodes have the same degree.""" degree = topo.degree(0) for rank in range(1, topo.number_of_nodes()): if topo.degree(rank) != degree: return False return True
def compute_centrality(graph: nx.DiGraph, data: geojson.feature.FeatureCollection, edge_map: Dict): for item in data['features']: edge = edge_map[item['properties']['id']] from_degree = graph.degree(edge[0]) to_degree = graph.degree(edge[1]) item['properties']["from_degree"] = from_degree item['properties']["to_degree"] = to_degree
def remove_unconnected_nodes(graph: nx.DiGraph): """Remove unconnected nodes and return the corresponding issues.""" removed_nodes = [node for node in graph.nodes() if graph.degree(node) == 0] removed_issues = {graph.node[node]['issue'] for node in removed_nodes} if removed_nodes: graph.remove_nodes_from(removed_nodes) return removed_issues
def create_node_data(node: NodeTy, commit: CommitRepoPair, cig: nx.DiGraph) -> ArcPlotNodeInfo: return { "info": commit.commit_hash.short_hash, "size": cig.degree(node), "fill_color": cig.out_degree(node), "line_color": cig.in_degree(node) }
def calc_node_size(graph: nx.DiGraph, node: tg.Text) -> int: """Size nodes by their degree. Use natrual log to make them less diversed multiply by SIZE_FACTOR to make them not too small plus 1 so that size won't be zero """ SIZE_FACTOR = 3 return math.floor(SIZE_FACTOR * math.log1p(graph.degree(node))) + 1
def _construct(self): """ We want to build the type of DFG that's used in "Automated Ident. of Crypto Primitives in Binary Code with Data Flow Graph Isomorphisms." Unlike that paper, however, we're building it on Vex IR instead of assembly instructions. """ cfg = self._cfg p = self.project dfgs = {} l.debug("Building Vex DFG...") for node in cfg.nodes(): try: if node.simprocedure_name == None: irsb = p.factory.block(node.addr).vex else: l.debug("Cannot process SimProcedures, ignoring %s" % node.simprocedure_name) continue except Exception as e: l.debug(e) continue tmpsnodes = {} storesnodes = {} putsnodes = {} statements = irsb.statements dfg = DiGraph() for stmt_idx, stmt in enumerate(statements): # We want to skip over certain types, such as Imarks if self._need_to_ignore(node.addr, stmt, stmt_idx): continue # break statement down into sub-expressions exprs = stmt.expressions stmt_node = stmt dfg.add_node(stmt) if stmt.tag == 'Ist_WrTmp': tmpsnodes[stmt.tmp] = stmt_node if exprs[0].tag == 'Iex_Binop': if exprs[1].tag == 'Iex_RdTmp': dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node) else: dfg.add_edge(exprs[1], stmt_node) if exprs[2].tag == 'Iex_RdTmp': dfg.add_edge(tmpsnodes[exprs[2].tmp], stmt_node) else: dfg.add_edge(exprs[2], stmt_node) elif exprs[0].tag == 'Iex_Unop': dfg.remove_node(stmt_node) if exprs[1].tag == 'Iex_RdTmp': tmpsnodes[stmt.tmp] = copy(tmpsnodes[exprs[1].tmp]) tmpsnodes[stmt.tmp].tmp = stmt.tmp else: tmpsnodes[stmt.tmp] = exprs[1] elif exprs[0].tag == 'Iex_RdTmp': tmpsnodes[stmt.tmp] = copy(tmpsnodes[exprs[0].tmp]) tmpsnodes[stmt.tmp].tmp = stmt.tmp elif exprs[0].tag == 'Iex_Get': if putsnodes.has_key(exprs[0].offset): dfg.add_edge(putsnodes[exprs[0].offset], stmt_node) if len(exprs) > 1 and exprs[1].tag == "Iex_RdTmp": dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node) elif len(exprs) > 1: dfg.add_edge(exprs[1], stmt_node) elif exprs[0].tag == 'Iex_Load': if exprs[1].tag == 'Iex_RdTmp': dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node) else: dfg.add_edge(exprs[1], stmt_node) else: # Take a guess by assuming exprs[0] is the op and any other expressions are args for e in exprs[1:]: if e.tag == 'Iex_RdTmp': dfg.add_edge(tmpsnodes[e.tmp], stmt_node) else: dfg.add_edge(e, stmt_node) elif stmt.tag == 'Ist_Store': if exprs[0].tag == 'Iex_RdTmp': dfg.add_edge(tmpsnodes[exprs[0].tmp], stmt_node) elif exprs[0].tag == 'Iex_Const': dfg.add_edge(exprs[0], stmt_node) if exprs[1].tag == 'Iex_RdTmp': dfg.add_edge(tmpsnodes[exprs[1].tmp], stmt_node) else: dfg.add_edge(exprs[1], stmt_node) elif stmt.tag == 'Ist_Put': if exprs[0].tag == 'Iex_RdTmp': dfg.add_edge(tmpsnodes[exprs[0].tmp], stmt_node) elif exprs[0].tag == 'Iex_Const': dfg.add_edge(exprs[0], stmt_node) putsnodes[stmt.offset] = stmt_node elif stmt.tag == 'Ist_Exit': if exprs[0].tag == 'Iex_RdTmp': dfg.add_edge(tmpsnodes[exprs[0].tmp], stmt_node) elif stmt.tag == 'Ist_Dirty': tmpsnodes[stmt.tmp] = stmt_node elif stmt.tag == 'Ist_CAS': tmpsnodes[stmt.oldLo] = stmt_node else: for e in stmt.expressions: if e.tag == 'Iex_RdTmp': dfg.add_edge(tmpsnodes[e.tmp], stmt_node) else: dfg.add_edge(e, stmt_node) for vtx in dfg.nodes(): if dfg.degree(vtx) == 0: dfg.remove_node(vtx) if dfg.size() > 0: dfgs[node.addr] = dfg return dfgs
class DepF : def __init__(self, field) : self.__offsets = [] self.__field = field ############ create depency field graph ######### # Initial values to build the graph (depth, width, cycles) self.__depth = 3 self.__width = 3 self.__cycles = 2 self.__G = DiGraph() G = self.__G G.add_node( self._new_node(G) ) # Create randomlu the graph without cycle self.__random( G, 0 ) # Find the best path to add cycles d = all_pairs_dijkstra_path_length( G ) l = list( reversed( sorted( d, key=lambda key: len(d[key]) ) ) ) for i in l : if self.__cycles == 0 : break d_tmp = sorted( d[i], key=lambda key: d[i][key] ) G.add_edge( d_tmp[-1], i) self.__cycles = self.__cycles - 1 print simple_cycles( G ) print G.node print G.edge print G.degree() ############ Add field <-> higher ##################################### # field F depends of the higher degree # F <---> sorted( G.degree(), key=lambda key : G.degree()[key] )[-1] ####################################################################### degree = G.degree() high_degree = sorted( degree, key=lambda key : degree[key] )[-1] # Link our protected field with the node which has the highest degree G.add_edge( field.get_name(), high_degree ) G.add_edge( high_degree, field.get_name() ) #draw_graphviz(G) #write_dot(G,'file.dot') def get_field(self) : return self.__field def add_offset(self, idx) : x = Offset( idx ) self.__offsets.append( x ) return x def run(self, _vm, _analysis, _vm_generate) : ############################################################### ## dict (method) of list ( offset / list (of instructions) ) # ## - insert an element # ## - modify the offset with the new insertion # ############################################################### list_OB = {} ############ Create dependencies fields ############ fields = { self.__field.get_name() : Field( _vm, _analysis, _vm_generate, self.__field, self, True ) } fields[ self.__field.get_name() ].run( self.__G.degree()[ self.__field.get_name() ] ) print self.__field.get_name(), ": REAL_FIELD ->", self.__G.predecessors( self.__field.get_name() ), self.__G.degree()[self.__field.get_name()] ############ Create the name, the initial value and all access of the field ############ for i in self.__G.node : # We have not yet add this new field if i not in fields : print i, "PRE ->", self.__G.predecessors( i ), self.__G.degree()[i] name, access_flag, descriptor = _analysis.get_like_field() _vm.insert_field( self.__field.get_class_name(), name, [ access_flag, descriptor ] ) fields[ i ] = Field( _vm, _analysis, _vm_generate, _vm.get_field_descriptor( self.__field.get_class_name(), name, descriptor ), self ) # degree of the field, prefix must be add if the protection is for a real field fields[ i ].run( self.__G.degree()[i] ) ########## Add all fields initialisation into the final list ############ for i in fields : print "FIELD ->", i, fields[ i ].show() x = fields[ i ].insert_init() if x != None : try : list_OB[ x[0] ].append( (x[1], x[2]) ) except KeyError : list_OB[ x[0] ] = [] list_OB[ x[0] ].append( (x[1], x[2]) ) ############ Create the depedency ############ ############################################################################ # Integer variables : # X -> Y # - a depth into the calcul # Y = { LV + LLV + NLV } x { &, -, +, |, *, /, ^ } x { &&, || } # if (Y != ?) { # X = ..... ; # } ############################################################################# # get a local variable # - used into a loop # - a parameter # - a new one ############################################################################ find = False print "F -->", self.__G.successors( self.__field.get_name() ) taint_field = _analysis.get_tainted_field( self.__field.get_class_name(), self.__field.get_name(), self.__field.get_descriptor() ) for path in taint_field.get_paths() : continue print "\t", path.get_access_flag(), "%s (%d-%d)" % (path.get_bb().get_name(), path.get_bb().get_start(), path.get_bb().get_end()) , path.get_bb().get_start() + path.get_idx(), x = _analysis.get( path.get_bb().get_method() ) bb = x.get_break_block( path.get_bb().get_start() + path.get_idx() ) print "\t\t", x.get_local_variables() if path.get_access_flag() == "R" and find == False : o = self.add_offset( _analysis.prev_free_block_offset( path.get_method(), bb.get_start() ) ) val = _analysis.next_free_block_offset( path.get_method(), o.get_idx() ) if o.get_idx() == -1 or val == -1 : raise("ooop") #try : # list_OB[ path.get_method() ].append( o, [ [ "iload_3" ], [ "iconst_0" ], [ "if_icmpge", val - o.get_idx() + 3 ] ] ) #except KeyError : # list_OB[ path.get_method() ] = [] # list_OB[ path.get_method() ].append( (o, [ [ "iload_3" ], [ "iconst_0" ], [ "if_icmpge", val - o.get_idx() + 3 ] ] ) ) find = True ##### Insert all modifications for m in list_OB : code = m.get_code() i = 0 while i < len( list_OB[ m ] ) : v = list_OB[ m ][ i ] print "INSERT ", v[0].get_idx(), v[1] size_r = code.inserts_at( code.get_relative_idx( v[0].get_idx() ), v[1] ) #code.show() j = i + 1 while j < len( list_OB[ m ] ) : v1 = list_OB[ m ][ j ] if v1[0].get_idx() >= v[0].get_idx() : v1[0].add_idx( size_r ) j = j + 1 i = i + 1 def _new_node(self, G) : return "X%d" % (len(G.node)) def _current_node(self, G) : return len(G.node) - 1 def __random(self, G, depth) : if depth >= self.__depth : return for i in range( random.randint(1, self.__width) ) : nd = self._new_node(G) G.add_edge( "X%d" % depth, nd ) self.__random( G, self._current_node(G) )
class ProcessGraph(object): """ Helper class to represent a materials synthesis action graph / quest in a format facilitating easy visualization in process form. This is also used by the SurfaceGenerator to create the textual description of a quest. """ def __init__(self) -> None: self.G = DiGraph() # For each material we hold a list of nodes representing it over time # as its state changes self.ent_states_map = {} def init_ent(self, e: Variable): """ Create and initialize new EntityNode. """ if not e in self.ent_states_map: self.ent_states_map[e] = [EntityNode(var=e, g=0)] def curr_state(self, e: Variable): """ Return the last (current) state for a given entity variable. """ return self.ent_states_map[e][-1] def add_new_state(self, e: Variable): """ Add new state for a given entity variable (except Operation entities). This creates a new EntityNode with g increased by 1. """ if KnowledgeBase.default().types.is_descendant_of(e.type, ['op']): return self.ent_states_map[e] else: new_state_node = EntityNode(var=e, g=(self.curr_state(e).g + 1)) self.ent_states_map[e].append(new_state_node) return new_state_node def topological_sort_actions(self): """ Return edges in "topological order", meaning we iterate over nodes in topological order and for each node iterate over its outgoing edges. """ edge_list = [] for n in nx.topological_sort(self.G): edge_list += edges_to_actions(self.G.out_edges(n, data=True)) return edge_list def get_incoming_actions(self, node: EntityNode) -> List[ActionEdge]: """ Return the incoming actions to a given EntityNode. """ return edges_to_actions(self.G.in_edges(node, data=True)) def get_actions_path(self, source: EntityNode, target: EntityNode) -> List[ActionEdge]: """ Return list of ActionEdges comprising the shortest path between a source and target node in the graph. """ action_path = [] path_nodes = nx.shortest_path(self.G, source, target) if len(path_nodes) < 2: return action_path edges = [e for e in zip(path_nodes[:-1], path_nodes[1:]) ] # get edges comprising path for u, v in edges: action_path.append( ActionEdge(action=self.G.edges()[(u, v)]['action'], source=u, target=v)) return action_path def get_start_material_nodes(self) -> List[EntityNode]: """ Get nodes of all materials at start of process (g=0). """ return [ n for n in self.G.nodes() if (n.var.type == 'm') and (n.g == 0) ] def get_start_material_vars(self) -> List[Variable]: """ Get variables of all materials at start of process (g=0). """ return [mat.var for mat in self.get_start_material_nodes()] def from_tw_actions(self, actions: Iterable[Action]) -> DiGraph: """ Construct internal graph given a TextWorld Quest. """ # TODO verify validity? from tw_textlabs.generator.quest_generator import convert_to_compact_action G = nx.DiGraph() action_vars = [convert_to_compact_action(a) for a in actions] for action_var in action_vars: action_name = action_var.name if action_name == 'op_o_obtain': # add change state var as arg to this action so it gets processed action_var.vars.insert(0, action_var.change_state_vars[0]) ents = action_var.vars for e in ents: self.init_ent(e) # ignore actions with single arg if len(ents) != 2: continue is_state_change_action = len(action_var.change_state_vars) > 0 if not is_state_change_action: G.add_edge(self.curr_state(ents[0]), self.curr_state(ents[1]), action=action_name) # assuming just one material changes state per action else: assert (len(action_var.change_state_vars) == 1) v = action_var.change_state_vars[0] assert (v == ents[0]) # second argument is the source for new one (such as device) source = self.curr_state(ents[1]) # Add a new state for the device, we assume that each operation # on a material is at a new state. if self.curr_state(v) in G.nodes(): # entity already exists in graph target = self.add_new_state(v) else: # entity doesn't exist yet in graph, connect init node target = self.curr_state(v) # Add 'result' type edge which corresponds to 'obtain' action G.add_edge(source, target, action='obtain') # add Generator dummy node to be source of graph (cosmetic) self.generator = EntityNode(var=Variable(name='START', type=GENERATOR_DUMMY_TYPE), g=0) G.add_node(self.generator) self.G = G materials = self.get_start_material_nodes() for m in materials: self.G.add_edge(self.generator, m, action='take') def get_op_type(self, node: EntityNode) -> str: assert (KnowledgeBase.default().types.is_descendant_of( node.var.type, 'tlq_op')) op_type_node = [ n for n in self.G.neighbors(node) if KnowledgeBase.default().types.is_descendant_of( n.var.type, 'toe') ] assert (len(op_type_node) == 1) op_type_node = op_type_node[0] return op_type_node.var.type def get_source_op(self, node: EntityNode) -> EntityNode: assert (KnowledgeBase.default().types.is_descendant_of( node.var.type, ['m', 'sa'])) # return first operation node producing target node for n in self.ent_states_map[node.var]: for s, t in self.G.in_edges(n): if s.var.type == 'tlq_op': return s return None def rename_edges(self, node: EntityNode, orig: str, target: str, incoming: bool = True) -> None: """ Rename a node's incoming or outgoing edges """ edges = [e for e in self.G.in_edges(node) ] if incoming else [e for e in self.G.out_edges(node)] for edge in edges: act = self.G.edges()[edge]['action'] self.G.edges()[edge]['action'] = act.replace(orig, target) def draw(self): """ Draw the ProcessGraph. """ connected_nodes = [n for n in self.G.nodes() if self.G.degree(n) > 0] connected_G = self.G.subgraph(connected_nodes) # Hierarchical layout pos = graphviz_layout(connected_G, prog='dot') edge_labels = dict([(( u, v, ), d['action']) for u, v, d in self.G.edges(data=True)]) node_labels = {n: str(n) for n in connected_G.nodes()} # Add dummy source node node_labels.update({self.generator: 'START'}) node_colors = [node_to_color(n) for n in connected_G.nodes()] nx.draw(connected_G, pos, labels=node_labels, with_labels=True, node_color=node_colors) nx.draw_networkx_edge_labels(connected_G, pos, edge_labels=edge_labels)
class Case3(Case): """ This class defines a radial network which will be studied further. It uses directed graph implemented in networkx. """ def __init__(self, name=None, base_mva=1.0, base_kv=20.0, buses=None, branches=None, generators=None): #Case.__init__(self,buses=buses,branches=branches) self.buses = buses self.base_mva = base_mva self.base_kv = base_kv self.base_z = base_kv**2 / base_mva self.branches = branches self.generators = generators self.G = DiGraph() self._buildGraph() self.rootbus = [b for b in self.buses if self.G.pred[b] == {}] self.forkbuses = [b for b in self.buses if self.G.degree()[b] > 2] self.normalbuses = [b for b in self.buses if self.G.degree()[b] == 2] self.endbuses = [b for b in self.buses if self.G.succ[b] == {}] self.updateTypeG() self.bus_level = shortest_path_length(self.G, source=self.rootbus[0]) self.updateBusConnBranch() self.buses_loadJ = self.updateBusLoadCurrent() self.identifyPVbuses() self.calculateZpos() for b in self.buses: b.base_mva = base_mva #draw(self.G) #pl.show() def _buildGraph(self): """ add buses as nodes and branches as edges. """ self.G.add_nodes_from([bus for bus in self.buses]) self.G.add_edges_from([(b.from_bus, b.to_bus) for b in self.branches]) def plotGraph(self): draw(self.G) show() return def plotVP(self): """ Plot voltage profile of distribution system""" V = [abs(b.E) for b in self.buses] name = [b.name for b in self.buses] plot(V) show() return def updateBusConnBranch(self): """ Updating to which branch a bus is connected. """ for a in self.branches: a.from_bus.connected_to_branch.append(a) a.to_bus.connected_from_branch.append(a) # as edges are added, there should be a procedure to update bus.connected_to and bus.connected_from def updateTypeG(self): for b in self.buses: if self.rootbus.__contains__(b): b.typeG = rootNode if self.forkbuses.__contains__(b): b.typeG = forkNode if self.endbuses.__contains__(b): b.typeG = endNode def updateBusLoadCurrent(self): return [b.calc_loadI() for b in self.buses] def getRoot(self): """ Get the root node of a network which is the substation node. """ rootNode = [n for n in self.G.node if self.G.pred[n] == {}] return rootNode[0] def updateBaseKV(self, aNode): """ update base_kv of aNode children. This is done after a voltage regulator changes its tap. """ for b in aNode.connected_to_branch: b.updateToBusBaseKV() if b.to_bus.connected_to_branch != []: self.updateBaseKV(b.to_bus) def identifyPVbuses(self): """ identify pv buses in the network. Important when creating the positive sequence matrix needed to update injected current """ self.pvbuses = [ g.bus for g in self.generators if g.type == 'PV' ] # a list of buses connected to generators controlled as PV bus. for b in self.pvbuses: b.type = 'PV' # update the bus.type according to its generator control mode. def calculateZpos(self): """ calculate positive impedance matrix for updating injected current at PV bus in BFS with DGs """ self.identifyPVbuses() pathToRoot = {} lines = {} Z1diag = {} pvbusindex = {} d = defaultdict(dict) npvb = len(self.pvbuses) # number of pv buses Z1 = zeros([npvb, npvb], dtype='complex') for pvb in self.pvbuses: pathToRoot[pvb] = shortest_path(self.G, self.rootbus[0], pvb) lines[pvb] = [ b.connected_from_branch[0] for b in pathToRoot[pvb] if b.connected_from_branch != [] ] self.pathToRoot = pathToRoot # a list of buses that incident with the path from a pv bus to root bus. self.lines = lines # a list of lines that forms the path from a pvbus to the root bus. Z_1 = zeros([npvb, npvb]) # prepare an all zero matrix with size npvb x npvb # This function requires a positif sequence of all lines in self.lines. The positive sequence impedance is available at line.Z012[0] # the situation is more intricate if there is a transformer within the path between the PV bus to root bus. If on of the transformer winding # is connected in delta, the path for positive sequence is disconnected. The implementation for this kind of situation is postponed. # the current implementation assumes that there is NO TRANSFORMER exist in the path. 27th May 2011. for pvb in self.pvbuses: Z1sum = complex(0, 0) for l in self.lines[pvb]: Z1sum = Z1sum + l.Z012[1][1] Z1diag[pvb] = Z1sum self.Z1diag = Z1diag # this is the diagonal element of positive impedance matrix, the output of this function. ind = 0 for pvb in self.pvbuses: pvbusindex[pvb] = ind ind = ind + 1 self.pvbusindex = pvbusindex for pvb in self.pvbuses: for pv in self.pvbuses: d[pvb][pv] = set(self.lines[pvb]).intersection( set(self.lines[pv])) Z1sum = 0 for l in d[pvb][pv]: Z1sum = Z1sum + l.Z012[1][1] Z1[pvbusindex[pvb], pvbusindex[pv]] = Z1sum self.d = d # this is a two dimensional dictionary. Its diagonal element contains list of lines on the path between a pv bus and root bus. # Its of diagonal contains the common lines between path to root of two pv buses. If we sum the positive sequence impedance # of each element, we have already the element of postive sequence impedance matrix we wish to calculate in this function. self.Z1 = Z1 self.X1 = imag(Z1)
class rxnGraph(object): def __init__(self, datatables, proteins=None, ligands=None, interactions=None): #, dockdf=None, seadf=None, rsimpn=None): # Construct graph self.graph = DiGraph() self.data = datatables self.proteins = proteins # list of nodes self.ligands = ligands # list of nodes self.protein_members = [] # subset of possible proteins, members in graph self.protein_choices = [] # subset of possible proteins, not members in graph self.ligand_members = [] # subset of possible ligands, members in graph self.ligand_choices = [] # subset of possible ligands, not members in graph self.transporter = None self.protein_nodes = [] # moveable nodes self.ligand_nodes = [] self.protein_allnodes = [] self.ligand_allnodes = [] self.fixed = [] self.keep = [] if proteins is not None and ligands is not None and interactions is not None: self.construct_graph(proteins, ligands, interactions) self.protein_members = proteins[:] self.ligand_members = ligands[:] def add_choices(self, protein_choices=[], ligand_choices=[]): self.ligand_choices = ligand_choices self.protein_choices = protein_choices def add_transporter(self, transporter): self.transporter = transporter #n = self.get_start() #self.transporter_substrate = self.graph.node[n]['label'] def copy(self): newGraph = rxnGraph(self.data, proteins=self.proteins[:], ligands = self.ligands[:]) #dockdf=self.dockdf, seadf=self.seadf, rsimpn=self.rsimpn) newGraph.graph = self.graph.copy() #newGraph.constraints = self.constraints newGraph.ligand_members = self.ligand_members[:] newGraph.protein_members = self.protein_members[:] newGraph.ligand_choices = self.ligand_choices[:] newGraph.protein_choices = self.protein_choices[:] newGraph.protein_nodes = self.protein_nodes newGraph.ligand_nodes = self.ligand_nodes newGraph.protein_allnodes = self.protein_allnodes newGraph.ligand_allnodes = self.ligand_allnodes newGraph.fixed = self.fixed newGraph.keep = self.keep newGraph.transporter = self.transporter return newGraph def construct_graph_from_path(self, proteins, ligands): #self.graph.add_nodes_from(ligands, bipartite=0) #self.graph.add_nodes_from(proteins, bipartite=1) interactions = [] for i in range(len(proteins)): interactions.append([ligands[i], proteins[i]]) if i > 0: interactions.append([proteins[i-1], ligands[i]]) if len(ligands) == len(proteins) + 1: interactions.append([proteins[-1], ligands[-1]]) self.construct_graph(proteins, ligands, interactions) def construct_graph(self, proteins, ligands, interactions): # interactions: [[protein, ligand], [protein, ligand]] self.graph.clear() labelposdict = {} for i, p in enumerate(proteins): self.graph.add_node(i, label=p, bipartite=1) labelposdict[p] = i self.protein_allnodes.append(i) for j, l in enumerate(ligands): self.graph.add_node(j + len(proteins), label=l, bipartite=0) labelposdict[l] = j + len(proteins) self.ligand_allnodes.append(j + len(proteins)) #for i in range(len(interactions)): # node1 = labelposdict[interactions[i][0]] # node2 = labelposdict[interactions[i][1]] # self.graph.add_edge(node1, node2) [self.graph.add_edge(labelposdict[interactions[i][0]], labelposdict[interactions[i][1]]) for i in range(len(interactions))] self.proteins = proteins self.ligands = ligands self.protein_members = proteins[:] self.ligand_members = ligands[:] self.ligand_nodes = self.ligand_allnodes[:] self.protein_nodes = self.protein_allnodes[:] for i in self.fixed: if i in self.protein_nodes: self.protein_nodes.remove(i) if i in self.ligand_nodes: self.ligand_nodes.remove(i) # updates the proteins and ligands in order # maybe not necessary def update_proteins_ligands(self): n = self.get_start() all_ligands = self.ligand_members[:] all_ligands.extend(self.ligand_choices[:]) all_proteins = self.protein_members[:] all_proteins.extend(self.protein_choices[:]) self.ligands = [] self.proteins = [] if n in all_proteins: self.proteins.append(n) elif n in all_ligands: self.ligands.append(n) while len(self.graph.successors(n)) > 0: n = self.graph.successors(n)[0] if n in all_proteins: self.proteins.append(n) elif n in all_ligands: self.ligands.append(n) # also, maybe not necessary def get_start(self): start = None for n in self.graph.nodes_iter(): if len(self.graph.predecessors(n)) < 1: start = n if start is None: print self.graph.edges() return start def get_single_predecessor_node(self, node): pred = self.graph.predecessors(node) if len(pred) > 0: return pred[0] else: return None def get_single_successor_node(self, node): succ = self.graph.successors(node) if len(succ) > 0: return succ[0] else: return None def move_remove_single(self): all = self.protein_members[:] all.extend(self.ligand_members[:]) node = self.choose_by_degree(1, all, 1) if node is not None: self.graph.remove_node(node) if node in self.ligand_members: self.ligand_members.remove(node) self.ligand_choices.append(node) elif node in self.protein_members: self.protein_members.remove(node) self.protein_choices.append(node) # Removes a pair of consecutive nodes, a protein node and a ligand node # Adds edge between predecessors to the successors of the removed nodes def move_remove_pair(self, pnode): lnodes = None if self.graph.in_edges(pnode) is not None and self.graph.out_edges(pnode) is not None: ppred = self.graph.predecessors(pnode) psucc = self.graph.successors(pnode) node_choices = [] if len(ppred) > 0: node_choices.append(ppred) if len(psucc) > 0: node_choices.append(psucc) if len(node_choices) > 0: #lnodes = rand.choice(node_choices) lnode = rand.choice(node_choices) if len(ppred) > 0 and len(psucc) > 0: for pp in ppred: for ps in ppred: self.graph.add_edge(pp, ps) self.graph.remove_node(pnode) self.protein_members.remove(pnode) self.protein_choices.append(pnode) if self.graph.in_edges(lnode) is not None and self.graph.out_edges(lnode) is not None: lpred = self.graph.predecessors(lnode) lsucc = self.graph.successors(lnode) if len(lpred) > 0 and len(lsucc) > 0: for lp in lpred: for ls in lsucc: self.graph.add_edge(lp, ls) self.graph.remove_node(lnode) self.ligand_members.remove(lnode) self.ligand_choices.append(lnode) # choose k nodes from a nodeset with a degree less than the limit def choose_by_degree_in(self, upperlim_deg, nodeset, k): choices = [] for node in nodeset: degree = self.graph.in_degree(node) if degree < upperlim_deg: choices.append(node) if len(choices) >= k: return rand.sample(choices, k) else: return None def choose_by_degree_out(self, upperlim_deg, nodeset, k): choices = [] for node in nodeset: degree = self.graph.out_degree(node) if degree < upperlim_deg: choices.append(node) if len(choices) >= k: return rand.sample(choices, k) else: return None def choose_by_degree(self, upperlim_deg, nodeset, k): choices = [] for node in nodeset: degree = self.graph.degree(node) if degree < upperlim_deg: choices.append(node) if len(choices) >= k: return rand.sample(choices, k) else: return None # Adds a pair of consecutive nodes to the end of the linear graph def move_add_pair(self, pnode, lnode): last = self.choose_by_degree_in(2, self.ligand_members, 1)[0] if last is not None: self.graph.add_node(pnode) self.graph.add_node(lnode) self.graph.add_edge(last, pnode) self.graph.add_edge(pnode, lnode) # Updates lists to reflect the current state of the full system self.ligand_members.append(lnode) self.ligand_choices.remove(lnode) self.protein_members.append(pnode) self.protein_choices.remove(pnode) def move_add_edge(self): if rand.choice([0, 1]): pnode = self.choose_by_degree_out(2, self.protein_members, 1) lnode = self.choose_by_degree_in(2, self.ligand_members, 1) if pnode is not None and lnode is not None: self.graph.add_edge(pnode[0], lnode[0]) else: pnode = self.choose_by_degree_in(2, self.protein_members, 1) lnode = self.choose_by_degree_out(2, self.ligand_members, 1) if pnode is not None and lnode is not None: self.graph.add_edge(lnode[0], pnode[0]) def move_swap_pair(self, swapLigands=True): if swapLigands: node1, node2 = rand.sample(self.ligand_nodes, 2) else: node1, node2 = rand.sample(self.protein_nodes, 2) if (node1 not in self.fixed) and (node2 not in self.fixed): label1 = self.graph.node[node1]['label'] label2 = self.graph.node[node2]['label'] self.graph.node[node1]['label'] = label2 self.graph.node[node2]['label'] = label1 def move_swap_neighboring_pair(self, swapLigands=True): if swapLigands: node = rand.choice(self.ligand_nodes) else: node = rand.choice(self.protein_nodes) predprot = self.get_single_predecessor_node(node) succprot = self.get_single_successor_node(node) options = [] if predprot is not None: predpred = self.get_single_predecessor_node(predprot) if predpred is not None: options.append(predpred) if succprot is not None: succsucc = self.get_single_successor_node(succprot) if succsucc is not None: options.append(succsucc) if self.fixed in options: options.remove(self.fixed) if len(options) > 0: swapnode = rand.choice(options) label1 = self.graph.node[node]['label'] label2 = self.graph.node[swapnode]['label'] self.graph.node[node]['label'] = label2 self.graph.node[swapnode]['label'] = label1 # assumes linear def move_lp_swap_pairs(self, pnode1, pnode2): lpred = self.graph.predecessors(pnode1) if not len(lpred) > 0: tempnode = pnode1 pnode1 = pnode2 pnode2 = tempnode lpred = self.graph.predecessors(pnode1) ppred = self.graph.predecessors(lpred[0]) lsucc = self.graph.successors(pnode1) if len(ppred) > 0: self.graph.remove_edge(ppred[0], lpred[0]) if len(lsucc) > 0: self.graph.remove_edge(pnode1, lsucc[0]) if len(ppred) > 0 and len(lsucc) > 0: self.graph.add_edge(ppred[0], lsucc[0]) lsucc2 = self.graph.successors(pnode2) if len(lsucc2) > 0: self.graph.remove_edge(pnode2, lsucc2[0]) self.graph.add_edge(pnode1, lsucc2[0]) self.graph.add_edge(pnode2, lpred[0]) def move_swap_in(self, swapLigands=True): if swapLigands: select = True while select or onodelabel in self.keep: onode = rand.choice(self.ligand_nodes) onodelabel = self.graph.node[onode]['label'] select = False newnodelabel = rand.choice(self.ligand_choices) else: onode = rand.choice(self.protein_nodes) onodelabel = self.graph.node[onode]['label'] newnodelabel = rand.choice(self.protein_choices) self.graph.node[onode]['label'] = newnodelabel if onode in self.ligand_allnodes: self.ligand_choices.remove(newnodelabel) self.ligand_choices.append(onodelabel) elif onode in self.protein_allnodes: self.protein_choices.append(onodelabel) self.protein_choices.remove(newnodelabel) self.protein_members.append(newnodelabel) self.protein_members.remove(onodelabel) # lnode is the new node to swap in # tested, doesn't help at all def move_swap_in_similar(self, lnode): npdt = np.dtype([('ligand', np.str_, 35), ('tc', np.float32)]) tcarray = [] tcsum = 0.0 for lig in self.ligand_members: tc = self.data.chsimdf.get_value(lig, lnode) tcarray.append((lig, tc)) nptc = np.array(tcarray, dtype = npdt) cutoff = rand.uniform(0, 0.5) bool = nptc['tc'] > cutoff if bool.any(): options = nptc['ligand'][bool] else: options = nptc['ligand'] outnode = rand.choice(options) pred = self.get_single_predecessor_node(outnode) succ = self.get_single_successor_node(outnode) self.graph.add_node(lnode) if pred is not None: self.graph.add_edge(pred, lnode) self.graph.remove_edge(pred, outnode) if succ is not None: self.graph.add_edge(lnode, succ) self.graph.remove_edge(outnode, succ) self.graph.remove_node(outnode) self.ligand_members.remove(outnode) self.ligand_choices.append(outnode) self.ligand_members.append(lnode) self.ligand_members.remove(lnode) def get_subgraph_successors(self, pnode, numnodes): currnode = pnode subgraph_nodes = [pnode] for n in range(numnodes - 1): if currnode is not None: nextnode = self.get_single_successor_node(currnode) if nextnode is not None: currnode = self.get_single_successor_node(nextnode) if currnode is self.fixed: break if currnode is not None: subgraph_nodes.append(currnode) else: return False return subgraph_nodes def get_subgraph_predecessors(self, pnode, numnodes): currnode = pnode subgraph_nodes = [pnode] for n in range(numnodes - 1): if currnode is not None: nextnode = self.get_single_predecessor_node(currnode) if nextnode is not None: currnode = self.get_single_successor_node(nextnode) if currnode is self.fixed: break if currnode is not None: subgraph_nodes.append(currnode) else: return False return subgraph_nodes.reverse() def move_subgraph(self): pnode = rand.choice(self.protein_nodes) numnodes = rand.choice(range(2, len(self.protein_nodes)/2 + 1)) subgraph_nodes = self.get_subgraph_successors(pnode, numnodes) if not subgraph_nodes: subgraph_nodes = self.get_subgraph_predecessors(pnode, numnodes) if subgraph_nodes: pnode = subgraph_nodes[0] qnode = subgraph_nodes[-1] lnode = self.get_single_predecessor_node(pnode) rchoices = self.protein_nodes[:] for s in subgraph_nodes: if s in rchoices: rchoices.remove(s) rnode = rand.choice(rchoices) # Remove edges if lnode is not None: lpred = self.get_single_predecessor_node(lnode) qsucc = self.get_single_successor_node(qnode) rsucc = self.get_single_successor_node(rnode) if lnode == rsucc: rchoices.remove(rnode) rnode = rand.choice(rchoices) rsucc = self.get_single_successor_node(rnode) if lpred is not None: self.graph.remove_edge(lpred, lnode) if qsucc is not None: self.graph.add_edge(lpred, qsucc) if qsucc is not None: self.graph.remove_edge(qnode, qsucc) if rsucc is not None: self.graph.remove_edge(rnode, rsucc) self.graph.add_edge(qnode, rsucc) self.graph.add_edge(rnode, lnode) def string_repr(self): string = '' for e in self.graph.edges(): string += "('%s', '%s'), " % (e[0], e[1]) return string.rstrip(', ') def string_interactions(self): string = '' for e in self.graph.edges(): string += "'%s', '%s', " % (e[0], e[1]) return string.rstrip(', ') def linear_to_string(self): string = '' if self.transporter is not None and self.transporter in self.data.dockdf.index: string += '%s -> ' % self.transporter n = self.get_start() string += self.graph.node[n]['label'] while len(self.graph.successors(n)) > 0: n = self.graph.successors(n)[0] label = self.graph.node[n]['label'] string += ' -> ' string += label return string def add_constraint(self, fun): self.constraints.append(fun) def test_constraints(self): for fun in self.constraints: if not fun(self): return False return True