Source code for molfoundry.dgvertexmapper

from typing import Iterator, Optional

import molfoundry.dg as dg
from .core import TSObject, coerce_type
from .graph import VertexFilter
from .vertexmap import VertexMap

_class_ref = TSObject.get_class_ref("DGVertexMapper")




class DGVertexMapper(TSObject):
    """Enumerates the atom-atom maps of a hyperedge.

    A derivation records one map per match, but relabeling a map by an
    automorphism of the educts or of the products gives another valid map, so the
    full set is an orbit whose size is the product of two automorphism group
    orders — easily thousands for ordinary molecules. Maps are therefore produced
    one at a time as they are iterated, and never more than ``limit`` of them::

        for m in DGVertexMapper(edge, limit=10):
            for entry in m:
                print(entry)

    No map is yielded twice, however it was reached. The enumeration is
    single-pass: a mapper that has run out stays exhausted, so build a new one to
    start over.

    When only some atoms matter, pass ``vertex_filter`` — a predicate on the
    :class:`~molfoundry.Vertex` atoms of the hyperedge's compounds. Only maps that
    differ on the accepted atoms are then enumerated, and every automorphism that
    fixes all of them is skipped, so the orbit shrinks to the accepted atoms' own
    symmetry. Tracing carbon through a hydrogen-rich network this way avoids the
    hydrogen permutations entirely — an ethanol whose 12-map orbit is all hydrogen
    relabelings collapses to a single map::

        for m in DGVertexMapper(edge, limit=1000,
                                vertex_filter=lambda v: v.atomId.symbol == 'C'):
            ...

    The maps produced are still complete (every atom is present); the atoms the
    filter rejects are simply frozen at one representative labeling rather than
    permuted. ``vertex_filter`` is evaluated at most once per atom.
    """

    def __init__(self, edge: "dg.DG.HyperEdge", limit: int = 1, vertex_filter: Optional[VertexFilter] = None) -> None:
        keep = None
        if vertex_filter is not None:
            # The TS side addresses an atom by (graphIndex, nodeIndex) in the DG.
            # Resolve that to the actual Vertex and defer to the user's predicate.
            # graphIndex is the atom's species index in dg.vertices and nodeIndex
            # its canonical vertex id, matching what VertexMap reports. The species
            # graphs are cached so only the vertex lookup repeats.
            vertices = edge.dg.vertices
            graph_cache = {}

            def keep(graph_index: int, node_index: int) -> bool:
                # The engine hands these across as JavaScript numbers, which arrive
                # as Python floats; coerce before indexing.
                graph_index = int(graph_index)
                node_index = int(node_index)
                graph = graph_cache.get(graph_index)
                if graph is None:
                    graph = vertices[graph_index].graph
                    graph_cache[graph_index] = graph
                return bool(vertex_filter(graph.vertices[node_index]))

        super().__init__(_class_ref, edge._ref, limit, keep)

    @property
    def limit(self) -> int:
        return int(self._ref.getLimit())

    @property
    def hyperEdge(self) -> "dg.DG.HyperEdge":
        return dg.DG.HyperEdge.wrap(self._ref.getHyperEdge())

    def __iter__(self) -> Iterator[VertexMap]:
        while True:
            nxt = coerce_type(self._ref.next())
            if nxt is None:
                return
            yield VertexMap.wrap(nxt)

    def __str__(self) -> str:
        return self._ref.toString()

    def __repr__(self) -> str:
        return self._ref.toString()