From 7d3203a69d59657ac420a2f363017c7e8ea35498 Mon Sep 17 00:00:00 2001 From: Fabian Gruenewald Date: Mon, 8 Jul 2024 16:57:13 +0200 Subject: [PATCH] read all fragments during init --- cgsmiles/resolve.py | 60 +++++++++++++++++++-------------------------- 1 file changed, 25 insertions(+), 35 deletions(-) diff --git a/cgsmiles/resolve.py b/cgsmiles/resolve.py index 950885e..e47dc4b 100644 --- a/cgsmiles/resolve.py +++ b/cgsmiles/resolve.py @@ -109,51 +109,49 @@ def __init__(self, self.fragment_dicts = fragment_dicts # this is the current meta_graph self.meta_graph = nx.Graph() - + # if the last resolution is all_atom + self.last_all_atom = last_all_atom + # what is the current resolution level + self.resolution_counter = 0 # here we figure out how many resolutions # we are dealing with elements = re.findall(r"\{[^\}]+\}", pattern) # case 1) we have a meta graph only described # and the fragment come from elsewhere - if len(elements) == 1 and self.fragment_dicts: + if self.fragment_dicts: + msg = ("You cannot provide a fragment dict and fragments defined as CGSmiles string." + "Instead first read all fragments using cgsmiles.read_fragments then provide " + "all fragments together to the MoleculeResolver") + if len(elements) != 1: raise IOError(msg) self.molecule = read_cgsmiles(elements[0]) - self.fragment_strs = None - # the number of resolutions available - self.resolutions = len(self.fragment_dicts) - # turn the framgent strings into an iterator - self.fragment_strs = None - self.fragment_dicts = self.fragment_dicts else: # case 2) # a meta_graph is provided which means we only # have fragments to deal with if meta_graph: - self.fragment_strs = elements + fragment_strs = elements self.molecule = meta_graph # case 3) a string containing both fragments and # the meta sequence is provided else: self.molecule = read_cgsmiles(elements[0]) - self.fragment_strs = elements[1:] + fragment_strs = elements[1:] - # the number of resolutions available - self.resolutions = len(self.fragment_strs) + # now we populate the fragment dicts + self.fragment_dicts = [] + for idx, fragment_str in enumerate(fragment_strs): + all_atom = (idx == len(fragment_strs) - 1 and self.last_all_atom) + f_dict = read_fragments(fragment_str, all_atom=all_atom) + self.fragment_dicts.append(f_dict) + + self.resolutions = len(self.fragment_dicts) # at this stage there are no atomnames for the nodes new_names = nx.get_node_attributes(self.molecule, "fragname") nx.set_node_attributes(self.meta_graph, new_names, "atomname") - # if the last resolution is all_atom - self.last_all_atom = last_all_atom - - # what is the current resolution - self.resolution_counter = 0 - - # the next resolution fragments - self.fragment_dict = {} - - def resolve_disconnected_molecule(self): + def resolve_disconnected_molecule(self, fragment_dict): """ Given a connected graph of nodes with associated fragment graphs generate a disconnected graph of the fragments and annotate @@ -162,7 +160,7 @@ def resolve_disconnected_molecule(self): """ for meta_node in self.meta_graph.nodes: fragname = self.meta_graph.nodes[meta_node]['fragname'] - fragment = self.fragment_dict[fragname] + fragment = fragment_dict[fragname] correspondence = merge_graphs(self.molecule, fragment) graph_frag = nx.Graph() @@ -250,18 +248,10 @@ def resolve(self): Resolve a CGSmiles string once and return the next resolution. """ # check if this is an all-atom level resolution - if self.resolution_counter == self.resolutions -1 and self.last_all_atom: - all_atom = True - else: - all_atom = False + all_atom = (self.resolution_counter == self.resolutions - 1 and self.last_all_atom) # get the next set of fragments - if self.fragment_strs: - fragment_str = self.fragment_strs[self.resolution_counter] - # read the fragment str and populate dict - self.fragment_dict = read_fragments(fragment_str, all_atom=all_atom) - else: - self.fragment_dict = self.fragment_dicts[self.resolution_counter ] + fragment_dict = self.fragment_dicts[self.resolution_counter] # set the previous molecule as meta_graph self.meta_graph = self.molecule @@ -274,7 +264,7 @@ def resolve(self): self.molecule = nx.Graph() # add disconnected fragments to graph - self.resolve_disconnected_molecule() + self.resolve_disconnected_molecule(fragment_dict) # connect valid bonding descriptors self.edges_from_bonding_descrpt(all_atom=all_atom) @@ -307,7 +297,7 @@ def resolve_iter(self): """ Iterator returning all resolutions in oder. """ - for _ in range(self.resolutions): + for r in range(self.resolutions): meta_graph, molecule = self.resolve() yield meta_graph, molecule