Merge pull request #16 from lamalab-org/dev

fix: stereo descriptor atom binding

Author: r-fedorov

src/bluenamer/assembly_parts.py

@@ -111,6 +111,7 @@ class AssemblyParts:
     parent_atom_symbols_by_locant: dict[str, str] = field(default_factory=dict)
     parent_atom_charges_by_locant: dict[str, int] = field(default_factory=dict)
     parent_bond_orders_by_locants: dict[tuple[str, str], int] = field(default_factory=dict)
+    parent_bond_ids_by_locants: dict[tuple[str, str], int] = field(default_factory=dict)
     name_atom_bindings: list[NameAtomBinding] = field(default_factory=list)
     name_token_spans: list[dict] = field(default_factory=list)
     name_rewrite_history: list[dict] = field(default_factory=list)

src/bluenamer/component_modifiers.py

@@ -10,6 +10,7 @@
 from .nomenclature import RULES
 from .perception import PerceivedGroup
 from .subgraph_tools import subgraph_component
+from .substituent_tokens import graph_bound_substituent_tokens
 from .trace_helpers import add_substituent_trace, bond_ids_within, decision_trace_data
 
 
@@ -156,8 +157,18 @@ def add_component_n_substituents(
                     mol, single_n, n_sub, sub_exclude, branch_namer, branch_decisions
                 )
                 if branch_name:
-                    branch_atoms = subgraph_component(mol, n_sub, sub_exclude | {single_n})
+                    branch_exclude = sub_exclude | {single_n}
+                    branch_atoms = subgraph_component(mol, n_sub, branch_exclude)
                     nested_decisions = decision_trace_data(branch_decisions)
