Source code for molfoundry.atomtracing.converters

"""Bridge derivation graphs into saturation-based atom tracing.

The saturation tracer (:mod:`molfoundry.atomtracing.saturate_tracer`) works on an
abstract input: a set of :class:`Template` compounds, each a bag of unique integer
atom *position IDs*, and a set of :class:`ReactionRule` atom-atom maps between those
positions. This module builds both directly from a :class:`~molfoundry.DG` (or a subset
of its hyperedges), using the atom-atom maps that :class:`~molfoundry.DGVertexMapper`
recovers from each derivation.

The central idea of the translation:

* Every DG vertex (a compound graph) becomes one :class:`Template`. Each traced atom
  of that compound is assigned a globally unique position ID. Two hyperedges that
  share a compound share its template and thus its position IDs, which is what lets
  the saturation stitch reactions together.
* Every hyperedge becomes one :class:`ReactionRule` per atom-atom map. A map is a
  bijection between the educt and product atoms; that bijection is exactly the
  ``educts`` / ``products`` slot wiring a :class:`ReactionRule` needs.

Only *conserved* atoms can be traced: an atom the rule creates has no provenance and
one it destroys has no future, and the tracer's rules are bijective by construction.
A hyperedge that creates or destroys a *traced* atom (one selected by ``vertex_filter``)
therefore cannot be represented and is skipped -- see ``on_unbalanced``. Restricting
the filter to an element that the network conserves (carbon is the usual choice) makes
this a non-issue: reactions balance carbon, so no carbon is ever created or destroyed.

Typical use::

    from molfoundry.atomtracing import DGAtomTracing, by_element

    inputs = DGAtomTracing.from_dg(dg, vertex_filter=by_element("C"))
    sat = inputs.build_tracer()
    sat.run()
"""

from collections import defaultdict
from dataclasses import dataclass
from typing import Dict, Iterable, List, Optional, Set, Tuple

import molfoundry.dg as _dg
from molfoundry.dgvertexmapper import DGVertexMapper, VertexMap
from molfoundry.graph import Graph, Vertex, VertexFilter
from .saturate_tracer import LabelSet, ReactionRule, SaturateTracer, Template


