Source code for molfoundry.atomtracing.saturate_tracer

from bisect import bisect_right
from collections import defaultdict
from operator import itemgetter
from typing import List, Tuple, Dict, Set, Optional, Callable, Final

type TemplateInstance = Tuple[Tuple[int, int], ...]

# Origin ID of an atom that is not tracked, and the symbol it is displayed as. The same value is
# reused for an origin that is present but deliberately not shown in a filtered projection: in a
# drawing, an atom whose origin is hidden and an atom that never had one carry the same claim.
UNTRACKED_ORIGIN: Final[int] = -1
UNTRACKED_ORIGIN_LABEL: Final[str] = '○'


[docs] class LabelSet: ids: frozenset[int] name: str def __init__(self, ids: Set[int], name: Optional[str] = None): self.ids = frozenset(ids) self.name = str(min(ids)) if name is None else name def __eq__(self, __value): if isinstance(__value, LabelSet): return self.ids == __value.ids return False def __hash__(self) -> int: return hash(self.ids)
[docs] class Template: ids: Tuple[int, ...] smiles: str | None name: str | None def __init__(self, ids: Tuple[int, ...] | Set[int] | List[int], name: Optional[str] = None, smiles: Optional[str] = None): self.ids = tuple(sorted(set(ids))) self._identity_instance = tuple((x, x) for x in self.ids) self.smiles = smiles self.name = name def __eq__(self, __value): if isinstance(__value, Template): return self.ids == __value.ids return False def __hash__(self) -> int: return hash(self.ids)
[docs] def identity_instance(self): return self._identity_instance
[docs] def unlabeled_instance(self): return tuple((UNTRACKED_ORIGIN, x) for x in self.ids)
def __str__(self): return str(self.ids) + (f', name={self.name}' if self.name else '') + \ (f', smiles={self.smiles}' if self.smiles else '')
[docs] class ReactionRule: educts: List[Dict[int, int]] products: List[Tuple[int, ...]] name: str | None color: str | None func: Callable[[Tuple[TemplateInstance, ...]], List[TemplateInstance]] educts_product_indices: List[Dict[int, Tuple[int, int]]] def __init__(self, educts: List[Dict[int, int]], products: List[Tuple[int, ...]], name: Optional[str] = None, color: Optional[str] = None): self.educts = educts self.products = products self.name = name self.color = color # Validation slot_indices = [v for e in educts for v in e.values()] slot_indices_set = set(slot_indices) if len(slot_indices_set) < len(slot_indices): raise ValueError('reaction slot indices "%s" must be unique' % slot_indices) if slot_indices_set != set(range(0, len(slot_indices))): raise ValueError('reaction slot indices "%s" must be in the range [0,n)' % slot_indices) available_product_slots = sum([len(p) for p in products]) if available_product_slots != len(slot_indices): raise ValueError('reaction defines %s slot indices for %s product slots which must be equal' % ( len(slot_indices), available_product_slots)) for product in products: if len(set(product)) < len(product): raise ValueError('product "%s" has duplicate ids' % product) cumulative = [0] for p in products: cumulative.append(cumulative[-1] + len(p)) def get_indices(v): product_index = bisect_right(cumulative, v) - 1 product_atom_index = v - cumulative[product_index] return product_index, product_atom_index self.educts_product_indices = [{k: get_indices(v) for k, v in educt.items()} for educt in educts] # As we assume all educt template instances to be sorted by position ID, we can precompute # educt+atom indices for each product position ID. The function executing a reaction rule # then only needs to generate product template instances in a single loop. products_metadata: List[List[Tuple[int, int, int]]] = [[] for _ in products] for j, source_specie in enumerate(educts): for k, source_id in enumerate(sorted(source_specie.keys())): target_species_index, target_species_offset = self.educts_product_indices[j][source_id] target_id = products[target_species_index][target_species_offset] products_metadata[target_species_index].append((j, k, target_id)) products_metadata = [sorted(p, key=lambda x: x[2]) for p in products_metadata] self.func = lambda educts: [ tuple((educts[m[0]][m[1]][0], m[2]) for m in p) for p in products_metadata ] def __eq__(self, other): if not isinstance(other, ReactionRule): return False self_educt_templates = sorted([tuple(sorted(e.keys())) for e in self.educts]) other_educt_templates = sorted([tuple(sorted(e.keys())) for e in other.educts]) if self_educt_templates != other_educt_templates: return False self_product_templates = sorted([tuple(sorted(p)) for p in self.products]) other_product_templates = sorted([tuple(sorted(p)) for p in other.products]) if self_product_templates != other_product_templates: return False def sorted_products_with_slots(products): """ Generate a list of product templates with slot indices. The first element of each tuple is the original slot index, the second element the new slot index, and the third element the template position ID. """ products_with_slot = [] counter = 0 for p in products: p_with_slots = sorted([(counter + i, p[i]) for i in range(len(p))], key=lambda x: x[1]) products_with_slot.append(tuple(p_with_slots)) counter += len(p) products_with_slot = sorted(products_with_slot, key=lambda x: list(map(itemgetter(1), x))) products_with_slots = [] counter = 0 for p in products_with_slot: products_with_slots.append(tuple((p[i][0], counter + i, p[i][1]) for i in range(len(p)))) counter += len(p) return products_with_slots def mapped_sorted_educts(educts, products_with_slots): product_slot_mapping = {} for p in products_with_slots: for x in p: product_slot_mapping[x[0]] = x[1] mapped_educts = [{k: product_slot_mapping[v] for k, v in e.items()} for e in educts] return sorted(mapped_educts, key=lambda x: (min(x.keys()), min(x.values()))) self_mapped_educts = mapped_sorted_educts(self.educts, sorted_products_with_slots(self.products)) other_mapped_educts = mapped_sorted_educts(other.educts, sorted_products_with_slots(other.products)) return self_mapped_educts == other_mapped_educts def __str__(self): return str(self.educts) + ' --> ' + str(self.products)
[docs] def reverse(self, name: Optional[str] = None): inverted_educts = [{k: -1 for k in p} for p in self.products] inverted_products = [] product_index_offset = 0 for educt in self.educts: keys = sorted(educt.keys()) inverted_products.append(tuple(keys)) for i, key in enumerate(keys): target_index = educt[key] new_target_index = product_index_offset + i target_index_counter = target_index target_product_index = 0 while target_index_counter >= len(self.products[target_product_index]): target_index_counter -= len(self.products[target_product_index]) target_product_index += 1 target_position_id = self.products[target_product_index][target_index_counter] inverted_educts[target_product_index][target_position_id] = new_target_index product_index_offset += len(keys) if name is None and self.name is not None: name = self.name + ' (reverse)' return ReactionRule(inverted_educts, inverted_products, name=name, color=self.color)
[docs] class SaturateTracer: def __init__(self, reaction_rules: List[ReactionRule], injected_template_instances: Optional[Set[TemplateInstance]] = None, custom_initial_configuration: Optional[Set[TemplateInstance]] = None, label_sets: Optional[Set[LabelSet]] = None, templates: Optional[Set[Template]] = None, remove_duplicate_rules: Optional[bool] = True, log_duplicate_rules: Optional[bool] = True, verbose: Optional[bool] = False): if label_sets is None: all_position_ids = set() for reaction_rule in reaction_rules: for t in reaction_rule.educts: all_position_ids.update(set(t.keys())) for t in reaction_rule.products: all_position_ids.update(set(t)) self.label_sets = {LabelSet({k}) for k in all_position_ids} else: self.label_sets = label_sets self.reaction_rules = reaction_rules self.verbose = verbose # Check for duplicate reaction rules. Each duplicate index is recorded once, against the # first rule it duplicates: collecting index *pairs* instead would list a rule occurring # three or more times several times over, and deleting per pair then removes the same # position twice and runs off the end of the list. duplicate_of: Dict[int, int] = {} for i in range(len(reaction_rules)): if i in duplicate_of: continue for j in range(i + 1, len(reaction_rules)): if j not in duplicate_of and reaction_rules[i] == reaction_rules[j]: duplicate_of[j] = i if remove_duplicate_rules: for j in sorted(duplicate_of, reverse=True): if log_duplicate_rules: print('[INFO] Removed duplicate reaction rule at index %s (%s): %s' % ( j, duplicate_of[j], reaction_rules[j])) del reaction_rules[j] elif log_duplicate_rules: for j in sorted(duplicate_of): print('[INFO] Found duplicate reaction rule at index %s and %s: %s' % ( duplicate_of[j], j, reaction_rules[j])) self.label_display_map = {min(s.ids): s.name for s in self.label_sets} self.templates = set() if templates is None else templates self.templates_lookup: Dict[Tuple[int, ...], Template] = {t.ids: t for t in self.templates} self.custom_initial_configuration = custom_initial_configuration self.injected_template_instances = injected_template_instances self._template_node_indices: Dict[Tuple[int, ...], Set[int]] = defaultdict(set) self.node_index_map: Dict[TemplateInstance, int] = {} self.nodes: List[TemplateInstance] = [] self.hyperedges: Set[Tuple[Tuple[int, ...], Tuple[int, ...], int]] = set() self.initial_nodes: List[TemplateInstance] = [] # Validate that every position ID belongs to exactly one template. Tracing identifies a # template by its set of position IDs, so two overlapping templates would be conflated and # silently corrupt the provenance. Templates are collected both from the given template set # and from the educts and products of the reaction rules, since a rule may reference a # template that was never declared explicitly. template_keys = {t.ids for t in self.templates} for reaction_rule in self.reaction_rules: for educt in reaction_rule.educts: template_keys.add(SaturateTracer._get_template_key(list(educt.keys()))) for product in reaction_rule.products: template_keys.add(SaturateTracer._get_template_key(product)) position_id_templates: Dict[int, List[Tuple[int, ...]]] = defaultdict(list) for template_key in template_keys: for position_id in template_key: position_id_templates[position_id].append(template_key) overlapping = {p: keys for p, keys in position_id_templates.items() if len(keys) > 1} if len(overlapping) > 0: details = '; '.join( 'Position ID %s is shared by templates %s' % ( position_id, ' and '.join(str(k) for k in sorted(keys))) for position_id, keys in sorted(overlapping.items()) ) raise ValueError('Position IDs must be unique to a single template: %s' % details) # Validate that every position ID belongs to at most one label set. A label set is # represented throughout the saturation by the minimum position ID of its set, so a position # ID shared by two sets would be mapped to whichever representative is assigned last, # silently merging the two labelings instead of keeping them apart. position_id_label_sets: Dict[int, List[LabelSet]] = defaultdict(list) for label_set in self.label_sets: for position_id in label_set.ids: position_id_label_sets[position_id].append(label_set) overlapping_label_sets = {p: sets for p, sets in position_id_label_sets.items() if len(sets) > 1} if len(overlapping_label_sets) > 0: details = '; '.join( 'Position ID %s is shared by label sets %s' % ( position_id, ' and '.join( '%s (%s)' % (tuple(sorted(s.ids)), s.name) for s in sorted(sets, key=lambda s: sorted(s.ids)))) for position_id, sets in sorted(overlapping_label_sets.items()) ) raise ValueError('Position IDs must be unique to a single label set: %s' % details)
[docs] def run(self, limit_hyperedges: Optional[int] = None): label_sets_mapping = {} for label_set in self.label_sets: if len(label_set.ids) == 0: continue # Use the lowest ID for all elements of the label set set_id = min(label_set.ids) for k in label_set.ids: label_sets_mapping[k] = set_id # Prepare all templates with default labeling unvisited_template_instances = set() if self.custom_initial_configuration is None: for reaction_rule in self.reaction_rules: for t in reaction_rule.educts: unvisited_template_instances.add(self._normalize_specie(tuple( (label_sets_mapping[k] if k in label_sets_mapping else UNTRACKED_ORIGIN, k) for k in t.keys() ))) for t in reaction_rule.products: unvisited_template_instances.add(self._normalize_specie(tuple( (label_sets_mapping[k] if k in label_sets_mapping else UNTRACKED_ORIGIN, k) for k in t ))) else: unvisited_template_instances.update({self._normalize_specie(t) for t in self.custom_initial_configuration}) if self.injected_template_instances is not None: for s in self.injected_template_instances: unvisited_template_instances.add(self._normalize_specie(s)) reaction_rules_species_keys = [ [SaturateTracer._get_template_key(list(educt.keys())) for educt in r.educts] for r in self.reaction_rules ] # Initialize nodes with initial configuration template instances for instance in unvisited_template_instances: self.node_index_map[instance] = len(self.nodes) self.nodes.append(instance) self.initial_nodes.append(instance) def recurse_unvisited_combination(path, requested_templates, unvisited_lookup: Dict[Tuple[int, ...], Set[TemplateInstance]], depth: int): """ Prepares all combinations of unvisited template instances for the requested keys. Where an educt template is currently not present in the unvisited template instances None will be added. """ if depth == len(requested_templates): if any(x is not None for x in path): yield path else: path[depth] = None yield from recurse_unvisited_combination(path, requested_templates, unvisited_lookup, depth + 1) key = requested_templates[depth] template_instances = unvisited_lookup.get(key) if template_instances is not None: for template_instance in template_instances: path[depth] = template_instance yield from recurse_unvisited_combination(path, requested_templates, unvisited_lookup, depth + 1) def recurse_finalize_combination(path, requested_templates, depth: int): """ Finalizes a template instance combination for the requested keys by generating all combinations of already known template instances for None elements in the combination. """ while depth < len(requested_templates) and path[depth] is not None: depth += 1 if depth == len(requested_templates): yield tuple(path) else: key = requested_templates[depth] template_instance_indices = self._template_node_indices[key] for template_instance_index in template_instance_indices: path[depth] = self.nodes[template_instance_index] yield from recurse_finalize_combination(path, requested_templates, depth + 1) # Reset back to None for other paths path[depth] = None _itemgetter1 = itemgetter(1) # Saturation-based enumeration loop of all possible template instances from reaction rules while len(unvisited_template_instances) > 0: if limit_hyperedges is not None and len(self.hyperedges) >= limit_hyperedges: break if self.verbose: print( f"[INFO] Unvisited instances: {len(unvisited_template_instances)}; Hypergraph nodes: {len(self.nodes)}, edges: {len(self.hyperedges)}") unvisited_lookup: Dict[Tuple[int, ...], Set[TemplateInstance]] = defaultdict(set) new_unvisited_template_instances = set() for instance in unvisited_template_instances: unvisited_lookup[tuple(map(_itemgetter1, instance))].add(instance) for i, t_sources in enumerate(reaction_rules_species_keys): if not any(x in unvisited_lookup for x in t_sources): continue rule_func = self.reaction_rules[i].func for pre_combination in recurse_unvisited_combination([None] * len(t_sources), t_sources, unvisited_lookup, 0): for combination in recurse_finalize_combination(pre_combination, t_sources, 0): # Generate products from educts product_template_instances = rule_func(combination) # Update graph nodes and unvisited instances for instance in product_template_instances: if instance not in self.node_index_map: self.node_index_map[instance] = len(self.nodes) self.nodes.append(instance) new_unvisited_template_instances.add(instance) # Update graph hyperedges self.hyperedges.add(( tuple(self.node_index_map[educt] for educt in combination), tuple(self.node_index_map[product] for product in product_template_instances), i )) # Add last set of unvisited instances to the saturated set of all instances for instance in unvisited_template_instances: key = tuple(map(_itemgetter1, instance)) self._template_node_indices[key].add(self.node_index_map[instance]) # Swap unvisited instances for next iteration unvisited_template_instances = new_unvisited_template_instances
@staticmethod def _normalize_specie(s) -> Tuple: return tuple(sorted(s, key=itemgetter(1))) @staticmethod def _get_template_key(species: List[int] | Tuple[int, ...] | Set[int]) -> Tuple[int, ...]: """ Get the sorted position ID tuple of the template. [1, 3, 2] --> (1, 2, 3) """ return tuple(sorted(species)) @staticmethod def _get_template_instance_key(species: List[Tuple[int, int]] | TemplateInstance | Set[Tuple[int, int]]) -> Tuple[ int, ...]: """ Get the template position ID tuple from a template instance [(1, 4), (3, 5)] --> (4, 5) """ return tuple(map(itemgetter(1), sorted(species, key=itemgetter(1))))
[docs] def get_all_origin_id_positions(self, origin_id: int) -> Set[int]: """ Find all position IDs where the origin ID has been traced to. """ return {atom[1] for node in self.nodes for atom in node if atom[0] == origin_id}
[docs] def get_all_origin_species_targets(self, species_origin_ids: Set[int]) -> Set[TemplateInstance]: """ Find all targets the origin specie(s) has been traced to. """ result = set() for node in self.nodes: if any(atom[0] in species_origin_ids for atom in node): result.add(tuple(atom if atom[0] in species_origin_ids else (UNTRACKED_ORIGIN, atom[1]) for atom in node)) return result
[docs] def get_all_position_id_labelings(self, position_id: int) -> Set[int]: """ Find all origin IDs of a specific position ID. """ result = set() for key in self._template_node_indices.keys(): if position_id in key: for node_index in self._template_node_indices[key]: for atom in self.nodes[node_index]: if atom[1] == position_id: result.add(atom[0]) break return result
[docs] def get_all_species_labelings(self, position_ids: Set[int]) -> Set[TemplateInstance]: """ Find all template instances of a specific template. """ s = self._template_node_indices.get(SaturateTracer._get_template_key(position_ids)) return {self.nodes[i] for i in s} if s is not None else set()
[docs] def get_all_origin_target_labelings(self, origin_position_ids: Set[int], target_position_ids: Set[int]) -> Set[ TemplateInstance]: """ Find all template instances of a specific target template and from a specific origin template. """ key = SaturateTracer._get_template_key(target_position_ids) if key in self._template_node_indices: return { tuple(atom if atom[0] in origin_position_ids else (UNTRACKED_ORIGIN, atom[1]) for atom in self.nodes[i]) for i in self._template_node_indices[key] if any(atom[0] in origin_position_ids for atom in self.nodes[i]) } return set()
[docs] def labeling_to_display_str(self, labeling: Tuple, unlabeled_text: Optional[str] = None) -> str: def format_cell(x) -> str: if x[0] == UNTRACKED_ORIGIN: return f"""({unlabeled_text or UNTRACKED_ORIGIN_LABEL}, {x[1]})""" return f"""({self.label_display_map[x[0]] if x[0] in self.label_display_map else x[0]}, {x[1]})""" return '(' + ', '.join([format_cell(x) for x in labeling]) + ')'