+                    emitted_tokens = graph_bound_substituent_tokens(
+                        mol,
+                        n_sub,
+                        branch_atoms,
+                        branch_name,
+                        single_n,
+                        branch_exclude,
+                        branch_namer,
+                    )
                     if _use_hydrazone_suffix_modifier(parts, principal_key):
                         parts.principal_suffix_modifiers.append(
                             SubstituentItem(
@@ -166,6 +177,7 @@ def add_component_n_substituents(
                                 atom_ids=branch_atoms,
                                 bond_ids=bond_ids_within(mol, branch_atoms | {single_n}),
                                 charge_atom_ids=_charged_atoms(mol, branch_atoms),
+                                emitted_tokens=emitted_tokens,
                                 trace_segments=branch_trace,
                                 nested_decisions=nested_decisions,
                                 substituent_tree=branch_tree,
@@ -181,6 +193,7 @@ def add_component_n_substituents(
                         _charged_atoms(mol, branch_atoms),
                         branch_trace,
                         nested_decisions,
+                        emitted_tokens,
                         substituent_tree=branch_tree,
                     )
             n_idx_global += 1

src/bluenamer/name_assembly.py

@@ -9,6 +9,7 @@
 from .assembly_parts import AssemblyParts, NameAtomBinding, NameTokenBinding
 from .molecule import Molecule
 from .name_bindings import ensure_name_atom_binding_tokens, postprocess_name_atom_bindings
+from .stereo_descriptors import is_searchable_stereo_token
 
 
 @dataclass(frozen=True)
@@ -922,14 +923,17 @@ def _native_token_spans(
         if binding.stage == "hydro" and binding.role == "indicated_hydrogen":
             spans.extend(_indicated_hydrogen_native_token_spans(search_text, binding_idx, binding))
             continue
+        if binding.stage == "assembly" and binding.role in {"absolute_stereo", "bond_stereo"}:
+            spans.extend(_absolute_stereo_native_token_spans(search_text, binding_idx, binding))
+            continue
         if not any(_is_locant_search_token(token_binding) for token_binding in binding.emitted_tokens):
             ordered_spans = _ordered_native_token_spans(text, search_text, binding_idx, binding)
             if ordered_spans:
                 spans.extend(ordered_spans)
                 continue
         for token_binding in binding.emitted_tokens:
             token = token_binding.text.strip().lower()
-            if not _native_token_is_searchable(token):
+            if not _native_token_is_searchable(token, token_binding):
                 continue
             pos = 0
             for found in _native_token_occurrences(text, search_text, token_binding, pos):
@@ -961,6 +965,30 @@ def _indicated_hydrogen_native_token_spans(
     return spans
 
 
+def _absolute_stereo_native_token_spans(
+    search_text: str,
+    binding_idx: int,
+    binding: NameAtomBinding,
+) -> list[tuple[int, int, int, NameTokenBinding]]:
+    locant_token = next((token for token in binding.emitted_tokens if token.token_kind == "locant"), None)
+    stereo_token = next((token for token in binding.emitted_tokens if token.token_kind == "stereo"), None)
+    if locant_token is None or stereo_token is None:
+        return []
+    locant_text = locant_token.text.lower()
+    stereo_text = stereo_token.text.lower()
+    pattern = f"{locant_text}{stereo_text}"
+    spans: list[tuple[int, int, int, NameTokenBinding]] = []
+    pos = 0
+    while True:
+        found = search_text.find(pattern, pos)
+        if found < 0:
+            break
+        spans.append((found, found + len(locant_text), binding_idx, locant_token))
+        spans.append((found + len(locant_text), found + len(pattern), binding_idx, stereo_token))
+        pos = found + 1
+    return spans
+
+
 def _ordered_native_token_spans(
     text: str,
     search_text: str,
@@ -973,7 +1001,7 @@ def _ordered_native_token_spans(
     cursor = 0
     for token_binding in binding.emitted_tokens:
         token = token_binding.text.strip().lower()
-        if not _native_token_is_searchable(token):
+        if not _native_token_is_searchable(token, token_binding):
             continue
         found = next(iter(_native_token_occurrences(text, search_text, token_binding, cursor)), -1)
         if found < 0:
@@ -1020,8 +1048,12 @@ def _is_standalone_locant_span(text: str, start: int, end: int) -> bool:
     return (not before or before in "-,( ") and (not after or after in "-,)' ")
 
 
-def _native_token_is_searchable(token: str) -> bool:
-    return bool(token) and (len(token) >= 2 or token.isdigit() or "," in token or token in _ELEMENT_LOCANT_TOKENS)
+def _native_token_is_searchable(token: str, token_binding: NameTokenBinding | None = None) -> bool:
+    if not token:
+        return False
+    if token_binding is not None and token_binding.token_kind == "stereo":
+        return is_searchable_stereo_token(token)
+    return len(token) >= 2 or token.isdigit() or "," in token or token in _ELEMENT_LOCANT_TOKENS
 
 
 def _direct_binding_spans(text: str, bindings: tuple[NameAtomBinding, ...]) -> list[tuple[int, int, int]]:
@@ -1684,6 +1716,13 @@ def _token_span_from_native_binding_group(
     end: int,
     matches: list[tuple[int, NameTokenBinding]],
 ) -> NameTokenSpan:
+    stereo_matches = [
+        (binding_idx, token_binding)
+        for binding_idx, token_binding in matches
+        if token_binding.source == "renderer_stereo"
+    ]
+    if stereo_matches and _should_prioritize_renderer_stereo(text, stereo_matches):
+        matches = stereo_matches
     atoms: set[int] = set()
     bonds: set[int] = set()
     charges: set[int] = set()
@@ -1727,6 +1766,17 @@ def _token_span_from_native_binding_group(
     )
 
 
+def _should_prioritize_renderer_stereo(_text: str, matches: list[tuple[int, NameTokenBinding]]) -> bool:
+    """Return whether a span is owned by an explicit stereo-renderer token."""
+
+    for _binding_idx, token_binding in matches:
+        if token_binding.token_kind == "stereo" and is_searchable_stereo_token(token_binding.text):
+            return True
+        if token_binding.token_kind == "locant" and token_binding.grammar_role in {"absolute_stereo", "bond_stereo"}:
+            return True
+    return False
+
+
 def _collapse_metadata(values: list[str], *, default: str) -> str:
     values = [value for value in values if value]
     if not values:

src/bluenamer/name_bindings.py

@@ -7,11 +7,17 @@
 from .nomenclature import RULES
 from .principal_suffixes import render_principal_suffix
 from .rules import bonds, stems
+from .stereo_descriptors import ABSOLUTE_STEREO_DESCRIPTORS, BOND_STEREO_DESCRIPTORS
 
 
 def refresh_name_atom_bindings(parts: AssemblyParts) -> list[NameAtomBinding]:
     """Populate structured bindings for the current assembly parts."""
 
+    preserved_assembly_stereo = [
+        binding
+        for binding in parts.name_atom_bindings
+        if binding.stage == "assembly" and binding.role in {"relative_stereo", "small_ring_stereo"}
+    ]
     bindings: list[NameAtomBinding] = []
     if parts.parent_atom_ids:
         bindings.append(
@@ -37,6 +43,40 @@ def refresh_name_atom_bindings(parts: AssemblyParts) -> list[NameAtomBinding]:
                 emitted_tokens=_hydro_operation_tokens(operation),
             )
         )
+    for locant, descriptor in parts.stereo_features:
+        if descriptor not in ABSOLUTE_STEREO_DESCRIPTORS or not locant:
+            continue
+        atom_idx = parts.parent_atom_ids_by_locant.get(str(locant))
+        if atom_idx is None:
+            continue
+        bindings.append(
+            NameAtomBinding(
+                stage="assembly",
+                role="absolute_stereo",
+                term=f"{locant}{descriptor}",
+                atom_ids={atom_idx},
+                locants=(str(locant),),
+                emitted_tokens=_absolute_stereo_tokens(str(locant), descriptor, atom_idx),
+            )
+        )
+    for locant, descriptor in parts.stereo_features:
+        if descriptor not in BOND_STEREO_DESCRIPTORS or not locant:
+            continue
+        _bond_locants, atom_ids, bond_ids = _bond_stereo_graph_scope(parts, str(locant))
+        if not atom_ids:
+            continue
+        bindings.append(
+            NameAtomBinding(
+                stage="assembly",
+                role="bond_stereo",
+                term=descriptor,
+                atom_ids=atom_ids,
+                bond_ids=bond_ids,
+                locants=(str(locant),),
+                emitted_tokens=_bond_stereo_tokens(str(locant), descriptor, atom_ids, bond_ids),
+            )
+        )
+    bindings.extend(preserved_assembly_stereo)
     if parts.front_modifiers:
         bindings.append(
             NameAtomBinding(
@@ -328,6 +368,8 @@ def _operation_emitted_tokens(binding: NameAtomBinding) -> tuple[NameTokenBindin
         return _locanted_operation_tokens(binding, charge_from_binding=True)
     if binding.stage == "charge" or binding.role == "parent_charge":
         return _charge_operation_tokens(binding)
+    if binding.stage == "assembly" and binding.role == "relative_stereo":
+        return _relative_stereo_tokens(binding)
     if binding.stage == "assembly" and "stereo" in binding.role:
         return _locanted_operation_tokens(binding)
     if binding.stage == "modifier":
@@ -430,6 +472,114 @@ def _hydro_operation_tokens(operation) -> tuple[NameTokenBinding, ...]:
     return tuple(tokens)
 
 
+def _absolute_stereo_tokens(locant: str, descriptor: str, atom_idx: int) -> tuple[NameTokenBinding, ...]:
+    """Return native tokens for an absolute stereochemical descriptor."""
+
+    return (
+        NameTokenBinding(
+            text=locant,
+            token_kind="locant",
+            ownership="exact",
+            confidence="exact",
+            source="renderer_stereo",
+            grammar_role="absolute_stereo",
+            binding_key="assembly:absolute_stereo",
+            atom_ids={atom_idx},
+            locants=(locant,),
+        ),
+        NameTokenBinding(
+            text=descriptor,
+            token_kind="stereo",
+            ownership="exact",
+            confidence="exact",
+            source="renderer_stereo",
+            grammar_role="absolute_stereo",
+            binding_key="assembly:absolute_stereo",
+            atom_ids={atom_idx},
+            locants=(locant,),
+        ),
+    )
+
+
+def _bond_stereo_graph_scope(parts: AssemblyParts, locant: str) -> tuple[tuple[str, ...], set[int], set[int]]:
+    """Return the graph scope for an E/Z descriptor rendered at a parent locant."""
+
+    for locant_pair, order in parts.parent_bond_orders_by_locants.items():
+        if locant not in locant_pair or order != 2:
+            continue
+        atom_ids = {
+            atom_idx
+            for pair_locant in locant_pair
+            if (atom_idx := parts.parent_atom_ids_by_locant.get(str(pair_locant))) is not None
+        }
+        if len(atom_ids) != 2:
+            continue
+        bond_id = parts.parent_bond_ids_by_locants.get(locant_pair)
+        return (
+            tuple(str(pair_locant) for pair_locant in locant_pair),
+            atom_ids,
+            {bond_id} if bond_id is not None else set(),
+        )
+    atom_idx = parts.parent_atom_ids_by_locant.get(locant)
+    return (locant,), {atom_idx} if atom_idx is not None else set(), set()
+
+
+def _bond_stereo_tokens(
+    locant: str,
+    descriptor: str,
+    atom_ids: set[int],
+    bond_ids: set[int],
+) -> tuple[NameTokenBinding, ...]:
+    """Return native tokens for an E/Z bond stereochemical descriptor."""
+
+    return (
+        NameTokenBinding(
+            text=locant,
+            token_kind="locant",
+            ownership="exact",
+            confidence="exact",
+            source="renderer_stereo",
+            grammar_role="bond_stereo",
+            binding_key="assembly:bond_stereo",
+            atom_ids=atom_ids,
+            bond_ids=bond_ids,
+            locants=(locant,),
+        ),
+        NameTokenBinding(
+            text=descriptor,
+            token_kind="stereo",
+            ownership="exact",
+            confidence="exact",
+            source="renderer_stereo",
+            grammar_role="bond_stereo",
+            binding_key="assembly:bond_stereo",
+            atom_ids=atom_ids,
+            bond_ids=bond_ids,
+            locants=(locant,),
+        ),
+    )
+
+
+def _relative_stereo_tokens(binding: NameAtomBinding) -> tuple[NameTokenBinding, ...]:
+    """Return native token metadata for relative cis/trans descriptors."""
+
+    return (
+        NameTokenBinding(
+            text=binding.term,
+            token_kind="stereo",
+            ownership="exact",
+            confidence="exact",
+            source="renderer_stereo",
+            grammar_role=binding.role,
+            binding_key=f"{binding.stage}:{binding.role}",
+            atom_ids=set(binding.atom_ids),
+            bond_ids=set(binding.bond_ids),
+            charge_atom_ids=set(binding.charge_atom_ids),
+            locants=tuple(binding.locants),
+        ),
+    )
+
+
 def _locanted_emitted_tokens(
     emitted_tokens: tuple[NameTokenBinding, ...],
     locants: tuple[str, ...],

src/bluenamer/parent_pipeline.py

@@ -144,7 +144,9 @@ def build_parent_parts(
                 neighbor_locant = str(get_loc(neighbor_idx))
                 bond = mol.get_bond(atom_idx, neighbor_idx)
                 if bond is not None:
-                    parts.parent_bond_orders_by_locants[tuple(sorted((locant, neighbor_locant)))] = bond.order
+                    locant_pair = tuple(sorted((locant, neighbor_locant)))
+                    parts.parent_bond_orders_by_locants[locant_pair] = bond.order
+                    parts.parent_bond_ids_by_locants[locant_pair] = bond.idx
     return parts
 
 

src/bluenamer/stereo_descriptors.py

@@ -0,0 +1,20 @@
+"""Shared stereochemical descriptor tokens used by renderers and metadata."""
+
+ABSOLUTE_STEREO_DESCRIPTORS = frozenset({"R", "S"})
+BOND_STEREO_DESCRIPTORS = frozenset({"E", "Z"})
+RELATIVE_STEREO_DESCRIPTORS = frozenset({"cis", "trans"})
+
+SEARCHABLE_STEREO_TOKENS = frozenset(
+    descriptor.lower()
+    for descriptor in (
+        *ABSOLUTE_STEREO_DESCRIPTORS,
+        *BOND_STEREO_DESCRIPTORS,
+        *RELATIVE_STEREO_DESCRIPTORS,
+    )
+)
+
+
+def is_searchable_stereo_token(text: str) -> bool:
+    """Return whether a renderer-emitted stereo token may be matched directly."""
+
+    return text.lower() in SEARCHABLE_STEREO_TOKENS