[docs] def by_element(*symbols: str) -> VertexFilter: """A :data:`VertexFilter` keeping only atoms of the given element symbol(s). ``by_element("C")`` traces carbon; ``by_element("C", "N")`` traces carbon and nitrogen. Vertices without a valid element (e.g. electrons in an EFG graph) never match. """ wanted = frozenset(symbols) def predicate(v: Vertex) -> bool: atom_id = v.atomId if not atom_id.valid: return False return atom_id.symbol in wanted return predicate
[docs] def heavy_atoms(v: Vertex) -> bool: """A :data:`VertexFilter` keeping every atom except hydrogen. Useful when a graph carries explicit hydrogens (as SMILES-derived graphs do) and you want to trace the heavy-atom skeleton without hydrogens exploding the maps. """ atom_id = v.atomId if not atom_id.valid: return False return atom_id.symbol != "H"
def _passes(vertex_filter: Optional[VertexFilter], v: Vertex) -> bool: return vertex_filter is None or vertex_filter(v)
[docs] @dataclass(frozen=True) class TracerInputs: """Everything needed to trace a network, produced by :class:`DGAtomTracing`. ``templates`` and ``rules`` are the two arguments a :class:`SaturateTracer` takes; the remaining fields expose the bookkeeping so callers can build label sets and initial configurations that target specific compounds or atoms, and can map traced positions back to the DG. """ #: One :class:`Template` per compound that carries at least one traced atom. templates: Set[Template] #: One :class:`ReactionRule` per hyperedge per atom-atom map that could be traced. rules: List[ReactionRule] #: DG-vertex graph ``id`` -> its :class:`Template`. template_by_graph_id: Dict[int, Template] #: ``(graph id, atom vertex id)`` -> the atom's global position ID. position_id_by_atom: Dict[Tuple[int, int], int] #: ``(hyperedge id, reason)`` for every hyperedge that produced no rule. skipped_edges: List[Tuple[int, str]]
[docs] def build_tracer(self, **kwargs) -> SaturateTracer: """Construct a :class:`SaturateTracer` from these templates and rules. Extra keyword arguments (``label_sets``, ``custom_initial_configuration``, ``verbose``, ...) are forwarded to :class:`SaturateTracer` unchanged. The rule list is copied first, because :class:`SaturateTracer` removes duplicate rules from the list it is given, and we do not want to mutate ``self.rules``. """ return SaturateTracer(list(self.rules), templates=self.templates, **kwargs)
[docs] def template_of(self, vertex: "Graph | _dg.DG.Vertex | Vertex") -> Optional[Template]: """The template of a compound, given its DG vertex, its graph, or one of its atom vertices. Returns ``None`` when that compound has no traced atoms.""" if isinstance(vertex, _dg.DG.Vertex): return self.template_by_graph_id.get(vertex.graph.id) if isinstance(vertex, Vertex): return self.template_by_graph_id.get(vertex.graph.id) return self.template_by_graph_id.get(vertex.id)
[docs] def position_id(self, atom: Vertex) -> Optional[int]: """The global position ID of an atom vertex, or ``None`` when it is not traced.""" return self.position_id_by_atom.get((atom.graph.id, atom.id))
[docs] def label_set_of(self, template: Template, name: Optional[str] = None) -> LabelSet: """A :class:`LabelSet` covering every position of ``template`` -- i.e. the compound uniformly labeled, all of its traced atoms one indistinguishable origin.""" return LabelSet(set(template.ids), name=name if name is not None else template.name)
[docs] class DGAtomTracing: """Static helpers translating a DG into atom-tracing inputs. :meth:`from_dg` converts a whole derivation graph; :meth:`from_hyperedges` converts an explicit subset of its hyperedges. Both return :class:`TracerInputs`. """
[docs] @staticmethod def from_dg(dg: "_dg.DG", *, vertex_filter: Optional[VertexFilter] = None, maps_per_edge: int = 65535, on_unbalanced: str = "skip", verbose: bool = False, ) -> TracerInputs: """Build :class:`TracerInputs` from every hyperedge of ``dg``. :param dg: the DG to translate; each hyperedge becomes one or more reaction rules and contributes its educt/product compounds as templates. :param vertex_filter: which atoms to trace. ``None`` traces every atom (hydrogens included); pass e.g. :func:`by_element` to restrict to an element. Only atoms this selects get position IDs, appear in templates, and constrain conservation. :param maps_per_edge: how many atom-atom maps to enumerate per hyperedge. One map is the canonical labeling; a symmetric molecule admits several maps that give genuinely different atom traces (e.g. the two ends of succinate), so raise this to capture that symmetry. Each map becomes a rule sharing the hyperedge's name; the tracer removes any that turn out identical. The orbit of maps can be large, hence the cap. ``vertex_filter`` is handed to the mapper, so the maps enumerated are only those distinct on the traced atoms: symmetry among untraced atoms (hydrogen relabelings, say) is pruned before a map is ever produced, not deduplicated afterward. :param on_unbalanced: what to do when a hyperedge creates or destroys a traced atom and so cannot be represented. ``"skip"`` (default) records it in ``skipped_edges`` and moves on; ``"error"`` raises :class:`ValueError`. :param verbose: print a line for every skipped hyperedge. """ return DGAtomTracing.from_hyperedges(dg.edges, vertex_filter=vertex_filter, maps_per_edge=maps_per_edge, on_unbalanced=on_unbalanced, verbose=verbose)
[docs] @staticmethod def from_hyperedges( edges: "Iterable[_dg.DG.HyperEdge]", *, vertex_filter: Optional[VertexFilter] = None, maps_per_edge: int = 65535, on_unbalanced: str = "skip", verbose: bool = False, ) -> TracerInputs: """Build :class:`TracerInputs` from a collection of hyperedges. :param edges: the hyperedges to translate; each becomes one or more reaction rules and contributes its educt/product compounds as templates. :param vertex_filter: which atoms to trace. ``None`` traces every atom (hydrogens included); pass e.g. :func:`by_element` to restrict to an element. Only atoms this selects get position IDs, appear in templates, and constrain conservation. :param maps_per_edge: how many atom-atom maps to enumerate per hyperedge. One map is the canonical labeling; a symmetric molecule admits several maps that give genuinely different atom traces (e.g. the two ends of succinate), so raise this to capture that symmetry. Each map becomes a rule sharing the hyperedge's name; the tracer removes any that turn out identical. The orbit of maps can be large, hence the cap. ``vertex_filter`` is handed to the mapper, so the maps enumerated are only those distinct on the traced atoms: symmetry among untraced atoms (hydrogen relabelings, say) is pruned before a map is ever produced, not deduplicated afterward. :param on_unbalanced: what to do when a hyperedge creates or destroys a traced atom and so cannot be represented. ``"skip"`` (default) records it in ``skipped_edges`` and moves on; ``"error"`` raises :class:`ValueError`. :param verbose: print a line for every skipped hyperedge. """ if maps_per_edge < 1: raise ValueError("maps_per_edge must be at least 1, got %s" % maps_per_edge) if on_unbalanced not in ("skip", "error"): raise ValueError('on_unbalanced must be "skip" or "error", got %r' % on_unbalanced) edges_list = list(edges) # 1. Collect every compound the edges touch, deduplicated by graph id, and give # each traced atom a globally unique position ID. Compounds are visited in a # deterministic order (by graph id, then atom vertex id) so the assignment is # reproducible across runs. involved: Dict[int, Graph] = {} for e in edges_list: for v in list(e.sources): g = v.graph involved.setdefault(g.id, g) for v in list(e.targets): g = v.graph involved.setdefault(g.id, g) position_id_by_atom: Dict[Tuple[int, int], int] = {} template_by_graph_id: Dict[int, Template] = {} templates: Set[Template] = set() next_position_id = 1 for graph_id in sorted(involved.keys()): g = involved[graph_id] atom_position_ids: Dict[int, int] = {} for v in g.vertices: if _passes(vertex_filter, v): position_id_by_atom[(graph_id, v.id)] = next_position_id atom_position_ids[v.id] = next_position_id next_position_id += 1 # A compound with no traced atoms (e.g. water when tracing carbon) is not a # template and never appears in a rule. if len(atom_position_ids) == 0: continue smiles = DGAtomTracing._template_smiles(g, atom_position_ids) template = Template(tuple(atom_position_ids.values()), name=g.name, smiles=smiles) templates.add(template) template_by_graph_id[graph_id] = template # 2. Turn each hyperedge's atom-atom map(s) into reaction rules. rules: List[ReactionRule] = [] skipped: List[Tuple[int, str]] = [] def record_skip(edge_id: int, reason: str) -> None: skipped.append((edge_id, reason)) if verbose: print("[INFO] atom tracing: skipped hyperedge %s: %s" % (edge_id, reason)) for e in edges_list: edge_id = e.id name = ';'.join(r.name for r in e.rules) if len(e.rules) > 0 else "e%s" % edge_id produced_any_map = False for vmap in DGVertexMapper(e, limit=maps_per_edge, vertex_filter=vertex_filter): produced_any_map = True rule, reason = DGAtomTracing._rule_from_map(vmap, position_id_by_atom, vertex_filter, name) if rule is not None: rules.append(rule) continue if reason == _NO_TRACED_ATOMS: # No atoms of interest take part in this reaction. This is a property # of the reaction, not of the particular map, so stop early: the # other maps would say the same. Not a skip worth recording. break # A traced atom is created or destroyed: the reaction is unrepresentable. # This too is map-independent, so record once and move to the next edge. if on_unbalanced == "error": raise ValueError("Hyperedge %s: %s" % (edge_id, reason)) record_skip(edge_id, reason) break if not produced_any_map: # No recorded atom-atom map (e.g. a derivation added without a rule). record_skip(edge_id, "no atom-atom map available") return TracerInputs( templates=templates, rules=rules, template_by_graph_id=template_by_graph_id, position_id_by_atom=position_id_by_atom, skipped_edges=skipped, )
@staticmethod def _rule_from_map( vmap: VertexMap, position_id_by_atom: Dict[Tuple[int, int], int], vertex_filter: Optional[VertexFilter], name: str, ) -> Tuple[Optional[ReactionRule], Optional[str]]: """Turn a single atom-atom map into a :class:`ReactionRule`. Returns ``(rule, None)`` on success, or ``(None, reason)`` when the map cannot be traced -- either because no traced atom takes part, or because a traced atom is created / destroyed / transmuted and so the educt-product bijection is incomplete. """ # A traced atom is identified by (graph id, species-copy instance, atom vertex id). # The position ID is copy-independent, but the bijection is not: two copies of one # compound have the same positions yet different fates. bijection: Dict[Tuple[int, int, int], Tuple[int, int, int]] = {} educt_atoms: Dict[Tuple[int, int], Set[int]] = defaultdict(set) product_atoms: Dict[Tuple[int, int], Set[int]] = defaultdict(set) created = 0 destroyed = 0 transmuted = 0 for entry in vmap.entries(): source = entry.source target = entry.target source_traced = source is not None and _passes(vertex_filter, source) target_traced = target is not None and _passes(vertex_filter, target) if source is not None and target is not None: if source_traced and target_traced: source_graph_id = source.graph.id target_graph_id = target.graph.id educt_key = (source_graph_id, entry.sourceInstance, source.id) product_key = (target_graph_id, entry.targetInstance, target.id) bijection[educt_key] = product_key educt_atoms[(source_graph_id, entry.sourceInstance)].add(source.id) product_atoms[(target_graph_id, entry.targetInstance)].add(target.id) elif source_traced or target_traced: # A conserved atom traced on one side but not the other would change # element across the reaction, which no atom-atom map should assert. transmuted += 1 elif target is not None: if target_traced: created += 1 elif source is not None: if source_traced: destroyed += 1 if created > 0 or destroyed > 0 or transmuted > 0: reason = "unbalanced traced atoms (created=%s, destroyed=%s, element-changed=%s)" % ( created, destroyed, transmuted, ) return None, reason if len(bijection) == 0: return None, _NO_TRACED_ATOMS # Products: one position tuple per product species copy. Their flattened order # defines the slot indices the educts refer to. products: List[Tuple[int, ...]] = [] slot_of_product_atom: Dict[Tuple[int, int, int], int] = {} slot = 0 for graph_id, instance in sorted(product_atoms.keys()): positions: List[int] = [] for vertex_id in sorted(product_atoms[(graph_id, instance)]): positions.append(position_id_by_atom[(graph_id, vertex_id)]) slot_of_product_atom[(graph_id, instance, vertex_id)] = slot slot += 1 products.append(tuple(positions)) # Educts: one {position ID: product slot} map per educt species copy, wiring # each educt atom to the product atom the bijection sends it to. educts: List[Dict[int, int]] = [] for graph_id, instance in sorted(educt_atoms.keys()): educt_map: Dict[int, int] = {} for vertex_id in educt_atoms[(graph_id, instance)]: position = position_id_by_atom[(graph_id, vertex_id)] educt_map[position] = slot_of_product_atom[bijection[(graph_id, instance, vertex_id)]] educts.append(educt_map) return ReactionRule(educts, products, name=name), None @staticmethod def _template_smiles(g: Graph, atom_position_ids: Dict[int, int]) -> str: """A SMILES for ``g`` whose atom-map numbers are the template position IDs. The graph is cloned before its external IDs are rewritten, so the DG's own graph is left untouched. Non-traced atoms have their external ID cleared so they emit no stray atom map. """ clone = g.clone() for v in clone.vertices: if v.id in atom_position_ids: v.setExternalId(atom_position_ids[v.id]) elif v.externalId > 0: v.removeExternalId() return clone.smilesWithExternalIds
# Sentinel reason meaning "this reaction has no traced atoms", handled specially by the # caller (stop enumerating this edge's maps, but do not record it as a skip). _NO_TRACED_ATOMS = "no traced atoms"